[RISCV] DAG combine (sra (shl X, 32), 32 - C) -> (shl (sext_inreg X, i32), C).

[craig.topper]

Summary

The sext_inreg can often be folded into an earlier instruction by
using a W instruction. The sext_inreg also works better with our ABI.

This is one of the steps to improving the generated code for this https://godbolt.org/z/hssn6sPco

[asb]

I think you need to push the commit with rv64i-shift-sext.ll?

[jrtc27]

Given this only deals with generic nodes, would this make sense to be a TLI-guided generic combine? (And maybe other targets could benefit from it, too, or already do this in target combines?)

[craig.topper]

In D128843#3620025, @jrtc27 wrote:

Given this only deals with generic nodes, would this make sense to be a TLI-guided generic combine? (And maybe other targets could benefit from it, too, or already do this in target combines?)

Looks like X86 has something similar in combineShiftRightArithmetic, but they handle a more generic form.

// fold (ashr (shl, a, [56,48,32,24,16]), SarConst)                            
// into (shl, (sext (a), [56,48,32,24,16] - SarConst)) or                      
// into (lshr, (sext (a), SarConst - [56,48,32,24,16]))                        
// depending on sign of (SarConst - [56,48,32,24,16])

For RISC-V we also support sext_inreg for i8 and i16 with Zbb. I'm not sure we want to convert to it naively since it won't always fold away and zext.h/zext.b aren't compressible. It probablys makes sense to convert i8/i16 if the input is a load even without Zbb. But that probably needs to create a sextload instead of creating a sext_inreg.

I have another patch I'm working on that depends on this one. So I'd like to consider making this target independent as a future pass. I'll add a FIXME.

[craig.topper]

Add FIXME to make this generic.

[craig.topper]

I'm probably going to end up relaxing the one use check here to be that all users of the shl X, 32 are SRA instructions so that we can end up with a sext_inreg shared by multiple shls.

[craig.topper]

My plan had been to solve https://godbolt.org/z/hssn6sPco by adding a fold for (add (shl X, 32), C<<32) -> (shl (add X, C), 32) and then let this patch optimize the shift with sra that might be after it. I thought the (add (shl X, 32), C<<32) fold could stand on its own as well, but in testing I found that some values of C cause an infinite loop with isDesirableToCommuteWithShift.

The larger (sra (add (shl X, 32), C << 32), 32-C1) -> (shl (sext_inreg (add X, C), i32), C1) pattern is probably always profitable since the sext_inreg will always becomes an ADDW.

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

8533

I typoed this comment

[craig.topper]

Fix mistake in comment

[asb]

LGTM, thanks!

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

 //===-- RISCVISelLowering.cpp - RISCV DAG Lowering Implementation  --------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
RISCVTargetLowering::RISCVTargetLowering(const TargetMachine &TM,
 
   setMinimumJumpTableEntries(5);
 
   // Jumps are expensive, compared to logic
   setJumpIsExpensive();
 
   setTargetDAGCombine({ISD::INTRINSIC_WO_CHAIN, ISD::ADD, ISD::SUB, ISD::AND,
                        ISD::OR, ISD::XOR});
+  if (Subtarget.is64Bit())
+    setTargetDAGCombine(ISD::SRA);
 
   if (Subtarget.hasStdExtF())
     setTargetDAGCombine({ISD::FADD, ISD::FMAXNUM, ISD::FMINNUM});
 
   if (Subtarget.hasStdExtZbp())
     setTargetDAGCombine({ISD::ROTL, ISD::ROTR});
 
   if (Subtarget.hasStdExtZbb())
if (NegAcc) {
     case RISCVISD::VFNMADD_VL: Opcode = RISCVISD::VFNMSUB_VL; break;
     case RISCVISD::VFNMSUB_VL: Opcode = RISCVISD::VFNMADD_VL; break;
     }
     // clang-format on
   }
 
   return Opcode;
 }
