Skip to content

[X86] Extend combinei64TruncSrlAdd to handle patterns with or and xor #128435

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

Merged
merged 8 commits into from
Mar 1, 2025
Merged

[X86] Extend combinei64TruncSrlAdd to handle patterns with or and xor #128435

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
fe8151f
[X86] Add pre-commit tests.
joaotgouveia Feb 23, 2025
5574dce
[X86] Extend `combinei64TruncSrlAdd` to handle patterns with `or` and…
joaotgouveia Feb 23, 2025
55189cc
[X86] Add extra test cases
joaotgouveia Feb 24, 2025
d2f27bf
[X86] Address reviews
joaotgouveia Feb 26, 2025
1ba8e32
[X86] Improve readability
joaotgouveia Feb 27, 2025
696f3a5
[X86] Remove getNumOperands
joaotgouveia Feb 27, 2025
7caceb1
[X86] Remove unnecessary switch case
joaotgouveia Feb 28, 2025
4388e73
[X86] Address reviews
joaotgouveia Mar 1, 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
77 changes: 49 additions & 28 deletions llvm/lib/Target/X86/X86ISelLowering.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53733,36 +53733,53 @@ static SDValue combineLRINT_LLRINT(SDNode *N, SelectionDAG &DAG,
return DAG.getNode(X86ISD::CVTP2SI, DL, VT, Src);
}

