[mlir][Linalg] Add comprehensive bufferization support for ConstantOp (13/n)

[nicolasvasilache]

Summary

ConstantOp are only supported in the ModulePass because they require a GlobalCreator object that must be constructed from a ModuleOp.
If the standlaone FunctionPass encounters a ConstantOp, bufferization fails.

[nicolasvasilache]

Add invalid test.

[silvas]

Can you just delete the function pass and have the module pass be the only registered pass?

[rriddle]

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

2187–2196

[bondhugula]

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

31–32

But for constants defined inside of functions, why should one fail bufferization and kill parallelism (using module pass instead of function pass)? Is the pass also transparently promoting the constants to global - something not immediately expected/implied from bufferization? And in any case, if this is what is being done, is there a motivation to even support the FunctionPass?

[nicolasvasilache]

I don't find that there is a clear unambiguous way of moving forward here, only somewhat unsatisfactory options.

Forcing everything to be a ModulePass is premature atm: others have expressed a wish for generalization to "more generic ops with regions" than just func, this is still unclear.

The parallelism granularity decision and implications is a particular tradeoff in performance vs simplicity/what pass writers have to care about vs how passes must be written.
I greatly appreciate the "parallelism for free" aspect that comes with FuncOp granularity so I don't imagine this is going to change (ever?) to give users more control as it would be a pretty big footgun.
I wonder however whether ModuleOp could gain a "synchronized region" for things that are not functions?
@herhut also mentioned wishes for transactional updates du ModuleOp.

The "transparently promoting the constants to global" part is debatable.
Atm, tensor constant bufferization is a specific independent ModuleOp pass.
Going to memref.global easily turns this into readonly memory and is pretty clean; the alternatives I see are either to unroll constants or plop them into a vector + vector.store, both of which have large drawbacks.

Another alternative that I'm leaning towards is to allow memref.global to be constructible within a FuncOp but illegal to lower to LLVM.
During conversion to LLVM, legalization would consist in hoisting to the first enclosing ModuleOp.

[silvas]

I would personally like to see this as a nice clean module pass, and then make targeted changes informed by concrete use cases (actual need for more parallelism, need to be applied locally within a particular kind of enclosing op, etc.). Really for more advanced use cases this is going to be accessed via a direct C++ API, so the pass structure we choose here should be clean and simple and have nice invariants that make it easy to reason about extracting a targeted C++ API for new use cases. For example, the C++ API could take a std::function<Value(ConstantOp op)> for creating memref values from constants, which for the module pass would use a GlobalCreator, but for other users would be implemented otherwise. I don't want to promote uses of ComprehensiveBufferize with this weird non-support for constants.

I don't think we want to allow memref.global inside a function. That just seems like a huge layering violation (e.g. it is a Symbol but a FuncOp is not a SymbolTable).

[nicolasvasilache]

After considering various options, @silvas was the one that made the most sense.

[nicolasvasilache]

Rebase on dropped func-bufferize pass.

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

 /// 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<CallOpInterface,
           tensor::CastOp,
+          ConstantOp,
           scf::ForOp,
           InitTensorOp,
           LinalgOp,
           ReturnOp,
           ExtractSliceOp,
           InsertSliceOp,
           VectorTransferOpInterface,
           scf::YieldOp>(op)
static OpResult getInplaceableOpResult(InsertSliceOp op, OpOperand &opOperand) {
   if (opOperand.get() != op.dest())
     return OpResult();
   return op->getResult(0);
 }
 
 /// Return the OpResult that may bufferize into the same buffer as `opOperand`
 /// when the op is bufferized inplace.
 /// Return null if no such result exists.
