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

joaotgouveia · 2025-02-23T21:02:53Z

As discussed in #126448, the fold implemented by #126448 / #128353 can be extended to operations other than add. This patch extends the fold performed by combinei64TruncSrlAdd to include or and xor (proof: https://alive2.llvm.org/ce/z/AXuaQu). There's no need to extend it to sub and and, as similar folds are already being performed for those operations.

CC: @phoebewang @RKSimon

… `xor`

llvmbot · 2025-02-23T21:03:29Z

@llvm/pr-subscribers-backend-x86

Author: João Gouveia (joaotgouveia)

Changes

As discussed in #126448, the fold implemented by #126448 / #128353 can be extended to operations other than add. This patch extends the fold performed by combinei64TruncSrlAdd to include or (proof: https://alive2.llvm.org/ce/z/SDZgos) and xor (proof: https://alive2.llvm.org/ce/z/fW8MSx). There's no need to extend it to sub and and, as similar folds are already being performed for those operations.

CC: @phoebewang @RKSimon

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

2 Files Affected:

(modified) llvm/lib/Target/X86/X86ISelLowering.cpp (+26-22)
(modified) llvm/test/CodeGen/X86/combine-i64-trunc-srl-add.ll (+57-2)

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index c146e1e6c0334..47dc9ffb4b24d 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -53733,36 +53733,42 @@ 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) {
+// 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
+// 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) {
   using namespace llvm::SDPatternMatch;
 
-  SDValue AddLhs;
-  APInt AddConst, SrlConst;
+  SDValue BinopLhs;
+  APInt BinopConst, 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)))))
+      !sd_match(
+          N,
+          m_AllOf(m_SpecificVT(MVT::i64),
+                  m_Srl(m_OneUse(m_AnyOf(
+                            m_Add(m_Value(BinopLhs), m_ConstInt(BinopConst)),
+                            m_Or(m_Value(BinopLhs), m_ConstInt(BinopConst)),
+                            m_Xor(m_Value(BinopLhs), m_ConstInt(BinopConst)))),
+                        m_ConstInt(SrlConst)))))
     return SDValue();
 
-  if (SrlConst.ule(32) || AddConst.countr_zero() < SrlConst.getZExtValue())
+  if (SrlConst.ule(32) || BinopConst.countr_zero() < SrlConst.getZExtValue())
     return SDValue();
 
-  SDValue AddLHSSrl =
-      DAG.getNode(ISD::SRL, DL, MVT::i64, AddLhs, N.getOperand(1));
-  SDValue Trunc = DAG.getNode(ISD::TRUNCATE, DL, VT, AddLHSSrl);
+  SDValue BinopLHSSrl =
+      DAG.getNode(ISD::SRL, DL, MVT::i64, BinopLhs, N.getOperand(1));
+  SDValue Trunc = DAG.getNode(ISD::TRUNCATE, DL, VT, BinopLHSSrl);
 
-  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);
+  APInt NewBinopConstVal = BinopConst.lshr(SrlConst).trunc(VT.getSizeInBits());
+  SDValue NewBinopConst = DAG.getConstant(NewBinopConstVal, DL, VT);
+  SDValue NewBinopNode =
+      DAG.getNode(N.getOperand(0).getOpcode(), DL, VT, Trunc, NewBinopConst);
 
   EVT CleanUpVT =
       EVT::getIntegerVT(*DAG.getContext(), 64 - SrlConst.getZExtValue());
-  return DAG.getZeroExtendInReg(NewAddNode, DL, CleanUpVT);
+  return DAG.getZeroExtendInReg(NewBinopNode, DL, CleanUpVT);
 }
 
 /// Attempt to pre-truncate inputs to arithmetic ops if it will simplify
