[mlir] Make MemRef element type extensible

[ftynse]

Summary

Historically, MemRef only supported a restricted list of element types that
were known to be storable in memory. This is unnecessarily restrictive given
the open nature of MLIR's type system. Allow types to opt into being used as
MemRef elements by implementing a type interface. For now, the interface is
merely a declaration with no methods. Later, methods to query, e.g., the type
size or whether a type can alias elements of another type may be added.

Harden the "standard"-to-LLVM conversion against memrefs with non-builtin
types.

See https://llvm.discourse.group/t/rfc-memref-of-custom-types/3558.

Depends On D103826

[ftynse]

cmake

[rriddle]

Can you add a link to the discussion for this?

mlir/docs/Dialects/Builtin.md

36

This looks wrong, BuiltinTypeInterfaces?

mlir/include/mlir/IR/BuiltinTypes.td

315

What about vector types implementing MemRefElementTypeInterface?

[ftynse]

Address review.

[ftynse]

In D103827#2803248, @rriddle wrote:

Can you add a link to the discussion for this?

https://llvm.discourse.group/t/rfc-memref-of-custom-types/3558

mlir/include/mlir/IR/BuiltinTypes.td

315

Only vector types or all allows builtin types? AFAICS, built-in types don't implement any interfaces yet so I hesitated to add one.

[rriddle]

In D103827#2803378, @ftynse wrote:

In D103827#2803248, @rriddle wrote:

Can you add a link to the discussion for this?

https://llvm.discourse.group/t/rfc-memref-of-custom-types/3558

Sorry, I meant can you add it in the commit description?

I honestly would have expected more situations that memref element type was being hardcoded.

mlir/include/mlir/IR/BuiltinTypes.td

252

I wonder if we should put these in a BuiltinTypeInterfaces file instead? I don't feel that strongly either way right now, but might if we end up putting this on the other builtin types (easier to ensure it's defined before used)

[rriddle]

mlir/include/mlir/IR/BuiltinTypes.td

315

They are all builtin types, so if it makes sense to have it on there I would just do it. Given they are all in the same lib, avoiding it doesn't really get us much (and will require hand coding them in anyways).

[ftynse]

In D103827#2804814, @rriddle wrote:

In D103827#2803378, @ftynse wrote:

In D103827#2803248, @rriddle wrote:

Can you add a link to the discussion for this?

https://llvm.discourse.group/t/rfc-memref-of-custom-types/3558

Sorry, I meant can you add it in the commit description?

I did. Arc doesn't seem to pick up amended commit messages, changed it through the web interface as well.

I honestly would have expected more situations that memref element type was being hardcoded.

Most of the in-tree uses have extra checks on the element type, e.g., SPIR-V only accepts memrefs of integers.

mlir/include/mlir/IR/BuiltinTypes.td

252

Will do in a separate commit in case we want to revert it.

[mehdi_amini]

I did. Arc doesn't seem to pick up amended commit messages, changed it through the web interface as well.

You need to pass the --verbatim option to arc.

mlir/docs/Dialects/Builtin.md

 
 ## Operations
 
 [include "Dialects/BuiltinOps.md"]
 
 ## Types
 
 [include "Dialects/BuiltinTypes.md"]
+
+## Type Interfaces
+
+[include "Dialects/BuiltinTypeInterfaces.md"]

mlir/include/mlir/IR/BuiltinTypes.h

 
 //===----------------------------------------------------------------------===//
 // Tablegen Type Declarations
 //===----------------------------------------------------------------------===//
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/IR/BuiltinTypes.h.inc"
 
