[MLIR][LINALG] Convert Linalg to Linalg

[dfki-ehna]

Summary

This is a basic pass to convert Linalg.GenericOp which works on tensors to use buffers instead.

[ftynse]

lib/Conversion is intended to for conversions from dialect A to dialect B to ensure it can depend on both without creating unnecessary dependencies A->B or B->A. Since this transformation stays within Linalg, it should go to lib/Dialect/Linalg/Transforms. Beyond being cleaner library-wise, it can have a more meaningful name than "linalg to linalg", e.g. "TensorsToBuffers" or "BufferAssignment".

[pifon2a]

just a suggestion: call the pass ConvertLinalgOnTensorsToBuffers?

[pifon2a]

mlir/lib/Conversion/LinalgToLinalg/LinalgToLinalg.cpp

74 ↗

(On Diff #260605)

shouldn't it be rewriter.replaceOp(op, <what to replace results with>)? Please, add a test with two chained Linalg ops.

103 ↗

(On Diff #260605)

why is the definition of runOnOperation() method not inside TensorLinalgToBufferLinalgPass? I don't think it improved readability here. Is there some guide that does not allow doing that?

[bondhugula]

In D78996#2007904, @pifon2a wrote:

just a suggestion: call the pass ConvertLinalgOnTensorsToBuffers?

+1.

[rriddle]

mlir/lib/Conversion/LinalgToLinalg/LinalgToLinalg.cpp

135 ↗

(On Diff #260605)

Just return applyFullConversion directly. failure converts to interrupt

[dfki-ehna]

Address the comments.

[dfki-ehna]

In D78996#2007630, @ftynse wrote:

lib/Conversion is intended to for conversions from dialect A to dialect B to ensure it can depend on both without creating unnecessary dependencies A->B or B->A. Since this transformation stays within Linalg, it should go to lib/Dialect/Linalg/Transforms. Beyond being cleaner library-wise, it can have a more meaningful name than "linalg to linalg", e.g. "TensorsToBuffers" or "BufferAssignment".

Thanks. The files were moved to the correct locations.

mlir/lib/Conversion/LinalgToLinalg/LinalgToLinalg.cpp

74 ↗

(On Diff #260605)

Thanks and resolved. We were replacing the results' uses without rewriter using replaceAllUsesWith which was unsafe.

103 ↗

(On Diff #260605)

Resolved. The reason for doing this was consistency with some other CPP files in the lib/conversion directory.

[ftynse]

Here's a more detailed review. I am not familiar enough with the buffer assignment mechanism that was introduced, deferring that to @pifon2a

mlir/include/mlir/Dialect/Linalg/Passes.h

20

These forward declarations look unnecessary in this file

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

45

Please, call .reserve on vectors before using push_back in a loop

48

rewriter.notifyMatchFailure

If this is an invariant of the op itself, it should be checked in the verifier and can be asserted to be correct here.

71

Nit: we prefer auto in places where it improves readability, i.e. if the type is too long to spell out (e.g., iterators, zippers) or if it is obvious due to a cast. I think Region & would be just fine here.

78

.addArgument (and any inplace modification) is not allowed in a conversion pattern because it's not undoable. You'll have to create a new block with rewriter.createBlock and move/clone the operations from the existing block.

87

LLVM prefers declaring functions static to putting them in an anonymous namespace, unlike classes.

94–95

Could you have using <foo> = NonVoidToVoidReturnOpConverter<mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp> and put just <foo> here to get less crazy formatting? Don't know what's the name for <foo> though.

103

Given that the pass is declaratively specified, use ConvertLinalgOnTensorsToBuffersPassBase<ConvertLinalgOnTensorToBuffers> instead. (also feel free to choose a shorter name for the anonymous class)

114

Style nit: let's put this lambda outside the function call

123

This looks like it would make impossible to have simple helper functions that have no Linalg in them.

func @square(%arg0: f32) {
  %0 = mulf %arg0, %arg0 : f32
  return %0 : f32
}

func @foo(...) {
  linalg.generic (...buffers...) {
    %0 = call @square(...) : (f32) -> f32
    linalg.yield %0 : f32
  }
}

127

Similar to the above, but I might understand that we want this pass to fail if there remains any function on tensors (maybe it doesn't use linalg at all...)

mlir/test/Dialect/Linalg/tensors-to-buffers.mlir

16

Do we actually care about lines being strictly consecutive, or just them being ordered suffices? If the latter, drop the -NEXT.

[ftynse]

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

78

I filed a feature request to make this easier, PR45738. But this needs to be correct in the meantime.

[dfki-ehna]

Address the comments.

[dfki-ehna]

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

123

Thanks for mentioning this. We addressed this issue with a quick fix to cover more test cases in this PR. However, the combination of FunctionAndBlockSignatureConverter and NonVoidToVoidReturnOpConverter of BufferAssignment would still fail in cases like the following:

func @function(%arg0: f32, %arg1: tensor<4xf32>) -> (f32, f32) {
  %0 = mulf %arg0, %arg0 : f32
  return %0, %0 : f32, f32
}

FunctionAndBlockSignatureConverter converts the function to:

func @function(%arg0: f32, %arg1: memref<4xf32>, f32, f32) {
  %0 = mulf %arg0, %arg0 : f32
  return %0, %0 : f32, f32
}

but the return operation remains unchanged and leaves the IR in a broken state.
These converters must be fixed more fundamentally and we are going to do so in the next PR.

[mravishankar]

Thanks. This is really useful!

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

123

Understand that this will be fixed later, but if a function is returning a scalar of scalarType, the modified function should have an argument that is memref<scalarType>. Otherwise I am not sure how to maintain the return semantics.

190

Nit: newline at end of file.

mlir/test/Dialect/Linalg/tensors-to-buffers.mlir

76

Nit: newline at end of file.

[ftynse]

Thanks!
Don't forget to update the commit message that still says "convert linalg to linalg".

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

108

Nit, it should now be possible do format it as

patterns->insert<FunctionAndBlockSignatureConverter,
                 GenericOpConverter,
                 ReturnOpConverter>(context, placer, converter);

123

Please add a big scary TODO it the comment that clearly describes the issue and you are good to go!

[dfki-ehna]

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

108

Clang-format would still fit GenericOpConverter in front of FunctionAndBlockSignatureConverter. So, I kept the clang-format off/on before and after.

123

Please look at our proposed solution for this issue (Update the FunctionAndBlockSignatureConverter and NonVoidToVoidReturnOpConverter of Buffer Assignment).

mlir/include/mlir/Dialect/Linalg/Passes.h

 
 #ifndef MLIR_DIALECT_LINALG_PASSES_H_
 #define MLIR_DIALECT_LINALG_PASSES_H_
 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
 
 namespace mlir {
 class FuncOp;
 class MLIRContext;
 class ModuleOp;
 template <typename T> class OperationPass;
 class OwningRewritePatternList;
 class Pass;
 
 std::unique_ptr<OperationPass<FuncOp>> createLinalgFusionPass();
 std::unique_ptr<Pass> createLinalgFusionOfTensorOpsPass();

 /// std.load/std.store accesses.
 std::unique_ptr<OperationPass<FuncOp>> createConvertLinalgToParallelLoopsPass();
 
 /// Create a pass to convert Linalg operations to affine.for loops and
 /// affine_load/affine_store accesses.
 /// Placeholder for now, this is NYI.
 std::unique_ptr<OperationPass<FuncOp>> createConvertLinalgToAffineLoopsPass();
 
+/// Create a pass to convert Linalg operations which work on tensors to use
+/// buffers instead.
+std::unique_ptr<OperationPass<ModuleOp>>
+createConvertLinalgOnTensorsToBuffersPass();
+
 /// Patterns for fusing linalg operation on tensors.
 void populateLinalgTensorOpsFusionPatterns(MLIRContext *context,
                                            OwningRewritePatternList &patterns);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_PASSES_H_

mlir/include/mlir/Dialect/Linalg/Passes.td

def LinalgLowerToAffineLoops : FunctionPass<"convert-linalg-to-affine-loops"> {
   let constructor = "mlir::createConvertLinalgToAffineLoopsPass()";
 }
 
 def LinalgLowerToLoops : FunctionPass<"convert-linalg-to-loops"> {
   let summary = "Lower the operations from the linalg dialect into loops";
   let constructor = "mlir::createConvertLinalgToLoopsPass()";
 }
 
+def LinalgOnTensorsToBuffers : Pass<"convert-linalg-on-tensors-to-buffers", "ModuleOp"> {
+  let summary = "Convert the Linalg operations which work on tensor-type "
+                "operands or results to use buffers instead";
+  let constructor = "mlir::createConvertLinalgOnTensorsToBuffersPass()";
+}
+
 def LinalgLowerToParallelLoops
     : FunctionPass<"convert-linalg-to-parallel-loops"> {
   let summary = "Lower the operations from the linalg dialect into parallel "
                 "loops";
   let constructor = "mlir::createConvertLinalgToParallelLoopsPass()";
 }
 
 def LinalgPromotion : FunctionPass<"linalg-promote-subviews"> {

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

 add_mlir_dialect_library(MLIRLinalgTransforms
   Fusion.cpp
   Interchange.cpp
   Loops.cpp
   Promotion.cpp
+  TensorsToBuffers.cpp
   Tiling.cpp
   Transforms.cpp
   Vectorization.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS

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

+//===- TensorsToBuffers.cpp - Transformation from tensors to buffers ------===//
+//
+// 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 the conversion from tensors to buffers on Linalg
+// operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/BufferPlacement.h"
+
+using namespace mlir;
+using ReturnOpConverter =
+    NonVoidToVoidReturnOpConverter<mlir::ReturnOp, mlir::ReturnOp,
+                                   linalg::CopyOp>;
+
+namespace {
+/// A pattern to convert Generic Linalg operations which work on tensors to
+/// use buffers. A buffer is allocated using BufferAssignmentPlacer for
+/// each operation result. BufferPlacement pass should be later used to move
+/// Alloc operations to the correct positions and insert the missing Dealloc
+/// operations in the correct places.
+class GenericOpConverter
+    : public BufferAssignmentOpConversionPattern<linalg::GenericOp> {
+public:
+  using BufferAssignmentOpConversionPattern<
+      linalg::GenericOp>::BufferAssignmentOpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    Location loc = op.getLoc();
+    ResultRange results = op.getOperation()->getResults();
+    SmallVector<Value, 2> newArgs, newResults;
+    newArgs.reserve(operands.size() + results.size());
+    newArgs.append(operands.begin(), operands.end());
+    newResults.reserve(results.size());
+
+    // Update all types to memref types.
+    for (auto result : results) {
+      auto type = result.getType().cast<ShapedType>();
+      assert(type && "tensor to buffer conversion expects ranked results");
+      if (!type.hasStaticShape())
+        return rewriter.notifyMatchFailure(
+            op, "dynamic shapes not currently supported");
+      auto memrefType = MemRefType::get(type.getShape(), type.getElementType());
+
+      // Compute alloc position and insert a custom allocation node.
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.restoreInsertionPoint(
+          bufferAssignment->computeAllocPosition(result));
+      auto alloc = rewriter.create<AllocOp>(loc, memrefType);
+      newArgs.push_back(alloc);
+      newResults.push_back(alloc);
+    }
+
+    // Generate a new linalg operation that works on buffers.
+    auto linalgOp = rewriter.create<linalg::GenericOp>(
+        loc, llvm::None, newArgs, rewriter.getI64IntegerAttr(operands.size()),
+        rewriter.getI64IntegerAttr(results.size()), op.indexing_maps(),
+        op.iterator_types(), op.docAttr(), op.library_callAttr());
+
+    // Create a new block in the region of the new Generic Op.
+    Block &oldBlock = op.getRegion().front();
+    Region &newRegion = linalgOp.region();
+    Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
+                                           oldBlock.getArgumentTypes());
+
+    // Add the result arguments to the new block.
+    for (auto result : newResults)
+      newBlock->addArgument(
+          result.getType().cast<ShapedType>().getElementType());
+
+    // Clone the body of the old block to the new block.
+    BlockAndValueMapping mapping;
+    for (unsigned i = 0; i < oldBlock.getNumArguments(); i++)
+      mapping.map(oldBlock.getArgument(i), newBlock->getArgument(i));
+    rewriter.setInsertionPointToEnd(newBlock);
+    for (auto &op : oldBlock.getOperations()) {
+      Operation *clonedOp = rewriter.clone(op, mapping);
+      mapping.map(op.getResults(), clonedOp->getResults());
+    }
+
+    // Replace the results of the old Generic Op with the results of the new
+    // one.
+    rewriter.replaceOp(op, newResults);
+    return success();
+  }
+};
+
+/// Populate the given list with patterns to convert Linalg operations on
+/// tensors to buffers.
+static void populateConvertLinalgOnTensorsToBuffersPattern(
+    MLIRContext *context, BufferAssignmentPlacer *placer,
+    TypeConverter *converter, OwningRewritePatternList *patterns) {
+  // clang-format off
+  patterns->insert<FunctionAndBlockSignatureConverter,
+                   GenericOpConverter,
+                   ReturnOpConverter>(context, placer, converter);
+  // clang-format on
+}
+
+/// Converts Linalg operations that work on tensor-type operands or results to
+/// work on buffers.
+struct ConvertLinalgOnTensorsToBuffers
+    : public LinalgOnTensorsToBuffersBase<ConvertLinalgOnTensorsToBuffers> {
+  void runOnOperation() override {
+    MLIRContext &context = getContext();
+    ConversionTarget target(context);
+    BufferAssignmentTypeConverter converter;
+
+    // Mark all Standard operations legal.
+    target.addLegalDialect<StandardOpsDialect>();
+
+    // Mark all Linalg operations illegal as long as they work on tensors.
+    auto isIllegalType = [&](Type type) { return !converter.isLegal(type); };
+    auto isLegalOperation = [&](Operation *op) {
+      return llvm::none_of(op->getOperandTypes(), isIllegalType) &&
+             llvm::none_of(op->getResultTypes(), isIllegalType);
+    };
+    target.addDynamicallyLegalDialect<linalg::LinalgDialect>(
+        Optional<ConversionTarget::DynamicLegalityCallbackFn>(
+            isLegalOperation));
+
+    // TODO: Considering the following dynamic legality checks, the current
+    // implementation of FunctionAndBlockSignatureConverter of Buffer Assignment
+    // will convert the function signature incorrectly. This converter moves
+    // all the return values of the function to the input argument list without
+    // considering the return value types and creates a void function. However,
+    // the NonVoidToVoidReturnOpConverter doesn't change the return operation if
+    // its operands are not tensors. The following example leaves the IR in a
+    // broken state.
+    //
+    // @function(%arg0: f32, %arg1: tensor<4xf32>) -> (f32, f32) {
+    //    %0 = mulf %arg0, %arg0 : f32
+    //    return %0, %0 : f32, f32
+    // }
+    //
+    // broken IR after conversion:
+    //
+    // func @function(%arg0: f32, %arg1: memref<4xf32>, f32, f32) {
+    //    %0 = mulf %arg0, %arg0 : f32
+    //    return %0, %0 : f32, f32
+    // }
+    //
+    // This issue must be fixed in FunctionAndBlockSignatureConverter and
+    // NonVoidToVoidReturnOpConverter.
+
+    // Mark Standard Return operations illegal as long as one operand is tensor.
+    target.addDynamicallyLegalOp<mlir::ReturnOp>([&](mlir::ReturnOp returnOp) {
+      return llvm::none_of(returnOp.getOperandTypes(), isIllegalType);
+    });
+
+    // Mark the function operation illegal as long as an argument is tensor.
+    target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
+      return converter.isSignatureLegal(funcOp.getType()) &&
+             llvm::none_of(funcOp.getType().getResults(),
+                           [&](Type type) { return type.isa<MemRefType>(); });
+    });
+
+    // Walk over all the functions to apply buffer assignment.
+    getOperation().walk([&](FuncOp function) {
+      OwningRewritePatternList patterns;
+      BufferAssignmentPlacer placer(function);
+      populateConvertLinalgOnTensorsToBuffersPattern(&context, &placer,
+                                                     &converter, &patterns);
+
+      // Applying full conversion
+      return WalkResult(
+          applyFullConversion(function, target, patterns, &converter));
+    });
+  }
+};
+} // end anonymous namespace
+
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::createConvertLinalgOnTensorsToBuffersPass() {
+  return std::make_unique<ConvertLinalgOnTensorsToBuffers>();
+}

mlir/test/Dialect/Linalg/tensors-to-buffers.mlir

+// RUN: mlir-opt -convert-linalg-on-tensors-to-buffers -buffer-placement -split-input-file %s | FileCheck %s -dump-input-on-failure
+
+#map0 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: func @muliple_results_generic_op
+func @muliple_results_generic_op(%arg0: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+    %0, %1 = linalg.generic {args_in = 1 : i64, args_out = 2 : i64, indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel"]} %arg0 {
+    ^bb0(%gen_arg1: f32):
+        %tmp1 = exp %gen_arg1 : f32
+        linalg.yield %tmp1, %tmp1 : f32, f32
+    }: tensor<4xf32> -> (tensor<4xf32>, tensor<4xf32>)
+    return %0, %1 : tensor<4xf32>, tensor<4xf32>
+}
+//      CHECK: (%[[NEW_ARG0:.*]]: [[TYPE:.*]], %[[ARG1_RESULT:.*]]: [[TYPE]], %[[ARG2_RESULT:.*]]: [[TYPE]])
+//      CHECK: %[[FIRST_ALLOC:.*]] = alloc() : [[TYPE]]
+//      CHECK: %[[SECOND_ALLOC:.*]] = alloc() : [[TYPE]]
+//      CHECK: linalg.generic
+// CHECK-SAME: %[[NEW_ARG0]], %[[FIRST_ALLOC]], %[[SECOND_ALLOC]]
+// CHECK-NEXT: ^{{[a-z0-9_]*}}
+// CHECK-SAME: %{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32
+// CHECK-NEXT: %{{.*}} = exp
+// CHECK-NEXT: linalg.yield
+// CHECK-NEXT: [[TYPE]], [[TYPE]], [[TYPE]]
+//      CHECK: linalg.copy(%[[FIRST_ALLOC]], %[[ARG1_RESULT]])
+//      CHECK: dealloc %[[FIRST_ALLOC]]
+//      CHECK: linalg.copy(%[[SECOND_ALLOC]], %[[ARG2_RESULT]])
+//      CHECK: dealloc %[[SECOND_ALLOC]]
+//      CHECK: return
+
+// -----
+
+#map0 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: func @chained_operations
+func @chained_operations(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+    %0 = linalg.generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0], iterator_types = ["parallel"]} %arg0 {
+    ^bb0(%gen_arg1: f32):
+        %tmp1 = exp %gen_arg1 : f32
+        linalg.yield %tmp1 : f32
+    }: tensor<4xf32> -> tensor<4xf32>
+    %1 = linalg.generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0], iterator_types = ["parallel"]} %0 {
+    ^bb0(%gen_arg2: f32):
+        %tmp2 = exp %gen_arg2 : f32
+        linalg.yield %tmp2 : f32
+    }: tensor<4xf32> -> tensor<4xf32>
+    return %1 : tensor<4xf32>
+}
+//      CHECK: (%[[NEW_ARG0:.*]]: [[TYPE:.*]], %[[ARG1_RESULT:.*]]: [[TYPE]])
+//      CHECK: %[[FIRST_ALLOC:.*]] = alloc() : [[TYPE]]
+//      CHECK: linalg.generic
+// CHECK-SAME: %[[NEW_ARG0]], %[[FIRST_ALLOC]]
+//      CHECK: ^{{[a-z0-9_]*}}
+// CHECK-SAME: %{{.*}}: f32, %{{.*}}: f32
+//      CHECK: [[TYPE]], [[TYPE]]
+//      CHECK: %[[SECOND_ALLOC:.*]] = alloc() : [[TYPE]]
+//      CHECK: linalg.generic
+// CHECK-SAME: %[[FIRST_ALLOC]], %[[SECOND_ALLOC]]
+//      CHECK: ^{{[a-z0-9_]*}}
+// CHECK-SAME: %{{.*}}: f32, %{{.*}}: f32
+//      CHECK: [[TYPE]], [[TYPE]]
+//      CHECK: dealloc %[[FIRST_ALLOC]]
+//      CHECK: linalg.copy(%[[SECOND_ALLOC]], %[[ARG1_RESULT]])
+//      CHECK: dealloc %[[SECOND_ALLOC]]
+//      CHECK: return
+
+// -----
+
+// CHECK-LABEL: func @no_linalg_op
+func @no_linalg_op(%arg0: f32) -> (f32, f32) {
+  %0 = mulf %arg0, %arg0 : f32
+  return %0, %0 : f32, f32
+}
+//      CHECK: (%[[NEW_ARG0:.*]]: [[TYPE:.*]]) -> ([[TYPE]], [[TYPE]])
+//      CHECK: %[[RESULT:.*]] = mulf %[[NEW_ARG0]], %[[NEW_ARG0]] : [[TYPE]]
+//      CHECK: return %[[RESULT]], %[[RESULT]] : [[TYPE]], [[TYPE]]