[OptimizeForJS] Optimize 64-bit div by constant

src/asmjs/shared-constants.cpp

@@ -87,6 +87,7 @@ cashew::IString WASM_FETCH_HIGH_BITS("__wasm_fetch_high_bits");
 cashew::IString INT64_TO_32_HIGH_BITS("i64toi32_i32$HIGH_BITS");
 cashew::IString WASM_NEAREST_F32("__wasm_nearest_f32");
 cashew::IString WASM_NEAREST_F64("__wasm_nearest_f64");
+cashew::IString WASM_I64_MUL_HIGH("__wasm_i64_mulh");
 cashew::IString WASM_I64_MUL("__wasm_i64_mul");
 cashew::IString WASM_I64_SDIV("__wasm_i64_sdiv");
 cashew::IString WASM_I64_UDIV("__wasm_i64_udiv");

src/asmjs/shared-constants.h

@@ -90,6 +90,7 @@ extern cashew::IString WASM_FETCH_HIGH_BITS;
 extern cashew::IString INT64_TO_32_HIGH_BITS;
 extern cashew::IString WASM_NEAREST_F32;
 extern cashew::IString WASM_NEAREST_F64;
+extern cashew::IString WASM_I64_MUL_HIGH;
 extern cashew::IString WASM_I64_MUL;
 extern cashew::IString WASM_I64_SDIV;
 extern cashew::IString WASM_I64_UDIV;

src/passes/OptimizeForJS.cpp

@@ -17,29 +17,94 @@
 #include <pass.h>
 #include <wasm.h>
 
+#include "abi/js.h"
+#include "asmjs/shared-constants.h"
+#include "passes/intrinsics-module.h"
 #include "wasm-builder.h"
+#include "wasm-s-parser.h"
 #include <ir/abstract.h>
 #include <ir/literal-utils.h>
 #include <ir/localize.h>
 #include <ir/match.h>
+#include <ir/module-utils.h>
+#include <limits>
+#include <support/bits.h>
+#include <support/div-by-const.h>
 
 namespace wasm {
 
 struct OptimizeForJSPass : public WalkerPass<PostWalker<OptimizeForJSPass>> {
-  bool isFunctionParallel() override { return true; }
+  bool requireMulhIntrinsic;
+
+  bool isFunctionParallel() override { return false; }
 
   Pass* create() override { return new OptimizeForJSPass; }
 
+  void doWalkModule(Module* module) {
+    super::doWalkModule(module);
+
+    if (requireMulhIntrinsic) {
+      Module intrinsics;
+      std::string input(IntrinsicsModuleWast);
+      SExpressionParser parser(const_cast<char*>(input.c_str()));
+      Element& root = *parser.root;
+      SExpressionWasmBuilder builder(intrinsics, *root[0], IRProfile::Normal);
+      auto* func = intrinsics.getFunction(WASM_I64_MUL_HIGH);
+      doWalkFunction(ModuleUtils::copyFunction(func, *module));
+    }
+  }
+
   void visitBinary(Binary* curr) {
     using namespace Abstract;
     using namespace Match;
     {
-      // Rewrite popcnt(x) == 1   ==>   !!x & !(x & (x - 1))
+      // popcnt(x) == 1   ==>   !!x & !(x & (x - 1))
       Expression* x;
       if (matches(curr, binary(Eq, unary(Popcnt, any(&x)), ival(1)))) {
         rewritePopcntEqualOne(x);
       }
     }
+    {
+      // i64(x) / C   ==>   mulh(x, M') >> S'
+      //   where M' and S' are magic constants
+      Const* c;
+      Expression* x;
+      if (matches(curr, binary(DivU, any(&x), i64(&c)))) {
+        requireMulhIntrinsic = true;
+        rewriteDivByConstU64(x, (uint64_t)c->value.geti64());
+      }
+    }
+    {
+      // i64(x) / C   ==>   mulh(x, M') >> S'
+      //   where M' and S' are magic constants
+      Const* c;
+      Expression* x;
+      if (matches(curr, binary(DivS, any(&x), i64(&c)))) {
+        requireMulhIntrinsic = true;
+        rewriteDivByConstS64(x, c->value.geti64());
+      }
+    }
+    // TODO: implement same approach for reminders
+    // {
+    //   // i64(x) % C   ==>   mulh(x, M') >> S'
+    //   //   where M' and S' are magic constants
+    //   Const* c;
+    //   Expression* x;
+    //   if (matches(curr, binary(RemU, any(&x), i64(&c)))) {
+    //     // requireMulhIntrinsic = true;
+    //     rewriteRemByConstU64(x, (uint64_t)c->value.geti64());
+    //   }
+    // }
+    // {
+    //   // i64(x) % C   ==>   mulh(x, M') >> S'
+    //   //   where M' and S' are magic constants
+    //   Const* c;
+    //   Expression* x;
+    //   if (matches(curr, binary(RemS, any(&x), i64(&c)))) {
+    //     // requireMulhIntrinsic = true;
+    //     rewriteRemByConstS64(x, c->value.geti64());
+    //   }
+    // }
   }
 
   void rewritePopcntEqualOne(Expression* expr) {
@@ -65,6 +130,226 @@ struct OptimizeForJSPass : public WalkerPass<PostWalker<OptimizeForJSPass>> {
             builder.makeLocalGet(temp.index, type),
             builder.makeConst(Literal::makeOne(type.getBasic())))))));
   }
