[mlir][linalg][bufferize][NFC] Specify bufferize traversal in `bufferize`

[springerm]

Summary

The interface method bufferize controls how (and it what order) nested ops are traversed. This simplifies bufferization of scf::ForOps and scf::IfOps, which used to need special rules in scf::YieldOp.

Depends On D114055

[nicolasvasilache]

Please improve the description as we discussed offline.

[springerm]

update

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

struct BufferizationState {
   DenseMap<FuncOp, FunctionType> bufferizedFunctionTypes;
 };
 
 /// Return the result buffer (memref) for a given OpResult (tensor). Allocate
 /// a new buffer and copy over data from the existing buffer if out-of-place
 /// bufferization is necessary.
 Value getResultBuffer(OpBuilder &b, OpResult result, BufferizationState &state);
 
+/// Bufferize all ops in the given region.
+LogicalResult bufferize(Region *region, BufferizationState &state);
+
+/// Bufferize all ops in the given block.
+LogicalResult bufferize(Block *block, BufferizationState &state);
+
 /// Bufferize the given op. If the op has no tensor OpOperands/OpResults, this
 /// function returns immediately. Otherwise, it calls the `bufferize` interface
 /// method of `BufferizableOpInterface`.
-LogicalResult bufferizeOp(Operation *op, BufferizationState &state);
+LogicalResult bufferize(Operation *op, BufferizationState &state);
 
 /// PostAnalysisSteps can be registered with `BufferizationOptions` and are
 /// executed after the analysis, but before bufferization. They can be used
 /// implement custom dialect-specific optimizations.
 struct PostAnalysisStep {
   virtual ~PostAnalysisStep() {}
 
   /// Run the post analysis step. This function may modify the IR, but must keep

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

if (!state.aliasInfo.isInPlace(result)) {
     return resultBuffer;
   }
 
   // Bufferizing in-place. No need to allocate a new buffer.
   return operandBuffer;
 }
 
 LogicalResult
