[mlir][bufferize] Better user control of layout maps

[springerm]

Summary

This changes replaces the fully-dynamic-layout-maps options (which was badly named) with two new options:

• unknown-type-conversion controls the layout maps on buffer types for which no layout map can be inferred.
• function-boundary-type-conversion controls the layout maps on buffer types inside of function signatures.

Depends On D125614

[pifon2a]

mlir/docs/Bufferization.md

344

nit: precise

mlir/docs/Bufferization.md

 All users of `%0` have fully dynamic layout maps. This ensures that the
 bufferized IR composes well with future bufferizations of `unbufferizable_op`
 (maybe bufferized by another pass), regardless of the exact memref type of the
 future bufferization. If the op turns out to be bufferized to an op with a
 simpler memref type (e.g., identity layout map), we expect that canonicalization
 patterns would clean up unnecessarily dynamic layout maps. (Some of these
 canonicalization patterns may not be implemented yet.)
 
-Note that One-Shot Bufferize always generates the most specific memref type when
-the entire IR is bufferizable. In that case, we do not have to rely on
-canonicalization patterns to clean up the bufferized IR.
-
-One-Shot Bufferize can be configured to always generate memref types with
-identity layout when the exact target memref type is not known via
-`fully-dynamic-layout-maps=0`. This can be useful for legacy code that cannot
-handle memref types with layout maps. Note that this leads to additional buffer
-copies when folding a `to_tensor`/`to_memref` pair with memref types that are
-not cast-compatible.
+One-Shot Bufferize tries to infer the most precise memref type when bufferizing
+an op. If the entire IR is bufferizable, we do not have to resort to
+conservatively use fully dynamic layout maps. In that case, we also do not have
+to rely on canonicalization patterns to clean up the bufferized IR.
+
+Note: There are some bufferizable ops for which a percise layout map cannot be
+inferred. E.g., a `tensor.cast` from a `tensor<*xf32>` to a `tensor<?x?xf32>`
+must be bufferized to a `memref.cast` with a memref type that has a fully
+dynamic layout map.
+
+One-Shot Bufferize has an option `unknown-type-conversion` to control the
+generation of layout maps when no precise layout can be inferred:
+
+*   `fully-dynamic-layout-map` uses fully dynamic layout maps and is the default
+    behavior. This composes well when IR is partially bufferized.
+*   `identity-layout-map` uses static identity layout maps. This option can be
+    useful for legacy code that cannot handle memref types with layout maps.
+    Note that this setting can lead to additional buffer copies when folding a
+    `to_tensor`/`to_memref` pair with memref types that are not cast-compatible.
+
+Note: The `unknown-type-conversion` option does not affect layout maps of
+function signatures. There is a separate `function-signature-type-conversion`
+option that controls layout maps of function parameters and function results.
 
 ## Extending One-Shot Bufferize
 
 Custom ops can be bufferized if they implement `BufferizableOpInterface`. Users
 must at least implement the following interface methods.
 
 *   `bufferizesToMemoryRead`: Return `true` if the buffer of the given tensor
     OpOperand is read.

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

struct OpFilterEntry {
 
     /// Filter type: A filter can either be a DENY filter or an ALLOW filter.
     enum FilterType : int8_t { DENY = 0, ALLOW = 1 };
 
     FilterFn fn;
     FilterType type;
   };
 
