Improving the `*.tensor.trace` op to carry shapes/encodings. (#15228)
This avoids a bunch of IR noise when lowering into the stream dialect
and forces all traced tensors to be transfered into staging resources.
This means we no longer require mapping (or worse, emulated mapping) on
any target in order to trace as we'll have the compiler implement the
readback asynchronously.
A change I held off on here but should be done soon is making the trace
ops return their tensors so that we can asynchronously schedule them.
Then instead of emitting runtime calls on staging resources we could
batch up the transfers into a ringbuffer and trace out without
introducing dispatch-to-dispatch host synchronization.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
index 0fc10cf..26dd4e9 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
@@ -191,6 +191,51 @@
}
//===----------------------------------------------------------------------===//
+// custom<ShapedOperandList>($values, type($values), $value_dims)
+//===----------------------------------------------------------------------===//
+// %value : type{%dynamic_dims}, ...
+
+static ParseResult parseShapedOperandList(
+ OpAsmParser &parser,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
+ SmallVectorImpl<Type> &valueTypes,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &valueDims) {
+ do {
+ values.emplace_back();
+ valueTypes.emplace_back();
+ if (failed(parser.parseOperand(values.back())) ||
+ failed(parser.parseColon()) ||
+ failed(parser.parseType(valueTypes.back())))
+ return failure();
+ if (int64_t dynamicDimCount =
+ cast<ShapedType>(valueTypes.back()).getNumDynamicDims()) {
+ if (failed(parser.parseOperandList(valueDims, dynamicDimCount,
+ AsmParser::Delimiter::Braces)))
+ return failure();
+ }
+ } while (succeeded(parser.parseOptionalComma()));
+ return success();
+}
+
+static void printShapedOperandList(OpAsmPrinter &p, Operation *op,
+ ValueRange values, TypeRange valueTypes,
+ ValueRange valueDims) {
+ llvm::interleaveComma(llvm::zip_equal(values, valueTypes), p, [&](auto it) {
+ auto [value, valueType] = it;
+ p << value;
+ p << " : ";
+ p << valueType;
+ if (int64_t dynamicDimCount =
+ cast<ShapedType>(valueType).getNumDynamicDims()) {
+ p << "{";
+ llvm::interleaveComma(valueDims.take_front(dynamicDimCount), p);
+ valueDims = valueDims.drop_front(dynamicDimCount);
+ p << "}";
+ }
+ });
+}
+
+//===----------------------------------------------------------------------===//
// custom<WorkgroupCountRegion>($body)
//===----------------------------------------------------------------------===//
@@ -1649,6 +1694,35 @@
}
//===----------------------------------------------------------------------===//
+// flow.tensor.trace
+//===----------------------------------------------------------------------===//
+
+LogicalResult TensorTraceOp::verify() {
+ TensorTraceOp op = *this;
+ if (failed(verifyOpDynamicDims(op, op.getValues(), op.getValueDims()))) {
+ return failure();
+ }
+ return success();
+}
+
+ValueRange TensorTraceOp::getOperandDynamicDims(unsigned idx) {
+ auto valueDims = getValueDims();
+ for (unsigned i = 0; i <= idx; ++i) {
+ auto valueType = cast<ShapedType>(getValues()[i].getType());
+ int64_t dynamicDimCount = valueType.getNumDynamicDims();
+ if (i == idx) {
+ return valueDims.take_front(dynamicDimCount);
+ }
+ valueDims = valueDims.drop_front(dynamicDimCount);
+ }
+ return ValueRange{};
+}
+
+ValueRange TensorTraceOp::getResultDynamicDims(unsigned idx) {
+ return ValueRange{};
+}
+
+//===----------------------------------------------------------------------===//
// Public methods
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td
index c15da61..e3d976c 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td
@@ -1509,21 +1509,37 @@
let hasFolder = 1;
}
-def FLOW_TensorTraceOp : FLOW_Op<"tensor.trace", []> {
- let summary = [{trace value(s) operation}];
+def FLOW_TensorTraceOp : FLOW_Op<"tensor.trace", [
+ AttrSizedOperandSegments,
+ DeclareOpInterfaceMethods<Util_ShapeAwareOp>,
+]> {
+ let summary = [{traces one or more tensor values at runtime}];
let description = [{
Traces out to a runtime trace sink (console, log file, etc) the given
- tensors and titles them with the given key. The key is informational only
- and useful for titling/marking specific sets of tensors for easier
- searching.
+ tensors. The key is arbitrary and can be used for identifying the set of
+ values being traced.
}];
let arguments = (ins
StrAttr:$key,
- Variadic<FLOW_Tensor>:$operands
+ Variadic<FLOW_Tensor>:$values,
+ FLOW_ShapeDynamicDims:$value_dims
);
- let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+ let assemblyFormat = [{
+ $key `=` `[`
+ custom<ShapedOperandList>($values, type($values), $value_dims)
+ `]` attr-dict-with-keyword
+ }];
+
+ let builders = [
+ OpBuilder<(ins "StringRef":$key, "ValueRange":$values), [{
+ build($_builder, $_state, key, values,
+ IREE::Util::buildDynamicDimsForValues($_state.location, values, $_builder));
+ }]>,
+ ];
+
+ let hasVerifier = 1;
}
//===---------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_ops.mlir b/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_ops.mlir
index 636b535..6cc909d 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_ops.mlir
@@ -180,3 +180,19 @@
%0 = flow.tensor.update %arg0, %arg1[%arg2, %arg3] : tensor<?x?xf32>{%c1, %c2} -> %arg1 as tensor<?x4xf32>{%c3}
return %0 : tensor<?x4xf32>
}
+
+// -----
+
+// CHECK-LABEL: @tensorTrace
+// CHECK-SAME: (%[[TENSOR0:.+]]: tensor<5xf32>, %[[TENSOR1:.+]]: tensor<?x3x?xi32>, %[[TENSOR1_DIM0:.+]]: index, %[[TENSOR1_DIM2:.+]]: index)
+func.func @tensorTrace(%tensor0: tensor<5xf32>, %tensor1: tensor<?x3x?xi32>, %tensor1_dim0: index, %tensor1_dim2: index) {
+ // CHECK: flow.tensor.trace "FOOBAR" = [
+ // CHECK-SAME: %[[TENSOR0]] : tensor<5xf32>,
+ // CHECK-SAME: %[[TENSOR1]] : tensor<?x3x?xi32>{%[[TENSOR1_DIM0]], %[[TENSOR1_DIM2]]}
+ // CHECK-SAME: ]
+ flow.tensor.trace "FOOBAR" = [
+ %tensor0 : tensor<5xf32>,
+ %tensor1 : tensor<?x3x?xi32>{%tensor1_dim0, %tensor1_dim2}
+ ]
+ return
+}
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/InjectDispatchTracing.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/InjectDispatchTracing.cpp
index 22c9c79..c70620b 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/InjectDispatchTracing.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/InjectDispatchTracing.cpp
@@ -46,14 +46,14 @@
// Input tensors:
OpBuilder builder(dispatchOp);
- builder.create<TensorTraceOp>(
+ builder.create<IREE::Flow::TensorTraceOp>(
dispatchOp.getLoc(),
builder.getStringAttr(entryPointName + " inputs"),
filterTensorValues(dispatchOp.getArguments()));
// Output tensors:
builder.setInsertionPointAfter(dispatchOp);
- builder.create<TensorTraceOp>(
+ builder.create<IREE::Flow::TensorTraceOp>(
dispatchOp.getLoc(),
builder.getStringAttr(entryPointName + " outputs"),
filterTensorValues(dispatchOp.getResults()));
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/InsertDispatchDebugTargets.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/InsertDispatchDebugTargets.cpp
index 8dbb63d..6186fd6 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/InsertDispatchDebugTargets.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/InsertDispatchDebugTargets.cpp
@@ -55,19 +55,20 @@
}
// Inserts flow.tensor.trace ops around the specified dispatch op.
-static void traceOpWithName(DispatchOp dispatchOp, std::string name) {
+static void traceOpWithName(IREE::Flow::DispatchOp dispatchOp,
+ std::string name) {
OpBuilder builder(dispatchOp);
// Input tensors:
- builder.create<TensorTraceOp>(
+ builder.create<IREE::Flow::TensorTraceOp>(
dispatchOp.getLoc(), builder.getStringAttr(name + " inputs"),
filterNonTensorValues(dispatchOp.getArguments()));
// Output tensors:
OpBuilder::InsertionGuard guard(builder);
builder.setInsertionPointAfter(dispatchOp);
- builder.create<TensorTraceOp>(dispatchOp.getLoc(),
- builder.getStringAttr(name + " outputs"),
- filterNonTensorValues(dispatchOp.getResults()));
+ builder.create<IREE::Flow::TensorTraceOp>(
+ dispatchOp.getLoc(), builder.getStringAttr(name + " outputs"),
+ filterNonTensorValues(dispatchOp.getResults()));
}
// Breaks the given function on the specified op by simply returning immediately
@@ -162,7 +163,8 @@
localTraceOrdinal = traceOrdinal;
auto &bodyRegion = op.getFunctionBody();
- auto dispatchOps = llvm::to_vector<8>(bodyRegion.getOps<DispatchOp>());
+ auto dispatchOps =
+ llvm::to_vector<8>(bodyRegion.getOps<IREE::Flow::DispatchOp>());
// Trace on a valid ordinal.
if (localTraceOrdinal >= 0 && localTraceOrdinal < dispatchOps.size()) {
@@ -222,8 +224,8 @@
// Find the target dispatch to break on and trace on all matching
// dispatches.
- DispatchOp breakTarget = DispatchOp();
- funcOp.walk([&](DispatchOp dispatchOp) {
+ IREE::Flow::DispatchOp breakTarget;
+ funcOp.walk([&](IREE::Flow::DispatchOp dispatchOp) {
std::string entryPointName =
dispatchOp.getEntryPoint().getRootReference().getValue().str();
for (FlatSymbolRefAttr nestedRef :
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/inject_dispatch_tracing.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/inject_dispatch_tracing.mlir
index 7471d55..c090fd6 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/inject_dispatch_tracing.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/inject_dispatch_tracing.mlir
@@ -4,10 +4,10 @@
// CHECK-SAME: (%[[ARG0:.+]]: tensor<4xf32>)
func.func @singleDispatch(%arg0: tensor<4xf32>) -> tensor<4xf32> {
%c4 = arith.constant 4 : index
- // CHECK: flow.tensor.trace {key = "ex::entry0 inputs"} %[[ARG0]] : tensor<4xf32>
+ // CHECK: flow.tensor.trace "ex::entry0 inputs" = [%[[ARG0]] : tensor<4xf32>]
// CHECK-NEXT: %[[RET0:.+]] = flow.dispatch @ex::@entry0[%c4](%[[ARG0]]) : (tensor<4xf32>) -> tensor<4xf32>
%0 = flow.dispatch @ex::@entry0[%c4](%arg0) : (tensor<4xf32>) -> tensor<4xf32>
- // CHECK-NEXT: flow.tensor.trace {key = "ex::entry0 outputs"} %[[RET0]] : tensor<4xf32>
+ // CHECK-NEXT: flow.tensor.trace "ex::entry0 outputs" = [%[[RET0]] : tensor<4xf32>]
// CHECK-NEXT: return %[[RET0]]
return %0 : tensor<4xf32>
}
@@ -19,15 +19,15 @@
func.func @multiDispatch(%arg0: tensor<4xf32>) -> tensor<4xf32> {
%c4 = arith.constant 4 : index
- // CHECK: flow.tensor.trace {key = "ex::entry0 inputs"} %[[ARG0]] : tensor<4xf32>
+ // CHECK: flow.tensor.trace "ex::entry0 inputs" = [%[[ARG0]] : tensor<4xf32>]
// CHECK-NEXT: %[[RET0:.+]] = flow.dispatch @ex::@entry0[%c4](%[[ARG0]]) : (tensor<4xf32>) -> tensor<4xf32>
%0 = flow.dispatch @ex::@entry0[%c4](%arg0) : (tensor<4xf32>) -> tensor<4xf32>
- // CHECK-NEXT: flow.tensor.trace {key = "ex::entry0 outputs"} %[[RET0]] : tensor<4xf32>
+ // CHECK-NEXT: flow.tensor.trace "ex::entry0 outputs" = [%[[RET0]] : tensor<4xf32>]
- // CHECK: flow.tensor.trace {key = "ex::entry1 inputs"} %[[RET0]] : tensor<4xf32>
+ // CHECK: flow.tensor.trace "ex::entry1 inputs" = [%[[RET0]] : tensor<4xf32>]
// CHECK-NEXT: %[[RET1:.+]] = flow.dispatch @ex::@entry1[%c4](%[[RET0]]) : (tensor<4xf32>) -> tensor<4xf32>
%1 = flow.dispatch @ex::@entry1[%c4](%0) : (tensor<4xf32>) -> tensor<4xf32>
- // CHECK-NEXT: flow.tensor.trace {key = "ex::entry1 outputs"} %[[RET1]] : tensor<4xf32>
+ // CHECK-NEXT: flow.tensor.trace "ex::entry1 outputs" = [%[[RET1]] : tensor<4xf32>]
// CHECK: return %[[RET1]]
return %1 : tensor<4xf32>
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/insert_dispatch_debug_markers.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/insert_dispatch_debug_markers.mlir
index f17537b..c15f268 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/insert_dispatch_debug_markers.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/insert_dispatch_debug_markers.mlir
@@ -2,17 +2,17 @@
// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(iree-flow-insert-debug-target-at-ordinal{break-debug-target=@target_func:0 trace-debug-target=@target_func:0})" %s | FileCheck %s --check-prefixes=ORDINAL_0
// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(iree-flow-insert-debug-target-at-symbol{break-debug-target=dispatch_1 trace-debug-target=dispatch_1[^0-9]})" %s | FileCheck %s --check-prefixes=CHECK,SYMBOL
-/// Multiple functions
+// Multiple functions.
// CHECK-LABEL: func.func @target_func
// ORDINAL_0-LABEL: func.func @target_func
func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
%c4 = arith.constant 4 : index
// CHECK: %[[D0:.+]] = flow.dispatch @dispatch_0::@dispatch_0_entry
- // ORDINAL_0: flow.tensor.trace {key = "dispatch_0::dispatch_0_entry::0 inputs"}
+ // ORDINAL_0: flow.tensor.trace "dispatch_0::dispatch_0_entry::0 inputs"
// ORDINAL_0-NEXT: %[[D0:.+]] = flow.dispatch @dispatch_0::@dispatch_0_entry
%0 = flow.dispatch @dispatch_0::@dispatch_0_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
- // ORDINAL_0-NEXT: flow.tensor.trace {key = "dispatch_0::dispatch_0_entry::0 outputs"}
+ // ORDINAL_0-NEXT: flow.tensor.trace "dispatch_0::dispatch_0_entry::0 outputs"
// CHECK: %[[D1:.+]] = flow.dispatch @dispatch_1::@dispatch_1_entry
%1 = flow.dispatch @dispatch_1::@dispatch_1_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
%2 = flow.dispatch @dispatch_2::@dispatch_2_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
@@ -22,23 +22,22 @@
return %3 : !hal.buffer_view
}
-
// CHECK-LABEL: func.func @other_func
func.func @other_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
%c4 = arith.constant 4 : index
- // CHECK: %[[D0:.+]] = flow.dispatch @dispatch_1::@dispatch_1_entry
- %0 = flow.dispatch @dispatch_1::@dispatch_1_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
+ // CHECK: %[[D3:.+]] = flow.dispatch @dispatch_3::@dispatch_3_entry
+ %0 = flow.dispatch @dispatch_3::@dispatch_3_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
- // CHECK: %[[D1:.+]] = flow.dispatch @dispatch_2::@dispatch_2_entry
- %1 = flow.dispatch @dispatch_2::@dispatch_2_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
- // CHECK: %[[D2:.+]] = flow.dispatch @dispatch_3::@dispatch_3_entry
- %2 = flow.dispatch @dispatch_3::@dispatch_3_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
+ // CHECK: %[[D4:.+]] = flow.dispatch @dispatch_4::@dispatch_4_entry
+ %1 = flow.dispatch @dispatch_4::@dispatch_4_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
+ // CHECK: %[[D5:.+]] = flow.dispatch @dispatch_5::@dispatch_5_entry
+ %2 = flow.dispatch @dispatch_5::@dispatch_5_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
- // ORDINAL: %[[ORIGINAL_EXPORT:.+]] = hal.tensor.export %[[D2]] : tensor<4xf32> -> !hal.buffer_view
- // SYMBOL: %[[BREAK_EXPORT:.+]] = hal.tensor.export %[[D0]] : tensor<4xf32> -> !hal.buffer_view
+ // ORDINAL: %[[ORIGINAL_EXPORT:.+]] = hal.tensor.export %[[D5]] : tensor<4xf32> -> !hal.buffer_view
+ // SYMBOL: %[[BREAK_EXPORT:.+]] = hal.tensor.export %[[D5]] : tensor<4xf32> -> !hal.buffer_view
%3 = hal.tensor.export %2 : tensor<4xf32> -> !hal.buffer_view
- /// Only break on the symbol as the ordinal specifies a different function
+ // Only break on the symbol as the ordinal specifies a different function.
// SYMBOL: return %[[BREAK_EXPORT]] : !hal.buffer_view
// ORDINAL: return %[[ORIGINAL_EXPORT]] : !hal.buffer_view
return %3 : !hal.buffer_view
@@ -46,7 +45,8 @@
// -----
-// Break on a dispatch with a different number of results
+// Break on a dispatch with a different number of results.
+
// CHECK-LABEL: func.func @target_func
func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
%c4 = arith.constant 4 : index
@@ -64,7 +64,8 @@
// -----
-// Break/trace on a dispatch not found in the target function should do nothing
+// Break/trace on a dispatch not found in the target function should do nothing.
+
// CHECK-LABEL: func.func @target_func
func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
%c4 = arith.constant 4 : index
@@ -78,7 +79,8 @@
// -----
-/// Combine tracing and breaking on the same dispatch
+// Combines tracing and breaking on the same dispatch.
+
// CHECK-LABEL: func.func @target_func
// CHECK-SAME: %[[ARG0:.+]]: tensor<4xf32>
func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
@@ -86,12 +88,12 @@
// CHECK: %[[D0:.+]] = flow.dispatch @dispatch_0::@dispatch_0_entry
%0 = flow.dispatch @dispatch_0::@dispatch_0_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
- // ORDINAL: flow.tensor.trace {key = "dispatch_1::dispatch_1_entry::1 inputs"} %[[ARG0]] : tensor<4xf32>
- // SYMBOL: flow.tensor.trace {key = "dispatch_1::dispatch_1_entry inputs"} %[[ARG0]] : tensor<4xf32>
+ // ORDINAL: flow.tensor.trace "dispatch_1::dispatch_1_entry::1 inputs" = [%[[ARG0]] : tensor<4xf32>]
+ // SYMBOL: flow.tensor.trace "dispatch_1::dispatch_1_entry inputs" = [%[[ARG0]] : tensor<4xf32>]
// CHECK: %[[D1:.+]] = flow.dispatch @dispatch_1::@dispatch_1_entry
%1 = flow.dispatch @dispatch_1::@dispatch_1_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
- // ORDINAL: flow.tensor.trace {key = "dispatch_1::dispatch_1_entry::1 outputs"} %[[D1]] : tensor<4xf32>
- // SYMBOL: flow.tensor.trace {key = "dispatch_1::dispatch_1_entry outputs"} %[[D1]] : tensor<4xf32>
+ // ORDINAL: flow.tensor.trace "dispatch_1::dispatch_1_entry::1 outputs" = [%[[D1]] : tensor<4xf32>]
+ // SYMBOL: flow.tensor.trace "dispatch_1::dispatch_1_entry outputs" = [%[[D1]] : tensor<4xf32>]
%2 = flow.dispatch @dispatch_2::@dispatch_2_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
%3 = hal.tensor.export %2 : tensor<4xf32> -> !hal.buffer_view
@@ -103,19 +105,21 @@
// -----
-/// Check regex matching on symbol
+// Checks regex matching on a dispatch symbol.
+
// CHECK-LABEL: func.func @target_func
func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
%c4 = arith.constant 4 : index
- // SYMBOL: flow.tensor.trace {key = "dispatch_1::dispatch_1_entry inputs"}
+
+ // SYMBOL: flow.tensor.trace "dispatch_1::dispatch_1_entry inputs"
// CHECK: flow.dispatch @dispatch_1::@dispatch_1_entry
%0 = flow.dispatch @dispatch_1::@dispatch_1_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
- // SYMBOL: flow.tensor.trace {key = "dispatch_1::dispatch_1_entry outputs"}
+ // SYMBOL: flow.tensor.trace "dispatch_1::dispatch_1_entry outputs"
- // SYMBOL-NOT: flow.tensor.trace {key = "dispatch_11::dispatch_11_entry inputs"}
+ // SYMBOL-NOT: flow.tensor.trace "dispatch_11::dispatch_11_entry inputs"
// CHECK: flow.dispatch @dispatch_11::@dispatch_11_entry
%1 = flow.dispatch @dispatch_11::@dispatch_11_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
- // SYMBOL-NOT: flow.tensor.trace {key = "dispatch_11::dispatch_11_entry outputs"}
+ // SYMBOL-NOT: flow.tensor.trace "dispatch_11::dispatch_11_entry outputs"
%2 = hal.tensor.export %1 : tensor<4xf32> -> !hal.buffer_view
return %2 : !hal.buffer_view
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
index 8d31a02..9c7ebe0 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
@@ -758,8 +758,22 @@
LogicalResult
matchAndRewrite(IREE::Stream::TensorTraceOp traceOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
+ SmallVector<Value> bufferViews;
+ auto resourceEncodingDims = adaptor.getResourceEncodingDims();
+ for (auto [resource, resourceSize, resourceEncoding] : llvm::zip_equal(
+ adaptor.getResources(), adaptor.getResourceSizes(),
+ adaptor.getResourceEncodings().getAsRange<TypeAttr>())) {
+ int64_t dynamicDimCount =
+ cast<ShapedType>(resourceEncoding.getValue()).getNumDynamicDims();
+ bufferViews.push_back(rewriter.create<IREE::Stream::TensorExportOp>(
+ traceOp.getLoc(), rewriter.getType<IREE::HAL::BufferViewType>(),
+ resource, resourceEncoding,
+ resourceEncodingDims.take_front(dynamicDimCount), resourceSize,
+ /*affinity=*/IREE::Stream::AffinityAttr{}));
+ resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
+ }
rewriter.replaceOpWithNewOp<IREE::HAL::BufferViewTraceOp>(
- traceOp, traceOp.getKeyAttr(), adaptor.getOperands());
+ traceOp, traceOp.getKeyAttr(), bufferViews);
return success();
}
};
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD.bazel
index 832da44..64493c0 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD.bazel
@@ -18,6 +18,7 @@
[
"channel_ops.mlir",
"cmd_ops.mlir",
+ "debug_ops.mlir",
"file_ops.mlir",
"resource_ops.mlir",
"timepoint_ops.mlir",
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt
index 5c25a89..fb7c40f 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt
@@ -16,6 +16,7 @@
SRCS
"channel_ops.mlir"
"cmd_ops.mlir"
+ "debug_ops.mlir"
"file_ops.mlir"
"resource_ops.mlir"
"timepoint_ops.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/debug_ops.mlir b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/debug_ops.mlir
new file mode 100644
index 0000000..a748326
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/debug_ops.mlir
@@ -0,0 +1,14 @@
+// RUN: iree-opt --split-input-file --iree-hal-conversion %s | FileCheck %s
+
+// CHECK-LABEL: @tensorTrace
+// CHECK-SAME: (%[[TENSOR0_BUFFER:.+]]: !hal.buffer, %[[TENSOR0_SIZE:.+]]: index, %[[TENSOR1_BUFFER:.+]]: !hal.buffer, %[[TENSOR1_SIZE:.+]]: index, %[[TENSOR1_DIM0:.+]]: index)
+func.func @tensorTrace(%tensor0: !stream.resource<staging>, %tensor0_size: index, %tensor1: !stream.resource<staging>, %tensor1_size: index, %tensor1_dim0: index) {
+ // CHECK-DAG: %[[TENSOR0:.+]] = hal.buffer_view.create buffer(%[[TENSOR0_BUFFER]] : !hal.buffer)[%c0{{.*}}, %[[TENSOR0_SIZE]]] shape([%c5, %c3])
+ // CHECK-DAG: %[[TENSOR1:.+]] = hal.buffer_view.create buffer(%[[TENSOR1_BUFFER]] : !hal.buffer)[%c0{{.*}}, %[[TENSOR1_SIZE]]] shape([%[[TENSOR1_DIM0]], %c5{{.*}}])
+ // CHECK: hal.buffer_view.trace "FOOBAR" = %[[TENSOR0]], %[[TENSOR1]] : !hal.buffer_view, !hal.buffer_view
+ stream.tensor.trace "FOOBAR" = [
+ %tensor0 : tensor<5x3xf32> in !stream.resource<staging>{%tensor0_size},
+ %tensor1 : tensor<?x5xf32>{%tensor1_dim0} in !stream.resource<staging>{%tensor1_size}
+ ]
+ return
+}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
index bc58f6f..e2f795e 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
@@ -880,6 +880,7 @@
);
let assemblyFormat = [{
+ $key `=`
$operands `:` type($operands)
attr-dict-with-keyword
}];
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp
index e7e2bed..e036755 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp
@@ -591,6 +591,9 @@
.Case([&](IREE::Stream::TensorExportOp op) {
removeAssumedBits(NOT_MUTATED | NOT_EXTERNAL);
})
+ .Case([&](IREE::Stream::TensorTraceOp op) {
+ removeAssumedBits(NOT_STAGING_READ);
+ })
.Case([&](IREE::Stream::AsyncCloneOp op) {
removeAssumedBits(NOT_TRANSFER_READ);
auto &resultUsage = solver.getElementFor<ValueResourceUsage>(
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
index 297eba6..4ae1571 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
@@ -244,30 +244,30 @@
LogicalResult
matchAndRewrite(IREE::Flow::TensorTraceOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- SmallVector<Value> exportedTensors;
+ SmallVector<Value> resources;
+ SmallVector<Value> resourceSizes;
+ SmallVector<Attribute> resourceEncodings;
for (auto [tensorOperand, resourceOperand] :
- llvm::zip_equal(op.getOperands(), adaptor.getOperands())) {
+ llvm::zip_equal(op.getValues(), adaptor.getValues())) {
auto source =
consumeTensorOperand(op.getLoc(), resourceOperand, rewriter);
- auto externalType = rewriter.getType<IREE::Stream::ResourceType>(
- IREE::Stream::Lifetime::External);
- auto exportSource = resourceOperand;
- if (source.resource.getType() != externalType) {
- exportSource = rewriter.create<IREE::Stream::AsyncTransferOp>(
- op.getLoc(), externalType, source.resource, source.resourceSize,
+ auto stagingType = rewriter.getType<IREE::Stream::ResourceType>(
+ IREE::Stream::Lifetime::Staging);
+ auto traceSource = source.resource;
+ if (source.resource.getType() != stagingType) {
+ traceSource = rewriter.create<IREE::Stream::AsyncTransferOp>(
+ op.getLoc(), stagingType, source.resource, source.resourceSize,
source.resourceSize,
/*source_affinity=*/getAffinityFor(op),
/*result_affinity=*/nullptr);
}
- auto dynamicDims = IREE::Util::buildDynamicDimsForValue(
- op.getLoc(), tensorOperand, rewriter);
- exportedTensors.push_back(rewriter.create<IREE::Stream::TensorExportOp>(
- op.getLoc(), tensorOperand.getType(), exportSource,
- TypeAttr::get(tensorOperand.getType()), dynamicDims,
- source.resourceSize, getAffinityFor(op)));
+ resources.push_back(traceSource);
+ resourceSizes.push_back(source.resourceSize);
+ resourceEncodings.push_back(TypeAttr::get(tensorOperand.getType()));
}
rewriter.replaceOpWithNewOp<IREE::Stream::TensorTraceOp>(
- op, adaptor.getKey(), exportedTensors);
+ op, adaptor.getKey(), resources, resourceSizes,
+ rewriter.getArrayAttr(resourceEncodings), adaptor.getValueDims());
return success();
}
};
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir
index d194376..4e2a972 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir
@@ -147,3 +147,21 @@
// CHECK: return %[[T2]]
return %0 : tensor<2x3xi32>
}
+
+// -----
+
+// CHECK-LABEL: @tensorTrace
+// CHECK-SAME: (%[[TENSOR0:.+]]: !stream.resource<*>, %[[TENSOR0_SIZE:.+]]: index, %[[TENSOR1:.+]]: !stream.resource<*>, %[[TENSOR1_SIZE:.+]]: index, %[[TENSOR1_DIM0:.+]]: index, %[[TENSOR1_DIM2:.+]]: index)
+func.func @tensorTrace(%tensor0: tensor<5xf32>, %tensor1: tensor<?x3x?xi32>, %tensor1_dim0: index, %tensor1_dim2: index) {
+ // CHECK-DAG: %[[TENSOR0_STAGED:.+]] = stream.async.transfer %[[TENSOR0]] : !stream.resource<*>{%[[TENSOR0_SIZE]]} -> !stream.resource<staging>{%[[TENSOR0_SIZE]]}
+ // CHECK-DAG: %[[TENSOR1_STAGED:.+]] = stream.async.transfer %[[TENSOR1]] : !stream.resource<*>{%[[TENSOR1_SIZE]]} -> !stream.resource<staging>{%[[TENSOR1_SIZE]]}
+ // CHECK: stream.tensor.trace "FOOBAR" = [
+ // CHECK-NEXT: %[[TENSOR0_STAGED]] : tensor<5xf32> in !stream.resource<staging>{%[[TENSOR0_SIZE]]},
+ // CHECK-NEXT: %[[TENSOR1_STAGED]] : tensor<?x3x?xi32>{%[[TENSOR1_DIM0]], %[[TENSOR1_DIM2]]} in !stream.resource<staging>{%[[TENSOR1_SIZE]]}
+ // CHECK-NEXT: ]
+ flow.tensor.trace "FOOBAR" = [
+ %tensor0 : tensor<5xf32>,
+ %tensor1 : tensor<?x3x?xi32>{%tensor1_dim0, %tensor1_dim2}
+ ]
+ return
+}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
index addada8..a294923 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
@@ -192,6 +192,20 @@
return success();
}
+static void eraseStreamRegionResults(Region ®ion,
+ ArrayRef<unsigned> excludedResultIndices) {
+ for (auto &block : region.getBlocks()) {
+ auto yieldOp = dyn_cast<IREE::Stream::YieldOp>(block.getTerminator());
+ if (!yieldOp)
+ continue;
+ llvm::SmallVector<Value> newOperands;
+ for (auto i : llvm::reverse(excludedResultIndices)) {
+ yieldOp.getResourceOperandsMutable().erase(i);
+ yieldOp.getResourceOperandSizesMutable().erase(i);
+ }
+ }
+}
+
// Computes the value access bits starting from |rootValue|.
// Traverses the IR graph along tied ops but does not handle branches.
static IREE::Util::ValueAccess computeValueAccess(Value rootValue) {
@@ -260,18 +274,82 @@
return access;
}
-static void eraseStreamRegionResults(Region ®ion,
- ArrayRef<unsigned> excludedResultIndices) {
- for (auto &block : region.getBlocks()) {
- auto yieldOp = dyn_cast<IREE::Stream::YieldOp>(block.getTerminator());
- if (!yieldOp)
- continue;
- llvm::SmallVector<Value> newOperands;
- for (auto i : llvm::reverse(excludedResultIndices)) {
- yieldOp.getResourceOperandsMutable().erase(i);
- yieldOp.getResourceOperandSizesMutable().erase(i);
+//===----------------------------------------------------------------------===//
+// custom<EncodedResourceOperands>(
+// $resources, type($resources), $resource_sizes,
+// $resource_encodings, $resource_encoding_dims)
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseEncodedResourceOperands(
+ OpAsmParser &parser,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resources,
+ SmallVectorImpl<Type> &resourceTypes,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceSizes,
+ ArrayAttr &resourceEncodings,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceEncodingDims) {
+ SmallVector<Attribute> resourceEncodingAttrs;
+ do {
+ resources.emplace_back();
+ TypeAttr resourceEncoding;
+ if (failed(parser.parseOperand(resources.back())) ||
+ failed(parser.parseColon()) ||
+ failed(parser.parseAttribute(resourceEncoding)))
+ return failure();
+ resourceEncodingAttrs.push_back(resourceEncoding);
+ if (int64_t dynamicDimCount =
+ cast<ShapedType>(resourceEncoding.getValue()).getNumDynamicDims()) {
+ if (failed(parser.parseOperandList(resourceEncodingDims, dynamicDimCount,
+ AsmParser::Delimiter::Braces)))
+ return failure();
}
- }
+ resourceTypes.emplace_back();
+ resourceSizes.emplace_back();
+ if (failed(parser.parseKeyword("in")) ||
+ failed(parseSizeAwareType(parser, resourceTypes.back(),
+ resourceSizes.back())))
+ return failure();
+ } while (succeeded(parser.parseOptionalComma()));
+ resourceEncodings = parser.getBuilder().getArrayAttr(resourceEncodingAttrs);
+ return success();
+}
+
+static void printEncodedResourceOperands(OpAsmPrinter &p, Operation *op,
+ ValueRange resources,
+ TypeRange resourceTypes,
+ ValueRange resourceSizes,
+ ArrayAttr resourceEncodings,
+ ValueRange resourceEncodingDims) {
+ p.increaseIndent();
+ p.printNewline();
+ llvm::interleave(
+ llvm::zip_equal(resources, resourceTypes, resourceSizes,
+ resourceEncodings.getAsValueRange<TypeAttr>()),
+ [&](auto it) {
+ auto [resource, resourceType, resourceSize, resourceEncoding] = it;
+ p << resource;
+ p << " : ";
+ p << resourceEncoding;
+ if (int64_t dynamicDimCount =
+ cast<ShapedType>(resourceEncoding).getNumDynamicDims()) {
+ p << "{";
+ llvm::interleaveComma(
+ resourceEncodingDims.take_front(dynamicDimCount), p);
+ resourceEncodingDims =
+ resourceEncodingDims.drop_front(dynamicDimCount);
+ p << "}";
+ }
+ p << " in ";
+ p << resourceType;
+ p << "{";
+ p << resourceSize;
+ p << "}";
+ },
+ [&]() {
+ p << ",";
+ p.printNewline();
+ });
+ p.decreaseIndent();
+ p.printNewline();
}
//===----------------------------------------------------------------------===//
@@ -1212,6 +1290,54 @@
}
//===----------------------------------------------------------------------===//
+// stream.tensor.trace
+//===----------------------------------------------------------------------===//
+
+LogicalResult TensorTraceOp::verify() {
+ TensorTraceOp op = *this;
+ if (op.getResources().size() != op.getResourceEncodings().size() ||
+ op.getResources().size() != op.getResourceSizes().size()) {
+ return op.emitOpError(
+ "each resource needs a matching resource encoding and size "
+ "(array length mismatch)");
+ }
+ auto resourceEncodingDims = op.getResourceEncodingDims();
+ for (auto [resource, resourceSize, resourceEncoding] :
+ llvm::zip_equal(op.getResources(), op.getResourceSizes(),
+ op.getResourceEncodings().getAsValueRange<TypeAttr>())) {
+ int64_t dynamicDimCount =
+ cast<ShapedType>(resourceEncoding).getNumDynamicDims();
+ if (failed(verifyOpDynamicDims(
+ op, resourceEncoding,
+ resourceEncodingDims.take_front(dynamicDimCount))) ||
+ failed(verifyOpValueSizes(op, resource, resourceSize))) {
+ return failure();
+ }
+ resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
+ }
+ return success();
+}
+
+ValueRange TensorTraceOp::getOperandDynamicDims(unsigned idx) {
+ auto resourceEncodings = getResourceEncodings().getValue();
+ auto resourceEncodingDims = getResourceEncodingDims();
+ for (unsigned i = 0; i <= idx; ++i) {
+ auto resourceEncoding = resourceEncodings[i].cast<TypeAttr>().getValue();
+ int64_t dynamicDimCount =
+ cast<ShapedType>(resourceEncoding).getNumDynamicDims();
+ if (i == idx) {
+ return resourceEncodingDims.take_front(dynamicDimCount);
+ }
+ resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
+ }
+ return ValueRange{};
+}
+
+ValueRange TensorTraceOp::getResultDynamicDims(unsigned idx) {
+ return ValueRange{};
+}
+
+//===----------------------------------------------------------------------===//
// stream.async.alloca
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
index 275c178..0d52718 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
@@ -1383,22 +1383,40 @@
let hasCanonicalizer = 1;
}
-def Stream_TensorTraceOp : Stream_Op<"tensor.trace", []> {
- let summary = [{trace value(s) operation}];
+def Stream_TensorTraceOp : Stream_Op<"tensor.trace", [
+ AttrSizedOperandSegments,
+ DeclareOpInterfaceMethods<Util_ShapeAwareOp>,
+ Util_SizeAwareOp,
+]> {
+ let summary = [{traces one or more tensor values at runtime}];
let description = [{
Traces out to a runtime trace sink (console, log file, etc) the given
- tensors and titles them with the given key. The key is informational only
- and useful for titling/marking specific sets of tensors for easier
- searching.
+ tensors. The key is arbitrary and can be used for identifying the set of
+ values being traced.
}];
let arguments = (ins
StrAttr:$key,
- Variadic<AnyTensor>:$operands
+ Variadic<Stream_StagingResource>:$resources,
+ Variadic<Stream_Size>:$resource_sizes,
+ TypeArrayAttr:$resource_encodings,
+ Stream_ShapeDynamicDims:$resource_encoding_dims
);
- // TODO(benvanik): stream.tensor.trace assembly format (custom).
- let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+ let assemblyFormat = [{
+ $key `=` `[`
+ custom<EncodedResourceOperands>(
+ $resources, type($resources), $resource_sizes,
+ $resource_encodings, $resource_encoding_dims)
+ `]` attr-dict-with-keyword
+ }];
+
+ let extraClassDeclaration = [{
+ Value getOperandSize(unsigned idx) { return getResourceSizes()[idx]; }
+ Value getResultSize(unsigned idx) { return {}; }
+ }];
+
+ let hasVerifier = 1;
}
} // OpGroupTensorOps
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_ops.mlir b/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_ops.mlir
index f650415..4cc586c 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_ops.mlir
@@ -135,3 +135,19 @@
%0 = stream.tensor.store %arg2, %arg0 : f32 -> tensor<f32> in %arg0 as !stream.resource<staging>{%arg1}
return %0 : !stream.resource<staging>
}
+
+// -----
+
+// CHECK-LABEL: @tensorTrace
+// CHECK-SAME: (%[[TENSOR0:.+]]: !stream.resource<staging>, %[[TENSOR0_SIZE:.+]]: index, %[[TENSOR1:.+]]: !stream.resource<staging>, %[[TENSOR1_SIZE:.+]]: index, %[[TENSOR1_DIM0:.+]]: index, %[[TENSOR1_DIM2:.+]]: index)
+func.func @tensorTrace(%tensor0: !stream.resource<staging>, %tensor0_size: index, %tensor1: !stream.resource<staging>, %tensor1_size: index, %tensor1_dim0: index, %tensor1_dim2: index) {
+ // CHECK: stream.tensor.trace "FOOBAR" = [
+ // CHECK-NEXT: %[[TENSOR0]] : tensor<5xf32> in !stream.resource<staging>{%[[TENSOR0_SIZE]]},
+ // CHECK-NEXT: %[[TENSOR1]] : tensor<?x3x?xi32>{%[[TENSOR1_DIM0]], %[[TENSOR1_DIM2]]} in !stream.resource<staging>{%[[TENSOR1_SIZE]]}
+ // CHECK-NEXT: ]
+ stream.tensor.trace "FOOBAR" = [
+ %tensor0 : tensor<5xf32> in !stream.resource<staging>{%tensor0_size},
+ %tensor1 : tensor<?x3x?xi32>{%tensor1_dim0, %tensor1_dim2} in !stream.resource<staging>{%tensor1_size}
+ ]
+ return
+}
diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.cpp b/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.cpp
index b13b245..b8cbc9c 100644
--- a/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.cpp
@@ -727,6 +727,15 @@
return dynamicDims;
}
+SmallVector<Value> buildDynamicDimsForValues(Location loc, ValueRange values,
+ OpBuilder &builder) {
+ SmallVector<Value> dynamicDims;
+ for (auto value : values) {
+ dynamicDims.append(buildDynamicDimsForValue(loc, value, builder));
+ }
+ return dynamicDims;
+}
+
static SmallVector<Value> buildShape(Location loc, ShapedType type,
ValueRange dynamicDims,
OpBuilder &builder) {
diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h b/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h
index 76bb1c1..f27ee11 100644
--- a/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h
+++ b/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h
@@ -170,6 +170,12 @@
SmallVector<Value> buildDynamicDimsForValue(Location loc, Value value,
OpBuilder &builder);
+// Returns dimension values for each dynamic dimension of the given |values|.
+// |values| must all have ShapedTypes. The returned value range will be empty if
+// all shapes are fully static.
+SmallVector<Value> buildDynamicDimsForValues(Location loc, ValueRange values,
+ OpBuilder &builder);
+
// Builds a ranked shape with all dimension values for the given operand.
SmallVector<Value> buildOperandShape(ShapeAwareOpInterface op,
unsigned operandIdx, OpBuilder &builder);
diff --git a/compiler/src/iree/compiler/InputConversion/Common/test/iree_import_public.mlir b/compiler/src/iree/compiler/InputConversion/Common/test/iree_import_public.mlir
index f9b3730..5aa7bbf 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/test/iree_import_public.mlir
+++ b/compiler/src/iree/compiler/InputConversion/Common/test/iree_import_public.mlir
@@ -62,7 +62,7 @@
// CHECK: %buffer = hal.buffer.subspan
// CHECK-SAME: <%arg0 : !hal.buffer>[%[[OFFSET]], %[[LENGTH]]] : !hal.buffer
-func.func @buffer_subspan(%arg0: !iree_input.buffer) -> !iree_input.buffer {
+func.func @buffer_subspan(%arg0: !iree_input.buffer) -> !iree_input.buffer {
%offset = arith.constant 100 : index
%length = arith.constant 200 : index
%buffer = iree_input.buffer.subspan<%arg0 : !iree_input.buffer>[%offset, %length] : !iree_input.buffer
@@ -275,9 +275,15 @@
// -----
// CHECK-LABEL: func.func @tensor_trace
-// CHECK: flow.tensor.trace {key = "FOOBAR"} %arg0, %arg1 : tensor<5xf32>, tensor<3xf32>
-func.func @tensor_trace(%arg0 : tensor<5xf32>, %arg1 : tensor<3xf32>) {
- iree_input.tensor.trace "FOOBAR" %arg0, %arg1 : tensor<5xf32>, tensor<3xf32>
+// CHECK: flow.tensor.trace "FOOBAR" = [
+// CHECK-SAME: %arg0 : tensor<5xf32>,
+// CHECK-SAME: %arg1 : tensor<?x3xf32>{%arg2}
+// CHECK-SAME: ]
+func.func @tensor_trace(%arg0: tensor<5xf32>, %arg1: tensor<?x3xf32>, %arg2: index) {
+ iree_input.tensor.trace "FOOBAR" = [
+ %arg0 : tensor<5xf32>,
+ %arg1 : tensor<?x3xf32>{%arg2}
+ ]
return
}
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/Patterns.cpp b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/Patterns.cpp
index fda46c4..f21e24c 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/Patterns.cpp
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/Patterns.cpp
@@ -374,8 +374,29 @@
LogicalResult
matchAndRewrite(IREE::Stream::TensorTraceOp traceOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
+ auto bufferType = rewriter.getType<IREE::HAL::BufferType>();
+ auto bufferViewType = rewriter.getType<IREE::HAL::BufferViewType>();
+ auto zero = rewriter.create<arith::ConstantIndexOp>(traceOp.getLoc(), 0);
+ auto resourceEncodingDims = adaptor.getResourceEncodingDims();
+ SmallVector<Value> bufferViews;
+ for (auto [resource, resourceSize, resourceEncoding] : llvm::zip_equal(
+ adaptor.getResources(), adaptor.getResourceSizes(),
+ adaptor.getResourceEncodings().getAsRange<TypeAttr>())) {
+ Value resourceBuffer = rewriter.create<IREE::HAL::Inline::BufferWrapOp>(
+ traceOp.getLoc(), bufferType, resource,
+ /*offset=*/
+ zero,
+ /*length=*/resourceSize);
+ int64_t dynamicDimCount =
+ cast<ShapedType>(resourceEncoding.getValue()).getNumDynamicDims();
+ bufferViews.push_back(rewriter.create<IREE::Stream::TensorExportOp>(
+ traceOp.getLoc(), bufferViewType, resourceBuffer, resourceEncoding,
+ resourceEncodingDims.take_front(dynamicDimCount), resourceSize,
+ /*affinity=*/IREE::Stream::AffinityAttr{}));
+ resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
+ }
rewriter.replaceOpWithNewOp<IREE::HAL::Inline::BufferViewTraceOp>(
- traceOp, traceOp.getKeyAttr(), adaptor.getOperands());
+ traceOp, traceOp.getKeyAttr(), bufferViews);
return success();
}
};
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/BUILD.bazel b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/BUILD.bazel
index db19be9..09466b8 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/BUILD.bazel
@@ -17,6 +17,7 @@
srcs = enforce_glob(
[
"cmd_ops.mlir",
+ "debug_ops.mlir",
"file_ops.mlir",
"resource_ops.mlir",
"timepoint_ops.mlir",
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/CMakeLists.txt b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/CMakeLists.txt
index 930427f..c3f6ba0 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/CMakeLists.txt
@@ -15,6 +15,7 @@
lit
SRCS
"cmd_ops.mlir"
+ "debug_ops.mlir"
"file_ops.mlir"
"resource_ops.mlir"
"timepoint_ops.mlir"
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/cmd_ops.mlir b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/cmd_ops.mlir
index 72549c1..cf6ff97 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/cmd_ops.mlir
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/cmd_ops.mlir
@@ -156,24 +156,3 @@
} => !stream.timepoint
return %timepoint : !stream.timepoint
}
-
-// -----
-
-// CHECK-LABEL: @trace_tensor
-// CHECK-SAME: %[[ARG0:.+]]: !hal.buffer
-func.func @trace_tensor(%arg0: !stream.resource<external>) -> () {
- // CHECK-DAG: %[[C180:.+]] = arith.constant 180 : index
- %c180 = arith.constant 180 : index
-
- // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
- // CHECK-DAG: %[[C45:.+]] = arith.constant 45 : index
- // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[C268435488:.+]] = arith.constant 268435488 : i32
- // CHECK-DAG: %[[C1_i32:.+]] = arith.constant 1 : i32
- // CHECK: %[[VIEW:.+]] = hal_inline.buffer_view.create buffer(%[[ARG0]] : !hal.buffer)[%[[C0]], %[[C180]]] shape([%[[C1]], %[[C45]]]) type(%[[C268435488]]) encoding(%[[C1_i32]]) : !hal.buffer_view
- %tensor = stream.tensor.export %arg0 : tensor<1x45xi32> in !stream.resource<external>{%c180} -> tensor<1x45xi32>
-
- // CHECK: hal_inline.buffer_view.trace %[[VIEW]] : !hal.buffer_view attributes {key = "whatevs"}
- stream.tensor.trace {key = "whatevs"} %tensor : tensor<1x45xi32>
- return
-}
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/debug_ops.mlir b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/debug_ops.mlir
new file mode 100644
index 0000000..24d3565
--- /dev/null
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/debug_ops.mlir
@@ -0,0 +1,16 @@
+// RUN: iree-opt --split-input-file --iree-hal-inline-conversion %s | FileCheck %s
+
+// CHECK-LABEL: @tensorTrace
+// CHECK-SAME: (%[[TENSOR0_STORAGE:.+]]: !util.buffer, %[[TENSOR0_SIZE:.+]]: index, %[[TENSOR1_STORAGE:.+]]: !util.buffer, %[[TENSOR1_SIZE:.+]]: index, %[[TENSOR1_DIM0:.+]]: index)
+func.func @tensorTrace(%tensor0: !stream.resource<staging>, %tensor0_size: index, %tensor1: !stream.resource<staging>, %tensor1_size: index, %tensor1_dim0: index) {
+ // CHECK-DAG: %[[TENSOR0_BUFFER:.+]] = hal_inline.buffer.wrap source(%[[TENSOR0_STORAGE]] : !util.buffer)[%c0, %[[TENSOR0_SIZE]]] : !hal.buffer
+ // CHECK-DAG: %[[TENSOR0:.+]] = hal_inline.buffer_view.create buffer(%[[TENSOR0_BUFFER]] : !hal.buffer)[%c0{{.*}}, %[[TENSOR0_SIZE]]] shape([%c5, %c3])
+ // CHECK-DAG: %[[TENSOR1_BUFFER:.+]] = hal_inline.buffer.wrap source(%[[TENSOR1_STORAGE]] : !util.buffer)[%c0, %[[TENSOR1_SIZE]]] : !hal.buffer
+ // CHECK-DAG: %[[TENSOR1:.+]] = hal_inline.buffer_view.create buffer(%[[TENSOR1_BUFFER]] : !hal.buffer)[%c0{{.*}}, %[[TENSOR1_SIZE]]] shape([%[[TENSOR1_DIM0]], %c5{{.*}}])
+ // CHECK: hal_inline.buffer_view.trace "FOOBAR" = %[[TENSOR0]], %[[TENSOR1]] : !hal.buffer_view, !hal.buffer_view
+ stream.tensor.trace "FOOBAR" = [
+ %tensor0 : tensor<5x3xf32> in !stream.resource<staging>{%tensor0_size},
+ %tensor1 : tensor<?x5xf32>{%tensor1_dim0} in !stream.resource<staging>{%tensor1_size}
+ ]
+ return
+}
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineOps.td b/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineOps.td
index 4587058..c93d1c8 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineOps.td
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineOps.td
@@ -423,6 +423,7 @@
);
let assemblyFormat = [{
+ $key `=`
$operands `:` type($operands)
attr-dict-with-keyword
}];
diff --git a/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/test/cmd_ops.mlir b/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/test/cmd_ops.mlir
index 3bde0c6..30d5fa4 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/test/cmd_ops.mlir
+++ b/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/test/cmd_ops.mlir
@@ -89,24 +89,3 @@
// CHECK: return %c0
return %fence : !stream.timepoint
}
-
-// -----
-
-// CHECK-LABEL: @trace_tensor
-// CHECK-SAME: %[[ARG0:.+]]: !hal.buffer
-func.func @trace_tensor(%arg0: !stream.resource<external>) -> () {
- // CHECK-DAG: %[[C180:.+]] = arith.constant 180 : index
- %c180 = arith.constant 180 : index
-
- // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
- // CHECK-DAG: %[[C45:.+]] = arith.constant 45 : index
- // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[C268435488:.+]] = arith.constant 268435488 : i32
- // CHECK-DAG: %[[C1_i32:.+]] = arith.constant 1 : i32
- // CHECK-DAG: %[[VIEW:.+]] = hal_inline.buffer_view.create buffer(%[[ARG0]] : !hal.buffer)[%[[C0]], %[[C180]]] shape([%[[C1]], %[[C45]]]) type(%[[C268435488]]) encoding(%[[C1_i32]]) : !hal.buffer_view
- %tensor = stream.tensor.export %arg0 : tensor<1x45xi32> in !stream.resource<external>{%c180} -> tensor<1x45xi32>
-
- // CHECK: hal_inline.buffer_view.trace %[[VIEW]] : !hal.buffer_view attributes {key = "whatevs"}
- stream.tensor.trace {key = "whatevs"} %tensor : tensor<1x45xi32>
- return
-}
diff --git a/llvm-external-projects/iree-dialects/BUILD.bazel b/llvm-external-projects/iree-dialects/BUILD.bazel
index 8b8c2d1..408b265 100644
--- a/llvm-external-projects/iree-dialects/BUILD.bazel
+++ b/llvm-external-projects/iree-dialects/BUILD.bazel
@@ -142,6 +142,7 @@
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:Support",
+ "@llvm-project//mlir:TensorDialect",
],
)
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/Input/InputOps.td b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/Input/InputOps.td
index a78a94b..a7053c6 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/Input/InputOps.td
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/Input/InputOps.td
@@ -681,21 +681,31 @@
];
}
-def IREEInput_TensorTraceOp : IREEInput_Op<"tensor.trace", []> {
- let summary = [{trace value(s) operation}];
+def IREEInput_TensorTraceOp : IREEInput_Op<"tensor.trace", [
+ AttrSizedOperandSegments,
+]> {
+ let summary = [{traces one or more tensor values at runtime}];
let description = [{
Traces out to a runtime trace sink (console, log file, etc) the given
- tensors and titles them with the given key. The key is informational only
- and useful for titling/marking specific sets of tensors for easier
- searching.
+ tensors. The key is arbitrary and can be used for identifying the set of
+ values being traced.
}];
let arguments = (ins
StrAttr:$key,
- Variadic<IREEInput_Tensor>:$operands
+ Variadic<IREEInput_Tensor>:$values,
+ IREEInput_ShapeDynamicDims:$value_dims
);
- let assemblyFormat = "$key attr-dict ($operands^ `:` type($operands))?";
+ let assemblyFormat = [{
+ $key `=` `[`
+ custom<ShapedOperandList>($values, type($values), $value_dims)
+ `]` attr-dict-with-keyword
+ }];
+
+ let builders = [
+ OpBuilder<(ins "StringRef":$key, "ValueRange":$values)>,
+ ];
}
} // OpGroupTensorOps
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/Input/CMakeLists.txt b/llvm-external-projects/iree-dialects/lib/Dialect/Input/CMakeLists.txt
index bd5df20..31fb4bf 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/Input/CMakeLists.txt
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/Input/CMakeLists.txt
@@ -14,6 +14,7 @@
MLIRIR
MLIRInferTypeOpInterface
MLIRSideEffectInterfaces
+ MLIRTensorDialect
)
iree_dialects_target_includes(IREEInputDialect)
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/Input/InputOps.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/Input/InputOps.cpp
index 3ed69a0..473e508 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/Input/InputOps.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/Input/InputOps.cpp
@@ -8,6 +8,7 @@
#include "iree-dialects/Dialect/Input/InputDialect.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
@@ -313,13 +314,57 @@
}
//===----------------------------------------------------------------------===//
+// custom<ShapedOperandList>($values, type($values), $value_dims)
+//===----------------------------------------------------------------------===//
+// %value : type{%dynamic_dims}, ...
+
+ParseResult parseShapedOperandList(
+ OpAsmParser &parser,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
+ SmallVectorImpl<Type> &valueTypes,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &valueDims) {
+ do {
+ values.emplace_back();
+ valueTypes.emplace_back();
+ if (failed(parser.parseOperand(values.back())) ||
+ failed(parser.parseColon()) ||
+ failed(parser.parseType(valueTypes.back())))
+ return failure();
+ if (int64_t dynamicDimCount =
+ cast<ShapedType>(valueTypes.back()).getNumDynamicDims()) {
+ if (failed(parser.parseOperandList(valueDims, dynamicDimCount,
+ AsmParser::Delimiter::Braces)))
+ return failure();
+ }
+ } while (succeeded(parser.parseOptionalComma()));
+ return success();
+}
+
+void printShapedOperandList(OpAsmPrinter &p, Operation *op, ValueRange values,
+ TypeRange valueTypes, ValueRange valueDims) {
+ llvm::interleaveComma(llvm::zip_equal(values, valueTypes), p, [&](auto it) {
+ auto [value, valueType] = it;
+ p << value;
+ p << " : ";
+ p << valueType;
+ if (int64_t dynamicDimCount =
+ cast<ShapedType>(valueType).getNumDynamicDims()) {
+ p << "{";
+ llvm::interleaveComma(valueDims.take_front(dynamicDimCount), p);
+ valueDims = valueDims.drop_front(dynamicDimCount);
+ p << "}";
+ }
+ });
+}
+
+//===----------------------------------------------------------------------===//
// custom<ShapedFunctionType>
//===----------------------------------------------------------------------===//
// (type, type{%dim0, %dim1}, type) -> (type{%dim2}, %operand4)
static ParseResult
-parseShapedOperandList(OpAsmParser &parser, SmallVectorImpl<Type> &types,
- SmallVectorImpl<OpAsmParser::UnresolvedOperand> &dims) {
+parseShapedTypeList(OpAsmParser &parser, SmallVectorImpl<Type> &types,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &dims) {
do {
Type type;
if (failed(parser.parseType(type)))
@@ -462,7 +507,7 @@
if (failed(parser.parseLParen()))
return failure();
if (failed(parser.parseOptionalRParen())) {
- if (failed(parseShapedOperandList(parser, operandTypes, operandDims)) ||
+ if (failed(parseShapedTypeList(parser, operandTypes, operandDims)) ||
failed(parser.parseRParen())) {
return failure();
}
@@ -692,6 +737,25 @@
}
//===----------------------------------------------------------------------===//
+// iree_input.tensor.trace
+//===----------------------------------------------------------------------===//
+
+void TensorTraceOp::build(OpBuilder &builder, OperationState &state,
+ StringRef key, ValueRange values) {
+ SmallVector<Value> dynamicDims;
+ for (auto value : values) {
+ auto valueType = cast<ShapedType>(value.getType());
+ for (unsigned i = 0; i < valueType.getRank(); ++i) {
+ if (valueType.isDynamicDim(i)) {
+ dynamicDims.push_back(
+ builder.createOrFold<tensor::DimOp>(state.location, value, i));
+ }
+ }
+ }
+ build(builder, state, key, values, dynamicDims);
+}
+
+//===----------------------------------------------------------------------===//
// iree_input.dispatch
//===----------------------------------------------------------------------===//