-static OpResult getInplaceableOpResult(tensor::CastOp op, OpOperand &opOperand) {
+static OpResult getInplaceableOpResult(tensor::CastOp op,
+                                       OpOperand &opOperand) {
   return op->getResult(0);
 }
 
 /// Return the OpResult that may bufferize into the same buffer as `opOperand`
 /// when the op is bufferized inplace.
 /// The inplace analysis uses this information along with interfering read
 /// analysis to determine which op results reuse the same buffer as some
 /// operand.

 /// bufferization aliasing behavior, which indicates that the op must allocate
 /// all of its tensor results.
 /// TODO: in the future this may need to evolve towards a list of OpOperand*.
 static Optional<OpOperand *> getAliasingOpOperand(OpResult result) {
   if (!hasKnownBufferizationAliasingBehavior(result.getDefiningOp()))
     return None;
   return TypeSwitch<Operation *, OpOperand *>(result.getDefiningOp())
       .Case([&](tensor::CastOp op) { return &op->getOpOperand(0); })
+      .Case([&](ConstantOp op) { return &op->getOpOperand(0); })
       .Case([&](LinalgOp op) {
         return op.getOutputTensorOperands()[result.getResultNumber()];
       })
       .Case([&](ExtractSliceOp op) { return &op->getOpOperand(0); })
       .Case([&](InsertSliceOp op) { return &op->getOpOperand(1); })
       .Case([&](vector::TransferWriteOp op) { return &op->getOpOperand(1); })
       // In the case of scf::ForOp, this currently assumes the iter_args / yield
       // are 1-1. This may fail and is verified at the end.

 /// TODO: in the future this may need to evolve towards a list of OpResult.
 static Optional<OpResult> getAliasingOpResult(OpOperand &opOperand) {
   if (!hasKnownBufferizationAliasingBehavior(opOperand.getOwner()))
     return None;
   return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
       // These terminators legitimately have no result.
       .Case<ReturnOp, linalg::YieldOp, scf::YieldOp>(
           [&](auto op) { return OpResult(); })
+      // ConstantOp is never inplaceable.
+      .Case([&](ConstantOp op) { return op->getResult(0); })
       // ExtractSliceOp is different: its result is not inplaceable on op.source
       // but when bufferized inplace, the result is an aliasing subregion of
       // op.source.
       .Case([&](ExtractSliceOp op) { return op->getResult(0); })
       // All other ops, including scf::ForOp, return the result of
       // `getInplaceableOpResult`.
       .Default(
           [&](Operation *op) { return getInplaceableOpResult(opOperand); });
Type memRefType =
                                               memorySpace);
   Value res = b.create<memref::CastOp>(castOp.getLoc(), memRefType,
                                        lookup(bvm, castOp.source()));
   aliasInfo.insertNewBufferEquivalence(res, castOp.getResult());
   map(bvm, castOp.getResult(), res);
   return success();
 }
 
+static LogicalResult bufferize(OpBuilder &b, ConstantOp constantOp,
+                               BlockAndValueMapping &bvm,
+                               BufferizationAliasInfo &aliasInfo,
+                               GlobalCreator &globals) {
+  if (!constantOp.getType().dyn_cast<RankedTensorType>())
+    return failure();
+
+  // Take a guard before anything else.
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPoint(constantOp);
+
+  auto globalMemref = globals.getGlobalFor(constantOp);
+  Value memref = b.create<memref::GetGlobalOp>(
+      constantOp.getLoc(), globalMemref.type(), globalMemref.getName());
+  aliasInfo.insertNewBufferEquivalence(memref, constantOp.getResult());
+  map(bvm, constantOp, memref);
+
+  return success();
+}
+
 /// DimOp tensor operand is modified inplace. This allows leaving dead
 /// tensors behind that will get DCE'd.
 static LogicalResult bufferize(OpBuilder &b, memref::DimOp dimOp,
                                BlockAndValueMapping &bvm,
                                BufferizationAliasInfo &aliasInfo) {
   if (dimOp.memrefOrTensor().getType().isa<RankedTensorType>()) {
     Value v = lookup(bvm, dimOp.memrefOrTensor());
     if (!v)

 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization entry-point for functions.
 //===----------------------------------------------------------------------===//
 
 static LogicalResult bufferizeFuncOpInternals(
     FuncOp funcOp, BlockAndValueMapping &bvm, BufferizationAliasInfo &aliasInfo,
-    DenseMap<FuncOp, FunctionType> &bufferizedFunctionTypes) {
+    DenseMap<FuncOp, FunctionType> &bufferizedFunctionTypes,
+    GlobalCreator *globals = nullptr) {
   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, bvm, aliasInfo)))
     return failure();
   // Walk in PreOrder to ensure ops with regions are handled before their body.
   // Since walk has to be PreOrder, we need to erase ops that require it
   // separately: this is the case for CallOp
   SmallVector<Operation *> toErase;
-  WalkResult result = funcOp.walk<WalkOrder::PreOrder>([&](Operation *op)
-                                                           -> WalkResult {
+  WalkResult result =
+      funcOp.walk<WalkOrder::PreOrder>([&](Operation *op) -> WalkResult {
         // clang-format off
     WalkResult result =
       TypeSwitch<Operation *, LogicalResult>(op)
       // Skip BufferCast and TensorLoad ops.
       .Case<memref::BufferCastOp,
             memref::TensorLoadOp>([&](auto) { return success(); })
       .Case<tensor::CastOp,
             memref::DimOp,
             scf::ForOp,
             InitTensorOp,
             LinalgOp,
             ReturnOp,
             ExtractSliceOp,
             InsertSliceOp,
             VectorTransferOpInterface,
             scf::YieldOp>([&](auto op) {
         LDBG("Begin bufferize:\n" << op << '\n');
         return bufferize(b, op, bvm, aliasInfo);
       })
       .Case([&](CallOpInterface op) {
         LDBG("Begin bufferize:\n" << op << '\n');
         return bufferize(b, op, bvm, aliasInfo, bufferizedFunctionTypes);
       })
+      .Case([&](ConstantOp op) {
+        if (!isaTensor(op.getResult().getType()))
+          return success();
+        LDBG("Begin bufferize:\n" << op << '\n');
+        if (!globals)
+          return LogicalResult(
+            op->emitError() << "Failed bufferizing tensor constant within a " <<
+              "function pass: cannot construct globals creator");
+        return bufferize(b, op, bvm, aliasInfo, *globals);
+      })
       .Default([&](Operation *op) {
         auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
         if (any_of(op->getOperandTypes(), isaTensor) ||
             any_of(op->getResultTypes(), isaTensor))
           return failure();
         return success();
       });
         // clang-format on
 
         // Register post-walk erasure, if necessary.
         if (isa<CallOpInterface>(op))
           if (llvm::any_of(op->getOperandTypes(), isaTensor) ||
               llvm::any_of(op->getResultTypes(), isaTensor))
             toErase.push_back(op);
 
         return result;
       });
   LDBG("End BufferizeFuncOpInternals:\n" << funcOp << '\n');
 
   for (Operation *op : toErase)
     op->erase();
 
   return failure(result.wasInterrupted());
 }
 
void LinalgComprehensiveModuleBufferize::runOnOperation() {
   ModuleOp moduleOp = getOperation();
 
   SmallVector<FuncOp> orderedFuncOps;
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
   DenseMap<FuncOp, FunctionType> bufferizedFunctionTypes;
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return signalPassFailure();
 
+  GlobalCreator globals(moduleOp);
   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) {
if (failed(inPlaceAnalysisFuncOpBody(funcOp, aliasInfo, domInfo))) {
       signalPassFailure();
       return;
     }
 
     // Bufferization phase.
     if (!testAnalysisOnly) {
       BlockAndValueMapping tensorToBufferMap;
       if (failed(bufferizeFuncOpInternals(funcOp, tensorToBufferMap, aliasInfo,
-                                          bufferizedFunctionTypes))) {
+                                          bufferizedFunctionTypes, &globals))) {
         signalPassFailure();
         return;
       }
     }
   }
   // Don't drop the attributes if we only want to report the analysis.
   if (testAnalysisOnly)
     return;

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

