[NVPTX] Compute 'rem' using the result of 'div', if possible.

[jlebar]

Summary

In isel, transform

Num % Den

into

Num - (Num / Den) * Den

if the result of Num / Den is already available.

[tra]

Nice.

[joerg]

We've been talking about this on IRC. It would be nice to have this as target-independent option, since on many platforms a generic rem is much slower than a div followed by a multiply/subtract. Only platforms with very slow hardware multiply want to use div and rem as separate instructions.

llvm/trunk/lib/Target/NVPTX/NVPTXISelLowering.cpp

NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
 
   // We have some custom DAG combine patterns for these nodes
   setTargetDAGCombine(ISD::ADD);
   setTargetDAGCombine(ISD::AND);
   setTargetDAGCombine(ISD::FADD);
   setTargetDAGCombine(ISD::MUL);
   setTargetDAGCombine(ISD::SHL);
   setTargetDAGCombine(ISD::SELECT);
+  setTargetDAGCombine(ISD::SREM);
+  setTargetDAGCombine(ISD::UREM);
 
   // Library functions.  These default to Expand, but we have instructions
   // for them.
   setOperationAction(ISD::FCEIL,  MVT::f32, Legal);
   setOperationAction(ISD::FCEIL,  MVT::f64, Legal);
   setOperationAction(ISD::FFLOOR, MVT::f32, Legal);
   setOperationAction(ISD::FFLOOR, MVT::f64, Legal);
   setOperationAction(ISD::FNEARBYINT, MVT::f32, Legal);
else
       IntrinsicId = IsMax ? Intrinsic::nvvm_max_ull : Intrinsic::nvvm_min_ull;
   }
 
   SDLoc DL(N);
   return DCI.DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, VT,
                          DCI.DAG.getConstant(IntrinsicId, DL, VT), LHS, RHS);
 }
 
+static SDValue PerformREMCombine(SDNode *N,
+                                 TargetLowering::DAGCombinerInfo &DCI,
+                                 CodeGenOpt::Level OptLevel) {
+  assert(N->getOpcode() == ISD::SREM || N->getOpcode() == ISD::UREM);
+
+  // Don't do anything at less than -O2.
+  if (OptLevel < CodeGenOpt::Default)
+    return SDValue();
+
+  SelectionDAG &DAG = DCI.DAG;
+  SDLoc DL(N);
+  EVT VT = N->getValueType(0);
+  bool IsSigned = N->getOpcode() == ISD::SREM;
+  unsigned DivOpc = IsSigned ? ISD::SDIV : ISD::UDIV;
+
+  const SDValue &Num = N->getOperand(0);
+  const SDValue &Den = N->getOperand(1);
+
+  for (const SDNode *U : Num->uses()) {
+    if (U->getOpcode() == DivOpc && U->getOperand(0) == Num &&
+        U->getOperand(1) == Den) {
+      // Num % Den -> Num - (Num / Den) * Den
+      return DAG.getNode(ISD::SUB, DL, VT, Num,
+                         DAG.getNode(ISD::MUL, DL, VT,
+                                     DAG.getNode(DivOpc, DL, VT, Num, Den),
+                                     Den));
+    }
+  }
+  return SDValue();
+}
+
 enum OperandSignedness {
   Signed = 0,
   Unsigned,
   Unknown
 };
 
 /// IsMulWideOperandDemotable - Checks if the provided DAG node is an operand
 /// that can be demoted to \p OptSize bits without loss of information. The
