[MLIR][Transform] Expose map layout option in `OneShotBufferizeOp`

[chelini]

Summary

Expose function-boundary-type-conversion in OneShotBufferizeOp. To
reuse options between passes and transform operations, create a
BufferizationEnums.td.

[springerm]

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

54

Can you make this an enum? Similar to OptionalAttr<MatchInterfaceEnum>:$interface, in Dialect/Linalg/TransformOps/LinalgTransformOps.td (MatchOp)? Then we get better IR verification.

[chelini]

Use Enum

[springerm]

Thanks!

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

 //===- BufferizableOpInterface.h - Bufferizable 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_IR_BUFFERIZABLEOPINTERFACE_H_
 #define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_
 
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/SetVector.h"
 
+#include "mlir/Dialect/Bufferization/IR/BufferizationEnums.h.inc"
+
 namespace mlir {
 class OpBuilder;
 
 namespace bufferization {
 
 class AnalysisState;
 class BufferizableOpInterface;
 struct DialectAnalysisState;
struct BufferizationOptions {
   /// Initializer function for dialect-specific analysis state.
   using DialectStateInitFn =
       std::function<std::unique_ptr<DialectAnalysisState>()>;
   /// Tensor -> MemRef type converter.
   /// Parameters: Value, memory space, bufferization options
   using UnknownTypeConverterFn = std::function<BaseMemRefType(
       Value, unsigned, const BufferizationOptions &)>;
 
-  enum class LayoutMapOption : int8_t {
-    InferLayoutMap = 0,
-    IdentityLayoutMap = 1,
-    FullyDynamicLayoutMap = 2
-  };
-
   BufferizationOptions();
 
   /// Try to cast the given op to BufferizableOpInterface if the op is allow
   /// listed.
   BufferizableOpInterface dynCastBufferizableOp(Operation *op) const;
 
   /// Try to cast the given value to BufferizableOpInterface if the op is allow
   /// listed.

 /// places.
 bool defaultIsRepetitiveRegion(BufferizableOpInterface bufferizableOp,
                                unsigned index);
 } // namespace detail
 
 } // namespace bufferization
 } // namespace mlir
 
+//===----------------------------------------------------------------------===//
+// Bufferization Interfaces
+//===----------------------------------------------------------------------===//
+
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td

+//===- BufferizationEnums.td - Bufferization enums ---------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the definition file for enums used in Bufferization.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BUFFERIZATION_ENUMS
+#define BUFFERIZATION_ENUMS
+
+include "mlir/IR/EnumAttr.td"
+
+def LayoutMapOption : I32EnumAttr<"LayoutMapOption", 
+                                  "option for map layout", [
+  I32EnumAttrCase<"InferLayoutMap", 0>,
+  I32EnumAttrCase<"IdentityLayoutMap", 1>,
+  I32EnumAttrCase<"FullyDynamicLayoutMap", 2>
+]> {
+  let cppNamespace = "::mlir::bufferization";
+}
+
+#endif // BUFFERIZATION_ENUMS

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

 add_mlir_dialect(BufferizationOps bufferization)
 add_mlir_doc(BufferizationOps BufferizationOps Dialects/ -gen-dialect-doc)
 add_mlir_interface(AllocationOpInterface)
 add_mlir_interface(BufferizableOpInterface)
