[IR] Allow llvm.ptrmask of vectors

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
+}