[RISCV] Custom-legalise i32 SDIV/UDIV/UREM on RV64M

[asb]

Summary

Follow the same custom legalisation strategy as used in D57085 for variable-length shifts (see that patch summary for more discussion). Although we may lose out on some late-stage DAG combines, I think this custom legalisation strategy is ultimately easier to reason about.

There are some codegen changes in rv64m-exhaustive-w-insts.ll but they are all neutral in terms of the number of instructions.

[efriedma]

Adding target-specific nodes for this is a little unfortunate, but seems reasonable.

lib/Target/RISCV/RISCVInstrInfoM.td

66

Since you have a custom DAG node, you might as well implement ComputeNumSignBitsForTargetNode instead of using a pattern like this.

[asb]

Update to implement ComputeNumSignBitsForTargetNode, simplifying the patterns in RISCVInstrInfoM.td. Thanks for the review feedback.

[efriedma]

LGTM

lib/Target/RISCV/RISCVISelLowering.h

enum NodeType : unsigned {
   MRET_FLAG,
   CALL,
   SELECT_CC,
   BuildPairF64,
   SplitF64,
   TAIL,
   SHLW,
   SRAW,
-  SRLW
+  SRLW,
+  DIVW,
+  DIVUW,
+  REMUW
 };
 }
 
 class RISCVTargetLowering : public TargetLowering {
   const RISCVSubtarget &Subtarget;
 
 public:
   explicit RISCVTargetLowering(const TargetMachine &TM,

lib/Target/RISCV/RISCVISelLowering.cpp

RISCVTargetLowering::RISCVTargetLowering(const TargetMachine &TM,
   setOperationAction(ISD::VAARG, MVT::Other, Expand);
   setOperationAction(ISD::VACOPY, MVT::Other, Expand);
   setOperationAction(ISD::VAEND, MVT::Other, Expand);
 
   for (auto VT : {MVT::i1, MVT::i8, MVT::i16})
     setOperationAction(ISD::SIGN_EXTEND_INREG, VT, Expand);
 
   if (Subtarget.is64Bit()) {
-    setTargetDAGCombine(ISD::ANY_EXTEND);
     setOperationAction(ISD::SHL, MVT::i32, Custom);
     setOperationAction(ISD::SRA, MVT::i32, Custom);
     setOperationAction(ISD::SRL, MVT::i32, Custom);
   }
 
   if (!Subtarget.hasStdExtM()) {
     setOperationAction(ISD::MUL, XLenVT, Expand);
     setOperationAction(ISD::MULHS, XLenVT, Expand);
     setOperationAction(ISD::MULHU, XLenVT, Expand);
     setOperationAction(ISD::SDIV, XLenVT, Expand);
     setOperationAction(ISD::UDIV, XLenVT, Expand);
     setOperationAction(ISD::SREM, XLenVT, Expand);
     setOperationAction(ISD::UREM, XLenVT, Expand);
   }
 
+  if (Subtarget.is64Bit() && Subtarget.hasStdExtM()) {
+    setOperationAction(ISD::SDIV, MVT::i32, Custom);
+    setOperationAction(ISD::UDIV, MVT::i32, Custom);
+    setOperationAction(ISD::UREM, MVT::i32, Custom);
+  }
+
   setOperationAction(ISD::SDIVREM, XLenVT, Expand);
   setOperationAction(ISD::UDIVREM, XLenVT, Expand);
   setOperationAction(ISD::SMUL_LOHI, XLenVT, Expand);
   setOperationAction(ISD::UMUL_LOHI, XLenVT, Expand);
 
   setOperationAction(ISD::SHL_PARTS, XLenVT, Expand);
   setOperationAction(ISD::SRL_PARTS, XLenVT, Expand);
   setOperationAction(ISD::SRA_PARTS, XLenVT, Expand);
static RISCVISD::NodeType getRISCVWOpcode(unsigned Opcode) {
   default:
     llvm_unreachable("Unexpected opcode");
   case ISD::SHL:
     return RISCVISD::SHLW;
   case ISD::SRA:
     return RISCVISD::SRAW;
   case ISD::SRL:
     return RISCVISD::SRLW;
+  case ISD::SDIV:
+    return RISCVISD::DIVW;
+  case ISD::UDIV:
+    return RISCVISD::DIVUW;
+  case ISD::UREM:
+    return RISCVISD::REMUW;
   }
 }
 
 // Converts the given 32-bit operation to a target-specific SelectionDAG node.
 // Because i32 isn't a legal type for RV64, these operations would otherwise
 // be promoted to i64, making it difficult to select the SLLW/DIVUW/.../*W
 // later one because the fact the operation was originally of type i32 is
 // lost.
void RISCVTargetLowering::ReplaceNodeResults(SDNode *N,
   case ISD::SRA:
   case ISD::SRL:
     assert(N->getValueType(0) == MVT::i32 && Subtarget.is64Bit() &&
            "Unexpected custom legalisation");
     if (N->getOperand(1).getOpcode() == ISD::Constant)
       return;
     Results.push_back(customLegalizeToWOp(N, DAG));
     break;
-  }
-}
-
-// Returns true if the given node is an sdiv, udiv, or urem with non-constant
-// operands.
-static bool isVariableSDivUDivURem(SDValue Val) {
-  switch (Val.getOpcode()) {
-  default:
-    return false;
   case ISD::SDIV:
   case ISD::UDIV:
   case ISD::UREM:
-    return Val.getOperand(0).getOpcode() != ISD::Constant &&
-           Val.getOperand(1).getOpcode() != ISD::Constant;
+    assert(N->getValueType(0) == MVT::i32 && Subtarget.is64Bit() &&
+           Subtarget.hasStdExtM() && "Unexpected custom legalisation");
+    if (N->getOperand(0).getOpcode() == ISD::Constant ||
+        N->getOperand(1).getOpcode() == ISD::Constant)
+      return;
+    Results.push_back(customLegalizeToWOp(N, DAG));
+    break;
   }
 }
 
 SDValue RISCVTargetLowering::PerformDAGCombine(SDNode *N,
                                                DAGCombinerInfo &DCI) const {
-  SelectionDAG &DAG = DCI.DAG;
-
   switch (N->getOpcode()) {
   default:
     break;
-  case ISD::ANY_EXTEND: {
-    // If any-extending an i32 sdiv/udiv/urem to i64, then instead sign-extend
-    // in order to increase the chance of being able to select the
-    // divw/divuw/remuw instructions.
-    SDValue Src = N->getOperand(0);
-    if (N->getValueType(0) != MVT::i64 || Src.getValueType() != MVT::i32)
-      break;
-    if (!(Subtarget.hasStdExtM() && isVariableSDivUDivURem(Src)))
-      break;
-    SDLoc DL(N);
-    // Don't add the new node to the DAGCombiner worklist, in order to avoid
-    // an infinite cycle due to SimplifyDemandedBits converting the
-    // SIGN_EXTEND back to ANY_EXTEND.
-    return DCI.CombineTo(N, DAG.getNode(ISD::SIGN_EXTEND, DL, MVT::i64, Src),
-                         false);
-  }
   case RISCVISD::SplitF64: {
     // If the input to SplitF64 is just BuildPairF64 then the operation is
     // redundant. Instead, use BuildPairF64's operands directly.
     SDValue Op0 = N->getOperand(0);
     if (Op0->getOpcode() != RISCVISD::BuildPairF64)
       break;
     return DCI.CombineTo(N, Op0.getOperand(0), Op0.getOperand(1));
   }
const char *RISCVTargetLowering::getTargetNodeName(unsigned Opcode) const {
   case RISCVISD::TAIL:
     return "RISCVISD::TAIL";
   case RISCVISD::SHLW:
     return "RISCVISD::SHLW";
   case RISCVISD::SRAW:
     return "RISCVISD::SRAW";
   case RISCVISD::SRLW:
     return "RISCVISD::SRLW";
+  case RISCVISD::DIVW:
+    return "RISCVISD::DIVW";
+  case RISCVISD::DIVUW:
+    return "RISCVISD::DIVUW";
+  case RISCVISD::REMUW:
+    return "RISCVISD::REMUW";
   }
   return nullptr;
 }
 
 std::pair<unsigned, const TargetRegisterClass *>
 RISCVTargetLowering::getRegForInlineAsmConstraint(const TargetRegisterInfo *TRI,
                                                   StringRef Constraint,
                                                   MVT VT) const {

lib/Target/RISCV/RISCVInstrInfoM.td

 //===-- RISCVInstrInfoM.td - RISC-V 'M' instructions -------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file describes the RISC-V instructions from the standard 'M', Integer
 // Multiplication and Division instruction set extension.
 //
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
+// RISC-V specific DAG Nodes.
+//===----------------------------------------------------------------------===//
+
+def riscv_divw  : SDNode<"RISCVISD::DIVW",  SDTIntBinOp>;
+def riscv_divuw : SDNode<"RISCVISD::DIVUW", SDTIntBinOp>;
+def riscv_remuw : SDNode<"RISCVISD::REMUW", SDTIntBinOp>;
+
+//===----------------------------------------------------------------------===//
 // Instructions
 //===----------------------------------------------------------------------===//
 
 let Predicates = [HasStdExtM] in {
 def MUL     : ALU_rr<0b0000001, 0b000, "mul">;
 def MULH    : ALU_rr<0b0000001, 0b001, "mulh">;
 def MULHSU  : ALU_rr<0b0000001, 0b010, "mulhsu">;
 def MULHU   : ALU_rr<0b0000001, 0b011, "mulhu">;

 def : PatGprGpr<udiv, DIVU>;
 def : PatGprGpr<srem, REM>;
 def : PatGprGpr<urem, REMU>;
 } // Predicates = [HasStdExtM]
 
 let Predicates = [HasStdExtM, IsRV64] in {
 def : Pat<(sext_inreg (mul GPR:$rs1, GPR:$rs2), i32),
           (MULW GPR:$rs1, GPR:$rs2)>;
-def : Pat<(sext_inreg (sdiv (sexti32 GPR:$rs1),
-                            (sexti32 GPR:$rs2)), i32),
+
+def : Pat<(riscv_divw GPR:$rs1, GPR:$rs2),
+          (DIVW GPR:$rs1, GPR:$rs2)>;
+def : Pat<(sext_inreg (riscv_divw GPR:$rs1, GPR:$rs2), i32),
           (DIVW GPR:$rs1, GPR:$rs2)>;
-def : Pat<(zexti32 (sdiv (sexti32 GPR:$rs1),
-                         (sexti32 GPR:$rs2))),
-          (SRLI (SLLI (DIVW GPR:$rs1, GPR:$rs2), 32), 32)>;
-def : Pat<(sext_inreg (udiv (zexti32 GPR:$rs1), (zexti32 GPR:$rs2)), i32),
+def : Pat<(riscv_divuw GPR:$rs1, GPR:$rs2),
           (DIVUW GPR:$rs1, GPR:$rs2)>;
-// It's cheaper to perform a divuw and zero-extend the result than to
-// zero-extend both inputs to a udiv.
-def : Pat<(udiv (and GPR:$rs1, 0xffffffff), (and GPR:$rs2, 0xffffffff)),
-          (SRLI (SLLI (DIVUW GPR:$rs1, GPR:$rs2), 32), 32)>;
+def : Pat<(sext_inreg (riscv_divuw GPR:$rs1, GPR:$rs2), i32),
+          (DIVUW GPR:$rs1, GPR:$rs2)>;
+def : Pat<(riscv_remuw GPR:$rs1, GPR:$rs2),
+          (REMUW GPR:$rs1, GPR:$rs2)>;
+def : Pat<(sext_inreg (riscv_remuw GPR:$rs1, GPR:$rs2), i32),
+          (REMUW GPR:$rs1, GPR:$rs2)>;
+
+// Handle the specific cases where using DIVU/REMU would be correct and result
+// in fewer instructions than emitting DIVUW/REMUW then zero-extending the
+// result.
+def : Pat<(zexti32 (riscv_divuw (zexti32 GPR:$rs1), (zexti32 GPR:$rs2))),
+          (DIVU GPR:$rs1, GPR:$rs2)>;
+def : Pat<(zexti32 (riscv_remuw (zexti32 GPR:$rs1), (zexti32 GPR:$rs2))),
+          (REMU GPR:$rs1, GPR:$rs2)>;
+
 // Although the sexti32 operands may not have originated from an i32 srem,
 // this pattern is safe as it is impossible for two sign extended inputs to
 // produce a result where res[63:32]=0 and res[31]=1.
 def : Pat<(srem (sexti32 GPR:$rs1), (sexti32 GPR:$rs2)),
           (REMW GPR:$rs1, GPR:$rs2)>;
 def : Pat<(sext_inreg (srem (sexti32 GPR:$rs1),
                             (sexti32 GPR:$rs2)), i32),
           (REMW GPR:$rs1, GPR:$rs2)>;
-def : Pat<(sext_inreg (urem (zexti32 GPR:$rs1), (zexti32 GPR:$rs2)), i32),
-          (REMUW GPR:$rs1, GPR:$rs2)>;
-// It's cheaper to perform a remuw and zero-extend the result than to
-// zero-extend both inputs to a urem.
-def : Pat<(urem (and GPR:$rs1, 0xffffffff), (and GPR:$rs2, 0xffffffff)),
-          (SRLI (SLLI (REMUW GPR:$rs1, GPR:$rs2), 32), 32)>;
 } // Predicates = [HasStdExtM, IsRV64]

test/CodeGen/RISCV/rv64m-exhaustive-w-insts.ll

 ; RV64IM-NEXT:    ret
   %1 = udiv i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_divuw_aext_zext(i32 %a, i32 zeroext %b) nounwind {
 ; RV64IM-LABEL: zext_divuw_aext_zext:
 ; RV64IM:       # %bb.0:
+; RV64IM-NEXT:    divuw a0, a0, a1
 ; RV64IM-NEXT:    slli a0, a0, 32
 ; RV64IM-NEXT:    srli a0, a0, 32
-; RV64IM-NEXT:    divu a0, a0, a1
 ; RV64IM-NEXT:    ret
   %1 = udiv i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_divuw_sext_aext(i32 signext %a, i32 %b) nounwind {
 ; RV64IM-LABEL: zext_divuw_sext_aext:
 ; RV64IM:       # %bb.0:

 ; RV64IM-NEXT:    ret
   %1 = udiv i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_divuw_sext_zext(i32 signext %a, i32 zeroext %b) nounwind {
 ; RV64IM-LABEL: zext_divuw_sext_zext:
 ; RV64IM:       # %bb.0:
+; RV64IM-NEXT:    divuw a0, a0, a1
 ; RV64IM-NEXT:    slli a0, a0, 32
 ; RV64IM-NEXT:    srli a0, a0, 32
-; RV64IM-NEXT:    divu a0, a0, a1
 ; RV64IM-NEXT:    ret
   %1 = udiv i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_divuw_zext_aext(i32 zeroext %a, i32 %b) nounwind {
 ; RV64IM-LABEL: zext_divuw_zext_aext:
 ; RV64IM:       # %bb.0:
-; RV64IM-NEXT:    slli a1, a1, 32
-; RV64IM-NEXT:    srli a1, a1, 32
-; RV64IM-NEXT:    divu a0, a0, a1
+; RV64IM-NEXT:    divuw a0, a0, a1
+; RV64IM-NEXT:    slli a0, a0, 32
+; RV64IM-NEXT:    srli a0, a0, 32
 ; RV64IM-NEXT:    ret
   %1 = udiv i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_divuw_zext_sext(i32 zeroext %a, i32 signext %b) nounwind {
 ; RV64IM-LABEL: zext_divuw_zext_sext:
 ; RV64IM:       # %bb.0:
-; RV64IM-NEXT:    slli a1, a1, 32
-; RV64IM-NEXT:    srli a1, a1, 32
-; RV64IM-NEXT:    divu a0, a0, a1
+; RV64IM-NEXT:    divuw a0, a0, a1
+; RV64IM-NEXT:    slli a0, a0, 32
+; RV64IM-NEXT:    srli a0, a0, 32
 ; RV64IM-NEXT:    ret
   %1 = udiv i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_divuw_zext_zext(i32 zeroext %a, i32 zeroext %b) nounwind {
 ; RV64IM-LABEL: zext_divuw_zext_zext:
 ; RV64IM:       # %bb.0:

 ; RV64IM-NEXT:    ret
   %1 = urem i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_remuw_aext_zext(i32 %a, i32 zeroext %b) nounwind {
 ; RV64IM-LABEL: zext_remuw_aext_zext:
 ; RV64IM:       # %bb.0:
+; RV64IM-NEXT:    remuw a0, a0, a1
 ; RV64IM-NEXT:    slli a0, a0, 32
 ; RV64IM-NEXT:    srli a0, a0, 32
-; RV64IM-NEXT:    remu a0, a0, a1
 ; RV64IM-NEXT:    ret
   %1 = urem i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_remuw_sext_aext(i32 signext %a, i32 %b) nounwind {
 ; RV64IM-LABEL: zext_remuw_sext_aext:
 ; RV64IM:       # %bb.0:

 ; RV64IM-NEXT:    ret
   %1 = urem i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_remuw_sext_zext(i32 signext %a, i32 zeroext %b) nounwind {
 ; RV64IM-LABEL: zext_remuw_sext_zext:
 ; RV64IM:       # %bb.0:
+; RV64IM-NEXT:    remuw a0, a0, a1
 ; RV64IM-NEXT:    slli a0, a0, 32
 ; RV64IM-NEXT:    srli a0, a0, 32
-; RV64IM-NEXT:    remu a0, a0, a1
 ; RV64IM-NEXT:    ret
   %1 = urem i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_remuw_zext_aext(i32 zeroext %a, i32 %b) nounwind {
 ; RV64IM-LABEL: zext_remuw_zext_aext:
 ; RV64IM:       # %bb.0:
-; RV64IM-NEXT:    slli a1, a1, 32
-; RV64IM-NEXT:    srli a1, a1, 32
-; RV64IM-NEXT:    remu a0, a0, a1
+; RV64IM-NEXT:    remuw a0, a0, a1
+; RV64IM-NEXT:    slli a0, a0, 32
+; RV64IM-NEXT:    srli a0, a0, 32
 ; RV64IM-NEXT:    ret
   %1 = urem i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_remuw_zext_sext(i32 zeroext %a, i32 signext %b) nounwind {
 ; RV64IM-LABEL: zext_remuw_zext_sext:
 ; RV64IM:       # %bb.0:
-; RV64IM-NEXT:    slli a1, a1, 32
-; RV64IM-NEXT:    srli a1, a1, 32
-; RV64IM-NEXT:    remu a0, a0, a1
+; RV64IM-NEXT:    remuw a0, a0, a1
+; RV64IM-NEXT:    slli a0, a0, 32
+; RV64IM-NEXT:    srli a0, a0, 32
 ; RV64IM-NEXT:    ret
   %1 = urem i32 %a, %b
   ret i32 %1
 }
 
 define zeroext i32 @zext_remuw_zext_zext(i32 zeroext %a, i32 zeroext %b) nounwind {
 ; RV64IM-LABEL: zext_remuw_zext_zext:
 ; RV64IM:       # %bb.0:
 ; RV64IM-NEXT:    remu a0, a0, a1
 ; RV64IM-NEXT:    ret
   %1 = urem i32 %a, %b
   ret i32 %1
 }