[IR] Allow llvm.ptrmask of vectors

[nikic]

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).

[llvmbot]

@llvm/pr-subscribers-llvm-ir
@llvm/pr-subscribers-backend-x86

@llvm/pr-subscribers-llvm-transforms

Changes

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).

---

Full diff: https://github.com/llvm/llvm-project/pull/67434.diff

5 Files Affected:

• (modified) llvm/docs/LangRef.rst (+2-1)
• (modified) llvm/include/llvm/IR/Intrinsics.td (+3-1)
• (modified) llvm/lib/IR/Verifier.cpp (+19)
• (modified) llvm/test/CodeGen/X86/lower-ptrmask.ll (+13)
• (modified) llvm/test/Transforms/InstCombine/consecutive-ptrmask.ll (+14)

diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index f542e70bcfee810..a0c1a43b096e02d 100644
--- a/llvm/docs/LangRef.rst
+++ b/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:
 """"""""""
diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
index e94b59508de7b5e..ab15b1f1e0ee888 100644
--- a/llvm/include/llvm/IR/Intrinsics.td
+++ b/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.
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 5dac691e17cd6ef..5a3328416db3eb0 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/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.
diff --git a/llvm/test/CodeGen/X86/lower-ptrmask.ll b/llvm/test/CodeGen/X86/lower-ptrmask.ll
index 8396ae5cb01eea5..185564e5a07ae5c 100644
--- a/llvm/test/CodeGen/X86/lower-ptrmask.ll
+++ b/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
+}
diff --git a/llvm/test/Transforms/InstCombine/consecutive-ptrmask.ll b/llvm/test/Transforms/InstCombine/consecutive-ptrmask.ll
index 904c758b99306f4..d2da3be3201cfc0 100644
--- a/llvm/test/Transforms/InstCombine/consecutive-ptrmask.ll
+++ b/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
+}

[goldsteinn]

Will this become an issue: https://reviews.llvm.org/D153773#inline-1486422
?

[nikic]

Will this become an issue: https://reviews.llvm.org/D153773#inline-1486422 ?

No, the final implementation used getScalarSizeInBits():

llvm-project/llvm/lib/Analysis/ValueTracking.cpp

Line 1636 in 50abfc4

Known2 = KnownBits(Mask->getType()->getScalarSizeInBits());

[fhahn]

LGTM, thanks!

I don't think there was any particular reason for limiting this to pointers only.