[mlir][bufferize] Better analysis for return values of CallOps

[springerm]

Summary

Support returning arbitrary tensors from functions. Even those that are not equivalent. To that end, additional information is gathered during the analysis phase. In particular, which function args are aliasing with which return values.

Also fix bugs in the current implementation when returning equivalent tensors. Various unit tests are added to ensure that we have better test coverage.

Note: Returning non-equivalent tensors is only allowed when allowReturnAllocs is enabled. This functionality is useful for unit testing and compatibility with other bufferizations such as the sparse compiler. This is also towards using ModuleBufferization as a replacement for --func-bufferize.

Depends On D123208

[springerm]

rebase

[springerm]

update

[springerm]

update

[springerm]

rebase

[nicolasvasilache]

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

117

Please add a meaningfully named type, e.g. using MeaningfulTypeName = int64_t

122–123

Please add detailed comment as to why this is needed and the operand / bbArg ordering assumption behind it.

126

This seems like an independent change that could be split out and reduce the load-bearing footprint of this PR.
Please feel free to land the NFC without further review.

470–471

This seems like an independent change that could be split out and reduce the load-bearing footprint of this PR.
Please feel free to land the PR without further review.

[springerm]

update

[springerm]

update

[springerm]

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

122–123

Added comment at the beginning of the struct.

[springerm]

rebase

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

public:
 
   /// Set the inPlace bufferization spec to true.
   /// Merge result's and operand's aliasing sets and iterate to a fixed point.
   void bufferizeInPlace(OpOperand &operand, AnalysisState &state);
 
   /// Set the inPlace bufferization spec to false.
   void bufferizeOutOfPlace(OpOperand &operand);
 
+  /// Return true if `v1` and `v2` may bufferize to aliasing buffers.
+  bool areAliasingBufferizedValues(Value v1, Value v2) const {
+    return aliasInfo.isEquivalent(v1, v2);
+  }
+
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const {
     return equivalentInfo.isEquivalent(v1, v2);
   }
 
   /// Union the alias sets of `v1` and `v2`.
   void unionAliasSets(Value v1, Value v2) { aliasInfo.unionSets(v1, v2); }
 

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

 // function analyzes the given module and determines the order of analysis and
 // bufferization: Functions that are called are processed before their
 // respective callers.
 //
 // After analyzing a FuncOp, additional information about its bbArgs is
 // gathered through PostAnalysisStepFns and stored in
 // `ModuleAnalysisState`.
 //
-// * `equivalentFuncOpBBArgsAnalysis` determines the equivalent bbArg for each
-//   tensor return value (if any).
+// * `aliasingFuncOpBBArgsAnalysis` determines the equivalent/aliasing bbArgs
+// for
+//   each tensor return value (if any).
 // * `funcOpBbArgReadWriteAnalysis` determines whether or not a tensor bbArg is
 //   read/written.
 //
 // Only tensors that are equivalent to some FuncOp bbArg may be returned.
 // Bufferization currently fails if other tensors (in particular tensors that
 // bufferize out-of-place and result in a new buffer allocation) are returned.
 // In the future, such allocations could be hoisted to the caller.
 //

 
 namespace {
 /// The state of analysis of a FuncOp.
 enum class FuncOpAnalysisState { NotAnalyzed, InProgress, Analyzed };
 
 /// Extra analysis state that is required for bufferization of function
 /// boundaries.
 struct ModuleAnalysisState : public DialectAnalysisState {
+  // Note: Function arguments and/or function return values may disappear during
+  // bufferization. Functions and their CallOps are analyzed and bufferized
+  // separately. To ensure that a CallOp analysis/bufferization can access an
+  // already bufferized function's analysis results, we store bbArg/return value
+  // indices instead of BlockArguments/OpOperand pointers.
+
   /// A set of block argument indices.
   using BbArgIndexSet = DenseSet<int64_t>;
 
   /// A mapping of indices to indices.
   using IndexMapping = DenseMap<int64_t, int64_t>;
 
+  /// A mapping of indices to a list of indices.
+  using IndexToIndexListMapping = DenseMap<int64_t, SmallVector<int64_t>>;
+
   /// A mapping of ReturnOp OpOperand indices to equivalent FuncOp BBArg
   /// indices.
   DenseMap<FuncOp, IndexMapping> equivalentFuncArgs;
 
+  /// A mapping of ReturnOp OpOperand indices to aliasing FuncOp BBArg indices.
+  DenseMap<FuncOp, IndexToIndexListMapping> aliasingFuncArgs;
+
+  /// A mapping of FuncOp BBArg indices to aliasing ReturnOp OpOperand indices.
+  DenseMap<FuncOp, IndexToIndexListMapping> aliasingReturnVals;
+
   /// A set of all read BlockArguments of FuncOps.
   DenseMap<FuncOp, BbArgIndexSet> readBbArgs;
 
   /// A set of all written-to BlockArguments of FuncOps.
   DenseMap<FuncOp, BbArgIndexSet> writtenBbArgs;
 
   /// Keep track of which FuncOps are fully analyzed or currently being
   /// analyzed.
   DenseMap<FuncOp, FuncOpAnalysisState> analyzedFuncOps;
 
   /// A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<FuncOp> orderedFuncOps;
 
   /// A mapping of FuncOps to their callers.
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
 
   /// This function is called right before analyzing the given FuncOp. It
   /// initializes the data structures for the FuncOp in this state object.
   void startFunctionAnalysis(FuncOp funcOp) {
     analyzedFuncOps[funcOp] = FuncOpAnalysisState::InProgress;
     auto createdEquiv = equivalentFuncArgs.try_emplace(funcOp, IndexMapping());
+    auto createdAliasingOperands =
+        aliasingFuncArgs.try_emplace(funcOp, IndexToIndexListMapping());
+    auto createdAliasingResults =
+        aliasingReturnVals.try_emplace(funcOp, IndexToIndexListMapping());
     auto createdRead = readBbArgs.try_emplace(funcOp, BbArgIndexSet());
     auto createdWritten = writtenBbArgs.try_emplace(funcOp, BbArgIndexSet());
     (void)createdEquiv;
+    (void)createdAliasingOperands;
+    (void)createdAliasingResults;
     (void)createdRead;
     (void)createdWritten;
 #ifndef NDEBUG
     assert(createdEquiv.second && "equivalence info exists already");
+    assert(createdAliasingOperands.second && "aliasing info exists already");
+    assert(createdAliasingResults.second && "aliasing info exists already");
     assert(createdRead.second && "bbarg access info exists already");
     assert(createdWritten.second && "bbarg access info exists already");
 #endif // NDEBUG
   }
 };
 } // namespace
 
 /// Get ModuleAnalysisState.
if (op->hasAttr(kEquivalentArgsAttr)) {
     equivBbArgs.append(op->getNumOperands(), -1);
   }
   equivBbArgs[returnVal.getOperandNumber()] = bbArg.getArgNumber();
 
   OpBuilder b(op->getContext());
   op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
 }
 
-/// Store function BlockArguments that are equivalent to a returned value in
-/// ModuleAnalysisState.
+/// Store function BlockArguments that are equivalent to/aliasing a returned
+/// value in ModuleAnalysisState.
 static LogicalResult