+  enum class LayoutMapOption : int8_t {
+    InferLayoutMap = 0,
+    IdentityLayoutMap = 1,
+    FullyDynamicLayoutMap = 2
+  };
+
   BufferizationOptions();
 
   /// Return `true` if the filter has at least one ALLOW rule.
   bool filterHasAllowRule() const {
     for (const OpFilterEntry &e : opFilter)
       if (e.type == OpFilterEntry::FilterType::ALLOW)
         return true;
     return false;
struct BufferizationOptions {
   /// bufferization.to_memref and bufferization.to_tensor ops are inserted at
   /// the boundaries.
   bool allowUnknownOps = false;
 
   /// Specifies whether function boundaries (ops in the func dialect) should be
   /// bufferized or not.
   bool bufferizeFunctionBoundaries = false;
 
+  /// This flag controls buffer types on function signatures.
+  ///
+  /// * 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. If
+  /// `promoteBufferResultsToOutParams` is set, `kInferMostPreciseLayoutMap` is
+  /// is an invalid option.
+  ///
+  /// 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;
+
+  /// This flag controls buffer types on unknown ops (to_memref wrappers) and in
+  /// other cases where a precise memref type cannot be inferred (e.g., the
+  /// bufferization of "tensor.cast").
+  ///
+  /// * InferLayoutMap: This option is invalid and cannot be used.
+  /// * FullyDynamicLayoutMap: Assume that unknown ops have results with fully
+  ///   dynamic layout maps after bufferization. This option is most efficient
+  ///   because any layout map can be casted to a fully dynamic one.
+  /// * IdentityLayoutMap: Assume that unknown ops have results with static
+  ///   identity layout (i.e., no layout map) after bufferization. This option
+  ///   introduces additional buffer allocs and copies if the unknown op is
+  ///   eventually bufferized to an op that returns a buffer with non-identity
+  ///   layout.
+  LayoutMapOption unknownTypeConversion =
+      LayoutMapOption::FullyDynamicLayoutMap;
+
   /// Specifies whether dealloc ops should be generated along with alloc ops. If
   /// not, new memory allocations will leak.
   bool createDeallocs = true;
 
   /// Seed for the analysis fuzzer. If set to `0`, the fuzzer is deactivated.
   /// Should be used only with `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
 
-  /// Specifies whether fully dynamic layout maps should be used on ranked
-  /// MemRef types. If false, MemRef types will have no layout maps.
-  bool fullyDynamicLayoutMaps = true;
-
   /// If set to `true`, does not modify the IR apart from adding attributes (for
   /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
 
   /// If set to `true`, the IR is annotated with details about RaW conflicts.
   /// For debugging only. Should be used together with `testAnalysisOnly`.
   bool printConflicts = false;
 
OpTy replaceOpWithNewBufferizedOp(RewriterBase &rewriter, Operation *op,
   return newOp;
 }
 
 /// Return a MemRefType to which the `tensorType` can be bufferized.
 ///
 /// If possible, op bufferization implementations should not use this function
 /// and instead infer precise memref types for tensor results by themselves.
 ///
-/// Unless a layout map was specified, `options` flags determine what kind of
-/// layout map will be used. For best composability (without copies), the fully
-/// dynamic layout map is used by default.
+/// Unless a layout map was specified, `options.unknownTypeConverter` determines
+/// what kind of layout map will be used. For best composability (without
+/// copies), the fully dynamic layout map is used by default.
 ///
 /// Note: Canonicalization patterns could clean up layout maps and infer more
 /// precise layout maps after bufferization. However, many possible
 /// canonicalizations are currently not implemented.
 BaseMemRefType getMemRefType(TensorType tensorType,
                              const BufferizationOptions &options,
                              MemRefLayoutAttrInterface layout = {},
                              Attribute memorySpace = {});
 
+/// Return a MemRef type with fully dynamic layout. If the given tensor type
+/// is unranked, return an unranked MemRef type.
+BaseMemRefType getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
+                                                   Attribute memorySpace = {});
+
+/// Return a MemRef type with a static identity layout (i.e., no layout map). If
+/// the given tensor type is unranked, return an unranked MemRef type.
+BaseMemRefType
+getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
+                                      Attribute memorySpace = {});
+
 /// Try to hoist all new buffer allocations until the next hoisting barrier.
 LogicalResult hoistBufferAllocations(Operation *op,
                                      const BufferizationOptions &options);
 
 /// Create alloc/dealloc ops as specified in the bufferization options. If
 /// `onlyLeakingAlloc`, only those buffer allocations are processed for which no
 /// buffer deallocation can be created. `changed` is set to `true` if the IR was
 /// modified.

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

