Skip to content

[InstCombine] Fold reconstruction across select #145102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

[InstCombine] Fold reconstruction across select #145102

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
683da5c
pre-commit test
macsencasaus Jun 17, 2025
6f0a946
[InstCombine] Fold reconstruction across select
macsencasaus Jun 20, 2025
9542631
address review comments
macsencasaus Jun 21, 2025
f5bcc06
match correct mask only
macsencasaus Jun 23, 2025
039432c
match constant V2 value only
Aug 7, 2025
1ec984a
match V1 as constant
Aug 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1349,6 +1349,40 @@ Value *InstCombinerImpl::SimplifySelectsFeedingBinaryOp(BinaryOperator &I,
return nullptr;
};

// Special case for reconstructing across a select:
// (Cond ? V1 : (X & Mask)) op
// zext (Cond ? V2 : trunc X)
// -> (Cond ? (V1 op zext V2) : ((X & Mask) op zext trunc X))
auto foldReconstruction = [&](Value *V1, Value *Masked,
Value *ZExtSel) -> Value * {
if (Opcode != Instruction::Or)
return nullptr;

Value *X;
if (!match(Masked, m_OneUse(m_And(m_Value(X), m_Constant()))))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if (!match(Masked, m_OneUse(m_And(m_Value(X), m_Constant()))))
if (!match(Masked, m_OneUse(m_And(m_Value(X), m_APInt(*C)))))

Missing check for the mask. It should be APInt::getBitsSetFrom(X->getType()->getScalarSizeInBits(), Trunc->getType()->getScalarSizeInBits()).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought of this, but it would fail to simplify a case like

llvm-project/llvm/test/Transforms/InstCombine/select-reconstruction.ll

Lines 40 to 57 in 6f0a946

define i40 @select_reconstruction_257_mask(i40 %arg0) {
; CHECK-LABEL: define i40 @select_reconstruction_257_mask(
; CHECK-SAME: i40 [[ARG0:%.*]]) {
; CHECK-NEXT: [[TMP1:%.*]] = trunc i40 [[ARG0]] to i8
; CHECK-NEXT: [[TMP2:%.*]] = icmp eq i8 [[TMP1]], 2
; CHECK-NEXT: [[TMP3:%.*]] = and i40 [[ARG0]], -257
; CHECK-NEXT: [[TMP4:%.*]] = select i1 [[TMP2]], i40 0, i40 [[TMP3]]
; CHECK-NEXT: ret i40 [[TMP4]]
;
%1 = trunc i40 %arg0 to i8
%2 = icmp eq i8 %1, 2
%3 = and i40 %arg0, -257
%4 = select i1 %2, i8 0, i8 %1
%5 = select i1 %2, i40 0, i40 %3
%6 = zext i8 %4 to i40
%7 = or disjoint i40 %5, %6
ret i40 %7
}

I can do it anyway if you believe this case is not relevant.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah it should work for all constants. It even holds if we replace the and with any other instructions. I'd like to focus on the original motivating issue if you don't have a better idea to generalize the pattern. With this constraint, we can simplify ((X & Mask) op zext trunc X) into X directly, instead of creating a temporary instruction and relying on later optimizations.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To generalize, this optimization already occurs when the condition is arbitrary:
https://godbolt.org/z/8EPEP81TP

It just doesn't work when the condition is an icmp because of this check:

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

Lines 178 to 186 in cfcb788

// We are casting a select. Try to fold the cast into the select if the
// select does not have a compare instruction with matching operand types
// or the select is likely better done in a narrow type.
// Creating a select with operands that are different sizes than its
// condition may inhibit other folds and lead to worse codegen.
auto *Cmp = dyn_cast<CmpInst>(Sel->getCondition());
if (!Cmp || Cmp->getOperand(0)->getType() != Sel->getType() ||
(CI.getOpcode() == Instruction::Trunc &&
shouldChangeType(CI.getSrcTy(), CI.getType()))) {

I didn't know how to implement the optimization from here, but this may be an area where a more general optimization might occur.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is exactly what I said in #145102 (comment).

return nullptr;

Value *V2, *Trunc;
if (!match(ZExtSel, m_ZExt(m_OneUse(m_Select(m_Specific(Cond), m_Value(V2),
m_Value(Trunc))))))
return nullptr;

if (!match(Trunc, m_Trunc(m_Specific(X))))
return nullptr;

Value *ZExtTrue = Builder.CreateZExt(V2, V1->getType());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This can abort the transform after creating a new instruction. We should not do this, as it risks infinite loops.

Do I see correctly that in the motivating cases (and all tests) V2 is actually a constant? In that case we can specialize to that.

Value *True;
if (!(True = simplifyBinOp(Opcode, V1, ZExtTrue, FMF, Q)))
True = Builder.CreateOr(V1, ZExtTrue);

Value *ZExtFalse = Builder.CreateZExt(Trunc, V1->getType());
Value *False;
if (!(False = simplifyBinOp(Opcode, Masked, ZExtFalse, FMF, Q)))
False = Builder.CreateOr(Masked, ZExtFalse);

return Builder.CreateSelect(Cond, True, False, I.getName());
};

if (LHSIsSelect && RHSIsSelect && A == D) {
// (A ? B : C) op (A ? E : F) -> A ? (B op E) : (C op F)
Cond = A;
Expand All @@ -1368,13 +1402,17 @@ Value *InstCombinerImpl::SimplifySelectsFeedingBinaryOp(BinaryOperator &I,
False = simplifyBinOp(Opcode, C, RHS, FMF, Q);
if (Value *NewSel = foldAddNegate(B, C, RHS))
return NewSel;
if (Value *NewSel = foldReconstruction(B, C, RHS))
return NewSel;
} else if (RHSIsSelect && RHS->hasOneUse()) {
// X op (D ? E : F) -> D ? (X op E) : (X op F)
Cond = D;
True = simplifyBinOp(Opcode, LHS, E, FMF, Q);
False = simplifyBinOp(Opcode, LHS, F, FMF, Q);
if (Value *NewSel = foldAddNegate(E, F, LHS))
return NewSel;
if (Value *NewSel = foldReconstruction(E, F, LHS))
return NewSel;
}

if (!True || !False)
Expand Down
116 changes: 116 additions & 0 deletions llvm/test/Transforms/InstCombine/select-reconstruction.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
; RUN: opt < %s -passes=instcombine -S | FileCheck %s

define i40 @select_reconstruction_i40(i40 %arg0) {
; CHECK-LABEL: define i40 @select_reconstruction_i40(
; CHECK-SAME: i40 [[ARG0:%.*]]) {
; CHECK-NEXT: [[TMP1:%.*]] = trunc i40 [[ARG0]] to i8
; CHECK-NEXT: [[TMP2:%.*]] = icmp eq i8 [[TMP1]], 2
; CHECK-NEXT: [[TMP3:%.*]] = select i1 [[TMP2]], i40 0, i40 [[ARG0]]
; CHECK-NEXT: ret i40 [[TMP3]]
;
%low = trunc i40 %arg0 to i8
%is_low_two = icmp eq i8 %low, 2
%high = and i40 %arg0, -256
%select_low = select i1 %is_low_two, i8 0, i8 %low
%select_high = select i1 %is_low_two, i40 0, i40 %high
%zext_low = zext i8 %select_low to i40
%recomb = or disjoint i40 %select_high, %zext_low
ret i40 %recomb
}

define i40 @select_reconstruction_any_cmp_val(i40 %arg0, i8 %arg1) {
; CHECK-LABEL: define i40 @select_reconstruction_any_cmp_val(
; CHECK-SAME: i40 [[ARG0:%.*]], i8 [[ARG1:%.*]]) {
; CHECK-NEXT: [[TMP1:%.*]] = trunc i40 [[ARG0]] to i8
; CHECK-NEXT: [[TMP2:%.*]] = icmp eq i8 [[ARG1]], [[TMP1]]
; CHECK-NEXT: [[TMP3:%.*]] = select i1 [[TMP2]], i40 0, i40 [[ARG0]]
; CHECK-NEXT: ret i40 [[TMP3]]
;
%low = trunc i40 %arg0 to i8
%is_low_arg1 = icmp eq i8 %low, %arg1
%high = and i40 %arg0, -256
%select_low = select i1 %is_low_arg1, i8 0, i8 %low
%select_high = select i1 %is_low_arg1, i40 0, i40 %high
%zext_low = zext i8 %select_low to i40
%recomb = or disjoint i40 %select_high, %zext_low
ret i40 %recomb
}

define i40 @select_reconstruction_257_mask(i40 %arg0) {
; CHECK-LABEL: define i40 @select_reconstruction_257_mask(
; CHECK-SAME: i40 [[ARG0:%.*]]) {
; CHECK-NEXT: [[TMP1:%.*]] = trunc i40 [[ARG0]] to i8
; CHECK-NEXT: [[TMP2:%.*]] = icmp eq i8 [[TMP1]], 2
; CHECK-NEXT: [[TMP3:%.*]] = and i40 [[ARG0]], -257
; CHECK-NEXT: [[TMP4:%.*]] = select i1 [[TMP2]], i40 0, i40 [[TMP3]]
; CHECK-NEXT: ret i40 [[TMP4]]
;
%low = trunc i40 %arg0 to i8
%is_low_two = icmp eq i8 %low, 2
%high = and i40 %arg0, -257
%select_low = select i1 %is_low_two, i8 0, i8 %low
%select_high = select i1 %is_low_two, i40 0, i40 %high
%zext_low = zext i8 %select_low to i40
%recomb = or disjoint i40 %select_high, %zext_low
ret i40 %recomb
}

define i40 @select_reconstruction_i16_mask(i40 %arg0) {
; CHECK-LABEL: define i40 @select_reconstruction_i16_mask(
; CHECK-SAME: i40 [[ARG0:%.*]]) {
; CHECK-NEXT: [[TMP1:%.*]] = trunc i40 [[ARG0]] to i16
; CHECK-NEXT: [[TMP2:%.*]] = icmp eq i16 [[TMP1]], 2
; CHECK-NEXT: [[TMP3:%.*]] = select i1 [[TMP2]], i40 0, i40 [[ARG0]]
; CHECK-NEXT: ret i40 [[TMP3]]
;
%low = trunc i40 %arg0 to i16
%is_low_two = icmp eq i16 %low, 2
%high = and i40 %arg0, -65536
%select_low = select i1 %is_low_two, i16 0, i16 %low
%select_high = select i1 %is_low_two, i40 0, i40 %high
%zext_low = zext i16 %select_low to i40
%recomb = or disjoint i40 %select_high, %zext_low
ret i40 %recomb
}

define <2 x i32> @select_reconstruction_vec_any_cmp_val(<2 x i32> %arg0, <2 x i8> %arg1) {
; CHECK-LABEL: define <2 x i32> @select_reconstruction_vec_any_cmp_val(
; CHECK-SAME: <2 x i32> [[ARG0:%.*]], <2 x i8> [[ARG1:%.*]]) {
; CHECK-NEXT: [[TMP1:%.*]] = trunc <2 x i32> [[ARG0]] to <2 x i8>
; CHECK-NEXT: [[TMP2:%.*]] = icmp eq <2 x i8> [[ARG1]], [[TMP1]]
; CHECK-NEXT: [[TMP7:%.*]] = select <2 x i1> [[TMP2]], <2 x i32> zeroinitializer, <2 x i32> [[ARG0]]
; CHECK-NEXT: ret <2 x i32> [[TMP7]]
;
%low = trunc <2 x i32> %arg0 to <2 x i8>
%is_low_arg1 = icmp eq <2 x i8> %low, %arg1
%high = and <2 x i32> %arg0, <i32 -256, i32 -256>
%select_low = select <2 x i1> %is_low_arg1, <2 x i8> <i8 0, i8 0>, <2 x i8> %low
%select_high = select <2 x i1> %is_low_arg1, <2 x i32> <i32 0, i32 0>, <2 x i32> %high
%zext_low = zext <2 x i8> %select_low to <2 x i32>
%recomb = or <2 x i32> %select_high, %zext_low
ret <2 x i32> %recomb
}

; negative test
define i40 @select_reconstruction_impure_i16_mask_and(i40 %arg0) {
; CHECK-LABEL: define i40 @select_reconstruction_impure_i16_mask_and(
; CHECK-SAME: i40 [[ARG0:%.*]]) {
; CHECK-NEXT: [[TMP1:%.*]] = trunc i40 [[ARG0]] to i16
; CHECK-NEXT: [[TMP2:%.*]] = icmp eq i16 [[TMP1]], 2
; CHECK-NEXT: [[TMP3:%.*]] = and i40 [[ARG0]], 180
; CHECK-NEXT: [[TMP4:%.*]] = select i1 [[TMP2]], i16 0, i16 [[TMP1]]
; CHECK-NEXT: [[TMP5:%.*]] = select i1 [[TMP2]], i40 0, i40 [[TMP3]]
; CHECK-NEXT: [[TMP6:%.*]] = zext i16 [[TMP4]] to i40
; CHECK-NEXT: [[TMP7:%.*]] = and i40 [[TMP5]], [[TMP6]]
; CHECK-NEXT: ret i40 [[TMP7]]
;
%low = trunc i40 %arg0 to i16
%is_low_two = icmp eq i16 %low, 2
%high = and i40 %arg0, -65356
%select_low = select i1 %is_low_two, i16 0, i16 %low
%select_high = select i1 %is_low_two, i40 0, i40 %high
%zext_low = zext i16 %select_low to i40
%recomb = and i40 %select_high, %zext_low
ret i40 %recomb
}