[mlir][bufferize] Add argument materialization for bufferization

[herhut]

Summary

This enables partial bufferization that includes function signatures.

[silvas]

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

36

FuncOpTypeConversionPattern will change the type of block arguments, which requires that all predecessor successor args are converted with a pattern like you have for ReturnOp.

See the rationale here: https://github.com/llvm/llvm-project/blob/25777080549bb62b6e46a1809f93257969f5dd53/mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td#L37

You will need to add a pattern that generically converts all BranchOpInterface ops and set up legality so that if we have a terminator that doesn't implement BranchOpInterface then conversion fails.

Once we do that, I think we can change func-bufferize to be a partial bufferization only, and have a "finalizing-bufferize" that just folds tensor_load/tensor_to_memref.

[silvas]

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

36

(This is actually how I wanted to originally design it (and it is a better design IMO), but I was too lazy / couldnt't figure out how to do the BranchOpInterface rewrite / legality; in retrospect I shouldn't have taken that shortcut)

[silvas]

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td

53

I was never happy with this "just documented" restriction. Using the design I described below, we can improve this situation, because then the finalizing bufferize is just folding tensor_to_memref/tensor_load which we can do while ensuring that no invalid IR is created!

[herhut]

Rework to support terminators that implement interfaces.

[herhut]

PTAL.

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

36

How about this version? If this works for you, I'll fix the documentation and add another test. Just wanted to get your opinion pre-thanksgiving so that I can go ahead with this (and unblock some other work).

We might still want to keep the tests and pass for the finalizing version (to make sure that use case works, as well) as many users might bufferize that way.

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

88

A cleanup would be to implement the interfaces or have a new one for "end of list" terminators. That's more a cleanup to generalize this, though.

[silvas]

I'm okay to submit this if you make func-bufferize only be a partial conversion (never a finalizing conversion), and add a separate finalizing-bufferize pass which folds away tensor_load/tensor_to_memref and emits an error if any remain.

Note that in that new world, mlir-opt -func-bufferize -finalizing-bufferize will reproduce the behavior of the current func-bufferize, so it is trivial for users to update; also, I wrote all the current users of func-bufferize so there is no need to worry about any other users ;)

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

91

This should just return failure() (or rewriter.notifyMatchFailure(op, "terminator does not implement required interfaces"))

The reason is that patterns should not take an opinion on emitting errors in cases like these. The errors should naturally come from how the conversion has the legality set up. For example, if a dialect conversion pass also adds a special pattern for a specific terminator, we don't want to be emitting errors from this pattern.

mlir/lib/Transforms/Bufferize.cpp

18

would prefer to call it materializeTensorLoad and for it to be a local variable in BufferizeTypeConverter::BufferizeTypeConverter just so it is closer to where it is used.

mlir/test/Dialect/Standard/func-bufferize-partial.mlir

2

nit: is --debug-only needed for the test?

[herhut]

Split out finalizing pass.

[rriddle]

mlir/include/mlir/Transforms/Passes.td

606 ↗

(On Diff #307790)

These should be sortsd.

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

67

IsKnownNonTerminator

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

47

Comments should be ///

59

Please use the proper terminator checks, isKnownNonTerminator

[herhut]

Frederik, can you take a look? Sean was happy for this to move forward with his comments addressed, which I did.

mlir/lib/Transforms/Bufferize.cpp

18

It is beyond my C++ voodoo skills to make that happen. If I use a local variable then the compiler complains

ERROR: third_party/llvm/llvm-project/llvm/include/llvm/ADT/STLExtras.h:118:59 type '(lambda at third_party/llvm/llvm-project/mlir/lib/Transforms/Bufferize.cpp:32:32) &' cannot be used prior to '::' because it has no members
struct function_traits : public function_traits<decltype(&T::operator())> {};

Using a static fixes this :)

[frgossen]

Sorry, I saw this only now.
Thank you!

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

62–63

I've seen this a couple of times now. Maybe it is worth moving this to the type converter.

69

nit: trivial braces

mlir/test/Dialect/Standard/func-bufferize-partial.mlir

17

nit: // ----- missing?

mlir/test/Dialect/Standard/func-bufferize.mlir

