[llvm] [InferAddressSpaces] Fix null PHI operand when an incoming flat expression is not rewritten (PR #205749)

[Fujun Han via llvm-commits]

================
@@ -1550,9 +1557,19 @@ bool InferAddressSpacesImpl::rewriteWithNewAddressSpaces(
 
     unsigned OperandNo = PoisonUse->getOperandNo();
     assert(isa<PoisonValue>(NewV->getOperand(OperandNo)));
-    WeakTrackingVH NewOp = ValueWithNewAddrSpace.lookup(PoisonUse->get());
-    assert(NewOp &&
-           "poison replacements in ValueWithNewAddrSpace shouldn't be null");
+    Value *NewOp = ValueWithNewAddrSpace.lookup(PoisonUse->get());
+    if (!NewOp) {
+      // The operand was never rewritten into the new address space (e.g. a flat
----------------
Peter9606 wrote:

Thanks for the detailed review. Let me add more context on where this comes from.

We first hit this crash in our own in-house compiler (its implementation is very similar to AMDGPU). The failure showd up when `infer-address-spaces` ran on IR that had already gone through earlier address-space-related lowering - not when running `infer-address-spaces` in isolation on hand-written IR.
On current upstream AMDGPU, I also cannot reproduce it with infer-address-sapces alone. The trigger I have locally is:
```bash
opt -S test.ll -passes='amdgpu-promote-kernel-arguments,infer-address-spaces'   
```
and the content of ```test.ll``` is as following:
```llvm
target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7"
target triple = "amdgcn-amd-amdhsa"

%"struct.fbgemm_gpu::Vec4T" = type { %"struct.fbgemm_gpu::Vec4BaseT" }
%"struct.fbgemm_gpu::Vec4BaseT" = type { %struct.float4 }
%struct.float4 = type { float, float, float, float }

define amdgpu_kernel void @test_kernel(ptr addrspace(0) %0) {
  br i1 false, label %4, label %2

2:
  %3 = getelementptr %"struct.fbgemm_gpu::Vec4T", ptr addrspace(0) null, i64 0
  br label %4

4:
  %5 = phi ptr addrspace(0) [ %3, %2 ], [ %0, %1 ]
  %6 = getelementptr i8, ptr addrspace(0) %5, i64 4
  ret void
}
```
Running `amdgpu-promote-kernel-arguments` alone on this IR shows that the pass does not create the problematic flat expression (%3 = gep ..., null, ... is already present in the input). It only inserts `flat -> global -> flat` addrspacecast pairs around the flat kernel argument and rewires the PHI's entry incoming from `%0` to `%.flat`:
```bash
opt -S test.ll -passes='amdgpu-promote-kernel-arguments' -S
```
```llvm
define amdgpu_kernel void @test_kernel(ptr %0) {
entry:
  %.global = addrspacecast ptr %0 to ptr addrspace(1)
  %.flat = addrspacecast ptr addrspace(1) %.global to ptr
  br i1 false, label %3, label %1

1:
  %2 = getelementptr %"struct.fbgemm_gpu::Vec4T", ptr null, i64 0
  br label %3

3:
  %4 = phi ptr [ %2, %1 ], [ %.flat, %entry ]
  %5 = getelementptr i8, ptr %4, i64 4
  ret void
}
```
The crash then happens on the next pass, when infer-address-spaces tries to rewrite merge PHI `%4` into `addrspace(1)` because the entry incoming value (`%.flat`, seeded by the `%.global`  cast inserted around `%0`) propagates global address space information, while the other incoming value `%2` remains an un-rewritten flat expression (a `getelementptr` based on a null pointer). That is the situation we see in our production.

The lit tests in this PR are a minimized version of that IR pattern so we can test the pass directly without pulling in the whole AMDGPU pipeline. They are meant to exercise the same rewrite/poison-fixup failure, not to claim this is the only way the bug can appear in the wild.

https://github.com/llvm/llvm-project/pull/205749