From 14882b847dff082de5cb894637edffb6570e16f1 Mon Sep 17 00:00:00 2001 From: Gregor Date: Wed, 22 Apr 2026 12:57:41 +0200 Subject: [PATCH] Revert "remove inverse.ts (move to feature/51x5-inverse)" This reverts commit a03a95039861fa9ec72c6a831d854ff519e7def3. --- src/51x5/inverse.ts | 246 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 246 insertions(+) create mode 100644 src/51x5/inverse.ts diff --git a/src/51x5/inverse.ts b/src/51x5/inverse.ts new file mode 100644 index 0000000..32bd82c --- /dev/null +++ b/src/51x5/inverse.ts @@ -0,0 +1,246 @@ +import { + $, + block, + br, + br_if, + call, + func, + i32, + i64, + if_, + Local, + local, + loop, + memory, + return_, + unreachable, + Func, +} from "wasmati"; +import { mod } from "../bigint/field-util.js"; +import { ImplicitMemory } from "../wasm/wasm-util.js"; +import { log2 } from "../util.js"; +import { bigintToData, createField } from "./field-base.js"; +import { mask51 } from "./common.js"; +import { fieldWithMethods } from "./field-single.js"; +import { multiplySingle } from "./fma.js"; + +// TODO: untested, WIP + +export { Inverse }; + +type Inverse = { + makeOdd: Func<[i32, i32], [i32]>; + inverse: Func<[i32, i32, i32], []>; +}; + +function Inverse(p: bigint, implicitMemory: ImplicitMemory): Inverse { + const Field = fieldWithMethods(createField(p, "single")); + const n = 5; + const multiply = multiplySingle(p, "single"); + + /** + * a core building block for montgomery inversion + * + * takes u, s < p. sets k=0. while u is even, update u /= 2 and s *= 2 and increment k++ + * at the end, u <- u/2^k, s <- s*2^k and the new u is odd + * returns k + * (the implementation shifts u >> k and s << k at once if k < w, and shifts by whole words until k < w) + * + * in the inversion algorithm it's guaranteed that s << k will remain < p, + * so everything holds modulo p + */ + const makeOdd = func( + { in: [i32, i32], locals: [i64, i32, i64, i64], out: [i32] }, + ([u, s], [k, k0, l, tmp]) => { + // k = count_trailing_zeros(u[0]) + let ui = Field.loadLimb(u, 0); + local.tee(k, i64.ctz(ui)); + i64.eqz(); + // if (k === 0) return; (the most common case) + if_(null, () => { + i32.const(0); + return_(); + }); + // while k === 64 (i.e., u[0] === 0), shift by whole words + // (note: u is not supposed to be 0, so u[0] = 0 implies that u is divisible by 2^w) + block(null, (block) => { + loop(null, (loop) => { + i64.ne(k, 64n); + br_if(block); + + // copy u[1],...,u[n-1] --> u[0],...,u[n-2] + local.get(u); + i32.add(u, 4); + i32.const((n - 1) * 4); + memory.copy(); + + // u[n-1] = 0 + Field.storeLimb(u, n - 1, 0n); + + // copy s[0],...,s[n-2] --> s[1],...,s[n-1] + i32.add(s, 4); + local.get(s); + i32.const((n - 1) * 4); + memory.copy(); + + // s[0] = 0 + Field.storeLimb(s, 0, 0n); + + local.set(k0, i32.add(k0, 51)); + local.set(k, i64.ctz(Field.loadLimb(u, 0))); + br(loop); + }); + }); + + // here we know that k \in 0,...,w-1 + // l = w - k + local.set(l, i64.sub(51n, k)); + + // u >> k + + // for (let i = 0; i < n-1; i++) { + // u[i] = (u[i] >> k) | ((u[i + 1] << l) & wordMax); + // } + // u[n-1] = u[n-1] >> k; + local.set(tmp, Field.loadLimb(u, 0)); + Field.forEach((i) => { + i64.shr_u(tmp, k); + if (i < n - 1) { + local.tee(tmp, Field.loadLimb(u, i + 1)); + i64.shl($, l); + i64.and($, mask51); + i64.or(); + } + Field.storeLimb(u, i, $); + }); + + // s << k + + // for (let i = n-1; i >= 0; i--) { + // s[i+1] = (s[i] >> l) | ((s[i+1] << k) & wordMax); + // } + // s[0] = (s[0] << k) & wordMax; + local.set(tmp, Field.loadLimb(s, n - 1)); + for (let i = n - 2; i >= 0; i--) { + i64.shl(tmp, k); + i64.and($, mask51); + local.tee(tmp, Field.loadLimb(s, i)); + i64.shr_u($, l); + i64.or(); + Field.storeLimb(s, i + 1, $); + } + i64.shl(tmp, k); + i64.and($, mask51); + Field.storeLimb(s, 0, $); + + // return k + i32.add(k0, i32.wrap_i64(k)); + } + ); + + // kaliski "almost inverse" algorithm + // this is modified from the algorithms in papers in that it + // * returns k-1 instead of k + // * returns r < p unconditionally + // * allows to batch left- / right-shifts + const almostInverse = func( + { in: [i32, i32, i32], locals: [i32, i32, i32], out: [i32] }, + ([u, r, a], [v, s, k]) => { + // setup locals + local.set(v, i32.add(u, Field.size)); + local.set(s, i32.add(v, Field.size)); + + // u = p, v = a, r = 0, s = 1 + Field.store(u, Field.P); + Field.copyInline(v, a); + Field.store(r, Field.Zero); + Field.store(s, Field.One); + + // main algorithm + call(makeOdd, [u, s]); + call(makeOdd, [v, r]); + local.set(k, i32.add()); + + block(null, (block) => { + loop(null, (loop) => { + call(Field.isGreater, [u, v]); + if_( + null, + () => { + call(Field.subCarry, [u, u, v]); + call(Field.addCarry, [r, r, s]); + call(makeOdd, [u, s]); + local.set(k, i32.add($, k)); + }, + () => { + call(Field.subCarry, [v, v, u]); + call(Field.addCarry, [s, s, r]); + call(Field.isZero, [v]); + br_if(block); + call(makeOdd, [v, r]); + local.set(k, i32.add($, k)); + } + ); + br(loop); + }); + }); + local.get(k); + } + ); + + // constants we store as global pointers + let K = n * 51; + let N = log2(p); + let R2corr = mod(1n << BigInt(4 * K - 2 * N + 1), p); + let r2corrGlobal = implicitMemory.data(bigintToData(R2corr)); + let pGlobal = implicitMemory.data(bigintToData(p)); + + /** + * montgomery multiplication r 2^l 2^(-K) + */ + function mulPow2(scratch: Local, r: Local, l: Local) { + // store 2^l in scratch + Field.setZero(scratch); + i64.store( + {}, + // limb index + i32.add(scratch, i32.div_u(l, 51)), + // bit vector + i64.shl(1n, i64.extend_i32_u(i32.rem_u(l, 51))) + ); + // multiply r by 2^l + call(multiply, [r, r, scratch]); + } + + /** + * montgomery inverse, a 2^K -> a^(-1) 2^K (mod p) + */ + const inverse = func( + { in: [i32, i32, i32], locals: [i32], out: [] }, + ([scratch, r, a], [k]) => { + call(Field.fullyReduce, [a]); + + // error if input is zero + call(Field.isZero, [a]); + if_(null, () => unreachable()); + + call(almostInverse, [scratch, r, a]); + local.set(k); + // don't have to reduce r here, because it's already < p + call(Field.subCarry, [r, pGlobal, r]); + // multiply by 2^(2N - k), where N = bit length of p + // we use k+1 here because that's the value the theory is about: + // N <= k+1 <= 2N, so that 0 <= 2N-(k+1) <= N, so that + // 1 <= 2^(2N-(k+1)) <= 2^N < 2p + // (in practice, k seems to be normally distributed around ~1.4N and never reach either N or 2N) + local.set(k, i32.sub(2 * N - 1, k)); + mulPow2(scratch, r, k); // * 2^(2N - (k+1)) * 2^(-K) + call(multiply, [r, r, r2corrGlobal]); // * 2^(4K - 2N + 1) * 2^(-K) + // = * 2^(2K - k) + // ^^^ transforms (a 2^K)^(-1) 2^k = a^(-1) 2^(-K+k) + // to a^(-1) 2^(-K+k + 2K -k) = a^(-1) 2^K = the montgomery representation of a^(-1) + } + ); + + return { makeOdd, inverse }; +}