[mlir][bufferization] Add OneShotBufferize transform op

[springerm]

Summary

This commit allows for One-Shot Bufferize to be used through the transform dialect. No op handle is currently returned for the bufferized IR.

[ftynse]

Nice, thank you!

I have one question about the targeting model. If the op requires some payload IR module to be passed as target (which also removes it from the mapping rendering it useless for further transformations), I feel like it should also return a handle to the bufferized module, otherwise we will have hard time composing with any transforms after the bufferization.

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

24

Nit: it's not very clear what actually gets bufferized given the target. Isn't bufferization a module-level transform meaning target should always be a module op?

31

39–40

I'm afraid you need to manually implement the MemoryEffectsOpInterface to also indicate a read/write effect on PayloadIRResource.

[springerm]

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

24

It is often applied to an entire module, but it can be applied to arbitrary operations. (Maybe the old bufferization did not allow this.)

39–40

Now that I'm looking at this again, this is actually not as precise as it could be.

If target is a ModuleOp, it survives bufferization. It is modified in-place. (Or to be precise, its regions are modified, but the op stays the same. This guaranteed and should be documented.) This is not necessarily the case for other ops. They may be updated in-place, but most likely they are deleted and replaced with something else.

But I think I cannot express this in TileOp::getEffects. I would have to check the type of target. If it is a ModuleOp, there should be no TransformMappingFree effect.

Adding a return value is difficult because there may not be one for certain ops. Maybe they disappear entirely for some reason. (A bufferization pattern could choose to do this.) If even they do not disappear, I currently do not get a handle to the new bufferized op from the bufferization.

Alternatively, we could also just limit this transform op to ModuleOps. Then we also don't need a return value because the handle can be reused. (Or we just return a new handle to the same ModuleOp.)

What do you think?

[ftynse]

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

39–40

We can also have two different ops or an attribute on this op to switch between the two cases:
(1) bufferizing the entire module, the handle is not freed
(2) bufferizing a piece of a module, the handle is freed (but the handle to the module is not, which lets us pdl_match again after bufferization).

[springerm]

address comments

[springerm]

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

39–40

added an attribute

mlir/include/mlir/Dialect/Bufferization/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(TransformOps)
 add_subdirectory(Transforms)

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

+//===- BufferizationTransformOps.h - Buff. transf. ops ----------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMOPS_BUFFERIZATIONTRANSFORMOPS_H
+#define MLIR_DIALECT_BUFFERIZATION_TRANSFORMOPS_BUFFERIZATIONTRANSFORMOPS_H
+
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/OpImplementation.h"
+
+//===----------------------------------------------------------------------===//
+// Bufferization Transform Operations
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h.inc"
+
+namespace mlir {
+class DialectRegistry;
+
+namespace bufferization {
+void registerTransformDialectExtension(DialectRegistry &registry);
+} // namespace bufferization
+} // namespace mlir
+
+#endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMOPS_BUFFERIZATIONTRANSFORMOPS_H

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

+//===- BufferizationTransformOps.td - Buff. transf. ops ----*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BUFFERIZATION_TRANSFORM_OPS
+#define BUFFERIZATION_TRANSFORM_OPS
+
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformEffects.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.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>]> {
+  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.
+
+    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.
+  }];
+
+  let arguments = (
+      ins PDL_Operation:$target,
+      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 assemblyFormat = "$target attr-dict";
+}
+
+#endif // BUFFERIZATION_TRANSFORM_OPS

mlir/include/mlir/Dialect/Bufferization/TransformOps/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS BufferizationTransformOps.td)
+mlir_tablegen(BufferizationTransformOps.h.inc -gen-op-decls)
+mlir_tablegen(BufferizationTransformOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRBufferizationTransformOpsIncGen)

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
 #include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
registry.insert<acc::OpenACCDialect,
                   sparse_tensor::SparseTensorDialect,
                   tensor::TensorDialect,
                   transform::TransformDialect,
                   tosa::TosaDialect,
                   x86vector::X86VectorDialect>();
   // clang-format on
 
   // Register all dialect extensions.
