[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

34

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

34

(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

41

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

34

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

20

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

622

These should be sortsd.

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

54

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

20

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

41–42

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

56

nit: trivial braces

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

17

nit: // ----- missing?

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

1

nit: Some are -, some are --

[rriddle]

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

41–42

It already exists, just pass the op itself.

[herhut]

Comments

[herhut]

Thanks for the reviews!

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

54

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

54

!isKnownTerminator?

[herhut]

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

54

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 branch operations and
+/// `return` to use operands that have been legalized by the conversion
+/// framework. This can only be done if the branch operation implements the
+/// BranchOpInterface. Only needed for partial conversions.
+void populateBranchOpInterfaceAndReturnOpTypeConversionPattern(
+    OwningRewritePatternList &patterns, MLIRContext *ctx,
+    TypeConverter &converter);
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_STANDARDOPS_TRANSFORMS_FUNCCONVERSIONS_H_

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:
-    - block arguments
-    - the result of tensor_load
-    Other values of tensor type should be eliminated by earlier
-    bufferization passes.
+    This pass must be done atomically because it changes func op signatures,
+    which requires atomically updating calls as well throughout the entire
+    module.
+
+    This pass also changes the type of block arguments, which requires that all
+    successor arguments of predecessors be converted. This is achieved by
+    rewriting terminators based on the information provided by the
+    `BranchOpInterface`.
+    As this pass rewrites function operations, it also rewrites the
+    corresponding return operations. Other return-like operations that
+    implement the `ReturnLike` trait are not rewritten in general, as they
+    require that the correspondign parent operation is also rewritten.
+    Finally, this pass fails for unknown terminators, as we cannot decide
+    whether they need rewriting. 
   }];
   let constructor = "mlir::createFuncBufferizePass()";
 }
 
 def TensorConstantBufferize : Pass<"tensor-constant-bufferize", "ModuleOp"> {
   let summary = "Bufferize tensor constants.";
   let description = [{
     This pass bufferizes tensor constants.

mlir/include/mlir/Transforms/Passes.h

 std::unique_ptr<Pass> createBufferLoopHoistingPass();
 
 /// Creates a pass that promotes heap-based allocations to stack-based ones.
 /// Only buffers smaller than the provided size are promoted.
 std::unique_ptr<Pass>
 createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes = 1024,
                                 unsigned bitwidthOfIndexType = 64);
 
+/// Creates a pass that finalizes a partial bufferization by removing remaining
+/// tensor_load and tensor_to_memref operations.
+std::unique_ptr<FunctionPass> createFinalizingBufferizePass();
+
 /// Creates a pass that converts memref function results to out-params.
 std::unique_ptr<Pass> createBufferResultsToOutParamsPass();
 
 /// Creates an instance of the Canonicalizer pass.
 std::unique_ptr<Pass> createCanonicalizerPass();
 
 /// Create a pass that removes unnecessary Copy operations.
 std::unique_ptr<Pass> createCopyRemovalPass();

mlir/include/mlir/Transforms/Passes.td

Option<"disableCanonicalization", "disable-simplify", "bool",
            /*default=*/"false",
            "Disable running simplifications during inlining">,
     Option<"maxInliningIterations", "max-iterations", "unsigned",
            /*default=*/"4",
            "Maximum number of iterations when inlining within an SCC">,
   ];
 }
 
