[mlir] Add func-bufferize pass.

[silvas]

Summary

This is the most basic possible finalizing bufferization pass, which I
also think is sufficient for most new use cases. The more concentrated
nature of this pass also greatly clarifies the invariants that it
requires on its input to safely transform the program (see the
pass description in Passes.td).

With this pass, I have now upstreamed practically all of the
bufferizations from npcomp (the exception being std.constant, which can
be upstreamed when std.global_memref lands:
https://llvm.discourse.group/t/rfc-global-variables-in-mlir/2076/16 )

[silvas]

rebase

[nicolasvasilache]

Loos good, thanks @silvas !

[silvas]

@dfki-mako , @herhut , I'd like for one of you to approve this one.

The overall picture is the users will be able to run pipelines like -scf-bufferize -std-bufferize -linalg-bufferize -func-bufferize. The func-bufferize will terminate the pass pipeline and "finalize" the conversion. That will then complete the picture for composable bufferizations.

With the changes discussed in https://llvm.discourse.group/t/rfc-simplifying-bufferizetypeconverter/2091/2 we can then really trim down the current TestBufferPlacementPreparation pass (possibly deleting it depending on how much we can simplify it).

[dfki-mako]

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

28

Question: std.func, std.call and std.return ops? The actual code inside the pass doesn't adapt any return ops, does it?

38

missing to -> to be converted?

42

Suggestion: bufferize -> bufferize step?

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

2

Bufferize.cpp -> FuncBufferize.cpp

[silvas]

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

28

Dialect conversion patterns are mostly about replacing results. Operand types change naturally as their definitions change type. So std.return is not explicitly handled or mentioned here, because it doesn't have results.

In the case of "non-finalizing" bufferize conversions, terminators like scf.yield need patterns so that they create materializations to the target type. But that is not here.

We could write this pass as a "non-finalizing" (also known as "composable") bufferize pass, in which case we would need it. However, as I describe below, because block arguments are rewritten, it is natural for this pass to be finalizing so that successor operands of terminators are all handled naturally. For example, we don't need patterns for std.br, std.cond_br or any other terminator that passes values to successors. If this was a "composable" bufferize pass, then we would need patterns for *every* possible terminator (and not all terminators implement BranchOpInterface for us to guarantee we can rewrite them generically).

38

I think it is right as-is. "requires that X be converted", for X = "all successor arguments of predecessors". (I think you would be correct if the "that" was removed).

[silvas]

Address comments

[pifon2a]

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

6

nit: remove CHECK for '}'. Here and below.

24

nit: reduce indentation to fit 80 chars? here and below

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();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_

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

 }
 
 def StdBufferize : FunctionPass<"std-bufferize"> {
   let summary = "Bufferize the std dialect";
   let constructor = "mlir::createStdBufferizePass()";
   let dependentDialects = ["scf::SCFDialect"];
 }
 
+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.
+
+    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.
+  }];
+  let constructor = "mlir::createFuncBufferizePass()";
+}
+
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES

mlir/include/mlir/Transforms/Bufferize.h

 };
 
 /// Marks ops used by bufferization for type conversion materializations as
 /// "legal" in the given ConversionTarget.
 ///
 /// This function should be called by all bufferization passes using
 /// BufferizeTypeConverter so that materializations work proprely. One exception
 /// is bufferization passes doing "full" conversions, where it can be desirable
-/// for even the materializations to remain illegal so that they are eliminated.
+/// for even the materializations to remain illegal so that they are eliminated,
+/// such as via the patterns in
+/// populateEliminateBufferizeMaterializationsPatterns.
 void populateBufferizeMaterializationLegality(ConversionTarget &target);
 
+/// Populate patterns to eliminate bufferize materializations.
+///
+/// In particular, these are the tensor_load/tensor_to_memref ops.
+void populateEliminateBufferizeMaterializationsPatterns(
+    MLIRContext *context, BufferizeTypeConverter &typeConverter,
+    OwningRewritePatternList &patterns);
+
 /// Helper conversion pattern that encapsulates a BufferizeTypeConverter
 /// instance.
 template <typename SourceOp>
 class BufferizeOpConversionPattern : public OpConversionPattern<SourceOp> {
 public:
   explicit BufferizeOpConversionPattern(MLIRContext *context,
                                         BufferizeTypeConverter &converter,
                                         PatternBenefit benefit = 1)

mlir/lib/Dialect/StandardOps/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRStandardOpsTransforms
   Bufferize.cpp
   ExpandAtomic.cpp
   ExpandMemRefReshape.cpp
   ExpandTanh.cpp
+  FuncBufferize.cpp
   FuncConversions.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps/Transforms
 
   DEPENDS
   MLIRStandardTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
   MLIRSCF
   MLIRStandard
   MLIRTransforms
   )

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

+//===- Bufferize.cpp - Bufferization for std ops --------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements bufferization of std.func's and std.call's.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/StandardOps/Transforms/FuncConversions.h"
+#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> {
+  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>();
+
+    // 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>();
+}

mlir/lib/Transforms/Bufferize.cpp