let description = [{
     With `dialect-filter`, bufferization can be restricted to a set of dialects.
     If no filter is specified, all ops that implement `BufferizableOpInterface`
     are bufferized. Ops from the `std` dialect are an exception: These ops are
     always ignored, even if no filter is specified. When specifying a dialect
     filter and `allow-unknown-ops` is not turned on, bufferization would fail
     when encountering an op that is not included in the filter (even if it is
     bufferizable).
 
+    One-Shot Bufferize will by default assume memref types with fully dynamic
+    layout maps when a precise layout cannot be inferred. E.g., this is the case
+    when wrapping a non-bufferizable op in to_memref/to_tensor ops. This
+    behavior can be overridden with `unknown-type-conversion`. Valid values are
+    `fully-dynamic-layout-map` and `identity-layout-map`.
+
+    Layout maps on function signatures can be controlled with a separate
+    `function-boundary-type-conversion` option, which can be set to
+    `infer-layout-map` in addition to the two possible values mentioned above.
+    When layout maps are referred, function return types may be more precise.
+    Function argument types cannot be inferred and have fully dynamic layout
+    maps in that case.
+
     For testing/debugging purposes, `test-analysis-only=1 print-conflicts=1`
     prints analysis results and explains why an OpOperand was decided to
     bufferize out-of-place. This is useful for understanding why One-Shot
     Bufferize chose to insert a certain buffer copy.
 
     `bufferize-function-boundaries` is an experimental flag for bufferizing
     `FuncOp`, `ReturnOp` and `CallOp`. This feature is still under development
     and supports only simple cases at the moment. In particular:
let options = [
     Option<"bufferizeFunctionBoundaries", "bufferize-function-boundaries",
            "bool", /*default=*/"0",
            "Bufferize function boundaries (experimental).">,
     Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
            "Specify if buffers should be deallocated. For compatibility with "
            "core bufferization passes.">,
     ListOption<"dialectFilter", "dialect-filter", "std::string",
                "Restrict bufferization to ops from these dialects.">,
-    Option<"fullyDynamicLayoutMaps", "fully-dynamic-layout-maps", "bool",
-           /*default=*/"true",
-           "Generate MemRef types with dynamic offset+strides by default.">,
+    Option<"functionBoundaryTypeConversion",
+           "function-boundary-type-conversion", "std::string",
+           /*default=*/"\"infer-layout-map\"",
+           "Controls layout maps when bufferizing function signatures.">,
     Option<"testAnalysisOnly", "test-analysis-only", "bool",
             /*default=*/"false",
            "Test only: Only run inplaceability analysis and annotate IR">,
     Option<"printConflicts", "print-conflicts", "bool",
             /*default=*/"false",
            "Test only: Annotate IR with RaW conflicts. Requires "
            "test-analysis-only.">,
     Option<"promoteBufferResultsToOutParams",
            "promote-buffer-results-to-out-params", "bool", /*default=*/"false",
            "Replace returned buffers (that were not dropped) with out params.">,
+    Option<"unknownTypeConversion", "unknown-type-conversion", "std::string",
+           /*default=*/"\"fully-dynamic-layout-map\"",
+           "Controls layout maps for non-inferrable memref types.">,
   ];
   let constructor = "mlir::bufferization::createOneShotBufferizePass()";
 }
 
 def PromoteBuffersToStack : Pass<"promote-buffers-to-stack", "func::FuncOp"> {
   let summary = "Promotes heap-based allocations to automatically managed "
                 "stack-based allocations";
   let description = [{

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

BaseMemRefType bufferization::getMemRefType(TensorType tensorType,
                                             Attribute memorySpace) {
   // Case 1: Unranked memref type.
   if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
     assert(!layout && "UnrankedTensorType cannot have a layout map");
     return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
                                    memorySpace);
   }
 
-  // Case 2: Ranked memref type with specified layout. If fully dynamic layout
-  // maps are not requested, generate a type with `layout`, which is empty (no
-  // layout map) by default.
+  // Case 2: Ranked memref type with specified layout.
   auto rankedTensorType = tensorType.cast<RankedTensorType>();
-  if (layout || !options.fullyDynamicLayoutMaps) {
+  if (layout) {
     return MemRefType::get(rankedTensorType.getShape(),
                            rankedTensorType.getElementType(), layout,
                            memorySpace);
   }
 
-  // Case 3: Ranked memref type with unspecified layout. Choose the most dynamic
-  // one.
-  // TODO: address space decisions to connect with the actual alloc.
+  // Case 3: Configured with "fully dynamic layout maps".
+  if (options.unknownTypeConversion ==
+      BufferizationOptions::LayoutMapOption::FullyDynamicLayoutMap)
+    return getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace);
+
+  // Case 4: Configured with "static identity layout maps".
+  if (options.unknownTypeConversion ==
+      BufferizationOptions::LayoutMapOption::IdentityLayoutMap)
+    return getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
+
+  llvm_unreachable("InferLayoutMap is an invalid option");
+}
+
+BaseMemRefType
+bufferization::getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
+                                                   Attribute memorySpace) {
+  // Case 1: Unranked memref type.
+  if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
+    return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
+                                   memorySpace);
+  }
+
+  // Case 2: Ranked memref type.
+  auto rankedTensorType = tensorType.cast<RankedTensorType>();
   int64_t dynamicOffset = ShapedType::kDynamicStrideOrOffset;
   SmallVector<int64_t> dynamicStrides(rankedTensorType.getRank(),
                                       ShapedType::kDynamicStrideOrOffset);
   AffineMap stridedLayout = makeStridedLinearLayoutMap(
       dynamicStrides, dynamicOffset, rankedTensorType.getContext());
   return MemRefType::get(rankedTensorType.getShape(),
                          rankedTensorType.getElementType(), stridedLayout,
                          memorySpace);
 }
