diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp index 0fc10cfe5b59..26dd4e9385a2 100644 --- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp +++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp @@ -190,6 +190,51 @@ static bool doesSliceSpanWholeTarget( return true; } +//===----------------------------------------------------------------------===// +// custom($values, type($values), $value_dims) +//===----------------------------------------------------------------------===// +// %value : type{%dynamic_dims}, ... + +static ParseResult parseShapedOperandList( + OpAsmParser &parser, + SmallVectorImpl &values, + SmallVectorImpl &valueTypes, + SmallVectorImpl &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(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(valueType).getNumDynamicDims()) { + p << "{"; + llvm::interleaveComma(valueDims.take_front(dynamicDimCount), p); + valueDims = valueDims.drop_front(dynamicDimCount); + p << "}"; + } + }); +} + //===----------------------------------------------------------------------===// // custom($body) //===----------------------------------------------------------------------===// @@ -1648,6 +1693,35 @@ SmallVector TensorUpdateOp::getTiedResultOperandIndices() { return {0}; // target } +//===----------------------------------------------------------------------===// +// 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(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 c15da61ddf19..e3d976c359ba 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 @@ def FLOW_TensorUpdateOp : FLOW_PureOp<"tensor.update", [ 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, +]> { + 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:$operands + Variadic:$values, + FLOW_ShapeDynamicDims:$value_dims ); - let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?"; + let assemblyFormat = [{ + $key `=` `[` + custom($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 636b535f9d71..6cc909d42975 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 @@ func.func @tensorUpdateDynamic(%arg0 : tensor, %arg1 : tensor, %0 = flow.tensor.update %arg0, %arg1[%arg2, %arg3] : tensor{%c1, %c2} -> %arg1 as tensor{%c3} return %0 : tensor } + +// ----- + +// CHECK-LABEL: @tensorTrace +// CHECK-SAME: (%[[TENSOR0:.+]]: tensor<5xf32>, %[[TENSOR1:.+]]: tensor, %[[TENSOR1_DIM0:.+]]: index, %[[TENSOR1_DIM2:.+]]: index) +func.func @tensorTrace(%tensor0: tensor<5xf32>, %tensor1: tensor, %tensor1_dim0: index, %tensor1_dim2: index) { + // CHECK: flow.tensor.trace "FOOBAR" = [ + // CHECK-SAME: %[[TENSOR0]] : tensor<5xf32>, + // CHECK-SAME: %[[TENSOR1]] : tensor{%[[TENSOR1_DIM0]], %[[TENSOR1_DIM2]]} + // CHECK-SAME: ] + flow.tensor.trace "FOOBAR" = [ + %tensor0 : tensor<5xf32>, + %tensor1 : tensor{%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 22c9c7902f03..c70620b5caaf 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 @@ class InjectDispatchTracingPass // Input tensors: OpBuilder builder(dispatchOp); - builder.create( + builder.create( dispatchOp.getLoc(), builder.getStringAttr(entryPointName + " inputs"), filterTensorValues(dispatchOp.getArguments())); // Output tensors: builder.setInsertionPointAfter(dispatchOp); - builder.create( + builder.create( 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 8dbb63d8c530..6186fd6ae33c 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 @@ getOrdinalFromDebugTarget(std::string marker) { } // 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( + builder.create( dispatchOp.getLoc(), builder.getStringAttr(name + " inputs"), filterNonTensorValues(dispatchOp.getArguments())); // Output tensors: OpBuilder::InsertionGuard guard(builder); builder.setInsertionPointAfter(dispatchOp); - builder.create(dispatchOp.getLoc(), - builder.getStringAttr(name + " outputs"), - filterNonTensorValues(dispatchOp.getResults())); + builder.create( + 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 @@ struct InsertDebugTargetAtOrdinalPass localTraceOrdinal = traceOrdinal; auto &bodyRegion = op.getFunctionBody(); - auto dispatchOps = llvm::to_vector<8>(bodyRegion.getOps()); + auto dispatchOps = + llvm::to_vector<8>(bodyRegion.getOps()); // Trace on a valid ordinal. if (localTraceOrdinal >= 0 && localTraceOrdinal < dispatchOps.size()) { @@ -222,8 +224,8 @@ struct InsertDebugTargetAtSymbolPass // 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 7471d553907c..c090fd6f38c5 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 @singleDispatch(%arg0: tensor<4xf32>) -> tensor<4xf32> { 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 f17537bbe2bb..c15f2686dde6 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 @@ func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view { 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 @@ func.func @other_func(%arg0: tensor<4xf32>) -> !hal.buffer_view { // ----- -// 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 @@ func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view { // ----- -// 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 @@ func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view { // ----- -/// 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 @@ func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view { // 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 @@ func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view { // ----- -/// 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 8d31a0282297..9c7ebe047a2e 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 @@ struct TensorTraceOpPattern LogicalResult matchAndRewrite(IREE::Stream::TensorTraceOp traceOp, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { + SmallVector bufferViews; + auto resourceEncodingDims = adaptor.getResourceEncodingDims(); + for (auto [resource, resourceSize, resourceEncoding] : llvm::zip_equal( + adaptor.getResources(), adaptor.getResourceSizes(), + adaptor.getResourceEncodings().getAsRange())) { + int64_t dynamicDimCount = + cast(resourceEncoding.getValue()).getNumDynamicDims(); + bufferViews.push_back(rewriter.create( + traceOp.getLoc(), rewriter.getType(), + resource, resourceEncoding, + resourceEncodingDims.take_front(dynamicDimCount), resourceSize, + /*affinity=*/IREE::Stream::AffinityAttr{})); + resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount); + } rewriter.replaceOpWithNewOp( - 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 832da446424f..64493c06838f 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 @@ iree_lit_test_suite( [ "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 5c25a89a7d7d..fb7c40faae4a 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 @@ iree_lit_test_suite( 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 000000000000..a748326cff22 --- /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, %tensor0_size: index, %tensor1: !stream.resource, %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{%tensor0_size}, + %tensor1 : tensor{%tensor1_dim0} in !stream.resource{%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 bc58f6fa2555..e2f795ec6d52 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 @@ def HAL_BufferViewTraceOp : HAL_Op<"buffer_view.trace", []> { ); 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 e7e2bedbf4b6..e036755f4358 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 @@ class ValueResourceUsage : public AbstractResourceUsage { .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( 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 297eba646a33..4ae15713d08b 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 @@ struct ConvertTensorTraceOp LogicalResult matchAndRewrite(IREE::Flow::TensorTraceOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - SmallVector exportedTensors; + SmallVector resources; + SmallVector resourceSizes; + SmallVector 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::Lifetime::External); - auto exportSource = resourceOperand; - if (source.resource.getType() != externalType) { - exportSource = rewriter.create( - op.getLoc(), externalType, source.resource, source.resourceSize, + auto stagingType = rewriter.getType( + IREE::Stream::Lifetime::Staging); + auto traceSource = source.resource; + if (source.resource.getType() != stagingType) { + traceSource = rewriter.create( + 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( - 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( - 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 d194376cf072..4e2a972a360b 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 @@ func.func @tensorStore(%target : tensor<2x3xi32>) -> tensor<2x3xi32> { // 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, %tensor1_dim0: index, %tensor1_dim2: index) { + // CHECK-DAG: %[[TENSOR0_STAGED:.+]] = stream.async.transfer %[[TENSOR0]] : !stream.resource<*>{%[[TENSOR0_SIZE]]} -> !stream.resource{%[[TENSOR0_SIZE]]} + // CHECK-DAG: %[[TENSOR1_STAGED:.+]] = stream.async.transfer %[[TENSOR1]] : !stream.resource<*>{%[[TENSOR1_SIZE]]} -> !stream.resource{%[[TENSOR1_SIZE]]} + // CHECK: stream.tensor.trace "FOOBAR" = [ + // CHECK-NEXT: %[[TENSOR0_STAGED]] : tensor<5xf32> in !stream.resource{%[[TENSOR0_SIZE]]}, + // CHECK-NEXT: %[[TENSOR1_STAGED]] : tensor{%[[TENSOR1_DIM0]], %[[TENSOR1_DIM2]]} in !stream.resource{%[[TENSOR1_SIZE]]} + // CHECK-NEXT: ] + flow.tensor.trace "FOOBAR" = [ + %tensor0 : tensor<5xf32>, + %tensor1 : tensor{%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 addada816aae..a294923d2133 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 @@ static LogicalResult verifyEscapingResources(Region ®ion, return success(); } +static void eraseStreamRegionResults(Region ®ion, + ArrayRef excludedResultIndices) { + for (auto &block : region.getBlocks()) { + auto yieldOp = dyn_cast(block.getTerminator()); + if (!yieldOp) + continue; + llvm::SmallVector 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 @@ static IREE::Util::ValueAccess computeValueAccess(Value rootValue) { return access; } -static void eraseStreamRegionResults(Region ®ion, - ArrayRef excludedResultIndices) { - for (auto &block : region.getBlocks()) { - auto yieldOp = dyn_cast(block.getTerminator()); - if (!yieldOp) - continue; - llvm::SmallVector newOperands; - for (auto i : llvm::reverse(excludedResultIndices)) { - yieldOp.getResourceOperandsMutable().erase(i); - yieldOp.getResourceOperandSizesMutable().erase(i); +//===----------------------------------------------------------------------===// +// custom( +// $resources, type($resources), $resource_sizes, +// $resource_encodings, $resource_encoding_dims) +//===----------------------------------------------------------------------===// + +static ParseResult parseEncodedResourceOperands( + OpAsmParser &parser, + SmallVectorImpl &resources, + SmallVectorImpl &resourceTypes, + SmallVectorImpl &resourceSizes, + ArrayAttr &resourceEncodings, + SmallVectorImpl &resourceEncodingDims) { + SmallVector 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(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()), + [&](auto it) { + auto [resource, resourceType, resourceSize, resourceEncoding] = it; + p << resource; + p << " : "; + p << resourceEncoding; + if (int64_t dynamicDimCount = + cast(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(); } //===----------------------------------------------------------------------===// @@ -1211,6 +1289,54 @@ SmallVector TensorStoreOp::getTiedResultOperandIndices() { return {0}; // target } +//===----------------------------------------------------------------------===// +// 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())) { + int64_t dynamicDimCount = + cast(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().getValue(); + int64_t dynamicDimCount = + cast(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 275c178ff812..0d5271860a6c 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 @@ def Stream_TensorStoreOp : Stream_PureOp<"tensor.store", [ 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_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:$operands + Variadic:$resources, + Variadic:$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( + $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 f65041548ba4..4cc586c74bfd 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 @@ func.func @tensorStoreRank0(%arg0: !stream.resource, %arg1: index, %arg %0 = stream.tensor.store %arg2, %arg0 : f32 -> tensor in %arg0 as !stream.resource{%arg1} return %0 : !stream.resource } + +// ----- + +// 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: !stream.resource, %tensor0_size: index, %tensor1: !stream.resource, %tensor1_size: index, %tensor1_dim0: index, %tensor1_dim2: index) { + // CHECK: stream.tensor.trace "FOOBAR" = [ + // CHECK-NEXT: %[[TENSOR0]] : tensor<5xf32> in !stream.resource{%[[TENSOR0_SIZE]]}, + // CHECK-NEXT: %[[TENSOR1]] : tensor{%[[TENSOR1_DIM0]], %[[TENSOR1_DIM2]]} in !stream.resource{%[[TENSOR1_SIZE]]} + // CHECK-NEXT: ] + stream.tensor.trace "FOOBAR" = [ + %tensor0 : tensor<5xf32> in !stream.resource{%tensor0_size}, + %tensor1 : tensor{%tensor1_dim0, %tensor1_dim2} in !stream.resource{%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 b13b245c3dce..b8cbc9c03e2b 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 @@ SmallVector buildDynamicDimsForValue(Location loc, Value value, return dynamicDims; } +SmallVector buildDynamicDimsForValues(Location loc, ValueRange values, + OpBuilder &builder) { + SmallVector dynamicDims; + for (auto value : values) { + dynamicDims.append(buildDynamicDimsForValue(loc, value, builder)); + } + return dynamicDims; +} + static SmallVector 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 76bb1c1f8209..f27ee119694f 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 @@ ValueRange findVariadicDynamicDims(unsigned idx, ValueRange values, SmallVector 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 buildDynamicDimsForValues(Location loc, ValueRange values, + OpBuilder &builder); + // Builds a ranked shape with all dimension values for the given operand. SmallVector 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 f9b373089ee2..5aa7bbfcdd53 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 @@ func.func @null_op() -> !iree_input.variant { // 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 @@ func.func @tensor_update(%arg0 : tensor, %arg1 : index, %arg2 : index, %a // ----- // 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{%arg2} +// CHECK-SAME: ] +func.func @tensor_trace(%arg0: tensor<5xf32>, %arg1: tensor, %arg2: index) { + iree_input.tensor.trace "FOOBAR" = [ + %arg0 : tensor<5xf32>, + %arg1 : tensor{%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 fda46c464872..f21e24cd9418 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 @@ struct TensorTraceOpPattern LogicalResult matchAndRewrite(IREE::Stream::TensorTraceOp traceOp, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { + auto bufferType = rewriter.getType(); + auto bufferViewType = rewriter.getType(); + auto zero = rewriter.create(traceOp.getLoc(), 0); + auto resourceEncodingDims = adaptor.getResourceEncodingDims(); + SmallVector bufferViews; + for (auto [resource, resourceSize, resourceEncoding] : llvm::zip_equal( + adaptor.getResources(), adaptor.getResourceSizes(), + adaptor.getResourceEncodings().getAsRange())) { + Value resourceBuffer = rewriter.create( + traceOp.getLoc(), bufferType, resource, + /*offset=*/ + zero, + /*length=*/resourceSize); + int64_t dynamicDimCount = + cast(resourceEncoding.getValue()).getNumDynamicDims(); + bufferViews.push_back(rewriter.create( + traceOp.getLoc(), bufferViewType, resourceBuffer, resourceEncoding, + resourceEncodingDims.take_front(dynamicDimCount), resourceSize, + /*affinity=*/IREE::Stream::AffinityAttr{})); + resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount); + } rewriter.replaceOpWithNewOp( - 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 db19be94a61d..09466b8d2ec6 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 @@ iree_lit_test_suite( 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 930427fdbd2f..c3f6ba0f2650 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 @@ iree_lit_test_suite( 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 72549c17eaca..cf6ff9750b45 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 @@ func.func @cmdCall(%arg0: !stream.resource, %arg1: i32, %arg2: !stream } => !stream.timepoint return %timepoint : !stream.timepoint } - -// ----- - -// CHECK-LABEL: @trace_tensor -// CHECK-SAME: %[[ARG0:.+]]: !hal.buffer -func.func @trace_tensor(%arg0: !stream.resource) -> () { - // 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{%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 000000000000..24d356555a75 --- /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, %tensor0_size: index, %tensor1: !stream.resource, %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{%tensor0_size}, + %tensor1 : tensor{%tensor1_dim0} in !stream.resource{%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 4587058ddc53..c93d1c81ae06 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 @@ def HALInline_BufferViewTraceOp : HALInline_Op<"buffer_view.trace", []> { ); 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 3bde0c6685a0..30d5fa4d8fd3 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 @@ func.func @cmdDispatch(%buffer0: !stream.resource, %buffer0_size: ind // CHECK: return %c0 return %fence : !stream.timepoint } - -// ----- - -// CHECK-LABEL: @trace_tensor -// CHECK-SAME: %[[ARG0:.+]]: !hal.buffer -func.func @trace_tensor(%arg0: !stream.resource) -> () { - // 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{%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 8b8c2d11e83d..408b2651e414 100644 --- a/llvm-external-projects/iree-dialects/BUILD.bazel +++ b/llvm-external-projects/iree-dialects/BUILD.bazel @@ -142,6 +142,7 @@ cc_library( "@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 a78a94b0b447..a7053c60171f 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_TensorUpdateOp : IREEInput_PureOp<"tensor.update", [ ]; } -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:$operands + Variadic:$values, + IREEInput_ShapeDynamicDims:$value_dims ); - let assemblyFormat = "$key attr-dict ($operands^ `:` type($operands))?"; + let assemblyFormat = [{ + $key `=` `[` + custom($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 bd5df2020cff..31fb4bfc8f8c 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 @@ add_mlir_library(IREEInputDialect 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 3ed69a006b5f..473e5082e1b5 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" @@ -312,14 +313,58 @@ static void printShapedTiedResult(OpAsmPrinter &p, TiedOpInterface op, printShapedTiedResult(p, op, resultType, resultDims); } +//===----------------------------------------------------------------------===// +// custom($values, type($values), $value_dims) +//===----------------------------------------------------------------------===// +// %value : type{%dynamic_dims}, ... + +ParseResult parseShapedOperandList( + OpAsmParser &parser, + SmallVectorImpl &values, + SmallVectorImpl &valueTypes, + SmallVectorImpl &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(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(valueType).getNumDynamicDims()) { + p << "{"; + llvm::interleaveComma(valueDims.take_front(dynamicDimCount), p); + valueDims = valueDims.drop_front(dynamicDimCount); + p << "}"; + } + }); +} + //===----------------------------------------------------------------------===// // custom //===----------------------------------------------------------------------===// // (type, type{%dim0, %dim1}, type) -> (type{%dim2}, %operand4) static ParseResult -parseShapedOperandList(OpAsmParser &parser, SmallVectorImpl &types, - SmallVectorImpl &dims) { +parseShapedTypeList(OpAsmParser &parser, SmallVectorImpl &types, + SmallVectorImpl &dims) { do { Type type; if (failed(parser.parseType(type))) @@ -462,7 +507,7 @@ static ParseResult parseShapedFunctionType( 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(); } @@ -691,6 +736,25 @@ SmallVector TensorUpdateOp::getTiedResultOperandIndices() { return {0}; // $target } +//===----------------------------------------------------------------------===// +// iree_input.tensor.trace +//===----------------------------------------------------------------------===// + +void TensorTraceOp::build(OpBuilder &builder, OperationState &state, + StringRef key, ValueRange values) { + SmallVector dynamicDims; + for (auto value : values) { + auto valueType = cast(value.getType()); + for (unsigned i = 0; i < valueType.getRank(); ++i) { + if (valueType.isDynamicDim(i)) { + dynamicDims.push_back( + builder.createOrFold(state.location, value, i)); + } + } + } + build(builder, state, key, values, dynamicDims); +} + //===----------------------------------------------------------------------===// // iree_input.dispatch //===----------------------------------------------------------------------===//