crypton-2.2.4: cbits/crypton_powm_ifma.py

#!/usr/bin/env python3
"""Generate cbits/crypton_powm_ifma.h: two almost-Montgomery multiplications
at once in radix 2^52 with AVX-512 IFMA on 256-bit registers, at 20, 30 and
40 digits.  Usage: python3 cbits/crypton_powm_ifma.py > cbits/crypton_powm_ifma.h

Gueron and Krasnov's method, a digit of b at a time.  Lane 0 is followed in
full in a scalar, so it is read before the high halves go in and they stay
off the chain the next digit waits on.
"""

SIZES = (20, 30, 40)

HEADER = """\
/* Generated by cbits/crypton_powm_ifma.py; do not edit.
 *
 * Almost-Montgomery multiplication in radix 2^52, two at once: r = a b
 * 2^-52D mod m, under 2m for a and b under 2m.  Digits are 52 bits in 64-bit
 * words, each number padded with zero words to a whole number of 256-bit
 * registers.  m is read from memory; at 20 digits a is held in registers,
 * above that it is read from memory too. */
#ifndef CRYPTON_POWM_IFMA_H
#define CRYPTON_POWM_IFMA_H

#include <stdint.h>
#include <immintrin.h>

#define IFMA_MASK52 ((UINT64_C(1) << 52) - 1)
#define IFMA_TARGET __attribute__((target("avx2,avx512f,avx512vl,avx512ifma")))
"""

FOOTER = """\
#endif
"""


def regs(d):
    return (d + 3) // 4


def gen_norm(R):
    o = []
    w = o.append
    xs = [f"X{j}" for j in range(R)]
    w(f"/* {4 * R} lanes to 52-bit digits: what is above 52 bits moves up a lane,")
    w(f" * and the one-bit carries that can then remain are resolved together */")
    w(f"IFMA_TARGET static inline void ifma_norm{R}(uint64_t *r, " +
      ", ".join(f"__m256i {x}" for x in xs) + ")")
    w("{")
    w("\tconst __m256i mk = _mm256_set1_epi64x((long long) IFMA_MASK52), one = _mm256_set1_epi64x(1);")
    w("\tconst __m256i zero = _mm256_setzero_si256();")
    w("\tuint64_t g = 0, p = 0, c;")
    for j in range(R):
        w(f"\t__m256i C{j} = _mm256_srli_epi64(X{j}, 52);")
    for j in range(R):
        prev = f"C{j - 1}" if j else "zero"
        w(f"\tX{j} = _mm256_add_epi64(_mm256_and_si256(X{j}, mk), _mm256_alignr_epi64(C{j}, {prev}, 3));")
    for j in range(R):
        w(f"\tg |= (uint64_t) _mm256_cmpgt_epu64_mask(X{j}, mk) << {4 * j};")
        w(f"\tp |= (uint64_t) _mm256_cmpeq_epu64_mask(X{j}, mk) << {4 * j};")
    w("\tc = ((g << 1) + p) ^ p;")
    for j in range(R):
        w(f"\tX{j} = _mm256_and_si256(_mm256_mask_add_epi64(X{j}, (__mmask8)(c >> {4 * j} & 15), X{j}, one), mk);")
        w(f"\t_mm256_storeu_si256((__m256i *)(r + {4 * j}), X{j});")
    w("}")
    w("")
    return "\n".join(o)


