[llvm] [InferAddressSpaces] Initialize op(uninit const, uninit const, ...) -> flat (PR #172143)

[Hongyu Chen via llvm-commits]

https://github.com/XChy created https://github.com/llvm/llvm-project/pull/172143

Fixes #171890
If the pointer operands of an instruction are all constants, we may infer the AS of the instruction as uninitialized finally. And the rewrite process will skip cloning the instruction, producing invalid IR.
This patch fixes it by inferring the AS of this kind of instruction as flat. Maybe we can fold the operator with all constants to get better performance, but I think this case is rare in the real world.

>From 4f50d25b0f307c19f7c6a9cc509a9607a6b13e7d Mon Sep 17 00:00:00 2001
From: XChy <xxs_chy at outlook.com>
Date: Sat, 13 Dec 2025 17:31:13 +0800
Subject: [PATCH] [InferAddressSpaces] Initialize op(uninit const, uninit
 const, ...) -> flat

---
 .../Transforms/Scalar/InferAddressSpaces.cpp  |  9 ++++-
 .../InferAddressSpaces/NVPTX/nullptr.ll       | 36 +++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/InferAddressSpaces/NVPTX/nullptr.ll

diff --git a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
index 352a1b331001a..d7caa60055750 100644
--- a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
+++ b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
@@ -1048,7 +1048,8 @@ bool InferAddressSpacesImpl::updateAddressSpace(
   } else {
     // Otherwise, infer the address space from its pointer operands.
     SmallVector<Constant *, 2> ConstantPtrOps;
-    for (Value *PtrOperand : getPointerOperands(V, *DL, TTI)) {
+    auto PtrOps = getPointerOperands(V, *DL, TTI);
+    for (Value *PtrOperand : PtrOps) {
       auto I = InferredAddrSpace.find(PtrOperand);
       unsigned OperandAS;
       if (I == InferredAddrSpace.end()) {
@@ -1079,12 +1080,18 @@ bool InferAddressSpacesImpl::updateAddressSpace(
       if (NewAS == FlatAddrSpace)
         break;
     }
+
     if (NewAS != FlatAddrSpace && NewAS != UninitializedAddressSpace) {
       if (any_of(ConstantPtrOps, [=](Constant *C) {
             return !isSafeToCastConstAddrSpace(C, NewAS);
           }))
         NewAS = FlatAddrSpace;
     }
+
+    // operator(uninit const, uninit const, ...) -> flat
+    if (NewAS == UninitializedAddressSpace &&
+        PtrOps.size() == ConstantPtrOps.size())
+      NewAS = FlatAddrSpace;
   }
 
   unsigned OldAS = InferredAddrSpace.lookup(&V);
diff --git a/llvm/test/Transforms/InferAddressSpaces/NVPTX/nullptr.ll b/llvm/test/Transforms/InferAddressSpaces/NVPTX/nullptr.ll
new file mode 100644
index 0000000000000..516d90d282f78
--- /dev/null
+++ b/llvm/test/Transforms/InferAddressSpaces/NVPTX/nullptr.ll
@@ -0,0 +1,36 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=infer-address-spaces %s | FileCheck %s
+
+target triple = "nvptx64-unknown-nvidiacl"
+
+define ptr @pr171890() {
+; CHECK-LABEL: define ptr @pr171890() {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[STACK:%.*]] = alloca i16, align 2
+; CHECK-NEXT:    [[TMP0:%.*]] = addrspacecast ptr [[STACK]] to ptr addrspace(5)
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr addrspace(5) [[TMP0]] to ptr
+; CHECK-NEXT:    br label %[[IF:.*]]
+; CHECK:       [[IF]]:
+; CHECK-NEXT:    [[NULLPHI:%.*]] = phi ptr [ null, %[[ENTRY]] ]
+; CHECK-NEXT:    br i1 false, label %[[THEN:.*]], label %[[ELSE:.*]]
+; CHECK:       [[ELSE]]:
+; CHECK-NEXT:    br label %[[THEN]]
+; CHECK:       [[THEN]]:
+; CHECK-NEXT:    [[PHI2:%.*]] = phi ptr [ [[TMP1]], %[[IF]] ], [ [[NULLPHI]], %[[ELSE]] ]
+; CHECK-NEXT:    ret ptr [[PHI2]]
+;
+entry:
+  %stack = alloca i16, align 2
+  br label %if
+
+if:
+  %nullphi = phi ptr [ null, %entry ]
+  br i1 false, label %then, label %else
+
+else:
+  br label %then
+
+then:
+  %phi2 = phi ptr [ %stack, %if ], [ %nullphi, %else ]
+  ret ptr %phi2
+}