1 ↗

(On Diff #307790)

nit: Some are -, some are --

[rriddle]

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

62–63

It already exists, just pass the op itself.

[herhut]

Comments

[herhut]

Thanks for the reviews!

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

67

This does not work for unknown operations. I have changed it now to first check isKnownNonTerminator to handle known operations efficiently and then also check for whether it is the last op in the block to handle unknown operations.

mlir/test/Dialect/Standard/func-bufferize-partial.mlir

17

It is not required, only for the last one as it will abort the conversion.

[rriddle]

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

67

!isKnownTerminator?

[herhut]

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

67

If I use !isKnownTerminator then this would declare unknown operations in terminator position as legal. But I want to fail the lowering if there is a block with an unknown terminator.

[herhut]

Do not rewrite ReturnLike ops in general.

[herhut]

I realized that we cannot rewrite all ReturnLike operations, as they require rewriting of their parent which cannot be done generically. PTAL the changes.

[frgossen]

Thanks!

mlir/include/mlir/Dialect/StandardOps/Transforms/FuncConversions.h

 class TypeConverter;
 
 /// Add a pattern to the given pattern list to convert the operand and result
 /// types of a CallOp with the given type converter.
 void populateCallOpTypeConversionPattern(OwningRewritePatternList &patterns,
                                          MLIRContext *ctx,
                                          TypeConverter &converter);
 
+/// Add a pattern to the given pattern list to rewrite the return op to use
+/// operands that have been legalized by the conversion framework. Only needed
+/// for partial conversions.
+void populateReturnOpTypeConversionPattern(OwningRewritePatternList &patterns,
+                                           MLIRContext *ctx,
+                                           TypeConverter &converter);
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_STANDARDOPS_TRANSFORMS_FUNCCONVERSIONS_H_

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.h

 void populateStdBufferizePatterns(MLIRContext *context,
                                   BufferizeTypeConverter &typeConverter,
                                   OwningRewritePatternList &patterns);
 
 /// Creates an instance of std bufferization pass.
 std::unique_ptr<Pass> createStdBufferizePass();
 
 /// Creates an instance of func bufferization pass.
-std::unique_ptr<Pass> createFuncBufferizePass();
+std::unique_ptr<Pass>
+createFuncBufferizePass(bool allowPartialBufferization = false);
 
 /// Creates an instance of tensor constant bufferization pass.
 std::unique_ptr<Pass> createTensorConstantBufferizePass();
 
 /// Creates an instance of the StdExpand pass that legalizes Std
 /// dialect ops to be convertible to LLVM. For example,
 /// `std.ceildivi_signed` gets transformed to a number of std operations,
 /// which can be lowered to LLVM; `memref_reshape` gets converted to

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td

 def StdExpandOps : FunctionPass<"std-expand"> {
   let summary = "Legalize std operations to be convertible to LLVM.";
   let constructor = "mlir::createStdExpandOpsPass()";
 }
 
 def FuncBufferize : Pass<"func-bufferize", "ModuleOp"> {
   let summary = "Bufferize func/call/return ops";
   let description = [{
-    A finalizing bufferize pass that bufferizes std.func and std.call ops.
+    A bufferize pass that bufferizes std.func and std.call ops.
 
     Because this pass updates std.func ops, it must be a module pass. It is
     useful to keep this pass separate from other bufferizations so that the
     other ones can be run at function-level in parallel.
 
     This pass must be done atomically for two reasons:
     1. This pass changes func op signatures, which requires atomically updating
        calls as well throughout the entire module.
     2. This pass changes the type of block arguments, which requires that all
        successor arguments of predecessors be converted. Terminators are not
        a closed universe (and need not implement BranchOpInterface), and so we
        cannot in general rewrite them.
 
-    Note, because this is a "finalizing" bufferize step, it can create
-    invalid IR because it will not create materializations. To avoid this
-    situation, the pass must only be run when the only SSA values of
-    tensor type are:
+    By default this is a "finalizing" bufferize step and it can create
+    invalid IR because it will not create materializations. 
+
+    A finalizing pass must only be run when the only SSA values of tensor type
+    are:
     - block arguments
     - the result of tensor_load
     Other values of tensor type should be eliminated by earlier
     bufferization passes.
+
+    The pass also supports partial bufferization of functions with the
+    limitations as described above. The only supported terminator in such a
+    setting is the return operation. Consequently, the partial bufferization
+    can only be used in the absence of controlflow (blocks and branch
+    operations).
   }];
   let constructor = "mlir::createFuncBufferizePass()";
+  let options = [
+    Option<"allowPartialBufferization", "allow-partial-bufferization",
+           "bool", /*default=*/"false", "Allow partial bufferization. "
+           "Value-based operations may remain.">,
+  ];
 }
 
 def TensorConstantBufferize : Pass<"tensor-constant-bufferize", "ModuleOp"> {
   let summary = "Bufferize tensor constants.";
   let description = [{
     This pass bufferizes tensor constants.
 
     This pass needs to be a module pass because it inserts std.global_memref
     ops into the module, which cannot be done safely from a function pass due to
     multi-threading. Most other bufferization passes can run in parallel at
     function granularity.
   }];
   let constructor = "mlir::createTensorConstantBufferizePass()";
 }
 
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/Transforms/Bufferize.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 namespace {
 struct FuncBufferizePass : public FuncBufferizeBase<FuncBufferizePass> {
+  FuncBufferizePass(bool allowPartialBufferization) {
+    this->allowPartialBufferization = allowPartialBufferization;
+  }
+
   void runOnOperation() override {
     auto module = getOperation();
     auto *context = &getContext();
 
     BufferizeTypeConverter typeConverter;
     OwningRewritePatternList patterns;
     ConversionTarget target(*context);
 
     populateFuncOpTypeConversionPattern(patterns, context, typeConverter);
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
       return typeConverter.isSignatureLegal(op.getType()) &&
              typeConverter.isLegal(&op.getBody());
     });
     populateCallOpTypeConversionPattern(patterns, context, typeConverter);
+    target.addDynamicallyLegalOp<CallOp>([&](CallOp op) {
+      return typeConverter.isLegal(op.getOperandTypes()) &&
+             typeConverter.isLegal(op.getResultTypes());
+    });
+
+    if (allowPartialBufferization) {
+      populateReturnOpTypeConversionPattern(patterns, context, typeConverter);
+      target.addLegalOp<TensorLoadOp, TensorToMemrefOp>();
+      target.addDynamicallyLegalOp<ReturnOp>([&](ReturnOp op) {
+        return typeConverter.isLegal(op.getOperandTypes());
+      });
+
+      if (failed(applyPartialConversion(module, target, std::move(patterns))))
+        signalPassFailure();
+      return;
+    }
+
     populateEliminateBufferizeMaterializationsPatterns(context, typeConverter,
                                                        patterns);
     target.addIllegalOp<TensorLoadOp, TensorToMemrefOp>();
 
     // If all result types are legal, and all block arguments are legal (ensured
     // by func conversion above), then all types in the program are legal.
     target.markUnknownOpDynamicallyLegal([&](Operation *op) {
       return typeConverter.isLegal(op->getResultTypes());
     });
 
     if (failed(applyFullConversion(module, target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
-std::unique_ptr<Pass> mlir::createFuncBufferizePass() {
-  return std::make_unique<FuncBufferizePass>();
+std::unique_ptr<Pass>
+mlir::createFuncBufferizePass(bool allowPartialBufferization) {
+  return std::make_unique<FuncBufferizePass>(allowPartialBufferization);
 }

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 namespace {
 // Converts the operand and result types of the Standard's CallOp, used together
 // with the FuncOpSignatureConversion.
 struct CallOpSignatureConversion : public OpConversionPattern<CallOp> {
-  CallOpSignatureConversion(MLIRContext *ctx, TypeConverter &converter)
-      : OpConversionPattern(ctx), converter(converter) {}
+  using OpConversionPattern<CallOp>::OpConversionPattern;
 
   /// Hook for derived classes to implement combined matching and rewriting.
   LogicalResult
   matchAndRewrite(CallOp callOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    FunctionType type = callOp.getCalleeType();
-
     // Convert the original function results.
     SmallVector<Type, 1> convertedResults;
-    if (failed(converter.convertTypes(type.getResults(), convertedResults)))
+    if (failed(typeConverter->convertTypes(callOp.getResultTypes(),
+                                           convertedResults)))
       return failure();
 
     // Substitute with the new result types from the corresponding FuncType
     // conversion.
     rewriter.replaceOpWithNewOp<CallOp>(callOp, callOp.callee(),
                                         convertedResults, operands);
     return success();
   }
-
-  /// The type converter to use when rewriting the signature.
-  TypeConverter &converter;
 };
 } // end anonymous namespace
 
 void mlir::populateCallOpTypeConversionPattern(
     OwningRewritePatternList &patterns, MLIRContext *ctx,
     TypeConverter &converter) {
-  patterns.insert<CallOpSignatureConversion>(ctx, converter);
+  patterns.insert<CallOpSignatureConversion>(converter, ctx);
+}
+
+namespace {
+// Only needed to support partial conversion of functions where this pattern
+// ensures that the return matches up with the new function signature.
+class ReturnOpTypeConversion : public OpConversionPattern<ReturnOp> {
+public:
+  using OpConversionPattern<ReturnOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ReturnOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    rewriter.replaceOpWithNewOp<ReturnOp>(op, operands);
+    return success();
+  }
+};
+} // end anonymous namespace
+
+void mlir::populateReturnOpTypeConversionPattern(
+    OwningRewritePatternList &patterns, MLIRContext *ctx,
+    TypeConverter &typeConverter) {
+  patterns.insert<ReturnOpTypeConversion>(typeConverter, ctx);
 }

mlir/lib/Transforms/Bufferize.cpp

 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // BufferizeTypeConverter
 //===----------------------------------------------------------------------===//
 
+static Value InsertTensorLoad(OpBuilder &builder, TensorType type,
+                              ValueRange inputs, Location loc) {
+  assert(inputs.size() == 1);
+  assert(inputs[0].getType().isa<BaseMemRefType>());
+  return builder.create<TensorLoadOp>(loc, type, inputs[0]);
+}
+
 /// Registers conversions into BufferizeTypeConverter
 BufferizeTypeConverter::BufferizeTypeConverter() {
   // Keep all types unchanged.
   addConversion([](Type type) { return type; });
   // Convert RankedTensorType to MemRefType.
   addConversion([](RankedTensorType type) -> Type {
     return MemRefType::get(type.getShape(), type.getElementType());
   });
   // Convert UnrankedTensorType to UnrankedMemRefType.
   addConversion([](UnrankedTensorType type) -> Type {
     return UnrankedMemRefType::get(type.getElementType(), 0);
   });
-  addSourceMaterialization([](OpBuilder &builder, TensorType type,
-                              ValueRange inputs, Location loc) -> Value {
-    assert(inputs.size() == 1);
-    assert(inputs[0].getType().isa<BaseMemRefType>());
-    return builder.create<TensorLoadOp>(loc, type, inputs[0]);
-  });
+  addArgumentMaterialization(InsertTensorLoad);
+  addSourceMaterialization(InsertTensorLoad);
   addTargetMaterialization([](OpBuilder &builder, BaseMemRefType type,
                               ValueRange inputs, Location loc) -> Value {
     assert(inputs.size() == 1);
     assert(inputs[0].getType().isa<TensorType>());
     return builder.create<TensorToMemrefOp>(loc, type, inputs[0]);
   });
 }
 

mlir/test/Dialect/Standard/func-bufferize-partial.mlir

+// RUN: mlir-opt %s -func-bufferize="allow-partial-bufferization=true" -split-input-file -verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL: func @partial()
+// CHECK-SAME: memref<f32>
+func @partial() -> tensor<f32> {
+  // CHECK-NEXT: %[[SRC:.*]] = "test.source"() : () -> tensor<f32>
+  // CHECK-NEXT: %[[MEM:.*]] = tensor_to_memref %[[SRC]] : memref<f32>
+  %0 = "test.source"() : () -> tensor<f32>
+  // CHECK-NEXT: return %[[MEM]] : memref<f32>
+  return %0 : tensor<f32>
+}