← Commits · 767b4275
767b4275313849b7d5bfe34347bac0c09294f3b5
diff --git a/test.c b/test.c
index 81de536..a8c7362 100644
--- a/test.c
+++ b/test.c
@@ -194,9 +194,9 @@ static void test_linear_backward(ut_dev dev) {
ut_linear_free(&l);
}
-static void test_softmax_lastdim(void) {
+static void test_softmax_lastdim(ut_dev dev) {
// last-dim softmax: rows sum to 1
- ut_tensor* t = ut_from_data(2, (int[]){2, 3}, (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f}, UT_CPU);
+ ut_tensor* t = ut_from_data(2, (int[]){2, 3}, (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f}, dev);
ut_tensor* s = ut_softmax(t, 1);
ut_sync_cpu(s);
// row 0: exp(1,2,3)/sum → [0.0900, 0.2447, 0.6652]
@@ -207,9 +207,9 @@ static void test_softmax_lastdim(void) {
ut_free(s);
}
-static void test_softmax_firstdim(void) {
+static void test_softmax_firstdim(ut_dev dev) {
// first-dim softmax: columns sum to 1
- ut_tensor* t = ut_from_data(2, (int[]){2, 3}, (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f}, UT_CPU);
+ ut_tensor* t = ut_from_data(2, (int[]){2, 3}, (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f}, dev);
ut_tensor* s = ut_softmax(t, 0);
ut_sync_cpu(s);
// dim=0 shrinks [2,3]→[2,3]: softmax each column pair, each col sums to 1
@@ -219,6 +219,50 @@ static void test_softmax_firstdim(void) {
ut_free(s);
}
+static void test_layernorm_forward(ut_dev dev) {
+ // d=2, x=[2,2]: [[1,2],[3,4]], weight=[1,1], bias=[0,0]
+ ut_layernorm ln = ut_layernorm_alloc(2, dev);
+ ut_tensor* x = ut_from_data(2, (int[]){2, 2}, (float[]){1.f, 2.f, 3.f, 4.f}, dev);
+ ut_tensor* out = ut_layernorm_forward(&ln, x, NULL);
+ ut_sync_cpu(out);
+ // each row normalised to mean≈0, std≈1
+ assert_data(out, ((float[]){-1.f, 1.f, -1.f, 1.f}), 0.01f);
+ ut_free(out);
+ ut_free(x);
+ ut_layernorm_free(&ln);
+}
+
+static void test_layernorm_backward(ut_dev dev) {
+ // d=3, x=[1,3]: [[1,2,3]], go=[1,0,0]
+ ut_layernorm ln = ut_layernorm_alloc(3, dev);
+ ut_tensor* x = ut_from_data(2, (int[]){1, 3}, (float[]){1.f, 2.f, 3.f}, dev);
+ ut_layernorm_cache c;
+ ut_tensor* out = ut_layernorm_forward(&ln, x, &c);
+ ut_sync_cpu(out);
+
+ ut_tensor* go = ut_from_data(2, (int[]){1, 3}, (float[]){1.f, 0.f, 0.f}, dev);
+ ut_tensor* dW = ut_alloc(1, (int[]){3}, dev);
+ ut_tensor* db = ut_alloc(1, (int[]){3}, dev);
+ memset(dW->data, 0, 12);
+ memset(db->data, 0, 12);
+
+ ut_tensor* dx = ut_layernorm_backward(&ln, &c, go, dW, db);
+ ut_sync_cpu(dx);
+ // per-row gradient sums to 0
+ assert_eq(dx->data[0] + dx->data[1] + dx->data[2], 0.f, 1e-4f);
+ // known values: dx ≈ [0.204, -0.408, 0.204]
+ assert_data(dx, ((float[]){0.2041f, -0.4083f, 0.2041f}), 1e-3f);
+
+ ut_free(dx);
+ ut_free(dW);
+ ut_free(db);
+ ut_free(go);
+ ut_layernorm_cache_free(&c);
+ ut_free(out);
+ ut_free(x);
+ ut_layernorm_free(&ln);
+}
+
static void test_sgd_momentum(void) {
ut_tensor* p = ut_alloc(1, (int[]){2}, UT_CPU);
ut_tensor* params[1] = {p};
@@ -246,17 +290,30 @@ int main() {
test_shape();
test_lifetime();
test_reshape();
+
test_elementwise(UT_CPU);
test_elementwise(UT_METAL);
+
test_matmul_2d(UT_CPU);
test_matmul_2d(UT_METAL);
test_matmul_3d(UT_CPU);
test_matmul_3d(UT_METAL);
+
test_linear_forward();
test_linear_backward(UT_CPU);
test_linear_backward(UT_METAL);
- test_softmax_lastdim();
- test_softmax_firstdim();
+
+ test_softmax_lastdim(UT_CPU);
+ test_softmax_lastdim(UT_METAL);
+
+ test_softmax_firstdim(UT_CPU);
+ test_softmax_firstdim(UT_METAL);
+
+ test_layernorm_forward(UT_CPU);
+ test_layernorm_forward(UT_METAL);
+ test_layernorm_backward(UT_CPU);
+ test_layernorm_backward(UT_METAL);
+
test_sgd_momentum();
return 0;
}
diff --git a/utensil.h b/utensil.h
index 7ebe5b2..bd3c05b 100644
--- a/utensil.h
+++ b/utensil.h
@@ -20,7 +20,7 @@ typedef struct ut_shape {
typedef struct ut_tensor {
ut_shape shape; // shape of the tensor
float* data; // pointer to the data buffer
- void* gpu_buf; /* MTLBuffer */
+ void* gpu_buf; // MTLBuffer
int rc; // reference count for memory management
struct ut_tensor* owner; // if this is a view, points to the owner tensor
ut_dev dev; // device where the tensor is allocated
@@ -35,6 +35,19 @@ typedef struct ut_linear {
bool has_bias;
} ut_linear;
+typedef struct ut_layernorm {
+ ut_tensor* weight; // [d] gain
+ ut_tensor* bias; // [d] shift
+ int d; // last dimension
+ float eps;
+} ut_layernorm;
+
+typedef struct ut_layernorm_cache {
+ ut_tensor* xnorm; // normalised x before affine
+ ut_tensor* mean; // per-row mean
+ ut_tensor* rstd; // per-row reciprocal std-dev
+} ut_layernorm_cache;
+
typedef struct ut_sgd {
ut_tensor** params; // pointers to model parameters (not owned)
ut_tensor** grads; // gradient accumulators (owned)
@@ -163,6 +176,36 @@ static const char* _mtl_src =
"for(int s=tgs/2;s>0;s>>=1){if(lid<s)sh[lid]+=sh[lid+s];"
"barrier(CLK_LOCAL_MEM_FENCE);}float gs_=sh[0];"
"for(int i=lid;i<C;i+=tgs)orow[i]/=gs_;}\n"
+ // layernorm forward: x[rows, d] -> out[rows, d]
+ // also writes xnorm[rows,d], rstd[rows] for backward.
+ // one threadgroup per row. p = {rows, d}.
+ "kernel void ln_fwd("
+ "device const float*x,device const float*w,device const float*b,"
+ "device float*out,device float*xn,device float*rstd,"
+ "constant int*p,"
+ "uint gid[[threadgroup_position_in_grid]],"
+ "uint lid[[thread_position_in_threadgroup]],"
+ "uint tgs[[threads_per_threadgroup]]){"
+ "int d=p[1];"
+ "device const float*row=x+gid*d;"
+ "device float*orow=out+gid*d;device float*xnrow=xn+gid*d;"
+ "threadgroup float sh[1024];"
+ // layernorm backward: given grad_out[rows,d], xnorm[rows,d], rstd[rows],
+ // weight[d] → dx[rows,d]. dW[d] and db[d] accumulated separately on CPU
+ // (only dx needs to be GPU-fast for the training loop hot path).
+ // one threadgroup per row. p = {rows, d}.
+ "kernel void ln_bwd("
+ "device const float*go,device const float*xn,device const float*rs_,"
+ "device const float*w,device float*dx,"
+ "constant int*p,"
+ "uint gid[[threadgroup_position_in_grid]],"
+ "uint lid[[thread_position_in_threadgroup]],"
+ "uint tgs[[threads_per_threadgroup]]){"
+ "int d=p[1];"
+ "device const float*gorow=go+gid*d;device const float*xnrow=xn+gid*d;"
+ "device float*dxrow=dx+gid*d;"
+ "float rs=rs_[gid];"
+ "threadgroup float sh[1024];"
"";
#define _MTL_MAX_PL 32 // maximum number of pipeline states
@@ -170,7 +213,7 @@ typedef struct {
void *device, *queue, *library, *pl[_MTL_MAX_PL];
const char* pl_name[_MTL_MAX_PL];
int pl_count;
- void* pending_cmd; /* open command buffer for batching; NULL if none */
+ void* pending_cmd; // open command buffer for batching; NULL if none
} _mtl_ctx_t;
typedef struct {
@@ -658,9 +701,9 @@ ut_tensor* ut_scale(ut_tensor* a, float s) {
return out;
}
-// ========================================================
+// =========================================================
// Matmul
-// ========================================================
+// =========================================================
static void gemm(const float* a, const float* b, float* c, int m, int n, int k, bool ta, bool tb) {
#ifdef __APPLE__
cblas_sgemm(CblasRowMajor, ta ? CblasTrans : CblasNoTrans, tb ? CblasTrans : CblasNoTrans, m, n,
@@ -755,9 +798,9 @@ static ut_tensor* ut_matmul_t(ut_tensor* a, ut_tensor* b, bool ta, bool tb) {
return c;
}
-// ========================================================
+// =========================================================
// Linear layer
-// ======================================================
+// =========================================================
ut_linear ut_linear_alloc(int in, int out, bool bias, ut_dev dev) {
ut_linear l = {.nin = in, .nout = out, .has_bias = bias};
l.weight = ut_randn(2, (int[]){in, out}, 0.f, sqrtf(2.f / (float)in), dev);
@@ -794,14 +837,14 @@ ut_tensor* ut_linear_backward(ut_linear* l, ut_tensor* x, ut_tensor* grad_out, u
ut_tensor* db) {
int B = x->shape.shape[0];
- /* dW += x^T @ grad_out — ta=true: x treated as [in, B] */
+ // dW += x^T @ grad_out — ta=true: x treated as [in, B]
ut_tensor* dWb = ut_matmul_t(x, grad_out, true, false);
ut_sync_cpu(dWb);
ut_sync_cpu(dW);
for (int i = 0; i < dW->shape.nelem; i++) dW->data[i] += dWb->data[i];
ut_free(dWb);
- /* db += sum(grad_out, axis=0) */
+ // db += sum(grad_out, axis=0)
if (l->has_bias && db) {
ut_sync_cpu(grad_out);
ut_sync_cpu(db);
@@ -809,7 +852,7 @@ ut_tensor* ut_linear_backward(ut_linear* l, ut_tensor* x, ut_tensor* grad_out, u
for (int j = 0; j < l->nout; j++) db->data[j] += grad_out->data[b * l->nout + j];
}
- /* dx = grad_out @ W^T — tb=true: weight treated as [in, out] → [out, in]^T */
+ // dx = grad_out @ W^T — tb=true: weight treated as [in, out] → [out, in]^T
return ut_matmul_t(grad_out, l->weight, false, true);
}
@@ -836,10 +879,177 @@ ut_tensor* ut_relu_backward(ut_tensor* grad_out, ut_tensor* fwd_input) {
}
return gi;
}
+// =========================================================
+// LayerNorm
+// =========================================================
+ut_layernorm ut_layernorm_alloc(int d, ut_dev dev) {
+ ut_layernorm l = {0};
+ l.d = d;
+ l.eps = 1e-5f;
+ int wd[1] = {d};
+ l.weight = ut_alloc(1, wd, UT_CPU);
+ l.bias = ut_alloc(1, wd, UT_CPU);
+ for (int i = 0; i < d; i++) l.weight->data[i] = 1.f;
+ if (dev == UT_METAL) {
+ ut_to_device(l.weight, UT_METAL);
+ ut_to_device(l.bias, UT_METAL);
+ }
+ return l;
+}
+
+// x: [..., d] — normalises the last dimension
+ut_tensor* ut_layernorm_forward(ut_layernorm* l, ut_tensor* x, ut_layernorm_cache* cache) {
+ int d = l->d;
+ int rows = x->shape.nelem / d;
-// ========================================================
+ // GPU fast-path
+ _mtl_ctx_t* _mc_ln = (_mtl_ctx_t*)ut_metal_ctx();
+ if (_mc_ln) {
+ ut_to_device(x, UT_METAL);
+ ut_to_device(l->weight, UT_METAL);
+ ut_to_device(l->bias, UT_METAL);
+
+ ut_tensor* out = ut_alloc(x->shape.ndim, x->shape.shape, UT_METAL);
+ int rd[1] = {rows};
+ ut_tensor* xnorm_t = ut_alloc(x->shape.ndim, x->shape.shape, UT_METAL);
+ ut_tensor* rstd_t = ut_alloc(1, rd, UT_METAL);
+
+ int tgsize = d < 64 ? 32 : (d < 256 ? 64 : (d < 512 ? 128 : 256));
+ if (tgsize > 1024) tgsize = 1024;
+ int p2[2] = {rows, d};
+ void* bufs[7] = {x->gpu_buf, l->weight->gpu_buf, l->bias->gpu_buf,
+ out->gpu_buf, xnorm_t->gpu_buf, rstd_t->gpu_buf};
+ _mtl_dispatch_rows(_mc_ln, "ln_fwd", bufs, 6, p2, (int)sizeof(p2), rows, tgsize);
+ out->dirty_cpu = true;
+ xnorm_t->dirty_cpu = true;
+ rstd_t->dirty_cpu = true;
+
+ if (cache) {
+ cache->xnorm = xnorm_t;
+ cache->rstd = rstd_t;
+ // allocate a dummy mean (not used by GPU backward but kept in struct)
+ cache->mean = ut_alloc(1, rd, UT_CPU);
+ } else {
+ ut_free(xnorm_t);
+ ut_free(rstd_t);
+ }
+ return out;
+ }
+
+ // CPU fallback
+ ut_sync_cpu(x);
+ ut_sync_cpu(l->weight);
+ ut_sync_cpu(l->bias);
+ ut_tensor* out = ut_alloc(x->shape.ndim, x->shape.shape, UT_CPU);
+ int rd[1] = {rows};
+ ut_tensor* mean_t = ut_alloc(1, rd, UT_CPU);
+ ut_tensor* rstd_t = ut_alloc(1, rd, UT_CPU);
+ ut_tensor* xnorm_t = ut_alloc(x->shape.ndim, x->shape.shape, UT_CPU);
+ for (int r = 0; r < rows; r++) {
+ const float* row = x->data + r * d;
+ float mu = 0;
+ for (int i = 0; i < d; i++) mu += row[i];
+ mu /= d;
+ float var = 0;
+ for (int i = 0; i < d; i++) {
+ float v = row[i] - mu;
+ var += v * v;
+ }
+ var /= d;
+ float rs = 1.f / sqrtf(var + l->eps);
+ mean_t->data[r] = mu;
+ rstd_t->data[r] = rs;
+ for (int i = 0; i < d; i++) {
+ float xn = (row[i] - mu) * rs;
+ xnorm_t->data[r * d + i] = xn;
+ out->data[r * d + i] = l->weight->data[i] * xn + l->bias->data[i];
+ }
+ }
+ if (cache) {
+ cache->xnorm = xnorm_t;
+ cache->mean = mean_t;
+ cache->rstd = rstd_t;
+ } else {
+ ut_free(xnorm_t);
+ ut_free(mean_t);
+ ut_free(rstd_t);
+ }
+ return out;
+}
+
+ut_tensor* ut_layernorm_backward(ut_layernorm* l, ut_layernorm_cache* cache, ut_tensor* grad_out,
+ ut_tensor* dW, ut_tensor* db) {
+ int d = l->d, rows = grad_out->shape.nelem / d;
+
+ // dW and db are always accumulated on CPU (small [d] vectors)
+ ut_sync_cpu(grad_out);
+ ut_sync_cpu(cache->xnorm);
+ ut_sync_cpu(dW);
+ ut_sync_cpu(db);
+ for (int r = 0; r < rows; r++) {
+ const float* go = grad_out->data + r * d;
+ const float* xn = cache->xnorm->data + r * d;
+ for (int i = 0; i < d; i++) {
+ dW->data[i] += go[i] * xn[i];
+ db->data[i] += go[i];
+ }
+ }
+ dW->dirty_gpu = true;
+ db->dirty_gpu = true;
+
+ // GPU fast-path for dx
+ _mtl_ctx_t* _mc_lnb = (_mtl_ctx_t*)ut_metal_ctx();
+ if (_mc_lnb) {
+ ut_to_device(grad_out, UT_METAL);
+ ut_to_device(cache->xnorm, UT_METAL);
+ ut_to_device(cache->rstd, UT_METAL);
+ ut_to_device(l->weight, UT_METAL);
+ ut_tensor* dx = ut_alloc(grad_out->shape.ndim, grad_out->shape.shape, UT_METAL);
+ int tgsize = d < 64 ? 32 : (d < 256 ? 64 : (d < 512 ? 128 : 256));
+ if (tgsize > 1024) tgsize = 1024;
+ int p2[2] = {rows, d};
+ void* bufs[5] = {grad_out->gpu_buf, cache->xnorm->gpu_buf, cache->rstd->gpu_buf,
+ l->weight->gpu_buf, dx->gpu_buf};
+ _mtl_dispatch_rows(_mc_lnb, "ln_bwd", bufs, 5, p2, (int)sizeof(p2), rows, tgsize);
+ dx->dirty_cpu = true;
+ return dx;
+ }
+
+ // CPU fallback
+ ut_sync_cpu(cache->rstd);
+ ut_sync_cpu(l->weight);
+ ut_tensor* dx = ut_alloc(grad_out->shape.ndim, grad_out->shape.shape, UT_CPU);
+ for (int r = 0; r < rows; r++) {
+ const float* go = grad_out->data + r * d;
+ const float* xn = cache->xnorm->data + r * d;
+ float rs = cache->rstd->data[r];
+ float sum_go_xn = 0, sum_go = 0;
+ for (int i = 0; i < d; i++) {
+ sum_go_xn += go[i] * l->weight->data[i] * xn[i];
+ sum_go += go[i] * l->weight->data[i];
+ }
+ float* dxr = dx->data + r * d;
+ for (int i = 0; i < d; i++)
+ dxr[i] = rs * (l->weight->data[i] * go[i] - (sum_go + xn[i] * sum_go_xn) / (float)d);
+ }
+ return dx;
+}
+
+void ut_layernorm_cache_free(ut_layernorm_cache* c) {
+ ut_free(c->xnorm);
+ ut_free(c->mean);
+ ut_free(c->rstd);
+ c->xnorm = c->mean = c->rstd = NULL;
+}
+void ut_layernorm_free(ut_layernorm* l) {
+ ut_free(l->weight);
+ ut_free(l->bias);
+ l->weight = l->bias = NULL;
+}
+
+// =========================================================
// Softmax
-// ========================================================
+// =========================================================
ut_tensor* ut_softmax(ut_tensor* t, int dim) {
int nd = t->shape.ndim, k = t->shape.shape[dim];
int outer = 1;
@@ -882,9 +1092,9 @@ ut_tensor* ut_softmax(ut_tensor* t, int dim) {
return out;
}
-// ========================================================
+// =========================================================
// Loss
-// ========================================================
+// =========================================================
float ut_cross_entropy(ut_tensor* logits, const int* labels, ut_tensor* grad_in) {
ut_sync_cpu(logits);
int B = logits->shape.shape[0], C = logits->shape.shape[1];
@@ -908,9 +1118,9 @@ float ut_cross_entropy(ut_tensor* logits, const int* labels, ut_tensor* grad_in)
return loss / (float)B;
}
-// ========================================================
+// =========================================================
// Optimisers
-// ========================================================
+// =========================================================
ut_sgd ut_sgd_alloc(ut_tensor** params, int n, float lr, float mom) {
ut_sgd o = {0};
o.nparams = n;