[llvm] [InstCombine] Fix a crash in `PointerReplacer` (PR #98987)

[Shilei Tian via llvm-commits]

https://github.com/shiltian updated https://github.com/llvm/llvm-project/pull/98987

>From 19c6b683096201c968ea5f254d734aca388d1e39 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Mon, 15 Jul 2024 23:54:31 -0400
Subject: [PATCH] [InstCombine] Fix a crash in `PointerReplacer`

A crash could happen in `PointerReplacer::replace` when constructing a new
select instruction and there is no replacement for one of its operand. This can
happen when the operand is a load instruction that has been replaced earlier
such that the operand itself is already the new value. In this case, it is not
in the replacement map and `getReplacement` simply returns nullptr.
---
 .../InstCombineLoadStoreAlloca.cpp            | 11 +++++--
 .../InstCombine/AMDGPU/select-from-load.ll    | 33 +++++++++++++++++++
 2 files changed, 41 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Transforms/InstCombine/AMDGPU/select-from-load.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
index 21d5e1dece024..3a47fd2474d64 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
@@ -394,9 +394,14 @@ void PointerReplacer::replace(Instruction *I) {
     NewI->setNoWrapFlags(GEP->getNoWrapFlags());
     WorkMap[GEP] = NewI;
   } else if (auto *SI = dyn_cast<SelectInst>(I)) {
-    auto *NewSI = SelectInst::Create(
-        SI->getCondition(), getReplacement(SI->getTrueValue()),
-        getReplacement(SI->getFalseValue()), SI->getName(), nullptr, SI);
+    Value *TrueValue = SI->getTrueValue();
+    Value *FalseValue = SI->getFalseValue();
+    if (Value *Replace = getReplacement(TrueValue))
+      TrueValue = Replace;
+    if (Value *Replace = getReplacement(FalseValue))
+      FalseValue = Replace;
+    auto *NewSI = SelectInst::Create(SI->getCondition(), TrueValue, FalseValue,
+                                     SI->getName(), nullptr, SI);
     IC.InsertNewInstWith(NewSI, SI->getIterator());
     NewSI->takeName(SI);
     WorkMap[SI] = NewSI;
diff --git a/llvm/test/Transforms/InstCombine/AMDGPU/select-from-load.ll b/llvm/test/Transforms/InstCombine/AMDGPU/select-from-load.ll
new file mode 100644
index 0000000000000..d9af665e663f3
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/AMDGPU/select-from-load.ll
@@ -0,0 +1,33 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=instcombine -S -o - %s | FileCheck %s
+; REQUIRES: amdgpu-registered-target
+
+target triple = "amdgcn-amd-amdhsa"
+
+%anon = type { i32, [8 x ptr], ptr }
+
+define void @foo(ptr addrspace(4) byref(%anon) align 8 %0) {
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP1:%.*]] = load ptr, ptr addrspace(4) [[TMP0:%.*]], align 8
+; CHECK-NEXT:    br label [[FOR_COND10:%.*]]
+; CHECK:       for.cond10:
+; CHECK-NEXT:    [[TMP2:%.*]] = load i64, ptr [[TMP1]], align 8
+; CHECK-NEXT:    store i64 [[TMP2]], ptr addrspace(1) null, align 8
+; CHECK-NEXT:    br label [[FOR_COND10]]
+;
+entry:
+  %coerce = alloca %anon, addrspace(5)
+  call void @llvm.memcpy.p5.p4.i64(ptr addrspace(5) %coerce, ptr addrspace(4) %0, i64 0, i1 false)
+  %asc = addrspacecast ptr addrspace(5) %coerce to ptr
+  %load = load ptr, ptr addrspace(5) %coerce
+  %retval.0.i = select i1 false, ptr %asc, ptr %load
+  br label %for.cond10
+
+for.cond10:                                       ; preds = %for.cond10, %entry
+  %3 = load i64, ptr %retval.0.i
+  store i64 %3, ptr addrspace(1) null
+  br label %for.cond10
+}
+
+declare void @llvm.memcpy.p5.p4.i64(ptr addrspace(5), ptr addrspace(4), i64, i1 immarg)