[mlir][Linalg] Add a ComprehensiveModuleBufferizePass and support for CallOp analysis(9/n)

[nicolasvasilache]

Summary

This revision adds the minimal plumbing to create a simple ComprehensiveModuleBufferizePass that can behave conservatively in the presence of CallOps.

A topological sort of caller/callee is performed and, if the call-graph is cycle-free, analysis can proceed.

[ftynse]

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

425

1908

Should we bail out entirely if the callable is not a FuncOp?

1916

Nit: I got confused by the combination of the name (callCounter) and behavior here. It seems to count calls to distinct functions, not all calls, i.e. a repeated call to the same function won't be counted twice.

1980–1981

I don't think we reset the IR on module failure, and the code just above this modifies function argument attributes.

[nicolasvasilache]

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

1916

spelled it out a little longer and dropped the doubly-nested walk.

1980–1981

Comment was stale, dropped it.

Re signalPassFailure, the doc says:

/// Signal that some invariant was broken when running. The IR is allowed to
/// be in an invalid state.
void signalPassFailure() { getPassState().irAndPassFailed.setInt(true); }

AFAIK no pass that signal failure tries to put the IR back in a valid state.

Added a post-pass cleanup for FuncOp bbargs indeed.

[nicolasvasilache]

Address comments.

[ftynse]

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

1980–1981

Yeah, I was just making sure that nobody will expect the IR to be rolled back.

[nicolasvasilache]

Properly return.

mlir/include/mlir/Dialect/Linalg/Passes.h

 std::unique_ptr<OperationPass<FuncOp>> createConvertLinalgToAffineLoopsPass();
 
 /// Create a pass that bufferizes the body of a FuncOp and tries to reuse the
 /// buffers for those arguments that:
 ///   a) have been annotated 'inplaceable' and
 ///   b) whose buffer uses would be free of memory hazards.
 std::unique_ptr<Pass> createLinalgComprehensiveFuncBufferizePass();
 
+/// This pass implements a cross-dialect bufferization approach and performs an
+/// analysis to determine which op operands and results may be bufferized in the
+/// same buffers. The analysis is performed on topologically sorted CallOp and
+/// FuncOp within a module. It provides analyses and bufferization across
+/// function boundaries. Within a single function body, the bufferization used
+/// is that provided by `LinalgComprehensiveFuncBufferizePass`.
+std::unique_ptr<Pass> createLinalgComprehensiveModuleBufferizePass();
+
 /// Create a pass to convert Linalg operations which work on tensors to use
 /// buffers instead.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgBufferizePass();
 
 /// Create a pass to conver named Linalg operations to Linalg generic
 /// operations.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgGeneralizationPass();
 

mlir/include/mlir/Dialect/Linalg/Passes.td

def LinalgComprehensiveFuncBufferize :
   let summary = "Bufferize (tensor into memref) the body of a FuncOp and try "
     "to reuse the buffers for those arguments that "
     "a) have been annotated 'inplaceable' and "
     "b) whose buffer uses would be free of memory hazards";
   let description = [{
     This pass implements a cross-dialect bufferization approach and performs an
     analysis to determine which op operands and results may be bufferized in the
     same buffers. The analysis is performed on SSA use-def chains starting from
-    function operands that are annotated with the 'inplaceable' attribute
+    function operands that are annotated with the 'inplaceable' attribute.
   }];
   let options = [
     Option<"testAnalysisOnly", "test-analysis-only", "bool",
             /*default=*/"false",
            "Only runs inplaceability analysis (for testing purposes only)">
   ];
   let constructor = "mlir::createLinalgComprehensiveFuncBufferizePass()";
 }
 