+
+// Combine (sra (shl X, 32), 32 - C) -> (shl (sext_inreg X, i32), C)
+// FIXME: Should this be a generic combine? There's a similar combine on X86.
+static SDValue performSRACombine(SDNode *N, SelectionDAG &DAG,
+                                 const RISCVSubtarget &Subtarget) {
+  assert(N->getOpcode() == ISD::SRA && "Unexpected opcode");
+
+  if (N->getValueType(0) != MVT::i64 || !Subtarget.is64Bit())
+    return SDValue();
+
+  auto *C = dyn_cast<ConstantSDNode>(N->getOperand(1));
+  if (!C || C->getZExtValue() >= 32)
+    return SDValue();
+
+  SDValue N0 = N->getOperand(0);
+  if (N0.getOpcode() != ISD::SHL || !N0.hasOneUse() ||
+      !isa<ConstantSDNode>(N0.getOperand(1)) ||
+      N0.getConstantOperandVal(1) != 32)
+    return SDValue();
+
+  SDLoc DL(N);
+  SDValue SExt = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, MVT::i64,
+                             N0.getOperand(0), DAG.getValueType(MVT::i32));
+  return DAG.getNode(ISD::SHL, DL, MVT::i64, SExt,
+                     DAG.getConstant(32 - C->getZExtValue(), DL, MVT::i64));
+}
+
 SDValue RISCVTargetLowering::PerformDAGCombine(SDNode *N,
                                                DAGCombinerInfo &DCI) const {
   SelectionDAG &DAG = DCI.DAG;
 
   // Helper to call SimplifyDemandedBits on an operand of N where only some low
   // bits are demanded. N will be added to the Worklist if it was not deleted.
   // Caller should return SDValue(N, 0) if this returns true.
   auto SimplifyDemandedLowBitsHelper = [&](unsigned OpNo, unsigned LowBits) {
if (ShAmt.getOpcode() == RISCVISD::SPLAT_VECTOR_SPLIT_I64_VL) {
       ShAmt = DAG.getNode(RISCVISD::VMV_V_X_VL, DL, VT, DAG.getUNDEF(VT),
                           ShAmt.getOperand(1), VL);
       return DAG.getNode(N->getOpcode(), DL, VT, N->getOperand(0), ShAmt,
                          N->getOperand(2), N->getOperand(3));
     }
     break;
   }
   case ISD::SRA:
+    if (SDValue V = performSRACombine(N, DAG, Subtarget))
+      return V;
+    LLVM_FALLTHROUGH;
   case ISD::SRL:
   case ISD::SHL: {
     SDValue ShAmt = N->getOperand(1);
     if (ShAmt.getOpcode() == RISCVISD::SPLAT_VECTOR_SPLIT_I64_VL) {
       // We don't need the upper 32 bits of a 64-bit element for a shift amount.
       SDLoc DL(N);
       EVT VT = N->getValueType(0);
       ShAmt = DAG.getNode(RISCVISD::VMV_V_X_VL, DL, VT, DAG.getUNDEF(VT),

llvm/test/CodeGen/RISCV/rv64i-shift-sext.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv64 -verify-machineinstrs < %s \
 ; RUN:   | FileCheck %s -check-prefix=RV64I
 
 ; Test that we turn (sra (shl X, 32), 32-C) into (slli (sext.w X), C)
 
 define i64 @test1(i64 %a) nounwind {
 ; RV64I-LABEL: test1:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    slli a0, a0, 32
-; RV64I-NEXT:    srai a0, a0, 30
+; RV64I-NEXT:    sext.w a0, a0
+; RV64I-NEXT:    slli a0, a0, 2
 ; RV64I-NEXT:    ret
   %1 = shl i64 %a, 32
   %2 = ashr i64 %1, 30
   ret i64 %2
 }
 
 define i64 @test2(i32 signext %a) nounwind {
 ; RV64I-LABEL: test2:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    slli a0, a0, 32
-; RV64I-NEXT:    srai a0, a0, 29
+; RV64I-NEXT:    slli a0, a0, 3
 ; RV64I-NEXT:    ret
   %1 = zext i32 %a to i64
   %2 = shl i64 %1, 32
   %3 = ashr i64 %2, 29
   ret i64 %3
 }
 
 define i64 @test3(i32* %a) nounwind {
 ; RV64I-LABEL: test3:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    lw a0, 0(a0)
-; RV64I-NEXT:    slli a0, a0, 32
-; RV64I-NEXT:    srai a0, a0, 28
+; RV64I-NEXT:    slli a0, a0, 4
 ; RV64I-NEXT:    ret
   %1 = load i32, i32* %a
   %2 = zext i32 %1 to i64
   %3 = shl i64 %2, 32
   %4 = ashr i64 %3, 28
   ret i64 %4
 }
 
 define i64 @test4(i32 signext %a, i32 signext %b) nounwind {
 ; RV64I-LABEL: test4:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    addw a0, a0, a1
-; RV64I-NEXT:    slli a0, a0, 32
-; RV64I-NEXT:    srai a0, a0, 2
+; RV64I-NEXT:    slli a0, a0, 30
 ; RV64I-NEXT:    ret
   %1 = add i32 %a, %b
   %2 = zext i32 %1 to i64
   %3 = shl i64 %2, 32
   %4 = ashr i64 %3, 2
   ret i64 %4
 }
 
 define i64 @test5(i32 signext %a, i32 signext %b) nounwind {
 ; RV64I-LABEL: test5:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    xor a0, a0, a1
-; RV64I-NEXT:    slli a0, a0, 32
-; RV64I-NEXT:    srai a0, a0, 1
+; RV64I-NEXT:    slli a0, a0, 31
 ; RV64I-NEXT:    ret
   %1 = xor i32 %a, %b
   %2 = zext i32 %1 to i64
   %3 = shl i64 %2, 32
   %4 = ashr i64 %3, 1
   ret i64 %4
 }
 
 define i64 @test6(i32 signext %a, i32 signext %b) nounwind {
 ; RV64I-LABEL: test6:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    sllw a0, a0, a1
-; RV64I-NEXT:    slli a0, a0, 32
-; RV64I-NEXT:    srai a0, a0, 16
+; RV64I-NEXT:    slli a0, a0, 16
 ; RV64I-NEXT:    ret
   %1 = shl i32 %a, %b
   %2 = zext i32 %1 to i64
   %3 = shl i64 %2, 32
   %4 = ashr i64 %3, 16
   ret i64 %4
 }