Adding optional target_storage to hal.tensor.export. This allows frontends to control the storage for exported tensors by tying arguments to results. Today it's required that these do not alias but we could do fun things in the future with providing slabs to place multiple results into. The native bindings have been updated to support iree.abi.output arg attrs indicating the result an argument provides storage for. Other bindings may decide this on their own however they want. There's still inefficiencies here (a copy is almost always performed) but this allows us to avoid the allocation and eliding the copy is something that we eventually need to do anyway. Output buffers can now be passed to tools as inputs with leading `&` (indicating storage reference), ex: `--function_input=&4x8xf32`.
diff --git a/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp b/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp index e2dcefd..a3873b8 100644 --- a/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp +++ b/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp
@@ -128,6 +128,25 @@ auto *entryBlock = wrapperFuncOp.addEntryBlock(); auto entryBuilder = OpBuilder::atBlockBegin(entryBlock); + // Build a map of result value to the argument that has its backing storage. + SmallVector<Value> resultStorages; + resultStorages.resize(resultTypes.size()); + for (unsigned i = 0; i < inputTypes.size(); ++i) { + auto outputAttr = + entryFuncOp.getArgAttrOfType<IntegerAttr>(i, "iree.abi.output"); + if (!outputAttr) continue; + // Today all outputs need to be a !hal.buffer - we could change this + // in the future to be something more generalized. + auto storageArg = entryBlock->getArgument(i); + if (!storageArg.getType().isa<IREE::HAL::BufferType>()) { + entryFuncOp.emitError() + << "storage argument " << i << " has an invalid type " + << storageArg.getType() << "; must be a !hal.buffer"; + return {}; + } + resultStorages[outputAttr.getInt()] = storageArg; + } + // Marshal arguments. SmallVector<Value> arguments; for (auto arg : llvm::enumerate(entryBlock->getArguments())) { @@ -152,8 +171,12 @@ auto oldType = entryFuncType.getResult(result.index()); auto newType = wrapperFuncType.getResult(result.index()); if (oldType.isa<TensorType>()) { - results.push_back(entryBuilder.createOrFold<IREE::HAL::TensorExportOp>( - entryFuncOp.getLoc(), newType, result.value())); + auto dynamicDims = IREE::Util::buildDynamicDimsForValue( + entryFuncOp.getLoc(), result.value(), entryBuilder); + results.push_back(entryBuilder.create<IREE::HAL::TensorExportOp>( + entryFuncOp.getLoc(), newType, result.value(), + TypeAttr::get(result.value().getType()), dynamicDims, + resultStorages[result.index()])); } else { results.push_back(result.value()); }
diff --git a/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points.mlir b/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points.mlir index 1762358..a56f0c3 100644 --- a/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points.mlir +++ b/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points.mlir
@@ -11,11 +11,11 @@ // CHECK-NEXT: %[[ARG0_TENSOR:.+]] = hal.tensor.import %[[ARG0]] : !hal.buffer_view -> tensor<?x8x8x3xf32>{%[[ARG0_DIM0]]} // CHECK-NEXT: %[[ARG1_DIM0:.+]] = hal.buffer_view.dim<%[[ARG1]] : !hal.buffer_view>[0] : index // CHECK-NEXT: %[[ARG1_TENSOR:.+]] = hal.tensor.import %[[ARG1]] : !hal.buffer_view -> tensor<?x8x8x3xf32>{%[[ARG1_DIM0]]} -// CHECK-NEXT: %[[RET_TENSOR:.+]]:2 = call @_dynamicEntry(%[[ARG0_TENSOR]], %[[ARG1_TENSOR]]) -// CHECK: %[[RET0_DIM0:.+]] = tensor.dim %[[RET_TENSOR]]#0, %c0{{.*}} : tensor<?x8x8x3xf32> -// CHECK-NEXT: %[[RET0_VIEW:.+]] = hal.tensor.export %[[RET_TENSOR]]#0 : tensor<?x8x8x3xf32>{%[[RET0_DIM0]]} -> !hal.buffer_view -// CHECK: %[[RET1_DIM0:.+]] = tensor.dim %[[RET_TENSOR]]#1, %c0{{.*}} : tensor<?x8x8x3xf32> -// CHECK-NEXT: %[[RET1_VIEW:.+]] = hal.tensor.export %[[RET_TENSOR]]#1 : tensor<?x8x8x3xf32>{%[[RET1_DIM0]]} -> !hal.buffer_view +// CHECK-NEXT: %[[RET_TENSORS:.+]]:2 = call @_dynamicEntry(%[[ARG0_TENSOR]], %[[ARG1_TENSOR]]) +// CHECK: %[[RET0_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#0, %c0{{.*}} : tensor<?x8x8x3xf32> +// CHECK-NEXT: %[[RET0_VIEW:.+]] = hal.tensor.export %[[RET_TENSORS]]#0 : tensor<?x8x8x3xf32>{%[[RET0_DIM0]]} -> !hal.buffer_view +// CHECK: %[[RET1_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#1, %c0{{.*}} : tensor<?x8x8x3xf32> +// CHECK-NEXT: %[[RET1_VIEW:.+]] = hal.tensor.export %[[RET_TENSORS]]#1 : tensor<?x8x8x3xf32>{%[[RET1_DIM0]]} -> !hal.buffer_view // CHECK-NEXT: return %[[RET0_VIEW]], %[[RET1_VIEW]] : !hal.buffer_view, !hal.buffer_view // CHECK-NEXT: } @@ -29,6 +29,34 @@ // ----- +// CHECK-LABEL: func @outputStorage( +// CHECK-SAME: %[[ARG0:.+]]: !hal.buffer_view, +// CHECK-SAME: %[[RET1_STORAGE:.+]]: !hal.buffer +// CHECK-SAME: -> ( +// CHECK-SAME: !hal.buffer_view, !hal.buffer_view +// CHECK-SAME: ) attributes { +// CHECK-SAME: iree.abi.stub +// CHECK-SAME: } { +// CHECK-NEXT: %[[ARG0_DIM0:.+]] = hal.buffer_view.dim<%[[ARG0]] : !hal.buffer_view>[0] : index +// CHECK-NEXT: %[[ARG0_TENSOR:.+]] = hal.tensor.import %[[ARG0]] : !hal.buffer_view -> tensor<?x8x8x3xf32>{%[[ARG0_DIM0]]} +// CHECK-NEXT: %[[RET_TENSORS:.+]]:2 = call @_outputStorage(%[[ARG0_TENSOR]], %[[RET1_STORAGE]]) +// CHECK: %[[RET0_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#0, %c0{{.*}} : tensor<?x8x8x3xf32> +// CHECK-NEXT: %[[RET0_VIEW:.+]] = hal.tensor.export %[[RET_TENSORS]]#0 : tensor<?x8x8x3xf32>{%[[RET0_DIM0]]} -> !hal.buffer_view +// CHECK: %[[RET1_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#1, %c0{{.*}} : tensor<?x8x8x3xf32> +// CHECK-NEXT: %[[RET1_VIEW:.+]] = hal.tensor.export %[[RET_TENSORS]]#1 into %[[RET1_STORAGE]] : tensor<?x8x8x3xf32>{%[[RET1_DIM0]]} -> !hal.buffer_view +// CHECK-NEXT: return %[[RET0_VIEW]], %[[RET1_VIEW]] : !hal.buffer_view, !hal.buffer_view +// CHECK-NEXT: } + +// CHECK-LABEL: func private @_outputStorage( +func @outputStorage(%arg0: tensor<?x8x8x3xf32>, %ret1: !hal.buffer {iree.abi.output = 1 : index}) -> + (tensor<?x8x8x3xf32>, tensor<?x8x8x3xf32>) { + %0 = "mhlo.add"(%arg0, %arg0) : (tensor<?x8x8x3xf32>, tensor<?x8x8x3xf32>) -> tensor<?x8x8x3xf32> + %1 = "mhlo.add"(%0, %arg0) : (tensor<?x8x8x3xf32>, tensor<?x8x8x3xf32>) -> tensor<?x8x8x3xf32> + return %0, %1 : tensor<?x8x8x3xf32>, tensor<?x8x8x3xf32> +} + +// ----- + // CHECK-LABEL: func @wrappedAlready // CHECK-SAME: (%arg0: !hal.buffer_view) -> !hal.buffer_view // CHECK-SAME: attributes {iree.abi.stub}
diff --git a/iree/compiler/Bindings/TFLite/Transforms/WrapEntryPoints.cpp b/iree/compiler/Bindings/TFLite/Transforms/WrapEntryPoints.cpp index 3ba6149..a74fa5c 100644 --- a/iree/compiler/Bindings/TFLite/Transforms/WrapEntryPoints.cpp +++ b/iree/compiler/Bindings/TFLite/Transforms/WrapEntryPoints.cpp
@@ -543,7 +543,7 @@ } callResults.push_back(entryBuilder.create<IREE::HAL::TensorExportOp>( result.getLoc(), bufferType, result, outputDynamicDims.tensorType, - dynamicDims)); + dynamicDims, /*target_storage=*/nullptr)); for (auto it : llvm::zip(dynamicDims, outputDynamicDims.globalOps)) { auto dynamicDim = std::get<0>(it); auto globalOp = std::get<1>(it);
diff --git a/iree/compiler/Dialect/HAL/IR/HALOps.cpp b/iree/compiler/Dialect/HAL/IR/HALOps.cpp index 76c96d2..7e6024c 100644 --- a/iree/compiler/Dialect/HAL/IR/HALOps.cpp +++ b/iree/compiler/Dialect/HAL/IR/HALOps.cpp
@@ -269,7 +269,7 @@ auto dynamicDims = IREE::Util::buildDynamicDimsForValue(result.location, source, builder); build(builder, result, resultType, source, TypeAttr::get(source.getType()), - dynamicDims); + dynamicDims, /*target_storage=*/nullptr); } Value TensorExportOp::getTiedResult(unsigned resultIndex) {
diff --git a/iree/compiler/Dialect/HAL/IR/HALOps.td b/iree/compiler/Dialect/HAL/IR/HALOps.td index 9d7d063..06df00b 100644 --- a/iree/compiler/Dialect/HAL/IR/HALOps.td +++ b/iree/compiler/Dialect/HAL/IR/HALOps.td
@@ -115,6 +115,7 @@ } def HAL_TensorExportOp : HAL_PureOp<"tensor.export", [ + AttrSizedOperandSegments, DeclareOpInterfaceMethods<Util_TiedOpInterface, [ "getTiedResult", "getTiedResultOperandIndex", @@ -135,19 +136,27 @@ dynamically shaped values must have the same number of dynamic dimensions. This allows for casting between rank-0 and rank-N types, different element types, etc. + + An optional `target_storage` buffer can be provided to hold the exported + result. The export will fail at runtime if the storage is null or if it has + insufficient capacity to store the output. The storage must be + device-visible and defined for transfer-target and dispatch usage. }]; let arguments = (ins AnyTensor:$source, TypeAttr:$source_encoding, - HAL_ShapeDynamicDims:$source_dims + HAL_ShapeDynamicDims:$source_dims, + Optional<HAL_Buffer>:$target_storage ); let results = (outs AnyTypeOf<[HAL_Buffer, HAL_BufferView]>:$target ); let assemblyFormat = [{ - $source `:` + $source + (`into` $target_storage^)? + `:` custom<TypeAlias>($source_encoding, type($source)) (`{` $source_dims^ `}`)? `->` type($target)
diff --git a/iree/compiler/Dialect/HAL/IR/test/tensor_ops.mlir b/iree/compiler/Dialect/HAL/IR/test/tensor_ops.mlir index 6498bc6..2baeaea 100644 --- a/iree/compiler/Dialect/HAL/IR/test/tensor_ops.mlir +++ b/iree/compiler/Dialect/HAL/IR/test/tensor_ops.mlir
@@ -24,3 +24,12 @@ %0 = hal.tensor.export %arg0 : tensor<?x3xf32> as tensor<?x3xi32>{%arg1} -> !hal.buffer_view return %0 : !hal.buffer_view } + +// ----- + +// CHECK-LABEL: @tensorExportInPlace +func @tensorExportInPlace(%arg0: tensor<?x3xi32>, %arg1 : index, %arg2: !hal.buffer) -> !hal.buffer_view { + // CHECK: hal.tensor.export %arg0 into %arg2 : tensor<?x3xf32> as tensor<?x3xi32>{%arg1} -> !hal.buffer_view + %0 = hal.tensor.export %arg0 into %arg2 : tensor<?x3xf32> as tensor<?x3xi32>{%arg1} -> !hal.buffer_view + return %0 : !hal.buffer_view +}
diff --git a/iree/compiler/Dialect/Stream/Conversion/HALToStream/ConvertHALToStream.cpp b/iree/compiler/Dialect/Stream/Conversion/HALToStream/ConvertHALToStream.cpp index fe93673..d8128c8 100644 --- a/iree/compiler/Dialect/Stream/Conversion/HALToStream/ConvertHALToStream.cpp +++ b/iree/compiler/Dialect/Stream/Conversion/HALToStream/ConvertHALToStream.cpp
@@ -141,18 +141,53 @@ auto externalType = rewriter.getType<IREE::Stream::ResourceType>( IREE::Stream::Lifetime::External); auto exportSource = adaptor.source(); - if (source.resource.getType() != externalType) { - exportSource = rewriter.create<IREE::Stream::AsyncTransferOp>( - op.getLoc(), externalType, source.resource, source.resourceSize, - source.resourceSize, - /*source_affinity=*/nullptr, - /*result_affinity=*/nullptr); + auto exportSize = source.resourceSize; + if (adaptor.target_storage()) { + // Query the target storage buffer length; we will only populate up to + // what is required for the output. + auto storageSize = + rewriter + .create<IREE::HAL::BufferLengthOp>( + op.getLoc(), rewriter.getIndexType(), op.target_storage()) + .result(); + + // Import the target storage as a resource that we can use as an update + // target. We overwrite the contents and just cast the storage to the + // target type so we know we can update it. + auto importOp = rewriter.create<IREE::Stream::TensorImportOp>( + op.getLoc(), externalType, adaptor.target_storage(), + TypeAttr::get(sourceType), adaptor.source_dims(), storageSize, + /*affinity=*/nullptr); + + // Copy the source value into the imported target storage. + auto zeroOffset = rewriter.create<arith::ConstantIndexOp>(op.getLoc(), 0); + auto updateOp = rewriter.create<IREE::Stream::AsyncUpdateOp>( + op.getLoc(), externalType, importOp.result(), importOp.result_size(), + zeroOffset, source.resourceSize, source.resource, source.resourceSize, + /*affinity=*/nullptr); + + // Export the updated resource. + // NOTE: the buffer size wrapped in the buffer view is the full size of + // the input buffer. This is so that we don't insert a data dependency on + // sparse operations or data-dependent dynamic shape dimensions. + exportSource = updateOp.result(); + exportSize = updateOp.target_size(); + } else { + // Exporting a produced value - transfer our source value to an externally + // usable resource and directly export it. This will cause an allocation. + if (source.resource.getType() != externalType) { + exportSource = rewriter.create<IREE::Stream::AsyncTransferOp>( + op.getLoc(), externalType, source.resource, source.resourceSize, + source.resourceSize, + /*source_affinity=*/nullptr, + /*result_affinity=*/nullptr); + } } // Export (stream resource to buffer view). rewriter.replaceOpWithNewOp<IREE::Stream::TensorExportOp>( op, targetType, exportSource, TypeAttr::get(sourceType), - adaptor.source_dims(), source.resourceSize, + adaptor.source_dims(), exportSize, /*affinity=*/nullptr); return success(); }
diff --git a/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/abi_ops.mlir b/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/abi_ops.mlir index d7d8849..5f8e181 100644 --- a/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/abi_ops.mlir +++ b/iree/compiler/Dialect/Stream/Conversion/HALToStream/test/abi_ops.mlir
@@ -47,3 +47,21 @@ // CHECK: return %[[RESULT]] return %0 : !hal.buffer_view } + +// ----- + +// CHECK-LABEL: @exportBufferViewInPlace +// CHECK-SAME: (%[[TENSOR:.+]]: !stream.resource<*>, %[[SIZE:.+]]: index, %[[DIM0:.+]]: index, %[[DIM1:.+]]: index, %[[STORAGE:.+]]: !hal.buffer) +func @exportBufferViewInPlace(%tensor: tensor<?x?x4xf32>, %dim0: index, %dim1: index, %storage: !hal.buffer) -> !hal.buffer_view { + // CHECK: %[[STORAGE_LENGTH:.+]] = hal.buffer.length<%[[STORAGE]] + // CHECK-NEXT: %[[STORAGE_IMPORT:.+]] = stream.tensor.import %[[STORAGE]] + // CHECK-SAME: : !hal.buffer -> tensor<?x?x4xf32>{%[[DIM0]], %[[DIM1]]} in !stream.resource<external>{%[[STORAGE_LENGTH]]} + // CHECK-NEXT: %[[STORAGE_UPDATE:.+]] = stream.async.update %[[TENSOR]], %[[STORAGE_IMPORT]][%c0 to %[[SIZE]]] + // CHECK-SAME: : !stream.resource<*>{%[[SIZE]]} -> %[[STORAGE_IMPORT]] as !stream.resource<external>{%[[STORAGE_LENGTH]]} + // CHECK-NEXT: %[[STORAGE_RESULT:.+]] = stream.tensor.export %[[STORAGE_UPDATE]] : + // CHECK-SAME: tensor<?x?x4xf32>{%[[DIM0]], %[[DIM1]]} in !stream.resource<external>{%[[STORAGE_LENGTH]]} + // CHECK-SAME: -> !hal.buffer_view + %0 = hal.tensor.export %tensor into %storage : tensor<?x?x4xf32>{%dim0, %dim1} -> !hal.buffer_view + // CHECK: return %[[STORAGE_RESULT]] + return %0 : !hal.buffer_view +}
diff --git a/iree/compiler/InputConversion/Common/IREEImportPublic.cpp b/iree/compiler/InputConversion/Common/IREEImportPublic.cpp index 895bf00..21096bb 100644 --- a/iree/compiler/InputConversion/Common/IREEImportPublic.cpp +++ b/iree/compiler/InputConversion/Common/IREEImportPublic.cpp
@@ -110,7 +110,8 @@ if (!resultType) return failure(); rewriter.replaceOpWithNewOp<IREE::HAL::TensorExportOp>( srcOp, resultType, adaptor.source(), - TypeAttr::get(adaptor.source().getType()), adaptor.source_dims()); + TypeAttr::get(adaptor.source().getType()), adaptor.source_dims(), + /*target_storage=*/nullptr); return success(); } };
diff --git a/iree/tools/utils/trace_replay.c b/iree/tools/utils/trace_replay.c index fc1a185..827ee84 100644 --- a/iree/tools/utils/trace_replay.c +++ b/iree/tools/utils/trace_replay.c
@@ -465,7 +465,7 @@ // contents: !!binary | // AACAPwAAAEAAAEBAAACAQA== // ``` -static iree_status_t iree_trace_replay_parse_hal_buffer( +static iree_status_t iree_trace_replay_parse_hal_buffer_contents( iree_trace_replay_t* replay, yaml_document_t* document, yaml_node_t* contents_node, iree_hal_element_type_t element_type, iree_hal_buffer_mapping_t* mapping) { @@ -651,9 +651,9 @@ iree_trace_replay_generation_params_t* params = (iree_trace_replay_generation_params_t*)user_data; if (params->contents_node) { - return iree_trace_replay_parse_hal_buffer(params->replay, params->document, - params->contents_node, - params->element_type, mapping); + return iree_trace_replay_parse_hal_buffer_contents( + params->replay, params->document, params->contents_node, + params->element_type, mapping); } else { return iree_trace_replay_generate_hal_buffer( params->replay, params->document, params->generator_node, @@ -661,6 +661,58 @@ } } +// Parses a !hal.buffer and appends it to |target_list|. +// +// ```yaml +// shape: +// - 4 +// element_type: 553648160 +// ``` +static iree_status_t iree_trace_replay_parse_hal_buffer( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* value_node, iree_vm_list_t* target_list) { + yaml_node_t* shape_node = NULL; + IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find( + document, value_node, iree_make_cstring_view("shape"), &shape_node)); + iree_hal_dim_t shape[16]; + iree_host_size_t shape_rank = 0; + IREE_RETURN_IF_ERROR(iree_trace_replay_parse_hal_shape( + replay, document, shape_node, IREE_ARRAYSIZE(shape), shape, &shape_rank)); + + yaml_node_t* element_type_node = NULL; + IREE_RETURN_IF_ERROR(iree_yaml_mapping_find( + document, value_node, iree_make_cstring_view("element_type"), + &element_type_node)); + iree_hal_element_type_t element_type = IREE_HAL_ELEMENT_TYPE_NONE; + IREE_RETURN_IF_ERROR(iree_trace_replay_parse_hal_element_type( + replay, document, element_type_node, &element_type)); + + yaml_node_t* encoding_type_node = NULL; + IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find( + document, value_node, iree_make_cstring_view("encoding_type"), + &encoding_type_node)); + iree_hal_encoding_type_t encoding_type = + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR; + IREE_RETURN_IF_ERROR(iree_trace_replay_parse_hal_encoding_type( + replay, document, encoding_type_node, &encoding_type)); + + iree_device_size_t allocation_size = 0; + IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_size( + shape, shape_rank, element_type, encoding_type, &allocation_size)); + + iree_hal_buffer_t* buffer = NULL; + IREE_RETURN_IF_ERROR(iree_hal_allocator_allocate_buffer( + iree_hal_device_allocator(replay->device), + IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, + IREE_HAL_BUFFER_USAGE_ALL, allocation_size, iree_const_byte_span_empty(), + &buffer)); + + iree_vm_ref_t buffer_ref = iree_hal_buffer_move_ref(buffer); + iree_status_t status = iree_vm_list_push_ref_move(target_list, &buffer_ref); + iree_vm_ref_release(&buffer_ref); + return status; +} + // Parses a !hal.buffer_view and appends it to |target_list|. // // ```yaml @@ -748,10 +800,31 @@ return status; } +// Parses a !hal.buffer in tensor form and appends it to |target_list|. +// The tensor form is used to size and initialize the buffer but then the +// metadata is thrown away. +// +// ```yaml +// !!hal.buffer 4xf32=[0 1 2 3] +// ``` +static iree_status_t iree_trace_replay_parse_inline_hal_buffer( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* value_node, iree_vm_list_t* target_list) { + iree_hal_buffer_view_t* buffer_view = NULL; + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_parse( + iree_yaml_node_as_string(value_node), + iree_hal_device_allocator(replay->device), &buffer_view)); + iree_vm_ref_t buffer_ref = + iree_hal_buffer_retain_ref(iree_hal_buffer_view_buffer(buffer_view)); + iree_status_t status = iree_vm_list_push_ref_move(target_list, &buffer_ref); + iree_hal_buffer_view_release(buffer_view); + return status; +} + // Parses a !hal.buffer_view in tensor form and appends it to |target_list|. // // ```yaml -// !tensor 4xf32=[0 1 2 3] +// !hal.buffer_view 4xf32=[0 1 2 3] // ``` static iree_status_t iree_trace_replay_parse_inline_hal_buffer_view( iree_trace_replay_t* replay, yaml_document_t* document, @@ -761,10 +834,7 @@ iree_yaml_node_as_string(value_node), iree_hal_device_allocator(replay->device), &buffer_view)); iree_vm_ref_t buffer_view_ref = iree_hal_buffer_view_move_ref(buffer_view); - iree_status_t status = - iree_vm_list_push_ref_move(target_list, &buffer_view_ref); - iree_vm_ref_release(&buffer_view_ref); - return status; + return iree_vm_list_push_ref_move(target_list, &buffer_view_ref); } // Parses a typed item from |value_node| and appends it to |target_list|. @@ -783,7 +853,10 @@ yaml_document_t* document, yaml_node_t* value_node, iree_vm_list_t* target_list) { - if (strcmp(value_node->tag, "!hal.buffer_view") == 0) { + if (strcmp(value_node->tag, "!hal.buffer") == 0) { + return iree_trace_replay_parse_inline_hal_buffer(replay, document, + value_node, target_list); + } else if (strcmp(value_node->tag, "!hal.buffer_view") == 0) { return iree_trace_replay_parse_inline_hal_buffer_view( replay, document, value_node, target_list); } @@ -801,6 +874,10 @@ } else if (iree_string_view_equal(type, iree_make_cstring_view("vm.list"))) { return iree_trace_replay_parse_vm_list(replay, document, value_node, target_list); + } else if (iree_string_view_equal(type, + iree_make_cstring_view("hal.buffer"))) { + return iree_trace_replay_parse_hal_buffer(replay, document, value_node, + target_list); } else if (iree_string_view_equal( type, iree_make_cstring_view("hal.buffer_view"))) { return iree_trace_replay_parse_hal_buffer_view(replay, document, value_node,
diff --git a/iree/tools/utils/vm_util.cc b/iree/tools/utils/vm_util.cc index f668aee..39e1331 100644 --- a/iree/tools/utils/vm_util.cc +++ b/iree/tools/utils/vm_util.cc
@@ -84,14 +84,27 @@ bool has_x = iree_string_view_find_char(input_view, 'x', 0) != IREE_STRING_VIEW_NPOS; if (has_equal || has_x) { - // Buffer view (either just a shape or a shape=value). + // Buffer view (either just a shape or a shape=value) or buffer. + bool is_storage_reference = iree_string_view_consume_prefix( + &input_view, iree_make_cstring_view("&")); iree_hal_buffer_view_t* buffer_view = nullptr; IREE_RETURN_IF_ERROR( iree_hal_buffer_view_parse(input_view, allocator, &buffer_view), "parsing value '%.*s'", (int)input_view.size, input_view.data); - auto buffer_view_ref = iree_hal_buffer_view_move_ref(buffer_view); - IREE_RETURN_IF_ERROR( - iree_vm_list_push_ref_move(variant_list.get(), &buffer_view_ref)); + if (is_storage_reference) { + // Storage buffer reference; just take the storage for the buffer view - + // it'll still have whatever contents were specified (or 0) but we'll + // discard the metadata. + auto buffer_ref = iree_hal_buffer_retain_ref( + iree_hal_buffer_view_buffer(buffer_view)); + iree_hal_buffer_view_release(buffer_view); + IREE_RETURN_IF_ERROR( + iree_vm_list_push_ref_move(variant_list.get(), &buffer_ref)); + } else { + auto buffer_view_ref = iree_hal_buffer_view_move_ref(buffer_view); + IREE_RETURN_IF_ERROR( + iree_vm_list_push_ref_move(variant_list.get(), &buffer_view_ref)); + } } else { // Scalar. bool has_dot = iree_string_view_find_char(input_view, '.', 0) !=
diff --git a/iree/tools/utils/vm_util_test.cc b/iree/tools/utils/vm_util_test.cc index 29df93b..d0477be 100644 --- a/iree/tools/utils/vm_util_test.cc +++ b/iree/tools/utils/vm_util_test.cc
@@ -38,6 +38,18 @@ }; TEST_F(VmUtilTest, ParsePrintBuffer) { + std::string buf_string = "&2x2xi32=[42 43][44 45]"; + vm::ref<iree_vm_list_t> variant_list; + IREE_ASSERT_OK(ParseToVariantList( + allocator_, std::vector<std::string>{buf_string}, &variant_list)); + std::stringstream os; + IREE_ASSERT_OK(PrintVariantList(variant_list.get(), &os)); + // TODO(benvanik): add a !hal.buffer printer. + EXPECT_EQ(os.str(), + std::string("result[0]: hal.buffer\n") + "(no printer)" + "\n"); +} + +TEST_F(VmUtilTest, ParsePrintBufferView) { std::string buf_string = "2x2xi32=[42 43][44 45]"; vm::ref<iree_vm_list_t> variant_list; IREE_ASSERT_OK(ParseToVariantList( @@ -58,7 +70,7 @@ EXPECT_EQ(os.str(), std::string("result[0]: i32=") + input_string + "\n"); } -TEST_F(VmUtilTest, ParsePrintRank0Buffer) { +TEST_F(VmUtilTest, ParsePrintRank0BufferView) { std::string buf_string = "i32=42"; vm::ref<iree_vm_list_t> variant_list; IREE_ASSERT_OK(ParseToVariantList( @@ -69,7 +81,7 @@ std::string("result[0]: hal.buffer_view\n") + buf_string + "\n"); } -TEST_F(VmUtilTest, ParsePrintMultipleBuffers) { +TEST_F(VmUtilTest, ParsePrintMultipleBufferViews) { std::string buf_string1 = "2x2xi32=[42 43][44 45]"; std::string buf_string2 = "2x3xf64=[1 2 3][4 5 6]"; vm::ref<iree_vm_list_t> variant_list;