[InstCombine] Fold select of undef to freeze

[nikic]

Fold select(c, undef, x) to freeze(x), in the case where we can't fold this to just x, because x is potentially poison (and does not imply c is poison).

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Nikita Popov (nikic)

Changes

Fold select(c, undef, x) to freeze(x), in the case where we can't fold this to just x, because x is potentially poison (and does not imply c is poison).

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp (+9)
• (modified) llvm/test/Transforms/InstCombine/select.ll (+4-4)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 27563d4414073..c457e719d3ba9 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -4016,5 +4016,14 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   if (CondVal->getType() == SI.getType() && isKnownInversion(FalseVal, TrueVal))
     return BinaryOperator::CreateXor(CondVal, FalseVal);
 
+  // select Cond, undef, X -> freeze(X)
+  // select Cond, X, undef -> freeze(X)
+  // The case where X is non-poison (or implies the condition is poison) is
+  // already handled by InstSimplify.
+  if (isa<UndefValue>(TrueVal))
+    return new FreezeInst(FalseVal);
+  if (isa<UndefValue>(FalseVal))
+    return new FreezeInst(TrueVal);
+
   return nullptr;
 }
diff --git a/llvm/test/Transforms/InstCombine/select.ll b/llvm/test/Transforms/InstCombine/select.ll
index b37e9175b26a5..9f150c6e57e20 100644
--- a/llvm/test/Transforms/InstCombine/select.ll
+++ b/llvm/test/Transforms/InstCombine/select.ll
@@ -2562,7 +2562,7 @@ exit:
 ; https://reviews.llvm.org/D83360
 define i32 @false_undef(i1 %cond, i32 %x) {
 ; CHECK-LABEL: @false_undef(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[COND:%.*]], i32 [[X:%.*]], i32 undef
+; CHECK-NEXT:    [[S:%.*]] = freeze i32 [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %s = select i1 %cond, i32 %x, i32 undef
@@ -2571,7 +2571,7 @@ define i32 @false_undef(i1 %cond, i32 %x) {
 
 define i32 @true_undef(i1 %cond, i32 %x) {
 ; CHECK-LABEL: @true_undef(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[COND:%.*]], i32 undef, i32 [[X:%.*]]
+; CHECK-NEXT:    [[S:%.*]] = freeze i32 [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %s = select i1 %cond, i32 undef, i32 %x
@@ -2580,7 +2580,7 @@ define i32 @true_undef(i1 %cond, i32 %x) {
 
 define <2 x i32> @false_undef_vec(i1 %cond, <2 x i32> %x) {
 ; CHECK-LABEL: @false_undef_vec(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[COND:%.*]], <2 x i32> [[X:%.*]], <2 x i32> undef
+; CHECK-NEXT:    [[S:%.*]] = freeze <2 x i32> [[X:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[S]]
 ;
   %s = select i1 %cond, <2 x i32> %x, <2 x i32> undef
@@ -2589,7 +2589,7 @@ define <2 x i32> @false_undef_vec(i1 %cond, <2 x i32> %x) {
 
 define <2 x i32> @true_undef_vec(i1 %cond, <2 x i32> %x) {
 ; CHECK-LABEL: @true_undef_vec(
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[COND:%.*]], <2 x i32> undef, <2 x i32> [[X:%.*]]
+; CHECK-NEXT:    [[S:%.*]] = freeze <2 x i32> [[X:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[S]]
 ;
   %s = select i1 %cond, <2 x i32> undef, <2 x i32> %x

[goldsteinn]

InstSimplify instead? (edit: obv not...)

But looks fine assuming dtcxzyw's benchmarks look fine.

[dtcxzyw]

I cannot tell whether this patch has a positive net effect. Some freezes are not eliminated in the end :(

[nikic]

Maybe we should wait on the nopoison attribute first? Though this will take time as it needs an LLVM upgrade in Rust first...