[mlir][LLVM] Verify correct pointer casts with `llvm.bitcast`

[zero9178]

Summary

llvm.bitcast has so far not had a verifier and this allowed various bugs to sneak into the codebase (including within tests!) which could only be caught once translated to actual LLVM IR. This patch fixes those problematic cases by now verifying bitcasts on pointers are done correctly.

Specifically, it verifies that if pointers are involved, that both result and source types are pointers, that this also applies to vector of pointers and that pointer casts are of the same address space.

The only thing left unverified is the general case of "source type size does not match result type size". I think this case is less trivial and more prone to false positives, so I did not yet implement it.

[Dinistro]

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

2336

Should we also check if the address-spaces of pointers in vectors match up, e.g., llvm.vec<4 x ptr<1>> and llvm.vec<4 x ptr<2>>, or is that checked elsewhere already?

[gysit]

Nice!

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

2336

I believe this is checked below

resultType.getAddressSpace() != sourceType.getAddressSpace()

since extractVectorElementType gets the element type of the vector.

2353

ultra nit: I would probably put a comma.

mlir/test/Dialect/LLVMIR/invalid.mlir

1327

nit: invalid_bitcast_addrspace?

[zero9178]

address review comments and add test to demonstrate invalid address casts on vector of pointers

[zero9178]

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

2336

I've added a test to show it works with vectors too

[Dinistro]

LGTM!

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

class LLVM_CastOp<string mnemonic, string instName, Type type,
   string mlirBuilder = [{
     $res = $_builder.create<$_qualCppClassName>(
       $_location, $_resultType, $arg);
   }];
 }
 def LLVM_BitcastOp : LLVM_CastOp<"bitcast", "BitCast",
                                  LLVM_AnyNonAggregate, LLVM_AnyNonAggregate> {
   let hasFolder = 1;
+  let hasVerifier = 1;
 }
 def LLVM_AddrSpaceCastOp : LLVM_CastOp<"addrspacecast", "AddrSpaceCast",
                                        LLVM_ScalarOrVectorOf<LLVM_AnyPointer>,
                                        LLVM_ScalarOrVectorOf<LLVM_AnyPointer>> {
   let hasFolder = 1;
 }
 def LLVM_IntToPtrOp : LLVM_CastOp<"inttoptr", "IntToPtr",
                                   LLVM_ScalarOrVectorOf<AnyInteger>,

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

if (getOrdering() == AtomicOrdering::not_atomic ||
       getOrdering() == AtomicOrdering::unordered ||
       getOrdering() == AtomicOrdering::monotonic)
     return emitOpError("can be given only acquire, release, acq_rel, "
                        "and seq_cst orderings");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// Folder for LLVM::BitcastOp
+// Folder and verifier for LLVM::BitcastOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult LLVM::BitcastOp::fold(FoldAdaptor adaptor) {
   // bitcast(x : T0, T0) -> x
   if (getArg().getType() == getType())
     return getArg();
   // bitcast(bitcast(x : T0, T1), T0) -> x
   if (auto prev = getArg().getDefiningOp<BitcastOp>())
     if (prev.getArg().getType() == getType())
       return prev.getArg();
   return {};
 }
 
+LogicalResult LLVM::BitcastOp::verify() {
+  auto resultType = extractVectorElementType(getResult().getType())
+                        .dyn_cast<LLVMPointerType>();
+  auto sourceType =
+      extractVectorElementType(getArg().getType()).dyn_cast<LLVMPointerType>();
+
+  // If one of the types is a pointer (or vector of pointers), then
+  // both source and result type have to be pointers.
+  if (static_cast<bool>(resultType) != static_cast<bool>(sourceType))
+    return emitOpError("can only cast pointers from and to pointers");
+
+  if (!resultType)
+    return success();
+
+  auto isVector = [](Type type) {
+    return type.isa<VectorType, LLVMScalableVectorType, LLVMFixedVectorType>();
+  };
+
+  // Due to bitcast requiring both operands to be of the same size, it is not
+  // possible for only one of the two to be a pointer of vectors.
+  if (isVector(getResult().getType()) && !isVector(getArg().getType()))
+    return emitOpError("cannot cast pointer to vector of pointers");
+
+  if (!isVector(getResult().getType()) && isVector(getArg().getType()))
+    return emitOpError("cannot cast vector of pointers to pointer");
+
+  // Bitcast cannot cast between pointers of different address spaces.
+  // 'llvm.addrspacecast' must be used for this purpose instead.
+  if (resultType.getAddressSpace() != sourceType.getAddressSpace())
+    return emitOpError("cannot cast pointers of different address spaces, "
+                       "use 'llvm.addrspacecast' instead");
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Folder for LLVM::AddrSpaceCastOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult LLVM::AddrSpaceCastOp::fold(FoldAdaptor adaptor) {
   // addrcast(x : T0, T0) -> x
   if (getArg().getType() == getType())
     return getArg();

mlir/test/Dialect/LLVMIR/invalid.mlir

func.func @extract_scalable_from_fixed_length_vector(%arg0 : vector<16xf32>) {
   %0 = llvm.intr.vector.extract %arg0[0] : vector<[8]xf32> from vector<16xf32>
 }
 
 // -----
 
 #void = #llvm.di_void_result_type
 // expected-error@below {{expected subroutine to have non-void argument types}}
 #void_argument_type = #llvm.di_subroutine_type<types = #void, #void>
+
+// -----
+
+func.func @invalid_bitcast_ptr_to_i64(%arg : !llvm.ptr) {
+  // expected-error@+1 {{can only cast pointers from and to pointers}}
+  %1 = llvm.bitcast %arg : !llvm.ptr to i64
+}
+
+// -----
+
+func.func @invalid_bitcast_i64_to_ptr() {
+  %0 = llvm.mlir.constant(2 : i64) : i64
+  // expected-error@+1 {{can only cast pointers from and to pointers}}
+  %1 = llvm.bitcast %0 : i64 to !llvm.ptr
+}
+
+// -----
+
+func.func @invalid_bitcast_vec_to_ptr(%arg : !llvm.vec<4 x ptr>) {
+  // expected-error@+1 {{cannot cast vector of pointers to pointer}}
+  %0 = llvm.bitcast %arg : !llvm.vec<4 x ptr> to !llvm.ptr
+}
+
+// -----
+
+func.func @invalid_bitcast_ptr_to_vec(%arg : !llvm.ptr) {
+  // expected-error@+1 {{cannot cast pointer to vector of pointers}}
+  %0 = llvm.bitcast %arg : !llvm.ptr to !llvm.vec<4 x ptr>
+}
+
+// -----
+
+func.func @invalid_bitcast_addr_cast(%arg : !llvm.ptr<1>) {
+  // expected-error@+1 {{cannot cast pointers of different address spaces, use 'llvm.addrspacecast' instead}}
+  %0 = llvm.bitcast %arg : !llvm.ptr<1> to !llvm.ptr
+}
+
+// -----
+
+func.func @invalid_bitcast_addr_cast_vec(%arg : !llvm.vec<4 x ptr<1>>) {
+  // expected-error@+1 {{cannot cast pointers of different address spaces, use 'llvm.addrspacecast' instead}}
+  %0 = llvm.bitcast %arg : !llvm.vec<4 x ptr<1>> to !llvm.vec<4 x ptr>
+}