[mlir][memref][transform] Add memref_to_llvm conversion patterns

[springerm]

Summary

These patterns are exposed via a new "apply_conversion_patterns" op.

Also provide a new type converter that converts from memref to LLVM types. Conversion patterns that lower to LLVM are special: they require an LLVMTypeConverter; a normal TypeConverter is not enough. This revision also adds a new interface method to pattern descriptor ops to verify that the default type converter of the enclosing "apply_conversion_patterns" op is compatible with the set of patterns. At the moment, a simple StringRef is used. This can evolve to a richer type in the future if needed.

[nicolasvasilache]

mlir/test/Dialect/MemRef/transform-ops.mlir

275

can we make the syntax a bit more idiomatic here? E.g.

with type_converter = {

}

[springerm]

address comments

[springerm]

The type converter related part are folded into D157369. The memref patterns op is no longer needed. We can use ConvertToLLVMPatternInterface. I'm going to add a transform op that queries that interface.

mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td

 #define MEMREF_TRANSFORM_OPS
 
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 include "mlir/Dialect/Transform/IR/TransformTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
+def ApplyMemrefToLLVMConversionPatternsOp : Op<Transform_Dialect,
+    "apply_conversion_patterns.memref.memref_to_llvm",
+    [DeclareOpInterfaceMethods<ConversionPatternDescriptorOpInterface,
+                               ["verifyTypeConverter"]>]> {
+  let description = [{
+    Collects patterns that convert memref dialect ops to LLVM dialect ops. These
+    patterns require an "LLVMTypeConverter".
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
+def MemrefToLLVMTypeConverterOp : Op<Transform_Dialect,
+    "apply_conversion_patterns.memref.memref_to_llvm_type_converter",
+    [DeclareOpInterfaceMethods<TypeConverterBuilderOpInterface,
+                               ["getTypeConverterType"]>]> {
+  let description = [{
+    This operation provides an "LLVMTypeConverter" that lowers memref types to
+    LLVM types.
+
+    The type converter can be customized as follows:
+    - `use_aligned_alloc`: Use aligned_alloc in place of malloc for heap
+      allocations.
+    - `index_bitwidth`: Bitwidth of the index type, "0" indicates the size of a
+      machine word.
+    - `use_generic_functions`: Use generic allocation and deallocation functions
+      instead of the classic "malloc", "aligned_alloc" and "free" functions.
+    - `use_opaque_pointers`: Generate LLVM IR using opaque pointers instead of
+      typed pointers.
+  }];
+
+  let arguments = (ins
+      DefaultValuedAttr<BoolAttr, "false">:$use_aligned_alloc,
+      DefaultValuedAttr<I64Attr, "0">:$index_bitwidth,
+      DefaultValuedAttr<BoolAttr, "false">:$use_generic_functions,
+      DefaultValuedAttr<BoolAttr, "false">:$use_opaque_pointers);
+  let assemblyFormat = "attr-dict";
+}
+
 def ApplyExpandOpsPatternsOp : Op<Transform_Dialect,
     "apply_patterns.memref.expand_ops",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
     Collects patterns to rewrite ops within the memref dialect.
 
     - Converts `atomic_rmw` that cannot be lowered to a simple atomic op with
       AtomicRMWOpLowering pattern, e.g. with "minf" or "maxf" attributes, to

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td

InterfaceMethod<
       }],
       /*returnType=*/"void",
       /*name=*/"populatePatterns",
       /*arguments=*/(ins "::mlir::RewritePatternSet &":$patterns)
     >,
   ];
 }
 