def gen_amm(d):
    R = regs(d)
    in_regs = R <= 5
    o = []
    w = o.append
    w(f"/* {d} digits, {R} registers a number */")
    w(f"IFMA_TARGET static void ifma_amm{d}_x2(uint64_t *r1, const uint64_t *a1, const uint64_t *b1,")
    w(f"                                    const uint64_t *m1, uint64_t k1,")
    w(f"                                    uint64_t *r2, const uint64_t *a2, const uint64_t *b2,")
    w(f"                                    const uint64_t *m2, uint64_t k2)")
    w("{")
    w("\tconst __m256i zero = _mm256_setzero_si256();")
    w("#define L(p, k) _mm256_loadu_si256((const __m256i *)((p) + 4 * (k)))")
    X = [f"X{j}" for j in range(R)]
    Y = [f"Y{j}" for j in range(R)]
    w("\t__m256i " + ", ".join(f"{v} = zero" for v in X + Y) + ";")
    if in_regs:
        w("\t__m256i " + ", ".join(f"A{j} = L(a1, {j})" for j in range(R)) + ";")
        w("\t__m256i " + ", ".join(f"B{j} = L(a2, {j})" for j in range(R)) + ";")
        A = [f"A{j}" for j in range(R)]
        B = [f"B{j}" for j in range(R)]
    else:
        A = [f"L(a1, {j})" for j in range(R)]
        B = [f"L(a2, {j})" for j in range(R)]
    M = [f"L(m1, {j})" for j in range(R)]
    N = [f"L(m2, {j})" for j in range(R)]
    w("\tuint64_t x1 = 0, x2 = 0;")
    w("\tint i;")
    w("")
    w(f"\tfor (i = 0; i < {d}; i++) {{")
    w("\t\tuint64_t bi = b1[i], ci = b2[i], y, z;")
    w("\t\t/* lane 0 in full: the products' high halves, which are lane 1's, too */")
    w("\t\tunsigned __int128 s1 = (unsigned __int128) a1[0] * bi + x1;")
    w("\t\tunsigned __int128 s2 = (unsigned __int128) a2[0] * ci + x2;")
    w("\t\t__m256i bv = _mm256_set1_epi64x((long long) bi), cv = _mm256_set1_epi64x((long long) ci);")
    w("\t\t__m256i yv, zv;")
    w("")
    w("\t\ty = ((uint64_t) s1 * k1) & IFMA_MASK52;")
    w("\t\tz = ((uint64_t) s2 * k2) & IFMA_MASK52;")
    w("\t\ts1 += (unsigned __int128) m1[0] * y;")
    w("\t\ts2 += (unsigned __int128) m2[0] * z;")
    w("\t\tyv = _mm256_set1_epi64x((long long) y);")
    w("\t\tzv = _mm256_set1_epi64x((long long) z);")
    for j in range(R):
        w(f"\t\tX{j} = _mm256_madd52lo_epu64(X{j}, {A[j]}, bv); Y{j} = _mm256_madd52lo_epu64(Y{j}, {B[j]}, cv);")
    for j in range(R):
        w(f"\t\tX{j} = _mm256_madd52lo_epu64(X{j}, {M[j]}, yv); Y{j} = _mm256_madd52lo_epu64(Y{j}, {N[j]}, zv);")
    w("\t\t/* down a lane: lane 0 is a multiple of 2^52, and s carries it */")
    for j in range(R):
        nx = f"X{j + 1}" if j + 1 < R else "zero"
        ny = f"Y{j + 1}" if j + 1 < R else "zero"
        w(f"\t\tX{j} = _mm256_alignr_epi64({nx}, X{j}, 1); Y{j} = _mm256_alignr_epi64({ny}, Y{j}, 1);")
    w("\t\tx1 = (uint64_t) _mm_cvtsi128_si64(_mm256_castsi256_si128(X0)) + (uint64_t)(s1 >> 52);")
    w("\t\tx2 = (uint64_t) _mm_cvtsi128_si64(_mm256_castsi256_si128(Y0)) + (uint64_t)(s2 >> 52);")
    for j in range(R):
        w(f"\t\tX{j} = _mm256_madd52hi_epu64(X{j}, {A[j]}, bv); Y{j} = _mm256_madd52hi_epu64(Y{j}, {B[j]}, cv);")
    for j in range(R):
        w(f"\t\tX{j} = _mm256_madd52hi_epu64(X{j}, {M[j]}, yv); Y{j} = _mm256_madd52hi_epu64(Y{j}, {N[j]}, zv);")
    w("\t}")
    w("#undef L")
    w("\t/* lane 0, with the carry it was never given, is what the scalar holds */")
    w("\tX0 = _mm256_mask_set1_epi64(X0, 1, (long long) x1);")
    w("\tY0 = _mm256_mask_set1_epi64(Y0, 1, (long long) x2);")
    w(f"\tifma_norm{R}(r1, " + ", ".join(X) + ");")
    w(f"\tifma_norm{R}(r2, " + ", ".join(Y) + ");")
    w("}")
    w("")
    return "\n".join(o)


def gen_sel(R):
    o = []
    w = o.append
    w(f"/* both tables' entries k1 and k2 of 32, every entry read and a mask keeping")
    w(f" * one; entries are {4 * R} words apart */")
    w(f"IFMA_TARGET static void ifma_sel{R}(uint64_t *o1, const uint64_t *t1, uint64_t k1,")
    w(f"                                uint64_t *o2, const uint64_t *t2, uint64_t k2)")
    w("{")
    w("\tconst __m256i one = _mm256_set1_epi64x(1);")
    w("\t__m256i v1 = _mm256_set1_epi64x((long long) k1), v2 = _mm256_set1_epi64x((long long) k2);")
    w("\t__m256i kv = _mm256_setzero_si256();")
    w("\t__m256i " + ", ".join(f"P{j} = kv, Q{j} = kv" for j in range(R)) + ";")
    w("\tint k;")
    w("\tfor (k = 0; k < 32; k++) {")
    w("\t\t__mmask8 s1 = _mm256_cmpeq_epi64_mask(v1, kv), s2 = _mm256_cmpeq_epi64_mask(v2, kv);")
    w(f"\t\tconst uint64_t *e1 = t1 + {4 * R} * k, *e2 = t2 + {4 * R} * k;")
    for j in range(R):
        w(f"\t\tP{j} = _mm256_mask_loadu_epi64(P{j}, s1, e1 + {4 * j}); Q{j} = _mm256_mask_loadu_epi64(Q{j}, s2, e2 + {4 * j});")
    w("\t\tkv = _mm256_add_epi64(kv, one);")
    w("\t}")
    for j in range(R):
        w(f"\t_mm256_storeu_si256((__m256i *)(o1 + {4 * j}), P{j}); _mm256_storeu_si256((__m256i *)(o2 + {4 * j}), Q{j});")
    w("}")
    w("")
    return "\n".join(o)


def main():
    print(HEADER)
    rs = sorted({regs(d) for d in SIZES})
    for R in rs:
        print(gen_norm(R))
        print(gen_sel(R))
    for d in SIZES:
        print(gen_amm(d))
    print(FOOTER, end="")


if __name__ == "__main__":
    main()