Replacing hal.tensor.export storage for hal.tensor.alias. (iree-org#1…

…7339) This fixes a design issue in the original `hal.tensor.export` optional storage feature that would lead to the export happening after any `hal.tensor.barrier` ops that may have been used on the source tensor. The new op is intended to be inserted prior to the barriers and can also be inserted elsewhere (not just at ABI boundaries). Minor improvements were required to folding of `stream.async.update` in order to ensure the aliased buffers are used in cases where barriers are present between producers and the alias ops consuming the values. iree-org#17135 made the folder too conservative and would result in all in-place operations of external values getting extra copies. Fixes iree-org#17316.
daveliddell · May 11, 2024 · d2dd9e2 · d2dd9e2
1 parent 5337bd7
commit d2dd9e2
Show file tree

Hide file tree

Showing 17 changed files with 483 additions and 123 deletions.
diff --git a/compiler/plugins/input/Torch/InputConversion/FuncConversion.cpp b/compiler/plugins/input/Torch/InputConversion/FuncConversion.cpp
@@ -258,14 +258,40 @@ LogicalResult ConvertedAsyncFunctionInfo::postProcess() {
   }
 
   // Emit the barrier and exports.
+  // If any of the exports are in-place we need to alias their storage to the
+  // provided buffers.
   Value coarseSignalFence =
       entryBlock->getArgument(entryBlock->getNumArguments() - 1);
   if (barrierInputs.empty()) {
     postambleBuilder.create<IREE::HAL::FenceSignalOp>(funcOp.getLoc(),
                                                       coarseSignalFence);
   } else {
+    SmallVector<Value> aliasedResults;
+    for (auto [barrierInput, meta] :
+         llvm::zip_equal(barrierInputs, barrierResultMeta)) {
+      Value exportStorage;
+      Type torchType;
+      int returnIndex;
+      std::tie(exportStorage, torchType, returnIndex) = meta;
+      if (exportStorage) {
+        // Use the wait fence indicating when the storage is available for
+        // mutation. We need to ensure that no writes are made to the storage
+        // until it indicates it's safe to do so.
+        auto waitSignalFences = getEnclosingWaitSignalFences(exportStorage);
+        assert(waitSignalFences && "async function missing fences");
+        Value waitFence = waitSignalFences->first;
+        auto barrierInputDims = IREE::Util::buildDynamicDimsForValue(
+            barrierInput.getLoc(), barrierInput, postambleBuilder);
+        aliasedResults.push_back(
+            postambleBuilder.create<IREE::HAL::TensorAliasOp>(
+                barrierInput.getLoc(), barrierInput.getType(), barrierInput,
+                barrierInputDims, exportStorage, waitFence));
+      } else {
+        aliasedResults.push_back(barrierInput);
+      }
+    }
     auto barrierOp = postambleBuilder.create<IREE::HAL::TensorBarrierOp>(
-        funcOp.getLoc(), barrierInputs, coarseSignalFence);
+        funcOp.getLoc(), aliasedResults, coarseSignalFence);
     for (auto [barrierResult, meta] :
          llvm::zip_equal(barrierOp.getResults(), barrierResultMeta)) {
       Value exportStorage;
@@ -275,13 +301,9 @@ LogicalResult ConvertedAsyncFunctionInfo::postProcess() {
       Value exportedValue = postambleBuilder.create<IREE::HAL::TensorExportOp>(
           funcOp.getLoc(),
           postambleBuilder.getType<IREE::HAL::BufferViewType>(), barrierResult,
-          TypeAttr::get(barrierResult.getType()), exportStorage, StringAttr());
+          TypeAttr::get(barrierResult.getType()), StringAttr());
       if (returnIndex >= 0) {
         newReturnOperands[returnIndex] = exportedValue;
-      } else {
-        // Don't drop it.
-        postambleBuilder.create<IREE::Util::OptimizationBarrierOp>(
-            funcOp.getLoc(), exportedValue);
       }
     }
   }

diff --git a/compiler/plugins/input/Torch/InputConversion/test/func_conversion.mlir b/compiler/plugins/input/Torch/InputConversion/test/func_conversion.mlir
@@ -73,11 +73,11 @@ func.func @main(%arg0: !torch.vtensor<[4,5],si32>) -> !torch.vtensor<[4,5],si32>
 //   CHECK-DAG: %[[TORCH_RESULT1:.+]] = torch.operator "mutate_inplace"(%[[TORCH_ARG1]])
 //   CHECK-DAG: %[[TENSOR_ARG0:.+]] = torch_c.to_builtin_tensor %[[TORCH_RESULT0]]
 //   CHECK-DAG: %[[TENSOR_ARG1:.+]] = torch_c.to_builtin_tensor %[[TORCH_RESULT1]]