+
+/// Return a MemRef type with a static identity layout (i.e., no layout map). If
+/// the given tensor type is unranked, return an unranked MemRef type.
+BaseMemRefType
+bufferization::getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
+                                                     Attribute memorySpace) {
+  // Case 1: Unranked memref type.
+  if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
+    return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
+                                   memorySpace);
+  }
+
+  // Case 2: Ranked memref type.
+  auto rankedTensorType = tensorType.cast<RankedTensorType>();
+  MemRefLayoutAttrInterface layout = {};
+  return MemRefType::get(rankedTensorType.getShape(),
+                         rankedTensorType.getElementType(), layout,
+                         memorySpace);
+}

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(std::string s) {
+  if (s == "fully-dynamic-layout-map")
+    return BufferizationOptions::LayoutMapOption::FullyDynamicLayoutMap;
+  if (s == "identity-layout-map")
+    return BufferizationOptions::LayoutMapOption::IdentityLayoutMap;
+  if (s == "infer-layout-map")
+    return BufferizationOptions::LayoutMapOption::InferLayoutMap;
+  llvm_unreachable("invalid layout map option");
+}
+
 struct OneShotBufferizePass
     : public OneShotBufferizeBase<OneShotBufferizePass> {
   OneShotBufferizePass() : OneShotBufferizeBase<OneShotBufferizePass>() {}
 
   explicit OneShotBufferizePass(const OneShotBufferizationOptions &options)
       : options(options) {}
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<bufferization::BufferizationDialect, memref::MemRefDialect>();
     registerAllocationOpInterfaceExternalModels(registry);
   }
 
   void runOnOperation() override {
     OneShotBufferizationOptions opt;
     if (!options) {
       // Make new bufferization options if none were provided when creating the
       // pass.
       opt.dropEquivalentFuncResults = dropEquivalentFuncResults;
       opt.allowReturnAllocs = allowReturnAllocs;
       opt.allowUnknownOps = allowUnknownOps;
       opt.alwaysAliasingWithDest = alwaysAliasingWithDest;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.createDeallocs = createDeallocs;
-      opt.fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
+      opt.functionBoundaryTypeConversion =
+          parseLayoutMapOption(functionBoundaryTypeConversion);
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
       opt.promoteBufferResultsToOutParams = promoteBufferResultsToOutParams;
+      opt.unknownTypeConversion = parseLayoutMapOption(unknownTypeConversion);
 
       BufferizationOptions::OpFilterEntry::FilterFn filterFn =
           [&](Operation *op) {
             // Filter may be specified via options.
             if (this->dialectFilter.hasValue())
               return llvm::find(this->dialectFilter,
                                 op->getDialect()->getNamespace()) !=
                      this->dialectFilter.end();
private:
   const BufferizationOptions &options;
 };
 } // namespace
 
 LogicalResult
 bufferization::bufferizeOp(Operation *op,
                            BufferizationState &bufferizationState) {
   const auto &options = bufferizationState.getOptions();
+  assert(options.unknownTypeConversion !=
+             BufferizationOptions::LayoutMapOption::InferLayoutMap &&
+         "invalid layout map option");
 
   // Keep track of to_memref ops.
   DenseSet<Operation *> toMemrefOps;
   op->walk([&](ToMemrefOp toMemrefOp) { toMemrefOps.insert(toMemrefOp); });
 
   // Gather all bufferizable ops in top-to-bottom order.
   //
   // We should ideally know the exact memref type of all operands when
   // bufferizing an op. (This is the case when bufferizing top-to-bottom.)
-  // Otherwise, we have to use a memref type with a fully dynamic layout map,
-  // which has to canonicalize away. This is less efficient.
-  //
-  // If "fullyDynamicLayoutMaps = false", we would have to insert buffer copies
-  // to fold ("finalize") to_memref(to_tensor(x)) ops with non-cast-compatible
-  // layout maps when doing a traversal other than top-to-bottom. These would
-  // not easily fold away.
+  // Otherwise, we have to use a memref type with a fully dynamic layout map to
+  // avoid copies. We are currently missing patterns for layout maps to
+  // canonicalize away (or canonicalize to more precise layouts).
   SmallVector<Operation *> worklist;
   op->walk<WalkOrder::PreOrder>([&](Operation *op) {
     if (hasTensorSemantics(op))
       worklist.push_back(op);
   });
 
   // Keep track of all erased ops.
   DenseSet<Operation *> erasedOps;
LogicalResult bufferization::bufferizeOp(Operation *op,
   AlwaysCopyAnalysisState state(options);
   return bufferizeOp(op, state);
 }
 
 BufferizationOptions bufferization::getPartialBufferizationOptions() {
   BufferizationOptions options;
   options.allowUnknownOps = true;
   options.createDeallocs = false;
-  options.fullyDynamicLayoutMaps = false;
+  options.unknownTypeConversion =
+      BufferizationOptions::LayoutMapOption::IdentityLayoutMap;
   return options;
 }

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, a fully dynamic map is
-/// used.
+/// specified by the user. If no layout map is specified, the default layout map
+/// (as per `options.functionBoundaryTypeConversion`) is used.
 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 = getMemRefType(tensorType, options);
+
+  BaseMemRefType memrefType;
+  if (options.functionBoundaryTypeConversion ==
+      BufferizationOptions::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>(
       index, BufferizationDialect::kBufferLayoutAttrName);
   if (!layoutAttr)
     return memrefType;
 
   auto rankedMemrefType = memrefType.dyn_cast<MemRefType>();
   assert(rankedMemrefType && "buffer layout not supported on unranked tensors");
#endif // NDEBUG
 
     // ReturnOps are bufferized as part of FuncOps.
     return failure();
   }
 };
 
 struct FuncOpInterface
     : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
-  /// Rewrite function bbArgs and return values into buffer form (using the
-  /// canonical memref layout for now). This function bufferizes the function
-  /// signature and the ReturnOp. When the entire function body has been
-  /// bufferized, function return types can be switched to more concise memref
-  /// types as part of `foldMemRefCasts`.
+  /// Rewrite function bbArgs and return values into buffer form. This function
+  /// bufferizes the function signature and the ReturnOp. When the entire
+  /// function body has been bufferized, function return types can be switched
+  /// to more concise memref types as part of `foldMemRefCasts`.
   ///
   /// When a tensor function argument is known to be equivalent to a tensor
   /// result, it is dropped from the return values.
   ///
   /// All function bbArgs are writable unless they are explicitly marked as
   /// read-only. Callers must insert copies when needed.
   ///
   /// Note: Returning a memref is possible, but corresponding CallOp
if (funcOp.getBody().empty()) {
       }
       funcOp.setType(FunctionType::get(op->getContext(), argTypes, retTypes));
       return success();
     }
 
     // TODO: Support functions with multiple returns.
     func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
     assert(returnOp && "expected func with single return op");