+  bufferization::registerTransformDialectExtension(registry);
   linalg::registerTransformDialectExtension(registry);
 
   // Register all external models.
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);

mlir/lib/Dialect/Bufferization/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(TransformOps)
 add_subdirectory(Transforms)

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/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/PDL/IR/PDL.h"
+#include "mlir/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+
+using namespace mlir;
+using namespace mlir::bufferization;
+using namespace mlir::transform;
+
+//===----------------------------------------------------------------------===//
+// OneShotBufferizeOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+transform::OneShotBufferizeOp::apply(TransformResults &transformResults,
+                                     TransformState &state) {
+  OneShotBufferizationOptions options;
+  options.allowReturnAllocs = getAllowReturnAllocs();
+  options.allowUnknownOps = getAllowUnknownOps();
+  options.bufferizeFunctionBoundaries = getBufferizeFunctionBoundaries();
+  options.createDeallocs = getCreateDeallocs();
+  options.testAnalysisOnly = getTestAnalysisOnly();
+  options.printConflicts = getPrintConflicts();
+
+  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  for (Operation *target : payloadOps) {
+    auto moduleOp = dyn_cast<ModuleOp>(target);
+    if (getTargetIsModule() && !moduleOp)
+      return emitError("expected ModuleOp target");
+    if (options.bufferizeFunctionBoundaries) {
+      if (!moduleOp)
+        return emitError("expected ModuleOp target");
+      if (failed(bufferization::runOneShotModuleBufferize(moduleOp, options)))
+        return emitError("bufferization failed");
+    } else {
+      if (failed(bufferization::runOneShotBufferize(target, options)))
+        return emitError("bufferization failed");
+    }
+  }
+
+  return success();
+}
+
+void transform::OneShotBufferizeOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  effects.emplace_back(MemoryEffects::Read::get(), getTarget(),
+                       TransformMappingResource::get());
+
+  // Handles that are not modules are not longer usable.
+  if (!getTargetIsModule())
+    effects.emplace_back(MemoryEffects::Free::get(), getTarget(),
+                         TransformMappingResource::get());
+}
+//===----------------------------------------------------------------------===//
+// Transform op registration
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// Registers new ops and declares PDL as dependent dialect since the additional
+/// ops are using PDL types for operands and results.
+class BufferizationTransformDialectExtension
+    : public transform::TransformDialectExtension<
+          BufferizationTransformDialectExtension> {
+public:
+  BufferizationTransformDialectExtension() {
+    declareDependentDialect<bufferization::BufferizationDialect>();
+    declareDependentDialect<pdl::PDLDialect>();
+    declareDependentDialect<memref::MemRefDialect>();
+    registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
+        >();
+  }
+};
+} // namespace
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
+
+void mlir::bufferization::registerTransformDialectExtension(
+    DialectRegistry &registry) {
+  registry.addExtensions<BufferizationTransformDialectExtension>();
+}

mlir/lib/Dialect/Bufferization/TransformOps/CMakeLists.txt