if (auto res = conversion(origin, converted))
       return res.getValue();
   return KeepAsFunctionResult;
 }
 
 void mlir::populateBufferizeMaterializationLegality(ConversionTarget &target) {
   target.addLegalOp<TensorLoadOp, TensorToMemrefOp>();
 }
 
+namespace {
+// In a finalizing bufferize conversion, we know that all tensors have been
+// converted to memrefs, thus, this op becomes an identity.
+class BufferizeTensorLoadOp : public OpConversionPattern<TensorLoadOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(TensorLoadOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    TensorLoadOp::Adaptor adaptor(operands);
+    rewriter.replaceOp(op, adaptor.memref());
+    return success();
+  }
+};
+} // namespace
+
+namespace {
+// In a finalizing bufferize conversion, we know that all tensors have been
+// converted to memrefs, thus, this op becomes an identity.
+class BufferizeTensorToMemrefOp : public OpConversionPattern<TensorToMemrefOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(TensorToMemrefOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    TensorToMemrefOp::Adaptor adaptor(operands);
+    rewriter.replaceOp(op, adaptor.tensor());
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateEliminateBufferizeMaterializationsPatterns(
+    MLIRContext *context, BufferizeTypeConverter &typeConverter,
+    OwningRewritePatternList &patterns) {
+  patterns.insert<BufferizeTensorLoadOp, BufferizeTensorToMemrefOp>(
+      typeConverter, context);
+}
+
 //===----------------------------------------------------------------------===//
 // BufferizeFuncOpConverter
 //===----------------------------------------------------------------------===//
 
 /// Performs the actual function signature rewriting step.
 LogicalResult BufferizeFuncOpConverter::matchAndRewrite(
     mlir::FuncOp funcOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {

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

+// RUN: mlir-opt %s -func-bufferize -split-input-file -verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL:   func @identity(
+// CHECK-SAME:                   %[[ARG:.*]]: memref<f32>) -> memref<f32> {
+// CHECK:           return %[[ARG]] : memref<f32>
+// CHECK:         }
+func @identity(%arg0: tensor<f32>) -> tensor<f32> {
+  return %arg0 : tensor<f32>
+}
+
+// CHECK-LABEL:   func @block_arguments(
+// CHECK-SAME:                          %[[ARG:.*]]: memref<f32>) -> memref<f32> {
+// CHECK:           br ^bb1(%[[ARG]] : memref<f32>)
+// CHECK:         ^bb1(%[[BBARG:.*]]: memref<f32>):
+// CHECK:           return %[[BBARG]] : memref<f32>
+// CHECK:         }
+func @block_arguments(%arg0: tensor<f32>) -> tensor<f32> {
+  br ^bb1(%arg0: tensor<f32>)
+^bb1(%bbarg: tensor<f32>):
+  return %bbarg : tensor<f32>
+}
+
+// CHECK-LABEL:   func @eliminate_target_materialization(
+// CHECK-SAME:                                           %[[ARG:.*]]: memref<f32>) -> memref<f32> {
+// CHECK:           return %[[ARG]] : memref<f32>
+// CHECK:         }
+func @eliminate_target_materialization(%arg0: tensor<f32>) -> memref<f32> {
+  %0 = tensor_to_memref %arg0 : memref<f32>
+  return %0 : memref<f32>
+}
+
+// CHECK-LABEL:   func @eliminate_source_materialization(
+// CHECK-SAME:                                           %[[ARG:.*]]: memref<f32>) -> memref<f32> {
+// CHECK:           return %[[ARG]] : memref<f32>
+// CHECK:         }
+func @eliminate_source_materialization(%arg0: memref<f32>) -> tensor<f32> {
+  %0 = tensor_load %arg0 : memref<f32>
+  return %0 : tensor<f32>
+}
+
+// CHECK-LABEL:   func @source() -> memref<f32>
+// CHECK-LABEL:   func @call_source() -> memref<f32> {
+// CHECK:           %[[RET:.*]] = call @source() : () -> memref<f32>
+// CHECK:           return %[[RET]] : memref<f32>
+// CHECK:         }
+func @source() -> tensor<f32>
+func @call_source() -> tensor<f32> {
+  %0 = call @source() : () -> tensor<f32>
+  return %0 : tensor<f32>
+}
+
+// CHECK-LABEL:   func @sink(memref<f32>)
+// CHECK-LABEL:   func @call_sink(
+// CHECK-SAME:                    %[[ARG:.*]]: memref<f32>) {
+// CHECK:           call @sink(%[[ARG]]) : (memref<f32>) -> ()
+// CHECK:           return
+// CHECK:         }
+func @sink(tensor<f32>)
+func @call_sink(%arg0: tensor<f32>) {
+  call @sink(%arg0) : (tensor<f32>) -> ()
+  return
+}
+
+// -----
+
+func @failed_to_legalize() -> tensor<f32> {
+  // expected-error @+1 {{failed to legalize operation 'test.source'}}
+  %0 = "test.source"() : () -> (tensor<f32>)
+  return %0 : tensor<f32>
+}