+//===----------------------------------------------------------------------===//
+// Tablegen Interface Declarations
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/BuiltinTypeInterfaces.h.inc"
+
 namespace mlir {
 //===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
 
 /// This is a builder type that keeps local references to arguments. Arguments
 /// that are passed into the builder must out-live the builder.
 class MemRefType::Builder {

 // Deferred Method Definitions
 //===----------------------------------------------------------------------===//
 
 inline bool BaseMemRefType::classof(Type type) {
   return type.isa<MemRefType, UnrankedMemRefType>();
 }
 
 inline bool BaseMemRefType::isValidElementType(Type type) {
-  return type.isIntOrIndexOrFloat() || type.isa<ComplexType, VectorType>();
+  return type.isIntOrIndexOrFloat() || type.isa<ComplexType, VectorType>() ||
+         type.isa<MemRefElementTypeInterface>();
 }
 
 inline bool FloatType::classof(Type type) {
   return type.isa<BFloat16Type, Float16Type, Float32Type, Float64Type,
                   Float80Type, Float128Type>();
 }
 
 inline FloatType FloatType::getBF16(MLIRContext *ctx) {

mlir/include/mlir/IR/BuiltinTypes.td

let extraClassDeclaration = [{
 
     /// Integer representation maximal bitwidth.
     /// Note: This is aligned with the maximum width of llvm::IntegerType.
     static constexpr unsigned kMaxWidth = (1 << 24) - 1;
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// MemRefElementTypeInterface
+//===----------------------------------------------------------------------===//
+
+def MemRefElementTypeInterface : TypeInterface<"MemRefElementTypeInterface"> {
+  let cppNamespace = "::mlir";
+  let description = [{
+    Indication that this type can be used as element in memref types.
+
+    Implementing this interface establishes a contract between this type and the
+    memref type indicating that this type can be used as element of ranked or
+    unranked memrefs. The type is expected to:
+
+      - model an entity stored in memory;
+      - have non-zero size.
+
+    For example, scalar values such as integers can implement this interface,
+    but indicator types such as `void` or `unit` should not.
+
+    The interface currently has no methods and is used by types to opt into
+    being memref elements. This may change in the future, in particular to
+    require types to provide their size or alignment given a data layout.
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
 
 def Builtin_MemRef : Builtin_Type<"MemRef", "BaseMemRefType"> {
   let summary = "Shaped reference to a region of memory";
   let description = [{
     Syntax:
 
let description = [{
     different types.
 
     A `memref` is allowed to have an unknown rank (e.g. `memref<*xf32>`). The
     purpose of unranked memrefs is to allow external library functions to
     receive memref arguments of any rank without versioning the functions based
     on the rank. Other uses of this type are disallowed or will have undefined
     behavior.
 
+    Are accepted as elements:
+
+    - built-in integer types;
+    - built-in index type;
+    - built-in floating point types;
+    - built-in vector types with elements of the above types;
+    - any other type implementing `MemRefElementTypeInterface`.
+
     ##### Codegen of Unranked Memref
 
     Using unranked memref in codegen besides the case mentioned above is highly
     discouraged. Codegen is concerned with generating loop nests and specialized
     instructions for high-performance, unranked memref is concerned with hiding
     the rank and thus, the number of enclosing loops required to iterate over
     the data. However, if there is a need to code-gen unranked memref, one
     possible path is to cast into a static ranked type based on the dynamic

mlir/include/mlir/IR/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS BuiltinOps.td)
 mlir_tablegen(BuiltinOps.h.inc -gen-op-decls)
 mlir_tablegen(BuiltinOps.cpp.inc -gen-op-defs)
 add_public_tablegen_target(MLIRBuiltinOpsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS BuiltinTypes.td)
 mlir_tablegen(BuiltinTypes.h.inc -gen-typedef-decls)
 mlir_tablegen(BuiltinTypes.cpp.inc -gen-typedef-defs)
+mlir_tablegen(BuiltinTypeInterfaces.h.inc -gen-type-interface-decls)
+mlir_tablegen(BuiltinTypeInterfaces.cpp.inc -gen-type-interface-defs)
 add_public_tablegen_target(MLIRBuiltinTypesIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TensorEncoding.td)
 mlir_tablegen(TensorEncInterfaces.h.inc -gen-attr-interface-decls)
 mlir_tablegen(TensorEncInterfaces.cpp.inc -gen-attr-interface-defs)
 add_public_tablegen_target(MLIRTensorEncodingIncGen)
 
 add_mlir_doc(BuiltinAttributes BuiltinAttributes Dialects/ -gen-attrdef-doc)
 add_mlir_doc(BuiltinLocationAttributes BuiltinLocationAttributes Dialects/ -gen-attrdef-doc)
 add_mlir_doc(BuiltinOps BuiltinOps Dialects/ -gen-op-doc)
 add_mlir_doc(BuiltinTypes BuiltinTypes Dialects/ -gen-typedef-doc)
+add_mlir_doc(BuiltinTypes BuiltinTypeInterfaces Dialects/ -gen-type-interface-docs)

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 
 /// Converts MemRefType to LLVMType. A MemRefType is converted to a struct that
 /// packs the descriptor fields as defined by `getMemRefDescriptorFields`.
 Type LLVMTypeConverter::convertMemRefType(MemRefType type) {
   // When converting a MemRefType to a struct with descriptor fields, do not
   // unpack the `sizes` and `strides` arrays.
   SmallVector<Type, 5> types =
       getMemRefDescriptorFields(type, /*unpackAggregates=*/false);
+  if (types.empty())
+    return {};
   return LLVM::LLVMStructType::getLiteral(&getContext(), types);
 }
 
 static constexpr unsigned kRankInUnrankedMemRefDescriptor = 0;
 static constexpr unsigned kPtrInUnrankedMemRefDescriptor = 1;
 
 /// Convert an unranked memref type into a list of non-aggregate LLVM IR types
 /// that will form the unranked memref descriptor. In particular, the fields
 /// for an unranked memref descriptor are:
 /// 1. index-typed rank, the dynamic rank of this MemRef
 /// 2. void* ptr, pointer to the static ranked MemRef descriptor. This will be
 ///    stack allocated (alloca) copy of a MemRef descriptor that got casted to
 ///    be unranked.
 SmallVector<Type, 2> LLVMTypeConverter::getUnrankedMemRefDescriptorFields() {
   return {getIndexType(),
           LLVM::LLVMPointerType::get(IntegerType::get(&getContext(), 8))};
 }
 
 Type LLVMTypeConverter::convertUnrankedMemRefType(UnrankedMemRefType type) {
+  if (!convertType(type.getElementType()))
+    return {};
   return LLVM::LLVMStructType::getLiteral(&getContext(),
                                           getUnrankedMemRefDescriptorFields());
 }
 
 /// Convert a memref type to a bare pointer to the memref element type.
 Type LLVMTypeConverter::convertMemRefToBarePtr(BaseMemRefType type) {
   if (type.isa<UnrankedMemRefType>())
     // Unranked memref is not supported in the bare pointer calling convention.

mlir/lib/IR/BuiltinTypes.cpp

 //===----------------------------------------------------------------------===//
 /// Tablegen Type Definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/IR/BuiltinTypes.cpp.inc"
 
 //===----------------------------------------------------------------------===//
+/// Tablegen Interface Definitions
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/BuiltinTypeInterfaces.cpp.inc"
+
+//===----------------------------------------------------------------------===//
 // BuiltinDialect
 //===----------------------------------------------------------------------===//
 
 void BuiltinDialect::registerTypes() {
   addTypes<
 #define GET_TYPEDEF_LIST
 #include "mlir/IR/BuiltinTypes.cpp.inc"
       >();

mlir/test/Conversion/StandardToLLVM/convert-static-memref-ops.mlir

func @address() {
     // CHECK: llvm.sub {{.*}} : i32
     // CHECK: llvm.add {{.*}} : i32
     // CHECK: llvm.urem {{.*}} : i32
     // CHECK: llvm.sub {{.*}} : i32
     // CHECK: llvm.inttoptr %{{.*}} : i32 to !llvm.ptr
     return
   }
 }
+
+// -----
+
+// Should not convert memrefs with unsupported types in any convention.
+
+// CHECK: @unsupported_memref_element_type
+// CHECK-SAME: memref<
+// CHECK-NOT: !llvm.struct
+// BAREPTR: @unsupported_memref_element_type
+// BAREPTR-SAME: memref<
+// BAREPTR-NOT: !llvm.ptr
+func private @unsupported_memref_element_type() -> memref<42 x !test.memref_element>
+
+// CHECK: @unsupported_unranked_memref_element_type
+// CHECK-SAME: memref<
+// CHECK-NOT: !llvm.struct
+// BAREPTR: @unsupported_unranked_memref_element_type
+// BAREPTR-SAME: memref<
+// BAREPTR-NOT: !llvm.ptr
+func private @unsupported_unranked_memref_element_type() -> memref<* x !test.memref_element>

mlir/test/Conversion/StandardToLLVM/invalid.mlir

 // RUN: mlir-opt %s -convert-std-to-llvm -verify-diagnostics -split-input-file
 
 // Should not crash on unsupported types in function signatures.
 func private @unsupported_signature() -> tensor<10 x i32>
 
 // -----
 
 func private @partially_supported_signature() -> (vector<10 x i32>, tensor<10 x i32>)
+

mlir/test/IR/parser.mlir

 func private @memref_with_index_elems(memref<1x?xindex>)
 
 // CHECK: func private @memref_with_complex_elems(memref<1x?xcomplex<f32>>)
 func private @memref_with_complex_elems(memref<1x?xcomplex<f32>>)
 
 // CHECK: func private @memref_with_vector_elems(memref<1x?xvector<10xf32>>)
 func private @memref_with_vector_elems(memref<1x?xvector<10xf32>>)
 
+// CHECK: func private @memref_with_custom_elem(memref<1x?x!test.memref_element>)
+func private @memref_with_custom_elem(memref<1x?x!test.memref_element>)
+
 // CHECK: func private @unranked_memref_with_complex_elems(memref<*xcomplex<f32>>)
 func private @unranked_memref_with_complex_elems(memref<*xcomplex<f32>>)
 
 // CHECK: func private @unranked_memref_with_index_elems(memref<*xindex>)
 func private @unranked_memref_with_index_elems(memref<*xindex>)
 
 // CHECK: func private @unranked_memref_with_vector_elems(memref<*xvector<10xf32>>)
 func private @unranked_memref_with_vector_elems(memref<*xvector<10xf32>>)

mlir/test/lib/Dialect/Test/CMakeLists.txt

 add_public_tablegen_target(MLIRTestInterfaceIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestAttrDefs.td)
 mlir_tablegen(TestAttrDefs.h.inc -gen-attrdef-decls)
 mlir_tablegen(TestAttrDefs.cpp.inc -gen-attrdef-defs)
 add_public_tablegen_target(MLIRTestAttrDefIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestTypeDefs.td)
-mlir_tablegen(TestTypeDefs.h.inc -gen-typedef-decls)
-mlir_tablegen(TestTypeDefs.cpp.inc -gen-typedef-defs)
+mlir_tablegen(TestTypeDefs.h.inc -gen-typedef-decls -typedefs-dialect=test)
+mlir_tablegen(TestTypeDefs.cpp.inc -gen-typedef-defs -typedefs-dialect=test)
 add_public_tablegen_target(MLIRTestTypeDefIncGen)
 
 
 set(LLVM_TARGET_DEFINITIONS TestOps.td)
 mlir_tablegen(TestOps.h.inc -gen-op-decls)
 mlir_tablegen(TestOps.cpp.inc -gen-op-defs)
 mlir_tablegen(TestOpsDialect.h.inc -gen-dialect-decls -dialect=test)
 mlir_tablegen(TestOpEnums.h.inc -gen-enum-decls)

mlir/test/lib/Dialect/Test/TestTypeDefs.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TEST_TYPEDEFS
 #define TEST_TYPEDEFS
 
 // To get the test dialect def.
 include "TestOps.td"
+include "mlir/IR/BuiltinTypes.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
 
 // All of the types will extend this class.
 class Test_Type<string name, list<Trait> traits = []>
   : TypeDef<Test_Dialect, name, traits>;
 
 def SimpleTypeA : Test_Type<"SimpleA"> {
   let mnemonic = "smpla";
]> {
   private:
     unsigned extractKind(DataLayoutEntryListRef params,
                          StringRef expectedKind) const;
 
   public:
   }];
 }
 
+def TestMemRefElementType : Test_Type<"TestMemRefElementType",
+                                      [MemRefElementTypeInterface]> {
+  let mnemonic = "memref_element";
+}
+
 #endif // TEST_TYPEDEFS