[InstCombine] Fold correlated poison-blocking logical selects - #221130
Open
JihyeonJeong129 wants to merge 2 commits into
Open
[InstCombine] Fold correlated poison-blocking logical selects#221130JihyeonJeong129 wants to merge 2 commits into
JihyeonJeong129 wants to merge 2 commits into
Conversation
|
Hello @JihyeonJeong129 👋 Thank you for submitting a Pull Request (PR) to the LLVM Project. Since this is your first PR, here are a few useful links covering our main contribution policies and review practices.
Please reply to this message to confirm that you have read these policies, especially the LLVM AI Tool Use Policy, and that any AI tool usage has been noted in the PR description. Frequently asked questionsHow do I add reviewers? This PR will be automatically labeled, and the relevant teams will be notified. For some parts of the project, reviewers may also be added automatically. You can also add reviewers manually using the Reviewers section on this page. If you cannot use that section, it is probably because you do not have write permissions for the repository. In that case, you can request a review by tagging reviewers in a comment using What if there are no comments? If you have not received any comments on your PR after a week, you can request a review by pinging the PR with a comment such as “Ping”. The common courtesy ping rate is once a week. Please remember that you are asking for volunteer time from other developers. Are any special GitHub settings required to contribute to LLVM? We only require contributors to have a public email address associated with their GitHub commits, see this section of LLVM Developer Policy for details. If you have questions, feel free to leave a comment on this PR, or ask on LLVM Discord or LLVM Discourse. Thank you, |
llvmorg-github-actions
Bot
added
llvm:instcombine
|
@llvm/pr-subscribers-llvm-transforms Author: Jeong Jihyeon (JihyeonJeong129) ChangesThis patch folds correlated poison-blocking logical select conditions of the following form: The transformation is implemented as part of Tests cover:
Testing: Fixes #198820 Full diff: https://github.com/llvm/llvm-project/pull/221130.diff 3 Files Affected:
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 81d949fd1fb28..8732f3c5c01eb 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -3626,6 +3626,11 @@ foldSelectOfSymmetricSelect(SelectInst &OuterSelVal,
/// and rewrite it as
/// %inner.sel = select i1 %cond.alternative, i8 %sel.outer.t, i8 %sel.inner.t
/// %sel.outer = select i1 %cond.inner, i8 %inner.sel, i8 %sel.inner.f
+///
+/// Also fold correlated poison-blocking conditions
+/// select (!A || B), (select (A && B), T, F), X
+/// into
+/// select A, (select B, T, X), F
static Instruction *foldNestedSelects(SelectInst &OuterSelVal,
InstCombiner::BuilderTy &Builder) {
// We must start with a `select`.
@@ -3662,6 +3667,33 @@ static Instruction *foldNestedSelects(SelectInst &OuterSelVal,
if (match(InnerSel.Cond, m_Not(m_Value(InnerSel.Cond))))
std::swap(InnerSel.TrueVal, InnerSel.FalseVal);
+ // Fold correlated poison-blocking logical selects:
+ // C = A && B
+ // G = !A || B
+ // select G, (select C, T, F), X
+ // --> select A, (select B, T, X), F
+ // Only the canonical poison-blocking select forms with this exact operand
+ // order are handled. Bitwise and/or, commuted operands and a 'not' with
+ // poison lanes may also be valid, but they have not been verified yet.
+ if (!IsAndVariant && isa<SelectInst>(OuterSel.Cond) &&
+ isa<SelectInst>(InnerSel.Cond)) {
+ Value *A, *B;
+ if (match(OuterSel.Cond,
+ m_LogicalOr(m_NotForbidPoison(m_Value(A)), m_Value(B))) &&
+ match(InnerSel.Cond, m_LogicalAnd(m_Specific(A), m_Specific(B)))) {
+ Value *NewInner = Builder.CreateSelectWithUnknownProfile(
+ B, InnerSel.TrueVal, OuterSel.FalseVal, DEBUG_TYPE);
+ NewInner->takeName(InnerSelVal);
+
+ auto *NewOuter = SelectInst::Create(A, NewInner, InnerSel.FalseVal);
+ setExplicitlyUnknownBranchWeightsIfProfiled(*NewOuter, DEBUG_TYPE,
+ OuterSelVal.getFunction());
+ if (auto *FPOp = dyn_cast<FPMathOperator>(&OuterSelVal))
+ NewOuter->setFastMathFlags(FPOp->getFastMathFlags());
+ return NewOuter;
+ }
+ }
+
Value *AltCond = nullptr;
auto matchOuterCond = [OuterSel, IsAndVariant, &AltCond](auto m_InnerCond) {
// An unsimplified select condition can match both LogicalAnd and LogicalOr
diff --git a/llvm/test/Transforms/InstCombine/nested-select.ll b/llvm/test/Transforms/InstCombine/nested-select.ll
index b8a7d4e8ab789..db4aefda6c753 100644
--- a/llvm/test/Transforms/InstCombine/nested-select.ll
+++ b/llvm/test/Transforms/InstCombine/nested-select.ll
@@ -29,6 +29,214 @@ define i8 @orcond(i1 %inner.cond, i1 %alt.cond, i8 %inner.sel.trueval, i8 %inner
ret i8 %outer.sel
}
+; Correlated poison-blocking logical conditions
+
+define i8 @correlated_poison_blocking_conditions(i1 %a, i1 %b, i8 %t, i8 %f, i8 %x) {
+; CHECK-LABEL: @correlated_poison_blocking_conditions(
+; CHECK-NEXT: [[MUX:%.*]] = select i1 [[B:%.*]], i8 [[T:%.*]], i8 [[X:%.*]]
+; CHECK-NEXT: [[RET:%.*]] = select i1 [[A:%.*]], i8 [[MUX]], i8 [[F:%.*]]
+; CHECK-NEXT: ret i8 [[RET]]
+;
+ %and = select i1 %a, i1 %b, i1 false
+ %not.a = xor i1 %a, true
+ %guard = select i1 %not.a, i1 true, i1 %b
+ %mux = select i1 %and, i8 %t, i8 %f
+ %ret = select i1 %guard, i8 %mux, i8 %x
+ ret i8 %ret
+}
+
+; Verify that foldSelectOfBools canonicalizes the guard before this fold.
+
+define i8 @correlated_poison_blocking_conditions_noncanonical_guard(i1 %a, i1 %b, i8 %t, i8 %f, i8 %x) {
+; CHECK-LABEL: @correlated_poison_blocking_conditions_noncanonical_guard(
+; CHECK-NEXT: [[MUX:%.*]] = select i1 [[B:%.*]], i8 [[T:%.*]], i8 [[X:%.*]]
+; CHECK-NEXT: [[RET:%.*]] = select i1 [[A:%.*]], i8 [[MUX]], i8 [[F:%.*]]
+; CHECK-NEXT: ret i8 [[RET]]
+;
+ %and = select i1 %a, i1 %b, i1 false
+ %guard = select i1 %a, i1 %b, i1 true
+ %mux = select i1 %and, i8 %t, i8 %f
+ %ret = select i1 %guard, i8 %mux, i8 %x
+ ret i8 %ret
+}
+
+define <2 x i8> @correlated_poison_blocking_conditions_vec(<2 x i1> %a, <2 x i1> %b, <2 x i8> %t, <2 x i8> %f, <2 x i8> %x) {
+; CHECK-LABEL: @correlated_poison_blocking_conditions_vec(
+; CHECK-NEXT: [[MUX:%.*]] = select <2 x i1> [[B:%.*]], <2 x i8> [[T:%.*]], <2 x i8> [[X:%.*]]
+; CHECK-NEXT: [[RET:%.*]] = select <2 x i1> [[A:%.*]], <2 x i8> [[MUX]], <2 x i8> [[F:%.*]]
+; CHECK-NEXT: ret <2 x i8> [[RET]]
+;
+ %and = select <2 x i1> %a, <2 x i1> %b, <2 x i1> zeroinitializer
+ %not.a = xor <2 x i1> %a, <i1 true, i1 true>
+ %guard = select <2 x i1> %not.a, <2 x i1> <i1 true, i1 true>, <2 x i1> %b
+ %mux = select <2 x i1> %and, <2 x i8> %t, <2 x i8> %f
+ %ret = select <2 x i1> %guard, <2 x i8> %mux, <2 x i8> %x
+ ret <2 x i8> %ret
+}
+
+; The existing profitability rule permits reassociation when the guard is
+; one-use, even if the inner mux has another use.
+
+define i8 @correlated_poison_blocking_conditions_mux_multiuse(i1 %a, i1 %b, i8 %t, i8 %f, i8 %x) {
+; CHECK-LABEL: @correlated_poison_blocking_conditions_mux_multiuse(
+; CHECK-NEXT: [[AND:%.*]] = select i1 [[A:%.*]], i1 [[B:%.*]], i1 false
+; CHECK-NEXT: [[OLD_MUX:%.*]] = select i1 [[AND]], i8 [[T:%.*]], i8 [[F:%.*]]
+; CHECK-NEXT: call void @use.i8(i8 [[OLD_MUX]])
+; CHECK-NEXT: [[MUX:%.*]] = select i1 [[B]], i8 [[T]], i8 [[X:%.*]]
+; CHECK-NEXT: [[RET:%.*]] = select i1 [[A]], i8 [[MUX]], i8 [[F]]
+; CHECK-NEXT: ret i8 [[RET]]
+;
+ %and = select i1 %a, i1 %b, i1 false
+ %not.a = xor i1 %a, true
+ %guard = select i1 %not.a, i1 true, i1 %b
+ %mux = select i1 %and, i8 %t, i8 %f
+ call void @use.i8(i8 %mux)
+ %ret = select i1 %guard, i8 %mux, i8 %x
+ ret i8 %ret
+}
+
+; Do not increase instruction count when neither the guard nor the mux is
+; one-use.
+
+define i8 @correlated_poison_blocking_conditions_both_multiuse(i1 %a, i1 %b, i8 %t, i8 %f, i8 %x) {
+; CHECK-LABEL: @correlated_poison_blocking_conditions_both_multiuse(
+; CHECK-NEXT: [[AND:%.*]] = select i1 [[A:%.*]], i1 [[B:%.*]], i1 false
+; CHECK-NEXT: [[NOT_A:%.*]] = xor i1 [[A]], true
+; CHECK-NEXT: [[GUARD:%.*]] = select i1 [[NOT_A]], i1 true, i1 [[B]]
+; CHECK-NEXT: call void @use.i1(i1 [[GUARD]])
+; CHECK-NEXT: [[MUX:%.*]] = select i1 [[AND]], i8 [[T:%.*]], i8 [[F:%.*]]
+; CHECK-NEXT: call void @use.i8(i8 [[MUX]])
+; CHECK-NEXT: [[RET:%.*]] = select i1 [[GUARD]], i8 [[MUX]], i8 [[X:%.*]]
+; CHECK-NEXT: ret i8 [[RET]]
+;
+ %and = select i1 %a, i1 %b, i1 false
+ %not.a = xor i1 %a, true
+ %guard = select i1 %not.a, i1 true, i1 %b
+ call void @use.i1(i1 %guard)
+ %mux = select i1 %and, i8 %t, i8 %f
+ call void @use.i8(i8 %mux)
+ %ret = select i1 %guard, i8 %mux, i8 %x
+ ret i8 %ret
+}
+
+; Bitwise or is not handled yet. It may also be valid, but it has not been
+; verified.
+
+define i8 @correlated_conditions_bitwise_or(i1 %a, i8 %b8, i8 %t, i8 %f, i8 %x) {
+; CHECK-LABEL: @correlated_conditions_bitwise_or(
+; CHECK-NEXT: [[B:%.*]] = trunc i8 [[B8:%.*]] to i1
+; CHECK-NEXT: [[AND:%.*]] = select i1 [[A:%.*]], i1 [[B]], i1 false
+; CHECK-NEXT: [[NOT_A:%.*]] = xor i1 [[A]], true
+; CHECK-NEXT: [[GUARD:%.*]] = or i1 [[NOT_A]], [[B]]
+; CHECK-NEXT: [[MUX:%.*]] = select i1 [[AND]], i8 [[T:%.*]], i8 [[F:%.*]]
+; CHECK-NEXT: [[RET:%.*]] = select i1 [[GUARD]], i8 [[MUX]], i8 [[X:%.*]]
+; CHECK-NEXT: ret i8 [[RET]]
+;
+ %b = trunc i8 %b8 to i1
+ %and = select i1 %a, i1 %b, i1 false
+ %not.a = xor i1 %a, true
+ %guard = or i1 %not.a, %b
+ %mux = select i1 %and, i8 %t, i8 %f
+ %ret = select i1 %guard, i8 %mux, i8 %x
+ ret i8 %ret
+}
+
+; Bitwise and is not handled yet. It may also be valid, but it has not been
+; verified.
+
+define i8 @correlated_conditions_bitwise_and(i1 %a, i8 %b8, i8 %t, i8 %f, i8 %x) {
+; CHECK-LABEL: @correlated_conditions_bitwise_and(
+; CHECK-NEXT: [[B:%.*]] = trunc i8 [[B8:%.*]] to i1
+; CHECK-NEXT: [[AND:%.*]] = and i1 [[A:%.*]], [[B]]
+; CHECK-NEXT: [[NOT_A:%.*]] = xor i1 [[A]], true
+; CHECK-NEXT: [[GUARD:%.*]] = select i1 [[NOT_A]], i1 true, i1 [[B]]
+; CHECK-NEXT: [[MUX:%.*]] = select i1 [[AND]], i8 [[T:%.*]], i8 [[F:%.*]]
+; CHECK-NEXT: [[RET:%.*]] = select i1 [[GUARD]], i8 [[MUX]], i8 [[X:%.*]]
+; CHECK-NEXT: ret i8 [[RET]]
+;
+ %b = trunc i8 %b8 to i1
+ %and = and i1 %a, %b
+ %not.a = xor i1 %a, true
+ %guard = select i1 %not.a, i1 true, i1 %b
+ %mux = select i1 %and, i8 %t, i8 %f
+ %ret = select i1 %guard, i8 %mux, i8 %x
+ ret i8 %ret
+}
+
+; Commuted AND operands are not handled yet. They may also be valid, but this
+; has not been verified.
+
+define i8 @correlated_conditions_swapped_and(i1 %a, i1 %b, i8 %t, i8 %f, i8 %x) {
+; CHECK-LABEL: @correlated_conditions_swapped_and(
+; CHECK-NEXT: [[AND:%.*]] = select i1 [[B:%.*]], i1 [[A:%.*]], i1 false
+; CHECK-NEXT: [[NOT_A:%.*]] = xor i1 [[A]], true
+; CHECK-NEXT: [[GUARD:%.*]] = select i1 [[NOT_A]], i1 true, i1 [[B]]
+; CHECK-NEXT: [[MUX:%.*]] = select i1 [[AND]], i8 [[T:%.*]], i8 [[F:%.*]]
+; CHECK-NEXT: [[RET:%.*]] = select i1 [[GUARD]], i8 [[MUX]], i8 [[X:%.*]]
+; CHECK-NEXT: ret i8 [[RET]]
+;
+ %and = select i1 %b, i1 %a, i1 false
+ %not.a = xor i1 %a, true
+ %guard = select i1 %not.a, i1 true, i1 %b
+ %mux = select i1 %and, i8 %t, i8 %f
+ %ret = select i1 %guard, i8 %mux, i8 %x
+ ret i8 %ret
+}
+
+; Commuted OR operands are not handled yet. They may also be valid, but this
+; has not been verified.
+
+define i8 @correlated_conditions_swapped_or(i1 %a, i1 %b, i8 %t, i8 %f, i8 %x) {
+; CHECK-LABEL: @correlated_conditions_swapped_or(
+; CHECK-NEXT: [[AND:%.*]] = select i1 [[A:%.*]], i1 [[B:%.*]], i1 false
+; CHECK-NEXT: [[NOT_A:%.*]] = xor i1 [[A]], true
+; CHECK-NEXT: [[GUARD:%.*]] = select i1 [[B]], i1 true, i1 [[NOT_A]]
+; CHECK-NEXT: [[MUX:%.*]] = select i1 [[AND]], i8 [[T:%.*]], i8 [[F:%.*]]
+; CHECK-NEXT: [[RET:%.*]] = select i1 [[GUARD]], i8 [[MUX]], i8 [[X:%.*]]
+; CHECK-NEXT: ret i8 [[RET]]
+;
+ %and = select i1 %a, i1 %b, i1 false
+ %not.a = xor i1 %a, true
+ %guard = select i1 %b, i1 true, i1 %not.a
+ %mux = select i1 %and, i8 %t, i8 %f
+ %ret = select i1 %guard, i8 %mux, i8 %x
+ ret i8 %ret
+}
+
+; A 'not' with a poison lane is not handled yet (m_NotForbidPoison). It may
+; also be valid, but it has not been verified.
+
+define <2 x i8> @correlated_conditions_poison_in_not(<2 x i1> %a, <2 x i1> %b, <2 x i8> %t, <2 x i8> %f, <2 x i8> %x) {
+; CHECK-LABEL: @correlated_conditions_poison_in_not(
+; CHECK-NEXT: [[AND:%.*]] = select <2 x i1> [[A:%.*]], <2 x i1> [[B:%.*]], <2 x i1> zeroinitializer
+; CHECK-NEXT: [[NOT_A:%.*]] = xor <2 x i1> [[A]], <i1 true, i1 poison>
+; CHECK-NEXT: [[GUARD:%.*]] = select <2 x i1> [[NOT_A]], <2 x i1> splat (i1 true), <2 x i1> [[B]]
+; CHECK-NEXT: [[MUX:%.*]] = select <2 x i1> [[AND]], <2 x i8> [[T:%.*]], <2 x i8> [[F:%.*]]
+; CHECK-NEXT: [[RET:%.*]] = select <2 x i1> [[GUARD]], <2 x i8> [[MUX]], <2 x i8> [[X:%.*]]
+; CHECK-NEXT: ret <2 x i8> [[RET]]
+;
+ %and = select <2 x i1> %a, <2 x i1> %b, <2 x i1> zeroinitializer
+ %not.a = xor <2 x i1> %a, <i1 true, i1 poison>
+ %guard = select <2 x i1> %not.a, <2 x i1> <i1 true, i1 true>, <2 x i1> %b
+ %mux = select <2 x i1> %and, <2 x i8> %t, <2 x i8> %f
+ %ret = select <2 x i1> %guard, <2 x i8> %mux, <2 x i8> %x
+ ret <2 x i8> %ret
+}
+
+define float @correlated_poison_blocking_conditions_fmf(i1 %a, i1 %b, float %t, float %f, float %x) {
+; CHECK-LABEL: @correlated_poison_blocking_conditions_fmf(
+; CHECK-NEXT: [[MUX:%.*]] = select i1 [[B:%.*]], float [[T:%.*]], float [[X:%.*]]
+; CHECK-NEXT: [[RET:%.*]] = select nnan i1 [[A:%.*]], float [[MUX]], float [[F:%.*]]
+; CHECK-NEXT: ret float [[RET]]
+;
+ %and = select i1 %a, i1 %b, i1 false
+ %not.a = xor i1 %a, true
+ %guard = select i1 %not.a, i1 true, i1 %b
+ %mux = select ninf i1 %and, float %t, float %f
+ %ret = select nnan i1 %guard, float %mux, float %x
+ ret float %ret
+}
+
; Extra use tests (basic test, no inversions)
define i8 @andcond.extrause0(i1 %inner.cond, i1 %alt.cond, i8 %inner.sel.trueval, i8 %inner.sel.falseval, i8 %outer.sel.trueval) {
diff --git a/llvm/test/Transforms/InstCombine/preserve-profile.ll b/llvm/test/Transforms/InstCombine/preserve-profile.ll
index 8cb3e685ae302..3cb8f9669f8cc 100644
--- a/llvm/test/Transforms/InstCombine/preserve-profile.ll
+++ b/llvm/test/Transforms/InstCombine/preserve-profile.ll
@@ -83,6 +83,21 @@ define i32 @add_zext_zext_i1(i1 %a) !prof !0 {
ret i32 %add
}
+define i32 @correlated_poison_blocking_selects(i1 %a, i1 %b, i32 %t, i32 %f, i32 %x) !prof !0 {
+; CHECK-LABEL: define i32 @correlated_poison_blocking_selects(
+; CHECK-SAME: i1 [[A:%.*]], i1 [[B:%.*]], i32 [[T:%.*]], i32 [[F:%.*]], i32 [[X:%.*]]) !prof [[PROF0]] {
+; CHECK-NEXT: [[MUX:%.*]] = select i1 [[B]], i32 [[T]], i32 [[X]], !prof [[PROF2]]
+; CHECK-NEXT: [[RET:%.*]] = select i1 [[A]], i32 [[MUX]], i32 [[F]], !prof [[PROF2]]
+; CHECK-NEXT: ret i32 [[RET]]
+;
+ %and = select i1 %a, i1 %b, i1 false
+ %not.a = xor i1 %a, true
+ %guard = select i1 %not.a, i1 true, i1 %b
+ %mux = select i1 %and, i32 %t, i32 %f, !prof !1
+ %ret = select i1 %guard, i32 %mux, i32 %x, !prof !1
+ ret i32 %ret
+}
+
define i32 @no_count_no_branch_weights(i1 %a) {
; CHECK-LABEL: define i32 @no_count_no_branch_weights(
; CHECK-SAME: i1 [[A:%.*]]) {
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This patch folds correlated poison-blocking logical select conditions of the following form:
The transformation is implemented as part of
foldNestedSelects()and uses its existing profitability rule to avoid increasing the instruction count.Tests cover:
foldSelectOfBools()andandorconditions excluded from the foldTesting:
Fixes #198820