[mlir] add an interface to support custom types in LLVM dialect pointers

[ftynse]

Summary

This may be necessary in partial multi-stage conversion when a container type
from dialect A containing types from dialect B goes through the conversion
where only dialect A is converted to the LLVM dialect. We will need to keep a
pointer-to-non-LLVM type in the IR until a further conversion can convert
dialect B types to LLVM types.

[wsmoses]

mlir/include/mlir/Dialect/LLVMIR/LLVMOpsInterfaces.td

26

Change on this line seems unnecessary. Probably best to preserve as is, unless the formatting is explicitly preferred.

[ftynse]

Rebase

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOpsInterfaces.td

26

It is preferred.

mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt

 mlir_tablegen(LLVMOpsDialect.h.inc -gen-dialect-decls)
 mlir_tablegen(LLVMOpsDialect.cpp.inc -gen-dialect-defs)
 mlir_tablegen(LLVMOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(LLVMOpsEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRLLVMOpsIncGen)
 
 add_mlir_doc(LLVMOps LLVMOps Dialects/ -gen-op-doc)
 
-add_mlir_interface(LLVMOpsInterfaces)
+set(LLVM_TARGET_DEFINITIONS LLVMOpsInterfaces.td)
+mlir_tablegen(LLVMOpsInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(LLVMOpsInterfaces.cpp.inc -gen-op-interface-defs)
+mlir_tablegen(LLVMTypeInterfaces.h.inc -gen-type-interface-decls)
+mlir_tablegen(LLVMTypeInterfaces.cpp.inc -gen-type-interface-defs)
+add_public_tablegen_target(MLIRLLVMOpsInterfacesIncGen)
 
 set(LLVM_TARGET_DEFINITIONS LLVMOps.td)
 mlir_tablegen(LLVMConversions.inc -gen-llvmir-conversions)
 mlir_tablegen(LLVMConversionEnumsToLLVM.inc -gen-enum-to-llvmir-conversions)
 mlir_tablegen(LLVMConversionEnumsFromLLVM.inc -gen-enum-from-llvmir-conversions)
 add_public_tablegen_target(MLIRLLVMConversionsIncGen)
 
 add_mlir_dialect(NVVMOps nvvm)

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

 def LLVM_AnyStruct : Type<CPred<"$_self.isa<::mlir::LLVM::LLVMStructType>()">,
                          "LLVM structure type">;
 
 // Type constraint accepting opaque LLVM structure type.
 def LLVM_OpaqueStruct : Type<
   And<[LLVM_AnyStruct.predicate,
        CPred<"$_self.cast<::mlir::LLVM::LLVMStructType>().isOpaque()">]>>;
 
+// Type constraint accepting types that implement that pointer element
+// interface.
+def LLVM_PointerElementType : Type<
+  CPred<"$_self.isa<::mlir::LLVM::PointerElementTypeInterface>()">,
+  "LLVM-compatible pointer element type">;
+
+
 // Type constraint accepting any LLVM type that can be loaded or stored, i.e. a
 // type that has size (not void, function or opaque struct type).
 def LLVM_LoadableType : Type<
-  And<[LLVM_PrimitiveType.predicate, Neg<LLVM_OpaqueStruct.predicate>]>,
+  Or<[And<[LLVM_PrimitiveType.predicate, Neg<LLVM_OpaqueStruct.predicate>]>,
+      LLVM_PointerElementType.predicate]>,
   "LLVM type with size">;
 
 // Type constraint accepting any LLVM aggregate type, i.e. structure or array.
 def LLVM_AnyAggregate : Type<
   CPred<"$_self.isa<::mlir::LLVM::LLVMStructType, "
                    "::mlir::LLVM::LLVMArrayType>()">,
   "LLVM aggregate type">;
 

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_GEPOp
   }];
 }
 
 def LLVM_LoadOp : LLVM_Op<"load">, MemoryOpWithAlignmentAndAttributes {
   let arguments = (ins LLVM_PointerTo<LLVM_LoadableType>:$addr,
                    OptionalAttr<SymbolRefArrayAttr>:$access_groups,
                    OptionalAttr<I64Attr>:$alignment, UnitAttr:$volatile_,
                    UnitAttr:$nontemporal);
-  let results = (outs LLVM_Type:$res);
+  let results = (outs LLVM_LoadableType:$res);
   string llvmBuilder = [{
     auto *inst = builder.CreateLoad(
       $addr->getType()->getPointerElementType(), $addr, $volatile_);
   }] # setAlignmentCode # setNonTemporalMetadataCode # setAccessGroupsMetadataCode # [{
     $res = inst;
   }];
   let builders = [
     OpBuilder<(ins "Value":$addr, CArg<"unsigned", "0">:$alignment,

mlir/include/mlir/Dialect/LLVMIR/LLVMOpsInterfaces.td

 def FastmathFlagsInterface : OpInterface<"FastmathFlagsInterface"> {
   let description = [{
     Access to op fastmath flags.
   }];
 
   let cppNamespace = "::mlir::LLVM";
 
   let methods = [
-    InterfaceMethod<"Get fastmath flags", "::mlir::LLVM::FastmathFlags", "fastmathFlags">,
+    InterfaceMethod<"Get fastmath flags", "::mlir::LLVM::FastmathFlags",
+                    "fastmathFlags">,
   ];
 }
 
+//===----------------------------------------------------------------------===//
+// LLVM dialect type interfaces.
+//===----------------------------------------------------------------------===//
+
+// An interface for LLVM pointer element types.
+def LLVM_PointerElementTypeInterface
+    : TypeInterface<"PointerElementTypeInterface"> {
+  let cppNamespace = "::mlir::LLVM";
+
+  let description = [{
+    An interface for types that are allowed as elements of LLVM pointer type.
+    Such types must have a size.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*description=*/"Returns the size of the type in bytes.",
+      /*retTy=*/"unsigned",
+      /*methodName=*/"getSizeInBytes",
+      /*args=*/(ins "const DataLayout &":$dataLayout),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return dataLayout.getTypeSize($_type);
+      }]
+    >
+  ];
+}
+
+
 #endif // LLVM_OPS_INTERFACES

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

 class LLVMDialect;
 
 namespace detail {
 struct LLVMFunctionTypeStorage;
 struct LLVMPointerTypeStorage;
 struct LLVMStructTypeStorage;
 struct LLVMTypeAndSizeStorage;
 } // namespace detail
+} // namespace LLVM
+} // namespace mlir
+
+#include "mlir/Dialect/LLVMIR/LLVMTypeInterfaces.h.inc"
+
+namespace mlir {
+namespace LLVM {
 
 //===----------------------------------------------------------------------===//
 // Trivial types.
 //===----------------------------------------------------------------------===//
 
 // Batch-define trivial types.
 #define DEFINE_TRIVIAL_LLVM_TYPE(ClassName)                                    \
   class ClassName : public Type::TypeBase<ClassName, Type, TypeStorage> {      \

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

LLVMFunctionType::verify(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Pointer type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMPointerType::isValidElementType(Type type) {
-  return !type.isa<LLVMVoidType, LLVMTokenType, LLVMMetadataType,
-                   LLVMLabelType>();
+  return isCompatibleType(type) ? !type.isa<LLVMVoidType, LLVMTokenType,
+                                            LLVMMetadataType, LLVMLabelType>()
+                                : type.isa<PointerElementTypeInterface>();
 }
 
 LLVMPointerType LLVMPointerType::get(Type pointee, unsigned addressSpace) {
   assert(pointee && "expected non-null subtype");
   return Base::get(pointee.getContext(), pointee, addressSpace);
 }
 
 LLVMPointerType
return llvm::TypeSwitch<Type, llvm::TypeSize>(type)
       .Default([](Type ty) {
         assert((ty.isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType,
                        LLVMTokenType, LLVMStructType, LLVMArrayType,
                        LLVMPointerType, LLVMFunctionType>()) &&
                "unexpected missing support for primitive type");
         return llvm::TypeSize::Fixed(0);
       });
 }
+
+#include "mlir/Dialect/LLVMIR/LLVMTypeInterfaces.cpp.inc"

mlir/test/Dialect/LLVMIR/types.mlir

 }
 
 func @verbose() {
   // CHECK: !llvm.struct<(i64, struct<(f32)>)>
   "some.op"() : () -> !llvm.struct<(i64, !llvm.struct<(f32)>)>
   return
 }
 
+// CHECK-LABEL: @ptr_elem_interface
+// CHECK-COUNT-3: !llvm.ptr<!test.smpla>
+func @ptr_elem_interface(%arg0: !llvm.ptr<!test.smpla>) {
+  %0 = llvm.load %arg0 : !llvm.ptr<!test.smpla>
+  llvm.store %0, %arg0 : !llvm.ptr<!test.smpla>
+  return
+}
+
 // -----
 
 // Check that type aliases can be used inside LLVM dialect types. Note that
 // currently they are _not_ printed back as this would require
 // DialectAsmPrinter to have a mechanism for querying the presence and
 // usability of an alias outside of its `printType` method.
 
 !baz = type i64

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

add_mlir_library(MLIRTestDialect
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces
   MLIRDerivedAttributeOpInterface
   MLIRDialect
   MLIRDLTI
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRLinalgTransforms
+  MLIRLLVMIR
   MLIRPass
   MLIRReduce
   MLIRStandard
   MLIRStandardOpsTransforms
   MLIRTensor
   MLIRTransformUtils
   MLIRTransforms
 )

mlir/test/lib/Dialect/Test/TestTypes.cpp

 //===- TestTypes.cpp - MLIR Test Dialect Types ------------------*- C++ -*-===//
 //
 // 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 contains types defined by the TestDialect for testing various
 // features of MLIR.
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestTypes.h"
 #include "TestDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Types.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;

 
 #define GET_TYPEDEF_CLASSES
 #include "TestTypeDefs.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
+namespace {
+
+struct PtrElementModel
+    : public LLVM::PointerElementTypeInterface::ExternalModel<PtrElementModel,
+                                                              SimpleAType> {};
+} // namespace
+
 void TestDialect::registerTypes() {
   addTypes<TestRecursiveType,
 #define GET_TYPEDEF_LIST
 #include "TestTypeDefs.cpp.inc"
            >();
+  SimpleAType::attachInterface<PtrElementModel>(*getContext());
 }
 
 static Type parseTestType(MLIRContext *ctxt, DialectAsmParser &parser,
                           SetVector<Type> &stack) {
   StringRef typeTag;
   if (failed(parser.parseKeyword(&typeTag)))
     return Type();
 

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

tbl_outs = [
         (
             ["-gen-op-interface-decls"],
             "include/mlir/Dialect/LLVMIR/LLVMOpsInterfaces.h.inc",
         ),
         (
             ["-gen-op-interface-defs"],
             "include/mlir/Dialect/LLVMIR/LLVMOpsInterfaces.cpp.inc",
         ),
+        (
+            ["-gen-type-interface-decls"],
+            "include/mlir/Dialect/LLVMIR/LLVMTypeInterfaces.h.inc",
+        ),
+        (
+            ["-gen-type-interface-defs"],
+            "include/mlir/Dialect/LLVMIR/LLVMTypeInterfaces.cpp.inc",
+        ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/LLVMIR/LLVMOpsInterfaces.td",
     deps = [":LLVMOpsTdFiles"],
 )
 
 gentbl_cc_library(
     name = "LLVMDialectAttributesIncGen",

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

deps = [
         "//mlir:ControlFlowInterfaces",
         "//mlir:CopyOpInterface",
         "//mlir:DLTIDialect",
         "//mlir:DataLayoutInterfaces",
         "//mlir:DerivedAttributeOpInterface",
         "//mlir:Dialect",
         "//mlir:IR",
         "//mlir:InferTypeOpInterface",
+        "//mlir:LLVMDialect",
         "//mlir:Pass",
         "//mlir:Reducer",
         "//mlir:SideEffects",
         "//mlir:StandardOps",
         "//mlir:StandardOpsTransforms",
         "//mlir:TensorDialect",
         "//mlir:TransformUtils",
         "//mlir:Transforms",