[InstCombine] Update the select operand when the cond is trunc and has the nuw or nsw property.

[c8ef]

This patch updates the select operand when the cond has the nuw or nsw property. Considering the semantics of the nuw and nsw flag, if there is no poison value in this expression, this code assumes that X can only be 0, 1 or -1.

close: #96765
alive2: https://alive2.llvm.org/ce/z/3n3n2Q

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: None (c8ef)

Changes

This patch adds a fold from select (trunc nuw X to i1), X, Y to select (trunc nuw X to i1), 1, Y. Considering the semantics of the nuw flag, if there is no poison value in this expression, this code assumes that X can only be 0 or 1.

Partially resolved: #96765

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp (+10)
• (added) llvm/test/Transforms/InstCombine/fold-select-trunc-if-nuw.ll (+35)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 18ffc209f259e0..49582004d9a4f6 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -4201,5 +4201,15 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
     }
   }
 
+  // select (trunc nuw X to i1), X, Y --> select (trunc nuw X to i1), 1, Y
+  if (auto *TI = dyn_cast<TruncInst>(CondVal)) {
+    Value *Trunc;
+    if (TI->hasNoUnsignedWrap() && match(CondVal, m_Trunc(m_Value(Trunc))) &&
+        match(TrueVal, m_Specific(Trunc))) {
+      return SelectInst::Create(
+          CondVal, ConstantInt::get(TrueVal->getType(), 1), FalseVal);
+    }
+  }
+
   return nullptr;
 }
diff --git a/llvm/test/Transforms/InstCombine/fold-select-trunc-if-nuw.ll b/llvm/test/Transforms/InstCombine/fold-select-trunc-if-nuw.ll
new file mode 100644
index 00000000000000..09fb2d77f9ef56
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/fold-select-trunc-if-nuw.ll
@@ -0,0 +1,35 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i8 @fold_select_trunc_if_nuw(i8 %x, i8 %y) {
+; CHECK-LABEL: @fold_select_trunc_if_nuw(
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nuw i8 [[X:%.*]] to i1
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[TRUNC]], i8 1, i8 [[Y:%.*]]
+; CHECK-NEXT:    ret i8 [[RET]]
+;
+  %trunc = trunc nuw i8 %x to i1
+  %ret = select i1 %trunc, i8 %x, i8 %y
+  ret i8 %ret
+}
+
+define i8 @fold_select_trunc_if_nuw_negative(i8 %x, i8 %y) {
+; CHECK-LABEL: @fold_select_trunc_if_nuw_negative(
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nsw i8 [[X:%.*]] to i1
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[TRUNC]], i8 [[X]], i8 [[Y:%.*]]
+; CHECK-NEXT:    ret i8 [[RET]]
+;
+  %trunc = trunc nsw i8 %x to i1
+  %ret = select i1 %trunc, i8 %x, i8 %y
+  ret i8 %ret
+}
+
+define <2 x i8> @fold_select_trunc_if_nuw_vector(<2 x i8> %x, <2 x i8> %y) {
+; CHECK-LABEL: @fold_select_trunc_if_nuw_vector(
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nuw <2 x i8> [[X:%.*]] to <2 x i1>
+; CHECK-NEXT:    [[RET:%.*]] = select <2 x i1> [[TRUNC]], <2 x i8> <i8 1, i8 1>, <2 x i8> [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i8> [[RET]]
+;
+  %trunc = trunc nuw <2 x i8> %x to <2 x i1>
+  %ret = select <2 x i1> %trunc, <2 x i8> %x, <2 x i8> %y
+  ret <2 x i8> %ret
+}

[nikic]

Please add an alive2 proof to the patch description.

Note that the same transform also works for trunc nsw and -1, though that's probably less practically relevant.

[nikic]

Note that the same transform also works for trunc nsw and -1, though that's probably less practically relevant.

There's also the variant where the trunc operand is used in FalseVal rather than TrueVal. In that case we can replace with zero.

[c8ef]

Please add an alive2 proof to the patch description.

Note that the same transform also works for trunc nsw and -1, though that's probably less practically relevant.

Added all four transformations and relevant alive2 proof.

[nikic]

LGTM

This could be generalized using simplifyWithOpReplaced, but probably not worthwhile. The non-trunc variant of that is:

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

Lines 3691 to 3702 in 99b85ca

	if (!isa<Constant>(CondVal) && SelType->isIntOrIntVectorTy() &&
	CondType->isVectorTy() == SelType->isVectorTy()) {
	if (Value *S = simplifyWithOpReplaced(TrueVal, CondVal,
	ConstantInt::getTrue(CondType), SQ,
	/* AllowRefinement */ true))
	return replaceOperand(SI, 1, S);
	if (Value *S = simplifyWithOpReplaced(FalseVal, CondVal,
	ConstantInt::getFalse(CondType), SQ,
	/* AllowRefinement */ true))
	return replaceOperand(SI, 2, S);
	}

[c8ef]

Would it be better to place this transformation at

llvm-project/llvm/lib/Analysis/InstructionSimplify.cpp

Line 4812 in 99b85ca

static Value *simplifySelectInst(Value *Cond, Value *TrueVal, Value *FalseVal,

? If yes, I'd like to relocate the implementation there.

[nikic]

Would it be better to place this transformation at

llvm-project/llvm/lib/Analysis/InstructionSimplify.cpp

Line 4812 in 99b85ca

static Value *simplifySelectInst(Value *Cond, Value *TrueVal, Value *FalseVal,

? If yes, I'd like to relocate the implementation there.

The transform can't go into InstSimplify because it modifies/creates an instruction.

[dtcxzyw]

LGTM.