[DAG] Lower frem of power-2 using div/trunk/mul+sub.

davemgreen · davemgreen · commit f6643e3c0db0 · 2024-05-05T22:32:27.000+01:00
If we are lowering a frem and the divisor is known to me an integer power-2, we can use the formula 'frem = x - trunc(x / d) * d'. This avoids the more expensive call to fmod. The results are identical as fmod so long as d is a power-2 (so the mul does not round incorrectly), and the sign of the return is either always positive or not important (nsz). Unfortunately Alive2 does not handle this well at the moment. I was using exhaustive checking to test this, hopefully I didn't make a mistake in it (https://gist.github.com/davemgreen/6078015f30d3bacd1e9572f8db5d4b64). I found this in cpythons implementation of float_pow. I currently added it as a DAG combine for frem with power-2 fp constants.
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -17,6 +17,7 @@
 
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
+#include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/IntervalMap.h"
@@ -17261,6 +17262,25 @@ SDValue DAGCombiner::visitFREM(SDNode *N) {
   if (SDValue NewSel = foldBinOpIntoSelect(N))
     return NewSel;
 
+  if (ConstantFPSDNode *C1 = isConstOrConstSplatFP(N1, true)) {
+    bool IsExact;
+    APSInt C1I(64, 0);
+    if (C1->getValueAPF().isInteger() && !C1->getValueAPF().isNegative() &&
+        C1->getValueAPF().convertToInteger(C1I, APFloat::rmTowardZero,
+                                           &IsExact) == APFloat::opOK &&
+        IsExact && isPowerOf2_64(C1I.getSExtValue()) &&
+        (Flags.hasNoSignedZeros() || N0.getOpcode() == ISD::FABS) &&
+        !TLI.isOperationLegal(ISD::FREM, VT) &&
+        TLI.isOperationLegalOrCustom(ISD::FMUL, VT) &&
+        TLI.isOperationLegalOrCustom(ISD::FDIV, VT) &&
+        TLI.isOperationLegalOrCustom(ISD::FTRUNC, VT)) {
+      SDValue Div = DAG.getNode(ISD::FDIV, SDLoc(N), VT, N0, N1);
+      SDValue Rnd = DAG.getNode(ISD::FTRUNC, SDLoc(N), VT, Div);
+      SDValue Mul = DAG.getNode(ISD::FMUL, SDLoc(N), VT, Rnd, N1);
+      return DAG.getNode(ISD::FSUB, SDLoc(N), VT, N0, Mul);
+    }
+  }
+
   return SDValue();
 }
 
