[mlir] Add BufferResultsToOutParams pass.

[silvas]

Summary

This pass allows removing getResultConversionKind from
BufferizeTypeConverter. This pass replaces the AppendToArgumentsList
functionality. As far as I could tell, the only use of this functionlity
is to perform the transformation that is implemented in this pass.

Future patches will remove the getResultConversionKind machinery from
BufferizeTypeConverter, but sending this patch for individual review for
clarity.

[dfki-mako]

mlir/include/mlir/Transforms/Passes.td

220

Just for my understanding: This pass basically simplifies the use of populateWithBufferizeOpConversionPatterns via a custom Rewriter instance for an input dialect, right? Do you plan to remove the conversion patterns in favor of this pass in the future?

My question is here: Why should this pass be executed after the entire bufferization process? Isn't it useful to apply this pass before applying the BufferDeallocation pass, which inserts all the necessary deallocations? Or do you refer to the first part of the "overall" bufferization process, which ports the input program into the "buffer world" (without taking deallocations into account)?

Regarding the Linalg dialect dependency: We are going to discuss possible solutions to remove the Std- and Linalg dialect dependencies from the BufferDeallocation pass in the public Discourse-Forums these days. I suspect that the result of this discussion can also be applied to this pass (in the future).

[silvas]

mlir/include/mlir/Transforms/Passes.td

220

This pass eliminates the need for the AppendToArgumentList conversion kind, which allows removing that complexity from BufferizeTypeConverter. (and together with removing the tuple stuff, BufferizeTypeConverter can just be a regular type converter; everything gets quite a bit simpler).

Regarding BufferDeallocation, I think we agree but just got confused by terminology. "bufferization" (as defined in Bufferize.h and being used here) is about applying conversions. So this would be applied right after applying the conversions, but before any buffer optimizations.

[dfki-mako]

I guess this approach introduces a limitation with respect to applicability to arbitrary dialects. We are now able to match ReturnOp and CallOps only while the previous version supported a generic customization to arbitrary dialects, right?
I would recommend that we should try to match CallOpInterface and ReturnLike implementations to make this code more generic in favor of standard CallOp and ReturnOp-based matchers from the Std dialect.

mlir/include/mlir/Transforms/Passes.td

220

AppendToArgumentList conversion kind, which allows removing that complexity from BufferizeTypeConverter

Sounds reasonable, since it simplifies the type converter to become "just a type converter" (as you said).

Regarding BufferDeallocation, I think we agree but just got confused by terminology. "bufferization"

Sounds pretty good to me, too. In this context, this pass should be the last one to be executed (as you mentioned) to "fix" all return ops and call sites (at least for statically shaped memrefs).

mlir/lib/Transforms/BufferResultsToOutParams.cpp

37

functionType (above) vs func.getType() (2 uses)

43

I am note sure about the braces here since it is a single statement. However, they improve readability in this case.

52

Suggestion: maybe we could change the function to receive an additional vector by reference and push the results into this vector instead of returning an empty vector in this case.

58

nit: maybe a tiny comment would be nice in this case
something like: // Updates all ReturnOps in the scope of the given FuncOp by either keeping them as return values or copying the associated buffer contents into the given out params.

70

Suggestion: maybe change t to outParamEntry. Regarding the braces: same comment as above.

79

nit: maybe another tiny comment would be nice in this case, as well
something like: // Updates all CallOps in the scope of the given ModuleOp by allocating temporary buffers for newly introduced out params

93

Maybe I am missing the point but can't we just check this in the first loop to avoid this additional cast?

111

nit: maybe change t to something more meaningful.

126

Suggestion: If we change the signature of the updateFuncOp function, we can avoid this variable. (see above)

mlir/test/Transforms/buffer-results-to-out-params.mlir

97

As you are updating the tests anyway, it would be really nice to break after 80 chars.

[silvas]

Address @dfki-mako comments.

[silvas]

Thanks Marcel for the very thorough review!

In D90071#2358875, @dfki-mako wrote:

I guess this approach introduces a limitation with respect to applicability to arbitrary dialects. We are now able to match ReturnOp and CallOps only while the previous version supported a generic customization to arbitrary dialects, right?
I would recommend that we should try to match CallOpInterface and ReturnLike implementations to make this code more generic in favor of standard CallOp and ReturnOp-based matchers from the Std dialect.

Good point! I really would like to write this pass as being generic as well. I tried to write this pass against CallOpInterface and ReturnLike and FunctionLike, but they currently don't have enough functionality for this transformation, in particular the parts where we need to mutate/recreate the ops.

