[Mem2Reg] Generate non-terminator unreachable for !noundef undef

[nikic]

When performing a load from uninitialized memory using !noundef, insert a non-terminator unreachable instruction, which will be converted to a proper unreachable by SimplifyCFG. This way we retain the fact that UB occurred on this code path.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Nikita Popov (nikic)

Changes

When performing a load from uninitialized memory using !noundef, insert a non-terminator unreachable instruction, which will be converted to a proper unreachable by SimplifyCFG. This way we retain the fact that UB occurred on this code path.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp (+9)
• (modified) llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll (+6-6)

diff --git a/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp b/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
index 40d0f6b75d69b..a1210ac0f88ac 100644
--- a/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
+++ b/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
@@ -453,6 +453,15 @@ static void addAssumeNonNull(AssumptionCache *AC, LoadInst *LI) {
 static void convertMetadataToAssumes(LoadInst *LI, Value *Val,
                                      const DataLayout &DL, AssumptionCache *AC,
                                      const DominatorTree *DT) {
+  if (isa<UndefValue>(Val) && LI->hasMetadata(LLVMContext::MD_noundef)) {
+    // Insert non-terminator unreachable.
+    LLVMContext &Ctx = LI->getContext();
+    new StoreInst(ConstantInt::getTrue(Ctx),
+                  PoisonValue::get(PointerType::getUnqual(Ctx)),
+                  /*isVolatile=*/false, Align(1), LI);
+    return;
+  }
+
   // If the load was marked as nonnull we don't want to lose that information
   // when we erase this Load. So we preserve it with an assume. As !nonnull
   // returns poison while assume violations are immediate undefined behavior,
diff --git a/llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll b/llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll
index e3c14cc662e2a..edc07864795e8 100644
--- a/llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll
+++ b/llvm/test/Transforms/Mem2Reg/preserve-nonnull-load-metadata.ll
@@ -143,6 +143,7 @@ fin:
 define ptr @no_store_single_load_noundef() {
 ; CHECK-LABEL: @no_store_single_load_noundef(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i1 true, ptr poison, align 1
 ; CHECK-NEXT:    ret ptr undef
 ;
 entry:
@@ -156,8 +157,10 @@ define ptr @no_store_multiple_loads_noundef(i1 %c) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[ELSE:%.*]]
 ; CHECK:       if:
+; CHECK-NEXT:    store i1 true, ptr poison, align 1
 ; CHECK-NEXT:    ret ptr undef
 ; CHECK:       else:
+; CHECK-NEXT:    store i1 true, ptr poison, align 1
 ; CHECK-NEXT:    ret ptr undef
 ;
 entry:
@@ -176,8 +179,7 @@ if:
 define ptr @no_store_single_load_nonnull_noundef() {
 ; CHECK-LABEL: @no_store_single_load_nonnull_noundef(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = icmp ne ptr undef, null
-; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP0]])
+; CHECK-NEXT:    store i1 true, ptr poison, align 1
 ; CHECK-NEXT:    ret ptr undef
 ;
 entry:
@@ -191,12 +193,10 @@ define ptr @no_store_multiple_loads_nonnull_noundef(i1 %c) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[ELSE:%.*]]
 ; CHECK:       if:
-; CHECK-NEXT:    [[TMP0:%.*]] = icmp ne ptr undef, null
-; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP0]])
+; CHECK-NEXT:    store i1 true, ptr poison, align 1
 ; CHECK-NEXT:    ret ptr undef
 ; CHECK:       else:
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne ptr undef, null
-; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP1]])
+; CHECK-NEXT:    store i1 true, ptr poison, align 1
 ; CHECK-NEXT:    ret ptr undef
 ;
 entry:

[dtcxzyw]

LGTM. Can we also do this for SROA?

[antoniofrighetto]

Can we also do this for SROA?

Do we have to, out of curiosity? Thought this was only being done here, exactly to relieve SROA from carrying that out.

[nikic]

LGTM. Can we also do this for SROA?

SROA uses mem2reg internally, so this mostly also covers SROA as well (apart from one edge case...)

[dianqk]

Why not use call void @llvm.assume(i1 false)? If I understand correctly, this should be the more canonical approach?

[nikic]

Why not use call void @llvm.assume(i1 false)? If I understand correctly, this should be the more canonical approach?

Store to poison is the canonical form for non-terminator unreachable. In fact, InstCombine will convert assume(false) to it: https://llvm.godbolt.org/z/rsfd7hzja

[dianqk]

LGTM. Thanks!