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

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

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D137833

Author: chelini

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

@@ -14,6 +14,8 @@
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/SetVector.h"
 
+#include "mlir/Dialect/Bufferization/IR/BufferizationEnums.h.inc"
+
 namespace mlir {
 class OpBuilder;
 
@@ -187,12 +189,6 @@ struct BufferizationOptions {
   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
@@ -585,6 +581,10 @@ bool defaultIsRepetitiveRegion(BufferizableOpInterface bufferizableOp,
 } // 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

@@ -0,0 +1,27 @@
+//===- 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

@@ -2,3 +2,9 @@ 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

@@ -9,6 +9,7 @@
 #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"

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

@@ -9,6 +9,7 @@
 #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"
@@ -42,6 +43,7 @@ def OneShotBufferizeOp
 
   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,
@@ -52,7 +54,10 @@ def OneShotBufferizeOp
 
   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

@@ -10,6 +10,7 @@ add_mlir_dialect_library(MLIRBufferizationDialect
   DEPENDS
   MLIRAllocationOpInterfaceIncGen
   MLIRBufferizationOpsIncGen
+  MLIRBufferizationEnumsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect

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

@@ -34,6 +34,9 @@ transform::OneShotBufferizeOp::apply(TransformResults &transformResults,
   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) {
@@ -94,6 +97,8 @@ class BufferizationTransformDialectExtension
 #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

@@ -163,14 +163,13 @@ struct FinalizingBufferizePass
   }
 };
 
-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");
 }
 
@@ -216,19 +215,17 @@ struct OneShotBufferizePass
       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 &&
-            "invalid layout map option");
+        assert(unknownTypeConversionOption ==
+                   LayoutMapOption::FullyDynamicLayoutMap &&
+               "invalid layout map option");
         return bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType,
                                                                   memorySpace);
       };

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

@@ -18,6 +18,7 @@ add_mlir_dialect_library(MLIRBufferizationTransforms
 
   DEPENDS
   MLIRBufferizationPassIncGen
+  MLIRBufferizationEnumsIncGen
 
   LINK_LIBS PUBLIC
   MLIRBufferizationDialect

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

@@ -69,7 +69,7 @@ getBufferizedFunctionArgType(FuncOp funcOp, int64_t index,
 
   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
@@ -471,7 +471,7 @@ struct FuncOpInterface
 
       BaseMemRefType resultType;
       if (options.functionBoundaryTypeConversion ==
-          BufferizationOptions::LayoutMapOption::IdentityLayoutMap) {
+          LayoutMapOption::IdentityLayoutMap) {
         resultType = getMemRefTypeWithStaticIdentityLayout(tensorType);
       } else {
         // Note: If `InferLayoutMap`, cast are later folded away.