For example, FunctionLike doesn't guarantee that the op has a FunctionType, so all the code here where we modify the function type can't be written generically (see this thread for a similar issue:
https://llvm.discourse.group/t/moving-erasearguments-from-funcop-to-functionlike/2016/14). In particular, even after the patches in that thread, we cannot append new result types. (also, I don't think that traits like FunctionLike / ReturnLike.

Also, note that FunctionLike is a trait, not an interface, so it cannot be accessed dynamically. One must already have an object of the concrete type to use it.

mlir/lib/Transforms/BufferResultsToOutParams.cpp

93

I conceptually associate this error with the creation of the AllocOp -- we don't care about this limitation elsewhere in the entire pass, and one can imagine ways to avoid this, which would involve changing this code, rather than the code above (which merely determines which results we need to convert to an out param -- not how we do it).

For example, we could require annotating functions with an attribute that somehow describes how to compute their out param sizes e.g. func @callee(...) -> (...) attributes {shape_transfer_function = @compute_callee_out_param_shapes}, or do so via some analysis, etc. In that circumstance, the static shape case is but one of multiple ways that we would know how to allocate the results.

Also this pass only runs once in the pipeline, so performance isn't that critical. For example, we are walking every op in the module just to find call ops, and we effectively do the same for return. If we really cared about performance, we would only do a single walk of the whole module, but that would make the code a lot more confusing.

111

Same comment as above, I think that succinctness here makes the code clearer, as the minimum meaningful name is just too long and distracting.

It's too bad that C++ doesn't let us do something like this:

for (auto [replaceMe, newValue] : llvm::zip(replaceWithNewCallResults, newCall.getResults()))
  replaceMe.replaceAllUsesWith(newValue)

I think that C++17 has a feature for this, but I don't know if we allow using it yet: https://en.cppreference.com/w/cpp/language/structured_binding

mlir/test/Transforms/buffer-results-to-out-params.mlir

97

I typically don't break in test cases so that it reflects the typical printed IR structure. The theory is that our brains are mostly tuned to parse non-line-broken IR since that is what we usually read (unlike in C++), so line-breaking it tends to be more confusing since there aren't clear style preferences and mental heuristics for parsing it.

[silvas]

Also, to answer your specific question, the previous version did not support a generic customization to arbitrary dialects. The FuncOp / CallOp patterns in the previous code are hard-coded (which as I described, cannot be removed with the current set of traits/interfaces). Only BufferizeReturnOpConverter converter is templatized, but without being able to convert generic function-like ops, that doesn't seem to be very useful (since converting the function-like ops is the hardest part).

[dfki-mako]

For example, FunctionLike doesn't guarantee that the op has a FunctionType, so all the code here where we modify the function type can't be written generically (see this thread for a similar issue:

Unfortunately, we cannot use this trait for the reasons you mentioned.

Also, to answer your specific question, the previous version did not support a generic customization to arbitrary dialects.

Yeah, you are absolutely right. Consequently, we will not lose any genericity by merging your CL, which is really great +1
Future extensions in the MLIR world might enable us to make this pass more generic with respect to FuncOps.

[dfki-mako]

mlir/lib/Transforms/BufferResultsToOutParams.cpp

93

Sounds reasonable. I guess we can safely keep it this way +1.

111

Yes, unfortunately that is also true. As far as I know, we shouldn't use this extension at the moment.

mlir/test/Transforms/buffer-results-to-out-params.mlir

97

+1

[dfki-mako]

Awesome work +1. I guess this should be ready to go after replacing some autos with their underlying types.

mlir/lib/Transforms/BufferResultsToOutParams.cpp

66

nit: auto -> OpOperand

88

nit: auto -> OpResult

96

nit: auto -> Value

[silvas]

address comments

mlir/include/mlir/Transforms/Passes.h

 /// Creates a pass that moves allocations upwards out of loops. This avoids
 /// reallocations inside of loops.
 std::unique_ptr<Pass> createBufferLoopHoistingPass();
 
 /// Creates a pass that promotes heap-based allocations to stack-based ones.
 std::unique_ptr<Pass>
 createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes = 1024);
 
+/// 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();
 
 /// Creates a pass to perform common sub expression elimination.
 std::unique_ptr<Pass> createCSEPass();

mlir/include/mlir/Transforms/Passes.td

def PromoteBuffersToStack : FunctionPass<"promote-buffers-to-stack"> {
   let constructor = "mlir::createPromoteBuffersToStackPass()";
   let options = [
     Option<"maxAllocSizeInBytes", "max-alloc-size-in-bytes", "unsigned",
            /*default=*/"1024",
            "Define the maximum size in bytes to promote allocations to stack.">,
   ];
 }
 
+def BufferResultsToOutParams : Pass<"buffer-results-to-out-params", "ModuleOp">  {
+  let summary = "Converts memref-typed function results to out-params";
+  let description = [{
+    Some calling conventions prefer to pass output memrefs as "out params". The
+    conversion to this calling convention must be done as an atomic
+    transformation of the entire program (hence this is a module pass).
+
+    For example, if a call is rewritten, the callee needs to be rewritten
+    otherwise the IR will end up invalid. Thus, this transformation
+    require an atomic change to the entire program (e.g. the whole module).
+
+    This pass is expected to run immediately after bufferization is finished.
+    At that point, tensor-typed results will have been converted to memref-typed
+    results, and can be consistently converted to out params.
+
+    All memref-typed results are appended to the function argument list.
+
+    The main issue with this pass (and the out-param calling convention) is that
+    buffers for results need to be allocated in the caller. This currently only
+    works for static shaped memrefs.
+  }];
+  let constructor = "mlir::createBufferResultsToOutParamsPass()";
+  let dependentDialects = ["linalg::LinalgDialect"];
+}
+
 def Canonicalizer : Pass<"canonicalize"> {
   let summary = "Canonicalize operations";
   let description = [{
     This pass performs various types of canonicalizations over a set of
     operations. See [Operation Canonicalization](Canonicalization.md) for more
     details.
   }];
   let constructor = "mlir::createCanonicalizerPass()";

mlir/lib/Transforms/BufferResultsToOutParams.cpp

+//===- BufferResultsToOutParams.cpp - Calling convention conversion -------===//
+//
+// 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 "PassDetail.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+
+using namespace mlir;
+
+// Update the func op and entry block.
+//
+// Returns the new appended entry block arguments.
+// If the function is external, an empty vector is returned.
+static SmallVector<BlockArgument, 6> updateFuncOp(FuncOp func) {
+  auto functionType = func.getType();
+
+  // Collect information about the results will become appended arguments.
+  SmallVector<Type, 6> erasedResultTypes;
+  SmallVector<unsigned, 6> erasedResultIndices;
+  for (auto resultType : llvm::enumerate(functionType.getResults())) {
+    if (resultType.value().isa<BaseMemRefType>()) {
+      erasedResultIndices.push_back(resultType.index());
+      erasedResultTypes.push_back(resultType.value());
+    }
+  }
+
+  // Add the new arguments to the function type.
+  auto newArgTypes = llvm::to_vector<6>(
+      llvm::concat<const Type>(func.getType().getInputs(), erasedResultTypes));
+  auto newFunctionType = FunctionType::get(
+      newArgTypes, func.getType().getResults(), func.getContext());
+  func.setType(newFunctionType);
+
+  // Transfer the result attributes to arg attributes.
+  for (int i = 0, e = erasedResultTypes.size(); i < e; i++) {
+    func.setArgAttrs(functionType.getNumInputs() + i,
+                     func.getResultAttrs(erasedResultIndices[i]));
+  }
+
+  // Erase the results.
+  func.eraseResults(erasedResultIndices);
+
+  // Add the new arguments to the entry block if the function is not external.
+  if (func.isExternal())
+    return {};
+  else
+    return llvm::to_vector<6>(func.front().addArguments(erasedResultTypes));
+}
+
+static void updateReturnOps(FuncOp func,
+                            ArrayRef<BlockArgument> appendedEntryArgs) {
+  func.walk([&](ReturnOp op) {
+    SmallVector<Value, 6> copyIntoOutParams;
+    SmallVector<Value, 6> keepAsReturnOperands;
+    for (auto operand : op.getOperands()) {
+      if (operand.getType().isa<BaseMemRefType>())
+        copyIntoOutParams.push_back(operand);
+      else
+        keepAsReturnOperands.push_back(operand);
+    }
+    OpBuilder builder(op);
+    for (auto t : llvm::zip(copyIntoOutParams, appendedEntryArgs)) {
+      builder.create<linalg::CopyOp>(op.getLoc(), std::get<0>(t),
+                                     std::get<1>(t));
+    }
+    builder.create<ReturnOp>(op.getLoc(), keepAsReturnOperands);
+    op.erase();
+  });
+}
+
+static LogicalResult updateCalls(ModuleOp module) {
+  bool didFail = false;
+  module.walk([&](CallOp op) {
+    SmallVector<Value, 6> replaceWithNewCallResults;
+    SmallVector<Value, 6> replaceWithOutParams;
+    for (auto result : op.getResults()) {
+      if (result.getType().isa<BaseMemRefType>())
+        replaceWithOutParams.push_back(result);
+      else
+        replaceWithNewCallResults.push_back(result);
+    }
+    SmallVector<Value, 6> outParams;
+    OpBuilder builder(op);
+    for (auto memref : replaceWithOutParams) {
+      if (!memref.getType().cast<BaseMemRefType>().hasStaticShape()) {
+        op.emitError()
+            << "cannot create out param for dynamically shaped result";
+        didFail = true;
+        return;
+      }
+      Value outParam = builder.create<AllocOp>(
+          op.getLoc(), memref.getType().cast<MemRefType>());
+      memref.replaceAllUsesWith(outParam);
+      outParams.push_back(outParam);
+    }
+
+    auto newOperands = llvm::to_vector<6>(op.getOperands());
+    newOperands.append(outParams.begin(), outParams.end());
+    auto newResultTypes = llvm::to_vector<6>(llvm::map_range(
+        replaceWithNewCallResults, [](Value v) { return v.getType(); }));
+    auto newCall = builder.create<CallOp>(op.getLoc(), op.calleeAttr(),
+                                          newResultTypes, newOperands);
+    for (auto t : llvm::zip(replaceWithNewCallResults, newCall.getResults()))
+      std::get<0>(t).replaceAllUsesWith(std::get<1>(t));
+    op.erase();
+  });
+
+  return failure(didFail);
+}
+
+namespace {
+struct BufferResultsToOutParamsPass
+    : BufferResultsToOutParamsBase<BufferResultsToOutParamsPass> {
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+
+    for (auto func : module.getOps<FuncOp>()) {
+      auto appendedEntryArgs = updateFuncOp(func);
+      if (func.isExternal())
+        continue;
+      updateReturnOps(func, appendedEntryArgs);
+    }
+    if (failed(updateCalls(module)))
+      return signalPassFailure();
+  }
+};
+} // end anonymous namespace
+
+std::unique_ptr<Pass> mlir::createBufferResultsToOutParamsPass() {
+  return std::make_unique<BufferResultsToOutParamsPass>();
+}

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   BufferDeallocation.cpp
   BufferOptimizations.cpp
+  BufferResultsToOutParams.cpp
   Bufferize.cpp
   Canonicalizer.cpp
   CopyRemoval.cpp
   CSE.cpp
   DialectConversion.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopCoalescing.cpp

mlir/test/Transforms/buffer-results-to-out-params.mlir

+// RUN: mlir-opt -buffer-results-to-out-params -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// CHECK-LABEL:   func @basic(
+// CHECK-SAME:                %[[ARG:.*]]: memref<f32>) {
+// CHECK:           %[[RESULT:.*]] = "test.source"() : () -> memref<f32>
+// CHECK:           linalg.copy(%[[RESULT]], %[[ARG]]) : memref<f32>, memref<f32>
+// CHECK:           return
+// CHECK:         }
+func @basic() -> (memref<f32>) {
+  %0 = "test.source"() : () -> (memref<f32>)
+  return %0 : memref<f32>
+}
+
+// CHECK-LABEL:   func @presence_of_existing_arguments(
+// CHECK-SAME:                                         %[[ARG0:.*]]: memref<1xf32>,
+// CHECK-SAME:                                         %[[ARG1:.*]]: memref<2xf32>) {
+// CHECK:           %[[RESULT:.*]] = "test.source"() : () -> memref<2xf32>
+// CHECK:           linalg.copy(%[[RESULT]], %[[ARG1]]) : memref<2xf32>, memref<2xf32>
+// CHECK:           return
+// CHECK:         }
+func @presence_of_existing_arguments(%arg0: memref<1xf32>) -> (memref<2xf32>) {
+  %0 = "test.source"() : () -> (memref<2xf32>)
+  return %0 : memref<2xf32>
+}
+
+// CHECK-LABEL:   func @multiple_results(
+// CHECK-SAME:                           %[[ARG0:.*]]: memref<1xf32>,
+// CHECK-SAME:                           %[[ARG1:.*]]: memref<2xf32>) {
+// CHECK:           %[[RESULTS:.*]]:2 = "test.source"() : () -> (memref<1xf32>, memref<2xf32>)
+// CHECK:           linalg.copy(%[[RESULTS]]#0, %[[ARG0]]) : memref<1xf32>, memref<1xf32>
+// CHECK:           linalg.copy(%[[RESULTS]]#1, %[[ARG1]]) : memref<2xf32>, memref<2xf32>
+// CHECK:           return
+// CHECK:         }
+func @multiple_results() -> (memref<1xf32>, memref<2xf32>) {
+  %0, %1 = "test.source"() : () -> (memref<1xf32>, memref<2xf32>)
+  return %0, %1 : memref<1xf32>, memref<2xf32>
+}
+
+// CHECK-LABEL:   func @non_memref_types(
+// CHECK-SAME:                           %[[OUTPARAM:.*]]: memref<f32>) -> (i1, i32) {
+// CHECK:           %[[RESULT1:.*]]:3 = "test.source"() : () -> (i1, memref<f32>, i32)
+// CHECK:           linalg.copy(%[[RESULT1]]#1, %[[OUTPARAM]]) : memref<f32>, memref<f32>
+// CHECK:           return %[[RESULT1]]#0, %[[RESULT1]]#2 : i1, i32
+// CHECK:         }
+func @non_memref_types() -> (i1, memref<f32>, i32) {
+  %0, %1, %2 = "test.source"() : () -> (i1, memref<f32>, i32)
+  return %0, %1, %2 : i1, memref<f32>, i32
+}
+
+// CHECK: func @external_function(memref<f32>)
+func @external_function() -> (memref<f32>)
+// CHECK: func @result_attrs(memref<f32> {test.some_attr})
+func @result_attrs() -> (memref<f32> {test.some_attr})
+// CHECK: func @mixed_result_attrs(memref<1xf32>, memref<2xf32> {test.some_attr}, memref<3xf32>)
+func @mixed_result_attrs() -> (memref<1xf32>, memref<2xf32> {test.some_attr}, memref<3xf32>)
+
+// -----
+
+// CHECK-LABEL: func @callee(memref<1xf32>)
+func @callee() -> memref<1xf32>
+
+// CHECK-LABEL:   func @call_basic() {
+// CHECK:           %[[OUTPARAM:.*]] = alloc() : memref<1xf32>
+// CHECK:           call @callee(%[[OUTPARAM]]) : (memref<1xf32>) -> ()
+// CHECK:           "test.sink"(%[[OUTPARAM]]) : (memref<1xf32>) -> ()
+// CHECK:           return
+// CHECK:         }
+func @call_basic() {
+  %0 = call @callee() : () -> memref<1xf32>
+  "test.sink"(%0) : (memref<1xf32>) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @callee(memref<1xf32>, memref<2xf32>)
+func @callee() -> (memref<1xf32>, memref<2xf32>)
+
+// CHECK-LABEL:   func @call_multiple_result() {
+// CHECK:           %[[RESULT0:.*]] = alloc() : memref<1xf32>
+// CHECK:           %[[RESULT1:.*]] = alloc() : memref<2xf32>
+// CHECK:           call @callee(%[[RESULT0]], %[[RESULT1]]) : (memref<1xf32>, memref<2xf32>) -> ()
+// CHECK:           "test.sink"(%[[RESULT0]], %[[RESULT1]]) : (memref<1xf32>, memref<2xf32>) -> ()
+// CHECK:         }
+func @call_multiple_result() {
+  %0, %1 = call @callee() : () -> (memref<1xf32>, memref<2xf32>)
+  "test.sink"(%0, %1) : (memref<1xf32>, memref<2xf32>) -> ()
+}
+
+// -----
+
+// CHECK-LABEL: func @callee(memref<1xf32>) -> (i1, i32)
+func @callee() -> (i1, memref<1xf32>, i32)
+
+// CHECK-LABEL:   func @call_non_memref_result() {
+// CHECK:           %[[RESULT0:.*]] = alloc() : memref<1xf32>
+// CHECK:           %[[NON_MEMREF_RESULTS:.*]]:2 = call @callee(%[[RESULT0]]) : (memref<1xf32>) -> (i1, i32)
+// CHECK:           "test.sink"(%[[NON_MEMREF_RESULTS]]#0, %[[RESULT0]], %[[NON_MEMREF_RESULTS]]#1) : (i1, memref<1xf32>, i32) -> ()
+// CHECK:         }
+func @call_non_memref_result() {
+  %0, %1, %2 = call @callee() : () -> (i1, memref<1xf32>, i32)
+  "test.sink"(%0, %1, %2) : (i1, memref<1xf32>, i32) -> ()
+}
+
+// -----
+
+func @callee() -> (memref<?xf32>)
+
+func @call_non_memref_result() {
+  // expected-error @+1 {{cannot create out param for dynamically shaped result}}
+  %0 = call @callee() : () -> (memref<?xf32>)
+  "test.sink"(%0) : (memref<?xf32>) -> ()
+}