← Commits · a912dcc7
a912dcc7011a47945acf67fac765af8e7c72353f
diff --git a/test.c b/test.c
index ef06357..81de536 100644
--- a/test.c
+++ b/test.c
@@ -153,18 +153,16 @@ static void test_linear_forward(void) {
// W[2,3], bias[3], x[4,2]
// x row 0 = [10,20] → [10*1+20*4, 10*2+20*5, 10*3+20*6] + [7,8,9]
ut_linear l = ut_linear_alloc(2, 3, true, UT_CPU);
- l.weight = ut_from_data(2, (int[]){2, 3},
- (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f}, UT_CPU);
+ l.weight = ut_from_data(2, (int[]){2, 3}, (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f}, UT_CPU);
l.bias = ut_from_data(1, (int[]){3}, (float[]){7.f, 8.f, 9.f}, UT_CPU);
- ut_tensor* x =
- ut_from_data(2, (int[]){4, 2},
- (float[]){10.f, 20.f, 30.f, 40.f, 50.f, 60.f, 70.f, 80.f}, UT_CPU);
+ ut_tensor* x = ut_from_data(2, (int[]){4, 2},
+ (float[]){10.f, 20.f, 30.f, 40.f, 50.f, 60.f, 70.f, 80.f}, UT_CPU);
ut_tensor* out = ut_linear_forward(&l, x);
ut_sync_cpu(out);
assert_data(out,
- ((float[]){97.f, 128.f, 159.f, 197.f, 268.f, 339.f, 297.f, 408.f, 519.f, 397.f,
- 548.f, 699.f}),
+ ((float[]){97.f, 128.f, 159.f, 197.f, 268.f, 339.f, 297.f, 408.f, 519.f, 397.f, 548.f,
+ 699.f}),
1e-4f);
ut_free(out);
ut_free(x);
@@ -177,11 +175,8 @@ static void test_linear_backward(ut_dev dev) {
l.weight = ut_from_data(2, (int[]){2, 1}, (float[]){3.f, 4.f}, dev);
l.bias = ut_alloc(1, (int[]){1}, UT_CPU);
- ut_tensor* x =
- ut_from_data(2, (int[]){3, 2},
- (float[]){1.f, 2.f, 5.f, 6.f, 7.f, 8.f}, dev);
- ut_tensor* go =
- ut_from_data(2, (int[]){3, 1}, (float[]){10.f, 20.f, 30.f}, dev);
+ ut_tensor* x = ut_from_data(2, (int[]){3, 2}, (float[]){1.f, 2.f, 5.f, 6.f, 7.f, 8.f}, dev);
+ ut_tensor* go = ut_from_data(2, (int[]){3, 1}, (float[]){10.f, 20.f, 30.f}, dev);
ut_tensor* dW = ut_alloc(2, (int[]){2, 1}, dev);
ut_tensor* db = ut_alloc(1, (int[]){1}, dev);
memset(dW->data, 0, (size_t)dW->shape.nelem * sizeof(float));
@@ -199,6 +194,31 @@ static void test_linear_backward(ut_dev dev) {
ut_linear_free(&l);
}
+static void test_softmax_lastdim(void) {
+ // 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* s = ut_softmax(t, 1);
+ ut_sync_cpu(s);
+ // row 0: exp(1,2,3)/sum → [0.0900, 0.2447, 0.6652]
+ // row 1: same offsets → [0.0900, 0.2447, 0.6652]
+ assert_data(s, ((float[]){0.090031f, 0.244728f, 0.665241f, 0.090031f, 0.244728f, 0.665241f}),
+ 1e-4f);
+ ut_free(t);
+ ut_free(s);
+}
+
+static void test_softmax_firstdim(void) {
+ // 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* 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
+ assert_data(s, ((float[]){0.047426f, 0.047426f, 0.047426f, 0.952574f, 0.952574f, 0.952574f}),
+ 1e-4f);
+ ut_free(t);
+ ut_free(s);
+}
+
static void test_sgd_momentum(void) {
ut_tensor* p = ut_alloc(1, (int[]){2}, UT_CPU);
ut_tensor* params[1] = {p};
@@ -235,6 +255,8 @@ int main() {
test_linear_forward();
test_linear_backward(UT_CPU);
test_linear_backward(UT_METAL);
+ test_softmax_lastdim();
+ test_softmax_firstdim();
test_sgd_momentum();
return 0;
}
diff --git a/utensil.h b/utensil.h
index 170b531..7ebe5b2 100644
--- a/utensil.h
+++ b/utensil.h
@@ -144,6 +144,25 @@ static const char* _mtl_src =
"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"
+ // row_softmax: one workgroup per row; local sh[] passed as last arg
+ // p={outer,C} global=outer*tgsize local=tgsize
+ "__kernel void row_softmax("
+ "__global const float*x,__global float*o,__global const int*p,"
+ "__local float* sh){"
+ "int gid=(int)get_group_id(0);int lid=(int)get_local_id(0);"
+ "int tgs=(int)get_local_size(0);int C=p[1];"
+ "__global const float*row=x+gid*C;__global float*orow=o+gid*C;"
+ "float mx=-1e38f;"
+ "for(int i=lid;i<C;i+=tgs){float v=row[i];if(v>mx)mx=v;}"
+ "sh[lid]=mx;barrier(CLK_LOCAL_MEM_FENCE);"
+ "for(int s=tgs/2;s>0;s>>=1){if(lid<s&&sh[lid+s]>sh[lid])sh[lid]=sh[lid+s];"
+ "barrier(CLK_LOCAL_MEM_FENCE);}float gmx=sh[0];"
+ "float loc=0.f;"
+ "for(int i=lid;i<C;i+=tgs){float e=exp(row[i]-gmx);orow[i]=e;loc+=e;}"
+ "sh[lid]=loc;barrier(CLK_LOCAL_MEM_FENCE);"
+ "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"
"";
#define _MTL_MAX_PL 32 // maximum number of pipeline states
@@ -221,7 +240,7 @@ static void* _mtl_get_pl(_mtl_ctx_t* c, const char* name) {
return ps;
}
-/* Get or create the batched command buffer. */
+// get or create the batched command buffer
static void* _mtl_get_cmd(_mtl_ctx_t* c) {
if (!c->pending_cmd) {
void* cb = _m0(c->queue, "commandBuffer");
@@ -254,6 +273,30 @@ static void _mtl_dispatch(_mtl_ctx_t* c, const char* kern, void** bufs, int nbuf
_v0(enc, "endEncoding");
}
+// dispatch a kernel that processes rows of a 2D tensor, with one threadgroup per row
+static void _mtl_dispatch_rows(_mtl_ctx_t* c, const char* kern, void** bufs, int nbufs,
+ const void* bytes, int blen, int rows, int tgsize) {
+ void* cb = _mtl_get_cmd(c);
+ void* enc = _m0(cb, "computeCommandEncoder");
+ void* ps = _mtl_get_pl(c, kern);
+ if (!ps) {
+ _v0(enc, "endEncoding");
+ return;
+ }
+ _v1(enc, "setComputePipelineState:", ps);
+ for (int i = 0; i < nbufs; i++)
+ ((void (*)(id, SEL, id, long, long))objc_msgSend)(
+ (id)enc, sel_getUid("setBuffer:offset:atIndex:"), (id)bufs[i], 0L, (long)i);
+ if (bytes)
+ ((void (*)(id, SEL, const void*, long, long))objc_msgSend)(
+ (id)enc, sel_getUid("setBytes:length:atIndex:"), bytes, (long)blen, (long)nbufs);
+ _msize_t grp = {(unsigned long)rows, 1, 1};
+ _msize_t thr = {(unsigned long)tgsize, 1, 1};
+ ((void (*)(id, SEL, _msize_t, _msize_t))objc_msgSend)(
+ (id)enc, sel_getUid("dispatchThreadgroups:threadsPerThreadgroup:"), grp, thr);
+ _v0(enc, "endEncoding");
+}
+
static void _mtl_matmul(_mtl_ctx_t* ctx, void* a, void* b, void* res, int m, int n, int k, bool ta,
bool tb) {
int ra = ta ? k : m, ca = ta ? m : k;
@@ -794,6 +837,51 @@ ut_tensor* ut_relu_backward(ut_tensor* grad_out, ut_tensor* fwd_input) {
return gi;
}
+// ========================================================
+// Softmax
+// ========================================================
+ut_tensor* ut_softmax(ut_tensor* t, int dim) {
+ int nd = t->shape.ndim, k = t->shape.shape[dim];
+ int outer = 1;
+ for (int i = 0; i < dim; i++) outer *= t->shape.shape[i];
+ int inner = 1;
+ for (int i = dim + 1; i < nd; i++) inner *= t->shape.shape[i];
+
+ // GPU fast-path: only when dim is last (inner==1)
+ _mtl_ctx_t* _mc_sm = (_mtl_ctx_t*)ut_metal_ctx();
+ if (_mc_sm && inner == 1) {
+ ut_to_device(t, UT_METAL);
+ ut_tensor* out = ut_alloc(nd, t->shape.shape, UT_METAL);
+ int tgsize = k < 64 ? 32 : (k < 256 ? 64 : (k < 512 ? 128 : 256));
+ if (tgsize > 1024) tgsize = 1024;
+ int p2[2] = {outer, k};
+ void* bufs[2] = {t->gpu_buf, out->gpu_buf};
+ _mtl_dispatch_rows(_mc_sm, "row_softmax", bufs, 2, p2, (int)sizeof(p2), outer, tgsize);
+ out->dirty_cpu = true;
+ return out;
+ }
+
+ // CPU fallback (also handles non-last-dim softmax)
+ ut_sync_cpu(t);
+ ut_tensor* out = ut_alloc(nd, t->shape.shape, UT_CPU);
+ for (int o = 0; o < outer; o++)
+ for (int in = 0; in < inner; in++) {
+ float mx = -FLT_MAX;
+ for (int d = 0; d < k; d++) {
+ float v = t->data[o * k * inner + d * inner + in];
+ if (v > mx) mx = v;
+ }
+ float sum = 0;
+ for (int d = 0; d < k; d++) {
+ float e = expf(t->data[o * k * inner + d * inner + in] - mx);
+ out->data[o * k * inner + d * inner + in] = e;
+ sum += e;
+ }
+ for (int d = 0; d < k; d++) out->data[o * k * inner + d * inner + in] /= sum;
+ }
+ return out;
+}
+
// ========================================================
// Loss
// ========================================================