[DAGCombiner][RISCV] Preserve disjoint flag in folding (shl (or x, c1), c2) -> (or (shl x, c2), c1 << c2)

[topperc]

Since we are shifting both inputs to the original Or by the same amount and inserting zeros, the result should still be disjoint.

[llvmbot]

@llvm/pr-subscribers-llvm-selectiondag

Author: Craig Topper (topperc)

Changes

Since we are shifting both inputs to the original Or by the same amount and inserting zeros, the result should still be disjoint.

---

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

2 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp (+5-1)
• (modified) llvm/test/CodeGen/RISCV/mem.ll (+1-2)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index bb44ac1fba486d..464e1becc0b8c3 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -10055,7 +10055,11 @@ SDValue DAGCombiner::visitSHL(SDNode *N) {
             DAG.FoldConstantArithmetic(ISD::SHL, SDLoc(N1), VT, {N01, N1})) {
       SDValue Shl0 = DAG.getNode(ISD::SHL, SDLoc(N0), VT, N0.getOperand(0), N1);
       AddToWorklist(Shl0.getNode());
-      return DAG.getNode(N0.getOpcode(), SDLoc(N), VT, Shl0, Shl1);
+      SDNodeFlags Flags;
+      // Preserve the disjoint flag for Or.
+      if (N0.getOpcode() == ISD::OR && N0->getFlags().hasDisjoint())
+        Flags.setDisjoint(true);
+      return DAG.getNode(N0.getOpcode(), SDLoc(N), VT, Shl0, Shl1, Flags);
     }
   }
 
diff --git a/llvm/test/CodeGen/RISCV/mem.ll b/llvm/test/CodeGen/RISCV/mem.ll
index 8f65973d4fde9b..7c98d4ae1b3f24 100644
--- a/llvm/test/CodeGen/RISCV/mem.ll
+++ b/llvm/test/CodeGen/RISCV/mem.ll
@@ -355,9 +355,8 @@ define i32 @disjoint_or_lw(ptr %a, i32 %off) nounwind {
 ; RV32I-LABEL: disjoint_or_lw:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    slli a1, a1, 2
-; RV32I-NEXT:    ori a1, a1, 12
 ; RV32I-NEXT:    add a0, a0, a1
-; RV32I-NEXT:    lw a0, 0(a0)
+; RV32I-NEXT:    lw a0, 12(a0)
 ; RV32I-NEXT:    ret
   %b = or disjoint i32 %off, 3
   %1 = getelementptr i32, ptr %a, i32 %b

[nikic]

LGTM. Proof: https://alive2.llvm.org/ce/z/_2Bzdd