Implement faster multiplication using Number Theoretic Transform

Performs ntt with three primes (29<<27|1, 26<<27|1, 24<<27|1)

Author: tompng

bigdecimal.gemspec

@@ -46,6 +46,7 @@ Gem::Specification.new do |s|
       ext/bigdecimal/feature.h
       ext/bigdecimal/missing.c
       ext/bigdecimal/missing.h
+      ext/bigdecimal/ntt.h
       ext/bigdecimal/missing/dtoa.c
       ext/bigdecimal/static_assert.h
     ]

ext/bigdecimal/bigdecimal.c

@@ -31,6 +31,12 @@
 #include "bits.h"
 #include "static_assert.h"
 
+#if SIZEOF_DECDIG == 4
+#define USE_NTT_MULTIPLICATION 1
+#include "ntt.h"
+#define NTT_MULTIPLICATION_THRESHOLD 100
+#endif
+
 #define BIGDECIMAL_VERSION "3.2.2"
 
 /* #define ENABLE_NUMERIC_STRING */
@@ -3251,6 +3257,25 @@ BigDecimal_vpmult(VALUE self, VALUE v) {
     RB_GC_GUARD(b.bigdecimal);
     return c.bigdecimal;
 }
+
+#if SIZEOF_DECDIG == 4
+VALUE
+BigDecimal_nttmult(VALUE self, VALUE v) {
+    BDVALUE a,b,c;
+    a = GetBDValueMust(self);
+    b = GetBDValueMust(v);
+    c = NewZeroWrap(1, VPMULT_RESULT_PREC(a.real, b.real) * BASE_FIG);
+    ntt_multiply(a.real->Prec, b.real->Prec, a.real->frac, b.real->frac, c.real->frac);
+    VpSetSign(c.real, a.real->sign * b.real->sign);
+    c.real->exponent = a.real->exponent + b.real->exponent;
+    c.real->Prec = a.real->Prec + b.real->Prec;
+    VpNmlz(c.real);
+    RB_GC_GUARD(a.bigdecimal);
+    RB_GC_GUARD(b.bigdecimal);
+    return c.bigdecimal;
+}
+#endif
+
 #endif /* BIGDECIMAL_USE_VP_TEST_METHODS */
 
 /* Document-class: BigDecimal
@@ -3623,6 +3648,9 @@ Init_bigdecimal(void)
 #ifdef BIGDECIMAL_USE_VP_TEST_METHODS
     rb_define_method(rb_cBigDecimal, "vpdivd", BigDecimal_vpdivd, 2);
     rb_define_method(rb_cBigDecimal, "vpmult", BigDecimal_vpmult, 1);
+#ifdef USE_NTT_MULTIPLICATION
+    rb_define_method(rb_cBigDecimal, "nttmult", BigDecimal_nttmult, 1);
+#endif
 #endif /* BIGDECIMAL_USE_VP_TEST_METHODS */
 
 #define ROUNDING_MODE(i, name, value) \
@@ -4926,6 +4954,15 @@ VpMult(Real *c, Real *a, Real *b)
         if (w) rbd_free_struct(c);
         return 0;
     }
+
+#ifdef USE_NTT_MULTIPLICATION
+    if (b->Prec >= NTT_MULTIPLICATION_THRESHOLD) {
+        ntt_multiply((uint32_t)a->Prec, (uint32_t)b->Prec, a->frac, b->frac, c->frac);
+        c->Prec = a->Prec + b->Prec;
+        goto Cleanup;
+    }
+#endif
+
     carry = 0;
     nc = ind_c = MxIndAB;
     memset(c->frac, 0, (nc + 1) * sizeof(DECDIG));        /* Initialize c  */
@@ -4972,6 +5009,8 @@ VpMult(Real *c, Real *a, Real *b)
 	    }
 	}
     }