-equivalentFuncOpBBArgsAnalysis(Operation *op, AnalysisState &state,
-                               BufferizationAliasInfo &aliasInfo,
-                               SmallVector<Operation *> &newOps) {
+aliasingFuncOpBBArgsAnalysis(Operation *op, AnalysisState &state,
   ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
 
   // Support only single return-terminated block in the function.
   auto funcOp = cast<FuncOp>(op);
   func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
   assert(returnOp && "expected func with single return op");
 
   for (OpOperand &returnVal : returnOp->getOpOperands())
     if (returnVal.get().getType().isa<RankedTensorType>())
       for (BlockArgument bbArg : funcOp.getArguments())
-        if (bbArg.getType().isa<RankedTensorType>())
+        if (bbArg.getType().isa<RankedTensorType>()) {
+          int64_t returnIdx = returnVal.getOperandNumber();
+          int64_t bbArgIdx = bbArg.getArgNumber();
           if (aliasInfo.areEquivalentBufferizedValues(returnVal.get(), bbArg)) {
-            moduleState
-                .equivalentFuncArgs[funcOp][returnVal.getOperandNumber()] =
-                bbArg.getArgNumber();
+            moduleState.equivalentFuncArgs[funcOp][returnIdx] = bbArgIdx;
             if (state.getOptions().testAnalysisOnly)
               annotateEquivalentReturnBbArg(returnVal, bbArg);
           }
+          if (aliasInfo.areAliasingBufferizedValues(returnVal.get(), bbArg)) {
+            moduleState.aliasingFuncArgs[funcOp][returnIdx].push_back(bbArgIdx);
+            moduleState.aliasingReturnVals[funcOp][bbArgIdx].push_back(
+                returnIdx);
+          }
+        }
 
   return success();
 }
 
 /// Return true if the buffer of the given tensor value is written to. Must not
 /// be called for values inside not yet analyzed functions. (Post-analysis
 /// steps do not have to be run yet, i.e., "in progress" is also OK.)
 static bool isValueWritten(Value value, const AnalysisState &state,
auto rewrite = [&](Type t) -> Type {
     return t;
   };
   auto argTypes = llvm::to_vector<4>(llvm::map_range(argumentTypes, rewrite));
   auto retTypes = llvm::to_vector<4>(llvm::map_range(resultTypes, rewrite));
   return FunctionType::get(ctx, argTypes, retTypes);
 }
 
 /// Gather equivalence info of CallOps.
-/// Note: This only adds new equivalence info if `funcOp` was already analyzed.
+/// Note: This only adds new equivalence info if the called function was already
+/// analyzed.
 // TODO: This does not handle cyclic function call graphs etc.
 static void equivalenceAnalysis(FuncOp funcOp,
                                 BufferizationAliasInfo &aliasInfo,
                                 ModuleAnalysisState &moduleState) {
   funcOp->walk([&](func::CallOp callOp) {
     FuncOp calledFunction = getCalledFunction(callOp);
     assert(calledFunction && "could not retrieved called FuncOp");
 
static LogicalResult bufferizeFuncOpBoundary(FuncOp funcOp,
   // ============================
   // The body of such functions is assumed opaque and we can't know the
   // bufferization contract they want to enforce atm.
   // As a consequence, only support functions that don't return any tensor atm.
   if (funcOp.getBody().empty()) {
     if (llvm::any_of(funcOp.getFunctionType().getResults(), isaTensor))
       return funcOp->emitError() << "cannot bufferize bodiless function that "
                                  << "returns a tensor";
     FunctionType bufferizedFuncType = getBufferizedFunctionType(
         funcOp.getContext(), funcOp.getFunctionType().getInputs(),
         funcOp.getFunctionType().getResults(), state.getOptions());
     funcOp.setType(bufferizedFuncType);
     return success();
   }
 
   // Support only single return-terminated block in the function.
   func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
   assert(returnOp && "expected func with single return op");
struct CallOpInterface
   }
 
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
     const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
+    if (getFuncOpAnalysisState(state, funcOp) !=
+        FuncOpAnalysisState::Analyzed) {
+      // FuncOp not analyzed yet. Any OpResult may be aliasing.
+      SmallVector<OpResult> result;
+      for (OpResult opResult : op->getOpResults())
+        if (opResult.getType().isa<TensorType>())
+          result.push_back(opResult);
+      return result;
+    }
 
+    // Get aliasing results from state.
+    auto aliasingReturnVals =
+        moduleState.aliasingReturnVals.lookup(funcOp).lookup(
+            opOperand.getOperandNumber());
     SmallVector<OpResult> result;
-    for (int64_t resultIdx = 0; resultIdx < callOp->getNumResults();
-         ++resultIdx)
-      if (Optional<int64_t> maybeArgNumber =
-              getEquivalentFuncArgIdx(funcOp, moduleState, resultIdx))
-        if (*maybeArgNumber == opOperand.getOperandNumber())
-          result.push_back(callOp->getOpResult(resultIdx));
-
+    for (int64_t resultIdx : aliasingReturnVals)
     return result;
   }
 
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
                        const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
     const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
+    if (getFuncOpAnalysisState(state, funcOp) !=
+        FuncOpAnalysisState::Analyzed) {
+      // FuncOp not analyzed yet. Any OpOperand may be aliasing.
+      SmallVector<OpOperand *> result;
+      for (OpOperand &opOperand : op->getOpOperands())
+        if (opOperand.get().getType().isa<TensorType>())
+          result.push_back(&opOperand);
+      return result;
+    }
 
-    // TODO: We should be looking for aliasing block arguments here. The current
-    // condition is actually stronger than neccesary. Once we check for aliasing
-    // block arguments, we may be multiple.
-    if (Optional<int64_t> maybeArgNumber = getEquivalentFuncArgIdx(
-            funcOp, moduleState, opResult.getResultNumber()))
-      return {&op->getOpOperand(*maybeArgNumber)};
-
-    // Note: Returning a non-equivalent tensor from a FuncOp is currently not
-    // supported an will fail bufferization.
-    return {};
+    // Get aliasing bbArgs from state.
+    auto aliasingFuncArgs = moduleState.aliasingFuncArgs.lookup(funcOp).lookup(
+        opResult.getResultNumber());
+    SmallVector<OpOperand *> result;
+    for (int64_t bbArgIdx : aliasingFuncArgs)
+      result.push_back(&callOp->getOpOperand(bbArgIdx));
+    return result;
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   /// In a first approximation, all the function arguments of a FuncOp are
   /// marked inplaceable. For now, it is the responsibility of the `callOp`
   /// bufferization to allow FuncOp that are inplaceable to write inPlace.
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     unsigned numResults = callOp.getNumResults();
     unsigned numOperands = callOp->getNumOperands();
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
     const ModuleAnalysisState &moduleState =
         getModuleAnalysisState(state.getAnalysisState());
+    const OneShotBufferizationOptions &options =
+        static_cast<const OneShotBufferizationOptions &>(state.getOptions());
 
     // Result types of the bufferized CallOp.
     SmallVector<Type> resultTypes;
     // Replacement values for the existing CallOp. These are usually the results
     // of the bufferized CallOp, unless a tensor result folds onto an operand.
     SmallVector<Value> replacementValues(numResults, Value());
     // For non-tensor results: A mapping from return val indices of the old
     // CallOp to return val indices of the bufferized CallOp.
for (const auto &it : llvm::enumerate(callOp.getResultTypes())) {
             state.getBuffer(rewriter, callOp->getOpOperand(*bbArgIdx));
         if (failed(bufferOrFailure))
           return failure();
         replacementValues[returnValIdx] = *bufferOrFailure;
         newOperands[*bbArgIdx] = *bufferOrFailure;
         continue;
       }
 
+      if (!options.allowReturnAllocs)
         return callOp->emitError(
             "call to FuncOp that returns non-equivalent tensors not supported");
+
+      // Returning a memref. This memref is not equivalent to any bbArg. It is
+      // likely a newly allocated buffer. We may want to hoist such allocations
+      // to the call site in the future.
+      retValMapping[returnValIdx] = resultTypes.size();
+      resultTypes.push_back(
+          funcOp.getFunctionType().getResult(resultTypes.size()));
     }
 
     // 2. Compute bufferized FunctionType.
     SmallVector<Type> argumentTypes{callOp->getOperandTypes()};
     // Get the bufferized FunctionType for funcOp or construct it if not yet
     // available.
     FunctionType bufferizedFuncType = getBufferizedFunctionType(
-        funcOp.getContext(), argumentTypes, resultTypes, state.getOptions());
+        funcOp.getContext(), argumentTypes, resultTypes, options);
 
     // 3. Rewrite tensor operands as memrefs based on `bufferizedFuncType`.
     for (OpOperand &opOperand : callOp->getOpOperands()) {
       unsigned idx = opOperand.getOperandNumber();
       Value tensorOperand = opOperand.get();
 
       // Non-tensor operands are just copied.
       if (!tensorOperand.getType().isa<TensorType>()) {
LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
   ModuleAnalysisState &moduleState = getModuleAnalysisState(analysisState);
   BufferizationAliasInfo &aliasInfo = analysisState.getAliasInfo();
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, moduleState.orderedFuncOps,
                                       moduleState.callerMap)))
     return failure();
 
   // Collect bbArg/return value information after the analysis.
-  options.addPostAnalysisStep(equivalentFuncOpBBArgsAnalysis);
+  options.addPostAnalysisStep(aliasingFuncOpBBArgsAnalysis);
   options.addPostAnalysisStep(funcOpBbArgReadWriteAnalysis);
 
   // Analyze ops.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Cross function boundary cases.
 //===----------------------------------------------------------------------===//
 