-// RUN: mlir-opt %s -linalg-comprehensive-func-bufferize=test-analysis-only -split-input-file -verify-diagnostics
-
-// -----
+// RUN: mlir-opt %s -linalg-comprehensive-func-bufferize -split-input-file -verify-diagnostics
 
 func @scf_for(%A : tensor<?xf32>,
               %B : tensor<?xf32> {linalg.inplaceable = true},
               %C : tensor<4xf32>,
               %lb : index, %ub : index, %step : index)
   -> (tensor<?xf32>, 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>
 
     // Throw a wrench in the system by swapping yielded values: this result in a
     // ping-pong of values at each iteration on which we currently want to fail.
 
     // expected-error @+1 {{Yield operand #1 does not bufferize to an equivalent buffer}}
     scf.yield %ttB, %ttA : tensor<?xf32>, tensor<?xf32>
   }
 
   return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
 }
 
+// -----
+
+func private @some_external_func(tensor<4xi32>)
+
+func @main() {
+  // expected-error @+1 {{Failed bufferizing tensor constant within a function pass: cannot construct globals creator}}
+  %A = constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+  call @some_external_func(%A) : (tensor<4xi32>) -> ()
+  return
+}

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

 // RUN: mlir-opt %s -linalg-comprehensive-module-bufferize -split-input-file | FileCheck %s
 
+//      CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
+//      CHECK: memref.global "private" constant @__constant_4xi32 : memref<4xi32> = dense<[1, 2, 3, 4]>
+//      CHECK: func private @some_external_func(memref<4xi32, #[[$DYN_1D_MAP]]>)
+func private @some_external_func(tensor<4xi32>)
+
+//      CHECK: func @main()
+func @main() {
+//      CHECK:   %[[A:.*]] = memref.get_global @__constant_4xi32 : memref<4xi32>
+  %A = constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+
+//      CHECK:   %[[B:.*]] = memref.cast %[[A]] : memref<4xi32> to memref<4xi32, #[[$DYN_1D_MAP]]>
+//      CHECK:   call @some_external_func(%[[B]]) : (memref<4xi32, #[[$DYN_1D_MAP]]>) -> ()
+  call @some_external_func(%A) : (tensor<4xi32>) -> ()
+
+  return
+}
+
 // -----
 
 //      CHECK: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
 //      CHECK:  func private @some_external_func(memref<?xf32, #[[$DYN_1D_MAP]]>)
 func private @some_external_func(tensor<?xf32>)
 
 //      CHECK:  func @scf_for_with_tensor_insert_slice(