[InferAddressSpaces] Extend undef pointer operand support to phi inst (#159548)

Previously undef pointer operand is only supported for select inst,
where undef in generic AS behaves like `take the other side`.

This PR extends the support to other instructions, e.g. phi inst. Defer
joining and inferring constant pointer operand until all other operand
AS states considered.

---------

Co-authored-by: Matt Arsenault <[email protected]>

Author: wenju-he

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

@@ -1002,70 +1002,49 @@ bool InferAddressSpacesImpl::updateAddressSpace(
   // isAddressExpression should guarantee that V is an operator or an argument.
   assert(isa<Operator>(V) || isa<Argument>(V));
 
-  if (isa<Operator>(V) &&
-      cast<Operator>(V).getOpcode() == Instruction::Select) {
-    const Operator &Op = cast<Operator>(V);
-    Value *Src0 = Op.getOperand(1);
-    Value *Src1 = Op.getOperand(2);
-
-    auto I = InferredAddrSpace.find(Src0);
-    unsigned Src0AS = (I != InferredAddrSpace.end())
-                          ? I->second
-                          : Src0->getType()->getPointerAddressSpace();
-
-    auto J = InferredAddrSpace.find(Src1);
-    unsigned Src1AS = (J != InferredAddrSpace.end())
-                          ? J->second
-                          : Src1->getType()->getPointerAddressSpace();
-
-    auto *C0 = dyn_cast<Constant>(Src0);
-    auto *C1 = dyn_cast<Constant>(Src1);
-
-    // If one of the inputs is a constant, we may be able to do a constant
-    // addrspacecast of it. Defer inferring the address space until the input
-    // address space is known.
-    if ((C1 && Src0AS == UninitializedAddressSpace) ||
-        (C0 && Src1AS == UninitializedAddressSpace))
-      return false;
-
-    if (C0 && isSafeToCastConstAddrSpace(C0, Src1AS))
-      NewAS = Src1AS;
-    else if (C1 && isSafeToCastConstAddrSpace(C1, Src0AS))
-      NewAS = Src0AS;
-    else
-      NewAS = joinAddressSpaces(Src0AS, Src1AS);
+  unsigned AS = TTI->getAssumedAddrSpace(&V);
+  if (AS != UninitializedAddressSpace) {
+    // Use the assumed address space directly.
+    NewAS = AS;
   } else {
-    unsigned AS = TTI->getAssumedAddrSpace(&V);
-    if (AS != UninitializedAddressSpace) {
-      // Use the assumed address space directly.
-      NewAS = AS;
-    } else {
-      // Otherwise, infer the address space from its pointer operands.
-      for (Value *PtrOperand : getPointerOperands(V, *DL, TTI)) {
-        auto I = InferredAddrSpace.find(PtrOperand);
-        unsigned OperandAS;
-        if (I == InferredAddrSpace.end()) {
-          OperandAS = PtrOperand->getType()->getPointerAddressSpace();
-          if (OperandAS == FlatAddrSpace) {
-            // Check AC for assumption dominating V.
-            unsigned AS = getPredicatedAddrSpace(*PtrOperand, &V);
-            if (AS != UninitializedAddressSpace) {
-              LLVM_DEBUG(dbgs()
-                         << "  deduce operand AS from the predicate addrspace "
-                         << AS << '\n');
-              OperandAS = AS;
-              // Record this use with the predicated AS.
-              PredicatedAS[std::make_pair(&V, PtrOperand)] = OperandAS;
-            }
+    // Otherwise, infer the address space from its pointer operands.
+    SmallVector<Constant *, 2> ConstantPtrOps;
+    for (Value *PtrOperand : getPointerOperands(V, *DL, TTI)) {
+      auto I = InferredAddrSpace.find(PtrOperand);
+      unsigned OperandAS;
+      if (I == InferredAddrSpace.end()) {
+        OperandAS = PtrOperand->getType()->getPointerAddressSpace();
+        if (auto *C = dyn_cast<Constant>(PtrOperand);
+            C && OperandAS == FlatAddrSpace) {
+          // Defer joining the address space of constant pointer operands.
+          ConstantPtrOps.push_back(C);
+          continue;
+        }
+        if (OperandAS == FlatAddrSpace) {
+          // Check AC for assumption dominating V.
+          unsigned AS = getPredicatedAddrSpace(*PtrOperand, &V);
+          if (AS != UninitializedAddressSpace) {
+            LLVM_DEBUG(dbgs()
+                       << "  deduce operand AS from the predicate addrspace "
+                       << AS << '\n');
+            OperandAS = AS;
+            // Record this use with the predicated AS.
+            PredicatedAS[std::make_pair(&V, PtrOperand)] = OperandAS;
           }
-        } else
-          OperandAS = I->second;
+        }
+      } else
+        OperandAS = I->second;
 
-        // join(flat, *) = flat. So we can break if NewAS is already flat.
-        NewAS = joinAddressSpaces(NewAS, OperandAS);
-        if (NewAS == FlatAddrSpace)
-          break;
-      }
+      // join(flat, *) = flat. So we can break if NewAS is already flat.
+      NewAS = joinAddressSpaces(NewAS, OperandAS);
+      if (NewAS == FlatAddrSpace)
+        break;
+    }
+    if (NewAS != FlatAddrSpace && NewAS != UninitializedAddressSpace) {
+      if (any_of(ConstantPtrOps, [=](Constant *C) {
+            return !isSafeToCastConstAddrSpace(C, NewAS);
+          }))
+        NewAS = FlatAddrSpace;
     }
   }
 

llvm/test/Transforms/InferAddressSpaces/AMDGPU/phi-poison.ll

@@ -0,0 +1,36 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -S -passes=infer-address-spaces %s | FileCheck %s
+
+define void @phi_poison(ptr addrspace(1) %arg, <2 x ptr addrspace(1)> %arg1) {
+; CHECK-LABEL: @phi_poison(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i1 @is_leader()
+; CHECK-NEXT:    br i1 [[TMP0]], label [[LEADER:%.*]], label [[MERGE:%.*]]
+; CHECK:       leader:
+; CHECK-NEXT:    br label [[MERGE]]
+; CHECK:       merge:
+; CHECK-NEXT:    [[I:%.*]] = phi ptr addrspace(1) [ [[ARG:%.*]], [[LEADER]] ], [ poison, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[I2:%.*]] = phi <2 x ptr addrspace(1)> [ [[ARG1:%.*]], [[LEADER]] ], [ poison, [[ENTRY]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast <2 x ptr addrspace(1)> [[I2]] to <2 x ptr>
+; CHECK-NEXT:    [[J:%.*]] = icmp eq ptr addrspace(1) [[I]], addrspacecast (ptr null to ptr addrspace(1))
+; CHECK-NEXT:    [[J1:%.*]] = icmp eq <2 x ptr> [[TMP1]], zeroinitializer
+; CHECK-NEXT:    ret void
+;
+entry:
+  %is = call i1 @is_leader()
+  br i1 %is, label %leader, label %merge
+
+leader:
+  %cast = addrspacecast ptr addrspace(1) %arg to ptr
+  %cast1 = addrspacecast <2 x ptr addrspace(1)> %arg1 to <2 x ptr>
+  br label %merge
+
+merge:
+  %i = phi ptr [%cast, %leader], [poison, %entry]
+  %i1 = phi <2 x ptr> [%cast1, %leader], [poison, %entry]
+  %j = icmp eq ptr %i, null
+  %j1 = icmp eq <2 x ptr> %i1, zeroinitializer
+  ret void
+}
+
+declare i1 @is_leader()