[llvm] [InferAS] infer noop inttoptr/ptrtoint in constant expression (PR #177809)

[Luo Yuanke via llvm-commits]

https://github.com/LuoYuanke updated https://github.com/llvm/llvm-project/pull/177809

>From 0c4aa2295c83b39a8adbe770461c86948fe77e75 Mon Sep 17 00:00:00 2001
From: Yuanke Luo <ykluo at birentech.com>
Date: Sun, 25 Jan 2026 08:42:16 +0800
Subject: [PATCH] [InferAS] infer noop inttoptr/ptrtoint in constant expression

The algorithm of infer address space is to update the address space
of a pointer and insert it's user to the work list, so that the
address space can be propagated to the pointer's user. However
inttoptr/ptrtoint is not associated with directly def/use, so the
address space inferred in ptrtoint may not be propagated to inttoptr.

This patch is to precede the ptrtoint before inttoptr in advance,
so that address space can be propagated from ptrtoint to inttoptr.
---
 .../Transforms/Scalar/InferAddressSpaces.cpp  | 27 ++++++++++++---
 .../NVPTX/inttoptr_constexpr.ll               | 34 +++++++++++++++++++
 2 files changed, 56 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Transforms/InferAddressSpaces/NVPTX/inttoptr_constexpr.ll

diff --git a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
index 03efc156df1e8..ff070c1e0de5e 100644
--- a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
+++ b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
@@ -521,12 +521,27 @@ void InferAddressSpacesImpl::appendsFlatAddressExpressionToPostorderStack(
     DenseSet<Value *> &Visited) const {
   assert(V->getType()->isPtrOrPtrVectorTy());
 
+  auto EmplaceBack = [&](Value *Val) {
+    // The inttoptr can't get the pointer operand of ptrtoint from its user
+    // operand. When the address space of ptrtoint changes, it won't insert
+    // inttoptr into WorkList during updating address space. We push its pointer
+    // operand ahead of inttoptr to make sure the pointer operand's address
+    // space can be updated first.
+    Operator *AddrExpr = dyn_cast<Operator>(Val);
+    if (AddrExpr && AddrExpr->getOpcode() == Instruction::IntToPtr) {
+      for (Value *PtrOperand : getPointerOperands(*AddrExpr, *DL, TTI)) {
+        appendsFlatAddressExpressionToPostorderStack(PtrOperand, PostorderStack,
+                                                     Visited);
+      }
+    }
+    PostorderStack.emplace_back(Val, false);
+  };
   // Generic addressing expressions may be hidden in nested constant
   // expressions.
   if (ConstantExpr *CE = dyn_cast<ConstantExpr>(V)) {
     // TODO: Look in non-address parts, like icmp operands.
     if (isAddressExpression(*CE, *DL, TTI) && Visited.insert(CE).second)
-      PostorderStack.emplace_back(CE, false);
+      EmplaceBack(CE);
 
     return;
   }
@@ -534,13 +549,13 @@ void InferAddressSpacesImpl::appendsFlatAddressExpressionToPostorderStack(
   if (V->getType()->getPointerAddressSpace() == FlatAddrSpace &&
       isAddressExpression(*V, *DL, TTI)) {
     if (Visited.insert(V).second) {
-      PostorderStack.emplace_back(V, false);
+      EmplaceBack(V);
 
       if (auto *Op = dyn_cast<Operator>(V))
         for (auto &O : Op->operands())
           if (ConstantExpr *CE = dyn_cast<ConstantExpr>(O))
             if (isAddressExpression(*CE, *DL, TTI) && Visited.insert(CE).second)
-              PostorderStack.emplace_back(CE, false);
+              EmplaceBack(CE);
     }
   }
 }
@@ -607,6 +622,7 @@ InferAddressSpacesImpl::collectFlatAddressExpressions(Function &F) const {
     if (PostorderStack.back().getInt()) {
       if (TopVal->getType()->getPointerAddressSpace() == FlatAddrSpace)
         Postorder.push_back(TopVal);
+
       PostorderStack.pop_back();
       continue;
     }
@@ -883,8 +899,9 @@ static Value *cloneConstantExprWithNewAddressSpace(
   if (CE->getOpcode() == Instruction::IntToPtr) {
     assert(isNoopPtrIntCastPair(cast<Operator>(CE), *DL, TTI));
     Constant *Src = cast<ConstantExpr>(CE->getOperand(0))->getOperand(0);
-    assert(Src->getType()->getPointerAddressSpace() == NewAddrSpace);
-    return Src;
+    if (Src->getType() == TargetType)
+      return Src;
+    return ConstantExpr::getAddrSpaceCast(Src, TargetType);
   }
 
   // Computes the operands of the new constant expression.
diff --git a/llvm/test/Transforms/InferAddressSpaces/NVPTX/inttoptr_constexpr.ll b/llvm/test/Transforms/InferAddressSpaces/NVPTX/inttoptr_constexpr.ll
new file mode 100644
index 0000000000000..5275b777a550c
--- /dev/null
+++ b/llvm/test/Transforms/InferAddressSpaces/NVPTX/inttoptr_constexpr.ll
@@ -0,0 +1,34 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -mtriple=nvptx64-nvidia-cuda -passes=infer-address-spaces %s | FileCheck %s
+
+ at g = addrspace(1) global i32 0, align 4
+
+define void @test() {
+; CHECK-LABEL: define void @test() {
+; CHECK-NEXT:    store i32 0, ptr inttoptr (i64 xor (i64 ptrtoint (ptr addrspacecast (ptr addrspace(1) @g to ptr) to i64), i64 7) to ptr), align 4
+; CHECK-NEXT:    ret void
+;
+  store i32 0, ptr inttoptr (i64
+    xor (i64
+      ptrtoint (ptr
+        addrspacecast (ptr addrspace(1) @g to ptr)
+      to i64),
+    i64 7)
+  to ptr)
+
+  ret void
+}
+
+define void @test1() {
+; CHECK-LABEL: define void @test1() {
+; CHECK-NEXT:    store i32 0, ptr addrspace(1) @g, align 4
+; CHECK-NEXT:    ret void
+;
+  store i32 0, ptr inttoptr (i64
+    ptrtoint (ptr
+      addrspacecast (ptr addrspace(1) @g to ptr)
+      to i64)
+    to ptr)
+
+  ret void
+}