[mlir][linalg][bufferize] Collect equivalent FuncOp BBArgs in PostAnalysisStep

[springerm]

Summary

Collect equivalent BBArgs right after the equivalence analysis of the FuncOp and before bufferizing. This is in preparation of decoupling bufferization from aliasInfo.

Depends On D114774

[nicolasvasilache]

This is bigger than what qualifies as NFC - in my book, could we add a specific test that prints what the analysis is doing and we only check on that ?

[springerm]

update

[springerm]

update

[springerm]

update

[springerm]

update

[springerm]

Added tests.

[nicolasvasilache]

Thanks!

[springerm]

no change

[springerm]

add missing CallOp equivalence analysis, fixes sandbox tests

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

struct PostAnalysisStep {
 
   /// Run the post analysis step. This function may modify the IR, but must keep
   /// `aliasInfo` (inside `state`) consistent. Newly created operations and
   /// operations that should be re-analyzed must be stored in `newOps`.
   virtual LogicalResult run(FuncOp funcOp, BufferizationState &state,
                             SmallVector<Operation *> &newOps) = 0;
 };
 
+using PostAnalysisStepList = std::vector<std::unique_ptr<PostAnalysisStep>>;
+
 /// Options for ComprehensiveBufferize.
 struct BufferizationOptions {
   BufferizationOptions();
 
   // BufferizationOptions cannot be copied.
   BufferizationOptions(const BufferizationOptions &other) = delete;
 
   /// Register a "post analysis" step. Such steps are executed after the
struct BufferizationOptions {
   /// Should be used only with `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
 
   /// 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;
 
   /// Registered post analysis steps.
-  std::vector<std::unique_ptr<PostAnalysisStep>> postAnalysisSteps;
+  PostAnalysisStepList postAnalysisSteps;
 };
 
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
   None,
   // TODO: ResultContainsOperand,
   // TODO: OperandContainsResult,
   Equivalent

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h

 
 namespace mlir {
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 struct BufferizationOptions;
 struct BufferizationState;
+struct PostAnalysisStep;
 
 /// Bufferize the given function. Does not bufferize the function boundary.
+/// Reuses an existing BufferizationState object.
 // TODO: This function is meant to be called from ModuleBufferize and not can
 // not yet be called standalone.
-LogicalResult runComprehensiveBufferize(FuncOp funcOp,
-                                        const BufferizationOptions &options,
-                                        BufferizationState &state);
+LogicalResult runComprehensiveBufferize(
+    FuncOp funcOp, const BufferizationOptions &options,
+    BufferizationState &state,
+    const std::vector<std::unique_ptr<PostAnalysisStep>> &extraSteps);
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H

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

for (OpResult opResult : op->getResults()) {
             setInPlaceFuncArgument(bbArg,
                                    aliasInfo.bufferizesToWritableMemory(bbArg));
     }
   });
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     FuncOp funcOp, const BufferizationOptions &options,
-    BufferizationState &state) {
+    BufferizationState &state, const PostAnalysisStepList &extraSteps) {
 
   DominanceInfo domInfo(funcOp);
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
   if (funcOp.body().empty())
     return success();
 
   if (failed(checkAliasInfoConsistency(funcOp, domInfo, aliasInfo)))
     return failure();
 
   // If the analysis fails, just return.
   Operation *op = funcOp.getOperation();
   if (failed(
           inPlaceAnalysis(op, aliasInfo, domInfo, options.analysisFuzzerSeed)))
     return failure();
   equivalenceAnalysis(op, aliasInfo);
 
-  for (const std::unique_ptr<PostAnalysisStep> &step :
-       options.postAnalysisSteps) {
+  auto runPostAnalysisSteps = [&](const PostAnalysisStepList &steps) {
+    for (const std::unique_ptr<PostAnalysisStep> &step : steps) {
       SmallVector<Operation *> newOps;
       if (failed(step->run(funcOp, state, newOps)))
         return failure();
       // Analyze ops that were created by the PostAnalysisStep.
       if (failed(inPlaceAnalysis(newOps, aliasInfo, domInfo)))
         return failure();
       equivalenceAnalysis(newOps, aliasInfo);
     }
+    return success();
+  };
+
+  if (failed(runPostAnalysisSteps(extraSteps)))
+    return failure();
+  if (failed(runPostAnalysisSteps(options.postAnalysisSteps)))
+    return failure();
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly) {
     annotateOpsWithBufferizationMarkers(funcOp, aliasInfo);
     return success();
   }
 
   // Bufferize all ops in funcOp.

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

 
 namespace {
 /// Extra bufferization state that is required for bufferization of function
 /// boundaries.
 struct ModuleBufferizationState : public DialectBufferizationState {
   /// A map for looking up bufferized function types.
   DenseMap<FuncOp, FunctionType> bufferizedFunctionTypes;
 
-  /// A mapping of return values to equivalent BlockArguments.
-  DenseMap<Value, BlockArgument> equivalentReturnValToBBArg;
+  /// A mapping of ReturnOp OpOperand indices to equivalent FuncOp BBArg
+  /// indices.
+  DenseMap<FuncOp, DenseMap<int64_t, int64_t>> equivalentFuncArgs;
 };
 } // namespace
 
 static ModuleBufferizationState &
 getModuleBufferizationState(BufferizationState &state) {
   return state.getDialectState<ModuleBufferizationState>(
       StandardOpsDialect::getDialectNamespace());
 }
 
+/// Return the unique ReturnOp that terminates `funcOp`.
+/// Return nullptr if there is no such unique ReturnOp.
+static ReturnOp getAssumedUniqueReturnOp(FuncOp funcOp) {
+  ReturnOp returnOp;
+  for (Block &b : funcOp.body()) {
+    if (auto candidateOp = dyn_cast<ReturnOp>(b.getTerminator())) {
+      if (returnOp)
+        return nullptr;
+      returnOp = candidateOp;
+    }
+  }
+  return returnOp;
+}
+
+namespace {
+/// Store function BlockArguments that are equivalent to a returned value in
+/// ModuleBufferizationState.
+struct EquivalentFuncOpBBArgsAnalysis : public PostAnalysisStep {
+  /// Annotate IR with the results of the analysis. For testing purposes only.
+  static void annotateReturnOp(OpOperand &returnVal, BlockArgument bbArg) {
+    const char *kEquivalentArgsAttr = "__equivalent_func_args__";
+    Operation *op = returnVal.getOwner();
+
+    SmallVector<int64_t> equivBbArgs;
+    if (op->hasAttr(kEquivalentArgsAttr)) {
+      auto attr = op->getAttr(kEquivalentArgsAttr).cast<ArrayAttr>();
+      equivBbArgs = llvm::to_vector<4>(llvm::map_range(attr, [](Attribute a) {
+        return a.cast<IntegerAttr>().getValue().getSExtValue();
+      }));
+    } else {
+      equivBbArgs.append(op->getNumOperands(), -1);
+    }
+    equivBbArgs[returnVal.getOperandNumber()] = bbArg.getArgNumber();
+
+    OpBuilder b(op->getContext());
+    op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
+  }
+
+  LogicalResult run(FuncOp funcOp, BufferizationState &state,
+                    SmallVector<Operation *> &newOps) override {
+    ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
+
+    // Support only single return-terminated block in the function.
+    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 (state.aliasInfo.areEquivalentBufferizedValues(returnVal.get(),
+                                                              bbArg)) {
+              moduleState
+                  .equivalentFuncArgs[funcOp][returnVal.getOperandNumber()] =
+                  bbArg.getArgNumber();
+              if (state.options.testAnalysisOnly)
+                annotateReturnOp(returnVal, bbArg);
+            }
+
+    return success();
+  }
+};
+} // namespace
+
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 /// If `value` is a memref::CastOp, return its source. Otherwise, return
 /// `value` directly.
 static Value getNonCastedValue(Value value) {
   while (auto castOp = value.getDefiningOp<memref::CastOp>())
     value = castOp.source();
   return value;

 static FuncOp getCalledFunction(CallOpInterface callOp) {
   SymbolRefAttr sym = callOp.getCallableForCallee().dyn_cast<SymbolRefAttr>();
   if (!sym)
     return nullptr;
   return dyn_cast_or_null<FuncOp>(
       SymbolTable::lookupNearestSymbolFrom(callOp, sym));
 }
 
-/// Return the unique ReturnOp that terminates `funcOp`.
-/// Return nullptr if there is no such unique ReturnOp.
-static ReturnOp getAssumedUniqueReturnOp(FuncOp funcOp) {
-  ReturnOp returnOp;
-  for (Block &b : funcOp.body()) {
-    if (auto candidateOp = dyn_cast<ReturnOp>(b.getTerminator())) {
-      if (returnOp)
-        return nullptr;
-      returnOp = candidateOp;
-    }
-  }
-  return returnOp;
-}
-
 /// Return the FunctionType with `argumentTypes` and `resultTypes` where each
 /// tensor is replaced by the corresponding buffer type.
 /// In order for all the callers to agree, this *must* bufferize to the most
 /// dynamic buffer type supported.
 /// A later pass across all CallOps in the module can decide whether to simplify
 /// the types of to version according to some cost model.
 static FunctionType getBufferizedFunctionType(MLIRContext *ctx,
                                               TypeRange argumentTypes,
static FunctionType getOrCreateBufferizedFunctionType(
 
   auto it2 = bufferizedFunctionTypes.try_emplace(
       funcOp, getBufferizedFunctionType(funcOp.getContext(), argumentTypes,
                                         resultTypes));
   LDBG("FT: " << funcOp.getType() << " -> " << it2.first->second << "\n");
   return it2.first->second;
 }
 
-/// Store function BlockArguments that are equivalent to a returned value in
-/// the given ModuleBufferizationState.
-static void populateEquivalentFuncOpBBArgs(FuncOp funcOp,
-                                           BufferizationState &state) {
-  ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
+/// Gather equivalence info of CallOps.
+/// Note: This only adds new equivalence info if `funcOp` was already analyzed.
+// TODO: This does not handle cyclic function call graphs etc.
+static void equivalenceAnalysis(FuncOp funcOp,
+                                BufferizationAliasInfo &aliasInfo,
+                                ModuleBufferizationState &moduleState) {
+  funcOp->walk([&](CallOp callOp) {
+    FuncOp calledFunction = getCalledFunction(callOp);
+    assert(calledFunction && "could not retrieved called FuncOp");
 
-  // Support only single return-terminated block in the function.
-  ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
-  assert(returnOp && "expected func with single return op");
+    // No equivalence info available for the called function.
+    if (!moduleState.equivalentFuncArgs.count(calledFunction))
+      return WalkResult::skip();
+
+    for (auto it : moduleState.equivalentFuncArgs[calledFunction]) {
+      int64_t returnIdx = it.first;
+      int64_t bbargIdx = it.second;
+      Value returnVal = callOp.getResult(returnIdx);
+      Value argVal = callOp->getOperand(bbargIdx);
+      aliasInfo.unionEquivalenceClasses(returnVal, argVal);
+    }
 
-  for (Value returnVal : returnOp.operands())
-    if (returnVal.getType().isa<RankedTensorType>())
-      for (BlockArgument bbArg : funcOp.getArguments())
-        if (bbArg.getType().isa<RankedTensorType>())
-          if (state.aliasInfo.areEquivalentBufferizedValues(returnVal, bbArg))
-            moduleState.equivalentReturnValToBBArg[returnVal] = bbArg;
+    return WalkResult::advance();
+  });
 }
 
 /// Rewrite the `funcOp` arguments analysis return values and terminator into
 /// buffer form (using the canonical memref layout for now), according to the
 /// inPlace-bufferizable information of the function arguments.
 ///
 /// This relies on a buffer equivalence analysis of each return operand. When a
 /// result buffer is equivalent to a BlockArgument of `funcOp`, it can be
for (OpOperand &returnOperand : returnOp->getOpOperands()) {
 
     // If not a renturn tensor type just forward it.
     if (!returnVal.getType().isa<RankedTensorType>()) {
       returnValues.push_back(returnVal);
       continue;
     }
 
     // If return operand is equivalent to some bbArg, no need to return it.
-    if (moduleState.equivalentReturnValToBBArg.count(returnVal))
+    if (moduleState.equivalentFuncArgs[funcOp].count(
+            returnOperand.getOperandNumber()))
       continue;
 
     // Cast values at the call site if necessary.
     returnValues.push_back(getNonCastedValue(state.lookupBuffer(returnVal)));
   }
 
   // 2. Rewrite the terminator without the inPlace bufferizable values.
   ValueRange retValues{returnValues};
if (funcOp.body().empty()) {
       for (OpOperand &returnOperand : returnOp->getOpOperands()) {
         Type returnType = returnOperand.get().getType();
         if (!isaTensor(returnType)) {
           resultTypes.push_back(returnType);
           continue;
         }
 
         // If return operand is equivalent to some bbArg, no need to return it.
-        Value returnVal = returnOperand.get();
-        if (moduleState.equivalentReturnValToBBArg.count(returnVal)) {
-          BlockArgument bbArg =
-              moduleState.equivalentReturnValToBBArg[returnVal];
+        if (moduleState.equivalentFuncArgs[funcOp].count(
+                returnOperand.getOperandNumber())) {
+          int64_t idx =
+              moduleState
+                  .equivalentFuncArgs[funcOp][returnOperand.getOperandNumber()];
           Value oldRes = callOp->getResult(returnOperand.getOperandNumber());
-          int64_t idx = bbArg.getArgNumber();
           Value buffer = state.lookupBuffer(callOp->getOperand(idx));
           // Add CallOp operand/result equivalence: this is interprocedural
           // info.
           state.aliasInfo.insertNewBufferEquivalence(oldRes, buffer);
           state.mapBuffer(oldRes, buffer);
           // Add a ToTensorOp to kill all uses of the CallOp return.
           // Replace all uses of the CallOp results so we can erase the CallOp.
           // This ToTensorOp must fold/DCE away or bufferization should be

 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     ModuleOp moduleOp, const BufferizationOptions &options) {
   SmallVector<FuncOp> orderedFuncOps;
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   BufferizationState state(moduleOp, options);
+  ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
   // Interestingly, all function args that are not visible outside of a module
   // can be fully bufferized inplace by guaranteeing the CallOp is bufferized
   // inplace. Therefore, we just bufferize funcOp as if none of its results were
   // inplaceable, detect which operands are cloned internally and decide what to
   // do at call sites.
   for (FuncOp funcOp : orderedFuncOps) {
for (FuncOp funcOp : orderedFuncOps) {
     // bufferizing to a writeable memory.
     for (BlockArgument bbArg : funcOp.getArguments()) {
       BoolAttr inplaceAttr = funcOp.getArgAttrOfType<BoolAttr>(
           bbArg.getArgNumber(), BufferizableOpInterface::kInplaceableAttrName);
       if (inplaceAttr && inplaceAttr.getValue())
         aliasInfo.setBufferizesToWritableMemory(bbArg);
     }
 
+    // Register extra post analysis steps. These cannot be stored in `options`
+    // because `options` is immutable.
+    PostAnalysisStepList extraSteps;
+    extraSteps.emplace_back(std::make_unique<EquivalentFuncOpBBArgsAnalysis>());
+
+    // Gather equivalence info for CallOps.
+    equivalenceAnalysis(funcOp, aliasInfo, moduleState);
+
     // Analyze and bufferize funcOp.
-    if (failed(runComprehensiveBufferize(funcOp, options, state)))
+    if (failed(runComprehensiveBufferize(funcOp, options, state, extraSteps)))
       return failure();
-
-    populateEquivalentFuncOpBBArgs(funcOp, state);
   }
 
   if (options.testAnalysisOnly)
     return success();
 
   for (FuncOp funcOp : orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.

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

 // CHECK-LABEL: func @insert_slice_fun
 func @insert_slice_fun(
     %A : tensor<?xf32>,
     %B : tensor<?xf32> {linalg.inplaceable = true},
     %C : tensor<4xf32>)
   -> (tensor<?xf32>, tensor<?xf32>)
 {
   // must bufferize out of place.
-  //     CHECK: tensor.insert_slice
+  //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
   %r0 = tensor.insert_slice %C into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
   // bufferizes inplace.
-  //     CHECK: tensor.insert_slice
+  //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %r1 = tensor.insert_slice %C into %B[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [-1, 1]}
   return %r0, %r1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @conflict_on_B
 func @conflict_on_B(
     %A : tensor<4x4xf32> {linalg.inplaceable = true},
func @conflict_on_B(
 
   // matmul output operand does not interferes with input operand.
   //     CHECK: linalg.matmul
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %E = linalg.matmul  ins(%A, %A: tensor<4x4xf32>, tensor<4x4xf32>)
                      outs(%B: tensor<4x4xf32>)
     -> tensor<4x4xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [-1, -1, 1]}
   return %C, %D, %E: tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>
 }
 
 //===----------------------------------------------------------------------===//
 // Length-1 producer-consumer cases.
 //===----------------------------------------------------------------------===//
 
 // -----
func @insert_slice_insert_slice(
   %r1 = tensor.insert_slice %r0 into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
   // CHECK: {__inplace_results_attr__ = ["false"]}
   %r2 = tensor.insert_slice %B3 into %B2[0][2][1] : tensor<2xf32> into tensor<4xf32>
 
   // CHECK: {__inplace_results_attr__ = ["false"]}
   %r3 = tensor.insert_slice %r2 into %B[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0, -1]}
   return %r1, %r3: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @extract_slice_nonmatching_insert_slice
 func @extract_slice_nonmatching_insert_slice(
     %A : tensor<?xf32> {linalg.inplaceable = true},
func @extract_slice_nonmatching_insert_slice(
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %r2 = tensor.extract_slice %B[0][4][1] : tensor<?xf32> to tensor<4xf32>
 
   // %r3 cannot bufferize inplace since %B is not inplaceable.
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
   %r3 = tensor.insert_slice %r2 into %B[%idx][4][1] : tensor<4xf32> into tensor<?xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0, -1]}
   return %r1, %r3: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @extract_slice_matching_insert_slice
 func @extract_slice_matching_insert_slice(
     %A : tensor<?xf32> {linalg.inplaceable = true},
func @extract_slice_matching_insert_slice(
   // tensor.insert_slice cannot bufferize inplace.
   // This should have been captured by a canonicalization pattern and it would
   // be unproductive to have special logic in bufferization to encode matching
   // insert_slice(extract_slice(A), A).
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
   %r3 = tensor.insert_slice %r2 into %B[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0, -1]}
   return %r1, %r3: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @read_of_matching_insert_slice_source
 func @read_of_matching_insert_slice_source(
     %A : tensor<?xf32> {linalg.inplaceable = true}, %idx : index, %idx2 : index)
func @read_of_matching_insert_slice_source(
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %1 = linalg.fill(%cst, %0) : f32, tensor<?xf32> -> tensor<?xf32>
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %2 = tensor.insert_slice %1 into %A[%idx][%idx][1] : tensor<?xf32> into tensor<?xf32>
 
   %3 = vector.transfer_read %1[%idx2], %cst2 : tensor<?xf32>, vector<5xf32>
+
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0, -1]}
   return %2, %3 : tensor<?xf32>, vector<5xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @read_of_matching_insert_slice_source_interleaved
 func @read_of_matching_insert_slice_source_interleaved(
     %A : tensor<?xf32> {linalg.inplaceable = true}, %idx : index, %idx2 : index,
func @read_of_matching_insert_slice_source_interleaved(
   %5 = linalg.fill(%cst, %4) : f32, tensor<?xf32> -> tensor<?xf32>
 
   %3 = vector.transfer_read %1[%idx2], %cst2 : tensor<?xf32>, vector<5xf32>
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %6 = tensor.insert_slice %5 into %2[%idx3][%idx3][1] : tensor<?xf32> into tensor<?xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0, -1]}
   return %6, %3 : tensor<?xf32>, vector<5xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @extract_slice_linalg_readonly_use
 func @extract_slice_linalg_readonly_use(
     %A : tensor<?x?xf32>,
func @extract_slice_linalg_readonly_use(
 
   // matmul output operand is inplaceable at the function boundary.
   //     CHECK: linalg.matmul
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %E = linalg.matmul  ins(%sA, %B: tensor<4x4xf32>, tensor<4x4xf32>)
                      outs(%C: tensor<4x4xf32>)
     -> tensor<4x4xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [-1, 2]}
   return %D, %E: tensor<4x4xf32>, tensor<4x4xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @extract_slice_to_linalg_write_use
 func @extract_slice_to_linalg_write_use(
     %A : tensor<4x4xf32>,
func @insert_slice_double_extract_slice(
   //      CHECK: tensor.extract_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %19 = tensor.extract_slice %18[0, 0] [%s1, %s2] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %20 = tensor.insert_slice %19 into %C[%s3, %s4] [%s1, %s2] [1, 1] : tensor<?x?xf32> into tensor<30x20xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [6]}
   return %20 : tensor<30x20xf32>
 }
 
 //===----------------------------------------------------------------------===//
 // Transitive cases
 //===----------------------------------------------------------------------===//
 
 // -----
func @nested_extract_slice_and_insert(
   // CHECK-NEXT: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %sC = tensor.extract_slice %C[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
   %ssC = tensor.extract_slice %sC[0, 0][%sz1, 4][1, 1] : tensor<?x?xf32> to tensor<?x4xf32>
   %FC = linalg.fill(%f0, %ssC) : f32, tensor<?x4xf32> -> tensor<?x4xf32>
   %rsC = tensor.insert_slice %FC into %sC[0, 0][%sz2, 4][1, 1] : tensor<?x4xf32> into tensor<?x?xf32>
   %rC = tensor.insert_slice %rsC into %C[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [-1, 1, 2]}
   return %rA, %rB, %rC: tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
 }
 
 //===----------------------------------------------------------------------===//
 // Simple loop cases
 //===----------------------------------------------------------------------===//
 
 // -----
func @scf_for_yield_only(%A : tensor<?xf32>,
 
   //      CHECK: scf.for
   // CHECK-NEXT: scf.yield
   // CHECK-NEXT: {__inplace_results_attr__ = ["true"]}
   %r1 = scf.for %i = %lb to %ub step %step iter_args(%t = %B) -> (tensor<?xf32>) {
     scf.yield %t : tensor<?xf32>
   }
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [-1, 1]}
   return %r0, %r1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_for_with_tensor.insert_slice
 func @scf_for_with_tensor.insert_slice(%A : tensor<?xf32>,
               %B : tensor<?xf32> {linalg.inplaceable = true},
func @scf_for_with_tensor.insert_slice(%A : tensor<?xf32>,
   %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
       -> (tensor<?xf32>, tensor<?xf32>)
   {
     %ttA = tensor.insert_slice %C into %tA[0][4][1] : tensor<4xf32> into tensor<?xf32>
     %ttB = tensor.insert_slice %C into %tB[0][4][1] : tensor<4xf32> into tensor<?xf32>
     scf.yield %ttA, %ttB : tensor<?xf32>, tensor<?xf32>
   }
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [-1, 1]}
   return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 func private @some_use(tensor<?xf32>) -> ()
 
 // CHECK-LABEL: func @scf_for_deps
func @scf_for_deps(%A : tensor<?xf32> {linalg.inplaceable = true},
   // CHECK-NEXT: {__inplace_results_attr__ = ["true"]}
   %r3 = linalg.tiled_loop (%i) = (%lb) to (%ub) step (%step)
         ins()
         outs(%t = %B: tensor<?xf32>) {
     call @some_use(%t) : (tensor<?xf32>) -> ()
     linalg.yield %t : tensor<?xf32>
   }
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0, 1]}
   return %r1, %r3: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Cross function boundary cases.
 //===----------------------------------------------------------------------===//
builtin.func @matmul_on_tensors(
   //      CHECK: linalg.matmul
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %sA = tensor.extract_slice %8[0, 0][256, 16][1, 1]: tensor<256x256xf32> to tensor<256x16xf32>
   %sB = tensor.extract_slice %11[0, 0][16, 256][1, 1]: tensor<256x256xf32> to tensor<16x256xf32>
   %r = linalg.matmul
          ins(%sA, %sB : tensor<256x16xf32>, tensor<16x256xf32>)
         outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [2]}
   return %r : tensor<256x256xf32>
 }
 
 // -----
 
 builtin.func @matmul_on_tensors(
     %arg0: tensor<518x518xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
     %arg1: tensor<518x518xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
builtin.func @matmul_on_tensors(
   //      CHECK: linalg.matmul
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %sA = tensor.extract_slice %10[0, 0][256, 16][1, 1]: tensor<256x256xf32> to tensor<256x16xf32>
   %sB = tensor.extract_slice %13[0, 0][16, 256][1, 1]: tensor<256x256xf32> to tensor<16x256xf32>
   %r = linalg.matmul
          ins(%sA, %sB : tensor<256x16xf32>, tensor<16x256xf32>)
         outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [2]}
   return %r : tensor<256x256xf32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Chain of tensor.insert_slice is better traversed in reverse order without
 // prioritizing  the tensor.insert_slice ops.
func @insert_slice_chain(
   %10 = tensor.extract_slice %8[32, 0] [30, 90] [1, 1] : tensor<62x90xf32> to tensor<30x90xf32>
   //      CHECK: vector.transfer_write
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %14 = vector.transfer_write %v2, %10[%c0, %c0] {in_bounds = [true, true]} : vector<30x90xf32>, tensor<30x90xf32>
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %15 = tensor.insert_slice %14 into %8[32, 0] [30, 90] [1, 1] : tensor<30x90xf32> into tensor<62x90xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [4]}
   return %15 : tensor<62x90xf32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Insert point issue cases.
 //===----------------------------------------------------------------------===//
func @ip(%t: tensor<10x20xf32> {linalg.inplaceable = true},
   %c257 = arith.constant 257 : index
   %r = scf.for %arg0 = %c0 to %c257 step %c256 iter_args(%arg1 = %t) -> (tensor<10x20xf32>) {
     %t1 = tensor.extract_slice %arg1[%x, 0] [5, %y] [1, 1] : tensor<10x20xf32> to tensor<5x?xf32>
     %t11 = tensor.extract_slice %t1[0, 0] [5, %y] [1, 1] : tensor<5x?xf32> to tensor<5x?xf32>
     %t2 = vector.transfer_write %v, %t11[%c0, %c0] : vector<5x6xf32>, tensor<5x?xf32>
     %t3 = tensor.insert_slice %t2 into %arg1[%x, 0] [5, %y] [1, 1] : tensor<5x?xf32> into tensor<10x20xf32>
     scf.yield %t3 : tensor<10x20xf32>
   }
+
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
  return %r : tensor<10x20xf32>
 }
 
 // -----
 
 #accesses = [
   affine_map<(i) -> (i)>,
   affine_map<(i) -> (i)>,
func @linalg_op_same_out_tensors(
 
   //      CHECK: linalg.generic
   // CHECK-SAME: {__inplace_results_attr__ = ["true", "false"]
   %o:2 = linalg.generic #trait ins(%t1 : tensor<?xf32>)
                                outs (%t2, %t2 : tensor<?xf32>, tensor<?xf32>) {
       ^bb(%0: f32, %1: f32, %2 : f32) :
         linalg.yield %0, %0 : f32, f32
     } -> (tensor<?xf32>, tensor<?xf32>)
+
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [1, -1]}
   return %o#0, %o#1 : tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @double_insert_slice_into_alias
 func @double_insert_slice_into_alias(
     %v1: vector<32x90xf32>,
func @double_insert_slice_into_alias(
   %10 = tensor.extract_slice %e[32, 0] [30, 90] [1, 1] : tensor<?x?xf32> to tensor<30x90xf32>
   //      CHECK: vector.transfer_write
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %14 = vector.transfer_write %v2, %10[%c0, %c0] {in_bounds = [true, true]} : vector<30x90xf32>, tensor<30x90xf32>
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %15 = tensor.insert_slice %14 into %e[32, 0] [30, 90] [1, 1] : tensor<30x90xf32> into tensor<?x?xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [2, -1]}
   return %8, %15 : tensor<62x90xf32>, tensor<?x?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @interleaved_extract_insert_slice_chain_1
 func @interleaved_extract_insert_slice_chain_1(
     %arg2: tensor<62x90xf32> {linalg.inplaceable = true})
func @interleaved_extract_insert_slice_chain_1(
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %8 = tensor.insert_slice %2 into %arg2[0, 0] [32, 90] [1, 1] : tensor<32x90xf32> into tensor<62x90xf32>
 
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %15 = tensor.insert_slice %10 into %8[32, 0] [30, 90] [1, 1] : tensor<30x90xf32> into tensor<62x90xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
   return %15 : tensor<62x90xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @interleaved_extract_insert_slice_chain_2
 func @interleaved_extract_insert_slice_chain_2(
     %arg2: tensor<62x90xf32> {linalg.inplaceable = true})
func @interleaved_extract_insert_slice_chain_2(
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %8 = tensor.insert_slice %2 into %arg2[0, 0] [32, 90] [1, 1] : tensor<32x90xf32> into tensor<62x90xf32>
 
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %15 = tensor.insert_slice %10 into %8[31, 0] [30, 90] [1, 1] : tensor<30x90xf32> into tensor<62x90xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
   return %15 : tensor<62x90xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @extract_once_insert_twice
 func @extract_once_insert_twice(
     %arg2: tensor<62x90xf32> {linalg.inplaceable = true})
   -> (tensor<62x90xf32>)
 {
   //      CHECK: tensor.extract_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["false"]
   %2 = tensor.extract_slice %arg2[0, 0] [32, 90] [1, 1] : tensor<62x90xf32> to tensor<32x90xf32>
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %8 = tensor.insert_slice %2 into %arg2[0, 0] [32, 90] [1, 1] : tensor<32x90xf32> into tensor<62x90xf32>
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %15 = tensor.insert_slice %2 into %8[15, 0] [32, 90] [1, 1] : tensor<32x90xf32> into tensor<62x90xf32>
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
   return %15 : tensor<62x90xf32>
 }
 
 // -----
 
 #accesses = [
   affine_map<(i) -> (i)>
 ]
func @non_reading_scf_for(%t1: tensor<?xf32> {linalg.inplaceable = true},
   // Use %t3 in some way without reading it, so that it does not get DCE'd.
     // CHECK:      linalg.generic
     // CHECK-SAME: __inplace_results_attr__ = ["true"]
   %o = linalg.generic #trait outs (%t3 : tensor<?xf32>) {
       ^bb(%0: f32) :
         linalg.yield %cst : f32
     } -> (tensor<?xf32>)
 
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0, -1]}
   return %o, %v3 : tensor<?xf32>, vector<5xf32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // InitTensorOp elimination
 //===----------------------------------------------------------------------===//
func @buffer_forwarding_conflict(%arg0: tensor<?xf32> {linalg.inplaceable = true}, %arg1: index) -> (tensor<?xf32>, tensor<?xf32>) {
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["false"]
   %2 = tensor.insert_slice %1 into %arg0[0] [%arg1] [1] : tensor<?xf32> into tensor<?xf32>
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %3 = tensor.insert_slice %1 into %arg0[42] [%arg1] [1] : tensor<?xf32> into tensor<?xf32>
+
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [-1, 0]}
   return %2, %3 : tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @buffer_forwarding_no_conflict
 func @buffer_forwarding_no_conflict(%arg0: tensor<?xf32> {linalg.inplaceable = true}, %arg1: index) -> (tensor<?xf32>, tensor<?xf32>) {
   %cst = arith.constant 0.000000e+00 : f32
   //      CHECK: tensor.extract_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   // Instead of allocating, share buffer with some inplace bufferization?
   %0 = linalg.init_tensor [%arg1] : tensor<?xf32>
 
   //      CHECK: linalg.fill
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %1 = linalg.fill(%cst, %0) : f32, tensor<?xf32> -> tensor<?xf32>
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %2 = tensor.insert_slice %1 into %arg0[42] [%arg1] [1] : tensor<?xf32> into tensor<?xf32>
+
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0, 0]}
   return %2, %2 : tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // scf.if cases
 //===----------------------------------------------------------------------===//
func @scf_if_inplace2(%t1: tensor<?xf32> {linalg.inplaceable = true},
   %r = scf.if %cond -> (tensor<?xf32>) {
     scf.yield %t1 : tensor<?xf32>
   } else {
     //      CHECK: vector.transfer_write
     // CHECK-SAME: {__inplace_results_attr__ = ["true"]
     %t2 = vector.transfer_write %v, %t1[%idx] : vector<5xf32>, tensor<?xf32>
     scf.yield %t2 : tensor<?xf32>
   }
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
   return %r : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_if_inplace3
 func @scf_if_inplace3(%t1: tensor<?xf32> {linalg.inplaceable = true},
                       %v1: vector<5xf32>, %v2: vector<5xf32>, %idx: index,
func @scf_if_in_place4(%t1: tensor<?xf32> {linalg.inplaceable = true},
   }
   %r_alias = scf.if %cond2 -> (tensor<?xf32>) {
     // Reading %r is OK. No conflict.
     scf.yield %r : tensor<?xf32>
   } else {
     scf.yield %r : tensor<?xf32>
   }
   %v2 = vector.transfer_read %r_alias[%idx], %cst : tensor<?xf32>, vector<10xf32>
+
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0, -1]}
   return %r_alias, %v2 : tensor<?xf32>, vector<10xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_if_inplace5
 func @scf_if_inplace5(%t1: tensor<?xf32> {linalg.inplaceable = true},
                       %idx: index, %cond: i1) -> tensor<?xf32> {
%r = scf.if %cond -> (tensor<?xf32>) {
     scf.yield %f : tensor<?xf32>
   }
 
   // Inserting into an equivalent tensor at the same offset. This bufferizes
   // inplace.
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %r2 = tensor.insert_slice %r into %t1[%idx][%idx][1] : tensor<?xf32> into tensor<?xf32>
+
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
   return %r2 : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_if_inplace6
 func @scf_if_inplace6(%t1: tensor<?xf32> {linalg.inplaceable = true},
                       %v1: vector<5xf32>, %v2: vector<5xf32>,
%r = scf.if %cond -> (tensor<?xf32>) {
     }
     scf.yield %t2 : tensor<?xf32>
   } else {
     //      CHECK: vector.transfer_write
     // CHECK-SAME: {__inplace_results_attr__ = ["true"]
     %t3 = vector.transfer_write %v2, %t1[%idx] : vector<5xf32>, tensor<?xf32>
     scf.yield %t3 : tensor<?xf32>
   }
+
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
   return %r : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_if_inplace7
 func @scf_if_inplace7(%t1: tensor<?xf32> {linalg.inplaceable = true},
                       %v1: vector<5xf32>, %v2: vector<5xf32>, %idx: index,
func @scf_if_out_of_place1b(%t1: tensor<?xf32> {linalg.inplaceable = true},
   // Reading from and writing to the same tensor via different args. This is a
   // conflict. In contrast to scf_if_out_of_place1a, the fact that %r aliases
   // with %t1 is only detected when analyzing the tensor.extract_slices. That's
   // why the tensor.insert_slice is inplace and the two extract_slices are
   // out-of-place.
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %r2 = tensor.insert_slice %r into %t1[%idx3][%idx3][1] : tensor<?xf32> into tensor<?xf32>
+
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
   return %r2 : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_if_out_of_place1c
 func @scf_if_out_of_place1c(%t1: tensor<?xf32> {linalg.inplaceable = true},
                             %idx: index, %idx2: index, %cond: i1) -> tensor<?xf32> {
   %r = scf.if %cond -> (tensor<?xf32>) {
     //      CHECK: tensor.extract_slice
     // CHECK-SAME: {__inplace_results_attr__ = ["false"]
     %e = tensor.extract_slice %t1[%idx][%idx][1] : tensor<?xf32> to tensor<?xf32>
     scf.yield %e : tensor<?xf32>
   } else {
     // TODO: This one could bufferize inplace, but the analysis is too restrictive.
     //      CHECK: tensor.extract_slice
     // CHECK-SAME: {__inplace_results_attr__ = ["false"]
     %f = tensor.extract_slice %t1[%idx2][%idx2][1] : tensor<?xf32> to tensor<?xf32>
     scf.yield %f : tensor<?xf32>
   }
 
   //      CHECK: tensor.insert_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]
   %r2 = tensor.insert_slice %r into %t1[%idx2][%idx2][1] : tensor<?xf32> into tensor<?xf32>
+
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
   return %r2 : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_if_out_of_place2
 func @scf_if_out_of_place2(%t1: tensor<?xf32> {linalg.inplaceable = true},
                            %v: vector<5xf32>, %idx: index,
func @to_memref_op_is_reading(%t1: tensor<?xf32> {linalg.inplaceable = true},
 
   // Write + read to/from same memref.
   %0 = bufferization.to_memref %t1 : memref<?xf32>
   vector.transfer_write %v1, %0[%idx1] : vector<5xf32>, memref<?xf32>
   %r2 = vector.transfer_read %0[%idx3], %cst : memref<?xf32>, vector<5xf32>
 
   return %r1, %r2 : vector<5xf32>, vector<5xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @inner_func
+func @inner_func(%t: tensor<?xf32>) -> tensor<?xf32> {
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
+  return %t : tensor<?xf32>
+}
+
+func @equivalent_func_arg(%c0: index, %c10: index, %c1: index, %t0: tensor<?xf32>) -> tensor<?xf32> {
+  // This test does not check IR. It just asserts there is no failure due to
+  // non-equivalent scf.for yield values.
+  %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {
+    %3 = call @inner_func(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+    scf.yield %3 : tensor<?xf32>
+  }
+  return %1: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @inner_func_2
+func @inner_func_2(%t: tensor<?xf32>) -> tensor<?xf32> {
+  %f = arith.constant 1.0 : f32
+  %c0 = arith.constant 0 : index
+  %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>
+  //      CHECK: return
+  // CHECK-SAME: {__equivalent_func_args__ = [0]}
+  return %0 : tensor<?xf32>
+}
+
+func @equivalent_func_arg_2(%c0: index, %c10: index, %c1: index, %t0: tensor<?xf32>) -> tensor<?xf32> {
+  // This test does not check IR. It just asserts there is no failure due to
+  // non-equivalent scf.for yield values.
+  %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {
+    %3 = call @inner_func_2(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+    scf.yield %3 : tensor<?xf32>
+  }
+  return %1: tensor<?xf32>
+}

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

 //  CHECK-SAME:     %[[t1:.*]]: memref<?xf32, {{.*}}>, %[[s:.*]]: f32, %[[i:.*]]: index
 func @insert_op(%t1 : tensor<?xf32> {linalg.inplaceable = true},
                 %s : f32, %i : index) -> tensor<?xf32> {
   // 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>) {
+    // TODO: There should be a memory copy here. This is a bug in CallOp
+    // bufferization.
+    // CHECK: call @inner_func_2(%[[arg0]])
+    %3 = call @inner_func_2(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+    scf.yield %t1 : tensor<?xf32>
+  }
+  return %1: tensor<?xf32>
+}