switch (N->getOpcode()) {
     case ISD::MUL:
       return PerformMULCombine(N, DCI, OptLevel);
     case ISD::SHL:
       return PerformSHLCombine(N, DCI, OptLevel);
     case ISD::AND:
       return PerformANDCombine(N, DCI);
     case ISD::SELECT:
       return PerformSELECTCombine(N, DCI);
+    case ISD::UREM:
+    case ISD::SREM:
+      return PerformREMCombine(N, DCI, OptLevel);
   }
   return SDValue();
 }
 
 /// ReplaceVectorLoad - Convert vector loads into multi-output scalar loads.
 static void ReplaceLoadVector(SDNode *N, SelectionDAG &DAG,
                               SmallVectorImpl<SDValue> &Results) {
   EVT ResVT = N->getValueType(0);

llvm/trunk/test/CodeGen/NVPTX/divrem-combine.ll

+; RUN: llc -O2 < %s -march=nvptx -mcpu=sm_35 | FileCheck %s --check-prefix=O2 --check-prefix=CHECK
+; RUN: llc -O0 < %s -march=nvptx -mcpu=sm_35 | FileCheck %s --check-prefix=O0 --check-prefix=CHECK
+
+; The following IR
+;
+;   quot = n / d
+;   rem  = n % d
+;
+; should be transformed into
+;
+;   quot = n / d
+;   rem = n - (n / d) * d
+;
+; during NVPTX isel, at -O2.  At -O0, we should leave it alone.
+
+; CHECK-LABEL: sdiv32(
+define void @sdiv32(i32 %n, i32 %d, i32* %quot_ret, i32* %rem_ret) {
+  ; CHECK: div.s32 [[quot:%r[0-9]+]], [[num:%r[0-9]+]], [[den:%r[0-9]+]];
+  %quot = sdiv i32 %n, %d
+
+  ; O0: rem.s32
+  ; (This is unfortunately order-sensitive, even though mul is commutative.)
+  ; O2: mul.lo.s32 [[mul:%r[0-9]+]], [[quot]], [[den]];
+  ; O2: sub.s32 [[rem:%r[0-9]+]], [[num]], [[mul]]
+  %rem = srem i32 %n, %d
+
+  ; O2: st{{.*}}[[quot]]
+  store i32 %quot, i32* %quot_ret
+  ; O2: st{{.*}}[[rem]]
+  store i32 %rem, i32* %rem_ret
+  ret void
+}
+
+; CHECK-LABEL: udiv32(
+define void @udiv32(i32 %n, i32 %d, i32* %quot_ret, i32* %rem_ret) {
+  ; CHECK: div.u32 [[quot:%r[0-9]+]], [[num:%r[0-9]+]], [[den:%r[0-9]+]];
+  %quot = udiv i32 %n, %d
+
+  ; O0: rem.u32
+
+  ; Selection DAG doesn't know whether this is signed or unsigned
+  ; multiplication and subtraction, but it doesn't make a difference either
+  ; way.
+  ; O2: mul.lo.{{u|s}}32 [[mul:%r[0-9]+]], [[quot]], [[den]];
+  ; O2: sub.{{u|s}}32 [[rem:%r[0-9]+]], [[num]], [[mul]]
+  %rem = urem i32 %n, %d
+
+  ; O2: st{{.*}}[[quot]]
+  store i32 %quot, i32* %quot_ret
+  ; O2: st{{.*}}[[rem]]
+  store i32 %rem, i32* %rem_ret
+  ret void
+}
+
+; Check that we don't perform this optimization if one operation is signed and
+; the other isn't.
+; CHECK-LABEL: mismatched_types1(
+define void @mismatched_types1(i32 %n, i32 %d, i32* %quot_ret, i32* %rem_ret) {
+  ; CHECK: div.u32
+  ; CHECK: rem.s32
+  %quot = udiv i32 %n, %d
+  %rem = srem i32 %n, %d
+  store i32 %quot, i32* %quot_ret
+  store i32 %rem, i32* %rem_ret
+  ret void
+}
+
+; CHECK-LABEL: mismatched_types2(
+define void @mismatched_types2(i32 %n, i32 %d, i32* %quot_ret, i32* %rem_ret) {
+  ; CHECK: div.s32
+  ; CHECK: rem.u32
+  %quot = sdiv i32 %n, %d
+  %rem = urem i32 %n, %d
+  store i32 %quot, i32* %quot_ret
+  store i32 %rem, i32* %rem_ret
+  ret void
+}
+
+; Check that we don't perform this optimization if the inputs to the div don't
+; match the inputs to the rem.
+; CHECK-LABEL: mismatched_inputs1(
+define void @mismatched_inputs1(i32 %n, i32 %d, i32* %quot_ret, i32* %rem_ret) {
+  ; CHECK: div.s32
+  ; CHECK: rem.s32
+  %quot = sdiv i32 %n, %d
+  %rem = srem i32 %d, %n
+  store i32 %quot, i32* %quot_ret
+  store i32 %rem, i32* %rem_ret
+  ret void
+}
+
+; CHECK-LABEL: mismatched_inputs2(
+define void @mismatched_inputs2(i32 %n1, i32 %n2, i32 %d, i32* %quot_ret, i32* %rem_ret) {
+  ; CHECK: div.s32
+  ; CHECK: rem.s32
+  %quot = sdiv i32 %n1, %d
+  %rem = srem i32 %n2, %d
+  store i32 %quot, i32* %quot_ret
+  store i32 %rem, i32* %rem_ret
+  ret void
+}
+
+; CHECK-LABEL: mismatched_inputs3(
+define void @mismatched_inputs3(i32 %n, i32 %d1, i32 %d2, i32* %quot_ret, i32* %rem_ret) {
+  ; CHECK: div.s32
+  ; CHECK: rem.s32
+  %quot = sdiv i32 %n, %d1
+  %rem = srem i32 %n, %d2
+  store i32 %quot, i32* %quot_ret
+  store i32 %rem, i32* %rem_ret
+  ret void
+}