crypton-2.2.4: cbits/crypton_mont.py
#!/usr/bin/env python3
"""Generate cbits/crypton_mont.h: Montgomery multiplication and squaring for
AArch64 at the limb counts of RSA's CRT halves, fully unrolled and
product-scanning.
Usage: python3 cbits/crypton_mont.py > cbits/crypton_mont.h
The result is a column at a time, with the reduction folded into the same
scan (the FIPS method), so the running sum stays in registers instead of
going through memory a row at a time. A column's products are spread over
three accumulators and added together at its end, since one accumulator's
carry chain is what holds the core back, not its multipliers. In a square,
each product off the diagonal is summed once, doubled, and added in.
"""
SIZES = (16, 24, 32)
HEADER = """\
/* Generated by cbits/crypton_mont.py; do not edit.
*
* Montgomery multiplication and squaring for AArch64 at 16, 24 and 32 limbs,
* the halves of RSA-2048, -3072 and -4096. Each answers what mont_mul and
* mont_sqr answer. */
#ifndef CRYPTON_MONT_H
#define CRYPTON_MONT_H
#include "crypton_bignum.h"
/* (t2:t1:t0) += x * y */
#define MONT_MAC(t0, t1, t2, x, y) \\
do { \\
limb_t lo_, hi_; \\
__asm__("mul\\t%0, %5, %6\\n\\t" \\
"umulh\\t%1, %5, %6\\n\\t" \\
"adds\\t%2, %2, %0\\n\\t" \\
"adcs\\t%3, %3, %1\\n\\t" \\
"adc\\t%4, %4, xzr" \\
: "=&r"(lo_), "=&r"(hi_), "+r"(t0), "+r"(t1), "+r"(t2) \\
: "r"(x), "r"(y) \\
: "cc"); \\
} while (0)
/* (s2:s1:s0) = x * y */
#define MONT_MUL(s0, s1, s2, x, y) \\
__asm__("mul\\t%0, %3, %4\\n\\t" \\
"umulh\\t%1, %3, %4\\n\\t" \\
"mov\\t%2, xzr" \\
: "=&r"(s0), "=&r"(s1), "=r"(s2) \\
: "r"(x), "r"(y))
/* (t2:t1:t0) += (s2:s1:s0) */
#define MONT_ADD(t0, t1, t2, s0, s1, s2) \\
__asm__("adds\\t%0, %0, %3\\n\\t" \\
"adcs\\t%1, %1, %4\\n\\t" \\
"adc\\t%2, %2, %5" \\
: "+r"(t0), "+r"(t1), "+r"(t2) \\
: "r"(s0), "r"(s1), "r"(s2) \\
: "cc")
/* (s2:s1:s0) *= 2 */
#define MONT_DBL(s0, s1, s2) \\
do { \\
s2 = (s2 << 1) | (s1 >> 63); \\
s1 = (s1 << 1) | (s0 >> 63); \\
s0 <<= 1; \\
} while (0)
"""
FOOTER = """\
#endif
"""
class Column:
"""Products added to accumulators, each set up by its first product."""
def __init__(self, w):
self.w = w
self.started = set()
def add(self, acc, x, y):
if acc == "t" or acc in self.started:
self.w(f"\tMONT_MAC({acc}0, {acc}1, {acc}2, {x}, {y});")
else:
self.w(f"\tMONT_MUL({acc}0, {acc}1, {acc}2, {x}, {y});")
self.started.add(acc)
def fold(self, into, acc):
if acc in self.started:
self.w(f"\tMONT_ADD({into}0, {into}1, {into}2, {acc}0, {acc}1, {acc}2);")
def column(w, singles, cross):
"""singles are added once, cross twice: s and c take the cross products,
t, v and w take the rest, and all of them end up in t."""
col = Column(w)
for k, (x, y) in enumerate(cross):
col.add("s" if k % 2 == 0 else "c", x, y)
for k, (x, y) in enumerate(singles):
col.add("tvw"[k % 3], x, y)
col.fold("t", "v")
col.fold("t", "w")
if cross:
col.fold("s", "c")
w("\tMONT_DBL(s0, s1, s2);")
w("\tMONT_ADD(t0, t1, t2, s0, s1, s2);")
def gen(n, square):
out = []
w = out.append
name = f"crypton_mont_{'sqr' if square else 'mul'}{n}"
args = "const limb_t *a," if square else "const limb_t *a, const limb_t *b,"
if square:
w(f"/* r = a * a * R^-1 mod m, for a under m; t is 2n limbs of scratch */")
else:
w(f"/* r = a * b * R^-1 mod m, for a and b under m; t is 2n limbs of scratch */")
w(f"static void {name}(limb_t *r, {args}")
w(f"{' ' * len(f'static void {name}(')}const limb_t *m, limb_t n0, limb_t *t)")
w("{")
w(f"\tlimb_t *u = t, *z = t + {n};")
w("\tlimb_t t0 = 0, t1 = 0, t2 = 0, v0, v1, v2, w0, w1, w2, borrow, take;")
if square:
w("\tlimb_t s0, s1, s2, c0, c1, c2;")
w("")
for i in range(2 * n - 1):
w(f"\t/* column {i} */")
singles, cross = [], []
for j in range(max(0, i - n + 1), min(i, n - 1) + 1):
k = i - j
if not square:
singles.append((f"a[{j}]", f"b[{k}]"))
elif j < k:
cross.append((f"a[{j}]", f"a[{k}]"))
elif j == k:
singles.append((f"a[{j}]", f"a[{j}]"))
for j in range(max(0, i - n + 1), min(i, n)):
singles.append((f"u[{j}]", f"m[{i - j}]"))
column(w, singles, cross)
if i < n:
w(f"\tu[{i}] = t0 * n0;")
w(f"\tMONT_MAC(t0, t1, t2, u[{i}], m[0]);")
else:
w(f"\tz[{i - n}] = t0;")
w("\tt0 = t1;")
w("\tt1 = t2;")
w("\tt2 = 0;")
w(f"\tz[{n - 1}] = t0;")
w("")
w("\t/* under 2m, with t1 the bit above n limbs: at most one subtraction */")
w(f"\tborrow = sub_n(r, z, m, {n});")
w("\ttake = t1 | (borrow ^ 1);")
w(f"\tselect_n(r, r, z, take & 1, {n});")
w("}")
w("")
return "\n".join(out)
def main():
print(HEADER)
for n in SIZES:
print(gen(n, False))
print(gen(n, True))
print(FOOTER, end="")
if __name__ == "__main__":
main()