[llvm] [InstCombine] Simplify the pointer operand of store if writing to null is UB (PR #127979)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/127979

Proof: https://alive2.llvm.org/ce/z/mzVj-u
I will add some follow-up patches to avoid duplicate code, support more memory instructions, and bypass gep instructions.
Before merging with the code in `visitLoadInst`, I need to investigate if the one-use check is necessary:
https://github.com/llvm/llvm-project/blob/44feae869570201e9920c26b151ce7ce24f0418d/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp#L1027-L1075

>From fa7fa3d33b3035b24179fe992ed302e7a75ec6b3 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Thu, 20 Feb 2025 17:21:15 +0800
Subject: [PATCH 1/2] [InstCombine] Add pre-commit tests. NFC.

---
 llvm/test/Transforms/InstCombine/store.ll | 44 +++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/store.ll b/llvm/test/Transforms/InstCombine/store.ll
index 673395464c85a..5caf43a3cd6e9 100644
--- a/llvm/test/Transforms/InstCombine/store.ll
+++ b/llvm/test/Transforms/InstCombine/store.ll
@@ -345,6 +345,50 @@ define void @store_to_readonly_noalias(ptr readonly noalias %0) {
   ret void
 }
 
+define void @store_select_with_null(i1 %cond, ptr %p) {
+; CHECK-LABEL: @store_select_with_null(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], ptr [[P:%.*]], ptr null
+; CHECK-NEXT:    store i32 0, ptr [[SEL]], align 4
+; CHECK-NEXT:    ret void
+;
+  %sel = select i1 %cond, ptr %p, ptr null
+  store i32 0, ptr %sel, align 4
+  ret void
+}
+
+define void @store_select_with_null_commuted(i1 %cond, ptr %p) {
+; CHECK-LABEL: @store_select_with_null_commuted(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], ptr null, ptr [[P:%.*]]
+; CHECK-NEXT:    store i32 0, ptr [[SEL]], align 4
+; CHECK-NEXT:    ret void
+;
+  %sel = select i1 %cond, ptr null, ptr %p
+  store i32 0, ptr %sel, align 4
+  ret void
+}
+
+define void @store_select_with_null_null_is_valid(i1 %cond, ptr %p) null_pointer_is_valid {
+; CHECK-LABEL: @store_select_with_null_null_is_valid(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], ptr [[P:%.*]], ptr null
+; CHECK-NEXT:    store i32 0, ptr [[SEL]], align 4
+; CHECK-NEXT:    ret void
+;
+  %sel = select i1 %cond, ptr %p, ptr null
+  store i32 0, ptr %sel, align 4
+  ret void
+}
+
+define void @store_select_with_unknown(i1 %cond, ptr %p, ptr %p2) {
+; CHECK-LABEL: @store_select_with_unknown(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], ptr [[P:%.*]], ptr [[P2:%.*]]
+; CHECK-NEXT:    store i32 0, ptr [[SEL]], align 4
+; CHECK-NEXT:    ret void
+;
+  %sel = select i1 %cond, ptr %p, ptr %p2
+  store i32 0, ptr %sel, align 4
+  ret void
+}
+
 !0 = !{!4, !4, i64 0}
 !1 = !{!"omnipotent char", !2}
 !2 = !{!"Simple C/C++ TBAA"}

>From 2d6240eaba08e068c2dbfac051faba3cac6cb81a Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Thu, 20 Feb 2025 17:34:40 +0800
Subject: [PATCH 2/2] [InstCombine] Simplify the pointer operand of store if
 writing to null is UB

---
 .../InstCombine/InstCombineLoadStoreAlloca.cpp     | 14 ++++++++++++++
 llvm/test/Transforms/InstCombine/store.ll          |  6 ++----
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
index 4b42e86e25161..d5534c15cca76 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
@@ -1437,6 +1437,20 @@ Instruction *InstCombinerImpl::visitStoreInst(StoreInst &SI) {
   if (isa<UndefValue>(Val))
     return eraseInstFromFunction(SI);
 
+  // TODO: Add a helper to simplify the pointer operand for all memory
+  // instructions.
+  // store val, (select (cond, null, P)) -> store val, P
+  // store val, (select (cond, P, null)) -> store val, P
+  if (!NullPointerIsDefined(SI.getFunction(), SI.getPointerAddressSpace())) {
+    if (SelectInst *Sel = dyn_cast<SelectInst>(Ptr)) {
+      if (isa<ConstantPointerNull>(Sel->getOperand(1)))
+        return replaceOperand(SI, 1, Sel->getOperand(2));
+
+      if (isa<ConstantPointerNull>(Sel->getOperand(2)))
+        return replaceOperand(SI, 1, Sel->getOperand(1));
+    }
+  }
+
   return nullptr;
 }
 
diff --git a/llvm/test/Transforms/InstCombine/store.ll b/llvm/test/Transforms/InstCombine/store.ll
index 5caf43a3cd6e9..0a2b0a5ee7987 100644
--- a/llvm/test/Transforms/InstCombine/store.ll
+++ b/llvm/test/Transforms/InstCombine/store.ll
@@ -347,8 +347,7 @@ define void @store_to_readonly_noalias(ptr readonly noalias %0) {
 
 define void @store_select_with_null(i1 %cond, ptr %p) {
 ; CHECK-LABEL: @store_select_with_null(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], ptr [[P:%.*]], ptr null
-; CHECK-NEXT:    store i32 0, ptr [[SEL]], align 4
+; CHECK-NEXT:    store i32 0, ptr [[SEL:%.*]], align 4
 ; CHECK-NEXT:    ret void
 ;
   %sel = select i1 %cond, ptr %p, ptr null
@@ -358,8 +357,7 @@ define void @store_select_with_null(i1 %cond, ptr %p) {
 
 define void @store_select_with_null_commuted(i1 %cond, ptr %p) {
 ; CHECK-LABEL: @store_select_with_null_commuted(
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], ptr null, ptr [[P:%.*]]
-; CHECK-NEXT:    store i32 0, ptr [[SEL]], align 4
+; CHECK-NEXT:    store i32 0, ptr [[SEL:%.*]], align 4
 ; CHECK-NEXT:    ret void
 ;
   %sel = select i1 %cond, ptr null, ptr %p