+def TypeConverterBuilderOpInterface
+    : OpInterface<"TypeConverterBuilderOpInterface"> {
+  let description = [{
+    This interface should be implemented by ops that specify a type converter
+    for a dialect conversion. Such ops can be used with
+    "apply_conversion_patterns".
+  }];
+
+  let cppNamespace = "::mlir::transform";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the type converter to be used with a dialect conversion.
+      }],
+      /*returnType=*/"std::unique_ptr<::mlir::TypeConverter>",
+      /*name=*/"getTypeConverter",
+      /*arguments=*/(ins)
+    >,
+    StaticInterfaceMethod<
+      /*desc=*/[{
+        Return the type of type converter that this `getTypeConverter` returns.
+        This function is used for op verification.
+      }],
+      /*returnType=*/"StringRef",
+      /*name=*/"getTypeConverterType",
+      /*arguments=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{ return "TypeConverter"; }]
+    >,
+  ];
+}
+
 def ConversionPatternDescriptorOpInterface
     : OpInterface<"ConversionPatternDescriptorOpInterface"> {
   let description = [{
     This interface should be implemented by ops that select conversion patterns
     of a `transform.apply_patterns` op. It provides a method to populate a
     rewrite pattern set with conversion patterns.
 
     Note: Non-conversion rewrite patterns should not be populated with
InterfaceMethod<
         "apply_conversion_patterns" op is used.
       }],
       /*returnType=*/"std::unique_ptr<::mlir::TypeConverter>",
       /*name=*/"getTypeConverter",
       /*arguments=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/"return nullptr;"
     >,
-  ];
-}
-
-def TypeConverterBuilderOpInterface
-    : OpInterface<"TypeConverterBuilderOpInterface"> {
-  let description = [{
-    This interface should be implemented by ops that specify a type converter
-    for a dialect conversion. Such ops can be used with
-    "apply_conversion_patterns".
-  }];
-
-  let cppNamespace = "::mlir::transform";
-
-  let methods = [
     InterfaceMethod<
       /*desc=*/[{
-        Return the type converter to be used with a dialect conversion.
+        Verify the default type converter that is provided by the enclosing
+        "apply_conversion_patterns" op.
       }],
-      /*returnType=*/"std::unique_ptr<::mlir::TypeConverter>",
-      /*name=*/"getTypeConverter",
-      /*arguments=*/(ins)
+      /*returnType=*/"::mlir::LogicalResult",
+      /*name=*/"verifyTypeConverter",
+      /*arguments=*/(ins "TypeConverterBuilderOpInterface":$builder),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/"return success();"
     >,
   ];
 }
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORM_INTERFACES_TD

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

def ApplyConversionPatternsOp : TransformDialectOp<"apply_conversion_patterns",
 
   let arguments = (ins TransformHandleTypeInterface:$target,
                        OptionalAttr<StrArrayAttr>:$legal_ops,
                        OptionalAttr<StrArrayAttr>:$illegal_ops,
                        OptionalAttr<StrArrayAttr>:$legal_dialects,
                        OptionalAttr<StrArrayAttr>:$illegal_dialects,
                        UnitAttr:$partialConversion);
   let results = (outs);
