skip to content
==> include/mesh/ref/gpt2_linear.h <==
1/* @ajax/mesh-ref-gpt2-linear: a linear layer's output row in the order the GPU's GEMM tile and GEMV
2 * give exactly (mesh-ref gpt2.rs `project_row`, `project_cols`): out[n] = (the sum over k ascending
3 * from +0 of x[k] w[k][n]) then the epilogue. Weights [k][nout] (HF Conv1D). Freestanding. */
4#ifndef MESH_REF_GPT2_LINEAR_H
5#define MESH_REF_GPT2_LINEAR_H
6
7#include <stddef.h>
8
9/* How an output element is finished; acc is the sequential dot. */
10typedef enum {
11    MESH_GPT2_EPI_BIAS,         /* acc + b */
12    MESH_GPT2_EPI_BIAS_GELU,    /* gelu(acc + b) */
13    MESH_GPT2_EPI_BIAS_RESIDUAL /* out + (acc + b): a residual add in place */
14} mesh_gpt2_epi;
15
16/* out[n] for n in 0..nout from x[0..kd) and w (kd x nout); out is read for the residual. */
17void mesh_gpt2_project_row(const float *x, size_t kd, const float *w, const float *b, size_t nout,
18                           mesh_gpt2_epi epi, float *out);
19/* project_row for the columns c0 .. c0 + ncols (out holds just those). */
20void mesh_gpt2_project_cols(const float *x, size_t kd, const float *w, const float *b, size_t nout,
21                            mesh_gpt2_epi epi, size_t c0, size_t ncols, float *out);
22
23#endif