[IR] Allow llvm.ptrmask of vectors (#67434)

llvm.ptrmask is currently limited to pointers only, and does not accept
vectors of pointers. This is an unnecessary limitation, especially as
the underlying instructions (getelementptr etc) do support vectors of
pointers.

We should relax this sooner rather than later, to avoid introducing code
that assumes non-vectors (#67166).

Author: nikic

llvm/docs/LangRef.rst

@@ -26865,7 +26865,8 @@ Syntax:
 Arguments:
 """"""""""
 
-The first argument is a pointer. The second argument is an integer.
+The first argument is a pointer or vector of pointers. The second argument is
+an integer or vector of integers.
 
 Overview:
 """"""""""
@@ -26880,7 +26881,7 @@ Semantics:
 
 The result of ``ptrmask(ptr, mask)`` is equivalent to
 ``getelementptr ptr, (ptrtoint(ptr) & mask) - ptrtoint(ptr)``. Both the returned
-pointer and the first argument are based on the same underlying object (for more
+pointer(s) and the first argument are based on the same underlying object (for more
 information on the *based on* terminology see
 :ref:`the pointer aliasing rules <pointeraliasing>`). If the bitwidth of the
 mask argument does not match the pointer size of the target, the mask is

llvm/include/llvm/IR/Intrinsics.td

@@ -1743,7 +1743,9 @@ def int_is_constant : DefaultAttrsIntrinsic<[llvm_i1_ty], [llvm_any_ty],
                                 "llvm.is.constant">;
 
 // Intrinsic to mask out bits of a pointer.
-def int_ptrmask: DefaultAttrsIntrinsic<[llvm_anyptr_ty], [LLVMMatchType<0>, llvm_anyint_ty],
+// First argument must be pointer or vector of pointer. This is checked by the
+// verifier.
+def int_ptrmask: DefaultAttrsIntrinsic<[llvm_any_ty], [LLVMMatchType<0>, llvm_anyint_ty],
                            [IntrNoMem, IntrSpeculatable, IntrWillReturn]>;
 
 // Intrinsic to wrap a thread local variable.

llvm/lib/IR/Verifier.cpp

@@ -5947,6 +5947,25 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
     break;
   case Intrinsic::experimental_convergence_loop:
     break;
+  case Intrinsic::ptrmask: {
+    Type *Ty0 = Call.getArgOperand(0)->getType();
+    Type *Ty1 = Call.getArgOperand(1)->getType();
+    Check(Ty0->isPtrOrPtrVectorTy(),
+          "llvm.ptrmask intrinsic first argument must be pointer or vector "
+          "of pointers",
+          &Call);
+    Check(
+        Ty0->isVectorTy() == Ty1->isVectorTy(),
+        "llvm.ptrmask intrinsic arguments must be both scalars or both vectors",
+        &Call);
+    if (Ty0->isVectorTy())
+      Check(cast<VectorType>(Ty0)->getElementCount() ==
+                cast<VectorType>(Ty1)->getElementCount(),
+            "llvm.ptrmask intrinsic arguments must have the same number of "
+            "elements",
+            &Call);
+    break;
+  }
   };
 
   // Verify that there aren't any unmediated control transfers between funclets.

llvm/test/CodeGen/X86/lower-ptrmask.ll

@@ -29,3 +29,16 @@ define ptr @test2(ptr %src) {
   %ptr = call ptr @llvm.ptrmask.p0.i32(ptr %src, i32 10000)
   ret ptr %ptr
 }
+
+declare <2 x ptr> @llvm.ptrmask.v2p0.v2i64(<2 x ptr>, <2 x i64>)
+
+; CHECK-LABEL: name: test3
+; CHECK: %0:vr128 = COPY $xmm0
+; CHECK-NEXT: %1:vr128 = PANDrm %0, $rip, 1, $noreg, %const.0, $noreg :: (load (s128) from constant-pool)
+; CHECK-NEXT: $xmm0 = COPY %1
+; CHECK-NEXT: RET 0, $xmm0
+
+define <2 x ptr> @test3(<2 x ptr> %src) {
+  %ptr = call <2 x ptr> @llvm.ptrmask.v2p0.v2i64(<2 x ptr> %src, <2 x i64> <i64 10000, i64 10000>)
+  ret <2 x ptr> %ptr
+}

llvm/test/Transforms/InstCombine/consecutive-ptrmask.ll

@@ -3,7 +3,9 @@
 
 declare ptr @llvm.ptrmask.p0.i64(ptr, i64)
 declare ptr @llvm.ptrmask.p0.i32(ptr, i32)
+declare <2 x ptr> @llvm.ptrmask.v2p0.v2i64(<2 x ptr>, <2 x i64>)
 declare void @use.ptr(ptr)
+
 define ptr @fold_2x(ptr %p, i64 %m0, i64 %m1) {
 ; CHECK-LABEL: define ptr @fold_2x
 ; CHECK-SAME: (ptr [[P:%.*]], i64 [[M0:%.*]], i64 [[M1:%.*]]) {
@@ -65,3 +67,15 @@ define ptr @fold_2x_fail_type_mismatch2(ptr %p, i64 %m0, i32 %m1) {
   %p1 = call ptr @llvm.ptrmask.p0.i32(ptr %p0, i32 %m1)
   ret ptr %p1
 }
+
+define <2 x ptr> @fold_2x_vec(<2 x ptr> %p, <2 x i64> %m0, <2 x i64> %m1) {
+; CHECK-LABEL: define <2 x ptr> @fold_2x_vec
+; CHECK-SAME: (<2 x ptr> [[P:%.*]], <2 x i64> [[M0:%.*]], <2 x i64> [[M1:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i64> [[M1]], [[M0]]
+; CHECK-NEXT:    [[P1:%.*]] = call <2 x ptr> @llvm.ptrmask.v2p0.v2i64(<2 x ptr> [[P]], <2 x i64> [[TMP1]])
+; CHECK-NEXT:    ret <2 x ptr> [[P1]]
+;
+  %p0 = call <2 x ptr> @llvm.ptrmask.v2p0.v2i64(<2 x ptr> %p, <2 x i64> %m0)
+  %p1 = call <2 x ptr> @llvm.ptrmask.v2p0.v2i64(<2 x ptr> %p0, <2 x i64> %m1)
+  ret <2 x ptr> %p1
+}

llvm/test/Verifier/ptrmask.ll

@@ -0,0 +1,34 @@
+; RUN: not llvm-as < %s 2>&1 | FileCheck %s
+
+declare float @llvm.ptrmask.f32.i64(float, i64)
+declare ptr @llvm.ptrmask.p0.v4i64(ptr, <4 x i64>)
+declare <2 x ptr> @llvm.ptrmask.v2p0.i64(<2 x ptr>, i64)
+declare <2 x ptr> @llvm.ptrmask.v2p0.v4i64(<2 x ptr>, <4 x i64>)
+
+; CHECK: llvm.ptrmask intrinsic first argument must be pointer or vector of pointers
+; CHECK-NEXT:  %1 = call float @llvm.ptrmask.f32.i64(float 0.000000e+00, i64 0)
+define void @not_ptr() {
+  call float @llvm.ptrmask.f32.i64(float 0.0, i64 0)
+  ret void
+}
+
+; CHECK: llvm.ptrmask intrinsic arguments must be both scalars or both vectors
+; CHECK: %1 = call ptr @llvm.ptrmask.p0.v4i64(ptr null, <4 x i64> zeroinitializer)
+define void @scalar_vector_mismatch_1() {
+  call ptr @llvm.ptrmask.p0.v4i64(ptr null, <4 x i64> zeroinitializer)
+  ret void
+}
+
+; CHECK: llvm.ptrmask intrinsic arguments must be both scalars or both vectors
+; CHECK: %1 = call <2 x ptr> @llvm.ptrmask.v2p0.i64(<2 x ptr> zeroinitializer, i64 0)
+define void @scalar_vector_mismatch_2() {
+  call <2 x ptr> @llvm.ptrmask.v2p0.i64(<2 x ptr> zeroinitializer, i64 0)
+  ret void
+}
+
+; CHECK: llvm.ptrmask intrinsic arguments must have the same number of elements
+; CHECK: %1 = call <2 x ptr> @llvm.ptrmask.v2p0.v4i64(<2 x ptr> zeroinitializer, <4 x i64> zeroinitializer)
+define void @vector_size_mismatch() {
+  call <2 x ptr> @llvm.ptrmask.v2p0.v4i64(<2 x ptr> zeroinitializer, <4 x i64> zeroinitializer)
+  ret void
+}