← Commits · d73c2124
d73c2124b69befb66de831290192737f4f063f67
diff --git a/test.c b/test.c
index f400325..5ca4674 100644
--- a/test.c
+++ b/test.c
@@ -334,6 +334,74 @@ static void test_conv1d_backward(ut_dev dev) {
ut_conv1d_free(&l);
}
+static void test_conv2d_forward(ut_dev dev) {
+ // N=2,in_c=1,out_c=2,kh=kw=2,s=1,p=0. Distinct per-batch/per-channel bias
+ // catches any N<->C mixup in the forward transpose or bias broadcast.
+ // batch0 = [[1,2,3],[4,5,6],[7,8,9]], batch1 = [[9,8,7],[6,5,4],[3,2,1]]
+ // oc0 W=[1,0;0,-1] bias=100, oc1 W=[0,1;1,0] bias=1000
+ ut_conv2d l = ut_conv2d_alloc(1, 2, 2, 2, 1, 0, true, dev);
+ ut_free(l.weight);
+ l.weight = ut_from_data(4, (int[]){2, 1, 2, 2},
+ (float[]){1.f, 0.f, 0.f, -1.f, 0.f, 1.f, 1.f, 0.f}, dev);
+ ut_free(l.bias);
+ l.bias = ut_from_data(1, (int[]){2}, (float[]){100.f, 1000.f}, dev);
+
+ ut_tensor* x = ut_from_data(
+ 4, (int[]){2, 1, 3, 3},
+ (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 9.f, 8.f, 7.f, 6.f, 5.f, 4.f, 3.f, 2.f,
+ 1.f},
+ dev);
+ ut_tensor* out = ut_conv2d_forward(&l, x, NULL);
+ assert(out->dev == dev); // bias-add must not silently move a CPU model onto Metal
+ ut_sync_cpu(out);
+ assert(out->shape.ndim == 4 && out->shape.shape[0] == 2 && out->shape.shape[1] == 2 &&
+ out->shape.shape[2] == 2 && out->shape.shape[3] == 2);
+ assert_data(out,
+ ((float[]){96.f, 96.f, 96.f, 96.f, 1006.f, 1008.f, 1012.f, 1014.f, 104.f, 104.f,
+ 104.f, 104.f, 1014.f, 1012.f, 1008.f, 1006.f}),
+ 1e-4f);
+
+ ut_free(out);
+ ut_free(x);
+ ut_conv2d_free(&l);
+}
+
+static void test_conv2d_backward(ut_dev dev) {
+ // in_c=out_c=1, kh=kw=2, s=1, p=0, no bias, one-hot grad_out.
+ // x=[[1,2,3],[4,5,6],[7,8,9]], W=[1,0;0,-1], go=[[1,0],[0,0]]
+ // dW = the input patch under go's hot position = [1,2;4,5]
+ // dx = W scattered back at that same position = [1,0,0;0,-1,0;0,0,0]
+ ut_conv2d l = ut_conv2d_alloc(1, 1, 2, 2, 1, 0, false, dev);
+ ut_free(l.weight);
+ l.weight = ut_from_data(4, (int[]){1, 1, 2, 2}, (float[]){1.f, 0.f, 0.f, -1.f}, dev);
+
+ ut_tensor* x = ut_from_data(4, (int[]){1, 1, 3, 3},
+ (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}, dev);
+ ut_conv2d_cache c;
+ ut_tensor* out = ut_conv2d_forward(&l, x, &c);
+ ut_free(out);
+
+ ut_tensor* go =
+ ut_from_data(4, (int[]){1, 1, 2, 2}, (float[]){1.f, 0.f, 0.f, 0.f}, dev);
+ ut_tensor* dW = ut_alloc(4, (int[]){1, 1, 2, 2}, dev);
+ memset(dW->data, 0, 16);
+
+ ut_tensor* dx = ut_conv2d_backward(&l, &c, go, dW, NULL);
+ assert(dx->dev == dev); // forced Metal promotion must not leak into a CPU model's device
+ ut_sync_cpu(dx);
+ ut_sync_cpu(dW);
+ assert(dx->shape.ndim == 4 && dx->shape.shape[2] == 3 && dx->shape.shape[3] == 3);
+ 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_conv2d_cache_free(&c);
+ ut_free(x);
+ ut_conv2d_free(&l);
+}
+
static void test_sgd_momentum(void) {
ut_tensor* p = ut_alloc(1, (int[]){2}, UT_CPU);
ut_tensor* params[1] = {p};
@@ -395,6 +463,11 @@ int main() {
test_conv1d_backward(UT_CPU);
test_conv1d_backward(UT_METAL);
+ test_conv2d_forward(UT_CPU);
+ test_conv2d_forward(UT_METAL);
+ test_conv2d_backward(UT_CPU);
+ test_conv2d_backward(UT_METAL);
+
test_sgd_momentum();
return 0;
}
diff --git a/utensil.h b/utensil.h
index 45a702b..9a94f44 100644
--- a/utensil.h
+++ b/utensil.h
@@ -58,6 +58,17 @@ typedef struct ut_conv1d_cache {
ut_tensor* col;
} ut_conv1d_cache;
+typedef struct ut_conv2d {
+ ut_tensor* weight; // [out_c, in_c, kh, kw]
+ ut_tensor* bias; // [out_c]
+ int in_c, out_c, kh, kw, stride, pad;
+} ut_conv2d;
+
+typedef struct ut_conv2d_cache {
+ ut_tensor* input;
+ ut_tensor* col;
+} ut_conv2d_cache;
+
typedef struct ut_sgd {
ut_tensor** params; // pointers to model parameters (not owned)
ut_tensor** grads; // gradient accumulators (owned)
@@ -654,7 +665,7 @@ ut_tensor* ut_transpose(ut_tensor* t, int dim0, int dim1) {
// 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)) {
+ if (_mc_t && t->dev == UT_METAL && (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];
@@ -667,7 +678,8 @@ ut_tensor* ut_transpose(ut_tensor* t, int dim0, int dim1) {
}
// GPU fast-path for 1<->2 dims
- if (_mc_t && t->shape.ndim >= 3 && ((dim0 == 1 && dim1 == 2) || (dim0 == 2 && dim1 == 1))) {
+ if (_mc_t && t->dev == UT_METAL && 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],
@@ -1330,6 +1342,120 @@ void ut_conv1d_free(ut_conv1d* l) {
if (l->bias) ut_free(l->bias);
l->weight = l->bias = NULL;
}
+// =========================================================
+// Conv2D
+// =========================================================
+ut_conv2d ut_conv2d_alloc(int in_c, int out_c, int kh, int kw, int stride, int pad, bool bias,
+ ut_dev dev) {
+ ut_conv2d l = {.in_c = in_c, .out_c = out_c, .kh = kh, .kw = kw, .stride = stride, .pad = pad};
+ float std = sqrtf(2.f / (float)(in_c * kh * kw));
+ l.weight = ut_randn(4, (int[]){out_c, in_c, kh, kw}, 0.f, std, dev);
+ if (bias) l.bias = ut_alloc(1, (int[]){out_c}, dev);
+ return l;
+}
+
+// x: [N, in_c, H, W]
+ut_tensor* ut_conv2d_forward(ut_conv2d* l, ut_tensor* x, ut_conv2d_cache* cache) {
+ int N = x->shape.shape[0], H = x->shape.shape[2], W = x->shape.shape[3];
+ int Ho = (H + 2 * l->pad - l->kh) / l->stride + 1;
+ int Wo = (W + 2 * l->pad - l->kw) / l->stride + 1;
+
+ 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.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_tensor* out = ut_transpose(out2, 0, 1); // [N, out_c, Ho, Wo]
+ ut_free(out2);
+
+ if (l->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 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);
+ out->dirty_cpu = true;
+ } else {
+ ut_sync_cpu(out);
+ ut_sync_cpu(l->bias);
+ for (int n = 0; n < N; n++)
+ for (int oc = 0; oc < l->out_c; oc++)
+ for (int h = 0; h < Ho; h++)
+ for (int w = 0; w < Wo; w++)
+ out->data[((n * l->out_c + oc) * Ho + h) * Wo + w] += l->bias->data[oc];
+ }
+ }
+
+ if (cache) {
+ cache->input = ut_retain(x);
+ cache->col = ut_retain(col);
+ }
+ ut_free(col);
+ return out;
+}
+
+ut_tensor* ut_conv2d_backward(ut_conv2d* l, ut_conv2d_cache* cache, ut_tensor* grad_out,
+ ut_tensor* dW, ut_tensor* db) {
+ ut_tensor* col = cache->col;
+ int N = cache->input->shape.shape[0];
+ int H = cache->input->shape.shape[2], W = cache->input->shape.shape[3];
+ int Ho = (H + 2 * l->pad - l->kh) / l->stride + 1;
+ int Wo = (W + 2 * l->pad - l->kw) / l->stride + 1;
+ ut_sync_cpu(grad_out);
+ 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* dWb = ut_matmul_t(go_t, col, false, true);
+ 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);
+
+ if (l->bias && db) {
+ ut_sync_cpu(db);
+ ut_sync_cpu(grad_out);
+ for (int n = 0; n < N; n++)
+ for (int oc = 0; oc < l->out_c; oc++)
+ for (int h = 0; h < Ho; h++)
+ for (int w = 0; w < Wo; w++)
+ 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.rc = 0x7fffffff;
+ ut_tensor* dcol = ut_matmul_t(&w_view, go_t, true, false);
+ ut_free(go_t);
+
+ ut_tensor* dx = ut_col2im(dcol, N, l->in_c, H, W, l->kh, l->kw, l->stride, l->pad);
+ ut_free(dcol);
+ return dx;
+}
+
+void ut_conv2d_cache_free(ut_conv2d_cache* c) {
+ ut_free(c->input);
+ ut_free(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);
+ l->weight = l->bias = NULL;
+}
// =========================================================
// Softmax