-  let regions = (region VariadicRegion<MaxSizedRegion<1>>:$regions);
+  let regions = (region
+      MaxSizedRegion<1>:$patterns,
+      VariadicRegion<MaxSizedRegion<1>>:$default_type_converter_region);
 
   let assemblyFormat = [{
-    `to` $target $regions attr-dict `:` type($target)
+    `to` $target $patterns
+    (`with` `type_converter` $default_type_converter_region^)?
+    attr-dict `:` type($target)
   }];
   let hasVerifier = 1;
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<(ins
         "Value":$target,
         CArg<"function_ref<void(OpBuilder &, Location)>", "nullptr">:
             $patternsBodyBuilder,
         CArg<"function_ref<void(OpBuilder &, Location)>", "nullptr">:
             $typeConverterBodyBuilder)>,
   ];
 
   let extraClassDeclaration = [{
-    ::mlir::Region &getPatterns() {
-      return getRegion(0);
-    }
-
     ::mlir::transform::TypeConverterBuilderOpInterface getDefaultTypeConverter() {
-      if (getNumRegions() < 2)
+      if (getDefaultTypeConverterRegion().size() == 0)
         return {};
       return ::llvm::cast<::mlir::transform::TypeConverterBuilderOpInterface>(
-          &getRegion(1).front().front());
+          &getDefaultTypeConverterRegion()[0].front().front());
     }
   }];
 }
 
 def ApplyDeadCodeEliminationOp : TransformDialectOp<"apply_dce",
     [TransformOpInterface, TransformEachOpTrait,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      ReportTrackingListenerFailuresOpTrait]> {

mlir/include/mlir/Transforms/DialectConversion.h

 //===----------------------------------------------------------------------===//
 // Type Conversion
 //===----------------------------------------------------------------------===//
 
 /// Type conversion class. Specific conversions and materializations can be
 /// registered using addConversion and addMaterialization, respectively.
 class TypeConverter {
 public:
+  virtual ~TypeConverter() = default;
+
   /// This class provides all of the information necessary to convert a type
   /// signature.
   class SignatureConversion {
   public:
     SignatureConversion(unsigned numOrigInputs)
         : remappedInputs(numOrigInputs) {}
 
     /// This struct represents a range of new types or a single value that

mlir/lib/Dialect/MemRef/TransformOps/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefTransformOps
   MemRefTransformOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef/TransformOps
 
   DEPENDS
   MLIRMemRefTransformOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRArithDialect
   MLIRIR
+  MLIRLLVMCommonConversion
+  MLIRLLVMDialect
   MLIRLoopLikeInterface
   MLIRMemRefDialect
+  MLIRMemRefToLLVM
   MLIRMemRefTransforms
   MLIRNVGPUDialect
   MLIRTransformDialect
   MLIRVectorDialect
 )

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

 //===- MemRefTransformOps.cpp - Implementation of Memref transform 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
+
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "memref-transforms"
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 
 //===----------------------------------------------------------------------===//
+// Apply...ConversionPatternsOp
+//===----------------------------------------------------------------------===//
+
+void transform::ApplyMemrefToLLVMConversionPatternsOp::populatePatterns(
+    TypeConverter &typeConverter, RewritePatternSet &patterns) {
+  populateFinalizeMemRefToLLVMConversionPatterns(
+      static_cast<LLVMTypeConverter &>(typeConverter), patterns);
+}
+
+LogicalResult
+transform::ApplyMemrefToLLVMConversionPatternsOp::verifyTypeConverter(
+    transform::TypeConverterBuilderOpInterface builder) {
+  if (builder.getTypeConverterType() != "LLVMTypeConverter")
+    return emitOpError("expected LLVMTypeConverter");
+  return success();
+}
+
+std::unique_ptr<TypeConverter>
+transform::MemrefToLLVMTypeConverterOp::getTypeConverter() {
+  LowerToLLVMOptions options(getContext());
+  options.allocLowering =
+      (getUseAlignedAlloc() ? LowerToLLVMOptions::AllocLowering::AlignedAlloc
+                            : LowerToLLVMOptions::AllocLowering::Malloc);
+  options.useGenericFunctions = getUseGenericFunctions();
+  options.useOpaquePointers = getUseOpaquePointers();
+
+  if (getIndexBitwidth() != kDeriveIndexBitwidthFromDataLayout)
+    options.overrideIndexBitwidth(getIndexBitwidth());
+
+  return std::make_unique<LLVMTypeConverter>(getContext(), options);
+}
+
+StringRef transform::MemrefToLLVMTypeConverterOp::getTypeConverterType() {
+  return "LLVMTypeConverter";
+}
+
+//===----------------------------------------------------------------------===//
 // Apply...PatternsOp
 //===----------------------------------------------------------------------===//
 
 void transform::ApplyExpandOpsPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   memref::populateExpandOpsPatterns(patterns);
 }
 
class MemRefTransformDialectExtension
     : public transform::TransformDialectExtension<
           MemRefTransformDialectExtension> {
 public:
   using Base::Base;
 
   void init() {
     declareGeneratedDialect<affine::AffineDialect>();
     declareGeneratedDialect<arith::ArithDialect>();
+    declareGeneratedDialect<LLVM::LLVMDialect>();
     declareGeneratedDialect<memref::MemRefDialect>();
     declareGeneratedDialect<nvgpu::NVGPUDialect>();
     declareGeneratedDialect<vector::VectorDialect>();
 
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp.inc"
         >();

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

for (Operation &op : getPatterns().front()) {
       }
     }
   }
   if (getNumRegions() == 2) {
     Region &typeConverterRegion = getRegion(1);
     if (!llvm::hasSingleElement(typeConverterRegion.front()))
       return emitOpError()
              << "expected exactly one op in default type converter region";
-    Operation *typeConverterOp = &typeConverterRegion.front().front();
-    if (!isa<transform::TypeConverterBuilderOpInterface>(typeConverterOp)) {
+    auto typeConverterOp = dyn_cast<transform::TypeConverterBuilderOpInterface>(
+        &typeConverterRegion.front().front());
+    if (!typeConverterOp) {
       InFlightDiagnostic diag = emitOpError()
                                 << "expected default converter child op to "
                                    "implement TypeConverterBuilderOpInterface";
       diag.attachNote(typeConverterOp->getLoc()) << "op without interface";
       return diag;
     }
+    // Check default type converter type.
+    if (!getPatterns().empty()) {
+      for (Operation &op : getPatterns().front()) {
+        auto descriptor =
+            cast<transform::ConversionPatternDescriptorOpInterface>(&op);
+        if (failed(descriptor.verifyTypeConverter(typeConverterOp)))
+          return failure();
+      }
+    }
   }
   if (!getLegalOps() && !getIllegalOps() && !getLegalDialects() &&
       !getIllegalDialects())
     return emitOpError() << "conversion target is not specified";
   return success();
 }
 
 void transform::ApplyConversionPatternsOp::getEffects(

mlir/test/Dialect/MemRef/transform-ops.mlir

 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["memref.alloc"]} in %arg1 : (!transform.any_op) -> !transform.op<"memref.alloc">
   %1 = transform.memref.multibuffer %0 {factor = 2 : i64, skip_analysis} : (!transform.op<"memref.alloc">) -> !transform.any_op
   // Verify that the returned handle is usable.
   transform.test_print_remark_at_operand %1, "transformed" : !transform.any_op
 }
