[llvm] [DAGCombiner] Expose CSE for shifts of logic operations (PR #213091)

Hao Ren via llvm-commits llvm-commits at lists.llvm.org
Thu Aug 6 13:48:43 PDT 2026


https://github.com/nvidia-moomoo updated https://github.com/llvm/llvm-project/pull/213091

>From 07715a6ca391659840bb039be0e9539266a8d9be Mon Sep 17 00:00:00 2001
From: Hao Ren <haor at nvidia.com>
Date: Thu, 6 Aug 2026 18:40:12 +0000
Subject: [PATCH 1/3] [DAGCombiner][NFC] Add pre-commit tests for
 shift-of-logic CSE

Capture current NVPTX codegen for shifts of AND/OR/XOR operations that share
an input and shift amount, ahead of the combine that exposes the shared shift
for CSE. Covers no-extend, zext, and sext variants, the address-arithmetic
(getelementptr) form, and an isolated negative case.
---
 llvm/test/CodeGen/NVPTX/shift-logic-cse.ll | 155 +++++++++++++++++++++
 1 file changed, 155 insertions(+)
 create mode 100644 llvm/test/CodeGen/NVPTX/shift-logic-cse.ll

diff --git a/llvm/test/CodeGen/NVPTX/shift-logic-cse.ll b/llvm/test/CodeGen/NVPTX/shift-logic-cse.ll
new file mode 100644
index 0000000000000..4a519b6fb7f52
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/shift-logic-cse.ll
@@ -0,0 +1,155 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc < %s -mtriple=nvptx64 | FileCheck %s
+
+; Three logic ops (and/or/xor) share %x and the same shift amount, with no
+; extend. Commuting the shifts through the logic ops exposes a single shifted
+; %x that is CSEd, replacing three shifts with one.
+define void @mixed_logic_siblings(i32 %x, ptr %p0, ptr %p1, ptr %p2) {
+; CHECK-LABEL: mixed_logic_siblings(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b32 %r<8>;
+; CHECK-NEXT:    .reg .b64 %rd<4>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    ld.param.b32 %r1, [mixed_logic_siblings_param_0];
+; CHECK-NEXT:    and.b32 %r2, %r1, 5;
+; CHECK-NEXT:    ld.param.b64 %rd1, [mixed_logic_siblings_param_1];
+; CHECK-NEXT:    shl.b32 %r3, %r2, 2;
+; CHECK-NEXT:    ld.param.b64 %rd2, [mixed_logic_siblings_param_2];
+; CHECK-NEXT:    st.b32 [%rd1], %r3;
+; CHECK-NEXT:    ld.param.b64 %rd3, [mixed_logic_siblings_param_3];
+; CHECK-NEXT:    shl.b32 %r4, %r1, 2;
+; CHECK-NEXT:    or.b32 %r5, %r4, 28;
+; CHECK-NEXT:    st.b32 [%rd2], %r5;
+; CHECK-NEXT:    xor.b32 %r6, %r1, 17;
+; CHECK-NEXT:    shl.b32 %r7, %r6, 2;
+; CHECK-NEXT:    st.b32 [%rd3], %r7;
+; CHECK-NEXT:    ret;
+  %and = and i32 %x, 5
+  %and.shl = shl i32 %and, 2
+  store i32 %and.shl, ptr %p0, align 4
+  %or = or i32 %x, 7
+  %or.shl = shl i32 %or, 2
+  store i32 %or.shl, ptr %p1, align 4
+  %xor = xor i32 %x, 17
+  %xor.shl = shl i32 %xor, 2
+  store i32 %xor.shl, ptr %p2, align 4
+  ret void
+}
+
+; zext siblings: shl (zext (and X, C)), S. The shared shl (zext X), S is
+; exposed and CSEd.
+define void @zext_and_siblings(i32 %x, ptr %p0, ptr %p1) {
+; CHECK-LABEL: zext_and_siblings(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b32 %r<4>;
+; CHECK-NEXT:    .reg .b64 %rd<5>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    ld.param.b32 %r1, [zext_and_siblings_param_0];
+; CHECK-NEXT:    and.b32 %r2, %r1, 5;
+; CHECK-NEXT:    ld.param.b64 %rd1, [zext_and_siblings_param_1];
+; CHECK-NEXT:    ld.param.b64 %rd2, [zext_and_siblings_param_2];
+; CHECK-NEXT:    mul.wide.u32 %rd3, %r2, 8;
+; CHECK-NEXT:    st.b64 [%rd1], %rd3;
+; CHECK-NEXT:    and.b32 %r3, %r1, 6;
+; CHECK-NEXT:    mul.wide.u32 %rd4, %r3, 8;
+; CHECK-NEXT:    st.b64 [%rd2], %rd4;
+; CHECK-NEXT:    ret;
+  %and0 = and i32 %x, 5
+  %z0 = zext i32 %and0 to i64
+  %s0 = shl i64 %z0, 3
+  store i64 %s0, ptr %p0, align 8
+  %and1 = and i32 %x, 6
+  %z1 = zext i32 %and1 to i64
+  %s1 = shl i64 %z1, 3
+  store i64 %s1, ptr %p1, align 8
+  ret void
+}
+
+; sext siblings: shl (sext (or X, C)), S.
+define void @sext_or_siblings(i32 %x, ptr %p0, ptr %p1) {
+; CHECK-LABEL: sext_or_siblings(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b32 %r<4>;
+; CHECK-NEXT:    .reg .b64 %rd<5>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    ld.param.b32 %r1, [sext_or_siblings_param_0];
+; CHECK-NEXT:    or.b32 %r2, %r1, 5;
+; CHECK-NEXT:    ld.param.b64 %rd1, [sext_or_siblings_param_1];
+; CHECK-NEXT:    ld.param.b64 %rd2, [sext_or_siblings_param_2];
+; CHECK-NEXT:    mul.wide.s32 %rd3, %r2, 4;
+; CHECK-NEXT:    st.b64 [%rd1], %rd3;
+; CHECK-NEXT:    or.b32 %r3, %r1, 6;
+; CHECK-NEXT:    mul.wide.s32 %rd4, %r3, 4;
+; CHECK-NEXT:    st.b64 [%rd2], %rd4;
+; CHECK-NEXT:    ret;
+  %or0 = or i32 %x, 5
+  %e0 = sext i32 %or0 to i64
+  %s0 = shl i64 %e0, 2
+  store i64 %s0, ptr %p0, align 8
+  %or1 = or i32 %x, 6
+  %e1 = sext i32 %or1 to i64
+  %s1 = shl i64 %e1, 2
+  store i64 %s1, ptr %p1, align 8
+  ret void
+}
+
+; The shift is implicit here: the getelementptr's element-size scaling lowers
+; to shl during SelectionDAG, so ordinary address arithmetic exercises the same
+; pattern. NVPTX does not fold the scaled index into the load, so the combine
+; still applies and the shared index computation is CSEd.
+define void @gep_index_siblings(i32 %base, ptr %data, ptr %p0, ptr %p1) {
+; CHECK-LABEL: gep_index_siblings(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b32 %r<6>;
+; CHECK-NEXT:    .reg .b64 %rd<8>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    ld.param.b32 %r1, [gep_index_siblings_param_0];
+; CHECK-NEXT:    xor.b32 %r2, %r1, 5;
+; CHECK-NEXT:    ld.param.b64 %rd1, [gep_index_siblings_param_1];
+; CHECK-NEXT:    ld.param.b64 %rd2, [gep_index_siblings_param_2];
+; CHECK-NEXT:    mul.wide.u32 %rd3, %r2, 4;
+; CHECK-NEXT:    add.s64 %rd4, %rd1, %rd3;
+; CHECK-NEXT:    ld.param.b64 %rd5, [gep_index_siblings_param_3];
+; CHECK-NEXT:    ld.b32 %r3, [%rd4];
+; CHECK-NEXT:    st.b32 [%rd2], %r3;
+; CHECK-NEXT:    xor.b32 %r4, %r1, 6;
+; CHECK-NEXT:    mul.wide.u32 %rd6, %r4, 4;
+; CHECK-NEXT:    add.s64 %rd7, %rd1, %rd6;
+; CHECK-NEXT:    ld.b32 %r5, [%rd7];
+; CHECK-NEXT:    st.b32 [%rd5], %r5;
+; CHECK-NEXT:    ret;
+  %x0 = xor i32 %base, 5
+  %i0 = zext i32 %x0 to i64
+  %a0 = getelementptr i32, ptr %data, i64 %i0
+  %v0 = load i32, ptr %a0, align 4
+  store i32 %v0, ptr %p0, align 4
+  %x1 = xor i32 %base, 6
+  %i1 = zext i32 %x1 to i64
+  %a1 = getelementptr i32, ptr %data, i64 %i1
+  %v1 = load i32, ptr %a1, align 4
+  store i32 %v1, ptr %p1, align 4
+  ret void
+}
+
+; Negative test: a lone shl of a logic op has no sibling to share a shift with,
+; so it is left alone. Preserving it lets targets that support it fold the shift
+; into an addressing mode; NVPTX simply keeps the single shift.
+define i32 @isolated_no_sibling(i32 %x) {
+; CHECK-LABEL: isolated_no_sibling(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b32 %r<4>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    ld.param.b32 %r1, [isolated_no_sibling_param_0];
+; CHECK-NEXT:    xor.b32 %r2, %r1, 5;
+; CHECK-NEXT:    shl.b32 %r3, %r2, 2;
+; CHECK-NEXT:    st.param.b32 [func_retval0], %r3;
+; CHECK-NEXT:    ret;
+  %xor = xor i32 %x, 5
+  %shl = shl i32 %xor, 2
+  ret i32 %shl
+}

>From ba9cf46eab958b2e4422c47e5636a7472528f9b2 Mon Sep 17 00:00:00 2001
From: Hao Ren <haor at nvidia.com>
Date: Fri, 24 Jul 2026 20:18:44 +0000
Subject: [PATCH 2/3] [DAGCombiner] Expose CSE for shifts of logic operations

When multiple AND, OR, or XOR operations with constant operands share the same input, commute outer SHLs through the operations so the common shifted input can be CSEd. Handle optional integer extends atomically across all eligible siblings.

Gate the combine behind a target hook because commuting the shifts can disrupt target-specific patterns such as addressing modes. Enable it for NVPTX, where reducing N shifts to one is profitable.
---
 llvm/include/llvm/CodeGen/TargetLowering.h    |   9 ++
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 132 ++++++++++++++++++
 llvm/lib/Target/NVPTX/NVPTXISelLowering.h     |   5 +
 llvm/test/CodeGen/NVPTX/shift-logic-cse.ll    |  71 +++++-----
 4 files changed, 179 insertions(+), 38 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/TargetLowering.h b/llvm/include/llvm/CodeGen/TargetLowering.h
index c663bb8ea65b7..bfd24f99d34d7 100644
--- a/llvm/include/llvm/CodeGen/TargetLowering.h
+++ b/llvm/include/llvm/CodeGen/TargetLowering.h
@@ -4661,6 +4661,15 @@ class LLVM_ABI TargetLowering : public TargetLoweringBase {
     return true;
   }
 
+  /// Return true if it is profitable to commute a constant SHL through
+  /// multiple bitwise logic operations that share an input. This can expose a
+  /// common shifted operand, but may also prevent shifts from folding into
+  /// addressing modes, so targets must opt in.
+  virtual bool shouldCommuteShiftOfLogicToExposeCSE(const SDNode *N,
+                                                    CombineLevel Level) const {
+    return false;
+  }
+
   /// GlobalISel - return true if it is profitable to move this shift by a
   /// constant amount through its operand, adjusting any immediate operands as
   /// necessary to preserve semantics. This transformation may not be desirable
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 232cd609251db..277ab25775433 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -590,6 +590,7 @@ namespace {
                                  EVT VT, SDValue N0, SDValue N1,
                                  SDNodeFlags Flags = SDNodeFlags());
 
+    SDValue foldShiftOfLogicWithSharedInput(SDNode *N);
     SDValue visitShiftByConstant(SDNode *N);
 
     SDValue foldSelectOfConstants(SDNode *N);
@@ -10776,6 +10777,133 @@ static SDValue combineShiftOfShiftedLogic(SDNode *Shift, SelectionDAG &DAG) {
                      LogicOp->getFlags());
 }
 
+/// Commute SHL with a bitwise logic operation when doing so exposes a common
+/// shifted operand. For example:
+///
+/// Before:
+///   N            = shl (zext (Logic X, C)), ShiftAmount
+///   SiblingShift = shl (zext (SiblingLogic X, SiblingC)), ShiftAmount
+///
+/// After:
+///   ShiftedX     = shl (zext X), ShiftAmount
+///   N            = Logic ShiftedX, ShiftedC
+///   SiblingShift = SiblingLogic ShiftedX, SiblingShiftedC
+///
+/// ShiftedC = (zext C) << ShiftAmount and SiblingShiftedC =
+/// (zext SiblingC) << ShiftAmount are folded constants. This replaces two
+/// variable shifts with the single shared ShiftedX. Requiring a matching
+/// sibling avoids disrupting isolated address calculations where a shift may be
+/// folded into the addressing mode.
+SDValue DAGCombiner::foldShiftOfLogicWithSharedInput(SDNode *N) {
+  struct ShiftOfLogic {
+    SDNode *Shift;
+    SDValue Logic;
+    SDValue X;
+    SDValue Constant;
+    unsigned ExtendOpcode;
+  };
+
+  // Match a logic operation, with an optional extension, inside a SHL.
+  auto matchShiftOfLogic = [&](SDNode *Shift, ShiftOfLogic &Match) {
+    if (Shift->getOpcode() != ISD::SHL)
+      return false;
+
+    SDValue ShiftInput = Shift->getOperand(0);
+    if (!ShiftInput.hasOneUse())
+      return false;
+
+    unsigned ExtendOpcode = 0;
+    SDValue Logic = ShiftInput;
+    if (ISD::isExtOpcode(ShiftInput.getOpcode())) {
+      ExtendOpcode = ShiftInput.getOpcode();
+      Logic = ShiftInput.getOperand(0);
+    }
+
+    SDValue X, C;
+    if (!sd_match(Logic, m_OneUse(m_BitwiseLogic(m_Value(X),
+                                                 m_Value(C, m_ConstInt())))))
+      return false;
+    ConstantSDNode *Constant = isConstOrConstSplat(C);
+    if (!Constant || Constant->isOpaque())
+      return false;
+
+    Match = {Shift, Logic, X, C, ExtendOpcode};
+    return true;
+  };
+
+  // Match N as the root shift-of-logic; bail if it does not fit the pattern or
+  // should not be commuted.
+  ShiftOfLogic Root;
+  if (!matchShiftOfLogic(N, Root) ||
+      !TLI.shouldCommuteShiftOfLogicToExposeCSE(N, Level))
+    return SDValue();
+
+  EVT VT = N->getValueType(0);
+  SDValue ShiftAmount = N->getOperand(1);
+
+  // Collect candidate shifts that share X. Reached through another user of X,
+  // the logic result feeds the shift directly or through an optional extend.
+  SmallVector<SDNode *, 4> CandidateShifts;
+  for (SDNode *CandidateLogic : Root.X->users()) {
+    if (CandidateLogic == Root.Logic.getNode())
+      continue;
+    for (SDNode *LogicUser : CandidateLogic->users()) {
+      if (ISD::isExtOpcode(LogicUser->getOpcode())) {
+        // shl (ext (logic X, C)): step through the extend to find the shift.
+        for (SDNode *ExtendUser : LogicUser->users())
+          if (ExtendUser->getOpcode() == ISD::SHL)
+            CandidateShifts.push_back(ExtendUser);
+      } else if (LogicUser->getOpcode() == ISD::SHL) {
+        // shl (logic X, C): the user is already the shift.
+        CandidateShifts.push_back(LogicUser);
+      }
+    }
+  }
+
+  // Verify each candidate against the root's pattern; those that match (same X,
+  // extend, type, and shift amount) become siblings.
+  SmallVector<ShiftOfLogic, 4> Siblings;
+  for (SDNode *CandidateShift : CandidateShifts) {
+    ShiftOfLogic Candidate;
+    if (!matchShiftOfLogic(CandidateShift, Candidate) ||
+        Candidate.X != Root.X || Candidate.ExtendOpcode != Root.ExtendOpcode ||
+        CandidateShift->getValueType(0) != VT ||
+        CandidateShift->getOperand(1) != ShiftAmount ||
+        !TLI.isDesirableToCommuteWithShift(CandidateShift, Level))
+      continue;
+    Siblings.push_back(Candidate);
+  }
+  if (Siblings.empty())
+    return SDValue();
+
+  // Build the shared shifted X once, then rewrite the root and every sibling
+  // into a logic op over it so the shift is CSE'd.
+  SDValue ShiftedX = DAG.getNode(
+      ISD::SHL, SDLoc(N), VT,
+      Root.ExtendOpcode ? DAG.getNode(Root.ExtendOpcode, SDLoc(N), VT, Root.X)
+                        : Root.X,
+      ShiftAmount);
+
+  // Rebuild the logic op from shared ShiftedX and a folded constant shift.
+  auto rebuild = [&](SDValue OldLogic, SDValue OldC, const SDLoc &NodeDL) {
+    SDValue NewC = OldC;
+    if (Root.ExtendOpcode)
+      NewC = DAG.getNode(Root.ExtendOpcode, NodeDL, VT, OldC);
+    SDValue ShiftedC = DAG.getNode(ISD::SHL, NodeDL, VT, NewC, ShiftAmount);
+    return DAG.getNode(OldLogic.getOpcode(), NodeDL, VT, ShiftedX, ShiftedC,
+                       OldLogic->getFlags());
+  };
+
+  for (const ShiftOfLogic &Sibling : Siblings) {
+    SDValue NewLogic =
+        rebuild(Sibling.Logic, Sibling.Constant, SDLoc(Sibling.Shift));
+    DAG.ReplaceAllUsesOfValueWith(SDValue(Sibling.Shift, 0), NewLogic);
+    deleteAndRecombine(Sibling.Shift);
+    AddToWorklist(NewLogic.getNode());
+  }
+  return rebuild(Root.Logic, Root.Constant, SDLoc(N));
+}
+
 /// Handle transforms common to the three shifts, when the shift amount is a
 /// constant.
 /// We are looking for: (shift being one of shl/sra/srl)
@@ -10798,6 +10926,10 @@ SDValue DAGCombiner::visitShiftByConstant(SDNode *N) {
   if (SDValue R = combineShiftOfShiftedLogic(N, DAG))
     return R;
 
+  // Fold sibling shifts of logic operations to expose a shared shifted input.
+  if (SDValue R = foldShiftOfLogicWithSharedInput(N))
+    return R;
+
   // We want to pull some binops through shifts, so that we have (and (shift))
   // instead of (shift (and)), likewise for add, or, xor, etc.  This sort of
   // thing happens with address calculations, so it's important to canonicalize
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.h b/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
index 42741ba98b438..92ff7415dd6f5 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
@@ -181,6 +181,11 @@ class NVPTXTargetLowering : public TargetLowering {
                                          TargetLoweringOpt &TLO,
                                          unsigned Depth = 0) const override;
 
+  bool shouldCommuteShiftOfLogicToExposeCSE(const SDNode *,
+                                            CombineLevel) const override {
+    return true;
+  }
+
 private:
   const NVPTXSubtarget &STI; // cache the subtarget here
   mutable unsigned GlobalUniqueCallSite;
diff --git a/llvm/test/CodeGen/NVPTX/shift-logic-cse.ll b/llvm/test/CodeGen/NVPTX/shift-logic-cse.ll
index 4a519b6fb7f52..1301ea220eea1 100644
--- a/llvm/test/CodeGen/NVPTX/shift-logic-cse.ll
+++ b/llvm/test/CodeGen/NVPTX/shift-logic-cse.ll
@@ -7,23 +7,21 @@
 define void @mixed_logic_siblings(i32 %x, ptr %p0, ptr %p1, ptr %p2) {
 ; CHECK-LABEL: mixed_logic_siblings(
 ; CHECK:       {
-; CHECK-NEXT:    .reg .b32 %r<8>;
+; CHECK-NEXT:    .reg .b32 %r<6>;
 ; CHECK-NEXT:    .reg .b64 %rd<4>;
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  // %bb.0:
 ; CHECK-NEXT:    ld.param.b32 %r1, [mixed_logic_siblings_param_0];
-; CHECK-NEXT:    and.b32 %r2, %r1, 5;
 ; CHECK-NEXT:    ld.param.b64 %rd1, [mixed_logic_siblings_param_1];
-; CHECK-NEXT:    shl.b32 %r3, %r2, 2;
+; CHECK-NEXT:    shl.b32 %r2, %r1, 2;
+; CHECK-NEXT:    and.b32 %r3, %r2, 20;
 ; CHECK-NEXT:    ld.param.b64 %rd2, [mixed_logic_siblings_param_2];
 ; CHECK-NEXT:    st.b32 [%rd1], %r3;
 ; CHECK-NEXT:    ld.param.b64 %rd3, [mixed_logic_siblings_param_3];
-; CHECK-NEXT:    shl.b32 %r4, %r1, 2;
-; CHECK-NEXT:    or.b32 %r5, %r4, 28;
-; CHECK-NEXT:    st.b32 [%rd2], %r5;
-; CHECK-NEXT:    xor.b32 %r6, %r1, 17;
-; CHECK-NEXT:    shl.b32 %r7, %r6, 2;
-; CHECK-NEXT:    st.b32 [%rd3], %r7;
+; CHECK-NEXT:    or.b32 %r4, %r2, 28;
+; CHECK-NEXT:    st.b32 [%rd2], %r4;
+; CHECK-NEXT:    xor.b32 %r5, %r2, 68;
+; CHECK-NEXT:    st.b32 [%rd3], %r5;
 ; CHECK-NEXT:    ret;
   %and = and i32 %x, 5
   %and.shl = shl i32 %and, 2
@@ -42,19 +40,18 @@ define void @mixed_logic_siblings(i32 %x, ptr %p0, ptr %p1, ptr %p2) {
 define void @zext_and_siblings(i32 %x, ptr %p0, ptr %p1) {
 ; CHECK-LABEL: zext_and_siblings(
 ; CHECK:       {
-; CHECK-NEXT:    .reg .b32 %r<4>;
-; CHECK-NEXT:    .reg .b64 %rd<5>;
+; CHECK-NEXT:    .reg .b32 %r<2>;
+; CHECK-NEXT:    .reg .b64 %rd<6>;
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  // %bb.0:
 ; CHECK-NEXT:    ld.param.b32 %r1, [zext_and_siblings_param_0];
-; CHECK-NEXT:    and.b32 %r2, %r1, 5;
 ; CHECK-NEXT:    ld.param.b64 %rd1, [zext_and_siblings_param_1];
 ; CHECK-NEXT:    ld.param.b64 %rd2, [zext_and_siblings_param_2];
-; CHECK-NEXT:    mul.wide.u32 %rd3, %r2, 8;
-; CHECK-NEXT:    st.b64 [%rd1], %rd3;
-; CHECK-NEXT:    and.b32 %r3, %r1, 6;
-; CHECK-NEXT:    mul.wide.u32 %rd4, %r3, 8;
-; CHECK-NEXT:    st.b64 [%rd2], %rd4;
+; CHECK-NEXT:    mul.wide.u32 %rd3, %r1, 8;
+; CHECK-NEXT:    and.b64 %rd4, %rd3, 40;
+; CHECK-NEXT:    st.b64 [%rd1], %rd4;
+; CHECK-NEXT:    and.b64 %rd5, %rd3, 48;
+; CHECK-NEXT:    st.b64 [%rd2], %rd5;
 ; CHECK-NEXT:    ret;
   %and0 = and i32 %x, 5
   %z0 = zext i32 %and0 to i64
@@ -71,19 +68,18 @@ define void @zext_and_siblings(i32 %x, ptr %p0, ptr %p1) {
 define void @sext_or_siblings(i32 %x, ptr %p0, ptr %p1) {
 ; CHECK-LABEL: sext_or_siblings(
 ; CHECK:       {
-; CHECK-NEXT:    .reg .b32 %r<4>;
-; CHECK-NEXT:    .reg .b64 %rd<5>;
+; CHECK-NEXT:    .reg .b32 %r<2>;
+; CHECK-NEXT:    .reg .b64 %rd<6>;
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  // %bb.0:
 ; CHECK-NEXT:    ld.param.b32 %r1, [sext_or_siblings_param_0];
-; CHECK-NEXT:    or.b32 %r2, %r1, 5;
 ; CHECK-NEXT:    ld.param.b64 %rd1, [sext_or_siblings_param_1];
 ; CHECK-NEXT:    ld.param.b64 %rd2, [sext_or_siblings_param_2];
-; CHECK-NEXT:    mul.wide.s32 %rd3, %r2, 4;
-; CHECK-NEXT:    st.b64 [%rd1], %rd3;
-; CHECK-NEXT:    or.b32 %r3, %r1, 6;
-; CHECK-NEXT:    mul.wide.s32 %rd4, %r3, 4;
-; CHECK-NEXT:    st.b64 [%rd2], %rd4;
+; CHECK-NEXT:    mul.wide.s32 %rd3, %r1, 4;
+; CHECK-NEXT:    or.b64 %rd4, %rd3, 20;
+; CHECK-NEXT:    st.b64 [%rd1], %rd4;
+; CHECK-NEXT:    or.b64 %rd5, %rd3, 24;
+; CHECK-NEXT:    st.b64 [%rd2], %rd5;
 ; CHECK-NEXT:    ret;
   %or0 = or i32 %x, 5
   %e0 = sext i32 %or0 to i64
@@ -103,24 +99,23 @@ define void @sext_or_siblings(i32 %x, ptr %p0, ptr %p1) {
 define void @gep_index_siblings(i32 %base, ptr %data, ptr %p0, ptr %p1) {
 ; CHECK-LABEL: gep_index_siblings(
 ; CHECK:       {
-; CHECK-NEXT:    .reg .b32 %r<6>;
-; CHECK-NEXT:    .reg .b64 %rd<8>;
+; CHECK-NEXT:    .reg .b32 %r<4>;
+; CHECK-NEXT:    .reg .b64 %rd<9>;
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  // %bb.0:
 ; CHECK-NEXT:    ld.param.b32 %r1, [gep_index_siblings_param_0];
-; CHECK-NEXT:    xor.b32 %r2, %r1, 5;
 ; CHECK-NEXT:    ld.param.b64 %rd1, [gep_index_siblings_param_1];
 ; CHECK-NEXT:    ld.param.b64 %rd2, [gep_index_siblings_param_2];
-; CHECK-NEXT:    mul.wide.u32 %rd3, %r2, 4;
-; CHECK-NEXT:    add.s64 %rd4, %rd1, %rd3;
-; CHECK-NEXT:    ld.param.b64 %rd5, [gep_index_siblings_param_3];
-; CHECK-NEXT:    ld.b32 %r3, [%rd4];
-; CHECK-NEXT:    st.b32 [%rd2], %r3;
-; CHECK-NEXT:    xor.b32 %r4, %r1, 6;
-; CHECK-NEXT:    mul.wide.u32 %rd6, %r4, 4;
-; CHECK-NEXT:    add.s64 %rd7, %rd1, %rd6;
-; CHECK-NEXT:    ld.b32 %r5, [%rd7];
-; CHECK-NEXT:    st.b32 [%rd5], %r5;
+; CHECK-NEXT:    mul.wide.u32 %rd3, %r1, 4;
+; CHECK-NEXT:    xor.b64 %rd4, %rd3, 20;
+; CHECK-NEXT:    add.s64 %rd5, %rd1, %rd4;
+; CHECK-NEXT:    ld.param.b64 %rd6, [gep_index_siblings_param_3];
+; CHECK-NEXT:    ld.b32 %r2, [%rd5];
+; CHECK-NEXT:    st.b32 [%rd2], %r2;
+; CHECK-NEXT:    xor.b64 %rd7, %rd3, 24;
+; CHECK-NEXT:    add.s64 %rd8, %rd1, %rd7;
+; CHECK-NEXT:    ld.b32 %r3, [%rd8];
+; CHECK-NEXT:    st.b32 [%rd6], %r3;
 ; CHECK-NEXT:    ret;
   %x0 = xor i32 %base, 5
   %i0 = zext i32 %x0 to i64

>From 766521f4da76ea5af21af04210726e22bd1af82e Mon Sep 17 00:00:00 2001
From: Hao Ren <haor at nvidia.com>
Date: Tue, 4 Aug 2026 17:55:35 +0000
Subject: [PATCH 3/3] [DAGCombiner] Preserve target shift-folding opportunities

---
 llvm/include/llvm/CodeGen/TargetLowering.h    |  9 ---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 71 +++++++++++++++----
 llvm/lib/Target/NVPTX/NVPTXISelLowering.h     |  5 --
 3 files changed, 56 insertions(+), 29 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/TargetLowering.h b/llvm/include/llvm/CodeGen/TargetLowering.h
index bfd24f99d34d7..c663bb8ea65b7 100644
--- a/llvm/include/llvm/CodeGen/TargetLowering.h
+++ b/llvm/include/llvm/CodeGen/TargetLowering.h
@@ -4661,15 +4661,6 @@ class LLVM_ABI TargetLowering : public TargetLoweringBase {
     return true;
   }
 
-  /// Return true if it is profitable to commute a constant SHL through
-  /// multiple bitwise logic operations that share an input. This can expose a
-  /// common shifted operand, but may also prevent shifts from folding into
-  /// addressing modes, so targets must opt in.
-  virtual bool shouldCommuteShiftOfLogicToExposeCSE(const SDNode *N,
-                                                    CombineLevel Level) const {
-    return false;
-  }
-
   /// GlobalISel - return true if it is profitable to move this shift by a
   /// constant amount through its operand, adjusting any immediate operands as
   /// necessary to preserve semantics. This transformation may not be desirable
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 277ab25775433..77a6ad731160e 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -2491,9 +2491,12 @@ static bool isTruncateOf(SelectionDAG &DAG, SDValue N, SDValue &Op,
 }
 
 /// Return true if 'Use' is a load or a store that uses N as its base pointer
-/// and that N may be folded in the load / store addressing mode.
-static bool canFoldInAddressingMode(SDNode *N, SDNode *Use, SelectionDAG &DAG,
-                                    const TargetLowering &TLI) {
+/// and that N may be folded in the load / store addressing mode. If IndexScale
+/// is specified, model N with a scaled-index operand.
+static bool
+canFoldInAddressingMode(SDNode *N, SDNode *Use, SelectionDAG &DAG,
+                        const TargetLowering &TLI,
+                        std::optional<int64_t> IndexScale = std::nullopt) {
   EVT VT;
   unsigned AS;
 
@@ -2525,7 +2528,9 @@ static bool canFoldInAddressingMode(SDNode *N, SDNode *Use, SelectionDAG &DAG,
   if (N->isAnyAdd()) {
     AM.HasBaseReg = true;
     ConstantSDNode *Offset = dyn_cast<ConstantSDNode>(N->getOperand(1));
-    if (Offset)
+    if (IndexScale)
+      AM.Scale = *IndexScale;
+    else if (Offset)
       // [reg +/- imm]
       AM.BaseOffs = Offset->getSExtValue();
     else
@@ -2548,6 +2553,36 @@ static bool canFoldInAddressingMode(SDNode *N, SDNode *Use, SelectionDAG &DAG,
                                    VT.getTypeForEVT(*DAG.getContext()), AS);
 }
 
+/// Return true if Shift can be absorbed as the scaled index of a load or
+/// store address.
+static bool canFoldShiftIntoAddressingMode(SDNode *Shift, SelectionDAG &DAG,
+                                           const TargetLowering &TLI) {
+  auto *ShiftAmount = dyn_cast<ConstantSDNode>(Shift->getOperand(1));
+  if (!ShiftAmount || ShiftAmount->getAPIntValue().uge(63))
+    return false;
+
+  int64_t Scale = int64_t{1} << ShiftAmount->getZExtValue();
+  auto canFoldAsIndex = [&](SDNode *Index) {
+    return any_of(Index->users(), [&](SDNode *User) {
+      if (!User->isAnyAdd())
+        return false;
+      return any_of(User->users(), [&](SDNode *MemUser) {
+        return canFoldInAddressingMode(User, MemUser, DAG, TLI, Scale);
+      });
+    });
+  };
+
+  if (canFoldAsIndex(Shift))
+    return true;
+  for (SDNode *User : Shift->users())
+    if ((User->getOpcode() == ISD::ANY_EXTEND ||
+         User->getOpcode() == ISD::SIGN_EXTEND ||
+         User->getOpcode() == ISD::ZERO_EXTEND) &&
+        canFoldAsIndex(User))
+      return true;
+  return false;
+}
+
 /// This inverts a canonicalization in IR that replaces a variable select arm
 /// with an identity constant. Codegen improves if we re-use the variable
 /// operand rather than load a constant. This can also be converted into a
@@ -10805,13 +10840,10 @@ SDValue DAGCombiner::foldShiftOfLogicWithSharedInput(SDNode *N) {
 
   // Match a logic operation, with an optional extension, inside a SHL.
   auto matchShiftOfLogic = [&](SDNode *Shift, ShiftOfLogic &Match) {
-    if (Shift->getOpcode() != ISD::SHL)
+    if (Shift->getOpcode() != ISD::SHL || !Shift->getOperand(0).hasOneUse())
       return false;
 
     SDValue ShiftInput = Shift->getOperand(0);
-    if (!ShiftInput.hasOneUse())
-      return false;
-
     unsigned ExtendOpcode = 0;
     SDValue Logic = ShiftInput;
     if (ISD::isExtOpcode(ShiftInput.getOpcode())) {
@@ -10823,19 +10855,24 @@ SDValue DAGCombiner::foldShiftOfLogicWithSharedInput(SDNode *N) {
     if (!sd_match(Logic, m_OneUse(m_BitwiseLogic(m_Value(X),
                                                  m_Value(C, m_ConstInt())))))
       return false;
-    ConstantSDNode *Constant = isConstOrConstSplat(C);
-    if (!Constant || Constant->isOpaque())
-      return false;
 
     Match = {Shift, Logic, X, C, ExtendOpcode};
     return true;
   };
 
+  // Leave shifts used by another logic operation for the existing logic-tree
+  // combines.
+  auto feedsBitwiseLogic = [](SDNode *Shift) {
+    return any_of(Shift->users(), [](SDNode *User) {
+      return ISD::isBitwiseLogicOp(User->getOpcode());
+    });
+  };
+
   // Match N as the root shift-of-logic; bail if it does not fit the pattern or
   // should not be commuted.
   ShiftOfLogic Root;
-  if (!matchShiftOfLogic(N, Root) ||
-      !TLI.shouldCommuteShiftOfLogicToExposeCSE(N, Level))
+  if (!matchShiftOfLogic(N, Root) || feedsBitwiseLogic(N) ||
+      canFoldShiftIntoAddressingMode(N, DAG, TLI))
     return SDValue();
 
   EVT VT = N->getValueType(0);
@@ -10868,9 +10905,13 @@ SDValue DAGCombiner::foldShiftOfLogicWithSharedInput(SDNode *N) {
     if (!matchShiftOfLogic(CandidateShift, Candidate) ||
         Candidate.X != Root.X || Candidate.ExtendOpcode != Root.ExtendOpcode ||
         CandidateShift->getValueType(0) != VT ||
-        CandidateShift->getOperand(1) != ShiftAmount ||
-        !TLI.isDesirableToCommuteWithShift(CandidateShift, Level))
+        CandidateShift->getOperand(1) != ShiftAmount)
       continue;
+    // Cancel the entire rewrite if a matched candidate must stay put.
+    if (feedsBitwiseLogic(CandidateShift) ||
+        canFoldShiftIntoAddressingMode(CandidateShift, DAG, TLI) ||
+        !TLI.isDesirableToCommuteWithShift(CandidateShift, Level))
+      return SDValue();
     Siblings.push_back(Candidate);
   }
   if (Siblings.empty())
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.h b/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
index 92ff7415dd6f5..42741ba98b438 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
@@ -181,11 +181,6 @@ class NVPTXTargetLowering : public TargetLowering {
                                          TargetLoweringOpt &TLO,
                                          unsigned Depth = 0) const override;
 
-  bool shouldCommuteShiftOfLogicToExposeCSE(const SDNode *,
-                                            CombineLevel) const override {
-    return true;
-  }
-
 private:
   const NVPTXSubtarget &STI; // cache the subtarget here
   mutable unsigned GlobalUniqueCallSite;



More information about the llvm-commits mailing list