Hugging Face's logo

Snapkitty
/
flash-attention-rtl

snapkitty
october-2026-drop
systemverilog
rtl
systolic-array
subleq
flash-attention
cocotb
hardware
nim
python
Model card Files
xet
 Community
flash-attention-rtl / cstack /tests /test_l1_uac.c
SNAPKITTYWEST's picture
SNAPKITTYWEST
October 2026 main drop: mirror from GitHub
2166d7e verified
Raw History Blame Contribute Delete
9.21 kB
#include "cs_test.h"
#include "l1_goldilocks.h"
#include "l1_pmat.h"
#include "l1_tensor.h"
#include <string.h>
__extension__ typedef unsigned __int128 u128;
static uint64_t rng_state = 0x123456789abcdef1ull;
static uint64_t rnd(void)
{
rng_state ^= rng_state << 13; rng_state ^= rng_state >> 7; rng_state ^= rng_state << 17;
return rng_state;
}
static gl_t rnd_gl(void) { return gl_reduce(rnd()); }
static void fill(pmat *m) { for (size_t i = 0; i < m->rows * m->cols; i++) m->d[i] = rnd_gl(); }
TEST_MAIN_BEGIN("l1 uac")
CHECK_EQ_STATUS(gl_selftest(), CS_OK);
/* edge values vs a __int128 reference */
const gl_t edges[] = {0, 1, 2, 0xFFFFFFFFull, 0x100000000ull, 0x100000001ull,
GL_P - 2, GL_P - 1, GL_P / 2, GL_P / 2 + 1, 0xFFFFFFFE00000001ull};
for (size_t i = 0; i < sizeof edges / sizeof edges[0]; i++)
for (size_t j = 0; j < sizeof edges / sizeof edges[0]; j++) {
gl_t a = edges[i], b = edges[j];
CHECK(gl_mul(a, b) == (gl_t)(((u128)a * b) % GL_P));
CHECK(gl_add(a, b) == (gl_t)(((u128)a + b) % GL_P));
CHECK(gl_sub(a, b) == (gl_t)(((u128)a + GL_P - b) % GL_P));
}
/* random vs reference */
int bad = 0;
for (int i = 0; i < 200000; i++) {
gl_t a = rnd_gl(), b = rnd_gl();
if (gl_mul(a, b) != (gl_t)(((u128)a * b) % GL_P)) bad++;
if (gl_add(a, b) != (gl_t)(((u128)a + b) % GL_P)) bad++;
if (gl_sub(a, b) != (gl_t)(((u128)a + GL_P - b) % GL_P)) bad++;
if (gl_add(a, gl_neg(a)) != 0) bad++;
}
CHECK(bad == 0);
/* field axioms on random triples */
bad = 0;
for (int i = 0; i < 20000; i++) {
gl_t a = rnd_gl(), b = rnd_gl(), c = rnd_gl();
if (gl_mul(a, gl_add(b, c)) != gl_add(gl_mul(a, b), gl_mul(a, c))) bad++; /* distributive */
if (gl_mul(gl_mul(a, b), c) != gl_mul(a, gl_mul(b, c))) bad++; /* associative */
gl_t inv;
if (a != 0 && (gl_inv(a, &inv) != CS_OK || gl_mul(a, inv) != 1)) bad++;
}
CHECK(bad == 0);
gl_t q;
CHECK_EQ_STATUS(gl_inv(0, &q), CS_ERR_SINGULAR);
CHECK_EQ_STATUS(gl_div(5, 0, &q), CS_ERR_SINGULAR);
CHECK_EQ_STATUS(gl_div(10, 2, &q), CS_OK);
CHECK(q == 5);
CHECK(gl_pow(3, 0) == 1 && gl_pow(0, 5) == 0 && gl_pow(GL_P - 1, 3) == GL_P - 1);
/* matrices */
pmat a, b, c, d, id;
CHECK_EQ_STATUS(pmat_new(&a, 0, 3), CS_ERR_ARG);
CHECK_EQ_STATUS(pmat_new(&a, (size_t)-1, 8), CS_ERR_NOMEM);
CHECK_EQ_STATUS(pmat_new(&a, 3, 4), CS_OK);
CHECK_EQ_STATUS(pmat_new(&b, 4, 2), CS_OK);
fill(&a); fill(&b);
CHECK_EQ_STATUS(pmat_mul(&a, &a, &c), CS_ERR_RANGE); /* 3x4 * 3x4 */
CHECK_EQ_STATUS(pmat_mul(&a, &b, &c), CS_OK);
CHECK(c.rows == 3 && c.cols == 2);
gl_t ref = 0;
for (size_t k = 0; k < 4; k++) ref = gl_add(ref, gl_mul(pmat_get(&a, 1, k), pmat_get(&b, k, 1)));
CHECK(pmat_get(&c, 1, 1) == ref);
/* (AB)^T = B^T A^T */
pmat abt, at, bt, btat;
CHECK_EQ_STATUS(pmat_transpose(&c, &abt), CS_OK);
CHECK_EQ_STATUS(pmat_transpose(&a, &at), CS_OK);
CHECK_EQ_STATUS(pmat_transpose(&b, &bt), CS_OK);
CHECK_EQ_STATUS(pmat_mul(&bt, &at, &btat), CS_OK);
CHECK(pmat_equal(&abt, &btat));
pmat_free(&abt); pmat_free(&at); pmat_free(&bt); pmat_free(&btat);
/* add / sub / scale */
pmat a2, s1, s2;
CHECK_EQ_STATUS(pmat_add(&a, &a, &a2), CS_OK);
CHECK_EQ_STATUS(pmat_scale(&a, 2, &s1), CS_OK);
CHECK(pmat_equal(&a2, &s1));
CHECK_EQ_STATUS(pmat_sub(&a2, &a, &s2), CS_OK);
CHECK(pmat_equal(&s2, &a));
CHECK_EQ_STATUS(pmat_add(&a, &b, &d), CS_ERR_RANGE);
pmat_free(&a2); pmat_free(&s1); pmat_free(&s2);
pmat_free(&a); pmat_free(&b); pmat_free(&c);
/* inverse / determinant */
int inv_bad = 0, det_bad = 0;
for (int trial = 0; trial < 40; trial++) {
size_t n = 1 + (size_t)(rnd() % 6);
pmat m, mi, prod, m2, prod2;
pmat_new(&m, n, n);
fill(&m);
gl_t det;
CHECK_EQ_STATUS(pmat_det(&m, &det), CS_OK);
if (pmat_inverse(&m, &mi) == CS_OK) {
pmat_mul(&m, &mi, &prod);
pmat_identity(&id, n);
if (!pmat_equal(&prod, &id)) inv_bad++;
pmat_mul(&mi, &m, &prod2);
if (!pmat_equal(&prod2, &id)) inv_bad++;
/* det(A) * det(A^-1) = 1 */
gl_t det_i;
pmat_det(&mi, &det_i);
if (gl_mul(det, det_i) != 1) det_bad++;
pmat_free(&mi); pmat_free(&prod); pmat_free(&prod2); pmat_free(&id);
}
/* det(A^T) = det(A) */
pmat_transpose(&m, &m2);
gl_t det_t;
pmat_det(&m2, &det_t);
if (det_t != det) det_bad++;
pmat_free(&m); pmat_free(&m2);
}
CHECK(inv_bad == 0 && det_bad == 0);
/* known 2x2: [[1,2],[3,4]] det = -2, inverse = [[-2,1],[3/2,-1/2]] */
pmat k, ki;
pmat_new(&k, 2, 2);
pmat_set(&k, 0, 0, 1); pmat_set(&k, 0, 1, 2); pmat_set(&k, 1, 0, 3); pmat_set(&k, 1, 1, 4);
gl_t kd;
CHECK_EQ_STATUS(pmat_det(&k, &kd), CS_OK);
CHECK(kd == gl_neg(2));
CHECK_EQ_STATUS(pmat_inverse(&k, &ki), CS_OK);
gl_t half;
gl_inv(2, &half);
CHECK(pmat_get(&ki, 0, 0) == gl_neg(2) && pmat_get(&ki, 0, 1) == 1);
CHECK(pmat_get(&ki, 1, 0) == gl_mul(3, half) && pmat_get(&ki, 1, 1) == gl_neg(half));
pmat_free(&ki);
/* singular: second row = 2 * first row */
pmat_set(&k, 1, 0, 2); pmat_set(&k, 1, 1, 4);
CHECK_EQ_STATUS(pmat_inverse(&k, &ki), CS_ERR_SINGULAR);
CHECK_EQ_STATUS(pmat_det(&k, &kd), CS_OK);
CHECK(kd == 0);
pmat_free(&k);
pmat rect;
pmat_new(&rect, 2, 3);
CHECK_EQ_STATUS(pmat_inverse(&rect, &ki), CS_ERR_RANGE);
pmat_free(&rect);
/* tensors */
size_t sh2[2] = {3, 4}, sh3[3] = {2, 3, 4};
tensor t;
CHECK_EQ_STATUS(tensor_new(&t, 0, sh2), CS_ERR_ARG);
CHECK_EQ_STATUS(tensor_new(&t, CS_TENSOR_MAX_RANK + 1, sh3), CS_ERR_ARG);
size_t zero_shape[2] = {3, 0};
CHECK_EQ_STATUS(tensor_new(&t, 2, zero_shape), CS_ERR_ARG);
CHECK_EQ_STATUS(tensor_new(&t, 3, sh3), CS_OK);
size_t idx[3] = {1, 2, 3};
CHECK_EQ_STATUS(tensor_set(&t, idx, GL_P + 5), CS_OK); /* reduced on set */
gl_t v;
CHECK_EQ_STATUS(tensor_get(&t, idx, &v), CS_OK);
CHECK(v == 5);
CHECK(t.d[1 * 12 + 2 * 4 + 3] == 5); /* row-major layout */
idx[1] = 3;
CHECK_EQ_STATUS(tensor_get(&t, idx, &v), CS_ERR_RANGE);
tensor_free(&t);
/* contraction of rank-2 tensors == matrix multiplication */
pmat ma, mb, mc, mc2;
pmat_new(&ma, 3, 5); pmat_new(&mb, 5, 2);
fill(&ma); fill(&mb);
tensor ta, tb, tc;
tensor_from_pmat(&ma, &ta);
tensor_from_pmat(&mb, &tb);
CHECK_EQ_STATUS(tensor_contract(&ta, 1, &tb, 0, &tc), CS_OK);
CHECK_EQ_STATUS(tensor_to_pmat(&tc, &mc), CS_OK);
CHECK_EQ_STATUS(pmat_mul(&ma, &mb, &mc2), CS_OK);
CHECK(pmat_equal(&mc, &mc2));
CHECK_EQ_STATUS(tensor_contract(&ta, 0, &tb, 0, &tc), CS_ERR_RANGE); /* 3 != 5 */
CHECK_EQ_STATUS(tensor_contract(&ta, 2, &tb, 0, &tc), CS_ERR_RANGE); /* bad axis */
tensor_free(&tc);
/* contracting over a different axis pair = multiplication by a transpose */
pmat mat_t;
pmat_transpose(&mb, &mat_t); /* 2x5 */
tensor tbt;
tensor_from_pmat(&mat_t, &tbt);
CHECK_EQ_STATUS(tensor_contract(&ta, 1, &tbt, 1, &tc), CS_OK);
pmat mc3;
tensor_to_pmat(&tc, &mc3);
CHECK(pmat_equal(&mc3, &mc2));
pmat_free(&mc3); pmat_free(&mat_t); tensor_free(&tbt); tensor_free(&tc);
/* outer product, hadamard, add */
size_t len3[1] = {3}, len2[1] = {2};
tensor x, y, o;
tensor_new(&x, 1, len3); tensor_new(&y, 1, len2);
for (size_t i = 0; i < 3; i++) x.d[i] = rnd_gl();
for (size_t i = 0; i < 2; i++) y.d[i] = rnd_gl();
CHECK_EQ_STATUS(tensor_outer(&x, &y, &o), CS_OK);
CHECK(o.rank == 2 && o.shape[0] == 3 && o.shape[1] == 2);
size_t oi[2] = {2, 1};
tensor_get(&o, oi, &v);
CHECK(v == gl_mul(x.d[2], y.d[1]));
/* contracting an outer product with a vector recovers a scaled vector */
tensor w, z;
tensor_new(&w, 1, len2);
w.d[0] = 1; w.d[1] = 0;
CHECK_EQ_STATUS(tensor_contract(&o, 1, &w, 0, &z), CS_OK);
CHECK(z.rank == 1 && z.d[0] == gl_mul(x.d[0], y.d[0]) && z.d[2] == gl_mul(x.d[2], y.d[0]));
CHECK_EQ_STATUS(tensor_hadamard(&x, &y, &z), CS_ERR_RANGE);
tensor x2, h;
CHECK_EQ_STATUS(tensor_add(&x, &x, &x2), CS_OK);
CHECK_EQ_STATUS(tensor_hadamard(&x, &x, &h), CS_OK);
CHECK(x2.d[1] == gl_add(x.d[1], x.d[1]) && h.d[1] == gl_mul(x.d[1], x.d[1]));
tensor big_a, big_b, big_o;
size_t s22[2] = {1, 1}, s222[3] = {1, 1, 1};
tensor_new(&big_a, 2, s22); tensor_new(&big_b, 3, s222);
CHECK_EQ_STATUS(tensor_outer(&big_a, &big_b, &big_o), CS_ERR_RANGE); /* rank 5 > max */
tensor_free(&big_a); tensor_free(&big_b);
tensor_free(&x); tensor_free(&y); tensor_free(&o); tensor_free(&w); tensor_free(&z);
tensor_free(&x2); tensor_free(&h);
tensor_free(&ta); tensor_free(&tb);
pmat_free(&ma); pmat_free(&mb); pmat_free(&mc); pmat_free(&mc2);
TEST_MAIN_END