-//      CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
-
-//      CHECK: memref.global "private" constant @__constant_4xi32 : memref<4xi32> = dense<[1, 2, 3, 4]>
-//      CHECK: func private @some_external_func(memref<4xi32, #[[$DYN_1D_MAP]]>)
-func private @some_external_func(tensor<4xi32>)
-
-//      CHECK: func @main()
-func @main() {
-//  CHECK-DAG:   %[[A:.*]] = memref.get_global @__constant_4xi32 : memref<4xi32>
-  %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
-
-//  CHECK-DAG:   %[[alloc:.*]] = memref.alloc
-//  CHECK-DAG:   %[[B:.*]] = memref.cast %[[alloc]] : memref<4xi32> to memref<4xi32, #[[$DYN_1D_MAP]]>
-//  CHECK-DAG:   memref.copy %[[A]], %[[alloc]]
-//      CHECK:   call @some_external_func(%[[B]]) : (memref<4xi32, #[[$DYN_1D_MAP]]>) -> ()
-  call @some_external_func(%A) : (tensor<4xi32>) -> ()
-
-//      CHECK: memref.dealloc %[[alloc]]
-  return
-}
-
-// -----
-
-//      CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
-
-//      CHECK: memref.global "private" constant @__constant_4xi32 : memref<4xi32> = dense<[1, 2, 3, 4]>
-//      CHECK: func private @some_external_func_within_scf_execute(memref<4xi32, #[[$DYN_1D_MAP]]>)
-func private @some_external_func_within_scf_execute(tensor<4xi32>)
-
-//      CHECK: func @main()
-func @main() {
-//  CHECK-DAG:   %[[A:.*]] = memref.get_global @__constant_4xi32 : memref<4xi32>
-  %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
-
-//  CHECK-DAG:   %[[alloc:.*]] = memref.alloc
-//  CHECK-DAG:   %[[B:.*]] = memref.cast %[[alloc]] : memref<4xi32> to memref<4xi32, #[[$DYN_1D_MAP]]>
-//  CHECK-DAG:   memref.copy %[[A]], %[[alloc]]
-//      CHECK:   call @some_external_func_within_scf_execute(%[[B]]) : (memref<4xi32, #[[$DYN_1D_MAP]]>) -> ()
-  scf.execute_region {
-    call @some_external_func_within_scf_execute(%A) : (tensor<4xi32>) -> ()
-    scf.yield
-  }
-
-//      CHECK:   memref.dealloc %[[alloc]]
-  return
-}
-
-// -----
-
-// CHECK: func private @external_func_with_return_val(memref<4xi32, #{{.*}}>) -> f32
-func private @external_func_with_return_val(tensor<4xi32>) -> f32
-
-// -----
-
-// CHECK-LABEL: func @execute_region_test(
-//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
-func @execute_region_test(%t1 : tensor<?xf32> {linalg.inplaceable = "true"})
-    -> (f32, tensor<?xf32>, f32)
-{
-  %f1 = arith.constant 0.0 : f32
-  %f2 = arith.constant 1.0 : f32
-  %idx = arith.constant 7 : index
-
-  // scf.execute_region is canonicalized away after bufferization. So just the
-  // memref.store is left over.
-
-  // CHECK: memref.store %{{.*}}, %[[m1]][%{{.*}}]
-  %0, %1, %2 = scf.execute_region -> (f32, tensor<?xf32>, f32) {
-    %t2 = tensor.insert %f2 into %t1[%idx] : tensor<?xf32>
-    scf.yield %f1, %t2, %f2 : f32, tensor<?xf32>, f32
-  }
-
-  // CHECK: return %{{.*}}, %{{.*}} : f32, f32
-  return %0, %1, %2 : f32, tensor<?xf32>, f32
-}
-
-// -----
-
 // CHECK-LABEL: func @execute_region_with_conflict(
 //  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
 func @execute_region_with_conflict(%t1 : tensor<?xf32> {linalg.inplaceable = "true"})
     -> (f32, tensor<?xf32>, f32)
 {
   %f1 = arith.constant 0.0 : f32
   %idx = arith.constant 7 : index
 
func @execute_region_with_conflict(%t1 : tensor<?xf32> {linalg.inplaceable = "true"})
   %3 = tensor.extract %t1[%idx] : tensor<?xf32>
 
   // CHECK: return %{{.*}}, %[[casted]], %[[load]] : f32, memref<?xf32, #{{.*}}>, f32
   return %0, %1, %3 : f32, tensor<?xf32>, f32
 }
 
 // -----
 