-//       CHECK: %[[BARRIER_RESULTS:.+]]:2 = hal.tensor.barrier join(%[[TENSOR_ARG1]], %[[TENSOR_ARG0]] : tensor<5x4xf32>, tensor<4x5xi32>) => %arg3 : !hal.fence
-//   CHECK-DAG: %[[EXPORT_RESULT1:.+]] = hal.tensor.export %[[BARRIER_RESULTS]]#0 into(%arg1 : !hal.buffer_view)
-//   CHECK-DAG: %[[UNUSED:.+]] = util.optimization_barrier %[[EXPORT_RESULT1]]
-//   CHECK-DAG: %[[EXPORT_RESULT0:.+]] = hal.tensor.export %[[BARRIER_RESULTS]]#1 :
-//       CHECK: util.return %[[EXPORT_RESULT0]]
+//       CHECK: %[[EXPORT_ALIAS1:.+]] = hal.tensor.alias wait(%arg2) => %[[TENSOR_ARG1]] : tensor<5x4xf32> to %arg1 : !hal.buffer_view
+//       CHECK: %[[BARRIER_RESULTS:.+]]:2 = hal.tensor.barrier join(%[[EXPORT_ALIAS1]], %[[TENSOR_ARG0]] : tensor<5x4xf32>, tensor<4x5xi32>) => %arg3 : !hal.fence
+//   CHECK-DAG: %[[EXPORT_RESULT0:.+]] = hal.tensor.export %[[BARRIER_RESULTS]]#0
+//   CHECK-DAG: %[[EXPORT_RESULT1:.+]] = hal.tensor.export %[[BARRIER_RESULTS]]#1
+//       CHECK: util.return %[[EXPORT_RESULT1]]
 builtin.module @mutable_input_overwrite_no_return {
 func.func @main(%arg0: !torch.vtensor<[4,5],si32>, %arg1: !torch.tensor<[5,4],f32>)
     -> (!torch.vtensor<[4,5],si32>) {
@@ -97,9 +97,10 @@ func.func @main(%arg0: !torch.vtensor<[4,5],si32>, %arg1: !torch.tensor<[5,4],f3
 // Not a good idea to do but legal. This verifies that if returning a mutated
 // tensor's intermediate value, you will get two exports, indicating a copy.
 // CHECK-LABEL: @mutable_input_overwrite_return_alias_copies
-//       CHECK: %[[BARRIER_RESULTS:.+]]:2 = hal.tensor.barrier join(%{{.*}}, %{{.*}} : tensor<5x4xf32>, tensor<5x4xf32>)
-//   CHECK-DAG: = hal.tensor.export %[[BARRIER_RESULTS]]#0 into(%arg0 : !hal.buffer_view)
-//   CHECK-DAG: = hal.tensor.export %[[BARRIER_RESULTS]]#1 :
+//       CHECK: %[[ALIASED:.+]] = hal.tensor.alias wait({{.+}}) => %{{.+}} : tensor<5x4xf32> to %arg0 : !hal.buffer_view
+//       CHECK: %[[BARRIER_RESULTS:.+]]:2 = hal.tensor.barrier join(%[[ALIASED]], %{{.*}} : tensor<5x4xf32>, tensor<5x4xf32>)
+//   CHECK-DAG: = hal.tensor.export %[[BARRIER_RESULTS]]#0
+//   CHECK-DAG: = hal.tensor.export %[[BARRIER_RESULTS]]#1
 builtin.module @mutable_input_overwrite_return_alias_copies {
 func.func @main(%arg0: !torch.tensor<[5,4],f32>) -> (!torch.vtensor<[5,4],f32>) {
   %0 = torch.copy.to_vtensor %arg0 : !torch.vtensor<[5,4],f32>

diff --git a/compiler/src/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp b/compiler/src/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp
@@ -560,6 +560,20 @@ createExportWrapperFunc(IREE::ABI::InvocationModel invocationModel,
                                                         exportOp, arguments);
   auto asyncResults = llvm::to_vector(callOp.getResults());
 
+  // Alias results to storage buffers if provided.
+  for (unsigned resultIndex = 0; resultIndex < asyncResults.size();
+       ++resultIndex) {
+    if (!resultStorages[resultIndex])
+      continue;
+    auto source = asyncResults[resultIndex];
+    auto sourceDims = IREE::Util::buildDynamicDimsForValue(
+        exportOp.getLoc(), source, entryBuilder);
+    auto aliasOp = entryBuilder.create<IREE::HAL::TensorAliasOp>(
+        exportOp.getLoc(), source.getType(), source, sourceDims,
+        resultStorages[resultIndex], waitFence);
+    asyncResults[resultIndex] = cast<OpResult>(aliasOp.getResult());
+  }
+
   // Insert a barrier if requested - all tensors will be calculated and the
   // fence will be signaled. Note that even if there are no tensor results we
   // need to signal the fence.
@@ -591,11 +605,9 @@ createExportWrapperFunc(IREE::ABI::InvocationModel invocationModel,
           resultIndex, "iree.abi.encoding");
       auto dynamicDims = IREE::Util::buildDynamicDimsForValue(
           result.getLoc(), result, entryBuilder);
-      auto resultStorage = resultStorages[resultIndex];
       results.push_back(entryBuilder.create<IREE::HAL::TensorExportOp>(
           result.getLoc(), newType, result,
           encoding ? encoding : TypeAttr::get(result.getType()), dynamicDims,
-          resultStorage,
           inferResultName(entryBuilder.getContext(), resultIndex,
                           exportOp.getResultAttrDict(resultIndex))));
     } else {

diff --git a/compiler/src/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points.mlir b/compiler/src/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points.mlir
@@ -105,10 +105,11 @@ util.func public @exportEncodings(%arg0: tensor<?x8x8x3xf32> {iree.abi.encoding
 //  CHECK-NEXT:   %[[ARG0_DIM0:.+]] = hal.buffer_view.dim<%[[ARG0]] : !hal.buffer_view>[0] : index
 //  CHECK-NEXT:   %[[ARG0_TENSOR:.+]] = hal.tensor.import %[[ARG0]] "input0" : !hal.buffer_view -> tensor<?x8x8x3xf32>{%[[ARG0_DIM0]]}
 //  CHECK-NEXT:   %[[RET_TENSORS:.+]]:2 = util.call @_outputStorage(%[[ARG0_TENSOR]], %[[RET1_STORAGE]])
-//       CHECK:   %[[RET0_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#0, %c0{{.*}} : tensor<?x8x8x3xf32>
+//   CHECK-DAG:   %[[RET1_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#1, %c0{{.*}} : tensor<?x8x8x3xf32>
+//   CHECK-DAG:   %[[RET1_ALIAS:.+]] = hal.tensor.alias %[[RET_TENSORS]]#1 : tensor<?x8x8x3xf32>{%[[RET1_DIM0]]} to %[[RET1_STORAGE]] : !hal.buffer
+//   CHECK-DAG:   %[[RET0_DIM0:.+]] = tensor.dim %[[RET_TENSORS]]#0, %c0{{.*}} : tensor<?x8x8x3xf32>
 //  CHECK-NEXT:   %[[RET0_VIEW:.+]] = hal.tensor.export %[[RET_TENSORS]]#0 "output0" : 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 "output1" into(%[[RET1_STORAGE]] : !hal.buffer) : tensor<?x8x8x3xf32>{%[[RET1_DIM0]]} -> !hal.buffer_view
+//  CHECK-NEXT:   %[[RET1_VIEW:.+]] = hal.tensor.export %[[RET1_ALIAS]] "output1" : tensor<?x8x8x3xf32>{%[[RET1_DIM0]]} -> !hal.buffer_view
 //  CHECK-NEXT:   util.return %[[RET0_VIEW]], %[[RET1_VIEW]] : !hal.buffer_view, !hal.buffer_view
 //  CHECK-NEXT: }
 

diff --git a/...er/src/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points_coarse_fences.mlir b/...er/src/iree/compiler/Bindings/Native/Transforms/test/wrap_entry_points_coarse_fences.mlir
@@ -102,6 +102,26 @@ util.func public @tensorResultOnly() -> tensor<4xf32> {
 
 // -----
 
+// CHECK-LABEL: util.func public @outputStorage
+//  CHECK-SAME: (%[[ARG0:.+]]: !hal.buffer_view, %[[ARG1:.+]]: !hal.buffer_view, %[[RET0:.+]]: !hal.buffer, %[[RET1:.+]]: !hal.buffer,
+//  CHECK-SAME:  %[[WAIT:.+]]: !hal.fence, %[[SIGNAL:.+]]: !hal.fence)
+//       CHECK: %[[RESULT_TENSORS:.+]]:2 = util.call @_outputStorage
+//   CHECK-DAG: %[[RESULT_ALIAS0:.+]] = hal.tensor.alias wait(%[[WAIT]]) => %[[RESULT_TENSORS]]#0 : tensor<4xf32> to %[[RET0]] : !hal.buffer
+//   CHECK-DAG: %[[RESULT_ALIAS1:.+]] = hal.tensor.alias wait(%[[WAIT]]) => %[[RESULT_TENSORS]]#1 : tensor<4xf32> to %[[RET1]] : !hal.buffer
+//   CHECK-DAG: %[[READY_RESULTS:.+]]:2 = hal.tensor.barrier join(%[[RESULT_ALIAS0]], %[[RESULT_ALIAS1]] : tensor<4xf32>, tensor<4xf32>) => %[[SIGNAL]] : !hal.fence
+//   CHECK-DAG: %[[EXPORT0:.+]] = hal.tensor.export %[[READY_RESULTS]]#0 "output0"
+//   CHECK-DAG: %[[EXPORT1:.+]] = hal.tensor.export %[[READY_RESULTS]]#1 "output1"
+//  CHECK-NEXT: util.return %[[EXPORT0]], %[[EXPORT1]]
+
+// CHECK-LABEL: util.func private @_outputStorage(
+util.func public @outputStorage(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>, %ret0: !hal.buffer {iree.abi.output = 0 : index}, %ret1: !hal.buffer {iree.abi.output = 1 : index}) -> (tensor<4xf32>, tensor<4xf32>) {
+  %0 = arith.addf %arg0, %arg1 : tensor<4xf32>
+  %1 = arith.addf %0, %arg0 : tensor<4xf32>
+  util.return %0, %1 : tensor<4xf32>, tensor<4xf32>
+}
+
+// -----
+
 // Tests that imported functions with the coarse-fences execution model
 // specified get wrapped with fences. Note that unlike exports controlled by
 // compiler flags imports only get the fences when explicitly specified so as

diff --git a/compiler/src/iree/compiler/Bindings/TFLite/Transforms/WrapEntryPoints.cpp b/compiler/src/iree/compiler/Bindings/TFLite/Transforms/WrapEntryPoints.cpp
@@ -539,7 +539,7 @@ class WrapEntryPointsPass
       }
       callResults.push_back(entryBuilder.create<IREE::HAL::TensorExportOp>(
           result.getLoc(), bufferType, result, outputDynamicDims.tensorType,
-          dynamicDims, /*target_storage=*/nullptr, /*name=*/nullptr));
+          dynamicDims, /*name=*/nullptr));
       for (auto [dynamicDim, globalOp] :
            llvm::zip_equal(dynamicDims, outputDynamicDims.globalOps)) {
         globalOp.createStoreOp(result.getLoc(), dynamicDim, entryBuilder);

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/ExportBenchmarkFuncs.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/ExportBenchmarkFuncs.cpp
@@ -147,19 +147,24 @@ static IREE::Util::GlobalOp createExportBufferGlobalOp(std::string name,
                                                        Explorer &explorer) {
   auto loc = arg.getLoc();
 
-  // Find a hal.tensor.export user.
-  IREE::HAL::TensorExportOp exportOp;
+  // Find a hal.tensor.export or alias user and extract the encoding.
+  Type sourceType;
   if (explorer.walkTransitiveUsers(arg, [&](Operation *op) -> WalkResult {
-        exportOp = dyn_cast<IREE::HAL::TensorExportOp>(op);
-        return exportOp ? WalkResult::interrupt() : WalkResult::advance();
+        if (auto aliasOp = dyn_cast<IREE::HAL::TensorAliasOp>(op)) {
+          sourceType = aliasOp.getResult().getType();
+          return WalkResult::interrupt();
+        } else if (auto exportOp = dyn_cast<IREE::HAL::TensorExportOp>(op)) {
+          sourceType = exportOp.getSourceEncoding();
+          return WalkResult::interrupt();
+        }
+        return WalkResult::advance();
       }) == TraversalResult::INCOMPLETE) {
     // Analysis failed to find an export op. User needs to rework their program.
     mlir::emitError(loc) << "unsupported dynamic buffer view export on " << arg;
     return {};
   }
 
-  // Extract the type, which must be a static tensor.
-  auto sourceType = exportOp.getSourceEncoding();
+  // The type must be a static tensor for this pass to work.
   auto tensorType = llvm::dyn_cast<RankedTensorType>(sourceType);
   if (!tensorType || !tensorType.hasStaticShape()) {
     mlir::emitError(loc) << "unsupported buffer view export tensor type on "

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/export_benchmark_funcs.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/export_benchmark_funcs.mlir
@@ -95,8 +95,9 @@ util.func public @importDynamicBufferView(%view: !hal.buffer_view) -> !hal.buffe
 util.func public @exportBufferViewInPlace(%view: !hal.buffer_view, %storage: !hal.buffer) -> !hal.buffer_view {
   %0 = hal.tensor.import %view : !hal.buffer_view -> tensor<4xi32>
   %1 = arith.muli %0, %0 : tensor<4xi32>
-  %2 = hal.tensor.export %1 into(%storage : !hal.buffer) : tensor<4xi32> -> !hal.buffer_view
-  util.return %2 : !hal.buffer_view
+  %2 = hal.tensor.alias %1 : tensor<4xi32> to %storage : !hal.buffer
+  %3 = hal.tensor.export %2 : tensor<4xi32> -> !hal.buffer_view
+  util.return %3 : !hal.buffer_view
 }
 
 //      CHECK: util.global private @[[GLOBAL_ARG0:.+]] {

diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.cpp b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.cpp
@@ -474,18 +474,9 @@ LogicalResult TensorImportOp::verify() {
 void TensorExportOp::build(OpBuilder &builder, OperationState &result,
                            Type resultType, Value source,
                            TypeAttr sourceEncoding, StringAttr name) {
-  build(builder, result, resultType, source, sourceEncoding,
-        /*targetStorage=*/nullptr, name);
-}
-
-void TensorExportOp::build(OpBuilder &builder, OperationState &result,
-                           Type resultType, Value source,
-                           TypeAttr sourceEncoding, Value targetStorage,
-                           StringAttr name) {
   auto dynamicDims =
       IREE::Util::buildDynamicDimsForValue(result.location, source, builder);
-  build(builder, result, resultType, source, sourceEncoding, dynamicDims,
-        targetStorage, name);
+  build(builder, result, resultType, source, sourceEncoding, dynamicDims, name);
 }
 
 Value TensorExportOp::getTiedResult(unsigned resultIndex) {
@@ -512,6 +503,33 @@ LogicalResult TensorExportOp::verify() {
                                         op.getSource().getType());
 }
 
+//===----------------------------------------------------------------------===//
+// hal.tensor.alias
+//===----------------------------------------------------------------------===//
+
+Value TensorAliasOp::getTiedResult(unsigned resultIndex) {
+  return IREE::Util::TiedOpInterface::findTiedBaseValue(getSource());
+}
+
+::std::optional<unsigned>
+TensorAliasOp::getTiedResultOperandIndex(unsigned resultIndex) {
+  return {0}; // source
+}
+
+SmallVector<int64_t> TensorAliasOp::getTiedResultOperandIndices() {
+  return {0}; // source
+}
+
+LogicalResult TensorAliasOp::verify() {
+  TensorAliasOp op = *this;
+  auto type = llvm::cast<TensorType>(op.getSource().getType());
+  if (type.getNumDynamicDims() != op.getSourceDims().size()) {
+    return op->emitOpError()
+           << "number of dynamic dims must match the operand type";
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // hal.tensor.barrier
 //===----------------------------------------------------------------------===//

diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
@@ -167,7 +167,6 @@ def HAL_TensorImportOp : HAL_PureOp<"tensor.import", [
 }
 
 def HAL_TensorExportOp : HAL_PureOp<"tensor.export", [
-  AttrSizedOperandSegments,
   DeclareOpInterfaceMethods<Util_TiedOpInterface, [
     "getTiedResult",
     "getTiedResultOperandIndex",
@@ -185,18 +184,12 @@ def HAL_TensorExportOp : HAL_PureOp<"tensor.export", [
     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,
-    Optional<AnyTypeOf<[HAL_Buffer, HAL_BufferView]>>:$target_storage,
     OptionalAttr<StrAttr>:$name
   );
   let results = (outs
@@ -206,7 +199,6 @@ def HAL_TensorExportOp : HAL_PureOp<"tensor.export", [
   let assemblyFormat = [{
     $source
     ($name^)?
-    (`into` `(` $target_storage^ `:` type($target_storage) `)`)?
     `:`
     custom<TypeAlias>($source_encoding, type($source)) (`{` $source_dims^ `}`)?
     `->`
@@ -221,13 +213,6 @@ def HAL_TensorExportOp : HAL_PureOp<"tensor.export", [
       "TypeAttr":$sourceEncoding,
       "StringAttr":$name
     )>,
-    OpBuilder<(ins
-      "Type":$resultType,
-      "Value":$source,
-      "TypeAttr":$sourceEncoding,
-      "Value":$targetStorage,
-      "StringAttr":$name
-    )>,
   ];
 
   let extraClassDeclaration = [{
@@ -240,6 +225,76 @@ def HAL_TensorExportOp : HAL_PureOp<"tensor.export", [
   let hasFolder = 1;
 }
 
+// TODO(#17328): specify an allocation policy to control behavior.
+def HAL_TensorAliasOp : HAL_PureOp<"tensor.alias", [
+  AllTypesMatch<["source", "result"]>,
+  AttrSizedOperandSegments,
+  DeclareOpInterfaceMethods<Util_TiedOpInterface, [
+    "getTiedResult",
+    "getTiedResultOperandIndex",
+    "getTiedResultOperandIndices",
+  ]>,
+  Util_ShapeAwareOp,
+]> {
+  let summary = [{hints that tensor storage should alias a HAL buffer view}];
+  let description = [{
+    Hints that the backing storage of an entire tensor aliases the given storage
+    buffer. There's no guarantee that the storage will alias and instead only
+    that the tensor contents will be written to the storage as if a copy had
+    occurred. This allows the compiler to avoid copies in the ideal case of a
+    producer that is able to produce directly into the target storage but still
+    handle cases where the producer is not able to be in-place.
+
+    The storage buffer provided must have sufficient space for the tensor once
+    encoded. Dynamically shaped tensors may not consume the entire provided
+    storage. If a buffer view is provided the metadata is ignored and only the
+    backing buffer is used.
+
+    An optional wait fence can be provided in cases where the storage is not
+    immediately available. Producers that may alias the storage will wait until
+    the storage is available before updating the contents.
+
+    Explicit aliasing side-steps any analysis that may be performed by the
+    compiler and requires users to guarantee that the safety of the aliasing.
+    Copy-on-write, alias analysis for overlap detection, and ordering via
+    use-def chains are all ignorant of the aliased buffer memory and only ensure
+    the compiler consumes or produces the aliased memory consistent with itself.
+
+    Example:
+    ```mlir
+    %init = tensor.empty
+    %value = linalg.generic ... outs(%init)
+    %aliased = hal.tensor.alias %value : tensor<...> to %buffer : !hal.buffer
+    ... linalg.generic ins(%aliased) ...
+    ```
+  }];
+
+  let arguments = (ins
+    AnyTensor:$source,
+    HAL_ShapeDynamicDims:$source_dims,
+    AnyTypeOf<[HAL_Buffer, HAL_BufferView]>:$storage,
+    Optional<HAL_Fence>:$wait_fence
+  );
+  let results = (outs
+    AnyTensor:$result
+  );
+
+  let assemblyFormat = [{
+    (`wait` `(` $wait_fence^ `)` `=` `` `>`)?
+    $source `:` type($source) (`{` $source_dims^ `}`)?
+    `to`
+    $storage `:` type($storage)
+    attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    ValueRange getOperandDynamicDims(unsigned idx) { return getSourceDims(); }
+    ValueRange getResultDynamicDims(unsigned idx) { return getSourceDims(); }
+  }];
+
+  let hasVerifier = 1;
+}
+
 def HAL_TensorBarrierOp : HAL_Op<"tensor.barrier", [
   AllTypesMatch<["sources", "results"]>,
   DeclareOpInterfaceMethods<Util_TiedOpInterface, [