test(math-cuda): pin the RPX kernel's canonicalisation loop with a deterministic witness (lane K, follow-up to #952)

crypto/math-cuda/tests/host_kat/rpx_canon_witness.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""Derives the "canonicalisation witness" row of `rpx_kat_vectors.h`.
+
+WHY. `rpx::permute` (kernels/rpx.cu) ends in a loop that canonicalises the
+state, which is what makes device digests byte-comparable to the host's. A
+known-answer check cannot see that loop unless some output lane is a raw twin
+(`value + p`, in `[p, 2^64)`) before it — a 2^-32 event per lane on random
+inputs. This script builds an input for which it is certain.
+
+HOW. The permutation's last operation is `out_i = add(m_i, ARK1[6][i])`, where
+`m_i` is the M-round MDS output. With `m_i` canonical and
+`m_i + ARK1[6][i] < 2^64`, the device `add` returns `m_i + ARK1[6][i]` as is; if
+that sum lies in `[p, 2^64)` it is the raw twin of `sum − p`. So choose the
+canonical MDS output `u` with `u_0 = p − ARK1[6][0] + 1` (raw `out_0 = p + 1`,
+field value 1), fill the other eleven lanes at random, invert the MDS to get the
+M-round input, and invert rounds 5..0 — `x^{1/7}` in `GF(p³)` for the E rounds,
+`x^7` / `MDS⁻¹` / `x^{1/7}` / `MDS⁻¹` for the FB rounds — to get the
+permutation input. `m_0` cannot itself be a twin (`u_0 + p > 2^64`), so the raw
+lane is deterministic whatever representation the earlier rounds happen to
+carry.
+
+TRUST. This is a THIRD transcription of the permutation, so it trusts nothing
+about itself: before printing, it reproduces every row of the header's Table 2
+forward and inverts each one back to its input. Run from anywhere:
+
+    python3 crypto/math-cuda/tests/host_kat/rpx_canon_witness.py
+
+The printed input goes into `prover/tests/rpx_host_kat_vectors.rs`
+(`permutation_inputs`, the row named "canonicalisation witness"); its output
+row comes from that generator, never from here.
+"""
+import pathlib
+import random
+import re
+
+REPO = pathlib.Path(__file__).resolve().parents[4]
+P = (1 << 64) - (1 << 32) + 1
+INV_ALPHA = 10540996611094048183  # rpo.rs:96
+assert (7 * INV_ALPHA) % (P - 1) == 1
+ROW = [7, 23, 8, 26, 13, 10, 9, 7, 6, 22, 21, 8]  # rpo.rs:114
+
+RPO_RS = (REPO / "prover/src/lfm/rpo.rs").read_text()
+
+
+def constant_table(name):
+    m = re.search(
+        r"pub const %s: \[\[u64; HASH_STATE_FELTS\]; NUM_ROUNDS\] = \[(.*?)\n\];" % name,
+        RPO_RS,
+        re.S,
+    )
+    rows = re.findall(r"\[\s*((?:\d+,\s*)+)\]", m.group(1))
+    vals = [[int(x) for x in re.findall(r"\d+", r)] for r in rows]
+    assert len(vals) == 7 and all(len(r) == 12 for r in vals), name
+    return vals
+
+
+ARK1, ARK2 = constant_table("ARK1"), constant_table("ARK2")
+
+
+# --- the field, the MDS and its inverse, the cubic extension -------------------------
+
+def mds(s):
+    return [sum(ROW[(j - i) % 12] * s[j] for j in range(12)) % P for i in range(12)]
+
+
+def matrix_inverse_mod_p(m):
+    n = len(m)
+    a = [row[:] + [1 if i == j else 0 for j in range(n)] for i, row in enumerate(m)]
+    for col in range(n):
+        piv = next(r for r in range(col, n) if a[r][col] % P)
+        a[col], a[piv] = a[piv], a[col]
+        inv = pow(a[col][col], P - 2, P)
+        a[col] = [(v * inv) % P for v in a[col]]
+        for r in range(n):
+            if r != col and a[r][col]:
+                f = a[r][col]
+                a[r] = [(vr - f * vc) % P for vr, vc in zip(a[r], a[col])]
+    return [row[n:] for row in a]
+
+
+MDS_INV = matrix_inverse_mod_p([[ROW[(j - i) % 12] for j in range(12)] for i in range(12)])
+
+
+def mds_inv(s):
+    return [sum(MDS_INV[i][j] * s[j] for j in range(12)) % P for i in range(12)]
+
+
+def ext_mul(a, b):  # rpx.rs:118-125, φ³ = φ + 1
+    return [
+        (a[0] * b[0] + a[1] * b[2] + a[2] * b[1]) % P,
+        (a[0] * b[1] + a[1] * b[0] + a[1] * b[2] + a[2] * b[1] + a[2] * b[2]) % P,
+        (a[0] * b[2] + a[1] * b[1] + a[2] * b[0] + a[2] * b[2]) % P,
+    ]
+
+
+def ext_pow(a, e):
+    r, b = [1, 0, 0], a[:]
+    while e:
+        if e & 1:
+            r = ext_mul(r, b)
+        b = ext_mul(b, b)
+        e >>= 1
+    return r
+
+
+EXT_INV7 = pow(7, -1, P**3 - 1)  # x ↦ x^7 permutes GF(p³) (rpx.rs tests), so this exists
+
+
+# --- the permutation, forward (rpx.rs:280-316) and inverse ----------------------------
+
+def add_constants(s, table, r, sign=1):
+    return [(v + sign * table[r][i]) % P for i, v in enumerate(s)]
+
+
+def fb_round(s, r):
+    s = add_constants(mds(s), ARK1, r)
+    s = mds([pow(v, 7, P) for v in s])
+    return [pow(v, INV_ALPHA, P) for v in add_constants(s, ARK2, r)]
+
+
+def ext_round(s, r):
+    s = add_constants(s, ARK1, r)
+    return sum((ext_pow(s[3 * e:3 * e + 3], 7) for e in range(4)), [])
+
+
+def final_round(s):
+    return add_constants(mds(s), ARK1, 6)
+
+
+def permute(s):
+    for r in range(6):
+        s = fb_round(s, r) if r % 2 == 0 else ext_round(s, r)
+    return final_round(s)
+
+
+def fb_round_inv(s, r):
+    s = add_constants([pow(v, 7, P) for v in s], ARK2, r, -1)
+    s = [pow(v, INV_ALPHA, P) for v in mds_inv(s)]
+    return mds_inv(add_constants(s, ARK1, r, -1))
+
+
+def ext_round_inv(s, r):
+    s = sum((ext_pow(s[3 * e:3 * e + 3], EXT_INV7) for e in range(4)), [])
+    return add_constants(s, ARK1, r, -1)
+
+
+def rounds_0_to_5_inv(t):
+    for r in (5, 4, 3, 2, 1, 0):
+        t = ext_round_inv(t, r) if r % 2 == 1 else fb_round_inv(t, r)
+    return t
+
+
+# --- self-check against the header's oracle table before trusting any of the above ----
+
+HEADER = (REPO / "crypto/math-cuda/tests/host_kat/rpx_kat_vectors.h").read_text()
+body = HEADER.split("RPX_PERMUTATION_VECTORS[NUM_RPX_PERMUTATION_VECTORS] = {")[1].split("};")[0]
+rows = re.findall(r'\{"([^"]*)",\s*\{([^}]*)\},\s*\{([^}]*)\}\}', body)
+assert len(rows) >= 8, "oracle table has %d rows" % len(rows)
+for name, inp, outp in rows:
+    x = [int(v) for v in re.findall(r"\d+", inp)]
+    y = [int(v) for v in re.findall(r"\d+", outp)]
+    assert permute(x) == y, "forward transcription disagrees with the oracle on %r" % name
+    t = rounds_0_to_5_inv(mds_inv(add_constants(y, ARK1, 6, -1)))
+    assert t == x, "inverse permutation does not round-trip on %r" % name
+print("self-check: %d/%d oracle rows reproduced forward and inverted back" % (len(rows), len(rows)))
+
+# --- the witness -------------------------------------------------------------------------
+
+c0 = ARK1[6][0]
+rng = random.Random(0x4B57)  # "KW"; one generator, eleven draws
+u = [P - c0 + 1] + [rng.randrange(P) for _ in range(11)]
+assert P - c0 <= u[0] < P - c0 + (1 << 32) - 1
+x = rounds_0_to_5_inv(mds_inv(u))
+y = permute(x)
+assert y[0] == 1 and y == add_constants(u, ARK1, 6)
+print("witness input :", ", ".join(str(v) for v in x))
+print("witness output:", ", ".join(str(v) for v in y), "   (lane 0 raw on device: %d = p + 1)" % (u[0] + c0))
+
+# The generator's hard-coded row must be exactly this input, or the header's
+# witness and this derivation have drifted apart.
+GENERATOR = (REPO / "prover/tests/rpx_host_kat_vectors.rs").read_text()
+block = GENERATOR.split('"canonicalisation witness"', 1)[1].split("]", 1)[0]
+assert [int(v) for v in re.findall(r"\d+", block)] == x, "the generator's witness row is not this derivation's"
+print("generator row check: prover/tests/rpx_host_kat_vectors.rs carries this exact input")

crypto/math-cuda/tests/host_kat/rpx_host_kat.cpp

@@ -382,6 +382,15 @@ void seven_fb_rounds_reproduce_the_miden_rpo_vectors() {
 
 // ===========================================================================
 // Layer 4 — ★ the Rust oracle.
+//
+// ⚠ Every comparison here is RAW: `s[i] == v.output[i]`, never
+// `canon(s[i]) == …`. The tables are canonical by construction (the generator
+// canonicalises), and `permute` ends in a canonicalisation loop that makes
+// digests byte-comparable to the host's; a check that canonicalised the kernel
+// side would pass with that loop deleted, and so would a raw check on outputs
+// that merely happen to be canonical — all but a 2^-32 slice per lane. The
+// "canonicalisation witness" row and `the_canonicalisation_loop_is_pinned…`
+// below are what make the loop observable.
 // ===========================================================================
 
 void rpx_permutation_matches_the_rust_oracle() {
@@ -402,12 +411,12 @@ void rpx_permutation_matches_the_rust_oracle() {
         saw_p_minus_one = saw_p_minus_one || all_pm1;
         rpx::permute(s);
         bool ok = true;
-        for (int i = 0; i < 12; ++i) ok = ok && canon(s[i]) == v.output[i];
+        for (int i = 0; i < 12; ++i) ok = ok && s[i] == v.output[i];
         if (!ok) {
             printf("FAIL rpx permutation vector %d (%s)\n", n, v.name);
             for (int i = 0; i < 12; ++i) {
-                if (canon(s[i]) != v.output[i]) {
-                    printf("  lane %2d got %llu want %llu\n", i, (unsigned long long)canon(s[i]),
+                if (s[i] != v.output[i]) {
+                    printf("  lane %2d got %llu (raw) want %llu\n", i, (unsigned long long)s[i],
                            (unsigned long long)v.output[i]);
                 }
             }
@@ -479,11 +488,11 @@ void leaf_sponge_matches_the_rust_oracle() {
         uint64_t got[4];
         rpx::sponge_leaf(v.felts, v.len, got);
         bool ok = true;
-        for (int d = 0; d < 4; ++d) ok = ok && canon(got[d]) == v.digest[d];
+        for (int d = 0; d < 4; ++d) ok = ok && got[d] == v.digest[d];
         if (!ok) {
             printf("FAIL rpx leaf vector len=%u\n  got  %llu %llu %llu %llu\n  want %llu %llu %llu %llu\n",
-                   v.len, (unsigned long long)canon(got[0]), (unsigned long long)canon(got[1]),
-                   (unsigned long long)canon(got[2]), (unsigned long long)canon(got[3]),
+                   v.len, (unsigned long long)got[0], (unsigned long long)got[1],
+                   (unsigned long long)got[2], (unsigned long long)got[3],
                    (unsigned long long)v.digest[0], (unsigned long long)v.digest[1],
                    (unsigned long long)v.digest[2], (unsigned long long)v.digest[3]);
             ++failures;
@@ -504,11 +513,11 @@ void parent_matches_the_rust_oracle() {
         uint64_t got[4];
         rpx::compress(v.left, v.right, got);
         bool ok = true;
-        for (int d = 0; d < 4; ++d) ok = ok && canon(got[d]) == v.digest[d];
+        for (int d = 0; d < 4; ++d) ok = ok && got[d] == v.digest[d];
         if (!ok) {
             printf("FAIL rpx parent vector %d (%s)\n  got  %llu %llu %llu %llu\n  want %llu %llu %llu %llu\n",
-                   n, v.name, (unsigned long long)canon(got[0]), (unsigned long long)canon(got[1]),
-                   (unsigned long long)canon(got[2]), (unsigned long long)canon(got[3]),
+                   n, v.name, (unsigned long long)got[0], (unsigned long long)got[1],
+                   (unsigned long long)got[2], (unsigned long long)got[3],
                    (unsigned long long)v.digest[0], (unsigned long long)v.digest[1],
                    (unsigned long long)v.digest[2], (unsigned long long)v.digest[3]);
             ++failures;
@@ -530,6 +539,53 @@ void parent_matches_the_rust_oracle() {
            NUM_RPX_PARENT_VECTORS);
 }
 
+// ★ The pin on the canonicalisation loop. The witness row's M-round MDS output
+// lane 0 is `p − ARK1[6][0] + 1`, so the device's final `add` returns the raw
+// twin `p + 1` for a field value of 1 — deterministically, since neither that
+// sum nor the MDS reduction can wrap there. Replaying the rounds without the
+// loop must therefore show a lane ≥ p (or the witness has gone stale and no
+// longer witnesses anything), and `permute` must then return the oracle's
+// canonical digits RAW — which a kernel without the loop cannot.
+void the_canonicalisation_loop_is_pinned_by_the_witness() {
+    const RpxPermutationVector *w = nullptr;
+    for (int n = 0; n < NUM_RPX_PERMUTATION_VECTORS; ++n) {
+        if (strcmp(RPX_PERMUTATION_VECTORS[n].name, "canonicalisation witness") == 0) {
+            w = &RPX_PERMUTATION_VECTORS[n];
+        }
+    }
+    check(w != nullptr,
+          "the permutation table must carry the 'canonicalisation witness' row (run the generator, see rpx_kat_vectors.h)");
+    if (w == nullptr) return;
+
+    uint64_t s[12];
+    memcpy(s, w->input, sizeof(s));
+    rpx::fb_round<0>(s);
+    rpx::ext_round<1>(s);
+    rpx::fb_round<2>(s);
+    rpx::ext_round<3>(s);
+    rpx::fb_round<4>(s);
+    rpx::ext_round<5>(s);
+    rpx::final_round<6>(s);
+    int twins = 0;
+    for (int i = 0; i < 12; ++i) twins += (s[i] >= P) ? 1 : 0;
+    check(twins > 0, "the witness must leave a raw lane >= p before the canonicalisation loop");
+    check(s[0] == P + 1, "the witness's lane 0 must be the raw twin p + 1 before the loop");
+    for (int i = 0; i < 12; ++i) {
+        check(canon(s[i]) == w->output[i], "the witness's field values must be the oracle's");
+    }
+
+    uint64_t full[12];
+    memcpy(full, w->input, sizeof(full));
+    rpx::permute(full);
+    const bool loop_present = memcmp(full, w->output, sizeof(full)) == 0;
+    check(loop_present,
+          "permute must return the witness's digits RAW — the canonicalisation loop is missing");
+    if (loop_present) {
+        printf("★ canonicalisation pin: witness leaves %d raw lane(s) >= p before the loop; permute() returns them canonical\n",
+               twins);
+    }
+}
+
 // ===========================================================================
 // Layer 5 — negative controls and the representation.
 // ===========================================================================
@@ -671,6 +727,7 @@ int main() {
     rpx_permutation_matches_the_rust_oracle();
     leaf_sponge_matches_the_rust_oracle();
     parent_matches_the_rust_oracle();
+    the_canonicalisation_loop_is_pinned_by_the_witness();
     printf("\n-- layer 5: negative controls and representation --\n");
     rpx_is_not_rpo();
     raw_and_canonical_inputs_agree_and_outputs_are_canonical();

crypto/math-cuda/tests/host_kat/rpx_kat_vectors.h

@@ -60,7 +60,10 @@ inline constexpr uint64_t MIDEN_HASH_ELEMENTS[NUM_MIDEN_HASH_ELEMENTS][4] = {
 // is self-contained. All values are canonical (`< p`).
 //
 //   Table 2 — the bare permutation: all-zero, all-(p−1), `0..12`, alternating,
-//             two one-hot lanes, four seeded random states.
+//             two one-hot lanes, four seeded random states, and the row named
+//             "canonicalisation witness" — an input whose output lane 0 is the
+//             raw twin `p + 1` before the kernel's final canonicalisation loop
+//             (derived by `rpx_canon_witness.py`; the harness replays it).
 //   Table 3 — the leaf sponge (`algebraic_commit::sponge_leaf`) at 0, 1, 7, 8,
 //             9, 16 and 17 felts; `felts[]` is zero beyond `len`.
 //   Table 4 — the parent `compress(l, r)`.
@@ -89,7 +92,7 @@ struct RpxParentVector {
 //   cargo test -p lambda-vm-prover --test rpx_host_kat_vectors -- --ignored --nocapture
 // (prover/tests/rpx_host_kat_vectors.rs). Paste verbatim; do not edit by hand.
 
-inline constexpr int NUM_RPX_PERMUTATION_VECTORS = 10;
+inline constexpr int NUM_RPX_PERMUTATION_VECTORS = 11;
 inline constexpr RpxPermutationVector RPX_PERMUTATION_VECTORS[NUM_RPX_PERMUTATION_VECTORS] = {
     {"all-zero",
      {0ull, 0ull, 0ull, 0ull, 0ull, 0ull, 0ull, 0ull, 0ull, 0ull, 0ull, 0ull},
@@ -121,6 +124,9 @@ inline constexpr RpxPermutationVector RPX_PERMUTATION_VECTORS[NUM_RPX_PERMUTATIO
     {"random #4",
      {389113379214421922ull, 1947929307647562990ull, 667333451960644926ull, 3487966933876559811ull, 4195385248066926332ull, 2153180418459341747ull, 2727969323864685845ull, 29633526854483411ull, 990649808851061115ull, 1355410330370587755ull, 11605520071788416946ull, 4884409355120715354ull},
      {7025469669435110295ull, 17270957437800346011ull, 13702589935335807876ull, 3666927270871270796ull, 16666721215101099684ull, 531487850530305024ull, 15550553335698242665ull, 8959489596577675281ull, 11020601500923732075ull, 16110845767020565054ull, 4778394010005480449ull, 7715575140819562371ull}},
+    {"canonicalisation witness",
+     {15055324559807314153ull, 10242425218814686878ull, 9326602342065331773ull, 15451135068213333861ull, 17942679252967467289ull, 9284164080268346300ull, 5090350781253234438ull, 9328738269791029498ull, 18385380985273671691ull, 3238854716908013220ull, 5495049682105235955ull, 15773368383738726538ull},
+     {1ull, 9023883145409261355ull, 5839950281880325605ull, 5697668523532261268ull, 13033383890974728246ull, 14801658261553133914ull, 3025695522291518949ull, 12907720598453111556ull, 14827640614007773288ull, 14642633917625231592ull, 3090884930034198616ull, 2894057710100710233ull}},
 };
 
 inline constexpr int NUM_RPX_LEAF_VECTORS = 7;