+
+set(LLVM_TARGET_DEFINITIONS BufferizationEnums.td)
+mlir_tablegen(BufferizationEnums.h.inc -gen-enum-decls)
+mlir_tablegen(BufferizationEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(MLIRBufferizationEnumsIncGen)
+add_dependencies(mlir-headers MLIRBufferizationEnumsIncGen)

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/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/OpImplementation.h"
 
 //===----------------------------------------------------------------------===//
 // Bufferization Transform Operations
 //===----------------------------------------------------------------------===//
 

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/Bufferization/IR/BufferizationEnums.td"
 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
let description = [{
     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,
+      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 assemblyFormat = "$target attr-dict";
+  let assemblyFormat = [{
+    (`layout` `{` $function_boundary_type_conversion^ `}`)?
+    $target attr-dict
+  }];
 }
 
 #endif // BUFFERIZATION_TRANSFORM_OPS

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

 add_mlir_dialect_library(MLIRBufferizationDialect
   AllocationOpInterface.cpp
   BufferizableOpInterface.cpp
   BufferizationOps.cpp
   BufferizationDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
 
   DEPENDS
   MLIRAllocationOpInterfaceIncGen
   MLIRBufferizationOpsIncGen
+  MLIRBufferizationEnumsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRDestinationStyleOpInterface
   MLIRDialect
   MLIRFuncDialect
   MLIRIR
   MLIRSparseTensorDialect
   MLIRTensorDialect
   MLIRMemRefDialect
   )

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

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();
+  if (getFunctionBoundaryTypeConversion().has_value())
+    options.functionBoundaryTypeConversion =
+        *getFunctionBoundaryTypeConversion();
 
   ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
   for (Operation *target : payloadOps) {
     auto moduleOp = dyn_cast<ModuleOp>(target);
     if (getTargetIsModule() && !moduleOp)
       return emitSilenceableError() << "expected ModuleOp target";
     if (options.bufferizeFunctionBoundaries) {
       if (!moduleOp)
#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
         >();
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
 
+#include "mlir/Dialect/Bufferization/IR/BufferizationEnums.cpp.inc"
+
 void mlir::bufferization::registerTransformDialectExtension(
     DialectRegistry &registry) {
   registry.addExtensions<BufferizationTransformDialectExtension>();
 }

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

void runOnOperation() override {
     target.markUnknownOpDynamicallyLegal(
         [&](Operation *op) { return typeConverter.isLegal(op); });
 
     if (failed(applyFullConversion(func, target, std::move(patterns))))
       signalPassFailure();
   }
 };
 
-static BufferizationOptions::LayoutMapOption
-parseLayoutMapOption(const std::string &s) {
+static LayoutMapOption parseLayoutMapOption(const std::string &s) {
   if (s == "fully-dynamic-layout-map")
-    return BufferizationOptions::LayoutMapOption::FullyDynamicLayoutMap;
+    return LayoutMapOption::FullyDynamicLayoutMap;
   if (s == "identity-layout-map")
-    return BufferizationOptions::LayoutMapOption::IdentityLayoutMap;
+    return LayoutMapOption::IdentityLayoutMap;
   if (s == "infer-layout-map")
-    return BufferizationOptions::LayoutMapOption::InferLayoutMap;
+    return LayoutMapOption::InferLayoutMap;
   llvm_unreachable("invalid layout map option");
 }
 
 static OneShotBufferizationOptions::AnalysisHeuristic
 parseHeuristicOption(const std::string &s) {
   if (s == "bottom-up")
     return OneShotBufferizationOptions::AnalysisHeuristic::BottomUp;
   if (s == "top-down")
if (!options) {
           parseLayoutMapOption(functionBoundaryTypeConversion);
       if (mustInferMemorySpace)
         opt.defaultMemorySpace = None;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
 
       // Configure type converter.
-      BufferizationOptions::LayoutMapOption unknownTypeConversionOption =
+      LayoutMapOption unknownTypeConversionOption =
           parseLayoutMapOption(unknownTypeConversion);
       opt.unknownTypeConverterFn = [=](Value value, unsigned memorySpace,
                                        const BufferizationOptions &options) {
         auto tensorType = value.getType().cast<TensorType>();
-        if (unknownTypeConversionOption ==
-            BufferizationOptions::LayoutMapOption::IdentityLayoutMap)
+        if (unknownTypeConversionOption == LayoutMapOption::IdentityLayoutMap)
           return bufferization::getMemRefTypeWithStaticIdentityLayout(
               tensorType, memorySpace);
-        assert(
-            unknownTypeConversionOption ==
-                BufferizationOptions::LayoutMapOption::FullyDynamicLayoutMap &&
+        assert(unknownTypeConversionOption ==
+                   LayoutMapOption::FullyDynamicLayoutMap &&
                "invalid layout map option");
         return bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType,
                                                                   memorySpace);
       };
 
       // Configure op filter.
       OpFilter::Entry::FilterFn filterFn = [&](Operation *op) {
         // Filter may be specified via options.
         if (this->dialectFilter.hasValue())

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

add_mlir_dialect_library(MLIRBufferizationTransforms
   OneShotModuleBufferize.cpp
   TensorCopyInsertion.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
 
   DEPENDS
   MLIRBufferizationPassIncGen
+  MLIRBufferizationEnumsIncGen
 
   LINK_LIBS PUBLIC
   MLIRBufferizationDialect
   MLIRControlFlowInterfaces
   MLIRFuncDialect
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRMemRefDialect
   MLIRPass
   MLIRTensorDialect
   MLIRTransforms
   MLIRViewLikeInterface
 )
 

mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp

 getBufferizedFunctionArgType(FuncOp funcOp, int64_t index,
                              const BufferizationOptions &options) {
   auto tensorType =
       funcOp.getFunctionType().getInput(index).dyn_cast<TensorType>();
   assert(tensorType && "expected TensorType");
 
   BaseMemRefType memrefType;
   if (options.functionBoundaryTypeConversion ==
-      BufferizationOptions::LayoutMapOption::IdentityLayoutMap) {
+      LayoutMapOption::IdentityLayoutMap) {
     memrefType = getMemRefTypeWithStaticIdentityLayout(tensorType);
   } else {
     // Note: Layout maps on function parameters cannot be inferred. The best we
     // can do at the moment is "fully dynamic".
     memrefType = getMemRefTypeWithFullyDynamicLayout(tensorType);
   }
 
   auto layoutAttr = funcOp.getArgAttrOfType<AffineMapAttr>(
for (OpOperand &returnOperand : returnOp->getOpOperands()) {
       // If not a tensor type just forward it.
       if (!tensorType) {
         returnValues.push_back(returnVal);
         continue;
       }
 
       BaseMemRefType resultType;
       if (options.functionBoundaryTypeConversion ==
-          BufferizationOptions::LayoutMapOption::IdentityLayoutMap) {
+          LayoutMapOption::IdentityLayoutMap) {
         resultType = getMemRefTypeWithStaticIdentityLayout(tensorType);
       } else {
         // Note: If `InferLayoutMap`, cast are later folded away.
         resultType = getMemRefTypeWithFullyDynamicLayout(tensorType);
       }
       Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
           loc, resultType, returnVal);
       returnValues.push_back(toMemrefOp);

mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp

LogicalResult mlir::bufferization::bufferizeModuleOp(
   // Bufferize functions.
   for (func::FuncOp funcOp : orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
     if (failed(bufferizeOp(funcOp, options, options.copyBeforeWrite)))
       return failure();
     // Change buffer return types to more precise layout maps.
     if (options.functionBoundaryTypeConversion ==
-        BufferizationOptions::LayoutMapOption::InferLayoutMap)
+        LayoutMapOption::InferLayoutMap)
       foldMemRefCasts(funcOp);
   }
 
   // Post-pass cleanup of function argument attributes.
   moduleOp.walk([&](func::FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationAttributes(bbArg);
   });

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

 static bufferization::OneShotBufferizationOptions
 getBufferizationOptions(bool analysisOnly) {
   using namespace bufferization;
   OneShotBufferizationOptions options;
   options.bufferizeFunctionBoundaries = true;
   // TODO(springerm): To spot memory leaks more easily, returning dense allocs
   // should be disallowed.
   options.allowReturnAllocs = true;
-  options.functionBoundaryTypeConversion =
-      BufferizationOptions::LayoutMapOption::IdentityLayoutMap;
+  options.functionBoundaryTypeConversion = LayoutMapOption::IdentityLayoutMap;
   options.unknownTypeConverterFn = [](Value value, unsigned memorySpace,
                                       const BufferizationOptions &options) {
     return getMemRefTypeWithStaticIdentityLayout(
         value.getType().cast<TensorType>(), memorySpace);
   };
   if (analysisOnly) {
     options.testAnalysisOnly = true;
     options.printConflicts = true;

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

func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
     // 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 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 }
+}
+
+// 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>)
+  %D = linalg.matmul ins(%A, %B: tensor<12x9xf32>, tensor<9x6xf32>) outs(%C: tensor<12x6xf32>) -> tensor<12x6xf32>
+  // CHECK: return %[[C]] : memref<12x6xf32>
+  return %D : tensor<12x6xf32>
+}