[mlir][bufferize] Simplify one_shot_bufferize transform op

[springerm]

Summary

Restrict the op to functions and modules. Such ops are modified in-place. The transform now consumes the handle and produces a new handle. The target_is_module attribute is no longer needed because a result handle is produced in either case.

Depends On: D147430

mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.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 OneShotBufferizeOp
     : Op<Transform_Dialect, "bufferization.one_shot_bufferize",
-        [DeclareOpInterfaceMethods<TransformOpInterface>,
-         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+        [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+         DeclareOpInterfaceMethods<TransformOpInterface>]> {
   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
     `one-shot-bufferize` pass. More information can be found in the pass
     documentation.
 
-    If `target_is_module` is set, `target` must be a module. In that case the
-    `target` handle can be reused by other transform ops. When bufferizing other
-    ops, the `target` handled is freed after bufferization and can no longer be
-    used.
+    The targeted ops must be modules or functions. This is because there is
+    always a single, bufferized replacement op for such targets.
 
     Note: Only ops that implement `BufferizableOpInterface` are bufferized. All
     other ops are ignored if `allow_unknown_ops`. If `allow_unknown_ops` is
     unset, this transform fails when an unknown/non-bufferizable op is found.
     Many ops implement `BufferizableOpInterface` via an external model. These
     external models must be registered when applying this transform op;
     otherwise, said ops would be considered non-bufferizable.
+
+    #### Return modes
+
+    This operation consumes the `target` handle and produces the `transformed`
+    handle.
   }];
 
   let arguments = (
-      ins PDL_Operation:$target,
+      ins TransformHandleTypeInterface:$target,
       OptionalAttr<LayoutMapOption>:$function_boundary_type_conversion,
       DefaultValuedAttr<BoolAttr, "false">:$allow_return_allocs,
       DefaultValuedAttr<BoolAttr, "false">:$allow_unknown_ops,
       DefaultValuedAttr<BoolAttr, "false">:$bufferize_function_boundaries,
       DefaultValuedAttr<BoolAttr, "true">:$create_deallocs,
-      DefaultValuedAttr<BoolAttr, "true">:$target_is_module,
       DefaultValuedAttr<BoolAttr, "false">:$test_analysis_only,
       DefaultValuedAttr<BoolAttr, "false">:$print_conflicts);
 