@@ -53810,11 +53816,9 @@ static SDValue combineTruncatedArithmetic(SDNode *N, SelectionDAG &DAG,
   if (!Src.hasOneUse())
     return SDValue();
 
-  if (SDValue R = combinei64TruncSrlAdd(Src, VT, DAG, DL))
+  if (SDValue R = combinei64TruncSrlBinop(Src, VT, DAG, DL))
     return R;
 
-  // Only support vector truncation for now.
-  // TODO: i64 scalar math would benefit as well.
   if (!VT.isVector())
     return SDValue();
 
diff --git a/llvm/test/CodeGen/X86/combine-i64-trunc-srl-add.ll b/llvm/test/CodeGen/X86/combine-i64-trunc-srl-add.ll
index 14992ca5bf488..ec29cf9d56c29 100644
--- a/llvm/test/CodeGen/X86/combine-i64-trunc-srl-add.ll
+++ b/llvm/test/CodeGen/X86/combine-i64-trunc-srl-add.ll
@@ -128,6 +128,61 @@ 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
+  %conv = trunc i64 %shr to i32
+  ret i32 %conv
+}
+
+define i32 @test_trunc_and(i64 %x) {
+; X64-LABEL: test_trunc_and:
+; 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(i64 %x) {
+; X64-LABEL: test_trunc_or:
+; 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(i64 %x) {
+; X64-LABEL: test_trunc_xor:
+; 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
+}
+
 ; Make sure we don't crash on this test case.
 
 define i32 @pr128158(i64 %x) {
@@ -137,10 +192,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:  .LBB13_1: # %for.body
 ; X64-NEXT:    # =>This Inner Loop Header: Depth=1
 ; X64-NEXT:    cmpl $9, %eax
-; X64-NEXT:    jb .LBB9_1
+; X64-NEXT:    jb .LBB13_1
 ; X64-NEXT:  # %bb.2: # %exit
 ; X64-NEXT:    xorl %eax, %eax
 ; X64-NEXT:    retq

phoebewang · 2025-02-24T00:32:41Z

llvm/lib/Target/X86/X86ISelLowering.cpp

    return SDValue();

-  if (SrlConst.ule(32) || AddConst.countr_zero() < SrlConst.getZExtValue())
+  if (SrlConst.ule(32) || BinopConst.countr_zero() < SrlConst.getZExtValue())


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.

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.

You are right. I mistook it with add.

llvm/lib/Target/X86/X86ISelLowering.cpp

phoebewang · 2025-02-24T00:36:57Z

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -53810,11 +53816,9 @@ static SDValue combineTruncatedArithmetic(SDNode *N, SelectionDAG &DAG,
  if (!Src.hasOneUse())
    return SDValue();

-  if (SDValue R = combinei64TruncSrlAdd(Src, VT, DAG, DL))
+  if (SDValue R = combinei64TruncSrlBinop(Src, VT, DAG, DL))


Does it work for MUL? Add test case if yes and exclude it explicitly in code if not?

We also have other Binop like min/max etc., so we should add switch check like below.

This exact fold doesn't work for MUL: https://alive2.llvm.org/ce/z/Xoz8hb (had to reduce the size of the ints because Alive was timing out).

phoebewang · 2025-02-24T00:43:28Z

llvm/lib/Target/X86/X86ISelLowering.cpp

+                            m_Add(m_Value(BinopLhs), m_ConstInt(BinopConst)),
+                            m_Or(m_Value(BinopLhs), m_ConstInt(BinopConst)),
+                            m_Xor(m_Value(BinopLhs), m_ConstInt(BinopConst)))),


Use m_Binary once we limited the binop to add/sub/and/or/xor and/or mul.

I think we can get rid of SDPatternMatch by checking VT == MVT::i32 && SrcVT == MVT::i64 && isa<ConstantSDNode>(N->getOperand(1)) && isa<ConstantSDNode>(Src.getOperand(1)) before calling combinei64TruncSrlBinop.

Wouldn't it be neater to move those checks inside combinei64TruncSrlBinop? We could use m_Binary to match the pattern and then switch case the opcode to perform the different checks required for AND/OR/XOR and ADD/SUB.

My assumption is when we have only one exact match, either m_Add or m_Binary, would be good. If neither match, we don't bother it, because the check for OR/XOR/ADD/SUB is enough.

I've changed the function according to your suggestions. It could probably be simplified if we always clean up using getZeroExtendInReg, as it seems that for OR and XOR the movz is not generated, regardless of whether we’re zero-extending or any-extending.

phoebewang · 2025-02-27T03:32:33Z

llvm/lib/Target/X86/X86ISelLowering.cpp

-  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 &&


Do not need to check NumOperands = 2.

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.

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

phoebewang · 2025-02-27T03:32:57Z

llvm/lib/Target/X86/X86ISelLowering.cpp

+      Src.getOperand(0).getNumOperands() == 2 &&
+      isa<ConstantSDNode>(Src.getOperand(1)) &&
+      isa<ConstantSDNode>(Src.getOperand(0).getOperand(1))) {
+    if (SDValue R = combinei64TruncSrlBinop(Src, VT, DAG, DL))


Maybe combinei64TruncSrlConstant?

phoebewang · 2025-02-27T03:34:27Z

llvm/lib/Target/X86/X86ISelLowering.cpp

+// 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


We can use add/or/xor to replace binop

phoebewang · 2025-02-27T03:36:47Z

llvm/lib/Target/X86/X86ISelLowering.cpp

+  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);


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

phoebewang · 2025-02-27T03:38:06Z

llvm/lib/Target/X86/X86ISelLowering.cpp

+  auto CleanUpFn = +[](SDValue Op, EVT CleanUpVT, EVT VT, SelectionDAG &DAG,
+                       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);
+  };


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

