← Commits · 89c030ec
89c030ec9ee9d3ef4f37752d0a8d97f3fef28d8c
diff --git a/utensil.h b/utensil.h
index 92b2542..4657e0a 100644
--- a/utensil.h
+++ b/utensil.h
@@ -63,6 +63,7 @@ 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;
+ void* _mpsg; // cached MPSGraph (forward + backward), lazily built; see _ut_conv2d_graph_t
} ut_conv2d;
typedef struct ut_conv2d_cache {
@@ -707,6 +708,160 @@ static void _mtl_buf_read(void* b, float* dst, size_t n) {
memcpy(dst, _p0(b, "contents"), n * sizeof(float));
}
+// =========================================================
+// MPSGraph conv2d (native Winograd/im2col-free conv, replaces the manual
+// im2col->MPSMatrix-matmul->col2im pipeline, which profiled ~10x slower)
+// =========================================================
+#define _MPS_F32 0x10000020UL
+
+static inline void* _mns(long v) { // [NSNumber numberWithLong:v]
+ return ((id (*)(id, SEL, long))objc_msgSend)((id)objc_getClass("NSNumber"),
+ sel_getUid("numberWithLong:"), v);
+}
+static void* _nsarr(void** objs, int n) {
+ return ((id (*)(id, SEL, const void*, unsigned long))objc_msgSend)(
+ (id)objc_getClass("NSArray"), sel_getUid("arrayWithObjects:count:"), objs, (unsigned long)n);
+}
+static void* _nsdict(void** objs, void** keys, int n) {
+ return ((id (*)(id, SEL, const void*, const void*, unsigned long))objc_msgSend)(
+ (id)objc_getClass("NSDictionary"), sel_getUid("dictionaryWithObjects:forKeys:count:"), objs,
+ keys, (unsigned long)n);
+}
+static void* _mps_shape4(int a, int b, int c, int d) {
+ return _nsarr((void*[]){_mns(a), _mns(b), _mns(c), _mns(d)}, 4);
+}
+static void* _mps_ph(void* graph, void* shape) {
+ return ((id (*)(id, SEL, id, unsigned long, id))objc_msgSend)(
+ (id)graph, sel_getUid("placeholderWithShape:dataType:name:"), (id)shape,
+ (unsigned long)_MPS_F32, (id)NULL);
+}
+// wrap a raw MTLBuffer as MPSGraphTensorData for feeding/reading a graph run
+static void* _mps_tdata(void* buf, void* shape) {
+ void* d = ((id (*)(id, SEL))objc_msgSend)((id)objc_getClass("MPSGraphTensorData"),
+ sel_getUid("alloc"));
+ return ((id (*)(id, SEL, id, id, unsigned long))objc_msgSend)(
+ (id)d, sel_getUid("initWithMTLBuffer:shape:dataType:"), (id)buf, (id)shape,
+ (unsigned long)_MPS_F32);
+}
+// run a graph, writing results directly into caller-owned MTLBuffers (no CPU round trip).
+// Builds and releases all the per-call ObjC wrapper objects (MPSGraphTensorData,
+// NSDictionary) itself, AND wraps the whole call in an autorelease pool.
+// This program never pushes a pool anywhere else, so without one here, every
+// object MPSGraph autoreleases internally while building/running the graph
+// (temporary NDArrays, encoder helpers, etc.) never gets drained -- it just
+// accumulates for the life of the process. That's invisible in host RSS (much
+// of it is GPU-backed) but very real: a few hundred batches of a many-layer
+// network (mobilenetv1) reliably got SIGKILLed by the OS before this fix, and
+// even the lighter resnet8 died a couple hundred batches into its 2nd epoch.
+static void _mps_run(_mtl_ctx_t* ctx, void* graph, void** feed_keys, void** feed_bufs,
+ void** feed_shapes, int nfeeds, void** res_keys, void** res_bufs,
+ void** res_shapes, int nres) {
+ extern void* objc_autoreleasePoolPush(void);
+ extern void objc_autoreleasePoolPop(void* ctx);
+ void* pool = objc_autoreleasePoolPush();
+
+ void* feed_objs[4];
+ void* res_objs[4];
+ for (int i = 0; i < nfeeds; i++) feed_objs[i] = _mps_tdata(feed_bufs[i], feed_shapes[i]);
+ for (int i = 0; i < nres; i++) res_objs[i] = _mps_tdata(res_bufs[i], res_shapes[i]);
+ void* feeds = _nsdict(feed_objs, feed_keys, nfeeds);
+ void* results = _nsdict(res_objs, res_keys, nres);
+ _mtl_flush(ctx); // ensure any pending raw-kernel dispatches complete before/order vs the graph
+ ((void (*)(id, SEL, id, id, id, id))objc_msgSend)(
+ (id)graph, sel_getUid("runWithMTLCommandQueue:feeds:targetOperations:resultsDictionary:"),
+ (id)ctx->queue, (id)feeds, (id)NULL, (id)results);
+ // feeds/results are autoreleased (NSDictionary factory ctor) -- the pool below
+ // owns releasing them now that one actually exists; only release the objects
+ // we explicitly alloc+init'd ourselves (MPSGraphTensorData), which the pool
+ // never touches.
+ for (int i = 0; i < nfeeds; i++) _v0(feed_objs[i], "release");
+ for (int i = 0; i < nres; i++) _v0(res_objs[i], "release");
+
+ objc_autoreleasePoolPop(pool);
+}
+
+// one compiled graph per ut_conv2d layer instance, holding forward conv +
+// both backward gradients (data, weights) so all three share placeholders;
+// rebuilt only if the batch shape (N) changes (e.g. last partial batch).
+typedef struct {
+ void* graph;
+ void* xph;
+ void* wph;
+ void* goph;
+ void* fwd_t;
+ void* dx_t;
+ void* dw_t;
+ void* xshape;
+ void* wshape;
+ void* oshape;
+ int N, C, H, W, Cout, KH, KW, stride, pad;
+} _ut_conv2d_graph_t;
+
+static _ut_conv2d_graph_t* _ut_conv2d_graph_build(int N, int C, int H, int W, int Cout, int KH,
+ int KW, int stride, int pad) {
+ _ut_conv2d_graph_t* g = (_ut_conv2d_graph_t*)calloc(1, sizeof(*g));
+ g->N = N, g->C = C, g->H = H, g->W = W, g->Cout = Cout, g->KH = KH, g->KW = KW,
+ g->stride = stride, g->pad = pad;
+ int Ho = (H + 2 * pad - KH) / stride + 1, Wo = (W + 2 * pad - KW) / stride + 1;
+
+ void* a = ((id (*)(id, SEL))objc_msgSend)((id)objc_getClass("MPSGraph"), sel_getUid("alloc"));
+ g->graph = _m0(a, "init");
+
+ g->xshape = _mps_shape4(N, C, H, W);
+ g->wshape = _mps_shape4(Cout, C, KH, KW);
+ g->oshape = _mps_shape4(N, Cout, Ho, Wo);
+ g->xph = _mps_ph(g->graph, g->xshape);
+ g->wph = _mps_ph(g->graph, g->wshape);
+ g->goph = _mps_ph(g->graph, g->oshape);
+
+ // descriptorWithStrideInX:strideInY:dilationRateInX:dilationRateInY:groups:
+ // paddingLeft:paddingRight:paddingTop:paddingBottom:paddingStyle:dataLayout:weightsLayout:
+ // paddingStyle=Explicit(0), dataLayout=NCHW(0), weightsLayout=OIHW(2)
+ void* desc = ((id (*)(id, SEL, unsigned long, unsigned long, unsigned long, unsigned long,
+ unsigned long, unsigned long, unsigned long, unsigned long, unsigned long,
+ unsigned long, unsigned long, unsigned long))objc_msgSend)(
+ (id)objc_getClass("MPSGraphConvolution2DOpDescriptor"),
+ sel_getUid("descriptorWithStrideInX:strideInY:dilationRateInX:dilationRateInY:groups:"
+ "paddingLeft:paddingRight:paddingTop:paddingBottom:paddingStyle:dataLayout:"
+ "weightsLayout:"),
+ (unsigned long)stride, (unsigned long)stride, 1UL, 1UL, 1UL, (unsigned long)pad,
+ (unsigned long)pad, (unsigned long)pad, (unsigned long)pad, 0UL, 0UL, 2UL);
+
+ g->fwd_t = ((id (*)(id, SEL, id, id, id, id))objc_msgSend)(
+ (id)g->graph, sel_getUid("convolution2DWithSourceTensor:weightsTensor:descriptor:name:"),
+ (id)g->xph, (id)g->wph, (id)desc, (id)NULL);
+ g->dx_t = ((id (*)(id, SEL, id, id, id, id, id))objc_msgSend)(
+ (id)g->graph,
+ sel_getUid("convolution2DDataGradientWithIncomingGradientTensor:weightsTensor:outputShape:"
+ "forwardConvolutionDescriptor:name:"),
+ (id)g->goph, (id)g->wph, (id)g->xshape, (id)desc, (id)NULL);
+ g->dw_t = ((id (*)(id, SEL, id, id, id, id, id))objc_msgSend)(
+ (id)g->graph,
+ sel_getUid("convolution2DWeightsGradientWithIncomingGradientTensor:sourceTensor:"
+ "outputShape:forwardConvolutionDescriptor:name:"),
+ (id)g->goph, (id)g->xph, (id)g->wshape, (id)desc, (id)NULL);
+ return g;
+}
+
+static void _ut_conv2d_graph_free(_ut_conv2d_graph_t* g) {
+ if (!g) return;
+ _v0(g->graph, "release");
+ free(g);
+}
+
+// lazily (re)build the cached graph for l, keyed on the current batch shape
+static _ut_conv2d_graph_t* _ut_conv2d_graph_get(void** slot, int N, int C, int H, int W, int Cout,
+ int KH, int KW, int stride, int pad) {
+ _ut_conv2d_graph_t* g = (_ut_conv2d_graph_t*)*slot;
+ if (g && g->N == N && g->C == C && g->H == H && g->W == W && g->Cout == Cout && g->KH == KH &&
+ g->KW == KW && g->stride == stride && g->pad == pad)
+ return g;
+ _ut_conv2d_graph_free(g);
+ g = _ut_conv2d_graph_build(N, C, H, W, Cout, KH, KW, stride, pad);
+ *slot = g;
+ return g;
+}
+
static void* _g_metal = NULL;
static int _g_metal_init = 0;
@@ -1617,6 +1772,38 @@ ut_tensor* ut_conv2d_forward(ut_conv2d* l, ut_tensor* x, ut_conv2d_cache* cache)
int Ho = (H + 2 * l->pad - l->kh) / l->stride + 1;
int Wo = (W + 2 * l->pad - l->kw) / l->stride + 1;
+#ifdef __APPLE__
+ // MPSGraph path: native convolution kernel (Winograd/implicit-GEMM, picked
+ // internally by MPS) instead of im2col+matmul -- ~10x faster in practice,
+ // see PERF.md. Falls through to the im2col path below for CPU tensors.
+ _mtl_ctx_t* mc0 = (_mtl_ctx_t*)ut_metal_ctx();
+ if (mc0 && x->dev == UT_METAL) {
+ ut_sync_gpu(x);
+ ut_to_device(l->weight, UT_METAL);
+ _ut_conv2d_graph_t* g = _ut_conv2d_graph_get(&l->_mpsg, N, l->in_c, H, W, l->out_c, l->kh,
+ l->kw, l->stride, l->pad);
+ ut_tensor* out = ut_alloc(4, (int[]){N, l->out_c, Ho, Wo}, UT_METAL);
+ _mps_run(mc0, g->graph, (void*[]){g->xph, g->wph},
+ (void*[]){x->gpu_buf, l->weight->gpu_buf}, (void*[]){g->xshape, g->wshape}, 2,
+ (void*[]){g->fwd_t}, (void*[]){out->gpu_buf}, (void*[]){g->oshape}, 1);
+ out->dirty_cpu = true;
+
+ if (l->bias) {
+ ut_to_device(l->bias, UT_METAL);
+ int params[3] = {N, l->out_c, Ho * Wo};
+ _mtl_dispatch(mc0, "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;
+ }
+
+ if (cache) {
+ cache->input = ut_retain(x);
+ cache->col = NULL; // no im2col buffer on this path -- signals backward to use MPSGraph too
+ }
+ return out;
+ }
+#endif
+
ut_tensor* col = ut_im2col(x, l->kh, l->kw, l->stride, l->pad);
ut_tensor w_view = *l->weight;
@@ -1695,6 +1882,33 @@ ut_tensor* ut_conv2d_backward(ut_conv2d* l, ut_conv2d_cache* cache, ut_tensor* g
}
}
+#ifdef __APPLE__
+ // MPSGraph path -- mirrors the forward path above; col is NULL exactly when
+ // forward used MPSGraph (see ut_conv2d_forward), so that flag also selects
+ // the matching backward here.
+ _mtl_ctx_t* mc1 = (_mtl_ctx_t*)ut_metal_ctx();
+ if (mc1 && grad_out->dev == UT_METAL && col == NULL) {
+ ut_tensor* x = cache->input;
+ ut_sync_gpu(grad_out);
+ ut_to_device(x, UT_METAL);
+ ut_to_device(l->weight, UT_METAL);
+ _ut_conv2d_graph_t* g = _ut_conv2d_graph_get(&l->_mpsg, N, l->in_c, H, W, l->out_c, l->kh,
+ l->kw, l->stride, l->pad);
+ ut_tensor* dx = ut_alloc(4, (int[]){N, l->in_c, H, W}, UT_METAL);
+ ut_tensor* dWb = ut_alloc(l->weight->shape.ndim, l->weight->shape.shape, UT_METAL);
+ _mps_run(mc1, g->graph, (void*[]){g->xph, g->wph, g->goph},
+ (void*[]){x->gpu_buf, l->weight->gpu_buf, grad_out->gpu_buf},
+ (void*[]){g->xshape, g->wshape, g->oshape}, 3, (void*[]){g->dx_t, g->dw_t},
+ (void*[]){dx->gpu_buf, dWb->gpu_buf}, (void*[]){g->xshape, g->wshape}, 2);
+ dx->dirty_cpu = dWb->dirty_cpu = 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);
+ return dx;
+ }
+#endif
+
// 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]
ut_reshape(go_t, 2, (int[]){l->out_c, N * Ho * Wo}); // [out_c, N*Ho*Wo]
@@ -1724,6 +1938,10 @@ void ut_conv2d_cache_free(ut_conv2d_cache* c) {
void ut_conv2d_free(ut_conv2d* l) {
ut_free_all(l->weight, l->bias);
l->weight = l->bias = NULL;
+#ifdef __APPLE__
+ _ut_conv2d_graph_free((_ut_conv2d_graph_t*)l->_mpsg);
+#endif
+ l->_mpsg = NULL;
}
// =========================================================