+def FinalizingBufferize : FunctionPass<"finalizing-bufferize"> {
+  let summary = "Finalize a partial bufferization";
+  let description = [{
+    A bufferize pass that finalizes a partial bufferization by removing
+    remaining `tensor_load` and `tensor_to_memref` operations.
+
+    The removal of those operations is only possible if the operations only
+    exist in pairs, i.e., all uses of `tensor_load` operations are
+    `tensor_to_memref` operations.
+
+    This pass will fail if not all operations can be removed or if any operation
+    with tensor typed operands remains.
+  }];
+  let constructor = "mlir::createFinalizingBufferizePass()";
+}
+
 def LocationSnapshot : Pass<"snapshot-op-locations"> {
   let summary = "Generate new locations from the current IR";
   let description = [{
     This pass allows for generating new locations from the IR during any stage
     of compilation, by snapshotting the IR to a file and using that file to
     generate new locations for the operations.
 
     Depending on the value of the `tag` option, different resulting locations
let description = [{
     }
     ```
 
     See [Symbols and SymbolTables](SymbolsAndSymbolTables.md) for more
     information on `Symbols`.
   }];
   let constructor = "mlir::createSymbolDCEPass()";
 }
 #endif // MLIR_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> {
+  using FuncBufferizeBase<FuncBufferizePass>::FuncBufferizeBase;
+
   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);
-    populateEliminateBufferizeMaterializationsPatterns(context, typeConverter,
-                                                       patterns);
-    target.addIllegalOp<TensorLoadOp, TensorToMemrefOp>();
+    target.addDynamicallyLegalOp<CallOp>(
+        [&](CallOp op) { return typeConverter.isLegal(op); });
 
-    // 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.
+    populateBranchOpInterfaceAndReturnOpTypeConversionPattern(patterns, context,
+                                                              typeConverter);
+    target.addLegalOp<ModuleOp, ModuleTerminatorOp, TensorLoadOp,
+                      TensorToMemrefOp>();
+    target.addDynamicallyLegalOp<ReturnOp>(
+        [&](ReturnOp op) { return typeConverter.isLegal(op); });
+    // Mark terminators as legal if they have the ReturnLike trait or
+    // implement the BranchOpInterface and have valid types. If they do not
+    // implement the trait or interface, mark them as illegal no matter what.
     target.markUnknownOpDynamicallyLegal([&](Operation *op) {
-      return typeConverter.isLegal(op->getResultTypes());
+      // If it is not a terminator, ignore it.
+      if (op->isKnownNonTerminator())
+        return true;
+      // If it is not the last operation in the block, also ignore it. We do
+      // this to handle unknown operations, as well.
+      Block *block = op->getBlock();
+      if (!block || &block->back() != op)
+        return true;
+      // ReturnLike operations have to be legalized with their parent. For
+      // return this is handled, for other ops they remain as is.
+      if (op->hasTrait<OpTrait::ReturnLike>())
+        return true;
+      // All successor operands of branch like operations must be rewritten.
+      if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
+        for (int p = 0, e = op->getBlock()->getNumSuccessors(); p < e; ++p) {
+          auto successorOperands = branchOp.getSuccessorOperands(p);
+          if (successorOperands.hasValue() &&
+              !typeConverter.isLegal(successorOperands.getValue().getTypes()))
+            return false;
+        }
+        return true;
+      }
+      return false;
     });
 
     if (failed(applyFullConversion(module, target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createFuncBufferizePass() {
   return std::make_unique<FuncBufferizePass>();
 }

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

 //===- FuncConversions.cpp - Standard Function conversions ----------------===//
 //
 // 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/StandardOps/Transforms/FuncConversions.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #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.
+/// 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 branch operation arguments matches up with the succesor
+/// block arguments.
+class BranchOpInterfaceTypeConversion : public ConversionPattern {
+public:
+  BranchOpInterfaceTypeConversion(TypeConverter &typeConverter,
+                                  MLIRContext *ctx)
+      : ConversionPattern(/*benefit=*/1, typeConverter, MatchAnyOpTypeTag()) {}
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    auto branchOp = dyn_cast<BranchOpInterface>(op);
+    if (!branchOp)
+      return failure();
+
+    // For a branch operation, only some operands go to the target blocks, so
+    // only rewrite those.
+    SmallVector<Value, 4> newOperands(op->operand_begin(), op->operand_end());
+    for (int succIdx = 0, succEnd = op->getBlock()->getNumSuccessors();
+         succIdx < succEnd; ++succIdx) {
+      auto successorOperands = branchOp.getSuccessorOperands(succIdx);
+      if (!successorOperands)
+        continue;
+      for (int idx = successorOperands->getBeginOperandIndex(),
+               eidx = idx + successorOperands->size();
+           idx < eidx; ++idx) {
+        newOperands[idx] = operands[idx];
+      }
+    }
+    rewriter.updateRootInPlace(
+        op, [newOperands, op]() { op->setOperands(newOperands); });
+    return success();
+  }
+};
+} // end anonymous namespace
+
+namespace {
+/// Only needed to support partial conversion of functions where this pattern
+/// ensures that the branch operation arguments matches up with the succesor
+/// block arguments.
+class ReturnOpTypeConversion : public OpConversionPattern<ReturnOp> {
+public:
+  using OpConversionPattern<ReturnOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ReturnOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    // For a return, all operands go to the results of the parent, so
+    // rewrite them all.
+    Operation *operation = op.getOperation();
+    rewriter.updateRootInPlace(
+        op, [operands, operation]() { operation->setOperands(operands); });
+    return success();
+  }
+};
+} // end anonymous namespace
+
+void mlir::populateBranchOpInterfaceAndReturnOpTypeConversionPattern(
+    OwningRewritePatternList &patterns, MLIRContext *ctx,
+    TypeConverter &typeConverter) {
+  patterns.insert<BranchOpInterfaceTypeConversion, ReturnOpTypeConversion>(
+      typeConverter, ctx);
 }

mlir/lib/Transforms/Bufferize.cpp

 //===- Bufferize.cpp - Bufferization utilities ----------------------------===//
 //
 // 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/Transforms/Bufferize.h"
+#include "PassDetail.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // BufferizeTypeConverter
 //===----------------------------------------------------------------------===//
 
+static Value materializeTensorLoad(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(materializeTensorLoad);
+  addSourceMaterialization(materializeTensorLoad);
   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]);
   });
 }
 

 } // namespace
 
 void mlir::populateEliminateBufferizeMaterializationsPatterns(
     MLIRContext *context, BufferizeTypeConverter &typeConverter,
     OwningRewritePatternList &patterns) {
   patterns.insert<BufferizeTensorLoadOp, BufferizeTensorToMemrefOp>(
       typeConverter, context);
 }
+
+namespace {
+struct FinalizingBufferizePass
+    : public FinalizingBufferizeBase<FinalizingBufferizePass> {
+  using FinalizingBufferizeBase<
+      FinalizingBufferizePass>::FinalizingBufferizeBase;
+
+  void runOnFunction() override {
+    auto func = getFunction();
+    auto *context = &getContext();
+
+    BufferizeTypeConverter typeConverter;
+    OwningRewritePatternList patterns;
+    ConversionTarget target(*context);
+
+    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(func, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<FunctionPass> mlir::createFinalizingBufferizePass() {
+  return std::make_unique<FinalizingBufferizePass>();
+}

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

+// RUN: mlir-opt %s -func-bufferize -split-input-file -verify-diagnostics --debug-only=dialect-conversion | FileCheck %s
+
+// CHECK-LABEL:   func @block_arguments(
+// CHECK-SAME:        %[[ARG:.*]]: memref<f32>) -> memref<f32> {
+// CHECK:           %[[T1:.*]] = tensor_load %[[ARG]] : memref<f32>
+// CHECK:           %[[M1:.*]] = tensor_to_memref %[[T1]] : memref<f32>
+// CHECK:           br ^bb1(%[[M1]] : memref<f32>)
+// CHECK:         ^bb1(%[[BBARG:.*]]: memref<f32>):
+// CHECK:           %[[T2:.*]] = tensor_load %[[BBARG]] : memref<f32>
+// CHECK:           %[[M2:.*]] = tensor_to_memref %[[T2]] : memref<f32>
+// CHECK:           return %[[M2]] : memref<f32>
+func @block_arguments(%arg0: tensor<f32>) -> tensor<f32> {
+  br ^bb1(%arg0: tensor<f32>)
+^bb1(%bbarg: tensor<f32>):
+  return %bbarg : tensor<f32>
+}
+
+// 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>
+}
+
+// CHECK-LABEL: func @region_op
+// CHECK-SAME: (%[[ARG0:.*]]: i1) -> memref<f32>
+func @region_op(%arg0: i1) -> tensor<f32> {
+  // CHECK-NEXT: %[[IF:.*]] = scf.if %[[ARG0]] -> (tensor<f32>)
+  %0 = scf.if %arg0 -> (tensor<f32>) {
+    // CHECK-NEXT: %[[SRC:.*]] = "test.source"() : () -> tensor<f32>
+    %1 = "test.source"() : () -> tensor<f32>
+    // CHECK-NEXT: scf.yield %[[SRC]] : tensor<f32>
+    scf.yield %1 : tensor<f32>
+  // CHECK-NEXT: else
+  } else {
+    // CHECK-NEXT: %[[OSRC:.*]] = "test.other_source"() : () -> tensor<f32>
+    %1 = "test.other_source"() : () -> tensor<f32>
+    // CHECK-NEXT: scf.yield %[[OSRC]] : tensor<f32>
+    scf.yield %1 : tensor<f32>
+  }
+  // CHECK: %[[MEM:.*]] = tensor_to_memref %[[IF]] : memref<f32>
+  // CHECK: return %[[MEM]] : memref<f32>
+  return %0 : tensor<f32>
+}
+
+// -----
+
+func @failed_to_legalize(%arg0: tensor<f32>) -> tensor<f32> {
+    %0 = constant true
+    cond_br %0, ^bb1(%arg0: tensor<f32>), ^bb2(%arg0: tensor<f32>)
+  ^bb1(%bbarg0: tensor<f32>):
+    // expected-error @+1 {{failed to legalize operation 'test.terminator'}}
+    "test.terminator"() : () -> ()
+  ^bb2(%bbarg1: tensor<f32>):
+    return %bbarg1 : tensor<f32>
+}

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

-// RUN: mlir-opt %s -func-bufferize -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: mlir-opt %s -func-bufferize -finalizing-bufferize -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL:   func @identity(
 // CHECK-SAME:        %[[ARG:.*]]: memref<f32>) -> memref<f32> {
 // CHECK:           return %[[ARG]] : memref<f32>
 func @identity(%arg0: tensor<f32>) -> tensor<f32> {
   return %arg0 : tensor<f32>
 }