RKSimon · 2025-02-27T11:10:36Z

llvm/lib/Target/X86/X86ISelLowering.cpp


+  APInt NewOpConstVal = OpConst.lshr(SrlConst).trunc(VT.getSizeInBits());


APInt::extractBits ?

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.

phoebewang · 2025-02-28T02:45:11Z

llvm/lib/Target/X86/X86ISelLowering.cpp

+  SDValue CleanUp = DAG.getAnyExtOrTrunc(NewOpNode, DL, CleanUpVT);
+  return DAG.getAnyExtOrTrunc(CleanUp, DL, VT);


Can we return DAG.getAnyExtOrTrunc(NewOpNode, DL, VT); directly?

That causes a crash because of an assert in SelectionDAG::ReplaceAllUsesWith.

I see the problem, so we should return NewOpNode directly because we don't need to clean up the upper bits.

Proof: https://alive2.llvm.org/ce/z/AXuaQu

phoebewang · 2025-02-28T02:47:26Z

llvm/lib/Target/X86/X86ISelLowering.cpp


-  if (SrlConst.ule(32) || AddConst.countr_zero() < SrlConst.getZExtValue())
+  switch (Opcode) {


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

phoebewang · 2025-02-28T13:13:59Z

llvm/test/CodeGen/X86/combine-i64-trunc-srl-add.ll

+; X64-NEXT:    movzwl %di, %eax
+; X64-NEXT:    retq
+  %sub = sub i64 %x, 3940649673949184
+  %shr = lshr i64 %sub, 48


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

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

I'd like change them all.

phoebewang · 2025-03-01T02:16:55Z

llvm/lib/Target/X86/X86ISelLowering.cpp

+  SDValue CleanUp = DAG.getAnyExtOrTrunc(NewOpNode, DL, CleanUpVT);
+  return DAG.getAnyExtOrTrunc(CleanUp, DL, VT);


return NewOpNode;

phoebewang

LGTM.

@phoebewang

… `xor` (llvm#128435) As discussed in llvm#126448, the fold implemented by llvm#126448 / llvm#128353 can be extended to operations other than `add`. This patch extends the fold performed by `combinei64TruncSrlAdd` to include `or` and `xor` (proof: https://alive2.llvm.org/ce/z/AXuaQu). There's no need to extend it to `sub` and `and`, as similar folds are already being performed for those operations. CC: @phoebewang @RKSimon

joaotgouveia added 2 commits February 23, 2025 18:46

[X86] Add pre-commit tests.

fe8151f

[X86] Extend combinei64TruncSrlAdd to handle patterns with or and…

5574dce

… `xor`

llvmbot added the backend:X86 label Feb 23, 2025

phoebewang reviewed Feb 24, 2025

View reviewed changes

llvm/lib/Target/X86/X86ISelLowering.cpp Outdated Show resolved Hide resolved

phoebewang reviewed Feb 24, 2025

View reviewed changes

joaotgouveia added 2 commits February 24, 2025 17:06

[X86] Add extra test cases

55189cc

[X86] Address reviews

d2f27bf

phoebewang reviewed Feb 27, 2025

View reviewed changes

phoebewang requested review from RKSimon and dtcxzyw February 27, 2025 03:38

dtcxzyw mentioned this pull request Feb 27, 2025

Fuzz PR128435 dtcxzyw/llvm-fuzz-service#61

Closed

[X86] Improve readability

1ba8e32

RKSimon reviewed Feb 27, 2025

View reviewed changes

[X86] Remove getNumOperands

696f3a5

phoebewang reviewed Feb 28, 2025

View reviewed changes

[X86] Remove unnecessary switch case

7caceb1

phoebewang reviewed Feb 28, 2025

View reviewed changes

phoebewang reviewed Mar 1, 2025

View reviewed changes

[X86] Address reviews

4388e73

phoebewang approved these changes Mar 1, 2025

View reviewed changes

phoebewang merged commit 0751418 into llvm:main Mar 1, 2025
11 checks passed

		// 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


		APInt NewOpConstVal = OpConst.lshr(SrlConst).trunc(VT.getSizeInBits());

		SDValue CleanUp = DAG.getAnyExtOrTrunc(NewOpNode, DL, CleanUpVT);
		return DAG.getAnyExtOrTrunc(CleanUp, DL, VT);


		if (SrlConst.ule(32) \|\| AddConst.countr_zero() < SrlConst.getZExtValue())
		switch (Opcode) {

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

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

Uh oh!

Conversation

joaotgouveia commented Feb 23, 2025 • edited by phoebewang Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Feb 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phoebewang Feb 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phoebewang left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

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

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

joaotgouveia commented Feb 23, 2025 •

edited by phoebewang

Loading

phoebewang Feb 28, 2025 •

edited

Loading