// Attempt to fold some (truncate (srl (add X, C1), C2)) patterns to
// (add (truncate (srl X, C2)), C1'). C1' will be smaller than C1 so we are able
// to avoid generating code with MOVABS and large constants in certain cases.
static SDValue combinei64TruncSrlAdd(SDValue N, EVT VT, SelectionDAG &DAG,
const SDLoc &DL) {
using namespace llvm::SDPatternMatch;

SDValue AddLhs;
APInt AddConst, SrlConst;
if (VT != MVT::i32 ||
!sd_match(N, m_AllOf(m_SpecificVT(MVT::i64),
m_Srl(m_OneUse(m_Add(m_Value(AddLhs),
m_ConstInt(AddConst))),
m_ConstInt(SrlConst)))))
return SDValue();
// Attempt to fold some (truncate (srl (binop X, C1), C2)) patterns to
// (binop (truncate (srl X, C2)), C1'). C1' will be smaller than C1 so we are
Copy link
Contributor

@phoebewang phoebewang Feb 27, 2025

Choose a reason for hiding this comment

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

We can use add/or/xor to replace binop

// able to avoid generating code with MOVABS and large constants in certain
// cases.
static SDValue combinei64TruncSrlBinop(SDValue N, EVT VT, SelectionDAG &DAG,
const SDLoc &DL) {

SDValue Binop = N.getOperand(0);
APInt BinopConst = Binop.getConstantOperandAPInt(1);
APInt SrlConst = N.getConstantOperandAPInt(1);
unsigned Opcode = Binop.getOpcode();

auto CleanUpFn = +[](SDValue Op, EVT CleanUpVT, EVT VT, SelectionDAG &DAG,
Copy link
Contributor

@phoebewang phoebewang Feb 28, 2025

Choose a reason for hiding this comment

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

No need to use switch since we have limit to ADD/OR/XOR in combineTruncatedArithmetic.

const SDLoc &DL) {
SDValue CleanUp = DAG.getAnyExtOrTrunc(Op, DL, CleanUpVT);
return DAG.getAnyExtOrTrunc(CleanUp, DL, VT);
};
auto ZeroExtCleanUp = +[](SDValue Op, EVT CleanUpVT, EVT VT,
SelectionDAG &DAG, const SDLoc &DL) {
return DAG.getZeroExtendInReg(Op, DL, CleanUpVT);
};
Copy link
Contributor

@phoebewang phoebewang Feb 27, 2025

Choose a reason for hiding this comment

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

The lambda makes the readability worse. I'd like a simple if/else or not use AnyExt.


Copy link
Contributor

@phoebewang phoebewang Feb 24, 2025

Choose a reason for hiding this comment

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

For or/xor they don't need the second condition and without the limition https://alive2.llvm.org/ce/z/uEAYxG, the and I showed in https://godbolt.org/z/1za939PKc should also work.

Copy link
Contributor Author

@joaotgouveia joaotgouveia Feb 24, 2025

Choose a reason for hiding this comment

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

I've added test cases that cover the scenario you mentioned for AND, as well as equivalent cases for OR and XOR. It seems the fold works as-is for these cases, though I'm not entirely sure why.

Copy link
Contributor

@phoebewang phoebewang Feb 25, 2025

Choose a reason for hiding this comment

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

You are right. I mistook it with add.

if (SrlConst.ule(32) || AddConst.countr_zero() < SrlConst.getZExtValue())
switch (Opcode) {
default:
return SDValue();
case ISD::ADD:
if (BinopConst.countr_zero() < SrlConst.getZExtValue())
return SDValue();
CleanUpFn = ZeroExtCleanUp;
[[fallthrough]];
case ISD::OR:
Copy link
Collaborator

@RKSimon RKSimon Feb 27, 2025

Choose a reason for hiding this comment

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

APInt::extractBits ?

Copy link
Contributor Author

@joaotgouveia joaotgouveia Feb 27, 2025

Choose a reason for hiding this comment

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

If I'm not mistaken, by using APInt::extractBits the constant transformation becomes
APInt NewOpConstVal = OpConst.extractBits(64 - SrlConst.getZExtValue(), SrlConst.getZExtValue()).zext(VT.getSizeInBits());. I believe the current implementation is simpler and easier to understand.

case ISD::XOR:
if (SrlConst.ule(32))
return SDValue();
break;
}

SDValue AddLHSSrl =
DAG.getNode(ISD::SRL, DL, MVT::i64, AddLhs, N.getOperand(1));
SDValue Trunc = DAG.getNode(ISD::TRUNCATE, DL, VT, AddLHSSrl);

APInt NewAddConstVal = AddConst.lshr(SrlConst).trunc(VT.getSizeInBits());
SDValue NewAddConst = DAG.getConstant(NewAddConstVal, DL, VT);
SDValue NewAddNode = DAG.getNode(ISD::ADD, DL, VT, Trunc, NewAddConst);
SDValue BinopLHSSrl =
DAG.getNode(ISD::SRL, DL, MVT::i64, Binop.getOperand(0), N.getOperand(1));
SDValue Trunc = DAG.getNode(ISD::TRUNCATE, DL, VT, BinopLHSSrl);

APInt NewBinopConstVal = BinopConst.lshr(SrlConst).trunc(VT.getSizeInBits());
SDValue NewBinopConst = DAG.getConstant(NewBinopConstVal, DL, VT);
SDValue NewBinopNode = DAG.getNode(Opcode, DL, VT, Trunc, NewBinopConst);
Copy link
Contributor

@phoebewang phoebewang Feb 27, 2025

Choose a reason for hiding this comment

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

I think we can remove Binop, it's a bit misleading that works for all binary operators.

EVT CleanUpVT =
EVT::getIntegerVT(*DAG.getContext(), 64 - SrlConst.getZExtValue());
return DAG.getZeroExtendInReg(NewAddNode, DL, CleanUpVT);
return CleanUpFn(NewBinopNode, CleanUpVT, VT, DAG, DL);
}

/// Attempt to pre-truncate inputs to arithmetic ops if it will simplify
Expand Down Expand Up @@ -53810,11 +53827,15 @@ static SDValue combineTruncatedArithmetic(SDNode *N, SelectionDAG &DAG,
if (!Src.hasOneUse())
return SDValue();

if (SDValue R = combinei64TruncSrlAdd(Src, VT, DAG, DL))
return R;
if (VT == MVT::i32 && SrcVT == MVT::i64 && SrcOpcode == ISD::SRL &&
Src.getOperand(0).getNumOperands() == 2 &&
Copy link
Contributor

@phoebewang phoebewang Feb 27, 2025

Choose a reason for hiding this comment

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

Do not need to check NumOperands = 2.

Copy link
Contributor Author

@joaotgouveia joaotgouveia Feb 27, 2025

Choose a reason for hiding this comment

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

Removing that check causes LLVM to crash on a few regression tests (CodeGen/X86/arg-cast.ll, CodeGen/X86/pr49162.ll, CodeGen/X86/vector-reduce-xor-bool.ll, and ExecutionEngine/MCJIT/non-extern-addend.ll). I might be mistaken, but at the time combineTruncatedArithmetic is called, we have no information about whether the first operand of Src is actually a binary operation. Calling Src.getOperand(0).getOperand(1) might cause us to attempt to access an invalid operand.

Copy link
Contributor

@phoebewang phoebewang Feb 27, 2025

Choose a reason for hiding this comment

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

Oh, right! We can check opcode = ADD/OR/AND before it then.

isa<ConstantSDNode>(Src.getOperand(1)) &&
isa<ConstantSDNode>(Src.getOperand(0).getOperand(1))) {
if (SDValue R = combinei64TruncSrlBinop(Src, VT, DAG, DL))
Copy link
Contributor

@phoebewang phoebewang Feb 27, 2025

Choose a reason for hiding this comment

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

Maybe combinei64TruncSrlConstant?

return R;
return SDValue();
}

// Only support vector truncation for now.
// TODO: i64 scalar math would benefit as well.
if (!VT.isVector())
return SDValue();

Expand Down
101 changes: 99 additions & 2 deletions llvm/test/CodeGen/X86/combine-i64-trunc-srl-add.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,103 @@ define i32 @test_trunc_add(i64 %x) {
ret i32 %conv
}

define i32 @test_trunc_sub(i64 %x) {
; X64-LABEL: test_trunc_sub:
; X64: # %bb.0:
; X64-NEXT: shrq $48, %rdi
; X64-NEXT: addl $65522, %edi # imm = 0xFFF2
; X64-NEXT: movzwl %di, %eax
; X64-NEXT: retq
%sub = sub i64 %x, 3940649673949184
%shr = lshr i64 %sub, 48
Copy link
Contributor

@phoebewang phoebewang Feb 28, 2025

Choose a reason for hiding this comment

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

Let's avoid to use 48 to test for the problem #128353 solving.

Copy link
Contributor Author

@joaotgouveia joaotgouveia Feb 28, 2025

Choose a reason for hiding this comment

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

Sure. Is changing one test enough, or should I change more than one?

Copy link
Contributor

@phoebewang phoebewang Mar 1, 2025

Choose a reason for hiding this comment

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

I'd like change them all.

%conv = trunc i64 %shr to i32
ret i32 %conv
}

define i32 @test_trunc_and_1(i64 %x) {
; X64-LABEL: test_trunc_and_1:
; X64: # %bb.0:
; X64-NEXT: movq %rdi, %rax
; X64-NEXT: shrq $48, %rax
; X64-NEXT: andl $14, %eax
; X64-NEXT: # kill: def $eax killed $eax killed $rax
; X64-NEXT: retq
%and = and i64 %x, 3940649673949184
%shr = lshr i64 %and, 48
%conv = trunc i64 %shr to i32
ret i32 %conv
}

define i32 @test_trunc_or_1(i64 %x) {
; X64-LABEL: test_trunc_or_1:
; X64: # %bb.0:
; X64-NEXT: movq %rdi, %rax
; X64-NEXT: shrq $48, %rax
; X64-NEXT: orl $14, %eax
; X64-NEXT: # kill: def $eax killed $eax killed $rax
; X64-NEXT: retq
%or = or i64 %x, 3940649673949184
%shr = lshr i64 %or, 48
%conv = trunc i64 %shr to i32
ret i32 %conv
}

define i32 @test_trunc_xor_1(i64 %x) {
; X64-LABEL: test_trunc_xor_1:
; X64: # %bb.0:
; X64-NEXT: movq %rdi, %rax
; X64-NEXT: shrq $48, %rax
; X64-NEXT: xorl $14, %eax
; X64-NEXT: # kill: def $eax killed $eax killed $rax
; X64-NEXT: retq
%xor = xor i64 %x, 3940649673949184
%shr = lshr i64 %xor, 48
%conv = trunc i64 %shr to i32
ret i32 %conv
}

define i32 @test_trunc_and_2(i64 %x) {
; X64-LABEL: test_trunc_and_2:
; X64: # %bb.0:
; X64-NEXT: movq %rdi, %rax
; X64-NEXT: shrq $48, %rax
; X64-NEXT: andl $13, %eax
; X64-NEXT: # kill: def $eax killed $eax killed $rax
; X64-NEXT: retq
%and = and i64 %x, 3940649673949183
%shr = lshr i64 %and, 48
%conv = trunc i64 %shr to i32
ret i32 %conv
}

define i32 @test_trunc_or_2(i64 %x) {
; X64-LABEL: test_trunc_or_2:
; X64: # %bb.0:
; X64-NEXT: movq %rdi, %rax
; X64-NEXT: shrq $48, %rax
; X64-NEXT: orl $13, %eax
; X64-NEXT: # kill: def $eax killed $eax killed $rax
; X64-NEXT: retq
%or = or i64 %x, 3940649673949183
%shr = lshr i64 %or, 48
%conv = trunc i64 %shr to i32
ret i32 %conv
}

define i32 @test_trunc_xor_2(i64 %x) {
; X64-LABEL: test_trunc_xor_2:
; X64: # %bb.0:
; X64-NEXT: movq %rdi, %rax
; X64-NEXT: shrq $48, %rax
; X64-NEXT: xorl $13, %eax
; X64-NEXT: # kill: def $eax killed $eax killed $rax
; X64-NEXT: retq
%xor = xor i64 %x, 3940649673949183
%shr = lshr i64 %xor, 48
%conv = trunc i64 %shr to i32
ret i32 %conv
}

; Make sure we don't crash on this test case.

define i32 @pr128158(i64 %x) {
Expand All @@ -137,10 +234,10 @@ define i32 @pr128158(i64 %x) {
; X64-NEXT: addq %rdi, %rax
; X64-NEXT: shrq $32, %rax
; X64-NEXT: .p2align 4
; X64-NEXT: .LBB9_1: # %for.body
; X64-NEXT: .LBB16_1: # %for.body
; X64-NEXT: # =>This Inner Loop Header: Depth=1
; X64-NEXT: cmpl $9, %eax
; X64-NEXT: jb .LBB9_1
; X64-NEXT: jb .LBB16_1
; X64-NEXT: # %bb.2: # %exit
; X64-NEXT: xorl %eax, %eax
; X64-NEXT: retq
Expand Down