crypton-2.2.3: cbits/crypton_mont_sqr.py

#!/usr/bin/env python3
"""Generate cbits/crypton_mont_sqr.h: Montgomery squaring for AArch64 at the
limb counts of RSA's CRT halves, fully unrolled and product-scanning.

Usage: python3 cbits/crypton_mont_sqr.py > cbits/crypton_mont_sqr.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 three registers instead
of going through memory a row at a time.  Each product off the diagonal is
summed once into a column of its own, doubled, and added in.
"""

SIZES = (16, 24, 32)

HEADER = """\
/* Generated by cbits/crypton_mont_sqr.py; do not edit.
 *
 * Montgomery squaring for AArch64 at 16, 24 and 32 limbs, the halves of
 * RSA-2048, -3072 and -4096.  Each answers what mont_sqr answers. */
#ifndef CRYPTON_MONT_SQR_H
#define CRYPTON_MONT_SQR_H

#include "crypton_bignum.h"

/* (t2:t1:t0) += x * y */
#define MONT_SQR_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_SQR_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) += 2 * (s2:s1:s0) */
#define MONT_SQR_ADD2(t0, t1, t2, s0, s1, s2)                           \\
	do {                                                                \\
		s2 = (s2 << 1) | (s1 >> 63);                                    \\
		s1 = (s1 << 1) | (s0 >> 63);                                    \\
		s0 <<= 1;                                                       \\
		__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");                                                \\
	} while (0)
"""

FOOTER = """\
#endif
"""


def gen(n):
    out = []
    w = out.append
    w(f"/* r = a * a * R^-1 mod m, for a under m; t is 2n limbs of scratch */")
    w(f"static void crypton_mont_sqr{n}(limb_t *r, const limb_t *a, const limb_t *m,")
    w(f"{' ' * len(f'static void crypton_mont_sqr{n}(')}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, s0, s1, s2, borrow, take;")
    w("")
    for i in range(2 * n - 1):
        w(f"\t/* column {i} */")
        pairs = [(j, i - j) for j in range(max(0, i - n + 1), n) if j < i - j]
        if pairs:
            j, k = pairs[0]
            w(f"\tMONT_SQR_MUL(s0, s1, s2, a[{j}], a[{k}]);")
            for (j, k) in pairs[1:]:
                w(f"\tMONT_SQR_MAC(s0, s1, s2, a[{j}], a[{k}]);")
            w("\tMONT_SQR_ADD2(t0, t1, t2, s0, s1, s2);")
        if i % 2 == 0:
            w(f"\tMONT_SQR_MAC(t0, t1, t2, a[{i // 2}], a[{i // 2}]);")
        for j in range(max(0, i - n + 1), min(i, n)):
            w(f"\tMONT_SQR_MAC(t0, t1, t2, u[{j}], m[{i - j}]);")
        if i < n:
            w(f"\tu[{i}] = t0 * n0;")
            w(f"\tMONT_SQR_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))
    print(FOOTER, end="")


if __name__ == "__main__":
    main()