[mlir] Decompose Bufferization Clone operation into Memref Alloc and Copy.

[dfki-jugr]

Summary

This patch introduces a new conversion to convert bufferization.clone operations
into a memref.alloc and a memref.copy operation. This transformation is needed to
transform all remaining clones which "survive" all previous transformations, before
a given program is lowered further (to LLVM e.g.). Otherwise, these operations
cannot be handled anymore and lead to compile errors.
See: https://llvm.discourse.group/t/bufferization-error-related-to-memref-clone/4665

[dpotop]

This is the first time I am solicited as a reviewer, probably because I raised the issue on discourse (https://llvm.discourse.group/t/bufferization-error-related-to-memref-clone/4665).
The patch implements a solution to the problem we signalled.
The code seems complete, and a test has been provided.
I am not yet competent enough to judge conformance to coding rules. If you think this should have made me resign as a reviewer, please let me know.

[herhut]

mlir/lib/Conversion/MemRefToMemRef/MemRefToMemRef.cpp

23 ↗

(On Diff #388458)

This only works for static shapes.

mlir/test/Conversion/MemRefToMemRef/memref-to-memref.mlir

4 ↗

(On Diff #388458)

Please add a test with dynamic shapes and unknown rank.

I think it is ok if this fails for unknown rank for now but at least dynamic shapes should be supported.

[pifon2a]

mlir/include/mlir/Conversion/Passes.td

355

Why is it not a "transform"? I thought that when we stay within a dialect, then it should be a transformation, not a conversion.

[dfki-jugr]

mlir/include/mlir/Conversion/Passes.td

355

The idea was, that we could simply rename it to "bufferize-to-memref", if a bufferize dialect will be created. In this case, it would be a conversion. But you are right, if we stay within a dialect it should be a transformation.

I suggest to move it to transforms, because this is the right place for this and both renaming and moving is almost the same amount of work.

[pifon2a]

mlir/include/mlir/Conversion/Passes.td

355

Since the bufferization dialect is almost landed, let's just wait a bit and not submit. Then you will be able to just rename it. In the meantime, you can address Stephan's comments.

[bondhugula]

mlir/test/Conversion/MemRefToMemRef/memref-to-memref.mlir

10–13 ↗

(On Diff #388458)

No numbered values in CHECK lines.

[dfki-jugr]

Add dynamic shaped allocs and test cases.
Addressed comments.

[dfki-jugr]

mlir/include/mlir/Conversion/Passes.td

355

Sounds good. Just tell me when it is landed.

[bondhugula]

mlir/include/mlir/Conversion/Passes.td

356

The description looks too generic -- perhaps rephrase to " ... to simpler ones in the MemRef ...".

358

I think you are missing the dependent dialects specification.

mlir/lib/Conversion/MemRefToMemRef/MemRefToMemRef.cpp

6 ↗

(On Diff #389208)

File-level summary comment here.

20 ↗

(On Diff #389208)

Summary doc comment.

[mehdi_amini]

mlir/include/mlir/Conversion/Passes.td

355

I think it did : https://reviews.llvm.org/D114552

[dfki-jugr]

All files are renamed to "bufferization to memref" after the bufferization
dialect has been created.
Addressed further comments.

[herhut]

mlir/include/mlir/Conversion/Passes.td

134

Does this need to be a module pass? This limits concurrency and I think the transformation done here is function local.

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

27

nit: transforms.

38

This limits composition. As an alternative, have the pattern fail and return a failure reason via the rewriter.

48

nit: would memrefType.isDynamicDim(i) work here, too?

50

Should this be a DimOp here?

63

Are these comments needed?

[dfki-jugr]

Addressed comments.

[pifon2a]

mlir/include/mlir/Conversion/Passes.td

137

nit: sort dialects

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

14–21

sort headers?

[herhut]

Please sort the headers, otherwise looks good. Thanks!

[mehdi_amini]

mlir/include/mlir/Conversion/Passes.td

135

Can you expand on the doc a bit, at least mention that it turns a clone into alloc+copy.
I'd also point to the end-to-end doc and mention that this pass alone won't resolve leaks.

[mehdi_amini]

mlir/include/mlir/Conversion/Passes.td

135

Ping on this?

[dfki-jugr]

mlir/include/mlir/Conversion/Passes.td

135

I think, this is not the right place to mention that a clone operation is converted into alloc+copy, since this is a general description of the conversion pass. All other conversion passes also don't have a description or list of all converted operations. If there are further op-conversions added, the description has to be changed again. So, I added these remarks to the conversion implementation in the cpp file.
However, if this pass will not resolve memory leaks in general (also using other conversions), we can add a remark here.

[mehdi_amini]

mlir/include/mlir/Conversion/Passes.td

135

this is not the right place to mention that a clone operation is converted into alloc+copy

Well this is all of what this pass does right now, and the bufferization stuff is complex enough that I'll insist on better documentation and cross referencing.

this is a general description of the conversion pass

My point is that this summary right now is basically useless as a doc: this is purely redundant with the pass mnemonic without adding any new information.

The doc gets published on the website: https://mlir.llvm.org/docs/Passes/#-convert-bufferization-to-memref-convert-operations-from-the-bufferization-dialect-to-the-memref-dialect ; So please fix and improve.

You only added a summary (which isn't published for some reason), but there should be a description fields which can contains multiple paragraph as needed.
If you need inspiration, https://mlir.llvm.org/docs/Passes/#-lower-affine-lower-affine-operations-to-a-combination-of-standard-and-scf-operations

Please also consider any such comment as blocking by default in the future.

[mehdi_amini]

mlir/include/mlir/Conversion/Passes.td

135

Ping here?

[dfki-jugr]

mlir/include/mlir/Conversion/Passes.td

135

Added documentation here:
https://reviews.llvm.org/D115326

[mehdi_amini]

mlir/include/mlir/Conversion/Passes.td

135

Thanks!

mlir/include/mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h

+//===- BufferizationToMemRef.h - Bufferization to MemRef 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_BUFFERIZATIONTOMEMREF_BUFFERIZATIONTOMEMREF_H
+#define MLIR_CONVERSION_BUFFERIZATIONTOMEMREF_BUFFERIZATIONTOMEMREF_H
+
+#include <memory>
+
+namespace mlir {
+class Pass;
+class RewritePatternSet;
+
+/// Collect a set of patterns to convert memory-related operations from the
+/// Bufferization dialect to the MemRef dialect.
+void populateBufferizationToMemRefConversionPatterns(
+    RewritePatternSet &patterns);
+
+std::unique_ptr<Pass> createBufferizationToMemRefPass();
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_BUFFERIZATIONTOMEMREF_BUFFERIZATIONTOMEMREF_H

mlir/include/mlir/Conversion/Passes.h

 //===- Passes.h - Conversion Pass Construction and Registration -----------===//
 //
 // 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_PASSES_H
 #define MLIR_CONVERSION_PASSES_H
 
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.h"
 #include "mlir/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.h"
 #include "mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h"
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
+#include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h"
 #include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h"
 #include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"

mlir/include/mlir/Conversion/Passes.td

def ConvertAsyncToLLVM : Pass<"convert-async-to-llvm", "ModuleOp"> {
   let constructor = "mlir::createConvertAsyncToLLVMPass()";
   let dependentDialects = [
     "arith::ArithmeticDialect",
     "LLVM::LLVMDialect",
   ];
 }
 
 //===----------------------------------------------------------------------===//
+// BufferizationToMemRef
+//===----------------------------------------------------------------------===//
+
+def ConvertBufferizationToMemRef : Pass<"convert-bufferization-to-memref"> {
+  let summary = "Convert operations from the Bufferization dialect to the "
+                "MemRef dialect";
+  let constructor = "mlir::createBufferizationToMemRefPass()";
+  let dependentDialects = ["memref::MemRefDialect", "arith::ArithmeticDialect"];
+}
+
+//===----------------------------------------------------------------------===//
 // ComplexToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertComplexToLLVM : Pass<"convert-complex-to-llvm", "ModuleOp"> {
   let summary = "Convert Complex dialect to LLVM dialect";
   let constructor = "mlir::createConvertComplexToLLVMPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];
 }
Option<"indexBitwidth", "index-bitwidth", "unsigned",
            "Bitwidth of the index type, 0 to use size of machine word">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefToSPIRV
 //===----------------------------------------------------------------------===//
 
 def ConvertMemRefToSPIRV : Pass<"convert-memref-to-spirv", "ModuleOp"> {
   let summary = "Convert MemRef dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertMemRefToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];
   let options = [
     Option<"boolNumBits", "bool-num-bits",
            "int", /*default=*/"8",
            "The number of bits to store a boolean value">
   ];
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

+//===- BufferizationToMemRef.cpp - Bufferization to MemRef 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements patterns to convert Bufferization dialect to MemRef
+// dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h"
+#include "../PassDetail.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+// The CloneOpConversion transforms all bufferization clone operations into
+// memref alloc and memref copy operations.
+struct CloneOpConversion : public OpConversionPattern<bufferization::CloneOp> {
+  using OpConversionPattern<bufferization::CloneOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(bufferization::CloneOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Check for unranked memref types which are currently not supported.
+    Type type = op.getType();
+    if (type.isa<UnrankedMemRefType>()) {
+      return rewriter.notifyMatchFailure(
+          op, "UnrankedMemRefType is not supported.");
+    }
+
+    // Transform a clone operation into alloc + copy operation and pay
+    // attention to the shape dimensions.
+    MemRefType memrefType = type.cast<MemRefType>();
+    Location loc = op->getLoc();
+    SmallVector<Value, 4> dynamicOperands;
+    for (int i = 0; i < memrefType.getRank(); ++i) {
+      if (!memrefType.isDynamicDim(i))
+        continue;
+      Value size = rewriter.createOrFold<arith::ConstantIndexOp>(loc, i);
+      Value dim = rewriter.createOrFold<memref::DimOp>(loc, op.input(), size);
+      dynamicOperands.push_back(dim);
+    }
+    Value alloc = rewriter.replaceOpWithNewOp<memref::AllocOp>(op, memrefType,
+                                                               dynamicOperands);
+    rewriter.create<memref::CopyOp>(loc, op.input(), alloc);
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateBufferizationToMemRefConversionPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<CloneOpConversion>(patterns.getContext());
+}
+
+namespace {
+struct BufferizationToMemRefPass
+    : public ConvertBufferizationToMemRefBase<BufferizationToMemRefPass> {
+  BufferizationToMemRefPass() = default;
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    populateBufferizationToMemRefConversionPatterns(patterns);
+
+    ConversionTarget target(getContext());
+    target.addLegalDialect<memref::MemRefDialect>();
+    target.addLegalOp<arith::ConstantOp>();
+    target.addIllegalDialect<bufferization::BufferizationDialect>();
+
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::createBufferizationToMemRefPass() {
+  return std::make_unique<BufferizationToMemRefPass>();
+}

mlir/lib/Conversion/BufferizationToMemRef/CMakeLists.txt

+add_mlir_conversion_library(MLIRBufferizationToMemRef
+  BufferizationToMemRef.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/BufferizationToMemRef
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRBufferization
+  )

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(ArithmeticToLLVM)
 add_subdirectory(ArithmeticToSPIRV)
 add_subdirectory(ArmNeon2dToIntr)
 add_subdirectory(AsyncToLLVM)
+add_subdirectory(BufferizationToMemRef)
 add_subdirectory(ComplexToLLVM)
 add_subdirectory(ComplexToStandard)
 add_subdirectory(GPUCommon)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)

mlir/test/Conversion/BufferizationToMemRef/bufferization-to-memref.mlir

+// RUN: mlir-opt -verify-diagnostics -convert-bufferization-to-memref -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @conversion_static
+func @conversion_static(%arg0 : memref<2xf32>) -> memref<2xf32> {
+    %0 = bufferization.clone %arg0 : memref<2xf32> to memref<2xf32>
+    memref.dealloc %arg0 : memref<2xf32>
+    return %0 : memref<2xf32>
+}
+
+// CHECK:      %[[ALLOC:.*]] = memref.alloc
+// CHECK-NEXT: memref.copy %[[ARG:.*]], %[[ALLOC]]
+// CHECK-NEXT: memref.dealloc %[[ARG]]
+// CHECK-NEXT: return %[[ALLOC]]
+
+// -----
+
+// CHECK-LABEL: @conversion_dynamic
+func @conversion_dynamic(%arg0 : memref<?xf32>) -> memref<?xf32> {
+    %1 = bufferization.clone %arg0 : memref<?xf32> to memref<?xf32>
+    memref.dealloc %arg0 : memref<?xf32>
+    return %1 : memref<?xf32>
+}
+
+// CHECK:      %[[CONST:.*]] = arith.constant
+// CHECK-NEXT: %[[DIM:.*]] = memref.dim %[[ARG:.*]], %[[CONST]]
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc(%[[DIM]])
+// CHECK-NEXT: memref.copy %[[ARG]], %[[ALLOC]]
+// CHECK-NEXT: memref.dealloc %[[ARG]]
+// CHECK-NEXT: return %[[ALLOC]]
+
+// -----
+
+func @conversion_unknown(%arg0 : memref<*xf32>) -> memref<*xf32> {
+// expected-error@+1 {{failed to legalize operation 'bufferization.clone' that was explicitly marked illegal}}
+    %1 = bufferization.clone %arg0 : memref<*xf32> to memref<*xf32>
+    memref.dealloc %arg0 : memref<*xf32>
+    return %1 : memref<*xf32>
+}