tritium demo mingpt
transcript ·
python ·
c ·
rust
/* TRITIUM-PROVENANCE sha256=6d7ce1815a5ea4eacf5df829ac23026deec566513c6bfb4d5c0fbfd4974218ee lang=c */
/* TRITIUM-HEIGHT high=0.13 mid=0.43 low=0.44 fns=15 */
#include <stdint.h>
#include <math.h>
// TRITIUM-PSEUDO fn=fmean
//| fn fmean(vals, n):
//| return x / (double)(n)
// TRITIUM-HEIGHT fn=fmean high=0.00 mid=0.41 low=0.59
double fmean(double* vals, int64_t n) {
int64_t i;
double x;
x = 0.0;
i = 0L;
while (i < n) {
x = x + vals[i];
i = (i + 1L);
}
return x / (double)(n);
}
// TRITIUM-PSEUDO fn=fsqrt
//| fn fsqrt(x):
//| return 0.0
//| return i_2
// TRITIUM-HEIGHT fn=fsqrt high=0.00 mid=0.69 low=0.31
double fsqrt(double x) {
int64_t i;
int fcmp4;
double i_2;
double x_2;
fcmp4 = x <= 0.0;
if (fcmp4) {
return 0.0;
} else {
i_2 = x;
i = 0L;
while (i < 8L) {
x_2 = 0.5 * (i_2 + (x / i_2));
i_2 = x_2;
i = (i + 1L);
}
return i_2;
}
}
// TRITIUM-PSEUDO fn=layer_norm
//| fn layer_norm(vals, n, out, eps):
//| fmean(...)
//| fsqrt(...)
//| return mu
// TRITIUM-HEIGHT fn=layer_norm high=0.20 mid=0.36 low=0.44
double layer_norm(double* vals, int64_t n, double* out, double eps) {
int64_t i;
double mu;
double x;
double x_2;
double x_3;
mu = fmean(vals, n);
x_2 = 0.0;
i = 0L;
while (i < n) {
x = vals[i] - mu;
x_2 = x_2 + (x * x);
i = (i + 1L);
}
x_2 = (x_2 / (double)(n));
x_3 = 1.0 / (fsqrt(x_2 + eps));
i = 0L;
while (i < n) {
out[i] = ((vals[i] - mu) * x_3);
i = (i + 1L);
}
return mu;
}
// TRITIUM-PSEUDO fn=ln_rows
//| fn ln_rows(x, n_tokens, d, out, tmp, nrm, eps):
//| layer_norm(...)
//| return x_2
// TRITIUM-HEIGHT fn=ln_rows high=0.11 mid=0.33 low=0.56
double ln_rows(double* x, int64_t n_tokens, int64_t d, double* out, double* tmp, double* nrm, double eps) {
int64_t t;
int64_t i;
double x_2;
t = 0L;
while (t < n_tokens) {
i = 0L;
while (i < d) {
tmp[i] = x[((t * d) + i)];
i = (i + 1L);
}
layer_norm(tmp, d, nrm, eps);
i = 0L;
while (i < d) {
out[((t * d) + i)] = nrm[i];
i = (i + 1L);
}
t = (t + 1L);
}
x_2 = (double)(n_tokens) * 1.0;
return x_2;
}
// TRITIUM-PSEUDO fn=gemv_f
//| fn gemv_f(a, x, n, m, out):
//| return n
// TRITIUM-HEIGHT fn=gemv_f high=0.00 mid=0.38 low=0.62
int64_t gemv_f(double* a, double* x, int64_t n, int64_t m, double* out) {
int64_t i;
int64_t j;
double x_2;
i = 0L;
while (i < n) {
x_2 = 0.0;
j = 0L;
while (j < m) {
x_2 = x_2 + (a[((i * m) + j)] * x[j]);
j = (j + 1L);
}
out[i] = x_2;
i = (i + 1L);
}
return n;
}
// TRITIUM-PSEUDO fn=qkv_seq
//| fn qkv_seq(x, n_tokens, d, w, out, tmp, row):
//| gemv_f(...)
//| return n_tokens
// TRITIUM-HEIGHT fn=qkv_seq high=0.11 mid=0.30 low=0.58
int64_t qkv_seq(double* x, int64_t n_tokens, int64_t d, double* w, double* out, double* tmp, double* row) {
int64_t t;
int64_t i;
t = 0L;
while (t < n_tokens) {
i = 0L;
while (i < d) {
tmp[i] = x[((t * d) + i)];
i = (i + 1L);
}
gemv_f(w, tmp, d, d, row);
i = 0L;
while (i < d) {
out[((t * d) + i)] = row[i];
i = (i + 1L);
}
t = (t + 1L);
}
return n_tokens;
}
// TRITIUM-PSEUDO fn=pack_heads
//| fn pack_heads(src, n_tokens, heads, dh, dest):
//| return (heads * n_tokens) * dh
// TRITIUM-HEIGHT fn=pack_heads high=0.00 mid=0.38 low=0.62
int64_t pack_heads(double* src, int64_t n_tokens, int64_t heads, int64_t dh, double* dest) {
int64_t h;
int64_t t;
int64_t u;
h = 0L;
while (h < heads) {
t = 0L;
while (t < n_tokens) {
u = 0L;
while (u < dh) {
dest[((((h * n_tokens) * dh) + (t * dh)) + u)] = src[((((t * heads) * dh) + (h * dh)) + u)];
u = (u + 1L);
}
t = (t + 1L);
}
h = (h + 1L);
}
return (heads * n_tokens) * dh;
}
// TRITIUM-PSEUDO fn=fexp
//| fn fexp(x):
//| return x_5 / x_4
//| return x_5 * x_4
// TRITIUM-HEIGHT fn=fexp high=0.00 mid=0.76 low=0.24
double fexp(double x) {
double ln2;
double inv_ln2;
int64_t i;
int64_t j;
int fcmp15;
int64_t ifexp_2;
double ifexp;
int64_t xi;
double x_2;
double x_5;
int64_t i_2;
double x_4;
double x_3;
ln2 = 0.6931471805599453;
inv_ln2 = 1.4426950408889634;
fcmp15 = x >= 0.0;
if (fcmp15) {
ifexp_2 = 0.5;
} else {
ifexp_2 = -0.5;
}
ifexp = ifexp_2;
xi = (int64_t)(((x * inv_ln2) + ifexp));
if (xi > 60L) {
xi = 60L;
}
if (xi < -60L) {
xi = -60L;
}
x_5 = 1.0;
x_2 = 1.0;
i = 1L;
while (i <= 14L) {
x_2 = (x_2 * (x - ((double)(xi) * ln2))) / (double)(i);
x_5 = x_5 + x_2;
i = (i + 1L);
}
x_4 = 1.0;
i_2 = xi;
if (i_2 < 0L) {
i_2 = (0L - i_2);
j = 0L;
while (j < i_2) {
x_4 = x_4 * 2.0;
j = (j + 1L);
}
return x_5 / x_4;
} else {
j = 0L;
while (j < i_2) {
x_3 = x_4 * 2.0;
x_4 = x_3;
j = (j + 1L);
}
return x_5 * x_4;
}
}
// TRITIUM-PSEUDO fn=mha_fwd
//| fn mha_fwd(q, k, v, scores, out, h, l, d, scale):
//| fexp(...)
//| return (h * l) * d
// TRITIUM-HEIGHT fn=mha_fwd high=0.04 mid=0.47 low=0.50
int64_t mha_fwd(double* q, double* k, double* v, double* scores, double* out, int64_t h, int64_t l, int64_t d, double scale) {
int64_t heap_dist;
int64_t i;
int64_t j;
int64_t t;
double x_2;
double i_2;
int fcmp96;
double e;
double x;
heap_dist = 0L;
while (heap_dist < h) {
i = 0L;
while (i < l) {
j = 0L;
while (j < l) {
x_2 = 0.0;
t = 0L;
while (t < d) {
x_2 = x_2 + (q[((((heap_dist * l) * d) + (i * d)) + t)] * k[((((heap_dist * l) * d) + (j * d)) + t)]);
t = (t + 1L);
}
scores[((i * l) + j)] = (x_2 * scale);
j = (j + 1L);
}
i = (i + 1L);
}
i = 0L;
while (i < l) {
i_2 = scores[(i * l)];
j = 1L;
while (j < l) {
fcmp96 = scores[((i * l) + j)] > i_2;
if (fcmp96) {
i_2 = scores[((i * l) + j)];
}
j = (j + 1L);
}
x = 0.0;
j = 0L;
while (j < l) {
e = fexp((scores[((i * l) + j)] - i_2));
scores[((i * l) + j)] = e;
x = x + e;
j = (j + 1L);
}
j = 0L;
while (j < l) {
scores[((i * l) + j)] = (scores[((i * l) + j)] / x);
j = (j + 1L);
}
i = (i + 1L);
}
i = 0L;
while (i < l) {
t = 0L;
while (t < d) {
x_2 = 0.0;
j = 0L;
while (j < l) {
x_2 = (x_2 + (scores[((i * l) + j)] * v[((((heap_dist * l) * d) + (j * d)) + t)]));
j = (j + 1L);
}
out[((((heap_dist * l) * d) + (i * d)) + t)] = x_2;
t = (t + 1L);
}
i = (i + 1L);
}
heap_dist = (heap_dist + 1L);
}
return (h * l) * d;
}
// TRITIUM-PSEUDO fn=unpack_heads
//| fn unpack_heads(src, n_tokens, heads, dh, dest):
//| return (n_tokens * heads) * dh
// TRITIUM-HEIGHT fn=unpack_heads high=0.00 mid=0.38 low=0.62
int64_t unpack_heads(double* src, int64_t n_tokens, int64_t heads, int64_t dh, double* dest) {
int64_t h;
int64_t t;
int64_t u;
h = 0L;
while (h < heads) {
t = 0L;
while (t < n_tokens) {
u = 0L;
while (u < dh) {
dest[((((t * heads) * dh) + (h * dh)) + u)] = src[((((h * n_tokens) * dh) + (t * dh)) + u)];
u = (u + 1L);
}
t = (t + 1L);
}
h = (h + 1L);
}
return (n_tokens * heads) * dh;
}
// TRITIUM-PSEUDO fn=residual_add
//| fn residual_add(a, b, n, out):
//| return n
// TRITIUM-HEIGHT fn=residual_add high=0.00 mid=0.23 low=0.77
int64_t residual_add(double* a, double* b, int64_t n, double* out) {
int64_t i;
i = 0L;
while (i < n) {
out[i] = (a[i] + b[i]);
i = (i + 1L);
}
return n;
}
// TRITIUM-PSEUDO fn=ftanh
//| fn ftanh(x):
//| fexp(...)
//| return x_2 / x_3
// TRITIUM-HEIGHT fn=ftanh high=0.19 mid=0.81 low=0.00 role=leaf_charset
double ftanh(double x) {
double i;
int fcmp5;
int fcmp7;
double e;
double x_2;
double x_3;
i = x;
fcmp5 = i > 20.0;
if (fcmp5) {
i = 20.0;
}
fcmp7 = i < -20.0;
if (fcmp7) {
i = -20.0;
}
e = fexp(i + i);
x_2 = e - 1.0;
x_3 = e + 1.0;
return x_2 / x_3;
}
// TRITIUM-PSEUDO fn=gelu_tanh
//| fn gelu_tanh(x):
//| ftanh(...)
//| return x_5 * x_4
// TRITIUM-HEIGHT fn=gelu_tanh high=0.43 mid=0.57 low=0.00
double gelu_tanh(double x) {
double x_5;
double x_2;
double x_3;
double x_4;
x_5 = 0.5 * x;
x_2 = 0.044715 * ((x * x) * x);
x_3 = 0.79788456 * (x + x_2);
x_4 = 1.0 + (ftanh(x_3));
return x_5 * x_4;
}
// TRITIUM-PSEUDO fn=mlp_gelu_vec
//| fn mlp_gelu_vec(x, w1, w2, d, hidden, hidden_out, out):
//| gemv_f(...)
//| gelu_tanh(...)
//| return d
// TRITIUM-HEIGHT fn=mlp_gelu_vec high=0.46 mid=0.15 low=0.39
int64_t mlp_gelu_vec(double* x, double* w1, double* w2, int64_t d, int64_t hidden, double* hidden_out, double* out) {
int64_t i;
gemv_f(w1, x, hidden, d, hidden_out);
i = 0L;
while (i < hidden) {
hidden_out[i] = (gelu_tanh(hidden_out[i]));
i = (i + 1L);
}
gemv_f(w2, hidden_out, d, hidden, out);
return d;
}
// TRITIUM-PSEUDO fn=transformer_block
//| fn transformer_block(x, n_tokens, d, heads, dh, hidden, scale, wq, wk, wv, wo, w1, w2, q, k, v, qh, kh, vh, scores, attention, flat, normed, normed2, tmp, nrm, row, hidden_out, mlp_out, projected, residual, out):
//| ln_rows(...)
//| qkv_seq(...)
//| pack_heads(...)
//| mha_fwd(...)
//| unpack_heads(...)
//| gemv_f(...)
//| residual_add(...)
//| mlp_gelu_vec(...)
//| return n_tokens
// TRITIUM-HEIGHT fn=transformer_block high=0.47 mid=0.18 low=0.34
int64_t transformer_block(double* x, int64_t n_tokens, int64_t d, int64_t heads, int64_t dh, int64_t hidden, double scale, double* wq, double* wk, double* wv, double* wo, double* w1, double* w2, double* q, double* k, double* v, double* qh, double* kh, double* vh, double* scores, double* attention, double* flat, double* normed, double* normed2, double* tmp, double* nrm, double* row, double* hidden_out, double* mlp_out, double* projected, double* residual, double* out) {
int64_t t;
int64_t i;
ln_rows(x, n_tokens, d, normed, tmp, nrm, 0.0);
qkv_seq(normed, n_tokens, d, wq, q, tmp, row);
qkv_seq(normed, n_tokens, d, wk, k, tmp, row);
qkv_seq(normed, n_tokens, d, wv, v, tmp, row);
pack_heads(q, n_tokens, heads, dh, qh);
pack_heads(k, n_tokens, heads, dh, kh);
pack_heads(v, n_tokens, heads, dh, vh);
mha_fwd(qh, kh, vh, scores, attention, heads, n_tokens, dh, scale);
unpack_heads(attention, n_tokens, heads, dh, flat);
t = 0L;
while (t < n_tokens) {
i = 0L;
while (i < d) {
tmp[i] = flat[((t * d) + i)];
i = (i + 1L);
}
gemv_f(wo, tmp, d, d, row);
i = 0L;
while (i < d) {
projected[((t * d) + i)] = row[i];
i = (i + 1L);
}
t = (t + 1L);
}
residual_add(x, projected, n_tokens * d, residual);
ln_rows(residual, n_tokens, d, normed2, tmp, nrm, 0.0);
t = 0L;
while (t < n_tokens) {
i = 0L;
while (i < d) {
tmp[i] = normed2[((t * d) + i)];
i = (i + 1L);
}
mlp_gelu_vec(tmp, w1, w2, d, hidden, hidden_out, row);
i = 0L;
while (i < d) {
mlp_out[((t * d) + i)] = row[i];
i = (i + 1L);
}
t = (t + 1L);
}
residual_add(residual, mlp_out, n_tokens * d, out);
return n_tokens;
}