+    Location loc = returnOp.getLoc();
 
     // 1. Rewrite the bbArgs. Turn every tensor bbArg into a memref bbArg.
     Block &frontBlock = funcOp.getBody().front();
     for (BlockArgument &bbArg : frontBlock.getArguments()) {
       auto tensorType = bbArg.getType().dyn_cast<TensorType>();
       // Non-tensor types stay the same.
       if (!tensorType)
         continue;
for (BlockArgument &bbArg : frontBlock.getArguments()) {
           use->set(toTensorOp);
       }
     }
 
     // 2. For each result, keep track of which inplace argument it reuses.
     SmallVector<Value> returnValues;
     for (OpOperand &returnOperand : returnOp->getOpOperands()) {
       Value returnVal = returnOperand.get();
+      auto tensorType = returnVal.getType().dyn_cast<TensorType>();
+      rewriter.setInsertionPoint(returnOp);
 
       // If not a tensor type just forward it.
-      if (!returnVal.getType().isa<RankedTensorType>()) {
+      if (!tensorType) {
         returnValues.push_back(returnVal);
         continue;
       }
 
       // If return operand is equivalent to some bbArg, no need to return it.
       if (options.dropEquivalentFuncResults) {
         if (Optional<int64_t> equivBbArgIdx = getEquivalentFuncArgIdx(
                 funcOp, funcState, returnOperand.getOperandNumber())) {
-          rewriter.setInsertionPoint(returnOp);
-          Location loc = returnOp.getLoc();
+          // TODO: Use memref type with fully dynamic layout map and add folder
+          // for memref.cast + memref.copy.
           Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
-              loc,
-              getMemRefType(returnVal.getType().cast<TensorType>(), options),
-              returnVal);
+              loc, getMemRefType(tensorType, options), returnVal);
           BlockArgument equivBbArg = funcOp.getArgument(*equivBbArgIdx);
           // Note: This copy will fold away. It must be inserted here to ensure
           // that `returnVal` still has at least one use and does not fold away.
           if (failed(
                   options.createMemCpy(rewriter, loc, toMemrefOp, equivBbArg)))
             return funcOp->emitError("could not generate copy for bbArg");
           continue;
         }
       }
 