diff --git a/llvm/test/CodeGen/AArch64/frem-power2.ll b/llvm/test/CodeGen/AArch64/frem-power2.ll
@@ -13,21 +13,40 @@ entry:
 }
 
 define float @frem2_nsz(float %x) {
-; CHECK-LABEL: frem2_nsz:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    fmov s1, #2.00000000
-; CHECK-NEXT:    b fmodf
+; CHECK-SD-LABEL: frem2_nsz:
+; CHECK-SD:       // %bb.0: // %entry
+; CHECK-SD-NEXT:    fmov s1, #2.00000000
+; CHECK-SD-NEXT:    fdiv s1, s0, s1
+; CHECK-SD-NEXT:    frintz s1, s1
+; CHECK-SD-NEXT:    fadd s1, s1, s1
+; CHECK-SD-NEXT:    fsub s0, s0, s1
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: frem2_nsz:
+; CHECK-GI:       // %bb.0: // %entry
+; CHECK-GI-NEXT:    fmov s1, #2.00000000
+; CHECK-GI-NEXT:    b fmodf
 entry:
   %fmod = frem nsz float %x, 2.0
   ret float %fmod
 }
 
 define float @frem2_abs(float %x) {
-; CHECK-LABEL: frem2_abs:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    fabs s0, s0
-; CHECK-NEXT:    fmov s1, #2.00000000
-; CHECK-NEXT:    b fmodf
+; CHECK-SD-LABEL: frem2_abs:
+; CHECK-SD:       // %bb.0: // %entry
+; CHECK-SD-NEXT:    fabs s0, s0
+; CHECK-SD-NEXT:    fmov s1, #2.00000000
+; CHECK-SD-NEXT:    fdiv s1, s0, s1
+; CHECK-SD-NEXT:    frintz s1, s1
+; CHECK-SD-NEXT:    fadd s1, s1, s1
+; CHECK-SD-NEXT:    fsub s0, s0, s1
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: frem2_abs:
+; CHECK-GI:       // %bb.0: // %entry
+; CHECK-GI-NEXT:    fabs s0, s0
+; CHECK-GI-NEXT:    fmov s1, #2.00000000
+; CHECK-GI-NEXT:    b fmodf
 entry:
   %a = tail call float @llvm.fabs.f32(float %x)
   %fmod = frem float %a, 2.0
@@ -37,14 +56,11 @@ entry:
 define half @hrem2_nsz(half %x) {
 ; CHECK-SD-LABEL: hrem2_nsz:
 ; CHECK-SD:       // %bb.0: // %entry
-; CHECK-SD-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
-; CHECK-SD-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-SD-NEXT:    .cfi_offset w30, -16
-; CHECK-SD-NEXT:    fcvt s0, h0
-; CHECK-SD-NEXT:    fmov s1, #2.00000000
-; CHECK-SD-NEXT:    bl fmodf
-; CHECK-SD-NEXT:    fcvt h0, s0
-; CHECK-SD-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-SD-NEXT:    fmov h1, #2.00000000
+; CHECK-SD-NEXT:    fdiv h1, h0, h1
+; CHECK-SD-NEXT:    frintz h1, h1
+; CHECK-SD-NEXT:    fadd h1, h1, h1
+; CHECK-SD-NEXT:    fsub h0, h0, h1
 ; CHECK-SD-NEXT:    ret
 ;
 ; CHECK-GI-LABEL: hrem2_nsz:
@@ -65,10 +81,19 @@ entry:
 }
 
 define double @drem2_nsz(double %x) {
-; CHECK-LABEL: drem2_nsz:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    fmov d1, #2.00000000
-; CHECK-NEXT:    b fmod
+; CHECK-SD-LABEL: drem2_nsz:
+; CHECK-SD:       // %bb.0: // %entry
+; CHECK-SD-NEXT:    fmov d1, #2.00000000
+; CHECK-SD-NEXT:    fdiv d1, d0, d1
+; CHECK-SD-NEXT:    frintz d1, d1
+; CHECK-SD-NEXT:    fadd d1, d1, d1
+; CHECK-SD-NEXT:    fsub d0, d0, d1
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: drem2_nsz:
+; CHECK-GI:       // %bb.0: // %entry
+; CHECK-GI-NEXT:    fmov d1, #2.00000000
+; CHECK-GI-NEXT:    b fmod
 entry:
   %fmod = frem nsz double %x, 2.0
   ret double %fmod
@@ -105,62 +130,115 @@ entry:
 }
 
 define float @frem4_abs(float %x) {
-; CHECK-LABEL: frem4_abs:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    fabs s0, s0
-; CHECK-NEXT:    fmov s1, #4.00000000
-; CHECK-NEXT:    b fmodf
+; CHECK-SD-LABEL: frem4_abs:
+; CHECK-SD:       // %bb.0: // %entry
+; CHECK-SD-NEXT:    fabs s0, s0
+; CHECK-SD-NEXT:    fmov s1, #4.00000000
+; CHECK-SD-NEXT:    fdiv s2, s0, s1
+; CHECK-SD-NEXT:    frintz s2, s2
+; CHECK-SD-NEXT:    fmul s1, s2, s1
+; CHECK-SD-NEXT:    fsub s0, s0, s1
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: frem4_abs:
+; CHECK-GI:       // %bb.0: // %entry
+; CHECK-GI-NEXT:    fabs s0, s0
+; CHECK-GI-NEXT:    fmov s1, #4.00000000
+; CHECK-GI-NEXT:    b fmodf
 entry:
   %a = tail call float @llvm.fabs.f32(float %x)
   %fmod = frem float %a, 4.0
   ret float %fmod
 }
 
 define float @frem16_abs(float %x) {
-; CHECK-LABEL: frem16_abs:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    fabs s0, s0
-; CHECK-NEXT:    fmov s1, #16.00000000
-; CHECK-NEXT:    b fmodf
+; CHECK-SD-LABEL: frem16_abs:
+; CHECK-SD:       // %bb.0: // %entry
+; CHECK-SD-NEXT:    fabs s0, s0
+; CHECK-SD-NEXT:    fmov s1, #16.00000000
+; CHECK-SD-NEXT:    fdiv s2, s0, s1
+; CHECK-SD-NEXT:    frintz s2, s2
+; CHECK-SD-NEXT:    fmul s1, s2, s1
+; CHECK-SD-NEXT:    fsub s0, s0, s1
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: frem16_abs:
+; CHECK-GI:       // %bb.0: // %entry
+; CHECK-GI-NEXT:    fabs s0, s0
+; CHECK-GI-NEXT:    fmov s1, #16.00000000
+; CHECK-GI-NEXT:    b fmodf
 entry:
   %a = tail call float @llvm.fabs.f32(float %x)
   %fmod = frem float %a, 16.0
   ret float %fmod
 }
 
 define float @frem4294967296_abs(float %x) {
-; CHECK-LABEL: frem4294967296_abs:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    fabs s0, s0
-; CHECK-NEXT:    mov w8, #1333788672 // =0x4f800000
-; CHECK-NEXT:    fmov s1, w8
-; CHECK-NEXT:    b fmodf
+; CHECK-SD-LABEL: frem4294967296_abs:
+; CHECK-SD:       // %bb.0: // %entry
+; CHECK-SD-NEXT:    fabs s0, s0
+; CHECK-SD-NEXT:    mov w8, #1333788672 // =0x4f800000
+; CHECK-SD-NEXT:    fmov s1, w8
+; CHECK-SD-NEXT:    fdiv s2, s0, s1
+; CHECK-SD-NEXT:    frintz s2, s2
+; CHECK-SD-NEXT:    fmul s1, s2, s1
+; CHECK-SD-NEXT:    fsub s0, s0, s1
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: frem4294967296_abs:
+; CHECK-GI:       // %bb.0: // %entry
+; CHECK-GI-NEXT:    fabs s0, s0
+; CHECK-GI-NEXT:    mov w8, #1333788672 // =0x4f800000
+; CHECK-GI-NEXT:    fmov s1, w8
+; CHECK-GI-NEXT:    b fmodf
 entry:
   %a = tail call float @llvm.fabs.f32(float %x)
   %fmod = frem float %a, 4294967296.0
   ret float %fmod
 }
 
 define float @frem1152921504606846976_abs(float %x) {
-; CHECK-LABEL: frem1152921504606846976_abs:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    fabs s0, s0
-; CHECK-NEXT:    mov w8, #1568669696 // =0x5d800000
-; CHECK-NEXT:    fmov s1, w8
-; CHECK-NEXT:    b fmodf
+; CHECK-SD-LABEL: frem1152921504606846976_abs:
+; CHECK-SD:       // %bb.0: // %entry
+; CHECK-SD-NEXT:    fabs s0, s0
+; CHECK-SD-NEXT:    mov w8, #1568669696 // =0x5d800000
+; CHECK-SD-NEXT:    fmov s1, w8
+; CHECK-SD-NEXT:    fdiv s2, s0, s1
+; CHECK-SD-NEXT:    frintz s2, s2
+; CHECK-SD-NEXT:    fmul s1, s2, s1
+; CHECK-SD-NEXT:    fsub s0, s0, s1
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: frem1152921504606846976_abs:
+; CHECK-GI:       // %bb.0: // %entry
+; CHECK-GI-NEXT:    fabs s0, s0
+; CHECK-GI-NEXT:    mov w8, #1568669696 // =0x5d800000
+; CHECK-GI-NEXT:    fmov s1, w8
+; CHECK-GI-NEXT:    b fmodf
 entry:
   %a = tail call float @llvm.fabs.f32(float %x)
   %fmod = frem float %a, 1152921504606846976.0
   ret float %fmod
 }
 
 define float @frem4611686018427387904_abs(float %x) {
-; CHECK-LABEL: frem4611686018427387904_abs:
-; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    fabs s0, s0
-; CHECK-NEXT:    mov w8, #1585446912 // =0x5e800000
-; CHECK-NEXT:    fmov s1, w8
-; CHECK-NEXT:    b fmodf
+; CHECK-SD-LABEL: frem4611686018427387904_abs:
+; CHECK-SD:       // %bb.0: // %entry
+; CHECK-SD-NEXT:    fabs s0, s0
+; CHECK-SD-NEXT:    mov w8, #1585446912 // =0x5e800000
+; CHECK-SD-NEXT:    fmov s1, w8
+; CHECK-SD-NEXT:    fdiv s2, s0, s1
+; CHECK-SD-NEXT:    frintz s2, s2
+; CHECK-SD-NEXT:    fmul s1, s2, s1
+; CHECK-SD-NEXT:    fsub s0, s0, s1
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: frem4611686018427387904_abs:
+; CHECK-GI:       // %bb.0: // %entry
+; CHECK-GI-NEXT:    fabs s0, s0
+; CHECK-GI-NEXT:    mov w8, #1585446912 // =0x5e800000
+; CHECK-GI-NEXT:    fmov s1, w8
+; CHECK-GI-NEXT:    b fmodf
 entry:
   %a = tail call float @llvm.fabs.f32(float %x)
   %fmod = frem float %a, 4611686018427387904.0
@@ -182,42 +260,11 @@ entry:
 define <4 x float> @frem2_nsz_vec(<4 x float> %x) {
 ; CHECK-SD-LABEL: frem2_nsz_vec:
 ; CHECK-SD:       // %bb.0: // %entry
-; CHECK-SD-NEXT:    sub sp, sp, #48
-; CHECK-SD-NEXT:    str x30, [sp, #32] // 8-byte Folded Spill
-; CHECK-SD-NEXT:    .cfi_def_cfa_offset 48
-; CHECK-SD-NEXT:    .cfi_offset w30, -16
-; CHECK-SD-NEXT:    str q0, [sp, #16] // 16-byte Folded Spill
-; CHECK-SD-NEXT:    mov s0, v0.s[1]
-; CHECK-SD-NEXT:    fmov s1, #2.00000000
-; CHECK-SD-NEXT:    bl fmodf
-; CHECK-SD-NEXT:    fmov s1, #2.00000000
-; CHECK-SD-NEXT:    // kill: def $s0 killed $s0 def $q0
-; CHECK-SD-NEXT:    str q0, [sp] // 16-byte Folded Spill
-; CHECK-SD-NEXT:    ldr q0, [sp, #16] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    // kill: def $s0 killed $s0 killed $q0
-; CHECK-SD-NEXT:    bl fmodf
-; CHECK-SD-NEXT:    ldr q1, [sp] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    // kill: def $s0 killed $s0 def $q0
-; CHECK-SD-NEXT:    mov v0.s[1], v1.s[0]
-; CHECK-SD-NEXT:    fmov s1, #2.00000000
-; CHECK-SD-NEXT:    str q0, [sp] // 16-byte Folded Spill
-; CHECK-SD-NEXT:    ldr q0, [sp, #16] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    mov s0, v0.s[2]
-; CHECK-SD-NEXT:    bl fmodf
-; CHECK-SD-NEXT:    ldr q1, [sp] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    // kill: def $s0 killed $s0 def $q0
-; CHECK-SD-NEXT:    mov v1.s[2], v0.s[0]
-; CHECK-SD-NEXT:    ldr q0, [sp, #16] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    mov s0, v0.s[3]
-; CHECK-SD-NEXT:    str q1, [sp] // 16-byte Folded Spill
-; CHECK-SD-NEXT:    fmov s1, #2.00000000
-; CHECK-SD-NEXT:    bl fmodf
-; CHECK-SD-NEXT:    ldr q1, [sp] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    // kill: def $s0 killed $s0 def $q0
-; CHECK-SD-NEXT:    ldr x30, [sp, #32] // 8-byte Folded Reload
-; CHECK-SD-NEXT:    mov v1.s[3], v0.s[0]
-; CHECK-SD-NEXT:    mov v0.16b, v1.16b
-; CHECK-SD-NEXT:    add sp, sp, #48
+; CHECK-SD-NEXT:    movi v1.4s, #64, lsl #24
+; CHECK-SD-NEXT:    fdiv v1.4s, v0.4s, v1.4s
+; CHECK-SD-NEXT:    frintz v1.4s, v1.4s
+; CHECK-SD-NEXT:    fadd v1.4s, v1.4s, v1.4s
+; CHECK-SD-NEXT:    fsub v0.4s, v0.4s, v1.4s
 ; CHECK-SD-NEXT:    ret
 ;
 ; CHECK-GI-LABEL: frem2_nsz_vec:
@@ -272,48 +319,13 @@ entry:
 define <4 x float> @frem1152921504606846976_absv(<4 x float> %x) {
 ; CHECK-SD-LABEL: frem1152921504606846976_absv:
 ; CHECK-SD:       // %bb.0: // %entry
-; CHECK-SD-NEXT:    sub sp, sp, #48
-; CHECK-SD-NEXT:    str d8, [sp, #32] // 8-byte Folded Spill
-; CHECK-SD-NEXT:    str x30, [sp, #40] // 8-byte Folded Spill
-; CHECK-SD-NEXT:    .cfi_def_cfa_offset 48
-; CHECK-SD-NEXT:    .cfi_offset w30, -8
-; CHECK-SD-NEXT:    .cfi_offset b8, -16
-; CHECK-SD-NEXT:    fabs v0.4s, v0.4s
 ; CHECK-SD-NEXT:    mov w8, #1568669696 // =0x5d800000
-; CHECK-SD-NEXT:    fmov s8, w8
-; CHECK-SD-NEXT:    str q0, [sp, #16] // 16-byte Folded Spill
-; CHECK-SD-NEXT:    mov s0, v0.s[1]
-; CHECK-SD-NEXT:    fmov s1, s8
-; CHECK-SD-NEXT:    bl fmodf
-; CHECK-SD-NEXT:    fmov s1, s8
-; CHECK-SD-NEXT:    // kill: def $s0 killed $s0 def $q0
-; CHECK-SD-NEXT:    str q0, [sp] // 16-byte Folded Spill
-; CHECK-SD-NEXT:    ldr q0, [sp, #16] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    // kill: def $s0 killed $s0 killed $q0
-; CHECK-SD-NEXT:    bl fmodf
-; CHECK-SD-NEXT:    ldr q1, [sp] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    // kill: def $s0 killed $s0 def $q0
-; CHECK-SD-NEXT:    mov v0.s[1], v1.s[0]
-; CHECK-SD-NEXT:    fmov s1, s8
-; CHECK-SD-NEXT:    str q0, [sp] // 16-byte Folded Spill
-; CHECK-SD-NEXT:    ldr q0, [sp, #16] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    mov s0, v0.s[2]
-; CHECK-SD-NEXT:    bl fmodf
-; CHECK-SD-NEXT:    ldr q1, [sp] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    // kill: def $s0 killed $s0 def $q0
-; CHECK-SD-NEXT:    mov v1.s[2], v0.s[0]
-; CHECK-SD-NEXT:    ldr q0, [sp, #16] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    mov s0, v0.s[3]
-; CHECK-SD-NEXT:    str q1, [sp] // 16-byte Folded Spill
-; CHECK-SD-NEXT:    fmov s1, s8
-; CHECK-SD-NEXT:    bl fmodf
-; CHECK-SD-NEXT:    ldr q1, [sp] // 16-byte Folded Reload
-; CHECK-SD-NEXT:    // kill: def $s0 killed $s0 def $q0
-; CHECK-SD-NEXT:    ldr x30, [sp, #40] // 8-byte Folded Reload
-; CHECK-SD-NEXT:    ldr d8, [sp, #32] // 8-byte Folded Reload
-; CHECK-SD-NEXT:    mov v1.s[3], v0.s[0]
-; CHECK-SD-NEXT:    mov v0.16b, v1.16b
-; CHECK-SD-NEXT:    add sp, sp, #48
+; CHECK-SD-NEXT:    fabs v0.4s, v0.4s
+; CHECK-SD-NEXT:    dup v1.4s, w8
+; CHECK-SD-NEXT:    fdiv v2.4s, v0.4s, v1.4s
+; CHECK-SD-NEXT:    frintz v2.4s, v2.4s
+; CHECK-SD-NEXT:    fmul v1.4s, v2.4s, v1.4s
+; CHECK-SD-NEXT:    fsub v0.4s, v0.4s, v1.4s
 ; CHECK-SD-NEXT:    ret
 ;
 ; CHECK-GI-LABEL: frem1152921504606846976_absv:
