[InstCombine] Simplify the pointer operand of store if writing to null is UB

[dtcxzyw]

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:

llvm-project/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

Lines 1027 to 1075 in 44feae8

	if (Op->hasOneUse()) {
	// Change select and PHI nodes to select values instead of addresses: this
	// helps alias analysis out a lot, allows many others simplifications, and
	// exposes redundancy in the code.
	//
	// Note that we cannot do the transformation unless we know that the
	// introduced loads cannot trap! Something like this is valid as long as
	// the condition is always false: load (select bool %C, int* null, int* %G),
	// but it would not be valid if we transformed it to load from null
	// unconditionally.
	//
	if (SelectInst *SI = dyn_cast<SelectInst>(Op)) {
	// load (select (Cond, &V1, &V2)) --> select(Cond, load &V1, load &V2).
	Align Alignment = LI.getAlign();
	if (isSafeToLoadUnconditionally(SI->getOperand(1), LI.getType(),
	Alignment, DL, SI) &&
	isSafeToLoadUnconditionally(SI->getOperand(2), LI.getType(),
	Alignment, DL, SI)) {
	LoadInst *V1 =
	Builder.CreateLoad(LI.getType(), SI->getOperand(1),
	SI->getOperand(1)->getName() + ".val");
	LoadInst *V2 =
	Builder.CreateLoad(LI.getType(), SI->getOperand(2),
	SI->getOperand(2)->getName() + ".val");
	assert(LI.isUnordered() && "implied by above");
	V1->setAlignment(Alignment);
	V1->setAtomic(LI.getOrdering(), LI.getSyncScopeID());
	V2->setAlignment(Alignment);
	V2->setAtomic(LI.getOrdering(), LI.getSyncScopeID());
	// It is safe to copy any metadata that does not trigger UB. Copy any
	// poison-generating metadata.
	V1->copyMetadata(LI, Metadata::PoisonGeneratingIDs);
	V2->copyMetadata(LI, Metadata::PoisonGeneratingIDs);
	return SelectInst::Create(SI->getCondition(), V1, V2);
	}
	// load (select (cond, null, P)) -> load P
	if (isa<ConstantPointerNull>(SI->getOperand(1)) &&
	!NullPointerIsDefined(SI->getFunction(),
	LI.getPointerAddressSpace()))
	return replaceOperand(LI, 0, SI->getOperand(2));
	// load (select (cond, P, null)) -> load P
	if (isa<ConstantPointerNull>(SI->getOperand(2)) &&
	!NullPointerIsDefined(SI->getFunction(),
	LI.getPointerAddressSpace()))
	return replaceOperand(LI, 0, SI->getOperand(1));
	}
	}

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

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:

llvm-project/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

Lines 1027 to 1075 in 44feae8

	if (Op->hasOneUse()) {
	// Change select and PHI nodes to select values instead of addresses: this
	// helps alias analysis out a lot, allows many others simplifications, and
	// exposes redundancy in the code.
	//
	// Note that we cannot do the transformation unless we know that the
	// introduced loads cannot trap! Something like this is valid as long as
	// the condition is always false: load (select bool %C, int* null, int* %G),
	// but it would not be valid if we transformed it to load from null
	// unconditionally.
	//
	if (SelectInst *SI = dyn_cast<SelectInst>(Op)) {
	// load (select (Cond, &V1, &V2)) --> select(Cond, load &V1, load &V2).
	Align Alignment = LI.getAlign();
	if (isSafeToLoadUnconditionally(SI->getOperand(1), LI.getType(),
	Alignment, DL, SI) &&
	isSafeToLoadUnconditionally(SI->getOperand(2), LI.getType(),
	Alignment, DL, SI)) {
	LoadInst *V1 =
	Builder.CreateLoad(LI.getType(), SI->getOperand(1),
	SI->getOperand(1)->getName() + ".val");
	LoadInst *V2 =
	Builder.CreateLoad(LI.getType(), SI->getOperand(2),
	SI->getOperand(2)->getName() + ".val");
	assert(LI.isUnordered() && "implied by above");
	V1->setAlignment(Alignment);
	V1->setAtomic(LI.getOrdering(), LI.getSyncScopeID());
	V2->setAlignment(Alignment);
	V2->setAtomic(LI.getOrdering(), LI.getSyncScopeID());
	// It is safe to copy any metadata that does not trigger UB. Copy any
	// poison-generating metadata.
	V1->copyMetadata(LI, Metadata::PoisonGeneratingIDs);
	V2->copyMetadata(LI, Metadata::PoisonGeneratingIDs);
	return SelectInst::Create(SI->getCondition(), V1, V2);
	}
	// load (select (cond, null, P)) -> load P
	if (isa<ConstantPointerNull>(SI->getOperand(1)) &&
	!NullPointerIsDefined(SI->getFunction(),
	LI.getPointerAddressSpace()))
	return replaceOperand(LI, 0, SI->getOperand(2));
	// load (select (cond, P, null)) -> load P
	if (isa<ConstantPointerNull>(SI->getOperand(2)) &&
	!NullPointerIsDefined(SI->getFunction(),
	LI.getPointerAddressSpace()))
	return replaceOperand(LI, 0, SI->getOperand(1));
	}
	}

---

Full diff: https://github.com/llvm/llvm-project/pull/127979.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp (+14)
• (modified) llvm/test/Transforms/InstCombine/store.ll (+42)

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 673395464c85a..0a2b0a5ee7987 100644
--- a/llvm/test/Transforms/InstCombine/store.ll
+++ b/llvm/test/Transforms/InstCombine/store.ll
@@ -345,6 +345,48 @@ 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:    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:    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"}

[nikic]

Before merging with the code in visitLoadInst, I need to investigate if the one-use check is necessary:

I think it's only needed for the first transform, but not the second...

[nikic]

LGTM