-//      CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
-
-//      CHECK:  func private @some_external_func(memref<?xf32, #[[$DYN_1D_MAP]]>)
-func private @some_external_func(tensor<?xf32>)
-
-//      CHECK:  func @scf_for_with_tensor_insert_slice(
-// CHECK-SAME:    %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]>
-// CHECK-SAME:    %[[B:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]>
-// CHECK-SAME:    %[[C:[a-zA-Z0-9]*]]: memref<4xf32, #[[$DYN_1D_MAP]]>
-func @scf_for_with_tensor_insert_slice(
-    %A : tensor<?xf32>, %B : tensor<?xf32>, %C : tensor<4xf32>,
-    %lb : index, %ub : index, %step : index)
-  -> (tensor<?xf32>, tensor<?xf32>)
-{
-  // CHECK-NEXT: scf.for
-  %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
-      -> (tensor<?xf32>, tensor<?xf32>)
-  {
-    // CHECK-NEXT:   %[[SVA:.*]] = memref.subview %[[A]]
-    // CHECK-NEXT:   memref.copy %[[C]], %[[SVA]] : memref<4xf32, #[[$DYN_1D_MAP]]> to memref<4xf32, #[[$DYN_1D_MAP]]>
-    %ttA = tensor.insert_slice %C into %tA[%i][4][1] : tensor<4xf32> into tensor<?xf32>
-
-    // CHECK-NEXT:   %[[SVB:.*]] = memref.subview %[[B]]
-    // CHECK-NEXT:   memref.copy %[[C]], %[[SVB]] : memref<4xf32, #[[$DYN_1D_MAP]]> to memref<4xf32, #[[$DYN_1D_MAP]]>
-    %ttB = tensor.insert_slice %C into %tB[%i][4][1] : tensor<4xf32> into tensor<?xf32>
-
-    // scf.yield is empty and is elided
-    //  CHECK-NOT:   scf.yield
-    scf.yield %ttA, %ttB : tensor<?xf32>, tensor<?xf32>
-  }
-
-  // Swaparoo requires bufferizing the whole function to figure out who's who.
-  return %r0#1, %r0#0: tensor<?xf32>, tensor<?xf32>
-}
-
-//      CHECK:  func @bar(
-// CHECK-SAME:    %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]>
-// CHECK-SAME:    %[[B:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]>
-// CHECK-SAME:    %[[C:[a-zA-Z0-9]*]]: memref<4xf32, #[[$DYN_1D_MAP]]>
-func @bar(
-    %A : tensor<?xf32> {linalg.inplaceable = true},
-    %B : tensor<?xf32> {linalg.inplaceable = true},
-    %C : tensor<4xf32> {linalg.inplaceable = true},
-    %lb : index, %ub : index, %step : index)
-  -> (tensor<?xf32>, tensor<?xf32>)
-{
-//  CHECK-DAG:   call @scf_for_with_tensor_insert_slice(%[[A]], %[[B]], %[[C]]
-  %r0:2 = call @scf_for_with_tensor_insert_slice(%A, %B, %C, %lb, %ub, %step) :
-      (tensor<?xf32>, tensor<?xf32>, tensor<4xf32>, index, index, index)
-        -> (tensor<?xf32>, tensor<?xf32>)
-
-  // %r0#0 requires a copy because we have no idea what the function is doing.
-//  CHECK-DAG:   %[[alloc:.*]] = memref.alloc
-//  CHECK-DAG:   %[[casted:.*]] = memref.cast %[[alloc]]
-//      CHECK:   memref.copy %[[B]], %[[alloc]]
-// CHECK-NEXT:   call @some_external_func(%[[casted]]) : (memref<?xf32, #[[$DYN_1D_MAP]]>) -> ()
-  call @some_external_func(%r0#0) : (tensor<?xf32>) -> ()
-
-//      CHECK:   return
-  return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
-}
-
-// -----
-
-//  CHECK-DAG: #[[$DYN_0D_MAP:.*]] = affine_map<()[s0] -> (s0)>
-//  CHECK-DAG: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
-
-//      CHECK:  func @init_and_dot(
-// CHECK-SAME:    %[[A:[a-zA-Z0-9]*]]: memref<64xf32, #[[$DYN_1D_MAP]]>
-// CHECK-SAME:    %[[B:[a-zA-Z0-9]*]]: memref<64xf32, #[[$DYN_1D_MAP]]>
-// CHECK-SAME:    %[[C:[a-zA-Z0-9]*]]: memref<f32, #[[$DYN_0D_MAP]]>
-func @init_and_dot(%a: tensor<64xf32>, %b: tensor<64xf32>, %c: tensor<f32>) -> tensor<f32> {
-  // CHECK-NEXT:   %[[C0:.*]] = arith.constant 0{{.*}} : f32
-  %v0 = arith.constant 0.0 : f32
-
-  // CHECK-NEXT:   linalg.fill ins(%[[C0]] : f32) outs(%[[C]] : memref<f32, #[[$DYN_0D_MAP]]>)
-  %d = linalg.fill ins(%v0 : f32) outs(%c : tensor<f32>) -> tensor<f32>
-
-  // CHECK-NEXT:   linalg.dot ins(%[[A]], %[[B]] : memref<64xf32, #[[$DYN_1D_MAP]]>, memref<64xf32, #[[$DYN_1D_MAP]]>) outs(%[[C]] : memref<f32, #[[$DYN_0D_MAP]]>)
-  %e = linalg.dot ins(%a, %b : tensor<64xf32>,tensor<64xf32>)
-    outs(%d: tensor<f32>) -> tensor<f32>
-
-  // CHECK-NEXT:   return
-  return %e : tensor<f32>
-}
-
-//      CHECK:  func @main()
-func @main() {
-  //  CHECK-DAG:   %[[C0:.*]] = arith.constant 0{{.*}} : f32
-  //  CHECK-DAG:   %[[C1:.*]] = arith.constant 1{{.*}} : f32
-  //  CHECK-DAG:   %[[C2:.*]] = arith.constant 2{{.*}} : f32
-  %v0 = arith.constant 0.0 : f32
-  %v1 = arith.constant 1.0 : f32
-  %v2 = arith.constant 2.0 : f32
-
-  // CHECK-NEXT:   %[[A:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
-  // CHECK-NEXT:   %[[B:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
-  // CHECK-NEXT:   %[[C:.*]] = memref.alloc() {alignment = 128 : i64} : memref<f32>
-  //  CHECK-DAG:   %[[cA:.*]] = memref.cast %[[A]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]>
-  //  CHECK-DAG:   %[[cB:.*]] = memref.cast %[[B]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]>
-  //  CHECK-DAG:   %[[cC:.*]] = memref.cast %[[C]] : memref<f32> to memref<f32, #[[$DYN_0D_MAP]]>
-  %A = linalg.init_tensor [64] : tensor<64xf32>
-  %B = linalg.init_tensor [64] : tensor<64xf32>
-  %C = linalg.init_tensor [] : tensor<f32>
-
-  //  CHECK-DAG:   linalg.fill ins(%[[C1]] : f32) outs(%[[A]] : memref<64xf32>)
-  //  CHECK-DAG:   linalg.fill ins(%[[C2]] : f32) outs(%[[B]] : memref<64xf32>)
-  //  CHECK-DAG:   linalg.fill ins(%[[C0]] : f32) outs(%[[C]] : memref<f32>)
-  %AA = linalg.fill ins(%v1 : f32) outs(%A : tensor<64xf32>) -> tensor<64xf32>
-  %BB = linalg.fill ins(%v2 : f32) outs(%B : tensor<64xf32>) -> tensor<64xf32>
-  %CC = linalg.fill ins(%v0 : f32) outs(%C : tensor<f32>) -> tensor<f32>
-
-  // CHECK-NEXT:   call @init_and_dot(%[[cA]], %[[cB]], %[[cC]])
-  %res = call @init_and_dot(%AA, %BB, %CC) :
-    (tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32>
-
-  // CHECK-NEXT:   %[[dC:.*]] = memref.cast %[[C]] : memref<f32> to memref<*xf32>
-  %res2 = tensor.cast %res: tensor<f32> to tensor<*xf32>
-
-  // CHECK-NEXT:   call @print_memref_f32(%[[dC]]) : (memref<*xf32>) -> ()
-  call @print_memref_f32(%res2) : (tensor<*xf32>) -> ()
-
-  // CHECK-DAG:   memref.dealloc %[[A]] : memref<64xf32>
-  // CHECK-DAG:   memref.dealloc %[[B]] : memref<64xf32>
-  // CHECK-DAG:   memref.dealloc %[[C]] : memref<f32>
-  // CHECK-NEXT:   return
-  return
-}
-
-//     CHECK:   func private @print_memref_f32(memref<*xf32>)
-func private @print_memref_f32(tensor<*xf32>)
-
-// -----
-
-// CHECK: #[[$DYNAMIC:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
-
-// CHECK: func private @external_func(memref<?xf32, #[[$DYNAMIC]]>)
-func private @external_func(tensor<?xf32>)
-
-//      CHECK: func @callee(
-// CHECK-SAME:   %[[A:[0-9a-zA-Z]*]]: memref<?xf32>
-// CHECK-SAME:   %[[B:[0-9a-zA-Z]*]]: memref<?xf32, #[[$DYNAMIC]]>
-// CHECK-SAME:   %[[C:[0-9a-zA-Z]*]]: memref<?xf32, #[[$DYNAMIC]]>
-func @callee(%A : tensor<?xf32> {linalg.buffer_layout = affine_map<(i)[s0, s1] -> (i)>},
-             %B : tensor<?xf32>,
-             %C : tensor<?xf32>) {
-// CHECK-NEXT: %[[CASTED:.*]] = memref.cast %[[A]] : memref<?xf32> to memref<?xf32, #[[$DYNAMIC]]>
-// CHECK-NEXT: call @external_func(%[[CASTED]]) : (memref<?xf32, #[[$DYNAMIC]]>) -> ()
-  call @external_func(%A) : (tensor<?xf32>) -> ()
-
-// CHECK-NEXT: call @external_func(%[[B]]) : (memref<?xf32, #[[$DYNAMIC]]>) -> ()
-  call @external_func(%B) : (tensor<?xf32>) -> ()
-
-// CHECK-NEXT: call @external_func(%[[C]]) : (memref<?xf32, #[[$DYNAMIC]]>) -> ()
-  call @external_func(%C) : (tensor<?xf32>) -> ()
-
-  return
-}
-
-//      CHECK: func @entry(
-// CHECK-SAME:   %[[A:[0-9a-zA-Z]*]]: memref<?xf32>
-// CHECK-SAME:   %[[B:[0-9a-zA-Z]*]]: memref<?xf32>
-// CHECK-SAME:   %[[C:[0-9a-zA-Z]*]]: memref<?xf32, #[[$DYNAMIC]]>
-func @entry(%A : tensor<?xf32> {linalg.buffer_layout = affine_map<(i)[s0, s1] -> (i)>, linalg.inplaceable = false},
-            %B : tensor<?xf32> {linalg.buffer_layout = affine_map<(i)[s0, s1] -> (i)>, linalg.inplaceable = false},
-            %C : tensor<?xf32> {linalg.inplaceable = false}) {
-// Note: `callee` does not write to its bbArg directly, but `external_func`
-// does. Inside `callee`, the writes via `external_func` do not cause a
-// conflict. However, inside `entry`, the writes do cause a conflict because
-// %A, %B and %C are not inplaceable. This test case shows that this kind of
-// conflict detection has a "transitive" nature.
-//      CHECK: %[[ALLOC_C:.*]] = memref.alloc
-//      CHECK: %[[CASTED_C:.*]] = memref.cast %[[ALLOC_C]]
-//      CHECK: %[[ALLOC_B:.*]] = memref.alloc
-//      CHECK: %[[CASTED_B:.*]] = memref.cast %[[ALLOC_B]]
-//      CHECK: %[[ALLOC_A:.*]] = memref.alloc
-//      CHECK: memref.copy %[[A]], %[[ALLOC_A]]
-//      CHECK: memref.copy %[[B]], %[[ALLOC_B]]
-//      CHECK: memref.copy %[[C]], %[[ALLOC_C]]
-//      CHECK: %[[CASTED_A:.*]] = memref.cast %[[ALLOC_A]]
-// CHECK-NEXT: call @callee(%[[CASTED_A]], %[[CASTED_B]], %[[CASTED_C]])
-  call @callee(%A, %B, %C) : (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) -> ()
-  return
-}
-
-// -----
-
 //      CHECK: func @matmul(
 // CHECK-SAME:   %[[A:[0-9a-zA-Z]*]]: memref<128x256xf32>
 // CHECK-SAME:   %[[B:[0-9a-zA-Z]*]]: memref<256x192xf32>
 // CHECK-SAME:   %[[C:[0-9a-zA-Z]*]]: memref<128x192xf32>
 func @matmul(
     %A: tensor<128x256xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
     %B: tensor<256x192xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
     %C: tensor<128x192xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = true})
func @insert_op(%t1 : tensor<?xf32> {linalg.inplaceable = true},
   // CHECK: memref.store %[[s]], %[[t1]][%[[i]]]
   %0 = tensor.insert %s into %t1[%i] : tensor<?xf32>
   // CHECK: return
   return %0 : tensor<?xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @inner_func(
-//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
-func @inner_func(%t: tensor<?xf32>) -> tensor<?xf32> {
-  %f = arith.constant 1.0 : f32
-  %c0 = arith.constant 0 : index
-  // CHECK: memref.store %{{.*}}, %[[arg0]]
-  %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>
-  return %0 : tensor<?xf32>
-}
-
-// CHECK-LABEL: func @equivalent_func_arg(
-//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
-func @equivalent_func_arg(%t0: tensor<?xf32> {linalg.inplaceable = true},
-                          %c0: index, %c10: index, %c1: index) -> tensor<?xf32> {
-  // CHECK-NOT: copy
-  %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {
-    // CHECK: call @inner_func(%[[arg0]])
-    %3 = call @inner_func(%t1) : (tensor<?xf32>) -> tensor<?xf32>
-    scf.yield %3 : tensor<?xf32>
-  }
-  return %1: tensor<?xf32>
-}
-
-// -----
-
-// CHECK-LABEL: func @inner_func_2(
-//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
-func @inner_func_2(%t: tensor<?xf32>) -> tensor<?xf32> {
-  %f = arith.constant 1.0 : f32
-  %c0 = arith.constant 0 : index
-  // CHECK: memref.store %{{.*}}, %[[arg0]]
-  %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>
-  return %0 : tensor<?xf32>
-}
-
-// CHECK-LABEL: func @equivalent_func_arg_2(
-//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
-func @equivalent_func_arg_2(%t0: tensor<?xf32> {linalg.inplaceable = true},
-                            %c0: index, %c10: index, %c1: index) -> tensor<?xf32> {
-  %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {
-    // CHECK: %[[alloc:.*]] = memref.alloc
-    // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
-    // CHECK: memref.copy %[[arg0]], %[[alloc]]
-    // CHECK: call @inner_func_2(%[[casted]])
-    %3 = call @inner_func_2(%t1) : (tensor<?xf32>) -> tensor<?xf32>
-    scf.yield %t1 : tensor<?xf32>
-  }
-  return %1: tensor<?xf32>
-}
-
-// -----
-
-// CHECK-LABEL: func @inner_func(
-//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
-func @inner_func(%t: tensor<?xf32>) -> (tensor<?xf32>, f32) {
-  // CHECK-NOT: copy
-  %f = arith.constant 1.0 : f32
-  %c0 = arith.constant 0 : index
-  %c1 = arith.constant 1 : index
-  // CHECK: memref.store %{{.*}}, %[[arg0]]
-  %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>
-  // CHECK: %[[load:.*]] = memref.load %[[arg0]]
-  %1 = tensor.extract %0[%c1] : tensor<?xf32>
-  // CHECK: return %[[load]] : f32
-  return %0, %1 : tensor<?xf32>, f32
-}
-
-// CHECK-LABEL: func @call_func_with_non_tensor_return(
-//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
-func @call_func_with_non_tensor_return(
-    %t0: tensor<?xf32> {linalg.inplaceable = true}) -> (f32, tensor<?xf32>) {
-  // CHECK-NOT: copy
-  // CHECK: %[[call:.*]] = call @inner_func(%[[arg0]])
-  %0, %1 = call @inner_func(%t0) : (tensor<?xf32>) -> (tensor<?xf32>, f32)
-  // CHECK: return %[[call]] : f32
-  return %1, %0 : f32, tensor<?xf32>
-}
-
-// -----
-
-// CHECK-LABEL: func @func_without_tensor_args
-func @func_without_tensor_args(%v : vector<10xf32>) -> () {
-  // CHECK: %[[alloc:.*]] = memref.alloc()
-  %0 = linalg.init_tensor[10] : tensor<10xf32>
-
-  %c0 = arith.constant 0 : index
-  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
-  %1 = vector.transfer_write %v, %0[%c0] : vector<10xf32>, tensor<10xf32>
-
-  %cst = arith.constant 0.0 : f32
-  // CHECK: vector.transfer_read %[[alloc]]
-  %r = vector.transfer_read %1[%c0], %cst : tensor<10xf32>, vector<11xf32>
-
-  vector.print %r : vector<11xf32>
-  return
-}
-
-// -----
-
-// CHECK-LABEL: func private @private_func
-func private @private_func(tensor<?xf32>) -> ()
-
-// CHECK-LABEL: func @empty_func()
-func @empty_func() -> () {
-  return
-}
-
-// -----
-
 func @gather_like(
     %arg0 : tensor<?x?xf32> {linalg.inplaceable = false},
     %arg1 : tensor<?xi32> {linalg.inplaceable = false},
     %arg2 : tensor<?x?xf32> {linalg.inplaceable = true}) -> tensor<?x?xf32> {
   %0 = linalg.generic {
       indexing_maps = [affine_map<(d0, d1) -> (d0)>,
                        affine_map<(d0, d1) -> (d0, d1)>],
       iterator_types = ["parallel", "parallel"]}

 //       CHECK:     memref.dealloc %[[for]]#0
 //       CHECK:     %[[r1:.*]] = memref.load %[[for]]#1
 //       CHECK:     memref.dealloc %[[for]]#1
   %f0 = tensor.extract %r0#0[%step] : tensor<?xf32>
   %f1 = tensor.extract %r0#1[%step] : tensor<?xf32>
 //       CHECK:     return %[[r0]], %[[r1]]
   return %f0, %f1: f32, f32
 }
+

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

+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize=allow-return-allocs -split-input-file | FileCheck %s
+
+// Run fuzzer with different seeds.
+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="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 -linalg-comprehensive-module-bufferize="allow-return-allocs fully-dynamic-layout-maps=0" -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]]
+func @create_tensor() -> tensor<10xf32> {
+  %0 = linalg.init_tensor [10] : tensor<10xf32>
+  return %0 : tensor<10xf32>
+}
+
+// CHECK: func @caller(
+// CHECK: %[[call:.*]] = call @create_tensor() : () -> memref<10xf32>
+// CHECK: %[[extracted:.*]] = memref.load %[[call]]
+// CHECK: return %[[extracted]]
+func @caller(%idx: index) -> f32 {
+  %0 = call @create_tensor() : () -> (tensor<10xf32>)
+  %1 = tensor.extract %0[%idx] : tensor<10xf32>
+  return %1 : f32
+}
+
+// -----
+
+// return_slice returns an aliasing tensor. In main, %t is overwritten (but not
+// read). This is a conflict because %0 is aliasing with %t. An alloc + copy is
+// needed.
+
+// CHECK-LABEL: func @return_slice(
+//   CHECK-NOT:   alloc
+//   CHECK-NOT:   copy
+//       CHECK:   memref.subview
+func @return_slice(%t: tensor<?xf32>, %sz: index) -> (tensor<?xf32>) {
+  %0 = tensor.extract_slice %t[4][%sz][1] : tensor<?xf32> to tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL: func @main(
+//  CHECK-SAME:     %[[t:.*]]: memref<?xf32
+//       CHECK:   %[[alloc:.*]] = memref.alloc
+//   CHECK-DAG:   memref.copy %[[t]], %[[alloc]]
+//   CHECK-DAG:   %[[casted:.*]] = memref.cast %[[alloc]]
+//       CHECK:   %[[call:.*]] = call @return_slice(%[[casted]]
+//       CHECK:   linalg.fill ins({{.*}}) outs(%[[t]]
+//       CHECK:   memref.load %[[call]]
+//       CHECK:   memref.load %[[t]]
+func @main(%t: tensor<?xf32>, %sz: index, %idx: index) -> (f32, f32) {
+  %cst = arith.constant 1.0 : f32
+  %0 = call @return_slice(%t, %sz) : (tensor<?xf32>, index) -> (tensor<?xf32>)
+  %filled = linalg.fill ins(%cst : f32) outs(%t : tensor<?xf32>) -> tensor<?xf32>
+  %r1 = tensor.extract %0[%idx] : tensor<?xf32>
+  %r2 = tensor.extract %filled[%idx] : tensor<?xf32>
+  return %r1, %r2 : f32, f32
+}

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

+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize -split-input-file | FileCheck %s
+
+// Run fuzzer with different seeds.
+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize="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 -linalg-comprehensive-module-bufferize="allow-return-allocs fully-dynamic-layout-maps=0" -split-input-file -o /dev/null
+
+// Bufferization of bodiless function with no tensor return value.
+
+// CHECK-LABEL: func private @private_func
+func private @private_func(tensor<?xf32>) -> ()
+
+// CHECK-LABEL: func @empty_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
+func private @external_func_with_return_val(tensor<4xi32>) -> f32
+
+// -----
+
+// CHECK-LABEL: func private @private_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
+//   CHECK-NOT: alloc
+//   CHECK-NOT: copy
+//       CHECK: call @private_func(%[[t]])
+func @main(%t: tensor<?xf32> {linalg.inplaceable = true}) -> (f32) {
+  %0 = call @private_func(%t) : (tensor<?xf32>) -> (f32)
+  return %0 : f32
+}
+
+// -----
+
+// CHECK-LABEL: func private @private_func
+func private @private_func(tensor<?xf32>) -> (f32)
+
+// private_func may modify the buffer arg, %t is not writable. A copy is needed.
+
+// CHECK-LABEL: func @main(
+//  CHECK-SAME:     %[[t:.*]]: memref<?xf32
+//       CHECK: %[[alloc:.*]] = memref.alloc
+//   CHECK-DAG: memref.copy %[[t]], %[[alloc]]
+//   CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]
+//       CHECK: call @private_func(%[[casted]])
+//       CHECK: memref.dealloc %[[alloc]]
+func @main(%t: tensor<?xf32> {linalg.inplaceable = false}) -> (f32) {
+  %0 = call @private_func(%t) : (tensor<?xf32>) -> (f32)
+  return %0 : f32
+}
+
+// -----
+
+// Test bufferization of a function without tensor args.
+
+// CHECK-LABEL: func @func_without_tensor_args
+func @func_without_tensor_args(%v : vector<10xf32>) -> () {
+  // CHECK: %[[alloc:.*]] = memref.alloc()
+  %0 = linalg.init_tensor[10] : tensor<10xf32>
+
+  %c0 = arith.constant 0 : index
+  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
+  %1 = vector.transfer_write %v, %0[%c0] : vector<10xf32>, tensor<10xf32>
+
+  %cst = arith.constant 0.0 : f32
+  // CHECK: vector.transfer_read %[[alloc]]
+  %r = vector.transfer_read %1[%c0], %cst : tensor<10xf32>, vector<11xf32>
+
+  vector.print %r : vector<11xf32>
+  return
+}
+
+// -----
+
+// Bufferization of a function that is reading and writing. %t0 is writable, so
+// no copy should be inserted.
+
+// CHECK-LABEL: func @inner_func(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
+func @inner_func(%t: tensor<?xf32>) -> (tensor<?xf32>, f32) {
+  // CHECK-NOT: copy
+  %f = arith.constant 1.0 : f32
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  // CHECK: memref.store %{{.*}}, %[[arg0]]
+  %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>
+  // CHECK: %[[load:.*]] = memref.load %[[arg0]]
+  %1 = tensor.extract %0[%c1] : tensor<?xf32>
+  // CHECK: return %[[load]] : f32
+  return %0, %1 : tensor<?xf32>, f32
+}
+
+// CHECK-LABEL: func @call_func_with_non_tensor_return(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
+func @call_func_with_non_tensor_return(
+    %t0: tensor<?xf32> {linalg.inplaceable = true}) -> (f32, tensor<?xf32>) {
+  // CHECK-NOT: alloc
+  // CHECK-NOT: copy
+  // CHECK: %[[call:.*]] = call @inner_func(%[[arg0]])
+  %0, %1 = call @inner_func(%t0) : (tensor<?xf32>) -> (tensor<?xf32>, f32)
+  // CHECK: return %[[call]] : f32
+  return %1, %0 : f32, tensor<?xf32>
+}
+
+// -----
+
+// Bufferization of a function that is reading and writing. %t0 is not writable,
+// so a copy is needed.
+
+// CHECK-LABEL: func @inner_func(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
+func @inner_func(%t: tensor<?xf32>) -> (tensor<?xf32>, f32) {
+  // CHECK-NOT: copy
+  %f = arith.constant 1.0 : f32
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  // CHECK: memref.store %{{.*}}, %[[arg0]]
+  %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>
+  // CHECK: %[[load:.*]] = memref.load %[[arg0]]
+  %1 = tensor.extract %0[%c1] : tensor<?xf32>
+  // CHECK: return %[[load]] : f32
+  return %0, %1 : tensor<?xf32>, f32
+}
+
+// CHECK-LABEL: func @call_func_with_non_tensor_return(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
+func @call_func_with_non_tensor_return(
+    %t0: tensor<?xf32> {linalg.inplaceable = false}) -> (f32, tensor<?xf32>) {
+  // CHECK: %[[alloc:.*]] = memref.alloc
+  // CHECK-DAG: memref.copy %[[arg0]], %[[alloc]]
+  // CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]
+  // CHECK: %[[call:.*]] = call @inner_func(%[[casted]])
+  %0, %1 = call @inner_func(%t0) : (tensor<?xf32>) -> (tensor<?xf32>, f32)
+
+  // Note: The tensor return value has folded away.
+  // CHECK: return %[[call]] : f32
+  return %1, %0 : f32, tensor<?xf32>
+}
+
+// -----
+
+// A chain of function calls. The last function f0 is potentially writing to the
+// buffer. This becomes a problem when bufferizing main and a copy must be
+// inserted then. (No copies in the other functions.)
+
+// CHECK-LABEL: func private @f0(
+func private @f0(tensor<?xf32>) -> (f32)
+
+// CHECK-LABEL: func @f1(
+//  CHECK-SAME:     %[[t1:.*]]: memref<?xf32
+//       CHECK:   %[[r1:.*]] = call @f0(%[[t1]])
+//       CHECK:   return %[[r1]]
+func @f1(%t: tensor<?xf32>) -> (f32) {
+  %0 = call @f0(%t) : (tensor<?xf32>) -> (f32)
+  return %0 : f32
+}
+
+// CHECK-LABEL: func @f2(
+//  CHECK-SAME:     %[[t2:.*]]: memref<?xf32
+//       CHECK:   %[[r2:.*]] = call @f1(%[[t2]])
+//       CHECK:   return %[[r2]]
+func @f2(%t: tensor<?xf32>) -> (f32) {
+  %0 = call @f1(%t) : (tensor<?xf32>) -> (f32)
+  return %0 : f32
+}
+
+// CHECK-LABEL: func @main(
+//  CHECK-SAME:     %[[t3:.*]]: memref<?xf32
+//       CHECK: %[[alloc:.*]] = memref.alloc
+//   CHECK-DAG: memref.copy %[[t3]], %[[alloc]]
+//   CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]
+//       CHECK: call @f2(%[[casted]])
+//       CHECK: memref.dealloc %[[alloc]]
+func @main(%t: tensor<?xf32> {linalg.inplaceable = false}) -> (f32) {
+  %0 = call @f2(%t) : (tensor<?xf32>) -> (f32)
+  return %0 : f32
+}
+
+// -----
+
+// This function does not read, just write. We need an alloc, but no copy.
+
+// CHECK-LABEL: func @does_not_read(
+//   CHECK-NOT:   alloc
+//   CHECK-NOT:   copy
+func @does_not_read(%t: tensor<?xf32>) -> tensor<?xf32> {
+  %f0 = arith.constant 0.0 : f32
+  %r = linalg.fill ins(%f0 : f32) outs(%t : tensor<?xf32>) -> tensor<?xf32>
+  return %r : tensor<?xf32>
+}
+
+// CHECK-LABEL: func @main(
+//  CHECK-SAME:     %[[t:.*]]: memref<?xf32
+//       CHECK:   %[[alloc:.*]] = memref.alloc
+//   CHECK-NOT:   copy
+//       CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
+//   CHECK-NOT:   copy
+//       CHECK:   call @does_not_read(%[[casted]])
+//       CHECK:   %[[r:.*]] = memref.load %[[alloc]]
+//       CHECK:   memref.dealloc %[[alloc]]
+func @main(%t: tensor<?xf32> {linalg.inplaceable = false}) -> f32 {
+  %0 = call @does_not_read(%t) : (tensor<?xf32>) -> (tensor<?xf32>)
+  %idx = arith.constant 4 : index
+  %r = tensor.extract %0[%idx] : tensor<?xf32>
+  return %r : f32
+}
+
+// -----
+
+// Alloc and copy must be inserted because the arith.constant is read-only.
+
+//      CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
+//      CHECK: memref.global "private" constant @__constant_4xi32 : memref<4xi32> = dense<[1, 2, 3, 4]>
+//      CHECK: func private @some_external_func(memref<4xi32, #[[$DYN_1D_MAP]]>)
+func private @some_external_func(tensor<4xi32>)
+
+//      CHECK: func @main()
+func @main() {
+//  CHECK-DAG:   %[[A:.*]] = memref.get_global @__constant_4xi32 : memref<4xi32>
+  %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+
+//  CHECK-DAG:   %[[alloc:.*]] = memref.alloc
+//  CHECK-DAG:   %[[B:.*]] = memref.cast %[[alloc]] : memref<4xi32> to memref<4xi32, #[[$DYN_1D_MAP]]>
+//  CHECK-DAG:   memref.copy %[[A]], %[[alloc]]
+//      CHECK:   call @some_external_func(%[[B]]) : (memref<4xi32, #[[$DYN_1D_MAP]]>) -> ()
+  call @some_external_func(%A) : (tensor<4xi32>) -> ()
+
+//      CHECK: memref.dealloc %[[alloc]]
+  return
+}
+
+// -----
+
+// Alloc and copy must be inserted because the arith.constant is read-only. The
+// function call is inside of an scf.execute_region.
+
+//      CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
+//      CHECK: memref.global "private" constant @__constant_4xi32 : memref<4xi32> = dense<[1, 2, 3, 4]>
+//      CHECK: func private @some_external_func_within_scf_execute(memref<4xi32, #[[$DYN_1D_MAP]]>)
+func private @some_external_func_within_scf_execute(tensor<4xi32>)
+
+//      CHECK: func @main()
+func @main() {
+//  CHECK-DAG:   %[[A:.*]] = memref.get_global @__constant_4xi32 : memref<4xi32>
+  %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+
+// Note: The scf.execute_region canonicalizes away.
+
+//  CHECK-DAG:   %[[alloc:.*]] = memref.alloc
+//  CHECK-DAG:   %[[B:.*]] = memref.cast %[[alloc]] : memref<4xi32> to memref<4xi32, #[[$DYN_1D_MAP]]>
+//  CHECK-DAG:   memref.copy %[[A]], %[[alloc]]
+//      CHECK:   call @some_external_func_within_scf_execute(%[[B]]) : (memref<4xi32, #[[$DYN_1D_MAP]]>) -> ()
+  scf.execute_region {
+    call @some_external_func_within_scf_execute(%A) : (tensor<4xi32>) -> ()
+    scf.yield
+  }
+
+//      CHECK:   memref.dealloc %[[alloc]]
+  return
+}
+
+// -----
+
+// A write inside an scf.execute_region. An equivalent tensor is yielded.
+
+// CHECK-LABEL: func @execute_region_test(
+//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
+func @execute_region_test(%t1 : tensor<?xf32>)
+    -> (f32, tensor<?xf32>, f32)
+{
+  %f1 = arith.constant 0.0 : f32
+  %f2 = arith.constant 1.0 : f32
+  %idx = arith.constant 7 : index
+
+  // scf.execute_region is canonicalized away after bufferization. So just the
+  // memref.store is left over.
+
+  // CHECK-NOT: alloc
+  // CHECK-NOT: copy
+  // CHECK: memref.store %{{.*}}, %[[m1]][%{{.*}}]
+  %0, %1, %2 = scf.execute_region -> (f32, tensor<?xf32>, f32) {
+    %t2 = tensor.insert %f2 into %t1[%idx] : tensor<?xf32>
+    scf.yield %f1, %t2, %f2 : f32, tensor<?xf32>, f32
+  }
+
+  // CHECK: return %{{.*}}, %{{.*}} : f32, f32
+  return %0, %1, %2 : f32, tensor<?xf32>, f32
+}
+
+// -----
+
+//      CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
+//      CHECK:  func private @some_external_func(memref<?xf32, #[[$DYN_1D_MAP]]>)
+func private @some_external_func(tensor<?xf32>)
+
+//      CHECK:  func @scf_for_with_tensor_insert_slice(
+// CHECK-SAME:    %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]>
+// CHECK-SAME:    %[[B:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]>
+// CHECK-SAME:    %[[C:[a-zA-Z0-9]*]]: memref<4xf32, #[[$DYN_1D_MAP]]>
+func @scf_for_with_tensor_insert_slice(
+    %A : tensor<?xf32>, %B : tensor<?xf32>, %C : tensor<4xf32>,
+    %lb : index, %ub : index, %step : index)
+  -> (tensor<?xf32>, tensor<?xf32>)
+{
+  // CHECK-NEXT: scf.for
+  %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
+      -> (tensor<?xf32>, tensor<?xf32>)
+  {
+    // CHECK-NEXT:   %[[SVA:.*]] = memref.subview %[[A]]
+    // CHECK-NEXT:   memref.copy %[[C]], %[[SVA]] : memref<4xf32, #[[$DYN_1D_MAP]]> to memref<4xf32, #[[$DYN_1D_MAP]]>
+    %ttA = tensor.insert_slice %C into %tA[%i][4][1] : tensor<4xf32> into tensor<?xf32>
+
+    // CHECK-NEXT:   %[[SVB:.*]] = memref.subview %[[B]]
+    // CHECK-NEXT:   memref.copy %[[C]], %[[SVB]] : memref<4xf32, #[[$DYN_1D_MAP]]> to memref<4xf32, #[[$DYN_1D_MAP]]>
+    %ttB = tensor.insert_slice %C into %tB[%i][4][1] : tensor<4xf32> into tensor<?xf32>
+
+    // scf.yield is empty and is elided
+    //  CHECK-NOT:   scf.yield
+    scf.yield %ttA, %ttB : tensor<?xf32>, tensor<?xf32>
+  }
+
+  // Swaparoo requires bufferizing the whole function to figure out who's who.
+  return %r0#1, %r0#0: tensor<?xf32>, tensor<?xf32>
+}
+
+//      CHECK:  func @bar(
+// CHECK-SAME:    %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]>
+// CHECK-SAME:    %[[B:[a-zA-Z0-9]*]]: memref<?xf32, #[[$DYN_1D_MAP]]>
+// CHECK-SAME:    %[[C:[a-zA-Z0-9]*]]: memref<4xf32, #[[$DYN_1D_MAP]]>
+func @bar(
+    %A : tensor<?xf32> {linalg.inplaceable = true},
+    %B : tensor<?xf32> {linalg.inplaceable = true},
+    %C : tensor<4xf32> {linalg.inplaceable = true},
+    %lb : index, %ub : index, %step : index)
+  -> (tensor<?xf32>, tensor<?xf32>)
+{
+//  CHECK-DAG:   call @scf_for_with_tensor_insert_slice(%[[A]], %[[B]], %[[C]]
+  %r0:2 = call @scf_for_with_tensor_insert_slice(%A, %B, %C, %lb, %ub, %step) :
+      (tensor<?xf32>, tensor<?xf32>, tensor<4xf32>, index, index, index)
+        -> (tensor<?xf32>, tensor<?xf32>)
+
+  // %r0#0 requires a copy because we have no idea what the function is doing.
+//  CHECK-DAG:   %[[alloc:.*]] = memref.alloc
+//  CHECK-DAG:   %[[casted:.*]] = memref.cast %[[alloc]]
+//      CHECK:   memref.copy %[[B]], %[[alloc]]
+// CHECK-NEXT:   call @some_external_func(%[[casted]]) : (memref<?xf32, #[[$DYN_1D_MAP]]>) -> ()
+  call @some_external_func(%r0#0) : (tensor<?xf32>) -> ()
+
+//      CHECK:   return
+  return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
+}
+
+// -----
+
+//  CHECK-DAG: #[[$DYN_0D_MAP:.*]] = affine_map<()[s0] -> (s0)>
+//  CHECK-DAG: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
+//      CHECK:  func @init_and_dot(
+// CHECK-SAME:    %[[A:[a-zA-Z0-9]*]]: memref<64xf32, #[[$DYN_1D_MAP]]>
+// CHECK-SAME:    %[[B:[a-zA-Z0-9]*]]: memref<64xf32, #[[$DYN_1D_MAP]]>
+// CHECK-SAME:    %[[C:[a-zA-Z0-9]*]]: memref<f32, #[[$DYN_0D_MAP]]>
+func @init_and_dot(%a: tensor<64xf32>, %b: tensor<64xf32>, %c: tensor<f32>) -> tensor<f32> {
+  // CHECK-NEXT:   %[[C0:.*]] = arith.constant 0{{.*}} : f32
+  %v0 = arith.constant 0.0 : f32
+
+  // CHECK-NEXT:   linalg.fill ins(%[[C0]] : f32) outs(%[[C]] : memref<f32, #[[$DYN_0D_MAP]]>)
+  %d = linalg.fill ins(%v0 : f32) outs(%c : tensor<f32>) -> tensor<f32>
+
+  // CHECK-NEXT:   linalg.dot ins(%[[A]], %[[B]] : memref<64xf32, #[[$DYN_1D_MAP]]>, memref<64xf32, #[[$DYN_1D_MAP]]>) outs(%[[C]] : memref<f32, #[[$DYN_0D_MAP]]>)
+  %e = linalg.dot ins(%a, %b : tensor<64xf32>,tensor<64xf32>)
+    outs(%d: tensor<f32>) -> tensor<f32>
+
+  // CHECK-NEXT:   return
+  return %e : tensor<f32>
+}
+
+//      CHECK:  func @main()
+func @main() {
+  //  CHECK-DAG:   %[[C0:.*]] = arith.constant 0{{.*}} : f32
+  //  CHECK-DAG:   %[[C1:.*]] = arith.constant 1{{.*}} : f32
+  //  CHECK-DAG:   %[[C2:.*]] = arith.constant 2{{.*}} : f32
+  %v0 = arith.constant 0.0 : f32
+  %v1 = arith.constant 1.0 : f32
+  %v2 = arith.constant 2.0 : f32
+
+  // CHECK-NEXT:   %[[A:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
+  // CHECK-NEXT:   %[[B:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
+  // CHECK-NEXT:   %[[C:.*]] = memref.alloc() {alignment = 128 : i64} : memref<f32>
+  //  CHECK-DAG:   %[[cA:.*]] = memref.cast %[[A]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]>
+  //  CHECK-DAG:   %[[cB:.*]] = memref.cast %[[B]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]>
+  //  CHECK-DAG:   %[[cC:.*]] = memref.cast %[[C]] : memref<f32> to memref<f32, #[[$DYN_0D_MAP]]>
+  %A = linalg.init_tensor [64] : tensor<64xf32>
+  %B = linalg.init_tensor [64] : tensor<64xf32>
+  %C = linalg.init_tensor [] : tensor<f32>
+
+  //  CHECK-DAG:   linalg.fill ins(%[[C1]] : f32) outs(%[[A]] : memref<64xf32>)
+  //  CHECK-DAG:   linalg.fill ins(%[[C2]] : f32) outs(%[[B]] : memref<64xf32>)
+  //  CHECK-DAG:   linalg.fill ins(%[[C0]] : f32) outs(%[[C]] : memref<f32>)
+  %AA = linalg.fill ins(%v1 : f32) outs(%A : tensor<64xf32>) -> tensor<64xf32>
+  %BB = linalg.fill ins(%v2 : f32) outs(%B : tensor<64xf32>) -> tensor<64xf32>
+  %CC = linalg.fill ins(%v0 : f32) outs(%C : tensor<f32>) -> tensor<f32>
+
+  // CHECK-NEXT:   call @init_and_dot(%[[cA]], %[[cB]], %[[cC]])
+  %res = call @init_and_dot(%AA, %BB, %CC) :
+    (tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32>
+
+  // CHECK-NEXT:   %[[dC:.*]] = memref.cast %[[C]] : memref<f32> to memref<*xf32>
+  %res2 = tensor.cast %res: tensor<f32> to tensor<*xf32>
+
+  // CHECK-NEXT:   call @print_memref_f32(%[[dC]]) : (memref<*xf32>) -> ()
+  call @print_memref_f32(%res2) : (tensor<*xf32>) -> ()
+
+  // CHECK-DAG:   memref.dealloc %[[A]] : memref<64xf32>
+  // CHECK-DAG:   memref.dealloc %[[B]] : memref<64xf32>
+  // CHECK-DAG:   memref.dealloc %[[C]] : memref<f32>
+  // CHECK-NEXT:   return
+  return
+}
+
+//     CHECK:   func private @print_memref_f32(memref<*xf32>)
+func private @print_memref_f32(tensor<*xf32>)
+
+// -----
+
+// CHECK: #[[$DYNAMIC:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
+// CHECK: func private @external_func(memref<?xf32, #[[$DYNAMIC]]>)
+func private @external_func(tensor<?xf32>)
+
+//      CHECK: func @callee(
+// CHECK-SAME:   %[[A:[0-9a-zA-Z]*]]: memref<?xf32>
+// CHECK-SAME:   %[[B:[0-9a-zA-Z]*]]: memref<?xf32, #[[$DYNAMIC]]>
+// CHECK-SAME:   %[[C:[0-9a-zA-Z]*]]: memref<?xf32, #[[$DYNAMIC]]>
+func @callee(%A : tensor<?xf32> {linalg.buffer_layout = affine_map<(i)[s0, s1] -> (i)>},
+             %B : tensor<?xf32>,
+             %C : tensor<?xf32>) {
+// CHECK-NEXT: %[[CASTED:.*]] = memref.cast %[[A]] : memref<?xf32> to memref<?xf32, #[[$DYNAMIC]]>
+// CHECK-NEXT: call @external_func(%[[CASTED]]) : (memref<?xf32, #[[$DYNAMIC]]>) -> ()
+  call @external_func(%A) : (tensor<?xf32>) -> ()
+
+// CHECK-NEXT: call @external_func(%[[B]]) : (memref<?xf32, #[[$DYNAMIC]]>) -> ()
+  call @external_func(%B) : (tensor<?xf32>) -> ()
+
+// CHECK-NEXT: call @external_func(%[[C]]) : (memref<?xf32, #[[$DYNAMIC]]>) -> ()
+  call @external_func(%C) : (tensor<?xf32>) -> ()
+
+  return
+}
+
+//      CHECK: func @entry(
+// CHECK-SAME:   %[[A:[0-9a-zA-Z]*]]: memref<?xf32>
+// CHECK-SAME:   %[[B:[0-9a-zA-Z]*]]: memref<?xf32>
+// CHECK-SAME:   %[[C:[0-9a-zA-Z]*]]: memref<?xf32, #[[$DYNAMIC]]>
+func @entry(%A : tensor<?xf32> {linalg.buffer_layout = affine_map<(i)[s0, s1] -> (i)>, linalg.inplaceable = false},
+            %B : tensor<?xf32> {linalg.buffer_layout = affine_map<(i)[s0, s1] -> (i)>, linalg.inplaceable = false},
+            %C : tensor<?xf32> {linalg.inplaceable = false}) {
+// Note: `callee` does not write to its bbArg directly, but `external_func`
+// does. Inside `callee`, the writes via `external_func` do not cause a
+// conflict. However, inside `entry`, the writes do cause a conflict because
+// %A, %B and %C are not inplaceable. This test case shows that this kind of
+// conflict detection has a "transitive" nature.
+//      CHECK: %[[ALLOC_C:.*]] = memref.alloc
+//      CHECK: %[[CASTED_C:.*]] = memref.cast %[[ALLOC_C]]
+//      CHECK: %[[ALLOC_B:.*]] = memref.alloc
+//      CHECK: %[[CASTED_B:.*]] = memref.cast %[[ALLOC_B]]
+//      CHECK: %[[ALLOC_A:.*]] = memref.alloc
+//      CHECK: memref.copy %[[A]], %[[ALLOC_A]]
+//      CHECK: memref.copy %[[B]], %[[ALLOC_B]]
+//      CHECK: memref.copy %[[C]], %[[ALLOC_C]]
+//      CHECK: %[[CASTED_A:.*]] = memref.cast %[[ALLOC_A]]
+// CHECK-NEXT: call @callee(%[[CASTED_A]], %[[CASTED_B]], %[[CASTED_C]])
+  call @callee(%A, %B, %C) : (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) -> ()
+  return
+}
+
+// -----
+
+// No alloc or copy inside of the loop.
+
+// CHECK-LABEL: func @inner_func(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
+func @inner_func(%t: tensor<?xf32>) -> tensor<?xf32> {
+  %f = arith.constant 1.0 : f32
+  %c0 = arith.constant 0 : index
+  // CHECK: memref.store %{{.*}}, %[[arg0]]
+  %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL: func @equivalent_func_arg(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
+func @equivalent_func_arg(%t0: tensor<?xf32> {linalg.inplaceable = true},
+                          %c0: index, %c10: index, %c1: index) -> tensor<?xf32> {
+  // CHECK-NOT: alloc
+  // CHECK-NOT: copy
+  %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {
+    // CHECK: call @inner_func(%[[arg0]])
+    %3 = call @inner_func(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+    scf.yield %3 : tensor<?xf32>
+  }
+  return %1: tensor<?xf32>
+}
+
+// -----
+
+// inner_func_2 modifies the bbArg, but the loop yields the original value. A
+// buffer copy must be inserted inside the loop.
+
+// CHECK-LABEL: func @inner_func_2(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
+func @inner_func_2(%t: tensor<?xf32>) -> tensor<?xf32> {
+  %f = arith.constant 1.0 : f32
+  %c0 = arith.constant 0 : index
+  // CHECK: memref.store %{{.*}}, %[[arg0]]
+  %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL: func @equivalent_func_arg_2(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
+func @equivalent_func_arg_2(%t0: tensor<?xf32> {linalg.inplaceable = true},
+                            %c0: index, %c10: index, %c1: index) -> tensor<?xf32> {
+  // CHECK: scf.for {{.*}} {
+  %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {
+    // CHECK: %[[alloc:.*]] = memref.alloc
+    // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
+    // CHECK: memref.copy %[[arg0]], %[[alloc]]
+    // CHECK: call @inner_func_2(%[[casted]])
+    // CHECK: memref.dealloc %[[alloc]]
+    // CHECK-NOT: scf.yield
+    %3 = call @inner_func_2(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+    scf.yield %t1 : tensor<?xf32>
+  }
+  return %1: tensor<?xf32>
+}