[mlir][linalg][bufferize][NFC] Move helper functions to op interface

[springerm]

Summary

Also enclose all bufferization code in a new namespace: comprehensive_bufferize

Depends On D113339

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

 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 class BlockAndValueMapping;
 
 namespace linalg {
+namespace comprehensive_bufferize {
+
 struct AllocationCallbacks;
 class BufferizationAliasInfo;
 
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
   None,
   // TODO: ResultContainsOperand,
   // TODO: OperandContainsResult,
   Equivalent
 };
+
+/// Determine which OpOperand* will alias with `result` if the op is bufferized
+/// in place. Return an empty vector if the op is not bufferizable.
+SmallVector<OpOperand *> getAliasingOpOperand(OpResult result);
+
+/// Determine which OpResult will alias with `opOperand` if the op is bufferized
+/// in place. Return an empty OpResult if the op is not bufferizable.
+OpResult getAliasingOpResult(OpOperand &opOperand);
+
+/// Return true if `opOperand` bufferizes to a memory read. Return `true` if the
+/// op is not bufferizable.
+bool bufferizesToMemoryRead(OpOperand &opOperand);
+
+/// Return true if `opOperand` bufferizes to a memory write. Return
+/// `true` if the op is not bufferizable.
+bool bufferizesToMemoryWrite(OpOperand &opOperand);
+
+/// Return true if `opOperand` does neither read nor write but bufferizes to an
+/// alias. Return false if the op is not bufferizable.
+bool bufferizesToAliasOnly(OpOperand &opOperand);
+
+/// Return true if the given value is read by an op that bufferizes to a memory
+/// read. Also takes into account ops that create an alias but do not read by
+/// themselves (e.g., ExtractSliceOp).
+bool isValueRead(Value value);
+
+/// Return the relationship between the operand and the its corresponding
+/// OpResult that it may alias with. Return None if the op is not bufferizable.
+BufferRelation bufferRelation(OpOperand &opOperand);
+
+} // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_

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

 
 include "mlir/IR/OpBase.td"
 
 def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
   let description = [{
     An op interface for Comprehensive Bufferization. Ops that implement this
     interface can be bufferized using Comprehensive Bufferization.
   }];
