[mlir][LLVM] Support opaque pointers in `llvm.mlir.addressof`

[zero9178]

Summary

The verifier of llvm.mlir.addressof did not properly account for opaque pointers, that is, the pointer type not having an element type equal to the type of the referenced global or function. This patch fixes that by skipping the test for the element type if the pointer is opaque.

[rriddle]

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

1657–1659

Can you pull the getType() into a separate variable?

It also feels like we could wrap both of these ifs within a check for !getType().isOpaque().

mlir/test/Dialect/LLVMIR/global.mlir

76

Why do we need this?

[zero9178]

Address review comments

[zero9178]

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

1657–1659

Still had to check for address space mismatch, regardless of whether the pointer is opaque, so I hoisted that check, gave it a more concrete error message and utilized an early return to skip the check of the element types

mlir/test/Dialect/LLVMIR/global.mlir

76

Think I had the intention of testing it remains a opaque pointer, but that is very much out of scope for this patch.

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

 
 // Type constraint accepting any LLVM floating point type.
 def LLVM_AnyFloat : Type<
   CPred<"::mlir::LLVM::isCompatibleFloatingPointType($_self)">,
   "floating point LLVM type">;
 
 // Type constraint accepting any LLVM pointer type.
 def LLVM_AnyPointer : Type<CPred<"$_self.isa<::mlir::LLVM::LLVMPointerType>()">,
-                          "LLVM pointer type">;
+                          "LLVM pointer type", "::mlir::LLVM::LLVMPointerType">;
 
 // Type constraint accepting LLVM pointer type with an additional constraint
 // on the element type.
 class LLVM_PointerTo<Type pointee> : Type<
   And<[LLVM_AnyPointer.predicate,
        Or<[CPred<"$_self.cast<::mlir::LLVM::LLVMPointerType>().isOpaque()">,
            SubstLeaves<
              "$_self",
              "$_self.cast<::mlir::LLVM::LLVMPointerType>().getElementType()",
              pointee.predicate>]>]>,
-  "LLVM pointer to " # pointee.summary>;
+  "LLVM pointer to " # pointee.summary, "::mlir::LLVM::LLVMPointerType">;
 
 // Type constraints accepting LLVM pointer type to integer of a specific width.
 class LLVM_IntPtrBase<int width, int addressSpace = 0> : Type<
   LLVM_PointerTo<I<width>>.predicate,
   "LLVM pointer to " # I<width>.summary>,
   BuildableType<"::mlir::LLVM::LLVMPointerType::get("
                 "::mlir::IntegerType::get($_builder.getContext(), "
                 # width #"), "# addressSpace #")">;

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