+
+// -----
+
+// CHECK-LABEL: func @lower_to_llvm
+//   CHECK-NOT:   memref.alloc
+//       CHECK:   llvm.call @malloc
+func.func @lower_to_llvm() {
+  %0 = memref.alloc() : memref<2048xi8>
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.apply_conversion_patterns to %0 {
+    transform.apply_conversion_patterns.memref.memref_to_llvm
+  } with type_converter {
+    transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter
+  } {legal_dialects = ["func", "llvm"]} : !transform.any_op
+}

mlir/test/Dialect/Transform/test-pattern-application.mlir

func.func @full_dialect_conversion() -> tensor<5xf32> {
   return %0 : tensor<5xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   transform.apply_conversion_patterns to %0 {
     transform.apply_conversion_patterns.transform.test_conversion_patterns
-  }, {
+  } with type_converter {
     transform.apply_conversion_patterns.transform.test_type_converter
   } {illegal_ops = ["test.foo"],
      legal_ops = ["func.func", "func.return", "test.new_op"]}
       : !transform.any_op
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-error @below{{conversion target is not specified}}
   transform.apply_conversion_patterns to %0 {
     transform.apply_conversion_patterns.transform.test_conversion_patterns
-  }, {
+  } with type_converter {
     transform.apply_conversion_patterns.transform.test_type_converter
   } : !transform.any_op
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-error @below{{pattern descriptor does not specify type converter and apply_conversion_patterns op has no default type converter}}
   transform.apply_conversion_patterns to %0 {
     // expected-note @below{{pattern descriptor op}}
     transform.apply_conversion_patterns.transform.test_conversion_patterns
   } {illegal_ops = ["test.foo"]} : !transform.any_op
 }

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     name = "MemRefTransformOps",
     srcs = glob(["lib/Dialect/MemRef/TransformOps/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/MemRef/TransformOps/*.h"]),
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":ArithDialect",
         ":IR",
+        ":LLVMCommonConversion",
+        ":LLVMDialect",
         ":LoopLikeInterface",
         ":MemRefDialect",
+        ":MemRefToLLVM",
         ":MemRefTransformOpsIncGen",
         ":MemRefTransforms",
         ":NVGPUDialect",
         ":SCFDialect",
         ":TransformDialect",
         ":TransformUtils",
         ":VectorDialect",
         "//llvm:Support",