← Commits · 399f0e0e
399f0e0ec3b14d897a25ecb67eb664085d518afe
diff --git a/utensil.h b/utensil.h
index d463e0e..57bb2f3 100644
--- a/utensil.h
+++ b/utensil.h
@@ -28,6 +28,21 @@ typedef struct ut_tensor {
bool dirty_gpu; // CPU is newer than GPU buffer
} ut_tensor;
+typedef struct ut_linear {
+ ut_tensor* weight; // [in, out]
+ ut_tensor* bias; // [out]
+ int nin, nout;
+ bool has_bias;
+} ut_linear;
+
+typedef struct ut_sgd {
+ ut_tensor** params; // pointers to model parameters (not owned)
+ ut_tensor** grads; // gradient accumulators (owned)
+ ut_tensor** velocity; // momentum buffers (owned, NULL if mom==0)
+ int nparams;
+ float lr, momentum;
+} ut_sgd;
+
// =========================================================
// Metal context management
// =========================================================
@@ -106,6 +121,29 @@ static const char* _mtl_src =
"float s=0;int ao=bat*M*K+m*K,bo=bat*K*N+n;"
"for(int k=0;k<K;k++)s+=a[ao+k]*b_[bo+k*N];"
"c[idx]=s;}\n"
+ // bias_add: p={N,C,HW}
+ "__kernel void bias_add(__global float*out,__global const float*bias,"
+ "__global const int*p){"
+ "int idx=(int)get_global_id(0);int tot=p[0]*p[1]*p[2];"
+ "if(idx>=tot)return;int c=(idx/p[2])%p[1];out[idx]+=bias[c];}\n"
+ // relu_bwd
+ "__kernel void relu_bwd(__global const float*go,__global const float*fwd,"
+ "__global float*gi,__global const int*p){"
+ "int idx=(int)get_global_id(0);if(idx<p[0])gi[idx]=fwd[idx]>0.f?go[idx]:0.f;}\n"
+ // transpose dims 0<->1: p={A,B,inner}
+ "__kernel void transpose_01(__global const float*in,__global float*out,"
+ "__global const int*p){"
+ "int idx=(int)get_global_id(0);"
+ "int A=p[0],B=p[1],inner=p[2];if(idx>=A*B*inner)return;"
+ "int k=idx%inner,t=idx/inner,a=t%A,b=t/A;"
+ "out[idx]=in[(a*B+b)*inner+k];}\n"
+ // transpose dims 1<->2: p={A,B,C,D}
+ "__kernel void transpose_12(__global const float*i,__global float*o,"
+ "__global const int*p){"
+ "int idx=(int)get_global_id(0);"
+ "int A=p[0],B=p[1],C=p[2],D=p[3];int n=A*B*C*D;if(idx>=n)return;"
+ "int d_=idx%D,t=idx/D,c=t%C,t2=t/C,b=t2%B,a=t2/B;"
+ "o[a*C*B*D+c*B*D+b*D+d_]=i[a*B*C*D+b*C*D+c*D+d_];}\n"
"";
#define _MTL_MAX_PL 32 // maximum number of pipeline states
@@ -423,6 +461,66 @@ ut_tensor* ut_view(ut_tensor* t, int ndim, const int* dim) {
return v;
}
+void ut_reshape(ut_tensor* t, int ndim, const int* dims) {
+ ut_shape ns = ut_shape_new(ndim, dims);
+ assert(ns.nelem == t->shape.nelem);
+ t->shape = ns;
+}
+
+ut_tensor* ut_transpose(ut_tensor* t, int dim0, int dim1) {
+ int new_shape[UT_MAX_DIMS];
+ memcpy(new_shape, t->shape.shape, (size_t)t->shape.ndim * sizeof(int));
+ int tmp = new_shape[dim0];
+ new_shape[dim0] = new_shape[dim1];
+ new_shape[dim1] = tmp;
+
+ // GPU fast-path for 0<->1 dims
+ _mtl_ctx_t* _mc_t = (_mtl_ctx_t*)ut_metal_ctx();
+ if (_mc_t && (dim0 == 0 || dim1 == 0) && (dim0 == 1 || dim1 == 1)) {
+ ut_to_device(t, UT_METAL);
+ ut_tensor* out_g = ut_alloc(t->shape.ndim, new_shape, UT_METAL);
+ int A = t->shape.shape[0], B = t->shape.shape[1];
+ int inner = t->shape.nelem / (A * B);
+ int _tp[3] = {A, B, inner};
+ void* _tb[2] = {t->gpu_buf, out_g->gpu_buf};
+ _mtl_dispatch(_mc_t, "transpose_01", _tb, 2, _tp, (int)sizeof(_tp), t->shape.nelem);
+ out_g->dirty_cpu = true;
+ return out_g;
+ }
+
+ // GPU fast-path for 1<->2 dims
+ if (_mc_t && t->shape.ndim >= 3 && ((dim0 == 1 && dim1 == 2) || (dim0 == 2 && dim1 == 1))) {
+ ut_to_device(t, UT_METAL);
+ ut_tensor* out_g = ut_alloc(t->shape.ndim, new_shape, UT_METAL);
+ int A = t->shape.shape[0], B = t->shape.shape[1], C = t->shape.shape[2],
+ D = t->shape.ndim >= 4 ? t->shape.shape[3] : 1;
+ int _tp[4] = {A, B, C, D};
+ void* _tb[2] = {t->gpu_buf, out_g->gpu_buf};
+ _mtl_dispatch(_mc_t, "transpose_12", _tb, 2, _tp, (int)sizeof(_tp), t->shape.nelem);
+ out_g->dirty_cpu = true;
+ return out_g;
+ }
+
+ ut_sync_cpu(t);
+ ut_tensor* out = ut_alloc(t->shape.ndim, new_shape, UT_CPU);
+ int nd = t->shape.ndim;
+ for (int i = 0; i < t->shape.nelem; i++) {
+ int idx_s[UT_MAX_DIMS] = {0};
+ int rem = i;
+ for (int d = nd - 1; d >= 0; d--) {
+ idx_s[d] = rem % t->shape.shape[d];
+ rem /= t->shape.shape[d];
+ }
+ int idx_d[UT_MAX_DIMS];
+ memcpy(idx_d, idx_s, sizeof(idx_d));
+ int sw = idx_d[dim0];
+ idx_d[dim0] = idx_d[dim1];
+ idx_d[dim1] = sw;
+ out->data[ut_index(out->shape, idx_d)] = t->data[i];
+ }
+ return out;
+}
+
// =========================================================
// Elementwise operations
// =========================================================
@@ -587,4 +685,195 @@ ut_tensor* ut_matmul(ut_tensor* a, ut_tensor* b) {
return NULL;
}
+static ut_tensor* ut_matmul_t(ut_tensor* a, ut_tensor* b, bool ta, bool tb) {
+ int ra = a->shape.shape[0], ca = a->shape.shape[1];
+ int rb = b->shape.shape[0], cb = b->shape.shape[1];
+ int m = ta ? ca : ra;
+ int k = ta ? ra : ca;
+ int n = tb ? rb : cb;
+ int kb = tb ? cb : rb;
+ if (k != kb) return NULL;
+ ut_dev dev = (a->dev == UT_METAL || b->dev == UT_METAL) ? UT_METAL : UT_CPU;
+ ut_sync_cpu(a);
+ ut_sync_cpu(b);
+ int cd[2] = {m, n};
+ ut_tensor* c = ut_alloc(2, cd, dev);
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ if (dev == UT_METAL && mc) {
+ ut_sync_gpu(a);
+ ut_sync_gpu(b);
+ _mtl_matmul(mc, a->gpu_buf, b->gpu_buf, c->gpu_buf, m, n, k, ta, tb);
+ c->dirty_cpu = true;
+ } else {
+ ut_sync_cpu(a);
+ ut_sync_cpu(b);
+ gemm(a->data, b->data, c->data, m, n, k, ta, 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);
+ if (bias) l.bias = ut_alloc(1, (int[]){out}, dev);
+ return l;
+}
+
+// out = x @ W + bias
+// x:[B,in], W:[in,out] → out:[B,out]
+ut_tensor* ut_linear_forward(ut_linear* l, ut_tensor* x) {
+ ut_tensor* out = ut_matmul(x, l->weight);
+ if (l->has_bias) {
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ if (out->dev == UT_METAL && mc) {
+ ut_sync_gpu(out);
+ ut_to_device(l->bias, UT_METAL);
+ int B = out->shape.shape[0];
+ int params[3] = {B, l->nout, 1};
+ void* bufs[2] = {out->gpu_buf, l->bias->gpu_buf};
+ _mtl_dispatch(mc, "bias_add", bufs, 2, params, (int)sizeof(params), B * l->nout);
+ out->dirty_cpu = true;
+ } else {
+ ut_sync_cpu(out);
+ ut_sync_cpu(l->bias);
+ int B = out->shape.shape[0];
+ for (int b = 0; b < B; b++)
+ for (int j = 0; j < l->nout; j++) out->data[b * l->nout + j] += l->bias->data[j];
+ }
+ }
+ return out;
+}
+
+ut_tensor* ut_linear_backward(ut_linear* l, ut_tensor* x, ut_tensor* grad_out, ut_tensor* dW,
+ ut_tensor* db) {
+ int B = x->shape.shape[0];
+
+ /* 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) */
+ if (l->has_bias && db) {
+ ut_sync_cpu(grad_out);
+ ut_sync_cpu(db);
+ for (int b = 0; b < B; b++)
+ 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 */
+ return ut_matmul_t(grad_out, l->weight, false, true);
+}
+
+void ut_linear_free(ut_linear* l) {
+ ut_free(l->weight);
+ if (l->bias) ut_free(l->bias);
+}
+
+ut_tensor* ut_relu_backward(ut_tensor* grad_out, ut_tensor* fwd_input) {
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ ut_tensor* gi = ut_alloc(grad_out->shape.ndim, grad_out->shape.shape, grad_out->dev);
+ if (grad_out->dev == UT_METAL && mc) {
+ ut_sync_gpu(grad_out);
+ ut_to_device(fwd_input, UT_METAL);
+ void* bufs[3] = {grad_out->gpu_buf, fwd_input->gpu_buf, gi->gpu_buf};
+ int n = grad_out->shape.nelem;
+ _mtl_dispatch(mc, "relu_bwd", bufs, 3, &n, sizeof(int), n);
+ gi->dirty_cpu = true;
+ } else {
+ ut_sync_cpu(grad_out);
+ ut_sync_cpu(fwd_input);
+ for (int i = 0; i < grad_out->shape.nelem; i++)
+ gi->data[i] = fwd_input->data[i] > 0 ? grad_out->data[i] : 0.f;
+ }
+ return gi;
+}
+
+// ========================================================
+// 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];
+ if (grad_in) { ut_sync_cpu(grad_in); }
+ float loss = 0;
+ for (int b = 0; b < B; b++) {
+ const float* row = logits->data + b * C;
+ float mx = row[0];
+ for (int c = 1; c < C; c++)
+ if (row[c] > mx) mx = row[c];
+ float sum = 0;
+ for (int c = 0; c < C; c++) sum += expf(row[c] - mx);
+ float log_sum = logf(sum) + mx;
+ loss += log_sum - row[labels[b]];
+ if (grad_in) {
+ float* g = grad_in->data + b * C;
+ for (int c = 0; c < C; c++) g[c] = expf(row[c] - log_sum) / (float)B;
+ g[labels[b]] -= 1.f / (float)B;
+ }
+ }
+ 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;
+ o.lr = lr;
+ o.momentum = mom;
+ o.params = malloc((size_t)n * sizeof(ut_tensor*));
+ o.grads = malloc((size_t)n * sizeof(ut_tensor*));
+ memcpy(o.params, params, (size_t)n * sizeof(ut_tensor*));
+ for (int i = 0; i < n; i++)
+ o.grads[i] = ut_alloc(params[i]->shape.ndim, params[i]->shape.shape, UT_CPU);
+ if (mom > 0) {
+ o.velocity = malloc((size_t)n * sizeof(ut_tensor*));
+ for (int i = 0; i < n; i++)
+ o.velocity[i] = ut_alloc(params[i]->shape.ndim, params[i]->shape.shape, UT_CPU);
+ }
+ return o;
+}
+
+void ut_sgd_zero(ut_sgd* o) {
+ for (int i = 0; i < o->nparams; i++)
+ memset(o->grads[i]->data, 0, (size_t)o->grads[i]->shape.nelem * sizeof(float));
+}
+
+void ut_sgd_step(ut_sgd* o, float clip) {
+ for (int i = 0; i < o->nparams; i++) {
+ ut_tensor *p = o->params[i], *g = o->grads[i];
+ ut_sync_cpu(p);
+ for (int j = 0; j < p->shape.nelem; j++) {
+ float gr = g->data[j];
+ if (gr > clip) gr = clip;
+ if (gr < -clip) gr = -clip;
+ if (o->momentum > 0) {
+ o->velocity[i]->data[j] = o->momentum * o->velocity[i]->data[j] - o->lr * gr;
+ p->data[j] += o->velocity[i]->data[j];
+ } else {
+ p->data[j] -= o->lr * gr;
+ }
+ }
+ p->dirty_gpu = true;
+ memset(g->data, 0, (size_t)g->shape.nelem * sizeof(float));
+ }
+}
+
+void ut_sgd_free(ut_sgd* o) {
+ for (int i = 0; i < o->nparams; i++) ut_free(o->grads[i]);
+ if (o->velocity) {
+ for (int i = 0; i < o->nparams; i++) ut_free(o->velocity[i]);
+ free(o->velocity);
+ }
+ free(o->params);
+ free(o->grads);
+}
+
#endif // UTENSIL_H