-      returnValues.push_back(*state.getBuffer(rewriter, returnOperand));
+      BaseMemRefType resultType;
+      if (options.functionBoundaryTypeConversion ==
+          BufferizationOptions::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);
     }
 
     // 3. Rewrite the terminator without the in-place bufferizable values.
     returnOp.operandsMutable().assign(returnValues);
 
     // 4. Rewrite the FuncOp type to buffer form.
     funcOp.setType(FunctionType::get(op->getContext(), argTypes,
                                      ValueRange(returnValues).getTypes()));

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

if (options.testAnalysisOnly)
     return success();
 
   // 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, bufferizationState)))
       return failure();
+    // Change buffer return types to more precise layout maps.
+    if (options.functionBoundaryTypeConversion ==
+        BufferizationOptions::LayoutMapOption::InferLayoutMap)
       foldMemRefCasts(funcOp);
   }
 
   // Check result.
   for (func::FuncOp funcOp : orderedFuncOps) {
     if (!options.allowReturnAllocs &&
         llvm::any_of(funcOp.getFunctionType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();
         })) {

mlir/test/Dialect/Arithmetic/one-shot-bufferize.mlir

 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
-// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs fully-dynamic-layout-maps=0 bufferize-function-boundaries" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // CHECK-LABEL: func @write_to_select_op_source
 //  CHECK-SAME:     %[[t1:.*]]: memref<?xf32, #{{.*}}>, %[[t2:.*]]: memref<?xf32, #{{.*}}>
 func.func @write_to_select_op_source(
     %t1 : tensor<?xf32> {bufferization.writable = true},
     %t2 : tensor<?xf32> {bufferization.writable = true},
     %c : i1)
   -> (tensor<?xf32>, tensor<?xf32>)

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-partial.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs allow-unknown-ops" -split-input-file | FileCheck %s
 
 // Test bufferization using memref types that have no layout map.
-// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs allow-unknown-ops fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs allow-unknown-ops unknown-type-conversion=identity-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="dialect-filter=tensor allow-unknown-ops allow-return-allocs" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-TENSOR
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="dialect-filter=scf allow-unknown-ops allow-return-allocs" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-SCF

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-allow-return-allocs.mlir

 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs" -split-input-file | FileCheck %s
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs drop-equivalent-func-results=false" -split-input-file | FileCheck %s --check-prefix=EQUIV
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs fully-dynamic-layout-maps=0" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map" -split-input-file -o /dev/null
 
 // Make sure that the returned buffer is not deallocated.
 // TODO: Such buffers currently leak. We need buffer hoisting / ref counting for
 // this in the future.
 
 // CHECK-LABEL: func @create_tensor() -> memref<10xf32> {
 //       CHECK:   %[[alloc:.*]] = memref.alloc
 //       CHECK:   return %[[alloc]]

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir

-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1" -split-input-file | FileCheck %s
+// Note: Default is function-boundary-type-conversion=infer-layout-map
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs" -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP-LABEL
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
+
+// Test bufferization using memref types that have fully dynamic layout maps.
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs function-boundary-type-conversion=fully-dynamic-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-FULLY-DYNAMIC-LAYOUT-MAP
+
 
 // Bufferization of bodiless function with no tensor return value.
 
-// CHECK-LABEL: func private @private_func
+// CHECK: #[[$map0:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+// CHECK: #[[$map1:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+// CHECK-LABEL: func private @private_func(memref<?xf32,
+//  CHECK-SAME:                                          #[[$map0]]>)
+// CHECK-NO-LAYOUT-MAP-LABEL: func private @private_func(memref<?xf32>)
 func.func private @private_func(tensor<?xf32>) -> ()
 
-// CHECK-LABEL: func @empty_func()
+// CHECK-LABEL: func private @private_func_2d(memref<?x?xf32,
+//  CHECK-SAME:                                               #[[$map1]]>)
+// CHECK-NO-LAYOUT-MAP-LABEL: func private @private_func_2d(memref<?x?xf32>)
+func.func private @private_func_2d(tensor<?x?xf32>) -> ()
+
+// CHECK-LABEL: func @empty_func() {
+// CHECK-NO-LAYOUT-MAP-LABEL: func @empty_func() {
+// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func @empty_func() {
 func.func @empty_func() -> () {
   return
 }
 
 // -----
 
 // A bodiless function that returns something that is not a tensor.
 
 // CHECK: func private @external_func_with_return_val(memref<4xi32, #{{.*}}>) -> f32
+// CHECK-FULLY-DYNAMIC-LAYOUT-MAP: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func private @external_func_with_return_val(memref<4xi32,
+// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-SAME: #[[$map1]]>
 func.func private @external_func_with_return_val(tensor<4xi32>) -> f32
 
 // -----
 
+// A function that returns a non-equivalent tensor with layout map.
+
+// CHECK: #[[$map2:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 10 + s0 + d1)>
+// CHECK-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32,
+//  CHECK-SAME:     #[[$map2]]> {
+//       CHECK:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<20x10xf32>
+//       CHECK:   %[[subview:.*]] = memref.subview {{.*}} : memref<20x10xf32> to memref<2x?xf32, #[[$map2]]>
+//       CHECK:   return %[[subview]]
+
+// CHECK-NO-LAYOUT-MAP: #[[$map2:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 10 + s0 + d1)>
+// CHECK-NO-LAYOUT-MAP-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32>
+//       CHECK-NO-LAYOUT-MAP:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<20x10xf32>
+//       CHECK-NO-LAYOUT-MAP:   %[[subview:.*]] = memref.subview {{.*}} : memref<20x10xf32> to memref<2x?xf32, #[[$map2]]>
+//       CHECK-NO-LAYOUT-MAP:   %[[alloc_no_layout:.*]] = memref.alloc(%{{.*}}) : memref<2x?xf32>
+//       CHECK-NO-LAYOUT-MAP:   memref.copy %[[subview]], %[[alloc_no_layout]]
+//       CHECK-NO-LAYOUT-MAP:   memref.dealloc %[[alloc]]
+//       CHECK-NO-LAYOUT-MAP:   return %[[alloc_no_layout]]
+
+// CHECK-FULLY-DYNAMIC-LAYOUT-MAP: #[[$map2a:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+// CHECK-FULLY-DYNAMIC-LAYOUT-MAP: #[[$map2b:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 10 + s0 + d1)>
+// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32,
+//  CHECK-FULLY-DYNAMIC-LAYOUT-MAP-SAME: #[[$map2a]]> {
+func.func @return_extract_slice(%idx: index, %sz: index) -> (tensor<2x?xf32>)
+{
+  %t = linalg.init_tensor [20, 10] : tensor<20x10xf32>
+  %0 = tensor.extract_slice %t[%idx, %idx][2, %sz][1, 1]
+      : tensor<20x10xf32> to tensor<2x?xf32>
+  return %0 : tensor<2x?xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func private @private_func
 func.func private @private_func(tensor<?xf32>) -> (f32)
 
 // private_func may modify the buffer arg, but that's OK because %t is writable.
 // No alloc/copy should be inserted.
 
 // CHECK-LABEL: func @main(
 //  CHECK-SAME:     %[[t:.*]]: memref<?xf32

mlir/test/Dialect/Linalg/one-shot-bufferize.mlir

 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
-// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs fully-dynamic-layout-maps=0 bufferize-function-boundaries" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
 
 // TODO: Some test cases from this file should be moved to other dialects.
 
 // CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
 // CHECK-LABEL: func @fill_inplace(
 //  CHECK-SAME:   %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
 // CHECK-NO-LAYOUT-MAP-LABEL: func @fill_inplace(%{{.*}}: memref<?xf32>) {

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
-// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs fully-dynamic-layout-maps=0 bufferize-function-boundaries" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
 // CHECK-LABEL: func @scf_for_yield_only(
 //  CHECK-SAME:   %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>,
 //  CHECK-SAME:   %[[t:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
 //  CHECK-SAME:   ) -> memref<?xf32> {
 func.func @scf_for_yield_only(

mlir/test/Dialect/Tensor/one-shot-bufferize.mlir

 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
-// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs fully-dynamic-layout-maps=0 bufferize-function-boundaries" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
 // CHECK-LABEL: func @insert_slice_fun
 //  CHECK-SAME:   %[[A0:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>,
 //  CHECK-SAME:   %[[A1:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>,
 //  CHECK-SAME:   %[[t0:[a-zA-Z0-9]*]]: memref<4xf32, #[[$map_1d_dyn]]>,
 //  CHECK-SAME:   %[[t1:[a-zA-Z0-9]*]]: memref<4xf32, #[[$map_1d_dyn]]>