[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

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

103

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

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

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

103

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 ↗

(On Diff #260890)

These forward declarations look unnecessary in this file

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

44 ↗

(On Diff #260890)

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

47 ↗

(On Diff #260890)

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.

70 ↗

(On Diff #260890)

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.

77 ↗

(On Diff #260890)

.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.

86 ↗

(On Diff #260890)

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

93–94 ↗

(On Diff #260890)

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.

102 ↗

(On Diff #260890)

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

113 ↗

(On Diff #260890)

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

122 ↗

(On Diff #260890)

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
  }
}

126 ↗

(On Diff #260890)

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

15 ↗

(On Diff #260890)

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

77 ↗

(On Diff #260890)

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

122 ↗

(On Diff #260890)

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

166 ↗

(On Diff #261235)

Nit: newline at end of file.

122 ↗

(On Diff #260890)

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.

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

76 ↗

(On Diff #261235)

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

107 ↗

(On Diff #261235)

Nit, it should now be possible do format it as

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

122 ↗

(On Diff #260890)

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

107 ↗

(On Diff #261235)

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

122 ↗

(On Diff #260890)

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

mlir/include/mlir/Conversion/LinalgToLinalg/LinalgToLinalg.h

+//===- LinalgToLinalg.h - Convert tensor Linalg to buffer linalg ----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_CONVERSION_LINALGTOLINALG_CONVERTLINALGTOLINALG_H
+#define MLIR_CONVERSION_LINALGTOLINALG_CONVERTLINALGTOLINALG_H
+
+#include <memory>
+
+namespace mlir {
+class BufferAssignmentPlacer;
+class MLIRContext;
+class ModuleOp;
+class OwningRewritePatternList;
+class TypeConverter;
+template <typename T>
+class OperationPass;
+
+/// Populate the given list with patterns that convert from tensor Linalg to
+/// buffer Linalg.
+void populateTensorLinalgToBufferLinalgConversionPattern(
+    MLIRContext *context, BufferAssignmentPlacer *placer,
+    TypeConverter *converter, OwningRewritePatternList *patterns);
+
+/// Create a pass to convert Linalg operations which work on tensors to use
+/// buffers instead.
+std::unique_ptr<OperationPass<ModuleOp>> createTensorLinalgToBufferLinalgPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_LINALGTOLINALG_CONVERTLINALGTOLINALG_H

mlir/include/mlir/Conversion/Passes.td

 def ConvertVulkanLaunchFuncToVulkanCalls
     : Pass<"launch-func-to-vulkan", "ModuleOp"> {
   let summary = "Convert vulkanLaunch external call to Vulkan runtime external "
                 "calls";
   let constructor = "mlir::createConvertVulkanLaunchFuncToVulkanCallsPass()";
 }
 
 //===----------------------------------------------------------------------===//
+// LinalgToLinalg
+//===----------------------------------------------------------------------===//
+
+def ConvertLinalgToLinalg : Pass<"convert-linalg-to-linalg", "ModuleOp"> {
+  let summary = "Convert the Linalg operations which work with tensor-type "
+                "operands or results to use buffers instead";
+  let constructor = "mlir::createTensorLinalgToBufferLinalgPass()";
+}
+
+//===----------------------------------------------------------------------===//
 // LinalgToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertLinalgToLLVM : Pass<"convert-linalg-to-llvm", "ModuleOp"> {
   let summary = "Convert the operations from the linalg dialect into the LLVM "
                 "dialect";
   let constructor = "mlir::createConvertLinalgToLLVMPass()";
 }

mlir/include/mlir/InitAllPasses.h

 
 #include "mlir/Conversion/AVX512ToLLVM/ConvertAVX512ToLLVM.h"
 #include "mlir/Conversion/GPUToCUDA/GPUToCUDAPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
+#include "mlir/Conversion/LinalgToLinalg/LinalgToLinalg.h"
 #include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h"
 #include "mlir/Conversion/LoopToStandard/ConvertLoopToStandard.h"
 #include "mlir/Conversion/LoopsToGPU/LoopsToGPUPass.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/GPU/Passes.h"

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(AVX512ToLLVM)
 add_subdirectory(GPUToCUDA)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)
+add_subdirectory(LinalgToLinalg)
 add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LoopsToGPU)
 add_subdirectory(LoopToStandard)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(StandardToStandard)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToLoops)

mlir/lib/Conversion/LinalgToLinalg/CMakeLists.txt

+add_mlir_conversion_library(MLIRLinalgToLinalg
+  LinalgToLinalg.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/LinalgToLinalg
+
+  DEPENDS
+  MLIRConversionPassIncGen
+  )
+
+target_link_libraries(MLIRLinalgToLinalg
+  PUBLIC
+  MLIRIR
+  MLIRLinalgOps
+  MLIRStandardOps
+  MLIRPass
+  )

mlir/lib/Conversion/LinalgToLinalg/LinalgToLinalg.cpp

+//===- LinalgToLinalg.cpp - Convert tensor Linalg to buffer linalg --------===//
+//
+// 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/Conversion/LinalgToLinalg/LinalgToLinalg.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/BufferPlacement.h"
+
+using namespace mlir;
+
+namespace {
+/// A pattern to convert Generic Linalg operations which work with tensors to
+/// use buffers instead. A buffer is allocated using BufferAssignmentPlacer for
+/// each operation result. BufferPlacement pass (--buffer-placement) will later
+/// move Alloc operations to the correct positions and insert the missing
+/// Dealloc operations in the correct place.
+class GenericOpConverter
+    : public BufferAssignmentOpConversionPattern<linalg::GenericOp> {
+public:
+  using BufferAssignmentOpConversionPattern<
+      linalg::GenericOp>::BufferAssignmentOpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    auto loc = op.getLoc();
+    SmallVector<Value, 4> args(operands.begin(), operands.end());
+
+    // Update all types to memref types.
+    auto results = op.getOperation()->getResults();
+    for (auto result : results) {
+      auto type = result.getType().cast<ShapedType>();
+      if (!type)
+        op.emitOpError()
+            << "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);
+      result.replaceAllUsesWith(alloc);
+      args.push_back(alloc);
+    }
+
+    // Generate a new linalg operation that works on buffers.
+    auto linalgOp = rewriter.create<linalg::GenericOp>(
+        loc, llvm::None, args, rewriter.getI64IntegerAttr(operands.size()),
+        rewriter.getI64IntegerAttr(results.size()), op.indexing_maps(),
+        op.iterator_types(), op.docAttr(), op.library_callAttr());
+
+    // Move regions from the old operation to the new one.
+    auto &region = linalgOp.region();
+    rewriter.inlineRegionBefore(op.region(), region, region.end());
+
+    // TODO: verify the internal memref-based linalg functionality.
+    auto &entryBlock = region.front();
+    for (auto result : results) {
+      auto type = result.getType().cast<ShapedType>();
+      entryBlock.addArgument(type.getElementType());
+    }
+
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+} // end anonymous namespace
+
+void mlir::populateTensorLinalgToBufferLinalgConversionPattern(
+    MLIRContext *context, BufferAssignmentPlacer *placer,
+    TypeConverter *converter, OwningRewritePatternList *patterns) {
+  // clang-format off
+  patterns->insert<
+                  FunctionAndBlockSignatureConverter,
+                  GenericOpConverter,
+                  NonVoidToVoidReturnOpConverter<
+                      mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>
+  >(context, placer, converter);
+  // clang-format on
+}
+
+namespace {
+/// Converts Linalg operations that work with tensor-type operands or results to
+/// work with memref-type ones instead.
+struct TensorLinalgToBufferLinalgPass
+    : mlir::PassWrapper<TensorLinalgToBufferLinalgPass,
+                        OperationPass<ModuleOp>> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void TensorLinalgToBufferLinalgPass::runOnOperation() {
+  auto &context = getContext();
+  ConversionTarget target(context);
+  BufferAssignmentTypeConverter converter;
+  // Make all linalg operations illegal as long as they work on tensors.
+  target.addLegalDialect<StandardOpsDialect>();
+  target.addDynamicallyLegalDialect<linalg::LinalgDialect>(
+      Optional<ConversionTarget::DynamicLegalityCallbackFn>([&](Operation *op) {
+        auto isIllegalType = [&](Type type) {
+          return !converter.isLegal(type);
+        };
+        return llvm::none_of(op->getOperandTypes(), isIllegalType) &&
+               llvm::none_of(op->getResultTypes(), isIllegalType);
+      }));
+
+  // Mark return operations illegal as long as they return values.
+  target.addDynamicallyLegalOp<mlir::ReturnOp>(
+      [](mlir::ReturnOp returnOp) { return returnOp.getNumOperands() == 0; });
+
+  // Mark the function whose arguments are in tensor-type illegal.
+  target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
+    return converter.isSignatureLegal(funcOp.getType());
+  });
+
+  // Walk over all the functions to apply buffer assignment.
+  getOperation().walk([&](FuncOp function) {
+    OwningRewritePatternList patterns;
+    BufferAssignmentPlacer placer(function);
+    populateTensorLinalgToBufferLinalgConversionPattern(&context, &placer,
+                                                        &converter, &patterns);
+
+    // Applying full conversion
+    return failed(applyFullConversion(function, target, patterns, &converter))
+               ? WalkResult::interrupt()
+               : WalkResult::advance();
+  });
+}
+
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::createTensorLinalgToBufferLinalgPass() {
+  return std::make_unique<TensorLinalgToBufferLinalgPass>();
+}
 No newline at end of file

mlir/test/Conversion/LinalgToLinalg/lingalg-to-linalg.mlir

+// RUN: mlir-opt -convert-linalg-to-linalg -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(%arg1: f32):
+        %1 = exp %arg1 : f32
+        linalg.yield %1, %1 : 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: %{{.*}}, %{{.*}}, %{{.*}}
+//      CHECK: ^{{[a-z0-9_]*}}
+// CHECK-SAME: %[[ARG0:.*]]: f32, %{{.*}}: f32, %{{.*}}: f32
+//      CHECK: %[[RESULT:[a-zA-Z0-9_]*]] = exp %[[ARG0]]
+//      CHECK: linalg.yield %[[RESULT]], %[[RESULT]]
+//      CHECK: [[TYPE]], [[TYPE]], [[TYPE]]
+// CHECK-NEXT: linalg.copy(%[[FIRST_ALLOC]], %[[ARG1_RESULT]])
+// CHECK-NEXT: dealloc %[[FIRST_ALLOC]]
+// CHECK-NEXT: linalg.copy(%[[SECOND_ALLOC]], %[[ARG2_RESULT]])
+// CHECK-NEXT: dealloc %[[SECOND_ALLOC]]
+// CHECK-NEXT: return
 No newline at end of file