[SelectionDAGBuilder] Use address width when lowering ptrtoaddr

[arichardson]

Instead of just deferring to ptrtoint, we should truncate to the index
width and then perform the ZextOrTrunc.

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

@llvm/pr-subscribers-llvm-selectiondag

Author: Alexander Richardson (arichardson)

Changes

Instead of just deferring to ptrtoint, we should truncate to the index
width and then perform the ZextOrTrunc.

---

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

2 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp (+14-1)
• (modified) llvm/test/CodeGen/AMDGPU/ptrtoint-ptrtoaddr-p8.ll (+5-1)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index e6651d000bd71..806bab5379bde 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -3878,7 +3878,20 @@ void SelectionDAGBuilder::visitSIToFP(const User &I) {
 }
 
 void SelectionDAGBuilder::visitPtrToAddr(const User &I) {
-  visitPtrToInt(I);
+  const auto &TLI = DAG.getTargetLoweringInfo();
+  const DataLayout &DL = DAG.getDataLayout();
+  LLVMContext &Ctx = *DAG.getContext();
+  // ptrtoaddr is equivalent to a truncate of ptrtoint to address/index width
+  SDValue N = getValue(I.getOperand(0));
+  Type *PtrTy = I.getOperand(0)->getType();
+  EVT AddrVT = EVT::getIntegerVT(Ctx, DL.getPointerAddressSizeInBits(PtrTy));
+  if (auto *VTy = dyn_cast<VectorType>(PtrTy)) {
+    Type *EltTy = VTy->getElementType();
+    AddrVT = EVT::getVectorVT(Ctx, AddrVT, VTy->getElementCount());
+  }
+  N = DAG.getPtrExtOrTrunc(N, getCurSDLoc(), AddrVT);
+  N = DAG.getZExtOrTrunc(N, getCurSDLoc(), TLI.getValueType(DL, I.getType()));
+  setValue(&I, N);
 }
 
 void SelectionDAGBuilder::visitPtrToInt(const User &I) {
diff --git a/llvm/test/CodeGen/AMDGPU/ptrtoint-ptrtoaddr-p8.ll b/llvm/test/CodeGen/AMDGPU/ptrtoint-ptrtoaddr-p8.ll
index 32b5d9441b61c..da4b531ab5b25 100644
--- a/llvm/test/CodeGen/AMDGPU/ptrtoint-ptrtoaddr-p8.ll
+++ b/llvm/test/CodeGen/AMDGPU/ptrtoint-ptrtoaddr-p8.ll
@@ -32,8 +32,9 @@ define <2 x i64> @ptrtoaddr_vec(<2 x ptr addrspace(8)> %ptr) {
 ; CHECK-LABEL: ptrtoaddr_vec:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
-; CHECK-NEXT:    v_mov_b32_e32 v3, v5
 ; CHECK-NEXT:    v_mov_b32_e32 v2, v4
+; CHECK-NEXT:    v_and_b32_e32 v1, 0xffff, v1
+; CHECK-NEXT:    v_and_b32_e32 v3, 0xffff, v5
 ; CHECK-NEXT:    s_setpc_b64 s[30:31]
   %ret = ptrtoaddr <2 x ptr addrspace(8)> %ptr to <2 x i64>
   ret <2 x i64> %ret
@@ -57,6 +58,9 @@ define i128 @ptrtoaddr_ext(ptr addrspace(8) %ptr) {
 ; CHECK-LABEL: ptrtoaddr_ext:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; CHECK-NEXT:    v_and_b32_e32 v1, 0xffff, v1
+; CHECK-NEXT:    v_mov_b32_e32 v2, 0
+; CHECK-NEXT:    v_mov_b32_e32 v3, 0
 ; CHECK-NEXT:    s_setpc_b64 s[30:31]
   %ret = ptrtoaddr ptr addrspace(8) %ptr to i128
   ret i128 %ret