-  let results = (outs);
+  let results = (outs TransformHandleTypeInterface:$transformed);
 
   let assemblyFormat = [{
     (`layout` `{` $function_boundary_type_conversion^ `}`)?
-    $target attr-dict
+    $target attr-dict `:` functional-type($target, results)
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // EmptyTensorToAllocTensorOp
 //===----------------------------------------------------------------------===//
 
 

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

 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.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"
+#include "mlir/IR/FunctionInterfaces.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 using namespace mlir::transform;
 
 //===----------------------------------------------------------------------===//
 // OneShotBufferizeOp
 //===----------------------------------------------------------------------===//
transform::OneShotBufferizeOp::apply(TransformResults &transformResults,
   options.testAnalysisOnly = getTestAnalysisOnly();
   options.printConflicts = getPrintConflicts();
   if (getFunctionBoundaryTypeConversion().has_value())
     options.functionBoundaryTypeConversion =
         *getFunctionBoundaryTypeConversion();
 
   ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
   for (Operation *target : payloadOps) {
+    if (!isa<ModuleOp, FunctionOpInterface>(target))
+      return emitSilenceableError() << "expected module or function target";
     auto moduleOp = dyn_cast<ModuleOp>(target);
-    if (getTargetIsModule() && !moduleOp)
-      return emitSilenceableError() << "expected ModuleOp target";
     if (options.bufferizeFunctionBoundaries) {
       if (!moduleOp)
-        return emitSilenceableError() << "expected ModuleOp target";
+        return emitSilenceableError() << "expected module target";
       if (failed(bufferization::runOneShotModuleBufferize(moduleOp, options)))
         return emitSilenceableError() << "bufferization failed";
     } else {
       if (failed(bufferization::runOneShotBufferize(target, options)))
         return emitSilenceableError() << "bufferization failed";
     }
   }
 
+  // This transform op is currently restricted to ModuleOps and function ops.
+  // Such ops are modified in-place.
+  transformResults.set(getTransformed().cast<OpResult>(), payloadOps);
   return DiagnosedSilenceableFailure::success();
 }
 
-void transform::OneShotBufferizeOp::getEffects(
-    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  // Handles that are not modules are not longer usable.
-  if (!getTargetIsModule()) {
-    consumesHandle(getTarget(), effects);
-  } else {
-    onlyReadsHandle(getTarget(), effects);
-  }
-
-  modifiesPayload(effects);
-}
-
 //===----------------------------------------------------------------------===//
 // EmptyTensorToAllocTensorOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 EmptyTensorToAllocTensorOp::applyToOne(tensor::EmptyOp target,
                                        ApplyToEachResultList &results,
                                        transform::TransformState &state) {

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

 // RUN: mlir-opt --test-transform-dialect-interpreter %s -split-input-file -verify-diagnostics | FileCheck %s
 
 // Test One-Shot Bufferize.
 
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!pdl.operation) -> !pdl.operation
-  transform.bufferization.one_shot_bufferize %0
-      {target_is_module = false}
+  %1 = transform.bufferization.one_shot_bufferize %0 : (!pdl.operation) -> !pdl.operation
 }
 
 // CHECK-LABEL: func @test_function(
 //  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
 func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
   %c0 = arith.constant 0 : index
 
   // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]

 
 // -----
 
 // Test analysis of One-Shot Bufferize only.
 
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!pdl.operation) -> !pdl.operation
-  transform.bufferization.one_shot_bufferize %0
-      {target_is_module = false, test_analysis_only = true}
+  %1 = transform.bufferization.one_shot_bufferize %0
+      {test_analysis_only = true} : (!pdl.operation) -> !pdl.operation
 }
 
 // CHECK-LABEL: func @test_function_analysis(
 //  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
 func.func @test_function_analysis(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
   %c0 = arith.constant 0 : index
   // CHECK: vector.transfer_write
   // CHECK-SAME: {__inplace_operands_attr__ = ["none", "false", "none"]}
   // CHECK-SAME: tensor<?xf32>
   %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
   return %0 : tensor<?xf32>
 }
 
 // -----
 
 // Test One-Shot Bufferize transform failure with an unknown op. This would be
 // allowed with `allow_unknown_ops`.
 
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   // expected-error @+1 {{bufferization failed}}
-  transform.bufferization.one_shot_bufferize %0 {target_is_module = false}
+  %1 = transform.bufferization.one_shot_bufferize %0 : (!pdl.operation) -> !pdl.operation
 }
 
 func.func @test_unknown_op_failure() -> (tensor<?xf32>) {
   // expected-error @+1 {{op was not bufferized}}
   %0 = "test.dummy_op"() : () -> (tensor<?xf32>)
   return %0 : tensor<?xf32>
 }
 
 // -----
 
-// Test One-Shot Bufferize transform failure with a module op.
-
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !pdl.operation):
   // %arg1 is the module
