← Commits · d754effd
d754effd22b17ae04a84de5e123690272ac8080a
diff --git a/test.c b/test.c
index 5ca4674..911349b 100644
--- a/test.c
+++ b/test.c
@@ -43,8 +43,7 @@ void test_reshape(void) {
ut_tensor* a = ut_randn(2, (int[]){2, 3}, 0.f, 1.f, UT_CPU);
ut_tensor* v = ut_view(a, 1, (int[]){6});
for (int i = 0; i < 6; i++) assert(v->data[i] == a->data[i]);
- ut_free(v);
- ut_free(a);
+ ut_free_all(v, a);
}
void test_elementwise(ut_dev dev) {
@@ -114,8 +113,7 @@ void test_elementwise(ut_dev dev) {
ut_free(c);
}
- ut_free(a);
- ut_free(b);
+ ut_free_all(a, b);
}
static void test_matmul_2d(ut_dev dev) {
@@ -129,9 +127,7 @@ static void test_matmul_2d(ut_dev dev) {
ut_sync_cpu(c);
assert(c->shape.ndim == 2 && c->shape.shape[0] == 2 && c->shape.shape[1] == 2);
assert_data(c, ((float[]){58.f, 64.f, 139.f, 154.f}), 1e-6f);
- ut_free(a);
- ut_free(b);
- ut_free(c);
+ ut_free_all(a, b, c);
}
static void test_matmul_3d(ut_dev dev) {
@@ -144,9 +140,7 @@ static void test_matmul_3d(ut_dev dev) {
ut_tensor* c = ut_matmul(a, b);
ut_sync_cpu(c);
assert_data(c, ((float[]){94.f, 100.f, 229.f, 244.f, 508.f, 532.f, 697.f, 730.f}), 1e-6f);
- ut_free(a);
- ut_free(b);
- ut_free(c);
+ ut_free_all(a, b, c);
}
static void test_linear_forward(void) {
@@ -164,8 +158,7 @@ static void test_linear_forward(void) {
((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);
+ ut_free_all(out, x);
ut_linear_free(&l);
}
@@ -186,11 +179,7 @@ static void test_linear_backward(ut_dev dev) {
ut_sync_cpu(db);
assert_eq(db->data[0], 60.f, 1e-4f);
- ut_free(dx);
- ut_free(dW);
- ut_free(db);
- ut_free(go);
- ut_free(x);
+ ut_free_all(dx, dW, db, go, x);
ut_linear_free(&l);
}
@@ -203,8 +192,7 @@ static void test_softmax_lastdim(ut_dev dev) {
// 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);
+ ut_free_all(t, s);
}
static void test_softmax_firstdim(ut_dev dev) {
@@ -215,8 +203,7 @@ static void test_softmax_firstdim(ut_dev dev) {
// 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);
+ ut_free_all(t, s);
}
static void test_layernorm_forward(ut_dev dev) {
@@ -227,8 +214,7 @@ static void test_layernorm_forward(ut_dev dev) {
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_free_all(out, x);
ut_layernorm_free(&ln);
}
@@ -253,13 +239,9 @@ static void test_layernorm_backward(ut_dev dev) {
// 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_free_all(dx, dW, db, go);
ut_layernorm_cache_free(&c);
- ut_free(out);
- ut_free(x);
+ ut_free_all(out, x);
ut_layernorm_free(&ln);
}
@@ -270,8 +252,7 @@ static void test_im2col(ut_dev dev) {
ut_sync_cpu(col);
assert(col->shape.ndim == 2 && col->shape.shape[0] == 3 && col->shape.shape[1] == 2);
assert_data(col, ((float[]){1.f, 2.f, 2.f, 3.f, 3.f, 4.f}), 1e-6f);
- ut_free(col);
- ut_free(x);
+ ut_free_all(col, x);
}
static void test_col2im(ut_dev dev) {
@@ -280,8 +261,7 @@ static void test_col2im(ut_dev dev) {
ut_tensor* dx = ut_col2im(col, 1, 1, 1, 4, 1, 3, 1, 0);
ut_sync_cpu(dx);
assert_data(dx, ((float[]){1.f, 4.f, 6.f, 4.f}), 1e-6f);
- ut_free(dx);
- ut_free(col);
+ ut_free_all(dx, col);
}
static void test_conv1d_forward(ut_dev dev) {
@@ -298,8 +278,7 @@ static void test_conv1d_forward(ut_dev dev) {
out->shape.shape[2] == 2);
assert_data(out, ((float[]){-2.f, -2.f}), 1e-4f);
- ut_free(out);
- ut_free(x);
+ ut_free_all(out, x);
ut_conv1d_free(&l);
}
@@ -326,9 +305,7 @@ static void test_conv1d_backward(ut_dev dev) {
assert(dx->shape.ndim == 3 && dx->shape.shape[2] == 4);
assert_data(dW, ((float[]){1.f, 2.f, 3.f}), 1e-4f);
- ut_free(dx);
- ut_free(dW);
- ut_free(go);
+ ut_free_all(dx, dW, go);
ut_conv1d_cache_free(&c);
ut_free(x);
ut_conv1d_free(&l);
@@ -361,8 +338,7 @@ static void test_conv2d_forward(ut_dev dev) {
104.f, 104.f, 1014.f, 1012.f, 1008.f, 1006.f}),
1e-4f);
- ut_free(out);
- ut_free(x);
+ ut_free_all(out, x);
ut_conv2d_free(&l);
}
@@ -394,9 +370,7 @@ static void test_conv2d_backward(ut_dev dev) {
assert_data(dW, ((float[]){1.f, 2.f, 4.f, 5.f}), 1e-4f);
assert_data(dx, ((float[]){1.f, 0.f, 0.f, 0.f, -1.f, 0.f, 0.f, 0.f, 0.f}), 1e-4f);
- ut_free(dx);
- ut_free(dW);
- ut_free(go);
+ ut_free_all(dx, dW, go);
ut_conv2d_cache_free(&c);
ut_free(x);
ut_conv2d_free(&l);
@@ -404,8 +378,7 @@ static void test_conv2d_backward(ut_dev dev) {
static void test_sgd_momentum(void) {
ut_tensor* p = ut_alloc(1, (int[]){2}, UT_CPU);
- ut_tensor* params[1] = {p};
- ut_sgd opt = ut_sgd_alloc(params, 1, 0.1f, 0.9f);
+ ut_sgd opt = ut_sgd_alloc((ut_tensor*[]){p}, 1, 0.1f, 0.9f);
ut_sgd_zero(&opt);
opt.grads[0]->data[0] = 1.f;
diff --git a/utensil.h b/utensil.h
index 9a94f44..34bd66b 100644
--- a/utensil.h
+++ b/utensil.h
@@ -625,7 +625,7 @@ ut_tensor* ut_clone(ut_tensor* t) {
}
void ut_free(ut_tensor* t) {
- if (--t->rc > 0) return;
+ if (!t || --t->rc > 0) return;
if (t->owner)
ut_free(t->owner);
else {
@@ -635,6 +635,12 @@ void ut_free(ut_tensor* t) {
free(t);
}
+#define ut_free_all(...) \
+ do { \
+ ut_tensor* _ts[] = {__VA_ARGS__}; \
+ for (size_t _i = 0; _i < sizeof(_ts) / sizeof(*_ts); _i++) ut_free(_ts[_i]); \
+ } while (0)
+
ut_tensor* ut_retain(ut_tensor* t) { return t->rc++, t; }
ut_tensor* ut_view(ut_tensor* t, int ndim, const int* dim) {
@@ -671,8 +677,8 @@ ut_tensor* ut_transpose(ut_tensor* t, int dim0, int dim1) {
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);
+ _mtl_dispatch(_mc_t, "transpose_01", (void*[]){t->gpu_buf, out_g->gpu_buf}, 2, _tp,
+ (int)sizeof(_tp), t->shape.nelem);
out_g->dirty_cpu = true;
return out_g;
}
@@ -685,8 +691,8 @@ ut_tensor* ut_transpose(ut_tensor* t, int dim0, int dim1) {
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);
+ _mtl_dispatch(_mc_t, "transpose_12", (void*[]){t->gpu_buf, out_g->gpu_buf}, 2, _tp,
+ (int)sizeof(_tp), t->shape.nelem);
out_g->dirty_cpu = true;
return out_g;
}
@@ -746,9 +752,8 @@ static ut_tensor* ew_unary(ut_tensor* a, const char* kern, void (*fn)(float*, co
ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape, a->dev);
if (a->dev == UT_METAL && mc) {
ut_sync_gpu(a);
- void* bufs[2] = {a->gpu_buf, out->gpu_buf};
int n = a->shape.nelem;
- _mtl_dispatch(mc, kern, bufs, 2, &n, sizeof(int), n);
+ _mtl_dispatch(mc, kern, (void*[]){a->gpu_buf, out->gpu_buf}, 2, &n, sizeof(int), n);
out->dirty_cpu = true;
} else {
@@ -765,9 +770,8 @@ static ut_tensor* ew_binary(ut_tensor* a, ut_tensor* b, const char* kern,
_mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape, dev);
if (dev == UT_METAL && mc) {
- void* bufs[3] = {a->gpu_buf, b->gpu_buf, out->gpu_buf};
int n = a->shape.nelem;
- _mtl_dispatch(mc, kern, bufs, 3, &n, sizeof(int), n);
+ _mtl_dispatch(mc, kern, (void*[]){a->gpu_buf, b->gpu_buf, out->gpu_buf}, 3, &n, sizeof(int), n);
out->dirty_cpu = true;
} else {
ut_sync_cpu(a);
@@ -791,12 +795,12 @@ ut_tensor* ut_scale(ut_tensor* a, float s) {
ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape, a->dev);
if (a->dev == UT_METAL && mc) {
ut_sync_gpu(a);
- void* bufs[2] = {a->gpu_buf, out->gpu_buf};
struct {
float s;
int n;
} args = {s, a->shape.nelem};
- _mtl_dispatch(mc, "bscale", bufs, 2, &args, (int)sizeof(args), a->shape.nelem);
+ _mtl_dispatch(mc, "bscale", (void*[]){a->gpu_buf, out->gpu_buf}, 2, &args, (int)sizeof(args),
+ a->shape.nelem);
out->dirty_cpu = true;
} else {
ut_sync_cpu(a);
@@ -831,8 +835,7 @@ ut_tensor* ut_matmul(ut_tensor* a, ut_tensor* b) {
ut_dev dev = (a->dev == UT_METAL || b->dev == UT_METAL) ? UT_METAL : UT_CPU;
ut_to_device(a, dev);
ut_to_device(b, dev);
- int cd[2] = {m, n};
- ut_tensor* c = ut_alloc(2, cd, dev);
+ ut_tensor* c = ut_alloc(2, (int[]){m, n}, dev);
_mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
if (dev == UT_METAL && mc) {
ut_sync_gpu(a);
@@ -859,9 +862,9 @@ ut_tensor* ut_matmul(ut_tensor* a, ut_tensor* b) {
ut_to_device(b, dev);
ut_tensor* c = ut_alloc(3, cd, UT_METAL);
int p[4] = {B, m, n, k};
- void* bufs[3] = {a->gpu_buf, b->gpu_buf, c->gpu_buf};
// dispatch B*M*N threads via a 3D grid encoded as 1D
- _mtl_dispatch(mc, "bmatmul", bufs, 3, p, (int)sizeof(p), B * m * n);
+ _mtl_dispatch(mc, "bmatmul", (void*[]){a->gpu_buf, b->gpu_buf, c->gpu_buf}, 3, p,
+ (int)sizeof(p), B * m * n);
c->dirty_cpu = true;
return c;
}
@@ -886,8 +889,7 @@ static ut_tensor* ut_matmul_t(ut_tensor* a, ut_tensor* b, bool ta, bool tb) {
ut_dev dev = (a->dev == UT_METAL || b->dev == UT_METAL) ? UT_METAL : UT_CPU;
ut_to_device(a, dev);
ut_to_device(b, dev);
- int cd[2] = {m, n};
- ut_tensor* c = ut_alloc(2, cd, dev);
+ ut_tensor* c = ut_alloc(2, (int[]){m, n}, dev);
_mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
if (dev == UT_METAL && mc) {
ut_sync_gpu(a);
@@ -923,8 +925,8 @@ ut_tensor* ut_linear_forward(ut_linear* l, ut_tensor* x) {
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);
+ _mtl_dispatch(mc, "bias_add", (void*[]){out->gpu_buf, l->bias->gpu_buf}, 2, params,
+ (int)sizeof(params), B * l->nout);
out->dirty_cpu = true;
} else {
ut_sync_cpu(out);
@@ -961,8 +963,8 @@ ut_tensor* ut_linear_backward(ut_linear* l, ut_tensor* x, ut_tensor* grad_out, u
}
void ut_linear_free(ut_linear* l) {
- ut_free(l->weight);
- if (l->bias) ut_free(l->bias);
+ ut_free_all(l->weight, l->bias);
+ l->weight = l->bias = NULL;
}
ut_tensor* ut_relu_backward(ut_tensor* grad_out, ut_tensor* fwd_input) {
@@ -971,9 +973,9 @@ ut_tensor* ut_relu_backward(ut_tensor* grad_out, ut_tensor* fwd_input) {
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);
+ _mtl_dispatch(mc, "relu_bwd", (void*[]){grad_out->gpu_buf, fwd_input->gpu_buf, gi->gpu_buf}, 3,
+ &n, sizeof(int), n);
gi->dirty_cpu = true;
} else {
ut_sync_cpu(grad_out);
@@ -1021,9 +1023,10 @@ ut_tensor* ut_layernorm_forward(ut_layernorm* l, ut_tensor* x, ut_layernorm_cach
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);
+ _mtl_dispatch_rows(_mc_ln, "ln_fwd",
+ (void*[]){x->gpu_buf, l->weight->gpu_buf, l->bias->gpu_buf, out->gpu_buf,
+ xnorm_t->gpu_buf, rstd_t->gpu_buf},
+ 6, p2, (int)sizeof(p2), rows, tgsize);
out->dirty_cpu = true;
xnorm_t->dirty_cpu = true;
rstd_t->dirty_cpu = true;
@@ -1112,9 +1115,10 @@ ut_tensor* ut_layernorm_backward(ut_layernorm* l, ut_layernorm_cache* cache, ut_
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);
+ _mtl_dispatch_rows(_mc_lnb, "ln_bwd",
+ (void*[]){grad_out->gpu_buf, cache->xnorm->gpu_buf, cache->rstd->gpu_buf,
+ l->weight->gpu_buf, dx->gpu_buf},
+ 5, p2, (int)sizeof(p2), rows, tgsize);
dx->dirty_cpu = true;
return dx;
}
@@ -1140,14 +1144,11 @@ ut_tensor* ut_layernorm_backward(ut_layernorm* l, ut_layernorm_cache* cache, ut_
}
void ut_layernorm_cache_free(ut_layernorm_cache* c) {
- ut_free(c->xnorm);
- ut_free(c->mean);
- ut_free(c->rstd);
+ ut_free_all(c->xnorm, c->mean, c->rstd);
c->xnorm = c->mean = c->rstd = NULL;
}
void ut_layernorm_free(ut_layernorm* l) {
- ut_free(l->weight);
- ut_free(l->bias);
+ ut_free_all(l->weight, l->bias);
l->weight = l->bias = NULL;
}
@@ -1166,8 +1167,8 @@ ut_tensor* ut_im2col(ut_tensor* x, int kh, int kw, int stride, int pad) {
if (mc && x->dev == UT_METAL) {
ut_tensor* col = ut_alloc(2, col_dims, UT_METAL);
int params[10] = {N, C, H, W, kh, kw, stride, pad, Ho, Wo};
- void* ib[2] = {x->gpu_buf, col->gpu_buf};
- _mtl_dispatch(mc, "im2col_k", ib, 2, params, (int)sizeof(params), C * kh * kw * N * Ho * Wo);
+ _mtl_dispatch(mc, "im2col_k", (void*[]){x->gpu_buf, col->gpu_buf}, 2, params,
+ (int)sizeof(params), C * kh * kw * N * Ho * Wo);
col->dirty_cpu = true;
return col;
}
@@ -1204,8 +1205,8 @@ ut_tensor* ut_col2im(ut_tensor* col, int N, int C, int H, int W, int kh, int kw,
ut_to_device(col, UT_METAL);
ut_tensor* out = ut_alloc(4, dims, UT_METAL);
int _cp[10] = {N, C, H, W, kh, kw, stride, pad, Ho, Wo};
- void* _cb[2] = {col->gpu_buf, out->gpu_buf};
- _mtl_dispatch(_mc_c, "col2im_k", _cb, 2, _cp, (int)sizeof(_cp), N * C * H * W);
+ _mtl_dispatch(_mc_c, "col2im_k", (void*[]){col->gpu_buf, out->gpu_buf}, 2, _cp,
+ (int)sizeof(_cp), N * C * H * W);
out->dirty_cpu = true;
return out;
}
@@ -1246,22 +1247,19 @@ ut_tensor* ut_conv1d_forward(ut_conv1d* l, ut_tensor* x, ut_conv1d_cache* cache)
int N = x->shape.shape[0], L = x->shape.shape[2];
int Lo = (L + 2 * l->pad - l->kw) / l->stride + 1;
// expand x to [N, C, 1, L]
- int xd4[4] = {N, l->in_c, 1, L};
- ut_reshape(x, 4, xd4); // temporarily treat as [N, in_c, 1, L]
+ ut_reshape(x, 4, (int[]){N, l->in_c, 1, L}); // temporarily treat as [N, in_c, 1, L]
// Use im2col with kh=1, stride_h=1, pad_h=0
ut_tensor* col = ut_im2col(x, 1, l->kw, l->stride, l->pad);
- int w2d[2] = {l->out_c, l->in_c * l->kw};
ut_tensor w_view = *l->weight;
- w_view.shape = ut_shape_new(2, w2d);
+ w_view.shape = ut_shape_new(2, (int[]){l->out_c, l->in_c * l->kw});
w_view.rc = 0x7fffffff;
ut_tensor* out2 = ut_matmul(&w_view, col); // [out_c, N*Lo]
ut_sync_cpu(out2);
// out2 is [out_c, N*Lo] — reshape to [out_c, N, Lo] then transpose(0,1)
- int cn[3] = {l->out_c, N, Lo};
- ut_reshape(out2, 3, cn);
+ ut_reshape(out2, 3, (int[]){l->out_c, N, Lo});
ut_tensor* out = ut_transpose(out2, 0, 1); // [N, out_c, Lo]
ut_free(out2);
@@ -1274,8 +1272,7 @@ ut_tensor* ut_conv1d_forward(ut_conv1d* l, ut_tensor* x, ut_conv1d_cache* cache)
out->data[(n * l->out_c + oc) * Lo + lp] += l->bias->data[oc];
}
// restore x to original 3D shape before cache retain
- int x3d[3] = {N, l->in_c, L};
- ut_reshape(x, 3, x3d);
+ ut_reshape(x, 3, (int[]){N, l->in_c, L});
if (cache) {
cache->input = ut_retain(x);
cache->col = ut_retain(col);
@@ -1296,8 +1293,7 @@ ut_tensor* ut_conv1d_backward(ut_conv1d* l, ut_conv1d_cache* cache, ut_tensor* g
// go_2d: [out_c, N*Lo] — transpose + reshape grad_out [N,out_c,Lo]
ut_tensor* go_t = ut_transpose(grad_out, 0, 1); // [out_c, N, Lo]
- int d2[2] = {l->out_c, N * Lo};
- ut_reshape(go_t, 2, d2); // [out_c, N*Lo]
+ ut_reshape(go_t, 2, (int[]){l->out_c, N * Lo}); // [out_c, N*Lo]
_mtl_ctx_t* _mc = (_mtl_ctx_t*)ut_metal_ctx();
if (_mc) ut_to_device(go_t, UT_METAL);
@@ -1317,9 +1313,8 @@ ut_tensor* ut_conv1d_backward(ut_conv1d* l, ut_conv1d_cache* cache, ut_tensor* g
db->data[oc] += grad_out->data[(n * l->out_c + oc) * Lo + lp];
}
- int w2d[2] = {l->out_c, l->in_c * l->kw};
ut_tensor w_view = *l->weight;
- w_view.shape = ut_shape_new(2, w2d);
+ w_view.shape = ut_shape_new(2, (int[]){l->out_c, l->in_c * l->kw});
w_view.rc = 0x7fffffff;
ut_tensor* dcol = ut_matmul_t(&w_view, go_t, true, false);
ut_free(go_t);
@@ -1327,19 +1322,16 @@ ut_tensor* ut_conv1d_backward(ut_conv1d* l, ut_conv1d_cache* cache, ut_tensor* g
// col2im back to [N, C, 1, L] then squeeze to [N, C, L]
ut_tensor* dx4 = ut_col2im(dcol, N, l->in_c, 1, L, 1, l->kw, l->stride, l->pad);
ut_free(dcol);
- int dx3d[3] = {N, l->in_c, L};
- ut_reshape(dx4, 3, dx3d);
+ ut_reshape(dx4, 3, (int[]){N, l->in_c, L});
return dx4;
}
void ut_conv1d_cache_free(ut_conv1d_cache* c) {
- ut_free(c->input);
- ut_free(c->col);
+ ut_free_all(c->input, c->col);
c->input = c->col = NULL;
}
void ut_conv1d_free(ut_conv1d* l) {
- ut_free(l->weight);
- if (l->bias) ut_free(l->bias);
+ ut_free_all(l->weight, l->bias);
l->weight = l->bias = NULL;
}
// =========================================================
@@ -1362,15 +1354,13 @@ ut_tensor* ut_conv2d_forward(ut_conv2d* l, ut_tensor* x, ut_conv2d_cache* cache)
ut_tensor* col = ut_im2col(x, l->kh, l->kw, l->stride, l->pad);
- int w2d[2] = {l->out_c, l->in_c * l->kh * l->kw};
ut_tensor w_view = *l->weight;
- w_view.shape = ut_shape_new(2, w2d);
+ w_view.shape = ut_shape_new(2, (int[]){l->out_c, l->in_c * l->kh * l->kw});
w_view.rc = 0x7fffffff;
ut_tensor* out2 = ut_matmul(&w_view, col); // [out_c, N*Ho*Wo]
// reshape to [out_c, N, Ho, Wo] then transpose(0,1) -> NCHW
- int cnhw[4] = {l->out_c, N, Ho, Wo};
- ut_reshape(out2, 4, cnhw);
+ ut_reshape(out2, 4, (int[]){l->out_c, N, Ho, Wo});
ut_tensor* out = ut_transpose(out2, 0, 1); // [N, out_c, Ho, Wo]
ut_free(out2);
@@ -1380,8 +1370,8 @@ ut_tensor* ut_conv2d_forward(ut_conv2d* l, ut_tensor* x, ut_conv2d_cache* cache)
ut_sync_gpu(out);
ut_to_device(l->bias, UT_METAL);
int params[3] = {N, l->out_c, Ho * Wo};
- void* bufs[2] = {out->gpu_buf, l->bias->gpu_buf};
- _mtl_dispatch(mc, "bias_add", bufs, 2, params, (int)sizeof(params), N * l->out_c * Ho * Wo);
+ _mtl_dispatch(mc, "bias_add", (void*[]){out->gpu_buf, l->bias->gpu_buf}, 2, params,
+ (int)sizeof(params), N * l->out_c * Ho * Wo);
out->dirty_cpu = true;
} else {
ut_sync_cpu(out);
@@ -1413,9 +1403,8 @@ ut_tensor* ut_conv2d_backward(ut_conv2d* l, ut_conv2d_cache* cache, ut_tensor* g
ut_sync_cpu(dW);
// go_2d: [out_c, N*Ho*Wo] — transpose + reshape grad_out [N,out_c,Ho,Wo]
- ut_tensor* go_t = ut_transpose(grad_out, 0, 1); // [out_c, N, Ho, Wo]
- int d2[2] = {l->out_c, N * Ho * Wo};
- ut_reshape(go_t, 2, d2); // [out_c, N*Ho*Wo]
+ ut_tensor* go_t = ut_transpose(grad_out, 0, 1); // [out_c, N, Ho, Wo]
+ ut_reshape(go_t, 2, (int[]){l->out_c, N * Ho * Wo}); // [out_c, N*Ho*Wo]
ut_tensor* dWb = ut_matmul_t(go_t, col, false, true);
ut_sync_cpu(dWb);
@@ -1433,9 +1422,8 @@ ut_tensor* ut_conv2d_backward(ut_conv2d* l, ut_conv2d_cache* cache, ut_tensor* g
db->data[oc] += grad_out->data[((n * l->out_c + oc) * Ho + h) * Wo + w];
}
- int w2d[2] = {l->out_c, l->in_c * l->kh * l->kw};
ut_tensor w_view = *l->weight;
- w_view.shape = ut_shape_new(2, w2d);
+ w_view.shape = ut_shape_new(2, (int[]){l->out_c, l->in_c * l->kh * l->kw});
w_view.rc = 0x7fffffff;
ut_tensor* dcol = ut_matmul_t(&w_view, go_t, true, false);
ut_free(go_t);
@@ -1446,14 +1434,12 @@ ut_tensor* ut_conv2d_backward(ut_conv2d* l, ut_conv2d_cache* cache, ut_tensor* g
}
void ut_conv2d_cache_free(ut_conv2d_cache* c) {
- ut_free(c->input);
- ut_free(c->col);
+ ut_free_all(c->input, c->col);
c->input = c->col = NULL;
}
void ut_conv2d_free(ut_conv2d* l) {
- ut_free(l->weight);
- if (l->bias) ut_free(l->bias);
+ ut_free_all(l->weight, l->bias);
l->weight = l->bias = NULL;
}
@@ -1475,8 +1461,8 @@ ut_tensor* ut_softmax(ut_tensor* t, int dim) {
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);
+ _mtl_dispatch_rows(_mc_sm, "row_softmax", (void*[]){t->gpu_buf, out->gpu_buf}, 2, p2,
+ (int)sizeof(p2), outer, tgsize);
out->dirty_cpu = true;
return out;
}