[llvm][SelectionDAG] Relax llvm.ptrmask's size check on arm64_32

[jroelofs]

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

[llvmbot]

@llvm/pr-subscribers-llvm-selectiondag

@llvm/pr-subscribers-backend-aarch64

Author: Jon Roelofs (jroelofs)

Changes

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

---

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

2 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp (+16-2)
• (added) llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll (+13)

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

[nikic]

This should be checking the pointer size, not the index size. The implementation below is only correct as long as these are the same. (The index size and the mask size are known to be the same by verifier enforcement.)

[nikic]

Drop the -stop-after=finalize-isel and use update_llc_test_checks.

[jroelofs]

Not sure why, but UTC doesn't produce meaningful CHECK lines on this one. Am I holding it wrong?

$ cat /Users/jonathan_roelofs/llvm-upstream/llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll
; RUN: llc -mtriple=arm64_32-apple-watchos %s -o - | FileCheck %s

define ptr @issue94075(ptr %p) {
entry:
  %rdar125263567 = call ptr @llvm.ptrmask.p0.i32(ptr %p, i32 4294967288)
  ret ptr %rdar125263567
}
$ ../llvm/utils/update_llc_test_checks.py --llc-binary=./bin/llc /Users/jonathan_roelofs/llvm-upstream/llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll -v
Scanning for RUN lines in test file: /Users/jonathan_roelofs/llvm-upstream/llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll
Found 1 RUN lines in /Users/jonathan_roelofs/llvm-upstream/llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll:
  RUN: llc -mtriple=arm64_32-apple-watchos %s -o - | FileCheck %s
Extracted LLC cmd: llc -mtriple=arm64_32-apple-watchos  -o -
Extracted FileCheck prefixes: ['CHECK']
Rewriting FileCheck prefixes: {'CHECK'}
Writing 10 lines to /Users/jonathan_roelofs/llvm-upstream/llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll...
$ cat /Users/jonathan_roelofs/llvm-upstream/llvm/test/CodeGen/AArch64/lower-ptrmask-arm64_32.ll
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
; RUN: llc -mtriple=arm64_32-apple-watchos %s -o - | FileCheck %s

define ptr @issue94075(ptr %p) {
entry:
  %rdar125263567 = call ptr @llvm.ptrmask.p0.i32(ptr %p, i32 4294967288)
  ret ptr %rdar125263567
}
;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
; CHECK: {{.*}}

[nikic]

Probably it doesn't support the arm64_32 triple right now.

[jroelofs]

ohh:

def get_run_handler(triple):
    target_handlers = {
        "i686": (scrub_asm_x86, ASM_FUNCTION_X86_RE),
        "x86": (scrub_asm_x86, ASM_FUNCTION_X86_RE),
        "i386": (scrub_asm_x86, ASM_FUNCTION_X86_RE),
        "arm64_32-apple-ios": (scrub_asm_arm_eabi, ASM_FUNCTION_AARCH64_DARWIN_RE),
        "arm64_32-apple-watchos2.0.0": (
            scrub_asm_arm_eabi,
            ASM_FUNCTION_AARCH64_DARWIN_RE,
        ),

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[nikic]

LGTM

[nikic]

Probably it doesn't support the arm64_32 triple right now.