-  transform.bufferization.one_shot_bufferize %arg1
+  %0 = transform.bufferization.one_shot_bufferize %arg1 : (!pdl.operation) -> !pdl.operation
 }
 
 module {
   // CHECK-LABEL: func @test_function(
   //  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
   func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
     %c0 = arith.constant 0 : index
 

 }
 
 // -----
 
 // Test we use identity layout at function boundaries.
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
-  transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %arg1 {
-    target_is_module = true,
-    bufferize_function_boundaries = true }
+  %0 = transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %arg1
+    { bufferize_function_boundaries = true } : (!pdl.operation) -> !pdl.operation
 }
 
 // CHECK: func.func @matmul(
 // CHECK-SAME:  %[[A:.*]]: memref<12x9xf32>,
 // CHECK-SAME:  %[[B:.*]]: memref<9x6xf32>,
 // CHECK-SAME:  %[[C:.*]]: memref<12x6xf32>) -> memref<12x6xf32> {
 func.func @matmul(%A: tensor<12x9xf32>, %B: tensor<9x6xf32>, %C: tensor<12x6xf32>) -> tensor<12x6xf32> {
   // CHECK: linalg.matmul ins(%[[A]], %[[B]] : memref<12x9xf32>, memref<9x6xf32>) outs(%[[C]] : memref<12x6xf32>)

mlir/test/Dialect/LLVM/transform-e2e.mlir

 }
 
 transform.sequence failures(propagate) {
 ^bb1(%module_op: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %module_op : (!pdl.operation) -> !pdl.operation
   %1, %loops:3 = transform.structured.tile %0 [2, 2, 2] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
   %2 = get_closest_isolated_parent %1 : (!pdl.operation) -> !pdl.operation
   transform.structured.vectorize %2
-  transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %module_op
-    {bufferize_function_boundaries = true}
+  %b = transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap}
+      %module_op {bufferize_function_boundaries = true}
+      : (!pdl.operation) -> !pdl.operation
 
-  %f = transform.structured.match ops{["func.func"]} in %module_op 
+  %f = transform.structured.match ops{["func.func"]} in %b
     : (!pdl.operation) -> !pdl.operation
 
   // TODO: group these lower-level controls into various properly named vector
   // lowering TD macros.
   %func = transform.vector.lower_contraction %f
     lowering_strategy = "outerproduct" 
       : (!pdl.operation) -> !pdl.operation
 

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

 }
 
 transform.sequence failures(propagate) {
 ^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
   // Make sure that One-Shot Bufferize can bufferize the rest.
-  transform.bufferization.one_shot_bufferize %arg1
+  %3 = transform.bufferization.one_shot_bufferize %arg1 : (!pdl.operation) -> !pdl.operation
 }
 
 // -----
 
 // CHECK-LABEL: func @materialization_of_bbarg(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
 //       CHECK:   %[[c0:.*]] = arith.constant 0 : index
 //       CHECK:   %[[dim:.*]] = tensor.dim %[[t]], %[[c0]]

 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["tensor.extract"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   %1 = test_produce_value_handle_to_argument_of_parent_block %0, 0 : (!pdl.operation) -> !transform.any_value
   %2 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
   // Make sure that One-Shot Bufferize can bufferize the rest.
-  transform.bufferization.one_shot_bufferize %arg1
+  %3 = transform.bufferization.one_shot_bufferize %arg1 : (!pdl.operation) -> !pdl.operation
 }
 
 // -----
 
 // CHECK-LABEL: func @materialization_of_opresult(
 //       CHECK:   %[[t:.*]] = "dummy.some_op"
 //       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4>
 //       CHECK:   memref.tensor_store %[[t]], %[[alloc]]

mlir/test/Dialect/Vector/transform-vector.mlir

 
 transform.sequence failures(propagate) {
 ^bb1(%module_op: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %module_op : (!pdl.operation) -> !pdl.operation
   %1, %loops:3 = transform.structured.tile %0 [8, 4, 2] 
     : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
   %2 = get_closest_isolated_parent %1 : (!pdl.operation) -> !pdl.operation
   transform.structured.vectorize %2
-  transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %module_op 
+  %b = transform.bufferization.one_shot_bufferize
+      layout{IdentityLayoutMap} %module_op
       {bufferize_function_boundaries = true, allow_return_allocs = true}
+      : (!pdl.operation) -> !pdl.operation
 
-  %f = transform.structured.match ops{["func.func"]} in %module_op 
+  %f = transform.structured.match ops{["func.func"]} in %b
     : (!pdl.operation) -> !pdl.operation
 
   // TODO: group these lower-level controls into various properly named vector
   // lowering TD macros.
   %func = transform.vector.lower_contraction %f
-    lowering_strategy = "outerproduct" 
+    lowering_strategy = "outerproduct"
       : (!pdl.operation) -> !pdl.operation
 
   %func_2 = transform.vector.apply_transfer_permutation_patterns %func
       : (!pdl.operation) -> !pdl.operation
 
   %func_3 = transform.vector.lower_multi_reduction %func_2
     lowering_strategy = "innerparallel"
       : (!pdl.operation) -> !pdl.operation