+
+Cleanup:
     VpNmlz(c);
     if (w != NULL) {        /* free work variable */
         VpAsgn(w, c, 10);

ext/bigdecimal/ntt.h

@@ -0,0 +1,200 @@
+// NTT (Number Theoretic Transform) implementation for BigDecimal multiplication
+
+#define NTT_PRIMITIVE_ROOT 17
+#define NTT_PRIME_BASE1 24
+#define NTT_PRIME_BASE2 26
+#define NTT_PRIME_BASE3 29
+#define NTT_PRIME_SHIFT 27
+#define NTT_PRIME1 (((uint32_t)NTT_PRIME_BASE1 << NTT_PRIME_SHIFT) | 1)
+#define NTT_PRIME2 (((uint32_t)NTT_PRIME_BASE2 << NTT_PRIME_SHIFT) | 1)
+#define NTT_PRIME3 (((uint32_t)NTT_PRIME_BASE3 << NTT_PRIME_SHIFT) | 1)
+#define MAX_NTT32_BITS 27
+#define NTT_DECDIG_BASE 1000000000
+
+// Calculates base**ex % mod
+static uint32_t
+mod_pow(uint32_t base, uint32_t ex, uint32_t mod) {
+    uint32_t res = 1;
+    uint32_t bit = 1;
+    while (true) {
+        if (ex & bit) {
+            ex ^= bit;
+            res = ((uint64_t)res * base) % mod;
+        }
+        if (!ex) break;
+        base = ((uint64_t)base * base) % mod;
+        bit <<= 1;
+    }
+    return res;
+}
+
+// Recursively performs butterfly operations of NTT
+static void
+ntt_recursive(int size_bits, uint32_t *input, uint32_t *output, uint32_t *tmp, int depth, uint32_t r, uint32_t prime) {
+    if (depth > 0) {
+        ntt_recursive(size_bits, input, tmp, output, depth - 1, ((uint64_t)r * r) % prime, prime);
+    } else {
+        tmp = input;
+    }
+    uint32_t size_half = (uint32_t)1 << (size_bits - 1);
+    uint32_t stride = (uint32_t)1 << (size_bits - depth - 1);
+    uint32_t n = size_half / stride;
+    uint32_t rn = 1, rm = prime - 1;
+    uint32_t idx = 0;
+    for (uint32_t i = 0; i < n; i++) {
+        uint32_t j = i * 2 * stride;
+        for (uint32_t k = 0; k < stride; k++, j++, idx++) {
+            uint32_t a = tmp[j], b = tmp[j + stride];
+            output[idx] = (a + (uint64_t)rn * b) % prime;
+            output[idx + size_half] = (a + (uint64_t)rm * b) % prime;
+        }
+        rn = ((uint64_t)rn * r) % prime;
+        rm = ((uint64_t)rm * r) % prime;
+    }
+}
+
+/* Perform NTT on input array.
+ * base, shift: Represent the prime number as (base << shift | 1)
+ * r_base: Primitive root of unity modulo prime
+ * size_bits: log2 of the size of the input array. Should be less or equal to shift
+ * input: input array of size 1 << size_bits
+ */
+static void
+ntt(int size_bits, uint32_t *input, uint32_t *output, uint32_t *tmp, int r_base, int base, int shift, int dir) {
+    uint32_t size = (uint32_t)1 << size_bits;
+    uint32_t prime = ((uint32_t)base << shift) | 1;
+
+    // rmax**(1 << shift) % prime == 1
+    // r**size % prime == 1
+    uint32_t rmax = mod_pow(r_base, base, prime);
+    uint32_t r = mod_pow(rmax, (uint32_t)1 << (shift - size_bits), prime);
+
+    if (dir < 0) r = mod_pow(r, prime - 2, prime);
+    ntt_recursive(size_bits, input, output, tmp, size_bits - 1, r, prime);
+    if (dir < 0) {
+        uint32_t n_inv = mod_pow((uint32_t)size, prime - 2, prime);
+        for (uint32_t i = 0; i < size; i++) {
+            output[i] = ((uint64_t)output[i] * n_inv) % prime;
+        }
+    }
+}
+
+/* Calculate c that satisfies: c % PRIME1 == mod1 && c % PRIME2 == mod2 && c % PRIME3 == mod3
+ * c = (mod1 * 35002755423056150739595925972 + mod2 * 14584479687667766215746868453 + mod3 * 37919651490985126265126719818) % (PRIME1 * PRIME2 * PRIME3)
+ */
+static inline void
+mod_restore_prime_24_26_29_shift_27(uint32_t mod1, uint32_t mod2, uint32_t mod3, uint32_t *digits) {
+    // Use mixed radix notation to eliminate modulo by PRIME1 * PRIME2 * PRIME3
+    // [DIG0, DIG1, DIG2] = DIG0 + DIG1 * PRIME1 + DIG2 * PRIME1 * PRIME2
+    // DIG0: 0...PRIME1, DIG1: 0...PRIME2, DIG2: 0...PRIME3
+    // 35002755423056150739595925972 = [1, 3489660916, 3113851359]
+    // 14584479687667766215746868453 = [0, 13, 1297437912]
+    // 37919651490985126265126719818 = [0, 0, 3373338954]
+    uint64_t c0 = mod1;
+    uint64_t c1 = (uint64_t)mod2 * 13 + (uint64_t)mod1 * 3489660916;
+    uint64_t c2 = (uint64_t)mod3 * 3373338954 % NTT_PRIME3 + (uint64_t)mod2 * 1297437912 % NTT_PRIME3 + (uint64_t)mod1 * 3113851359 % NTT_PRIME3;
+    c2 += c1 / NTT_PRIME2;
+    c1 %= NTT_PRIME2;
+    c2 %= NTT_PRIME3;
+    // Base conversion
+    c1 += c2 % NTT_DECDIG_BASE * NTT_PRIME2;
+    c0 += c1 % NTT_DECDIG_BASE * NTT_PRIME1;
+    c1 /= NTT_DECDIG_BASE;
+    digits[0] = c0 % NTT_DECDIG_BASE;
+    c0 /= NTT_DECDIG_BASE;
+    c1 += c2 / NTT_DECDIG_BASE % NTT_DECDIG_BASE * NTT_PRIME2;
+    c0 += c1 % NTT_DECDIG_BASE * NTT_PRIME1;
+    c1 /= NTT_DECDIG_BASE;
+    digits[1] = c0 % NTT_DECDIG_BASE;
+    c0 = c0 / NTT_DECDIG_BASE + c1 % NTT_DECDIG_BASE * NTT_PRIME1;
+    digits[2] = c0 % NTT_DECDIG_BASE;
+    digits[3] = (c0 / NTT_DECDIG_BASE + c1 / NTT_DECDIG_BASE % NTT_DECDIG_BASE * NTT_PRIME1) % NTT_DECDIG_BASE;
+}
+
+/*
+ * NTT multiplication
+ * Uses three NTTs with mod (24 << 27 | 1), (26 << 27 | 1), and (29 << 27 | 1)
+ */
+static void
+ntt_multiply(size_t a_size, size_t b_size, uint32_t *a, uint32_t *b, uint32_t *c) {
+    if (a_size < b_size) {
+      ntt_multiply(b_size, a_size, b, a, c);
+      return;
+    }
+
+    int b_bits = 0;
+    while (((uint32_t)1 << b_bits) < (uint32_t)b_size) b_bits++;
+    int ntt_size_bits = b_bits + 1;
+    if (ntt_size_bits > MAX_NTT32_BITS) {
+      rb_raise(rb_eArgError, "Multiply size too large");
+    }
+
+    // To calculate large_a * small_b faster, split into several batches.
+    uint32_t ntt_size = (uint32_t)1 << ntt_size_bits;
+    uint32_t batch_size = ntt_size - (uint32_t)b_size;
+    uint32_t batch_count = (uint32_t)((a_size + batch_size - 1) / batch_size);
+
+    uint32_t *ntt1 = ruby_xcalloc(sizeof(uint32_t), ntt_size);
+    uint32_t *ntt2 = ruby_xcalloc(sizeof(uint32_t), ntt_size);
+    uint32_t *ntt3 = ruby_xcalloc(sizeof(uint32_t), ntt_size);
+    uint32_t *tmp1 = ruby_xcalloc(sizeof(uint32_t), ntt_size);
+    uint32_t *tmp2 = ruby_xcalloc(sizeof(uint32_t), ntt_size);
+    uint32_t *tmp3 = ruby_xcalloc(sizeof(uint32_t), ntt_size);
+    uint32_t *conv1 = ruby_xcalloc(sizeof(uint32_t), ntt_size);
+    uint32_t *conv2 = ruby_xcalloc(sizeof(uint32_t), ntt_size);
+    uint32_t *conv3 = ruby_xcalloc(sizeof(uint32_t), ntt_size);
+
+    // Calculate NTT for b in three primes. Result is reused for each batch of a.
+    memcpy(tmp1, b, b_size * sizeof(uint32_t));
+    memset(tmp1 + b_size, 0, (ntt_size - b_size) * sizeof(uint32_t));
+    ntt(ntt_size_bits, tmp1, ntt1, tmp2, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE1, NTT_PRIME_SHIFT, +1);
+    ntt(ntt_size_bits, tmp1, ntt2, tmp2, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE2, NTT_PRIME_SHIFT, +1);
+    ntt(ntt_size_bits, tmp1, ntt3, tmp2, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE3, NTT_PRIME_SHIFT, +1);
+
+    memset(c, 0, (a_size + b_size) * sizeof(uint32_t));
+    for (uint32_t idx = 0; idx < batch_count; idx++) {
+        if (idx == batch_count - 1) {
+            uint32_t len = (uint32_t)a_size - idx * batch_size;
+            memcpy(tmp1, a + idx * batch_size, len * sizeof(uint32_t));
+            memset(tmp1 + len, 0, (ntt_size - len) * sizeof(uint32_t));
+        } else {
+            memcpy(tmp1, a + idx * batch_size, batch_size * sizeof(uint32_t));
+        }
+        // Calculate convolution for this batch in three primes
+        ntt(ntt_size_bits, tmp1, tmp2, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE1, NTT_PRIME_SHIFT, +1);
+        for (uint32_t i = 0; i < ntt_size; i++) tmp2[i] = ((uint64_t)tmp2[i] * ntt1[i]) % NTT_PRIME1;
+        ntt(ntt_size_bits, tmp2, conv1, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE1, NTT_PRIME_SHIFT, -1);
+        ntt(ntt_size_bits, tmp1, tmp2, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE2, NTT_PRIME_SHIFT, +1);
+        for (uint32_t i = 0; i < ntt_size; i++) tmp2[i] = ((uint64_t)tmp2[i] * ntt2[i]) % NTT_PRIME2;
+        ntt(ntt_size_bits, tmp2, conv2, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE2, NTT_PRIME_SHIFT, -1);
+        ntt(ntt_size_bits, tmp1, tmp2, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE3, NTT_PRIME_SHIFT, +1);
+        for (uint32_t i = 0; i < ntt_size; i++) tmp2[i] = ((uint64_t)tmp2[i] * ntt3[i]) % NTT_PRIME3;
+        ntt(ntt_size_bits, tmp2, conv3, tmp3, NTT_PRIMITIVE_ROOT, NTT_PRIME_BASE3, NTT_PRIME_SHIFT, -1);
+
+        // Restore the original convolution value from three convolutions calculated in three primes
+        for (uint32_t i = 0; i < ntt_size; i++) {
+            uint32_t dig[4];
+            mod_restore_prime_24_26_29_shift_27(conv1[i], conv2[i], conv3[i], dig);
+            for (int j = 0; j < 4; j++) {
+                // Maximum overlap(4) * maximum_value(999999999) does not overflow 32-bit integer.
+                // Index check: if dig[j] is non-zero, assign index is within valid range.
+                if (dig[j]) c[idx * batch_size + i + 1 - j] += dig[j];
+            }
+        }
+    }
+    uint32_t carry = 0;
+    for (int32_t i = (uint32_t)(a_size + b_size - 1); i >= 0; i--) {
+        uint32_t v = c[i] + carry;
+        c[i] = v % NTT_DECDIG_BASE;
+        carry = v / NTT_DECDIG_BASE;
+    }
+    ruby_xfree(ntt1);
+    ruby_xfree(ntt2);
+    ruby_xfree(ntt3);
+    ruby_xfree(tmp1);
+    ruby_xfree(tmp2);
+    ruby_xfree(tmp3);
+    ruby_xfree(conv1);
+    ruby_xfree(conv2);
+    ruby_xfree(conv3);
+}

test/bigdecimal/test_vp_operation.rb

@@ -13,6 +13,10 @@ def setup
     end
   end
 
+  def ntt_mult_available?
+    BASE_FIG == 9
+  end
+
   def test_vpmult
     assert_equal(BigDecimal('121932631112635269'), BigDecimal('123456789').vpmult(BigDecimal('987654321')))
     assert_equal(BigDecimal('12193263.1112635269'), BigDecimal('123.456789').vpmult(BigDecimal('98765.4321')))
@@ -21,6 +25,15 @@ def test_vpmult
     assert_equal(BigDecimal("#{x * y}e-300"), BigDecimal("#{x}e-100").vpmult(BigDecimal("#{y}e-200")))
   end
 
+  def test_nttmult
+    omit 'NTT multiplication is only available for 32-bit DECDIG' unless ntt_mult_available?
+    [*1..32].repeated_permutation(2) do |a, b|
+      x = BigDecimal(10 ** (BASE_FIG * a) / 7)
+      y = BigDecimal(10 ** (BASE_FIG * b) / 13)
+      assert_equal(x.to_i * y.to_i, x.nttmult(y))
+    end
+  end
+
   def test_vpdivd
     # a[0] > b[0]
     # XXXX_YYYY_ZZZZ / 1111 #=> 000X_000Y_000Z