[mlir][bufferization] Make function boundary type convertion logic dynamic.

[olegshyshkov]

Summary

Having to choose from only static or dynamic layout for all function is limiting.

[olegshyshkov]

clang-format

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

 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/SetVector.h"
 #include <optional>
 
 #include "mlir/Dialect/Bufferization/IR/BufferizationEnums.h.inc"
 
 namespace mlir {
 class OpBuilder;
+namespace func {
+class FuncOp;
+}
 
 namespace bufferization {
 
 class AnalysisState;
 class BufferizableOpInterface;
 
 /// Specifies a fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
struct BufferizationOptions {
   using DeallocationFn =
       std::function<LogicalResult(OpBuilder &, Location, Value)>;
   /// Memcpy function: Generate a memcpy between two buffers.
   using MemCpyFn =
       std::function<LogicalResult(OpBuilder &, Location, Value, Value)>;
   /// Initializer function for analysis state.
   using AnalysisStateInitFn = std::function<void(AnalysisState &)>;
   /// Tensor -> MemRef type converter.
+  /// Parameters: Value, memory space, func op, bufferization options
+  using FunctionArgTypeConverterFn =
+      std::function<BaseMemRefType(TensorType, Attribute memorySpace,
+                                   func::FuncOp, const BufferizationOptions &)>;
+  /// Tensor -> MemRef type converter.
   /// Parameters: Value, memory space, bufferization options
   using UnknownTypeConverterFn = std::function<BaseMemRefType(
       Value, Attribute memorySpace, const BufferizationOptions &)>;
 
   BufferizationOptions();
 
   /// Try to cast the given op to BufferizableOpInterface if the op is allow
   /// listed.
struct BufferizationOptions {
   /// with buffer copies.
   ///
   /// Note: Deactivating this flag can lead to incorrect bufferization results
   /// when used incorrectly. This flag is useful with
   /// `AlwaysCopyAnalysisState` which bufferizes all writing tensor
   /// OpOperands out-of-place.
   bool enforceAliasingInvariants = true;
 
-  /// This flag controls buffer types on function signatures.
+  /// This function controls buffer types on function signatures. Sets
+  /// `functionArgTypeConverterFn` and `inferFunctionResultLayout` accordingly.
   ///
   /// * InferLayoutMap: All function parameter types have a fully dynamic layout
   ///   map, but function result types are inferred from the body of the
   ///   function.
   /// * FullyDynamicLayoutMap: All function parameter types and result types
   ///   have a fully dynamic layout map. This option is most efficient because
   ///   any layout map can be casted to a fully dynamic one.
   /// * IdentityLayoutMap: All function parameter types and result types have a
   ///   static identity layout (i.e., no layout map). This option may introduce
   ///   additional buffer allocs and copies because layout maps cannot be casted
   ///   away.
   ///
-  /// If `bufferizeFunctionBoundaries` is not set, this flag has no effect.
-  ///
   /// Note: Inferred layout maps may not be desireable when interacting with
   /// external functions, because the generated function signatures will be less
   /// predictable.
-  LayoutMapOption functionBoundaryTypeConversion =
-      LayoutMapOption::InferLayoutMap;
+  void setFunctionBoundaryTypeConversion(LayoutMapOption layoutMapOption);
+
+  /// Type converter from tensors to memrefs. This type converter is used to
+  /// determine bufferized function argument types. By default, a type
+  /// converter that returns a memref type with a fully dynamic layout map is
+  /// used.
+  ///
+  /// If `bufferizeFunctionBoundaries` is not set, this function isn't used.
+  FunctionArgTypeConverterFn functionArgTypeConverterFn = nullptr;
+
+  /// If true, function result types are inferred from the body of the function.
+  /// Otherwise, function result type is determined by
+  /// `functionArgTypeConverterFn`.
+  ///
+  /// If `bufferizeFunctionBoundaries` is not set, this flag has no effect.
+  bool inferFunctionResultLayout = true;
 
   /// Type converter from tensors to memrefs. This type converter is used if no
   /// memref type could be inferred during bufferization. By default, a type
   /// converter that returns a memref type with a fully dynamic layout map is
   /// used.
   UnknownTypeConverterFn unknownTypeConverterFn = nullptr;
 
   /// Specifies whether dealloc ops should be generated along with alloc ops. If

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

bool OpFilter::isOpAllowed(Operation *op) const {
   }
   return isAllowed;
 }
 
 //===----------------------------------------------------------------------===//
 // BufferizationOptions
 //===----------------------------------------------------------------------===//
 
+namespace {
+
+/// Default function arg type converter: Use a fully dynamic layout map.
+BaseMemRefType
+defaultFunctionArgTypeConverter(TensorType type, Attribute memorySpace,
+                                func::FuncOp funcOp,
+                                const BufferizationOptions &options) {
+  return getMemRefTypeWithFullyDynamicLayout(type, memorySpace);
+}
 /// Default unknown type converter: Use a fully dynamic layout map.
-static BaseMemRefType
+BaseMemRefType
 defaultUnknownTypeConverter(Value value, Attribute memorySpace,
                             const BufferizationOptions &options) {
   return getMemRefTypeWithFullyDynamicLayout(value.getType().cast<TensorType>(),
                                              memorySpace);
 }
 
+}; // namespace
+
 // Default constructor for BufferizationOptions.
 BufferizationOptions::BufferizationOptions()
-    : unknownTypeConverterFn(defaultUnknownTypeConverter) {}
+    : functionArgTypeConverterFn(defaultFunctionArgTypeConverter),
+      unknownTypeConverterFn(defaultUnknownTypeConverter) {}
 
 bool BufferizationOptions::isOpAllowed(Operation *op) const {
   // Special case: If function boundary bufferization is deactivated, do not
   // allow ops that belong to the `func` dialect.
   bool isFuncBoundaryOp = isa_and_nonnull<func::FuncDialect>(op->getDialect());
   if (!bufferizeFunctionBoundaries && isFuncBoundaryOp)
     return false;
 

 BufferizableOpInterface
 BufferizationOptions::dynCastBufferizableOp(Value value) const {
   if (auto bufferizableOp = value.getDefiningOp<BufferizableOpInterface>())
     if (isOpAllowed(bufferizableOp.getOperation()))
       return bufferizableOp;
   return nullptr;
 }
 
+void BufferizationOptions::setFunctionBoundaryTypeConversion(
+    LayoutMapOption layoutMapOption) {
+  functionArgTypeConverterFn = [=](TensorType tensorType, Attribute memorySpace,
+                                   func::FuncOp funcOp,
+                                   const BufferizationOptions &options) {
+    if (layoutMapOption == LayoutMapOption::IdentityLayoutMap)
+      return bufferization::getMemRefTypeWithStaticIdentityLayout(tensorType,
+                                                                  memorySpace);
+    return bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType,
+                                                              memorySpace);
+  };
+  inferFunctionResultLayout =
+      layoutMapOption == LayoutMapOption::InferLayoutMap;
+}
+
 //===----------------------------------------------------------------------===//
 // Helper functions for BufferizableOpInterface
 //===----------------------------------------------------------------------===//
 
 static void setInsertionPointAfter(OpBuilder &b, Value value) {
   if (auto bbArg = value.dyn_cast<BlockArgument>()) {
     b.setInsertionPointToStart(bbArg.getOwner());
   } else {

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

transform::OneShotBufferizeOp::apply(TransformResults &transformResults,
   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();
+    options.setFunctionBoundaryTypeConversion(
+        *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 (options.bufferizeFunctionBoundaries) {
       if (!moduleOp)

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

if (!options) {
       // Make new bufferization options if none were provided when creating the
       // pass.
       opt.allowReturnAllocs = allowReturnAllocs;
       opt.allowUnknownOps = allowUnknownOps;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.analysisHeuristic = parseHeuristicOption(analysisHeuristic);
       opt.copyBeforeWrite = copyBeforeWrite;
       opt.createDeallocs = createDeallocs;
-      opt.functionBoundaryTypeConversion =
-          parseLayoutMapOption(functionBoundaryTypeConversion);
+      opt.setFunctionBoundaryTypeConversion(
+          parseLayoutMapOption(functionBoundaryTypeConversion));
       if (mustInferMemorySpace)
         opt.defaultMemorySpace = std::nullopt;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
       opt.noAnalysisFuncFilter = noAnalysisFuncFilter;
 
       // Configure type converter.

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

if (auto candidateOp = dyn_cast<func::ReturnOp>(b.getTerminator())) {
       returnOp = candidateOp;
     }
   }
   return returnOp;
 }
 
 /// Return the index-th bufferized function argument type. This assumes that the
 /// specified argument is a tensor. If the tensor is ranked, a layout map may be
-/// specified by the user. If no layout map is specified, the default layout map
-/// (as per `options.functionBoundaryTypeConversion`) is used.
+/// specified by the user (as per `options.functionArgTypeConverterFn`).
 static BaseMemRefType
 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 ==
-      LayoutMapOption::IdentityLayoutMap) {
-    memrefType = getMemRefTypeWithStaticIdentityLayout(
-        tensorType, *options.defaultMemorySpace);
-  } else {
-    // Note: Layout maps on function parameters cannot be inferred. The best we
-    // can do at the moment is "fully dynamic".
-    memrefType = getMemRefTypeWithFullyDynamicLayout(
-        tensorType, *options.defaultMemorySpace);
-  }
+  BaseMemRefType memrefType = options.functionArgTypeConverterFn(
+      tensorType, *options.defaultMemorySpace, funcOp, options);
 
   auto layoutAttr = funcOp.getArgAttrOfType<AffineMapAttr>(
       index, BufferizationDialect::kBufferLayoutAttrName);
   if (!layoutAttr)
     return memrefType;
 
   auto rankedMemrefType = memrefType.dyn_cast<MemRefType>();
   assert(rankedMemrefType && "buffer layout not supported on unranked tensors");
for (OpOperand &returnOperand : returnOp->getOpOperands()) {
       rewriter.setInsertionPoint(returnOp);
 
       // If not a tensor type just forward it.
       if (!tensorType) {
         returnValues.push_back(returnVal);
         continue;
       }
 
-      BaseMemRefType resultType;
-      if (options.functionBoundaryTypeConversion ==
-          LayoutMapOption::IdentityLayoutMap) {
-        resultType = getMemRefTypeWithStaticIdentityLayout(
-            tensorType, *options.defaultMemorySpace);
-      } else {
-        // Note: If `InferLayoutMap`, cast are later folded away.
-        resultType = getMemRefTypeWithFullyDynamicLayout(
-            tensorType, *options.defaultMemorySpace);
-      }
+      // Note: If `inferFunctionResultLayout = true`, cast are later folded
+      // away.
+      BaseMemRefType resultType = options.functionArgTypeConverterFn(
+          tensorType, *options.defaultMemorySpace, funcOp, options);
       Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
           loc, resultType, returnVal);
       returnValues.push_back(toMemrefOp);
     }
 
     // 3. Rewrite the terminator without the in-place bufferizable values.
     returnOp.getOperandsMutable().assign(returnValues);
 

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

for (func::FuncOp funcOp : orderedFuncOps) {
     // would be invalidated.
     bool copyBeforeWrite =
         options.copyBeforeWrite ||
         llvm::is_contained(options.noAnalysisFuncFilter, funcOp.getSymName());
     if (failed(bufferizeOp(funcOp, options, copyBeforeWrite,
                            /*opFilter=*/nullptr, statistics)))
       return failure();
     // Change buffer return types to more precise layout maps.
-    if (options.functionBoundaryTypeConversion ==
-        LayoutMapOption::InferLayoutMap)
+    if (options.inferFunctionResultLayout)
       foldMemRefCasts(funcOp);
   }
 
   // Post-pass cleanup of function argument attributes.
   removeBufferizationAttributesInModule(moduleOp);
 
   return success();
 }

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 = LayoutMapOption::IdentityLayoutMap;
+  options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);
   options.unknownTypeConverterFn = [](Value value, Attribute memorySpace,
                                       const BufferizationOptions &options) {
     return getMemRefTypeWithStaticIdentityLayout(
         value.getType().cast<TensorType>(), memorySpace);
   };
   if (analysisOnly) {
     options.testAnalysisOnly = true;
     options.printConflicts = true;