+def LinalgComprehensiveModuleBufferize :
+    Pass<"linalg-comprehensive-module-bufferize", "ModuleOp"> {
+  let summary = "Bufferize (tensor into memref) for a Module.";
+  let description = [{
+    This pass implements a cross-dialect bufferization approach and performs an
+    analysis to determine which op operands and results may be bufferized in the
+    same buffers. The analysis is performed on topologically sorted CallOp and
+    FuncOp within a module. It provides analyses and bufferization across
+    function boundaries. Within a single function body, the bufferization used
+    is that provided by `-linalg-comprehensive-func-bufferize`.
+  }];
+  let options = [
+    Option<"testAnalysisOnly", "test-analysis-only", "bool",
+            /*default=*/"false",
+           "Only runs inplaceability analysis (for testing purposes only)">
+  ];
+  let constructor = "mlir::createLinalgComprehensiveModuleBufferizePass()";
+}
+
 def LinalgFoldUnitExtentDims : FunctionPass<"linalg-fold-unit-extent-dims"> {
   let summary = "Remove unit-extent dimension in Linalg ops on tensors";
   let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()";
   let options = [
     Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool",
             /*default=*/"false",
            "Only folds the one-trip loops from Linalg ops on tensors "
            "(for testing purposes only)">

mlir/include/mlir/IR/FunctionSupport.h

void setArgAttr(unsigned index, StringRef name, Attribute value) {
     setArgAttr(index, Identifier::get(name, this->getOperation()->getContext()),
                value);
   }
 
   /// Remove the attribute 'name' from the argument at 'index'. Return the
   /// attribute that was erased, or nullptr if there was no attribute with such
   /// name.
   Attribute removeArgAttr(unsigned index, Identifier name);
+  Attribute removeArgAttr(unsigned index, StringRef name) {
+    return removeArgAttr(
+        index, Identifier::get(name, this->getOperation()->getContext()));
+  }
 
   //===--------------------------------------------------------------------===//
   // Result Attributes
   //===--------------------------------------------------------------------===//
 
   /// FunctionLike operations allow for attaching attributes to each of the
   /// respective function results. These result attributes are stored as
   /// DictionaryAttrs in the main operation attribute dictionary. The name of

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

 // of ops that have well-defined inplace semantics.
 // This pass caters to high-performance codegen where buffer reuse is deemed
 // critical: the pass should fail if the bufferized form of the function needs
 // to return any buffer.
 // Generic control-flow and branching are unsupported.
 // Composability with extensible set of ops is not a first-class concern.
 //
 // Bufferization occurs by:
-//  a. performing an inPlace analysis `inPlaceAnalysisFuncOpInternals`
+//  a. performing an inPlace analysis `inPlaceAnalysisFuncOpBody`
 //     which marks each operation within the function with the
 //     `kInPlaceResultsAttrName` attribute.
 //  b. traversing each operation in the function and rewriting it in
 //     buffer form and keeping a BlockAndValueMapping mapping of the
 //     rewrites. New allocations are introduced during this step.
 //     TODO: Allocation + depending op hoisting to outermost enclosing
 //     sequential scope.
 //  c. at the end of this bufferization, 3 cases may occur:

 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
 
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X)
 
 //===----------------------------------------------------------------------===//
+// Generic helpers.
+//===----------------------------------------------------------------------===//
+
+/// Return the FuncOp called by `callOp`.
+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));
+}
+
+//===----------------------------------------------------------------------===//
 // Bufferization-specific BlockAndValueMapping support with debugging.
 //===----------------------------------------------------------------------===//
 
 /// Wrapper for better debugging.
 static void map(BlockAndValueMapping &bvm, ValueRange keys, ValueRange values) {
   assert(!keys.empty() && "Unexpected empty keys");
   LDBG("Map: " << keys.front() << " to " << values.front() << '\n');
   return bvm.map(keys, values);
if (auto bbArg = key.dyn_cast<BlockArgument>()) {
     if (isa<FuncOp>(key.getParentBlock()->getParentOp()))
       parentOp = key.getParentBlock()->getParentOp();
     else
       parentOp = key.getParentBlock()->getParentOp()->getParentOfType<FuncOp>();
   } else {
     parentOp = key.getDefiningOp()->getParentOfType<FuncOp>();
   }
   LDBG("In func:\n" << *parentOp << "NO VALUE FOR KEY: " << key << '\n');
+  (void)parentOp;
   return Value();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific attribute manipulation.
 // These could be simplified with helper structs on the side, for now attributes
 // allow simple embedding in the IR which simplifies testing.
 // This could also be folded in BufferizationAliasInfo or a Bufferizer class
static InPlaceSpec getInPlace(BlockArgument bbArg) {
   //   2. Or the matching iter operand is bufferized inplace and bbArg just
   //      bufferizes to that too.
   if (auto forOp = dyn_cast<scf::ForOp>(bbArg.getOwner()->getParentOp()))
     return InPlaceSpec::True;
   // Unknown cases.
   return InPlaceSpec::None;
 }
 
+/// Set the attribute that triggers inplace bufferization on a FuncOp argument
+/// `bbArg`.
+static void
+setInPlaceFuncArgument(BlockArgument bbArg,
+                       InPlaceSpec inPlaceSpec = InPlaceSpec::True) {
+  auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
+  funcOp.setArgAttr(
+      bbArg.getArgNumber(), LinalgDialect::kInplaceableAttrName,
+      BoolAttr::get(bbArg.getContext(), inPlaceSpec == InPlaceSpec::True));
+}
+
+/// Remove the attribute that triggers inplace bufferization on a FuncOp
+/// argument `bbArg`.
+static void removeInPlaceFuncArgument(BlockArgument bbArg) {
+  auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
+  funcOp.removeArgAttr(bbArg.getArgNumber(),
+                       LinalgDialect::kInplaceableAttrName);
+}
+
 LLVM_ATTRIBUTE_UNUSED static InPlaceSpec getInPlace(Value v) {
   if (auto bbArg = v.dyn_cast<BlockArgument>())
     return getInPlace(bbArg);
   return getInPlace(v.cast<OpResult>());
 }
 
 //===----------------------------------------------------------------------===//
 // Op-specific semantics helper to retrieve matching inplaceable result.

 //===----------------------------------------------------------------------===//
 
 /// Return `true` if the op is explicitly supported by bufferization or if it
 /// has no result tensors.
 /// Other cases must be conservative.
 static bool hasKnownBufferizationAliasingBehavior(Operation *op) {
   return
       // clang-format off
-      isa<scf::ForOp,
+      isa<CallOpInterface,
+          scf::ForOp,
           LinalgOp,
           ReturnOp,
           ExtractSliceOp,
           InsertSliceOp,
           VectorTransferOpInterface,
           scf::YieldOp>(op)
       // clang-format on
       || (none_of(op->getResultTypes(),
return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
         .Case<scf::ForOp,
               LinalgOp,
               InsertSliceOp,
               VectorTransferOpInterface>(
             [&](auto op) { return getInplaceableOpResult(op, opOperand); })
         // ExtractSliceOp is special, when bufferized inplace it just returns an
         // alias to its operand. Its result is never inplaceable on its operand.
         .Case([&](ExtractSliceOp op) { return OpResult(); })
+        // CallOpInterface is special, it needs to wait for the callee to be
+        // bufferized and needs to inspect the BufferAliasInfo object. It can't
+        // make a proper determination by itself and needs to be conservative.
+        .Case([&](CallOpInterface op) { return OpResult(); })
         // Other ops.
         .Default([&](Operation *op) { return OpResult(); });
   // clang-format on
 }
 
 /// Determine which OpOperand* will alias with `result` if the op is bufferized
 /// in place.
 /// Return None if the owner of `opOperand` does not have known
static bool bufferizesToMemoryRead(OpOperand &opOperand) {
   // ExtractSliceOp alone doesn't bufferize to a memory read, one of its uses
   // may.
   if (isa<ExtractSliceOp>(opOperand.getOwner()))
     return false;
   // scf::ForOp alone doesn't bufferize to a memory read, one of the uses of its
   // matching bbArg may.
   if (isa<scf::ForOp>(opOperand.getOwner()))
     return false;
+  // CallOpInterface alone doesn't bufferize to a memory read, one of the uses
+  // of the matching bbArg may. It is the responsibility of the caller to
+  // inspect bbArgs. In the absence of a BufferizationAliasInfo, we need to be
+  // conservative.
+  if (auto callOp = dyn_cast<CallOpInterface>(opOperand.getOwner()))
+    return true;
   if (auto linalgOp = dyn_cast<LinalgOp>(opOperand.getOwner()))
     return linalgOp.isInputTensor(&opOperand) ||
            linalgOp.isInitTensor(&opOperand);
   // All other cases are considered to bufferize to memory reads.
   // In particular, terminators are often the last use and need to be considered
   // as reads to return the proper value and avoid WAW clobbers.
   return true;
 }
 
 /// Return true if `opOperand` bufferizes to a memory write.
 /// If inPlaceSpec is different from InPlaceSpec::None, additionally require the
 /// write to match the inplace specification.
 static bool
 bufferizesToMemoryWrite(OpOperand &opOperand,
                         InPlaceSpec inPlaceSpec = InPlaceSpec::None) {
+  // These terminators are not writes.
+  if (isa<ReturnOp, linalg::YieldOp, scf::YieldOp>(opOperand.getOwner()))
+    return false;
+  // ExtractSliceOp alone doesn't bufferize to a memory write, one of its uses
+  // may.
+  if (isa<ExtractSliceOp>(opOperand.getOwner()))
+    return false;
+  // CallOpInterface alone doesn't bufferize to a memory write, one of the uses
+  // of the matching bbArg may. It is the responsibility of the caller to
+  // inspect bbArgs. In the absence of a BufferizationAliasInfo, we need to be
+  // conservative.
+  if (auto callOp = dyn_cast<CallOpInterface>(opOperand.getOwner()))
+    return true;
   Optional<OpResult> maybeOpResult = getAliasingOpResult(opOperand);
   // Unknown op that returns a tensor. The inplace analysis does not support
   // it. Conservatively return true.
   if (!maybeOpResult)
     return true;
   // Supported op without a matching result for opOperand (e.g. ReturnOp).
   // This does not bufferize to a write.
   if (!*maybeOpResult)
     return false;
-  // These terminators are not writes.
-  if (isa<ReturnOp, linalg::YieldOp, scf::YieldOp>(opOperand.getOwner()))
-    return false;
-  // ExtractSliceOp alone doesn't bufferize to a memory write, one of its uses
-  // may.
-  if (maybeOpResult->getDefiningOp<ExtractSliceOp>())
-    return false;
   // If we have a matching OpResult, this is a write.
   // Additionally allow to restrict to only inPlace write, if so specified.
   return inPlaceSpec == InPlaceSpec::None ||
          getInPlace(*maybeOpResult) == inPlaceSpec;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific alias analysis.
public:
   /// Specify fine-grain relationship between buffers to enable more analysis.
   enum class BufferRelation {
     None,
     // TODO: ResultContainsOperand,
     // TODO: OperandContainsResult,
     Equivalent
   };
 
-  explicit BufferizationAliasInfo(FuncOp funcOp);
+  explicit BufferizationAliasInfo(Operation *rootOp);
+
+  /// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
+  /// beginning the alias and equivalence sets only contain `v` itself.
+  void createAliasInfoEntry(Value v);
 
   /// Return true if the buffer to which `operand` would bufferize aliases a
   /// buffer that is known to not be writeable. This implies that the matching
   /// OpResult cannot be bufferized inplace.
   bool aliasesNonWriteableBuffer(OpOperand &operand) const;
 
   /// Return true if the buffer to which `operand` would bufferize is equivalent
   /// to some use that would bufferize to a write to a buffer.
private:
   /// "equivalent" buffer relationships.
   DenseMap<Value, DenseSet<Value>> aliasInfo;
 
   /// Auxiliary structure to store all the equivalent buffer classes.
   llvm::EquivalenceClasses<ValueWrapper> equivalentInfo;
 };
 } // namespace
 
-BufferizationAliasInfo::BufferizationAliasInfo(FuncOp funcOp) {
-  funcOp.walk([&](Operation *op) {
-    for (Value v : op->getResults()) {
-      if (!v.getType().isa<TensorType>())
-        continue;
-      assert(getInPlace(v) == InPlaceSpec::None &&
-             "unexpected inplace in analysis.");
-      DenseSet<Value> selfSet;
-      selfSet.insert(v);
-      aliasInfo.try_emplace(v, selfSet);
-      equivalentInfo.insert(v);
-    }
-    for (Region &r : op->getRegions()) {
-      for (Block &b : r.getBlocks()) {
-        for (auto bbArg : b.getArguments()) {
-          if (!bbArg.getType().isa<TensorType>())
-            continue;
-          DenseSet<Value> selfSet;
-          selfSet.insert(bbArg);
-          aliasInfo.try_emplace(bbArg, selfSet);
-          equivalentInfo.insert(bbArg);
-        }
-      }
-    }
-  });
-}
+BufferizationAliasInfo::BufferizationAliasInfo(Operation *rootOp) {
+  rootOp->walk([&](Operation *op) {
+    for (Value v : op->getResults())
+      if (v.getType().isa<TensorType>())
+        createAliasInfoEntry(v);
+    for (Region &r : op->getRegions())
+      for (Block &b : r.getBlocks())
+        for (auto bbArg : b.getArguments())
+          if (bbArg.getType().isa<TensorType>())
+            createAliasInfoEntry(bbArg);
+  });
+}
+
+/// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
+/// beginning the alias and equivalence sets only contain `v` itself.
+void BufferizationAliasInfo::createAliasInfoEntry(Value v) {
 
 /// Return true if the buffer to which `operand` would bufferize aliases a
 /// buffer that is known to not be writeable. This implies that the matching
 /// OpResult cannot be bufferized inplace.
 bool BufferizationAliasInfo::aliasesNonWriteableBuffer(
     OpOperand &operand) const {
   LDBG("----Start aliasesNonWriteableBuffer\n");
   LDBG("-------for operand #" << operand.getOperandNumber() << ": "
for (OpOperand &operand : yieldOp->getOpOperands()) {
     }
   }
 
   return success();
 }
 
 /// Analyze the `funcOp` body to determine which OpResults are inplaceable.
 static LogicalResult
-inPlaceAnalysisFuncOpInternals(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
-                               const DominanceInfo &domInfo) {
+inPlaceAnalysisFuncOpBody(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
   LLVM_DEBUG(llvm::dbgs() << "\n\n");
   LDBG("Begin InPlaceAnalysisFuncOpInternals:\n" << funcOp << '\n');
   assert(funcOp && funcOp->getNumRegions() > 0 && !funcOp.body().empty() &&
          "expected a funcOp definition with a body");
 
   // Collect ops so we can build our own traversal.
   SmallVector<ExtractSliceOp> extractSliceOps;
   SmallVector<InsertSliceOp> insertSliceOps;

 void LinalgComprehensiveFuncBufferize::runOnFunction() {
   auto funcOp = getFunction();
 
   // Analysis phase.
   DominanceInfo domInfo(funcOp);
   BufferizationAliasInfo aliasInfo(funcOp);
   // If the analysis fails, just return. This is expected to reset the IR and no
   // single OpResult should be marked inPlace.
-  if (failed(inPlaceAnalysisFuncOpInternals(funcOp, aliasInfo, domInfo))) {
+  if (failed(inPlaceAnalysisFuncOpBody(funcOp, aliasInfo, domInfo))) {
     signalPassFailure();
     return;
   }
 
   if (testAnalysisOnly)
     return;
 
   // Bufferization phase.
   BlockAndValueMapping bvm;
   if (failed(bufferizeFuncOpInternals(funcOp, bvm, aliasInfo)))
     signalPassFailure();
 
   // Post-pass cleanup of inplaceable attributes.
   funcOp.walk([&](Operation *op) { op->removeAttr(kInPlaceResultsAttrName); });
 }
 
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveFuncBufferizePass() {
   return std::make_unique<LinalgComprehensiveFuncBufferize>();
 }
+
+//===----------------------------------------------------------------------===//
+// Bufferization entry-point for modules.
+//===----------------------------------------------------------------------===//
+
+/// Store all functions of the `moduleOp` in `orderedFuncOps`, sorted by
+/// callee-caller order (i.e. callees without callers first).
+/// Store the map of FuncOp to all its callers in `callerMap`.
+/// Return `failure()` if a cycle of calls is detected or if we are unable to
+/// retrieve the called FuncOp from any CallOpInterface.
+static LogicalResult
+getFuncOpsOrderedByCalls(ModuleOp moduleOp,
+                         SmallVectorImpl<FuncOp> &orderedFuncOps,
+                         DenseMap<FuncOp, DenseSet<Operation *>> &callerMap) {
+  // For each FuncOp, the set of functions called by it (i.e. the union of
+  // symbols of all nested CallOpInterfaceOp).
+  DenseMap<FuncOp, DenseSet<FuncOp>> calledBy;
+  // For each FuncOp, the number of CallOpInterface it contains.
+  DenseMap<FuncOp, unsigned> numberCallOpsContainedInFuncOp;
+  WalkResult res = moduleOp.walk([&](FuncOp funcOp) {
+    numberCallOpsContainedInFuncOp[funcOp] = 0;
+    return funcOp.walk([&](CallOpInterface callOp) {
+      FuncOp calledFunction = getCalledFunction(callOp);
+      if (!calledFunction)
+        return WalkResult::interrupt();
+      auto it = callerMap.try_emplace(calledFunction, DenseSet<Operation *>{});
+      it.first->getSecond().insert(callOp);
+      if (calledBy[calledFunction].count(funcOp) == 0) {
+        calledBy[calledFunction].insert(funcOp);
+        numberCallOpsContainedInFuncOp[funcOp]++;
+      }
+      return WalkResult::advance();
+    });
+  });
+  if (res.wasInterrupted())
+    return failure();
+  // Iteratively remove function operation that do not call any of the
+  // functions remaining in the callCounter map and add them to the worklist.
+  while (!numberCallOpsContainedInFuncOp.empty()) {
+    auto it = llvm::find_if(numberCallOpsContainedInFuncOp,
+                            [](auto entry) { return entry.getSecond() == 0; });
+    if (it == numberCallOpsContainedInFuncOp.end())
+      return moduleOp.emitOpError(
+          "expected callgraph to be free of circular dependencies.");
+    orderedFuncOps.push_back(it->getFirst());
+    for (auto callee : calledBy[it->getFirst()])
+      numberCallOpsContainedInFuncOp[callee]--;
+    numberCallOpsContainedInFuncOp.erase(it);
+  }
+  return success();
+}
+
+namespace {
+struct LinalgComprehensiveModuleBufferize
+    : public LinalgComprehensiveModuleBufferizeBase<
+          LinalgComprehensiveModuleBufferize> {
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<linalg::LinalgDialect, memref::MemRefDialect>();
+  }
+};
+} // end namespace
+
+void LinalgComprehensiveModuleBufferize::runOnOperation() {
+  ModuleOp moduleOp = getOperation();
+
+  SmallVector<FuncOp> orderedFuncOps;
+  DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
+  if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
+    return signalPassFailure();
+
+  DominanceInfo domInfo(moduleOp);
+  BufferizationAliasInfo aliasInfo(moduleOp);
+  // 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) {
+    // No body => no analysis.
+    if (funcOp.body().empty())
+      continue;
+
+    // In a first approximation:
+    // =========================
+    // If the function is called, we can allocate on the caller side which lets
+    // us force inplace arguments at function boundaries.
+    // TODO: do not rely on this behavior.
+    if (callerMap.find(funcOp) != callerMap.end())
+      for (BlockArgument bbArg : funcOp.getArguments())
+        if (bbArg.getType().isa<TensorType>())
+          setInPlaceFuncArgument(bbArg);
+
+    // If the analysis fails, just return.
+    if (failed(inPlaceAnalysisFuncOpBody(funcOp, aliasInfo, domInfo))) {
+      signalPassFailure();
+      return;
+    }
+
+    // TODO: Bufferization phase.
+  }
+  // Don't drop the attributes if we only want to report the analysis.
+  if (testAnalysisOnly)
+    return;
+
+  // Post-pass cleanup of inplaceable attributes.
+  moduleOp.walk(
+      [&](Operation *op) { op->removeAttr(kInPlaceResultsAttrName); });
+  moduleOp.walk([&](FuncOp op) {
+    for (BlockArgument bbArg : op.getArguments())
+      removeInPlaceFuncArgument(bbArg);
+  });
+}
+
+std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass() {
+  return std::make_unique<LinalgComprehensiveModuleBufferize>();
+}

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

+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize=test-analysis-only -split-input-file | FileCheck %s
+
+func private @foo(tensor<64xf32>)
+
+// CHECK-LABEL: dependence_through_call
+func @dependence_through_call(%I : tensor<64xf32> {linalg.inplaceable = true}) {
+  %f1 = constant 1.000000e+00 : f32
+  %f2 = constant 2.000000e+00 : f32
+
+  // 2. %B already bufferizes inplace, %A would alias and have a different
+  // value. The calls to `foo` are determined to read conservatively, so %A
+  // cannot bufferize inplace.
+  //     CHECK: fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
+  %A = linalg.fill(%f1, %I) : f32, tensor<64xf32> -> tensor<64xf32>
+
+  // 1. Bufferizes inplace: no alias to %A is yet possible.
+  //     CHECK: fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %B = linalg.fill(%f2, %I) : f32, tensor<64xf32> -> tensor<64xf32>
+
+  call @foo(%A) : (tensor<64xf32>) -> ()
+  call @foo(%B) : (tensor<64xf32>) -> ()
+
+  return
+}
+
+// -----
+
+func private @foo(tensor<64xf32>)
+
+func private @bar(%A : tensor<64xf32>) {
+  call @foo(%A) : (tensor<64xf32>) -> ()
+  return
+}
+
+func @read_dependence_through_scf_and_call(
+    %I : tensor<64xf32> {linalg.inplaceable = true},
+    %I2 : tensor<64xf32> {linalg.inplaceable = true}) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c10 = constant 10 : index
+  %f1 = constant 1.000000e+00 : f32
+  %f2 = constant 2.000000e+00 : f32
+
+  // 5. %B bufferizes inplace, %A would alias and have a different value.
+  // The calls to `foo` are determined to read conservatively, so %A cannot
+  // bufferize inplace.
+  //     CHECK: fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
+  %A = linalg.fill(%f1, %I) : f32, tensor<64xf32> -> tensor<64xf32>
+
+  // 4. Bufferizes inplace: no alias to %A is yet possible.
+  //     CHECK: fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %B = linalg.fill(%f2, %I) : f32, tensor<64xf32> -> tensor<64xf32>
+
+  // 3. Does not read or write, bufferizes inplace.
+  //     CHECK: scf.for
+  //     CHECK: {__inplace_results_attr__ = ["true", "true"]}
+  %r:2 = scf.for %i = %c0 to %c10 step %c1 iter_args(%0 = %A, %1 = %B)
+    -> (tensor<64xf32>, tensor<64xf32>)
+  {
+    scf.yield %0, %1 : tensor<64xf32>, tensor<64xf32>
+  }
+  call @foo(%r#0) : (tensor<64xf32>) -> ()
+  call @foo(%r#1) : (tensor<64xf32>) -> ()
+
+  // 2. %B2 already bufferizes inplace, %A2 would alias and have a different
+  // value. The calls to `foo` are determined to read conservatively, so %A2
+  // cannot bufferize inplace.
+  //     CHECK: fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
+  %A2 = linalg.fill(%f1, %I2) : f32, tensor<64xf32> -> tensor<64xf32>
+
+  // 1. Bufferizes inplace: no alias to %A2 is yet possible.
+  //     CHECK: fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %B2 = linalg.fill(%f2, %I2) : f32, tensor<64xf32> -> tensor<64xf32>
+
+  call @bar(%A2) : (tensor<64xf32>) -> ()
+  call @bar(%B2) : (tensor<64xf32>) -> ()
+  return
+}

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

+// RUN: mlir-opt %s -linalg-comprehensive-module-bufferize -split-input-file -verify-diagnostics
+
+// -----
+
+// expected-error @-3 {{expected callgraph to be free of circular dependencies}}
+
+func @foo() {
+  call @bar() : () -> ()
+  return
+}
+
+func @bar() {
+  call @foo() : () -> ()
+  return
+}