+add_mlir_dialect_library(MLIRBufferizationTransformOps
+  BufferizationTransformOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization/TransformOps
+
+  DEPENDS
+  MLIRBufferizationTransformOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRBufferization
+  MLIRBufferizationTransforms
+  MLIRParser
+  MLIRPDL
+  MLIRSideEffectInterfaces
+  MLIRTransformDialect
+  )

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.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  sequence %arg0 {
+    ^bb0(%arg1: !pdl.operation):
+      %0 = pdl_match @pdl_target in %arg1
+      transform.bufferization.one_shot_bufferize %0
+          {target_is_module = false}
+  }
+
+  pdl.pattern @pdl_target : benefit(1) {
+    %0 = operation "func.func"
+    rewrite %0 with "transform.dialect"
+  }
+}
+
+// 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]]
+  // CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]
+  // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
+  // CHECK: memref.copy %[[A_memref]], %[[alloc]]
+  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
+  // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]
+  %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
+
+  // CHECK: memref.dealloc %[[alloc]]
+  // CHECK: return %[[res_tensor]]
+  return %0 : tensor<?xf32>
+}
+
+// -----
+
+// Test analysis of One-Shot Bufferize only.
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  sequence %arg0 {
+    ^bb0(%arg1: !pdl.operation):
+      %0 = pdl_match @pdl_target in %arg1
+      transform.bufferization.one_shot_bufferize %0
+          {target_is_module = false, test_analysis_only = true}
+  }
+
+  pdl.pattern @pdl_target : benefit(1) {
+    %0 = operation "func.func"
+    rewrite %0 with "transform.dialect"
+  }
+}
+
+// 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.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  sequence %arg0 {
+    ^bb0(%arg1: !pdl.operation):
+      %0 = pdl_match @pdl_target in %arg1
+      // expected-error @+1 {{bufferization failed}}
+      transform.bufferization.one_shot_bufferize %0 {target_is_module = false}
+  }
+
+  pdl.pattern @pdl_target : benefit(1) {
+    %0 = operation "func.func"
+    rewrite %0 with "transform.dialect"
+  }
+}
+
+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.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  sequence %arg0 {
+    ^bb0(%arg1: !pdl.operation):
+      // %arg1 is the module
+      transform.bufferization.one_shot_bufferize %arg1
+  }
+}
+
+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
+
+    // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
+    // CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]
+    // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
+    // CHECK: memref.copy %[[A_memref]], %[[alloc]]
+    // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
+    // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]
+    %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
+
+    // CHECK: memref.dealloc %[[alloc]]
+    // CHECK: return %[[res_tensor]]
+    return %0 : tensor<?xf32>
+  }
+}
 No newline at end of file

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":ArmNeon",
         ":ArmSVE",
         ":ArmSVETransforms",
         ":Async",
         ":AsyncPassIncGen",
         ":AsyncToLLVM",
         ":AsyncTransforms",
         ":BufferizationDialect",
+        ":BufferizationTransformOps",
         ":BufferizationTransforms",
         ":ComplexDialect",
         ":ComplexToLLVM",
         ":ComplexToLibm",
         ":ControlFlowOps",
         ":ConversionPasses",
         ":DLTIDialect",
         ":EmitC",
tbl_outs = [
             "include/mlir/Dialect/Bufferization/IR/BufferizationOpsDialect.cpp.inc",
         ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Bufferization/IR/BufferizationBase.td",
     deps = [":BufferizationOpsTdFiles"],
 )
 
+td_library(
+    name = "BufferizationTransformOpsTdFiles",
+    srcs = [
+        "include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td",
+    ],
+    includes = ["include"],
+    deps = [
+        ":PDLDialectTdFiles",
+        ":TransformDialectTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "BufferizationTransformOpsIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-op-decls"],
+            "include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h.inc",
+        ),
+        (
+            ["-gen-op-defs"],
+            "include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td",
+    deps = [
+        ":BufferizationTransformOpsTdFiles",
+    ],
+)
+
+cc_library(
+    name = "BufferizationTransformOps",
+    srcs = [
+        "lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp",
+    ],
+    hdrs = [
+        "include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h",
+    ],
+    includes = ["include"],
+    deps = [
+        ":BufferizationDialect",
+        ":BufferizationTransformOpsIncGen",
+        ":BufferizationTransforms",
+        ":IR",
+        ":MemRefDialect",
+        ":PDLDialect",
+        ":Parser",
+        ":SideEffectInterfaces",
+        ":TransformDialect",
+        "//llvm:Support",
+    ],
+)
+
 gentbl_cc_library(
     name = "BufferizationOpsIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             ["-gen-op-decls"],
             "include/mlir/Dialect/Bufferization/IR/BufferizationOps.h.inc",
         ),