-  let cppNamespace = "::mlir::linalg";
+  let cppNamespace = "::mlir::linalg::comprehensive_bufferize";
   let methods = [
       InterfaceMethod<
         /*desc=*/[{
           Return `true` if the given OpOperand bufferizes to a memory read. This
           method will never be called on OpOperands that do not have a tensor
           type.
 
           Note: It is always safe to consider an OpOperand as a memory read,

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

 namespace mlir {
 
 class DominanceInfo;
 class FuncOp;
 class GlobalCreator;
 class ModuleOp;
 
 namespace linalg {
+namespace comprehensive_bufferize {
 
 // TODO: from some HW description.
 static constexpr int64_t kBufferAlignments = 128;
 
 /// The BufferizationAliasInfo class maintains a list of buffer aliases and
 /// equivalence classes to support bufferization.
 /// ExtractSliceOps have special behavior, they act as a level of indirection
 /// for bufferization. They don't create reads or writes themselves and analysis
struct BufferizationOptions {
   bool allowReturnMemref = false;
   unsigned analysisFuzzerSeed = 0;
   bool testAnalysisOnly = false;
 };
 
 LogicalResult runComprehensiveBufferize(ModuleOp moduleOp,
                                         const BufferizationOptions &options);
 
+} // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H

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

 //===- BufferizableOpInterface.cpp - Comprehensive Bufferize --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/IR/Operation.h"
 
 namespace mlir {
 namespace linalg {
+namespace comprehensive_bufferize {
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp.inc"
 
+} // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
+
+using namespace mlir;
+using namespace linalg::comprehensive_bufferize;
+
+//===----------------------------------------------------------------------===//
+// Helper functions for BufferizableOpInterface
+//===----------------------------------------------------------------------===//
+
+/// Determine which OpOperand* will alias with `result` if the op is bufferized
+/// in place. Return an empty vector if the op is not bufferizable.
+SmallVector<OpOperand *>
+mlir::linalg::comprehensive_bufferize::getAliasingOpOperand(OpResult result) {
+  if (Operation *op = result.getDefiningOp())
+    if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+      return bufferizableOp.getAliasingOpOperand(result);
+  return {};
+}
+
+/// Determine which OpResult will alias with `opOperand` if the op is bufferized
+/// in place. Return an empty OpResult if the op is not bufferizable.
+OpResult mlir::linalg::comprehensive_bufferize::getAliasingOpResult(
+    OpOperand &opOperand) {
+  if (auto bufferizableOp =
+          dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
+    return bufferizableOp.getAliasingOpResult(opOperand);
+  return OpResult();
+}
+
+/// Return true if `opOperand` bufferizes to a memory read. Return `true` if the
+/// op is not bufferizable.
+bool mlir::linalg::comprehensive_bufferize::bufferizesToMemoryRead(
+    OpOperand &opOperand) {
+  if (auto bufferizableOp =
+          dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
+    return bufferizableOp.bufferizesToMemoryRead(opOperand);
+
+  // Unknown op that returns a tensor. The inplace analysis does not support it.
+  // Conservatively return true.
+  return true;
+}
+
+/// Return true if `opOperand` bufferizes to a memory write. Return
+/// `true` if the op is not bufferizable.
+bool mlir::linalg::comprehensive_bufferize::bufferizesToMemoryWrite(
+    OpOperand &opOperand) {
+  if (auto bufferizableOp =
+          dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
+    return bufferizableOp.bufferizesToMemoryWrite(opOperand);
+
+  // Unknown op that returns a tensor. The inplace analysis does not support it.
+  // Conservatively return true.
+  return true;
+}
+
+/// Return true if `opOperand` does neither read nor write but bufferizes to an
+/// alias. Return false if the op is not bufferizable.
+bool mlir::linalg::comprehensive_bufferize::bufferizesToAliasOnly(
+    OpOperand &opOperand) {
+  if (auto bufferizableOp =
+          dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
+    return bufferizableOp.bufferizesToAliasOnly(opOperand);
+
+  // Unknown op that returns a tensor. The inplace analysis does not support it.
+  // Conservatively return false.
+  return false;
+}
+
+/// Return true if the given value is read by an op that bufferizes to a memory
+/// read. Also takes into account ops that create an alias but do not read by
+/// themselves (e.g., ExtractSliceOp).
+bool mlir::linalg::comprehensive_bufferize::isValueRead(Value value) {
+  SmallVector<OpOperand *> workingSet;
+  for (OpOperand &use : value.getUses())
+    workingSet.push_back(&use);
+
+  while (!workingSet.empty()) {
+    OpOperand *uMaybeReading = workingSet.pop_back_val();
+    // Skip over all ops that neither read nor write (but create an alias).
+    if (bufferizesToAliasOnly(*uMaybeReading))
+      for (OpOperand &use : getAliasingOpResult(*uMaybeReading).getUses())
+        workingSet.push_back(&use);
+    if (bufferizesToMemoryRead(*uMaybeReading))
+      return true;
+  }
+
+  return false;
+}
+
+/// Return the relationship between the operand and the its corresponding
+/// OpResult that it may alias with. Return None if the op is not bufferizable.
+BufferRelation
+mlir::linalg::comprehensive_bufferize::bufferRelation(OpOperand &opOperand) {
+  if (auto bufferizableOp =
+          dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
+    return bufferizableOp.bufferRelation(opOperand);
+
+  // Unknown op that returns a tensor. The inplace analysis does not support it.
+  // Conservatively return None.
+  return BufferRelation::None;
+}

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

 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 
 #define DEBUG_TYPE "comprehensive-module-bufferize"
 
 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
+using namespace comprehensive_bufferize;
 
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X)
 
 // Forward declarations.
 static std::string printOperationInfo(Operation *, bool prefix = true);
 static std::string printValueInfo(Value, bool prefix = true);
 
static std::string printValueInfo(Value value, bool prefix) {
   llvm::raw_string_ostream os(result);
   AsmState state(value.getParentRegion()->getParentOfType<mlir::FuncOp>());
   os << value;
   printTensorOrBufferInfo("\n\t - ", value, state, os);
   return result;
 }
 
 //===----------------------------------------------------------------------===//
-// Helper functions for BufferizableOpInterface
-//===----------------------------------------------------------------------===//
-
-/// Determine which OpOperand* will alias with `result` if the op is bufferized
-/// in place. Return an empty vector if the op is not bufferizable.
-static SmallVector<OpOperand *> getAliasingOpOperand(OpResult result) {
-  if (Operation *op = result.getDefiningOp())
-    if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
-      return bufferizableOp.getAliasingOpOperand(result);
-  return {};
-}
-
-/// Determine which OpResult will alias with `opOperand` if the op is bufferized
-/// in place. Return an empty OpResult if the op is not bufferizable.
-static OpResult getAliasingOpResult(OpOperand &opOperand) {
-  if (auto bufferizableOp =
-          dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
-    return bufferizableOp.getAliasingOpResult(opOperand);
-  return OpResult();
-}
-
-/// Return true if `opOperand` bufferizes to a memory read. Return `true` if the
-/// op is not bufferizable.
-static bool bufferizesToMemoryRead(OpOperand &opOperand) {
-  if (auto bufferizableOp =
-          dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
-    return bufferizableOp.bufferizesToMemoryRead(opOperand);
-
-  // Unknown op that returns a tensor. The inplace analysis does not support it.
-  // Conservatively return true.
-  return true;
-}
-
-/// Return true if `opOperand` bufferizes to a memory write. Return
-/// `true` if the op is not bufferizable.
-static bool bufferizesToMemoryWrite(OpOperand &opOperand) {
-  if (auto bufferizableOp =
-          dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
-    return bufferizableOp.bufferizesToMemoryWrite(opOperand);
-
-  // Unknown op that returns a tensor. The inplace analysis does not support it.
-  // Conservatively return true.
-  return true;
-}
-
-/// Return true if `opOperand` does neither read nor write but bufferizes to an
-/// alias. Return false if the op is not bufferizable.
-static bool bufferizesToAliasOnly(OpOperand &opOperand) {
-  if (auto bufferizableOp =
-          dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
-    return bufferizableOp.bufferizesToAliasOnly(opOperand);
-
-  // Unknown op that returns a tensor. The inplace analysis does not support it.
-  // Conservatively return false.
-  return false;
-}
-
-/// Return true if the given value is read by an op that bufferizes to a memory
-/// read. Also takes into account ops that create an alias but do not read by
-/// themselves (e.g., ExtractSliceOp).
-static bool isValueRead(Value value) {
-  SmallVector<OpOperand *> workingSet;
-  for (OpOperand &use : value.getUses())
-    workingSet.push_back(&use);
-
-  while (!workingSet.empty()) {
-    OpOperand *uMaybeReading = workingSet.pop_back_val();
-    // Skip over all ops that neither read nor write (but create an alias).
-    if (bufferizesToAliasOnly(*uMaybeReading))
-      for (OpOperand &use : getAliasingOpResult(*uMaybeReading).getUses())
-        workingSet.push_back(&use);
-    if (bufferizesToMemoryRead(*uMaybeReading))
-      return true;
-  }
-
-  return false;
-}
-
-/// Return the relationship between the operand and the its corresponding
-/// OpResult that it may alias with. Return None if the op is not bufferizable.
-static BufferRelation bufferRelation(OpOperand &opOperand) {
-  if (auto bufferizableOp =
-          dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
-    return bufferizableOp.bufferRelation(opOperand);
-
-  // Unknown op that returns a tensor. The inplace analysis does not support it.
-  // Conservatively return None.
-  return BufferRelation::None;
-}
-
-//===----------------------------------------------------------------------===//
 // Bufferization-specific alias analysis.
 //===----------------------------------------------------------------------===//
 
 /// Return true if the (ExtractSliceOp, InsertSliceOp) pair match (i.e.
 /// equivalent operand / result and same offset/sizes/strides specification).
 ///
 /// This is one particular type of relationship between ops on tensors that
 /// reduce to an equivalence on buffers. This should be generalized and
bufferizableInPlaceAnalysis(OpOperand &operand,
   if (OpResult result = bufferizableOp.getAliasingOpResult(operand))
     return bufferizableInPlaceAnalysisImpl(operand, result, aliasInfo, domInfo);
   return success();
 }
 
 /// Analyze the `ops` to determine which OpResults are inplaceable. Walk ops in
 /// reverse and bufferize ops greedily. This is a good starter heuristic.
 /// ExtractSliceOps are interleaved with other ops in traversal order.
-LogicalResult mlir::linalg::inPlaceAnalysis(SmallVector<Operation *> &ops,
-                                            BufferizationAliasInfo &aliasInfo,
-                                            const DominanceInfo &domInfo,
-                                            unsigned analysisFuzzerSeed) {
+LogicalResult mlir::linalg::comprehensive_bufferize::inPlaceAnalysis(
+    SmallVector<Operation *> &ops, BufferizationAliasInfo &aliasInfo,
+    const DominanceInfo &domInfo, unsigned analysisFuzzerSeed) {
   if (analysisFuzzerSeed) {
     // This is a fuzzer. For testing purposes only. Randomize the order in which
     // operations are analyzed. The bufferization quality is likely worse, but
     // we want to make sure that no assertions are triggered anywhere.
     std::mt19937 g(analysisFuzzerSeed);
     llvm::shuffle(ops.begin(), ops.end(), g);
   }
 
inPlaceAnalysisFuncOpBody(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
 
   return res;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization entry-point for functions.
 //===----------------------------------------------------------------------===//
 
-Optional<Value>
-mlir::linalg::defaultAllocationFn(OpBuilder &b, Location loc, MemRefType type,
-                                  const SmallVector<Value> &dynShape) {
+Optional<Value> mlir::linalg::comprehensive_bufferize::defaultAllocationFn(
+    OpBuilder &b, Location loc, MemRefType type,
   Value allocated = b.create<memref::AllocOp>(
       loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
   return allocated;
 }
 
-void mlir::linalg::defaultDeallocationFn(OpBuilder &b, Location loc,
-                                         Value allocatedBuffer) {
+void mlir::linalg::comprehensive_bufferize::defaultDeallocationFn(
+    OpBuilder &b, Location loc, Value allocatedBuffer) {
   b.create<memref::DeallocOp>(loc, allocatedBuffer);
 }
 
-void mlir::linalg::defaultMemCpyFn(OpBuilder &b, Location loc, Value from,
+void mlir::linalg::comprehensive_bufferize::defaultMemCpyFn(OpBuilder &b,
+                                                            Location loc,
+                                                            Value from,
                                                             Value to) {
   b.create<CopyOp>(loc, from, to);
 }
 
-LogicalResult mlir::linalg::bufferizeOp(
+LogicalResult mlir::linalg::comprehensive_bufferize::bufferizeOp(
     Operation *op, BlockAndValueMapping &bvm, BufferizationAliasInfo &aliasInfo,
     AllocationCallbacks allocationFns,
     DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes) {
   OpBuilder b(op->getContext());
 
   // CallOps are handled separately.
   if (auto callOp = dyn_cast<CallOpInterface>(op)) {
     LDBG("Begin bufferize:\n" << callOp << '\n');
static void layoutPostProcessing(ModuleOp moduleOp) {
   }
 }
 
 /// Try to eliminate InitTensorOps inside funcOp. An InitTensorOp is replaced
 /// with the the result of `rewriteFunc` if it is anchored on a matching
 /// OpOperand. "Anchored" means that there is a path on the reverse SSA use-def
 /// chain, starting from the OpOperand and always following the aliasing
 /// OpOperand, that eventually ends at a single InitTensorOp.
-LogicalResult mlir::linalg::initTensorElimination(
+LogicalResult mlir::linalg::comprehensive_bufferize::initTensorElimination(
     FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo,
     std::function<bool(OpOperand &)> anchorMatchFunc,
     std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
     bool skipAnalysis) {
   OpBuilder b(funcOp->getContext());
 
   WalkResult status = funcOp->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {

 /// * The InsertSliceOp was decided to bufferize inplace.
 /// * On the reverse use-def chain path from the InsertSliceOp to the
 ///   InitTensorOp, all ops were decided to bufferize inplace and the buffer
 ///   relation is "equivalent" (TODO: can be relaxed if needed).
 /// * The reverse use-def chain has exactly one end, which is the InitTensorOp.
 ///
 /// Note that the newly inserted ExtractSliceOp may have to bufferize
 /// out-of-place due to RaW conflicts.
-LogicalResult mlir::linalg::eliminateInsertSliceAnchoredInitTensorOps(
-    FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo) {
+LogicalResult mlir::linalg::comprehensive_bufferize::
+    eliminateInsertSliceAnchoredInitTensorOps(FuncOp funcOp,
+                                              BufferizationAliasInfo &aliasInfo,
+                                              DominanceInfo &domInfo) {
   return initTensorElimination(
       funcOp, aliasInfo, domInfo,
       [](OpOperand &operand) {
         auto insertSliceOp = dyn_cast<InsertSliceOp>(operand.getOwner());
         if (!insertSliceOp)
           return false;
         // Only inplace bufferized InsertSliceOps are eligible.
         if (getInPlace(insertSliceOp->getOpResult(0)) != InPlaceSpec::True)
           return false;
         return &operand == &insertSliceOp->getOpOperand(0) /*source*/;
       },
       [](OpBuilder &b, Location loc, OpOperand &operand) {
         auto insertSliceOp = cast<InsertSliceOp>(operand.getOwner());
         auto extractOp = b.create<tensor::ExtractSliceOp>(
             loc, insertSliceOp.dest(), insertSliceOp.getMixedOffsets(),
             insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
         return extractOp.result();
       });
 }
 
-LogicalResult
-mlir::linalg::runComprehensiveBufferize(ModuleOp moduleOp,
-                                        const BufferizationOptions &options) {
+LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
+    ModuleOp moduleOp, const BufferizationOptions &options) {
   SmallVector<FuncOp> orderedFuncOps;
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
   DenseMap<FuncOp, FunctionType> bufferizedFunctionTypes;
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   DominanceInfo domInfo(moduleOp);
   BufferizationAliasInfo aliasInfo(moduleOp);

 // BufferizableOpInterface Implementations
 //===----------------------------------------------------------------------===//
 
 // TODO: Move these to a different file and BUILD target, so that they are
 // decoupled from ComprehensiveBufferize.
 
 namespace mlir {
 namespace linalg {
+namespace comprehensive_bufferize {
 namespace arith_ext {
 
 struct ConstantOpInterface
     : public BufferizableOpInterface::ExternalModel<ConstantOpInterface,
                                                     arith::ConstantOp> {
   SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
                                                 OpResult opResult) const {
     return {};
void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   // not possible to attach an external interface to an existing interface.
   // Therefore, attach the `BufferizableOpInterface` to all ops one-by-one.
   LinalgOpInterfaceHelper<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >::registerOpInterface(registry);
 }
 
+} // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir

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

 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
+using namespace mlir::linalg::comprehensive_bufferize;
 
 namespace {
 struct LinalgComprehensiveModuleBufferize
     : public LinalgComprehensiveModuleBufferizeBase<
           LinalgComprehensiveModuleBufferize> {
   LinalgComprehensiveModuleBufferize() {}
 
   LinalgComprehensiveModuleBufferize(