+
+  void rewriteDivByConstU64(Expression* dividend, uint64_t divisor) {
+    // skip if divisor is power of two or negative value
+    // also skip if dividend is const. All this handled by optimize
+    // instructions pass.
+    if (Bits::isPowerOf2(divisor) || int64_t(divisor) < 0LL ||
+        dividend->is<Const>()) {
+      return;
+    }
+
+    Builder builder(*getModule());
+
+    if (divisor == 0) {
+      // dividend / 0   ->   0
+      //
+      // This valid case in JavaScript
+      replaceCurrent(builder.makeConst(uint64_t(0)));
+      return;
+    }
+
+    const unsigned shift = Bits::countTrailingZeroes(divisor);
+
+    Type type = dividend->type;
+    Localizer temp(dividend, getFunction(), getModule());
+    Index tempIndex = temp.index;
+
+    uint64_t shiftedDivisor = divisor;
+    Expression* shiftedDividend;
+
+    if (shift) {
+      shiftedDivisor >>= shift;
+      shiftedDividend =
+        builder.makeBinary(ShrUInt64,
+                           builder.makeLocalGet(tempIndex, type),
+                           builder.makeConst(uint64_t(shift)));
+    } else {
+      shiftedDividend = builder.makeLocalGet(tempIndex, type);
+    }
+
+    const auto payload = unsignedDivisionByConstant(shiftedDivisor, shift);
+
+    // quotient = mulh(dividend, M')
+    Expression* quotient =
+      builder.makeCall(WASM_I64_MUL_HIGH,
+                       {shiftedDividend, builder.makeConst(payload.multiplier)},
+                       type);
+
+    if (payload.add) {
+      // t1 = dividend - quotient
+      // t2 = (t1 >> 1) + quotient
+      // res = t2 >> (S' - 1)
+      assert(payload.shift > 0);
+      quotient = builder.makeBinary(
+        ShrUInt64,
+        builder.makeBinary(
+          AddInt64,
+          builder.makeBinary(
+            ShrUInt64,
+            builder.makeBinary(
+              SubInt64, builder.makeLocalGet(tempIndex, type), quotient),
+            builder.makeConst(uint64_t(1))),
+          builder.makeLocalGet(tempIndex, type)),
+        builder.makeConst(uint64_t(payload.shift - 1)));
+    } else {
+      // res = quot >> shift
+      quotient = builder.makeBinary(
+        ShrUInt64, quotient, builder.makeConst(uint64_t(payload.shift)));
+    }
+
+    // use following control flow logic:
+    //
+    // if (!(dividend >> 32)) { // if high word is empty
+    //   return i64(i32(dividend) / i32(C)) // or 0 if C > 2 ** 32
+    // } else {
+    //   return mulh(dividend, M') >> S'
+    // }
+    Expression* quotient32;
+    Expression* cond = builder.makeUnary(
+      EqZInt64,
+      builder.makeBinary(
+        ShrUInt64, temp.expr, builder.makeConst(uint64_t(32))));
+
+    if (divisor <= (uint64_t)std::numeric_limits<uint32_t>::max()) {
+      // i64(i32(dividend) / i32(C))
+      quotient32 = builder.makeUnary(
+        ExtendUInt32,
+        builder.makeBinary(
+          DivUInt32,
+          builder.makeUnary(WrapInt64, builder.makeLocalGet(tempIndex, type)),
+          builder.makeConst(uint32_t(divisor))));
+    } else {
+      // i32(dividend) / C, where C > 2 ** 32   ->   0
+      quotient32 = builder.makeConst(uint64_t(0));
+    }
+
+    replaceCurrent(builder.makeIf(cond, quotient32, quotient));
+  }
+
+  void rewriteDivByConstS64(Expression* dividend, int64_t divisor) {
+    if (divisor == std::numeric_limits<int64_t>::min() ||
+        dividend->is<Const>()) {
+      return;
+    }
+
+    if (divisor == 1LL) {
+      // dividend / 1   ->   dividend
+      replaceCurrent(dividend);
+      return;
+    }
+
+    Builder builder(*getModule());
+
+    if (divisor == 0LL) {
+      // dividend / 0   ->   0
+      //
+      // This valid case in JavaScript
+      replaceCurrent(builder.makeConst(uint64_t(0)));
+      return;
+    }
+
+    if (divisor == -1LL) {
+      // dividend / -1   ->   0 - dividend
+      //
+      // Note: i64.min / -1 special case leads to overflow (trap) in WebAssembly
+      // but valid in JavaScript (0 - i64.min ->  i64.min).
+      replaceCurrent(
+        builder.makeBinary(SubInt64, builder.makeConst(uint64_t(0)), dividend));
+      return;
+    }
+
+    Localizer temp(dividend, getFunction(), getModule());
+    Type type = dividend->type;
+    Index tempIndex = temp.index;
+
+    int64_t absoluteDivisor = std::abs(divisor);
+
+    if (Bits::isPowerOf2(absoluteDivisor)) {
+      // dividend / +C_pot   ->
+      //   +((x < 0 ? (x + (C_pot - 1)) : x) >> ctz(abs(C_pot)))
+      //
+      // dividend / -C_pot   ->
+      //   -((x < 0 ? (x + (C_pot - 1)) : x) >> ctz(abs(C_pot)))
+
+      // x < 0
+      Expression* cond =
+        builder.makeBinary(LtSInt64, temp.expr, builder.makeConst(int64_t(0)));
+      Expression* ifTrue =
+        builder.makeBinary(AddInt64,
+                           builder.makeLocalGet(tempIndex, type),
+                           builder.makeConst(int64_t(absoluteDivisor - 1LL)));
+      Expression* ifFalse = builder.makeLocalGet(tempIndex, type);
+
+      Expression* quotient = builder.makeBinary(
+        ShrSInt64,
+        builder.makeSelect(cond, ifTrue, ifFalse),
+        builder.makeConst(int64_t(Bits::countTrailingZeroes(absoluteDivisor))));
+
+      if (divisor < 0) {
+        quotient =
+          builder.makeBinary(SubInt64, builder.makeConst(int64_t(0)), quotient);
+      }
+
+      replaceCurrent(quotient);
+      return;
+    }
+
+    const auto payload = signedDivisionByConstant(uint64_t(divisor));
+
+    // quotient = mulh(dividend, M')
+    Expression* quotient =
+      builder.makeCall(WASM_I64_MUL_HIGH,
+                       {builder.makeLocalGet(tempIndex, type),
+                        builder.makeConst(payload.multiplier)},
+                       type);
+
+    if (divisor > 0 && int64_t(payload.multiplier) < 0) {
+      quotient = builder.makeBinary(
+        AddInt64, quotient, builder.makeLocalGet(tempIndex, type));
+    } else if (divisor < 0 && int64_t(payload.multiplier) > 0) {
+      quotient = builder.makeBinary(
+        SubInt64, quotient, builder.makeLocalGet(tempIndex, type));
+    }
+
+    quotient = builder.makeBinary(
+      AddInt64,
+      builder.makeBinary(
+        ShrSInt64, quotient, builder.makeConst(int64_t(payload.shift))),
+      builder.makeBinary(ShrUInt64,
+                         builder.makeLocalGet(tempIndex, type),
+                         builder.makeConst(int64_t(63))));
+
+    // use following control flow logic:
+    //
+    // if (!(dividend >> 32)) { // if high word is empry
+    //   return i64(i32(dividend) / i32(C)) // or 0 if C > 2 ** 32
+    // } else {
+    //   return mulh(dividend, M') >> S'
+    // }
+    Expression* quotient32;
+    Expression* cond = builder.makeUnary(
+      EqZInt64,
+      builder.makeBinary(
+        ShrUInt64, temp.expr, builder.makeConst(uint64_t(32))));
+
+    if ((uint64_t)std::abs(divisor) <=
+        (uint64_t)std::numeric_limits<uint32_t>::max()) {
+      // i64(i32(dividend) / i32(C))
+      quotient32 = builder.makeUnary(
+        ExtendSInt32,
+        builder.makeBinary(
+          DivSInt32,
+          builder.makeUnary(WrapInt64, builder.makeLocalGet(tempIndex, type)),
+          builder.makeConst(int32_t(divisor))));
+    } else {
+      // i32(dividend) / C, where C > 2 ** 32   ->   0
+      quotient32 = builder.makeConst(int64_t(0));
+    }
+
+    replaceCurrent(builder.makeIf(cond, quotient32, quotient));
+  }
 };
 
 Pass* createOptimizeForJSPass() { return new OptimizeForJSPass(); }