def UnnamedAddr : LLVM_EnumAttr<
     "::llvm::GlobalValue::UnnamedAddr",
     "LLVM GlobalValue UnnamedAddr",
     [UnnamedAddrNone, UnnamedAddrLocal, UnnamedAddrGlobal]> {
   let cppNamespace = "::mlir::LLVM";
 }
 
 def LLVM_AddressOfOp : LLVM_Op<"mlir.addressof", [NoSideEffect]> {
   let arguments = (ins FlatSymbolRefAttr:$global_name);
-  let results = (outs LLVM_Type:$res);
+  let results = (outs LLVM_AnyPointer:$res);
 
   let summary = "Creates a pointer pointing to a global or a function";
 
   let description = [{
     Creates an SSA value containing a pointer to a global variable or constant
     defined by `llvm.mlir.global`. The global value can be defined after its
     first referenced. If the global value is a constant, storing into it is not
     allowed.

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

 
 LogicalResult AddressOfOp::verify() {
   auto global = getGlobal();
   auto function = getFunction();
   if (!global && !function)
     return emitOpError(
         "must reference a global defined by 'llvm.mlir.global' or 'llvm.func'");
 
-  if (global &&
-      LLVM::LLVMPointerType::get(global.getType(), global.getAddrSpace()) !=
-          getResult().getType())
+  if (global && (global.getAddrSpace() != getType().getAddressSpace() ||
+                 (!getType().isOpaque() &&
+                  getType().getElementType() != global.getType())))
     return emitOpError(
         "the type must be a pointer to the type of the referenced global");
 
-  if (function && LLVM::LLVMPointerType::get(function.getFunctionType()) !=
-                      getResult().getType())
+  if (function && !getType().isOpaque() &&
+      getType().getElementType() != function.getFunctionType())
     return emitOpError(
         "the type must be a pointer to the type of the referenced function");
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Builder, printer and verifier for LLVM::GlobalOp.

mlir/test/Dialect/LLVMIR/global.mlir

 // CHECK: llvm.mlir.global linkonce_odr
 llvm.mlir.global linkonce_odr @linkonce_odr() : i64
 // CHECK: llvm.mlir.global weak_odr
 llvm.mlir.global weak_odr @weak_odr() : i64
 // CHECK: llvm.mlir.global external @has_thr_local(42 : i64) {thr_local} : i64
 llvm.mlir.global external @has_thr_local(42 : i64) {thr_local} : i64
 // CHECK: llvm.mlir.global external @has_dso_local(42 : i64) {dso_local} : i64
 llvm.mlir.global external @has_dso_local(42 : i64) {dso_local} : i64
+// CHECK: llvm.mlir.global external @has_addr_space(32 : i64) {addr_space = 3 : i32} : i64
+llvm.mlir.global external @has_addr_space(32 : i64) {addr_space = 3: i32} : i64
 
 // CHECK-LABEL: references
 func.func @references() {
   // CHECK: llvm.mlir.addressof @global : !llvm.ptr<i64>
   %0 = llvm.mlir.addressof @global : !llvm.ptr<i64>
 
   // CHECK: llvm.mlir.addressof @".string" : !llvm.ptr<array<6 x i8>>
   %1 = llvm.mlir.addressof @".string" : !llvm.ptr<array<6 x i8>>
 
+  // CHECK: llvm.mlir.addressof @global : !llvm.ptr
+  // CHECK-NOT: <i64>
+  %2 = llvm.mlir.addressof @global : !llvm.ptr
+
+  // CHECK: llvm.mlir.addressof @has_addr_space : !llvm.ptr<3>
+  %3 = llvm.mlir.addressof @has_addr_space : !llvm.ptr<3>
+
   llvm.return
 }
 
 // CHECK: llvm.mlir.global private local_unnamed_addr constant @local(42 : i64) : i64
 llvm.mlir.global private local_unnamed_addr constant @local(42 : i64) : i64
 
 // CHECK: llvm.mlir.global private unnamed_addr constant @foo(42 : i64) : i64
 llvm.mlir.global private unnamed_addr constant @foo(42 : i64) : i64

 
 // -----
 
 llvm.mlir.global internal @g(32 : i64) {addr_space = 3: i32} : i64
 func.func @mismatch_addr_space() {
   // expected-error @+1 {{op the type must be a pointer to the type of the referenced global}}
   llvm.mlir.addressof @g : !llvm.ptr<i64, 4>
 }
+// -----
+
+llvm.mlir.global internal @g(32 : i64) {addr_space = 3: i32} : i64
+
+func.func @mismatch_addr_space_opaque() {
+  // expected-error @+1 {{op the type must be a pointer to the type of the referenced global}}
+  llvm.mlir.addressof @g : !llvm.ptr<4>
+}
 
 // -----
 
 llvm.func @ctor() {
   llvm.return
 }
 
 // CHECK: llvm.mlir.global_ctors {ctors = [@ctor], priorities = [0 : i32]}

mlir/test/Dialect/LLVMIR/invalid.mlir

func.func @invalid_vector_type_5(%a : vector<4xf32>, %idx : i32) -> vector<4xf32> {
   // expected-error@+1 {{'llvm.insertelement' op Type mismatch: cannot insert 'vector<4xf32>' into 'vector<4xf32>'}}
   %b = "llvm.insertelement"(%a, %a, %idx) : (vector<4xf32>, vector<4xf32>, i32) -> vector<4xf32>
   return %b : vector<4xf32>
 }
 
 // -----
 
 func.func @null_non_llvm_type() {
-  // expected-error@+1 {{must be LLVM pointer type, but got 'i32'}}
+  // expected-error@+1 {{custom op 'llvm.mlir.null' invalid kind of type specified}}
   llvm.mlir.null : i32
 }
 
 // -----
 
 func.func @nvvm_invalid_shfl_pred_1(%arg0 : i32, %arg1 : i32, %arg2 : i32, %arg3 : i32) {
   // expected-error@+1 {{expected return type to be a two-element struct with i1 as the second element}}
   %0 = nvvm.shfl.sync bfly %arg0, %arg3, %arg1, %arg2 {return_value_and_is_valid} : i32 -> i32