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

[Craig Topper via llvm-commits]

https://github.com/topperc created https://github.com/llvm/llvm-project/pull/76860

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

>From 815739c7956ea3f61836d586d8738199c71cb33d Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Wed, 3 Jan 2024 12:24:00 -0800
Subject: [PATCH] [DAGCombiner][RISCV] Preserve disjoint flag in folding (shl
 (or x, c1), c2) -> (or (shl x, c2), c1 << c2)

Since we are shifting both inputs to the original Or by the same
amount and inserting zeros, the result should still be disjoint.
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 6 +++++-
 llvm/test/CodeGen/RISCV/mem.ll                | 3 +--
 2 files changed, 6 insertions(+), 3 deletions(-)

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