bigmul

bigmul.c (9397B)

---

 #include <rcx/all.h>
 #include <rcx/bench.h>
 #include <stdio.h>
 #include <unistd.h>
 
 // TODO: Add (slow) fallback for __builtin_{add,sub}cl when not using clang or
 // GCC 14. Maybe that should go in rcx.
 
 // This power of 2 results in the lowest run-time (on the hardware on which the
 // benchmarks were run). This must be at least 4 (see the comments in
 // karatsuba).
 #define KARATSUBA_THRESH 32
 
 
 /* ----- Wide u64 math ----- */
 
 inline void
 mul64(u64 *rh, u64 *rl, u64 x, u64 y) {
 	u128 r = (u128)x * (u128)y;
 	*rh = r >> 64;
 	*rl = r;
 }
 
 inline void
 fmaa64(u64 *rh, u64 *rl, u64 w, u64 x, u64 y, u64 z) {
 	u64 h0, h1, h2, l;
 	mul64(&h0, &l, w, x);               // h0:l = w * x
 	l = __builtin_addcl(l, y, 0, &h1);  // h1:l = l + y
 	l = __builtin_addcl(l, z, 0, &h2);  // h2:l = l + z
 	*rh = h0 + h1 + h2;
 	*rl = l;
 }
 
 
 /* ----- Big nat math ----- */
 
 // Precondition: m >= n
 u64
 add(u64 *r, u64 *x, usize m, u64 *y, usize n) {
 	u64 c = 0;
 	usize i = 0;
 
 	for (; i + 3 < n; i += 4) {
 		r[i + 0] = __builtin_addcl(x[i + 0], y[i + 0], c, &c);
 		r[i + 1] = __builtin_addcl(x[i + 1], y[i + 1], c, &c);
 		r[i + 2] = __builtin_addcl(x[i + 2], y[i + 2], c, &c);
 		r[i + 3] = __builtin_addcl(x[i + 3], y[i + 3], c, &c);
 	}
 	for (; i < n; i++)
 		r[i] = __builtin_addcl(x[i], y[i], c, &c);
 
 	for (; i + 3 < m; i += 4) {
 		r[i + 0] = __builtin_addcl(x[i + 0], 0, c, &c);
 		r[i + 1] = __builtin_addcl(x[i + 1], 0, c, &c);
 		r[i + 2] = __builtin_addcl(x[i + 2], 0, c, &c);
 		r[i + 3] = __builtin_addcl(x[i + 3], 0, c, &c);
 	}
 	for (; i < m; i++)
 		r[i] = __builtin_addcl(x[i], 0, c, &c);
 
 	return c;
 }
 
 // Precondition: m >= n
 // TODO: sub is not commutative like add, so we need a "bus" operation
 // ("sub" backwards) for when m < n.
 u64
 sub(u64 *r, u64 *x, usize m, u64 *y, usize n) {
 	u64 c = 0;
 	usize i = 0;
 
 	for (; i + 3 < n; i += 4) {
 		r[i + 0] = __builtin_subcl(x[i + 0], y[i + 0], c, &c);
 		r[i + 1] = __builtin_subcl(x[i + 1], y[i + 1], c, &c);
 		r[i + 2] = __builtin_subcl(x[i + 2], y[i + 2], c, &c);
 		r[i + 3] = __builtin_subcl(x[i + 3], y[i + 3], c, &c);
 	}
 	for (; i < n; i++)
 		r[i] = __builtin_subcl(x[i], y[i], c, &c);
 
 	for (; i + 3 < m; i += 4) {
 		r[i + 0] = __builtin_subcl(x[i + 0], 0, c, &c);
 		r[i + 1] = __builtin_subcl(x[i + 1], 0, c, &c);
 		r[i + 2] = __builtin_subcl(x[i + 2], 0, c, &c);
 		r[i + 3] = __builtin_subcl(x[i + 3], 0, c, &c);
 	}
 	for (; i < m; i++)
 		r[i] = __builtin_subcl(x[i], 0, c, &c);
 
 	return c;
 }
 
 // Precondition: capacity(r) >= m + n, capacity(x) = m, capactiy(y) = n
 // Precondition: r is disjoint with x and y
 void
 fma_quadratic(u64 *r, u64 *x, usize m, u64 *y, usize n) {
 	for (usize j = 0; j < n; j++) {
 		u64 c = 0;
 		for (usize i = 0; i < m; i++)
 			fmaa64(&c, &r[i + j], x[i], y[j], r[i + j], c);
 		r[m + j] = c;
 	}
 }
 
 // Precondition: capacity(r) >= 2 * n, capacity(x) = capactiy(y) = n
 // Precondition: r is disjoint with x and y
 // Precondition: capacity(tt) >= 2 * kk, where kk := 2 * (n - n / 2 + 1)
 void
 karatsuba(u64 *r, u64 *x, u64 *y, usize n, u64 *tt) {
 	/* We seek to multiply x and y, which each have n "words" (digits of base
 	 * b := 2^64), obtaining the full 2*n word product. For this, we let
 	 *     l := floor(n / 2)   and   h := ceil(n / 2)
 	 * and partition x and y into low parts xl and yl with l words and high
 	 * parts xh and yh with h words, such that
 	 *     x = xh * b^l + xl   and   y = yh * b^l + yl.
 	 * Then
 	 *     x * y = s * b^(2*l) + t * b^l + u,
 	 * where
 	 *     s := xh * yh,   t := xh * yl + xl * yh,   and   u := xl * yl.
 	 * Thus, we could multiply x and y by recursively evaluating the four
 	 * products in the definition of s, t, and u (the products involving b^i
 	 * are just bit shifts). However, this would result in a time complexity
 	 * of O(n^2), the same asymptotic performance as the naive "doubly-nested
 	 * for-loop" algorithm. Instead, we exploit the following identity, whose
 	 * significance for multiplication algorithms was first noticed by Anatoly
 	 * Karatsuba in 1960:
 	 *     t = p * q - u - s
 	 * where
 	 *     p := xl + xh   and   q := yl + yh.
 	 * Computing t in its latter form saves one multiplication at the expense
 	 * of a few additions/subtractions (which have O(n) time complexity),
 	 * thereby reducing the time complexity to O(n^(lg 3)).
 	 *
 	 * Let |z| denote the number of words in a number z. For well-founded
 	 * recusion, we need n to strictly decrease in each of the three recursive
 	 * calls. When we compute u = xl * yl and s = xh * yh, this is true as long
 	 * as n >= 2, for then |xl|,|yl| <= l < n and |yh|,|yh| <= h < n. For the
 	 * computation of p * q in t, on the other hand, we have
 	 *     |p|,|q|  = |xl + xh|,|yl + yh|   (definition of p and q)
 	 *             <= max(l, h) + 1         (addition adds at most 1 word)
 	 *              = h + 1                 (h >= l)
 	 *              = ceil(n / 2) + 1       (definition of h)
 	 * and as long as n >= 4, this quantity is strictly less than n. It
 	 * therefore suffices to separate the case n < 4 for the recursion
 	 * basis. */
 
 	// 1. Basis
 	if (n < KARATSUBA_THRESH) {
 		fma_quadratic(r, x, n, y, n);
 		return;
 	}
 
 	// 2. Compute l, h, and k, and their doubles ll, hh, and kk
 	// The significance of these quantities is as follows:
 	//   - l is the max width of xl and yl
 	//   - h is the max width of xh and yh
 	//   - k is the max width of p and q
 	//   - ll is the max width of u
 	//   - hh is the max width of s
 	//   - kk is the max width of t
 	usize l = n / 2, ll = 2 * l;
 	usize h = n - l, hh = 2 * h;
 	usize k = h + 1, kk = 2 * k;
 
 	// 3. Split x and y
 	u64 *xh = x + l, *xl = x;
 	u64 *yh = y + l, *yl = y;
 
 	// 4. Assign blocks of memory for intermediate results
 	// Note that we use the output buffer r as temporary storage for p and q.
 	// We also store u and s directly in r at the appropriate offsets (free bit
 	// shifts!), such that p and q overlap with u and s, but that's ok, because
 	// we're done with p and q by the time we calculate u and s.
 	u64 *p = r;
 	u64 *q = r + k;
 	u64 *t0 = tt;       // Storage for t in this invocation
 	u64 *t1 = tt + kk;  // Storage for t in future (recursive) invocations
 	u64 *u = r;
 	u64 *s = r + ll;
 
 	// 5. Arithmetic
 	p[k] = add(p, xh, h, xl, l);        // p = xh + xl
 	q[k] = add(q, yh, h, yl, l);        // q = yh + yl
 	karatsuba(t0, p, q, k, t1);         // t = p * q
 	karatsuba(u, xl, yl, l, t1);        // u = xl * yl
 	karatsuba(s, xh, yh, h, t1);        // s = xh * yh
 	sub(t0, t0, kk, u, ll);             // t -= u  [1]
 	sub(t0, t0, kk, s, hh);             // t -= s  [1]
 	add(r + l, r + l, hh + l, t0, kk);  // r[l..] += t  [2] [3]
 	// [1]: The borrow outs are guaranteed to be 0, because t0 - u - s must
 	//      be positive.
 	// [2]: The carry out is guaranteed to be 0, because the full product
 	//      x * y must fit in 2 * n words.
 	// [3]: The add precondition hh + l >= kk is satisfied here as long as
 	//      n >= 4, and it is, because n < 4 is the recursion basis.
 }
 
 // Precondition: capacity(r) >= m + n, capacity(x) = m, capactiy(y) = n
 // Precondition: r is disjoint with x and y
 // Precondition: m >= n
 void
 fma_karatsuba(u64 *r, u64 *x, usize m, u64 *y, usize n) {
 	// TODO: Accept capacity of r as an argument, and use excess memory for
 	// scratch, if it's big enough, instead of allocating.
 	usize firstkk = 2 * (n - n / 2 + 1);
 	u64 *scratch = r_eallocn(2 * n + 2 * firstkk, sizeof *scratch);
 	u64 *prod = scratch;
 	u64 *tt = scratch + 2 * n;
 
 	for (;;) {
 		for (; m >= n; r += n, x += n, m -= n) {
 			karatsuba(prod, x, y, n, tt);
 			add(r, prod, 2 * n, r, n);
 		}
 
 		if (m == 0) break;
 
 		if (m < KARATSUBA_THRESH) {
 			fma_quadratic(r, x, m, y, n);
 			break;
 		}
 
 		u64 *t0 = x; x = y; y = t0;   // Swap x and y
 		usize t1 = m; m = n; n = t1;  // Swap m and n
 	}
 
 	free(scratch);
 }
 
 // Precondition: capacity(r) >= m + n, capacity(x) = m, capactiy(y) = n
 // Precondition: r is disjoint with x and y
 void
 mul_karatsuba(u64 *r, u64 *x, usize m, u64 *y, usize n) {
 	if (m < n) {
 		u64 *t0 = x; x = y; y = t0;   // Swap x and y
 		usize t1 = m; m = n; n = t1;  // Swap m and n
 	}
 
 	memset(r, 0, (m + n) * sizeof *r);
 	(n < KARATSUBA_THRESH ? fma_quadratic : fma_karatsuba)(r, x, m, y, n);
 }
 
 
 /* ----- Benchmarks ----- */
 
 u64 x[4096];
 u64 y[4096];
 u64 r[8192];
 
 NOINLINE void
 bench_karatsuba(u64 l, u64 n) {
 	r_bench_start();
 	for (u64 i = 0; i < n; i++) mul_karatsuba(r, x, l, y, l);
 	r_bench_stop();
 }
 
 void bench_karatsuba16(u64 n)   { bench_karatsuba(16, n); }
 void bench_karatsuba32(u64 n)   { bench_karatsuba(32, n); }
 void bench_karatsuba64(u64 n)   { bench_karatsuba(64, n); }
 void bench_karatsuba128(u64 n)  { bench_karatsuba(128, n); }
 void bench_karatsuba256(u64 n)  { bench_karatsuba(256, n); }
 void bench_karatsuba512(u64 n)  { bench_karatsuba(512, n); }
 void bench_karatsuba1024(u64 n) { bench_karatsuba(1024, n); }
 void bench_karatsuba2048(u64 n) { bench_karatsuba(2048, n); }
 void bench_karatsuba4096(u64 n) { bench_karatsuba(4096, n); }
 
 int
 main(void) {
 	for (usize i = 0; i < LEN(x); i++) x[i] = r_prand64();
 	for (usize i = 0; i < LEN(y); i++) y[i] = r_prand64();
 
 	r_bench(bench_karatsuba16,   1000);
 	r_bench(bench_karatsuba32,   1000);
 	r_bench(bench_karatsuba64,   1000);
 	r_bench(bench_karatsuba128,  1000);
 	r_bench(bench_karatsuba256,  1000);
 	r_bench(bench_karatsuba512,  1000);
 	r_bench(bench_karatsuba1024, 1000);
 	r_bench(bench_karatsuba2048, 1000);
 	r_bench(bench_karatsuba4096, 1000);
 }