[mlir][bufferization] Add bufferization.eliminate_empty_tensors transform op

[springerm]

Summary

Also add a test case that shows how to efficiently bufferize a tensor.pad op to an allocation by wrapping it in a bufferization.alloc_tensor op.

Depends On: D144409

[springerm]

rebase

[nicolasvasilache]

can this land independently of its parents in the stack ?

mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td

 include "mlir/Dialect/Bufferization/IR/BufferizationEnums.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 include "mlir/Dialect/Transform/IR/TransformTypes.td"
 include "mlir/Dialect/PDL/IR/PDLTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
+def Transform_EmptyOp : Transform_ConcreteOpType<"tensor.empty">;
+def Transform_AllocTensorOp : Transform_ConcreteOpType<"bufferization.alloc_tensor">;
+
+//===----------------------------------------------------------------------===//
+// OneShotBufferizeOp
+//===----------------------------------------------------------------------===//
+
 def OneShotBufferizeOp
     : Op<Transform_Dialect, "bufferization.one_shot_bufferize",
         [DeclareOpInterfaceMethods<TransformOpInterface>,
          DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let description = [{
     Indicates that the given `target` op should be bufferized with One-Shot
     Bufferize. The bufferization can be configured with various attributes that
     corresponding to options in `BufferizationOptions` and the
def OneShotBufferizeOp
 
   let assemblyFormat = [{
     (`layout` `{` $function_boundary_type_conversion^ `}`)?
     $target attr-dict
   }];
 }
 
 //===----------------------------------------------------------------------===//
-// EmptyTensorToAllocTensorOp
+// EliminateEmptyTensorsOp
 //===----------------------------------------------------------------------===//
 
+def EliminateEmptyTensorsOp
+    : Op<Transform_Dialect, "bufferization.eliminate_empty_tensors",
+        [DeclareOpInterfaceMethods<TransformOpInterface>,
+         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let description = [{
+    Try to eliminate all `tensor.empty` ops within the targeted op by replacing
+    them with a destination tensor.
 
-def Transform_EmptyOp : Transform_ConcreteOpType<"tensor.empty">;
-def Transform_AllocTensorOp : Transform_ConcreteOpType<"bufferization.alloc_tensor">;
+    `tensor.empty` ops cannot be bufferizes. They can either be converted to
+    `bufferization.alloc_tensor` or replaced with another tensor (via this
+    transform). `tensor.empty` does not specify the contents of the returned
+    tensor so their results can be replaced with arbitrary tensor values as long
+    as the dimensions match.
+
+    This transform looks for `tensor.empty` ops where the SSA use-def chain of
+    the result ends in a supported "anchor op" (always following the aliasing
+    OpOperand/OpResult chain). Currently supported anchor ops are:
+    - `tensor.insert_slice`
+    - `bufferization.yield` (inside `bufferization.alloc_tensor`)
+
+    Example:
+
+    ```
+    %0 = tensor.empty() : tensor<5xf32>
+    %1 = linalg.fill ... outs(%0)
+    %2 = tensor.insert_slice %1 into %t[1][5][1]
+    ```
+
+    Is rewritten with:
+    ```
+    %0 = tensor.extract_slice %t[1][5][1]
+    %1 = linalg.fill ... outs(%0)
+    %2 = tensor.insert_slice %1 into %t[1][5][1]
+    ```
+
+    The above example can bufferize without an allocation (in the absence of
+    other conflicts) because there is no longer a `tensor.empty` op.
+
+    See `-eliminate-empty-tensors` for more details.
+
+    #### Return modes
+
+    This transform reads the target handle and modifies the payload. It does
+    not produce any handle.
+  }];
+
+  let arguments = (ins PDL_Operation:$target);
+
+  let results = (outs);
+
+  let assemblyFormat = "$target attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
+// EmptyTensorToAllocTensorOp
+//===----------------------------------------------------------------------===//
 
 def EmptyTensorToAllocTensorOp
     : Op<Transform_Dialect, "bufferization.empty_tensor_to_alloc_tensor",
         [FunctionalStyleTransformOpTrait,
          MemoryEffectsOpInterface,
          TransformOpInterface,
          TransformEachOpTrait]> {
   let description = [{

mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp

 //===- BufferizationTransformOps.h - Bufferization transform ops ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
 
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
void transform::OneShotBufferizeOp::getEffects(
   } else {
     onlyReadsHandle(getTarget(), effects);
   }
 
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
+// EliminateEmptyTensorsOp
+//===----------------------------------------------------------------------===//
+
+void transform::EliminateEmptyTensorsOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsHandle(getTarget(), effects);
+  modifiesPayload(effects);
+}
+
+DiagnosedSilenceableFailure
+transform::EliminateEmptyTensorsOp::apply(TransformResults &transformResults,
+                                          TransformState &state) {
+  IRRewriter rewriter(getContext());
+  OneShotBufferizationOptions options;
+  options.allowReturnAllocs = true;
+
+  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  for (Operation *target : payloadOps) {
+    OneShotAnalysisState state(target, options);
+    if (failed(analyzeOp(target, state)))
+      return mlir::emitSilenceableFailure(target->getLoc())
+             << "failed to analyze op";
+    if (failed(bufferization::insertSliceAnchoredEmptyTensorEliminationStep(
+            rewriter, target, state)))
+      return mlir::emitSilenceableFailure(target->getLoc())
+             << "failed to eliminate insert_slice anchored tensor.empty ops";
+    if (failed(bufferization::allocTensorAnchoredEmptyTensorEliminationStep(
+            rewriter, target, state)))
+      return mlir::emitSilenceableFailure(target->getLoc())
+             << "failed to eliminate alloc_tensor anchored tensor.empty ops";
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
 // EmptyTensorToAllocTensorOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 EmptyTensorToAllocTensorOp::applyToOne(tensor::EmptyOp target,
                                        ApplyToEachResultList &results,
                                        transform::TransformState &state) {
   IRRewriter rewriter(target->getContext());

mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir

 }
 
 // Expect `bufferization.empty_tensor_to_alloc_tensor` to replace the tensor.empty.
 func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {
   // CHECK: bufferization.alloc_tensor
   %0 = tensor.empty() : tensor<2x2xf32>
   return %0 : tensor<2x2xf32>
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  transform.bufferization.eliminate_empty_tensors %0
+}
+
+// CHECK-LABEL: func @empty_tensor_elimination(
+//       CHECK:   tensor.extract_slice
+//       CHECK:   linalg.fill
+//       CHECK:   tensor.insert_slice
+func.func @empty_tensor_elimination(
+    %t: tensor<10xf32>, %f: f32) -> tensor<10xf32> {
+  %0 = tensor.empty() : tensor<5xf32>
+  %1 = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
+  %2 = tensor.insert_slice %1 into %t [1][5][1]
+      : tensor<5xf32> into tensor<10xf32>
+  return %2 : tensor<10xf32>
+}

mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir

 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   %1 = transform.get_result %0[0] : (!pdl.operation) -> !transform.any_value
   %2 = transform.structured.bufferize_to_allocation %1
 }
 
 // -----
 
+// CHECK-LABEL: func @tensor_pad_wrapped_and_eliminated(
+//       CHECK:   bufferization.alloc_tensor({{.*}}) init {
+//  CHECK-NEXT:   ^{{.*}}(%[[bbarg:.*]]: tensor<?x?xindex>):
+//       CHECK:     %[[filled:.*]] = linalg.fill {{.*}} outs(%[[bbarg]] : tensor<?x?xindex>)
+//       CHECK:     %[[inserted:.*]] = tensor.insert_slice %{{.*}} into %[[filled]]
+//       CHECK:     bufferization.yield %[[inserted]]
+func.func @tensor_pad_wrapped_and_eliminated(
+    %t: tensor<?x10xindex>, %l2: index, %h1: index, %h2: index) -> tensor<?x?xindex> {
+  %0 = tensor.pad %t low[5, %l2] high[%h1, %h2] {
+  ^bb0(%arg0: index, %arg1: index):
+    %c = arith.constant 50 : index
+    tensor.yield %c : index
+  } : tensor<?x10xindex> to tensor<?x?xindex>
+  return %0 : tensor<?x?xindex>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %p0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!pdl.operation) -> !transform.any_op
+  %1 = transform.get_result %p0[0] : (!transform.any_op) -> !transform.any_value
+  // Wrap the tensor.pad op inside a new bufferization.alloc_tensor op.
+  %2 = transform.structured.bufferize_to_allocation %1
+  %alloc_tensor_op = transform.get_defining_op %2 : (!transform.any_value) ->!transform.any_op
+
+  %p1 = transform.structured.match ops{["tensor.pad"]} in %alloc_tensor_op : (!transform.any_op) -> !transform.any_op
+  %3 = transform.structured.rewrite_in_destination_passing_style %p1 : (!transform.any_op) -> !transform.any_op
+  transform.bufferization.eliminate_empty_tensors %arg1
+}
+
+// -----
+
 // CHECK-LABEL: func @materialization_of_bbarg(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
 //       CHECK:   %[[c0:.*]] = arith.constant 0 : index
 //       CHECK:   %[[dim:.*]] = tensor.dim %[[t]], %[[c0]]
 //       CHECK:   %[[alloc:.*]] = bufferization.alloc_tensor(%[[dim]]) copy(%[[t]]) {memory_space = 4 : i64} : tensor<?x10xindex>
 //       CHECK:   %[[r:.*]] = tensor.extract %[[alloc]]
 //       CHECK:   return %[[r]]
 func.func @materialization_of_bbarg(%t: tensor<?x10xindex>, %idx: index) -> index {