-mlir::linalg::comprehensive_bufferize::bufferizeOp(Operation *op,
-                                                   BufferizationState &state) {
+mlir::linalg::comprehensive_bufferize::bufferize(Region *region,
+                                                 BufferizationState &state) {
+  for (Block &block : *region)
+    if (failed(bufferize(&block, state)))
+      return failure();
+  return success();
+}
+
+LogicalResult
+mlir::linalg::comprehensive_bufferize::bufferize(Block *block,
+                                                 BufferizationState &state) {
+  for (Operation &op : *block)
+    if (failed(bufferize(&op, state)))
+      return failure();
+  return success();
+}
+
+LogicalResult
+mlir::linalg::comprehensive_bufferize::bufferize(Operation *op,
   OpBuilder b(op->getContext());
 
   // Skip BufferCast and TensorLoad ops.
   if (isa<memref::BufferCastOp, memref::TensorLoadOp>(op))
     return success();
 
   // Bufferize using `BufferizableOpInterface`.
   b.setInsertionPoint(op);
   if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
     return bufferizableOp.bufferize(b, state);
 
   // Check if op has tensor results or operands.
   auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
   if (hasTensorResult && hasTensorOperand)
     return op->emitError() << "unsupported op with tensors";
 
+  // Bufferize all nested ops.
+  for (Region &region : op->getRegions())
+    if (failed(bufferize(&region, state)))
+      return failure();
+
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific BlockAndValueMapping support with debugging.
 //===----------------------------------------------------------------------===//
 
 /// Wrapper for better debugging.

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

static Value createNewAllocDeallocPairForShapedValue(
   return casted;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization as simple BlockAndValueMapping rewrites.
 //===----------------------------------------------------------------------===//
 
 /// FuncOp always creates TensorToMemRef ops.
-static LogicalResult bufferize(OpBuilder &b, FuncOp funcOp,
-                               BufferizationState &state) {
+static LogicalResult bufferizeFuncOp(FuncOp funcOp, BufferizationState &state) {
   // Take a guard before anything else.
-  OpBuilder::InsertionGuard g(b);
+  OpBuilder b(funcOp->getContext());
   b.setInsertionPointToStart(&funcOp.body().front());
+
+  // Create BufferCastOps for function args.
   for (auto bbArg : funcOp.getArguments()) {
     auto tensorType = bbArg.getType().dyn_cast<TensorType>();
     if (!tensorType)
       continue;
     auto rankedTensorType = tensorType.dyn_cast<RankedTensorType>();
     // Cast the tensor to the most dynamic buffer possible. Further
     // canonicalizations will clean up.
     Type memRefType = rankedTensorType
                           ? getDynamicMemRefType(rankedTensorType)
                           : getContiguousOrUnrankedMemRefType(tensorType);
     Value bufferCast =
         b.create<memref::BufferCastOp>(funcOp.getLoc(), memRefType, bbArg);
     state.aliasInfo.insertNewBufferEquivalence(bufferCast, bbArg);
     state.mapBuffer(bbArg, bufferCast);
   }
-  return success();
+
+  // Bufferize function body.
+  return bufferize(&funcOp.body(), state);
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
 /// Determine if `operand` can be bufferized in-place with `result`.
 static LogicalResult
inPlaceAnalysisFuncOpBody(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
   LogicalResult res =
       inPlaceAnalysis(ops, aliasInfo, domInfo, analysisFuzzerSeed);
   LDBG("End InPlaceAnalysisFuncOpInternals:\n" << funcOp << '\n');
 
   return res;
 }
 
 //===----------------------------------------------------------------------===//
-// Bufferization entry-point for functions.
-//===----------------------------------------------------------------------===//
-
-static LogicalResult bufferizeFuncOpInternals(FuncOp funcOp,
-                                              BufferizationState &state) {
-  LLVM_DEBUG(llvm::dbgs() << "\n\n");
-  LDBG("Begin BufferizeFuncOpInternals:\n" << funcOp << '\n');
-  OpBuilder b(funcOp->getContext());
-
-  // Start by bufferizing `funcOp` arguments.
-  if (failed(bufferize(b, funcOp, state)))
-    return failure();
-
-  auto walkFunc = [&](Operation *op) -> WalkResult {
-    if (failed(bufferizeOp(op, state)))
-      return failure();
-    return success();
-  };
-
-  // Bufferize ops pre-order, i.e., bufferize ops first, then their children.
-  // This is needed for ops with blocks that have BlockArguments. These must be
-  // mapped before bufferizing the children.
-  if (funcOp.walk<WalkOrder::PreOrder>(walkFunc).wasInterrupted())
-    return failure();
-
-  LDBG("End BufferizeFuncOpInternals:\n" << funcOp << '\n');
-
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
 // Bufferization entry-point for modules.
 //===----------------------------------------------------------------------===//
 
 /// Return the op with Allocate MemoryEffect if `v` is equivalent to such an
 /// an op. Return null otherwise.
 static Operation *getEquivalentAlloc(Value value,
                                      const BufferizationAliasInfo &aliasInfo) {
   Operation *res = nullptr;
for (const std::unique_ptr<PostAnalysisStep> &step :
       // Analyze ops that were created by the PostAnalysisStep.
       if (failed(inPlaceAnalysis(newOps, aliasInfo, domInfo)))
         return failure();
     }
 
     // Bufferization phase.
     if (!options.testAnalysisOnly) {
       // Bufferize all ops in funcOp.
-      if (failed(bufferizeFuncOpInternals(funcOp, state)))
+      if (failed(bufferizeFuncOp(funcOp, state)))
         return failure();
 
       // Erase all obsolete ops.
       state.eraseObsoleteOps();
     }
   }
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly) {
bool mustBufferizeInPlace(Operation *op, OpResult opResult) const {
     // ExecuteRegionOp results always bufferize in-place. Since they have no
     // OpOperands, they are mostly ignored by the analysis once alias sets are
     // set up.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
+    auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
     bool hasTensorReturnType = any_of(
         op->getResultTypes(), [](Type t) { return t.isa<TensorType>(); });
     if (hasTensorReturnType)
       return op->emitError(
           "scf.execute_region with tensor result not supported");
-    return success();
+    return comprehensive_bufferize::bufferize(&executeRegionOp.region(), state);
   }
 };
 
 struct IfOpInterface
     : public BufferizableOpInterface::ExternalModel<IfOpInterface, scf::IfOp> {
   SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
                                                 OpResult opResult) const {
     // IfOps do not have tensor OpOperands. The yielded value can be any SSA
struct IfOpInterface
   bool mustBufferizeInPlace(Operation *op, OpResult opResult) const {
     // IfOp results always bufferize in-place. Since they have no OpOperands,
     // they are mostly ignored by the analysis once alias sets are set up.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
-    // scf::IfOp is bufferized after scf::YieldOp in the else branch.
-    return success();
-  }
-};
+    auto ifOp = cast<scf::IfOp>(op);
 
-/// Bufferize the scf::IfOp. This function is called after the YieldOp was
-/// bufferized.
-static LogicalResult bufferizeIfOp(scf::IfOp ifOp, OpBuilder &b,
-                                   BufferizationState &state) {
-  // Take a guard before anything else.
-  OpBuilder::InsertionGuard g(b);
-  b.setInsertionPoint(ifOp);
+    // Bufferize then/else blocks.
+    if (failed(comprehensive_bufferize::bufferize(ifOp.thenBlock(), state)))
+      return failure();
+    if (failed(comprehensive_bufferize::bufferize(ifOp.elseBlock(), state)))
+      return failure();
 
     for (OpResult opResult : ifOp->getResults()) {
       if (!opResult.getType().isa<TensorType>())
         continue;
       // TODO: Atm we bail on unranked TensorType because we don't know how to
       // alloc an UnrankedMemRefType + its underlying ranked MemRefType.
       assert(opResult.getType().isa<RankedTensorType>() &&
              "unsupported unranked tensor");
 
       Value resultBuffer = getResultBuffer(b, opResult, state);
       if (!resultBuffer)
         return failure();
 
       state.aliasInfo.createAliasInfoEntry(resultBuffer);
       state.mapBuffer(opResult, resultBuffer);
     }
 
     return success();
   }
+};
 
 struct ForOpInterface
     : public BufferizableOpInterface::ExternalModel<ForOpInterface,
                                                     scf::ForOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     // scf::ForOp alone doesn't bufferize to a memory read, one of the uses of
     // its matching bbArg may.
     auto forOp = cast<scf::ForOp>(op);
bool isWritable(Operation *op, Value value) const {
     //      alloc + optional copy makes the bbArg itself inplaceable.
     //   2. Or the matching iter operand is bufferized inplace and bbArg just
     //      bufferizes to that too.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
-    // Note: This method is just setting up the mappings for the block arguments
-    // and the result buffer. The op is bufferized after the scf::YieldOp.
-
     auto forOp = cast<scf::ForOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
 
     for (OpResult opResult : forOp->getResults()) {
       if (!opResult.getType().isa<TensorType>())
         continue;
for (OpResult opResult : forOp->getResults()) {
       OpOperand &opOperand = forOp.getOpOperandForResult(opResult);
       BlockArgument bbArg = forOp.getRegionIterArgForOpOperand(opOperand);
       state.aliasInfo.createAliasInfoEntry(resultBuffer);
       state.aliasInfo.insertNewBufferEquivalence(bbArg, resultBuffer);
       state.mapBuffer(bbArg, resultBuffer);
       state.mapBuffer(opResult, resultBuffer);
     }
 
-    return success();
-  }
-};
+    // Bufferize loop body.
+    if (failed(comprehensive_bufferize::bufferize(&forOp.region(), state)))
+      return failure();
 
-/// Bufferize the scf::ForOp. This function is called after the YieldOp was
-/// bufferized.
-static LogicalResult bufferizeForOp(scf::ForOp forOp, OpBuilder &b,
-                                    BufferizationState &state) {
+    // Finish bufferizing scf::ForOp.
     auto yieldOp = cast<scf::YieldOp>(&forOp.region().front().back());
     for (OpOperand &operand : yieldOp->getOpOperands()) {
       auto tensorType = operand.get().getType().dyn_cast<TensorType>();
       if (!tensorType)
         continue;
 
       OpOperand &forOperand = forOp.getOpOperandForResult(
           forOp->getResult(operand.getOperandNumber()));
       auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
       Value yieldedBuffer = state.lookupBuffer(operand.get());
       Value bbArgBuffer = state.lookupBuffer(bbArg);
       if (!state.aliasInfo.areEquivalentBufferizedValues(yieldedBuffer,
                                                          bbArgBuffer)) {
         // TODO: this could get resolved with copies but it can also turn into
         // swaps so we need to be careful about order of copies.
         return yieldOp->emitError()
                << "Yield operand #" << operand.getOperandNumber()
                << " does not bufferize to an equivalent buffer to the matching"
                << " enclosing scf::for operand";
       }
 
       // Buffers are equivalent so the work is already done and we just yield
       // the bbArg so that it later canonicalizes away.
       operand.set(bbArg);
     }
     return success();
   }
+};
 
 struct YieldOpInterface
     : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
                                                     scf::YieldOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto yieldOp = cast<scf::YieldOp>(op);
-
-    if (auto execOp = dyn_cast<scf::ExecuteRegionOp>(yieldOp->getParentOp())) {
-      if (execOp->getNumResults() != 0)
-        return execOp->emitError(
-            "expected result-less scf.execute_region containing op");
-      return success();
-    }
-
-    // Bufferize scf::IfOp after bufferizing the scf::YieldOp in the else
-    // branch.
-    if (auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp())) {
-      if (ifOp.elseYield() != yieldOp)
-        return success();
-      return bufferizeIfOp(ifOp, b, state);
-    }
-
-    // Bufferize scf::ForOp after bufferizing the scf::YieldOp.
-    if (auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp()))
-      return bufferizeForOp(forOp, b, state);
-
-    return yieldOp->emitError("expected scf::ForOp parent for scf::YieldOp");
-  }
+    if (!isa<scf::ExecuteRegionOp, scf::IfOp, scf::ForOp>(
+            yieldOp->getParentOp()))
+      return yieldOp->emitError("unsupported scf::YieldOp parent");
 };
 
 } // namespace scf_ext
 
 namespace std_ext {
 
 struct CallOpInterface
     : public BufferizableOpInterface::ExternalModel<CallOpInterface, CallOp> {

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

LogicalResult bufferize(Operation *op, OpBuilder &b,
     // TODO: Helper method to avoid leaking impl details.
     tiledLoopOp->setAttr(
         TiledLoopOp::getOperandSegmentSizeAttr(),
         b.getI32VectorAttr(
             {numLoops, numLoops, numLoops,
              static_cast<int>(oldInputs.size()) + numNewInputBuffers,
              static_cast<int>(oldOutputs.size()) + numNewOutputBuffers}));
 
-    return success();
+    // Bufferize loop body.
+    return comprehensive_bufferize::bufferize(&tiledLoopOp.region(), state);
   }
 };
 
 struct YieldOpInterface
     : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
                                                     linalg::YieldOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return true;