[llvm] [llvm][SelectionDAG] Relax llvm.ptrmask's size check on arm64_32 (PR #94125)

[Jon Roelofs via llvm-commits]

https://github.com/jroelofs created https://github.com/llvm/llvm-project/pull/94125

Since pointers in memory, as well as the index type are both 32 bits, but in registers pointers are 64 bits, the mask generated by llvm.ptrmask needs to be zero-extended.

Fixes: #94075
Fixes: rdar://125263567

>From 1eff480ed1d00d62d71cdb9d8564103d3bf5f55b Mon Sep 17 00:00:00 2001
From: Jon Roelofs <jonathan_roelofs at apple.com>
Date: Sat, 1 Jun 2024 14:20:51 -0700
Subject: [PATCH] [llvm][SelectionDAG] Relax llvm.ptrmask's size check on
 arm64_32

Since pointers in memory, as well as the index type are both 32 bits, but in
registers pointers are 64 bits, the mask generated by llvm.ptrmask needs to be
zero-extended.

Fixes: #94075
Fixes: rdar://125263567
---
 .../SelectionDAG/SelectionDAGBuilder.cpp       | 18 ++++++++++++++++--
 .../CodeGen/AArch64/lower-ptrmask-arm64_32.ll  | 13 +++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 85e4cc3b82e6e..7ab02409fb06d 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -7839,12 +7839,26 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
     return;
   }
   case Intrinsic::ptrmask: {
+    unsigned PtrBits =
+        DAG.getDataLayout().getIndexTypeSizeInBits(I.getOperand(0)->getType());
+    unsigned MaskBits = I.getOperand(1)->getType()->getScalarSizeInBits();
+    (void)MaskBits;
+    assert(PtrBits == MaskBits &&
+           "llvm.ptrmask intrinsic second argument bitwidth must match pointer "
+           "index type size of first argument");
+
     SDValue Ptr = getValue(I.getOperand(0));
     SDValue Mask = getValue(I.getOperand(1));
 
     EVT PtrVT = Ptr.getValueType();
-    assert(PtrVT == Mask.getValueType() &&
-           "Pointers with different index type are not supported by SDAG");
+
+    // On arm64_32, pointers are 32 bits when stored in memory, but
+    // zero-extended to 64 bits when in registers.  Thus the index type is 32
+    // bits, but the mask here must be zero-extended up to 64 bits to match the
+    // pointer.
+    if (PtrBits < Ptr.getValueType().getFixedSizeInBits())
+      Mask = DAG.getZExtOrTrunc(Mask, sdl, PtrVT);
+
     setValue(&I, DAG.getNode(ISD::AND, sdl, PtrVT, Ptr, Mask));
     return;
   }
diff --git a/llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll b/llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll
new file mode 100644
index 0000000000000..aceb7c38e24b2
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll
@@ -0,0 +1,13 @@
+; RUN: llc -mtriple=arm64_32-apple-watchos -stop-after=finalize-isel %s -o - | FileCheck %s
+
+define ptr @issue94075(ptr %p) {
+entry:
+  %rdar125263567 = call ptr @llvm.ptrmask.p0.i32(ptr %p, i32 4294967288)
+  ret ptr %rdar125263567
+}
+
+; CHECK-LABEL: name: issue94075
+; CHECK:         %0:gpr64 = COPY $x0
+; CHECK-NEXT:    %1:gpr64sp = ANDXri %0, 8028
+; CHECK-NEXT:    $x0 = COPY %1
+; CHECK-NEXT:    RET_ReallyLR implicit $x0
\ No newline at end of file