← Commits · c4a33fc2
c4a33fc2dca85b9f9142513e7cd314722b68c462
diff --git a/utensil.h b/utensil.h
index b88604e..be3cd8e 100644
--- a/utensil.h
+++ b/utensil.h
@@ -316,6 +316,101 @@ static const char* _mtl_src =
// dx
"for(int i=(int)lid;i<d;i+=(int)tgs)"
"dxrow[i]=rs*(w[i]*gorow[i]-(sum_go_w+xnrow[i]*sum_go_w_xn)/(float)d);}"
+ "struct BNParams{int N,C,HW;float eps,mom;};\n"
+ // batchnorm2d forward (training): x[N,C,H,W] -> out[N,C,H,W], one threadgroup
+ // per channel, reducing over N*HW elements. Writes xn (pre-affine) and rstd
+ // for backward, and updates running_mean/running_var in place (unbiased var).
+ "kernel void bn2d_fwd_train("
+ "device const float*x,device const float*w,device const float*b,"
+ "device float*rm,device float*rv,"
+ "device float*out,device float*xn,device float*rstd,"
+ "constant BNParams&p,"
+ "uint gid[[threadgroup_position_in_grid]],"
+ "uint lid[[thread_position_in_threadgroup]],"
+ "uint tgs[[threads_per_threadgroup]]){"
+ "int C=p.C,HW=p.HW,M=p.N*HW,c=(int)gid;"
+ "threadgroup float sh[1024];"
+ "float ls=0;for(int i=(int)lid;i<M;i+=(int)tgs){int n=i/HW,hw=i%HW;ls+=x[(n*C+c)*HW+hw];}"
+ "sh[lid]=ls;threadgroup_barrier(mem_flags::mem_threadgroup);"
+ "for(uint "
+ "s=tgs/"
+ "2;s>0;s>>=1){if(lid<s)sh[lid]+=sh[lid+s];threadgroup_barrier(mem_flags::mem_threadgroup);}"
+ "float mu=sh[0]/(float)M;"
+ "ls=0;for(int i=(int)lid;i<M;i+=(int)tgs){int n=i/HW,hw=i%HW;float "
+ "v=x[(n*C+c)*HW+hw]-mu;ls+=v*v;}"
+ "sh[lid]=ls;threadgroup_barrier(mem_flags::mem_threadgroup);"
+ "for(uint "
+ "s=tgs/"
+ "2;s>0;s>>=1){if(lid<s)sh[lid]+=sh[lid+s];threadgroup_barrier(mem_flags::mem_threadgroup);}"
+ "float var=sh[0]/(float)M;float rs=rsqrt(var+p.eps);"
+ "if(lid==0){"
+ "rstd[c]=rs;"
+ "float vu=M>1?var*(float)M/(float)(M-1):var;"
+ "rm[c]=(1.f-p.mom)*rm[c]+p.mom*mu;"
+ "rv[c]=(1.f-p.mom)*rv[c]+p.mom*vu;}"
+ "float wc=w[c],bc=b[c];"
+ "for(int i=(int)lid;i<M;i+=(int)tgs){"
+ "int n=i/HW,hw=i%HW,idx=(n*C+c)*HW+hw;"
+ "float xni=(x[idx]-mu)*rs;xn[idx]=xni;out[idx]=wc*xni+bc;}}\n"
+ // batchnorm2d forward (eval): elementwise using the running stats, no reduction
+ "kernel void bn2d_fwd_eval("
+ "device const float*x,device const float*w,device const float*b,"
+ "device const float*rm,device const float*rv,device float*out,"
+ "constant BNParams&p,uint idx[[thread_position_in_grid]]){"
+ "int C=p.C,HW=p.HW,tot=p.N*C*HW;if((int)idx>=tot)return;"
+ "int c=((int)idx/HW)%C;float rs=rsqrt(rv[c]+p.eps);"
+ "out[idx]=w[c]*(x[idx]-rm[c])*rs+b[c];}\n"
+ // batchnorm2d backward: dx only (dW/db stay CPU-accumulated, like LayerNorm's
+ // ln_bwd). One threadgroup per channel, reducing over N*HW elements.
+ // also emits dW_partial[C]/db_partial[C] (this channel's contribution to
+ // this batch) so the caller only has to sync [C] floats back to CPU for
+ // the dW/db accumulation, instead of the full [N,C,H,W] go/xn tensors.
+ "kernel void bn2d_bwd("
+ "device const float*go,device const float*xn,device const float*rs_,"
+ "device const float*w,device float*dx,device float*dwp,device float*dbp,"
+ "constant BNParams&p,"
+ "uint gid[[threadgroup_position_in_grid]],"
+ "uint lid[[thread_position_in_threadgroup]],"
+ "uint tgs[[threads_per_threadgroup]]){"
+ "int C=p.C,HW=p.HW,M=p.N*HW,c=(int)gid;"
+ "threadgroup float sh[1024];"
+ "float ls=0;for(int i=(int)lid;i<M;i+=(int)tgs){int n=i/HW,hw=i%HW;ls+=go[(n*C+c)*HW+hw];}"
+ "sh[lid]=ls;threadgroup_barrier(mem_flags::mem_threadgroup);"
+ "for(uint "
+ "s=tgs/"
+ "2;s>0;s>>=1){if(lid<s)sh[lid]+=sh[lid+s];threadgroup_barrier(mem_flags::mem_threadgroup);}"
+ "float sum_go=sh[0];"
+ "ls=0;for(int i=(int)lid;i<M;i+=(int)tgs){int "
+ "n=i/HW,hw=i%HW,idx=(n*C+c)*HW+hw;ls+=go[idx]*xn[idx];}"
+ "sh[lid]=ls;threadgroup_barrier(mem_flags::mem_threadgroup);"
+ "for(uint "
+ "s=tgs/"
+ "2;s>0;s>>=1){if(lid<s)sh[lid]+=sh[lid+s];threadgroup_barrier(mem_flags::mem_threadgroup);}"
+ "float sum_go_xn=sh[0];"
+ "if(lid==0){dwp[c]=sum_go_xn;dbp[c]=sum_go;}"
+ "float rs=rs_[c],wc=w[c];"
+ "for(int i=(int)lid;i<M;i+=(int)tgs){"
+ "int n=i/HW,hw=i%HW,idx=(n*C+c)*HW+hw;"
+ "dx[idx]=rs*wc*(go[idx]-(sum_go+xn[idx]*sum_go_xn)/(float)M);}}\n"
+ // global avgpool forward: x[N,C,H,W] -> out[N,C], one threadgroup per (n,c)
+ // pair reducing a contiguous HW-sized block (rows = N*C)
+ "kernel void gap_fwd(device const float*x,device float*out,constant int&HW,"
+ "uint gid[[threadgroup_position_in_grid]],"
+ "uint lid[[thread_position_in_threadgroup]],"
+ "uint tgs[[threads_per_threadgroup]]){"
+ "device const float*row=x+gid*HW;"
+ "threadgroup float sh[1024];"
+ "float ls=0;for(int i=(int)lid;i<HW;i+=(int)tgs)ls+=row[i];"
+ "sh[lid]=ls;threadgroup_barrier(mem_flags::mem_threadgroup);"
+ "for(uint s=tgs/2;s>0;s>>=1){if(lid<s)sh[lid]+=sh[lid+s];"
+ "threadgroup_barrier(mem_flags::mem_threadgroup);}"
+ "if(lid==0)out[gid]=sh[0]/(float)HW;}\n"
+ // global avgpool backward: broadcast go[N,C]/HW back to dx[N,C,H,W]
+ "struct GAPParams{int HW,tot;};\n"
+ "kernel void gap_bwd(device const float*go,device float*dx,constant GAPParams&p,"
+ "uint idx[[thread_position_in_grid]]){"
+ "if((int)idx>=p.tot)return;"
+ "dx[idx]=go[(int)idx/p.HW]/(float)p.HW;}\n"
// col2im (gather, no atomics): col[C*kh*kw,N*Ho*Wo] -> dx[N,C,H,W]
// p = {N,C,H,W,kh,kw,stride,pad,Ho,Wo}
"kernel void col2im_k("
@@ -1517,7 +1612,54 @@ ut_batchnorm2d ut_batchnorm2d_alloc(int c, ut_dev dev) {
ut_tensor* ut_batchnorm2d_forward(ut_batchnorm2d* l, ut_tensor* x, bool training,
ut_batchnorm2d_cache* cache) {
int N = x->shape.shape[0], C = l->c, H = x->shape.shape[2], W = x->shape.shape[3];
- int M = N * H * W;
+ int HW = H * W, M = N * HW;
+
+ // GPU fast-path
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ if (mc && x->dev == UT_METAL) {
+ ut_sync_gpu(x);
+ ut_to_device(l->weight, UT_METAL);
+ ut_to_device(l->bias, UT_METAL);
+ ut_to_device(l->running_mean, UT_METAL);
+ ut_to_device(l->running_var, UT_METAL);
+ ut_tensor* out = ut_alloc(4, x->shape.shape, UT_METAL);
+ struct {
+ int N, C, HW;
+ float eps, mom;
+ } params = {N, C, HW, l->eps, l->momentum};
+
+ if (training) {
+ ut_tensor* xnorm_t = ut_alloc(4, x->shape.shape, UT_METAL);
+ ut_tensor* rstd_t = ut_alloc(1, (int[]){C}, UT_METAL);
+ int tgsize = M < 64 ? 32 : (M < 256 ? 64 : (M < 512 ? 128 : 256));
+ if (tgsize > 1024) tgsize = 1024;
+ _mtl_dispatch_rows(
+ mc, "bn2d_fwd_train",
+ (void*[]){x->gpu_buf, l->weight->gpu_buf, l->bias->gpu_buf, l->running_mean->gpu_buf,
+ l->running_var->gpu_buf, out->gpu_buf, xnorm_t->gpu_buf, rstd_t->gpu_buf},
+ 8, ¶ms, (int)sizeof(params), C, tgsize);
+ out->dirty_cpu = true;
+ l->running_mean->dirty_cpu = true;
+ l->running_var->dirty_cpu = true;
+ xnorm_t->dirty_cpu = true;
+ rstd_t->dirty_cpu = true;
+ if (cache) {
+ cache->xnorm = xnorm_t;
+ cache->rstd = rstd_t;
+ } else {
+ ut_free_all(xnorm_t, rstd_t);
+ }
+ } else {
+ _mtl_dispatch(mc, "bn2d_fwd_eval",
+ (void*[]){x->gpu_buf, l->weight->gpu_buf, l->bias->gpu_buf,
+ l->running_mean->gpu_buf, l->running_var->gpu_buf, out->gpu_buf},
+ 6, ¶ms, (int)sizeof(params), N * C * HW);
+ out->dirty_cpu = true;
+ }
+ return out;
+ }
+
+ // CPU fallback
ut_sync_cpu(x);
ut_sync_cpu(l->weight);
ut_sync_cpu(l->bias);
@@ -1533,16 +1675,14 @@ ut_tensor* ut_batchnorm2d_forward(ut_batchnorm2d* l, ut_tensor* x, bool training
if (training) {
float sum = 0;
for (int n = 0; n < N; n++)
- for (int h = 0; h < H; h++)
- for (int w = 0; w < W; w++) sum += x->data[((n * C + c) * H + h) * W + w];
+ for (int hw = 0; hw < HW; hw++) sum += x->data[(n * C + c) * HW + hw];
mu = sum / (float)M;
float var = 0;
for (int n = 0; n < N; n++)
- for (int h = 0; h < H; h++)
- for (int w = 0; w < W; w++) {
- float diff = x->data[((n * C + c) * H + h) * W + w] - mu;
- var += diff * diff;
- }
+ for (int hw = 0; hw < HW; hw++) {
+ float diff = x->data[(n * C + c) * HW + hw] - mu;
+ var += diff * diff;
+ }
var /= (float)M;
rs = 1.f / sqrtf(var + l->eps);
float var_unbiased = M > 1 ? var * (float)M / (float)(M - 1) : var;
@@ -1554,13 +1694,12 @@ ut_tensor* ut_batchnorm2d_forward(ut_batchnorm2d* l, ut_tensor* x, bool training
rs = 1.f / sqrtf(l->running_var->data[c] + l->eps);
}
for (int n = 0; n < N; n++)
- for (int h = 0; h < H; h++)
- for (int w = 0; w < W; w++) {
- int idx = ((n * C + c) * H + h) * W + w;
- float xn = (x->data[idx] - mu) * rs;
- out->data[idx] = l->weight->data[c] * xn + l->bias->data[c];
- if (cache) xnorm_t->data[idx] = xn;
- }
+ for (int hw = 0; hw < HW; hw++) {
+ int idx = (n * C + c) * HW + hw;
+ float xn = (x->data[idx] - mu) * rs;
+ out->data[idx] = l->weight->data[c] * xn + l->bias->data[c];
+ if (cache) xnorm_t->data[idx] = xn;
+ }
if (cache) rstd_t->data[c] = rs;
}
out->dirty_gpu = true;
@@ -1579,34 +1718,71 @@ ut_tensor* ut_batchnorm2d_backward(ut_batchnorm2d* l, ut_batchnorm2d_cache* cach
ut_tensor* grad_out, ut_tensor* dW, ut_tensor* db) {
int N = grad_out->shape.shape[0], C = l->c, H = grad_out->shape.shape[2],
W = grad_out->shape.shape[3];
- int M = N * H * W;
+ int HW = H * W, M = N * HW;
+
+ // GPU fast-path
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ if (mc && grad_out->dev == UT_METAL) {
+ ut_sync_gpu(grad_out);
+ ut_to_device(cache->xnorm, UT_METAL);
+ ut_to_device(cache->rstd, UT_METAL);
+ ut_to_device(l->weight, UT_METAL);
+ ut_tensor* dx = ut_alloc(4, grad_out->shape.shape, UT_METAL);
+ ut_tensor* dwp = ut_alloc(1, (int[]){C}, UT_METAL);
+ ut_tensor* dbp = ut_alloc(1, (int[]){C}, UT_METAL);
+ struct {
+ int N, C, HW;
+ float eps, mom;
+ } params = {N, C, HW, l->eps, l->momentum};
+ int tgsize = M < 64 ? 32 : (M < 256 ? 64 : (M < 512 ? 128 : 256));
+ if (tgsize > 1024) tgsize = 1024;
+ _mtl_dispatch_rows(mc, "bn2d_bwd",
+ (void*[]){grad_out->gpu_buf, cache->xnorm->gpu_buf, cache->rstd->gpu_buf,
+ l->weight->gpu_buf, dx->gpu_buf, dwp->gpu_buf, dbp->gpu_buf},
+ 7, ¶ms, (int)sizeof(params), C, tgsize);
+ dx->dirty_cpu = true;
+ dwp->dirty_cpu = true;
+ dbp->dirty_cpu = true;
+
+ ut_sync_cpu(dwp);
+ ut_sync_cpu(dbp);
+ ut_sync_cpu(dW);
+ ut_sync_cpu(db);
+ for (int c = 0; c < C; c++) {
+ dW->data[c] += dwp->data[c];
+ db->data[c] += dbp->data[c];
+ }
+ dW->dirty_gpu = true;
+ db->dirty_gpu = true;
+ ut_free_all(dwp, dbp);
+ return dx;
+ }
+
+ // CPU fallback (dW/db and dx together, same as before)
ut_sync_cpu(grad_out);
ut_sync_cpu(cache->xnorm);
+ ut_sync_cpu(cache->rstd);
+ ut_sync_cpu(l->weight);
ut_sync_cpu(dW);
ut_sync_cpu(db);
- ut_sync_cpu(l->weight);
-
- ut_tensor* dx = ut_alloc(4, grad_out->shape.shape, grad_out->dev);
+ ut_tensor* dx = ut_alloc(4, grad_out->shape.shape, UT_CPU);
for (int c = 0; c < C; c++) {
float sum_go = 0, sum_go_xn = 0;
for (int n = 0; n < N; n++)
- for (int h = 0; h < H; h++)
- for (int w = 0; w < W; w++) {
- int idx = ((n * C + c) * H + h) * W + w;
- sum_go += grad_out->data[idx];
- sum_go_xn += grad_out->data[idx] * cache->xnorm->data[idx];
- }
+ for (int hw = 0; hw < HW; hw++) {
+ int idx = (n * C + c) * HW + hw;
+ sum_go += grad_out->data[idx];
+ sum_go_xn += grad_out->data[idx] * cache->xnorm->data[idx];
+ }
dW->data[c] += sum_go_xn;
db->data[c] += sum_go;
-
float rs = cache->rstd->data[c], wgt = l->weight->data[c];
for (int n = 0; n < N; n++)
- for (int h = 0; h < H; h++)
- for (int w = 0; w < W; w++) {
- int idx = ((n * C + c) * H + h) * W + w;
- float go = grad_out->data[idx], xn = cache->xnorm->data[idx];
- dx->data[idx] = rs * wgt * (go - (sum_go + xn * sum_go_xn) / (float)M);
- }
+ for (int hw = 0; hw < HW; hw++) {
+ int idx = (n * C + c) * HW + hw;
+ float go = grad_out->data[idx], xn = cache->xnorm->data[idx];
+ dx->data[idx] = rs * wgt * (go - (sum_go + xn * sum_go_xn) / (float)M);
+ }
}
dW->dirty_gpu = true;
db->dirty_gpu = true;
@@ -1631,6 +1807,19 @@ void ut_batchnorm2d_free(ut_batchnorm2d* l) {
// only needs H,W, which the caller already has from x's own shape.
ut_tensor* ut_global_avgpool2d(ut_tensor* x) {
int N = x->shape.shape[0], C = x->shape.shape[1], HW = x->shape.shape[2] * x->shape.shape[3];
+
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ if (mc && x->dev == UT_METAL) {
+ ut_sync_gpu(x);
+ ut_tensor* out = ut_alloc(2, (int[]){N, C}, UT_METAL);
+ int tgsize = HW < 64 ? 32 : (HW < 256 ? 64 : (HW < 512 ? 128 : 256));
+ if (tgsize > 1024) tgsize = 1024;
+ _mtl_dispatch_rows(mc, "gap_fwd", (void*[]){x->gpu_buf, out->gpu_buf}, 2, &HW, sizeof(int),
+ N * C, tgsize);
+ out->dirty_cpu = true;
+ return out;
+ }
+
ut_sync_cpu(x);
ut_tensor* out = ut_alloc(2, (int[]){N, C}, x->dev);
for (int n = 0; n < N; n++)
@@ -1645,6 +1834,20 @@ ut_tensor* ut_global_avgpool2d(ut_tensor* x) {
ut_tensor* ut_global_avgpool2d_backward(ut_tensor* grad_out, int H, int W) {
int N = grad_out->shape.shape[0], C = grad_out->shape.shape[1], HW = H * W;
+
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ if (mc && grad_out->dev == UT_METAL) {
+ ut_sync_gpu(grad_out);
+ ut_tensor* dx = ut_alloc(4, (int[]){N, C, H, W}, UT_METAL);
+ struct {
+ int HW, tot;
+ } params = {HW, N * C * HW};
+ _mtl_dispatch(mc, "gap_bwd", (void*[]){grad_out->gpu_buf, dx->gpu_buf}, 2, ¶ms,
+ (int)sizeof(params), params.tot);
+ dx->dirty_cpu = true;
+ return dx;
+ }
+
ut_sync_cpu(grad_out);
ut_tensor* dx = ut_alloc(4, (int[]){N, C, H, W}, grad_out->dev);
for (int n = 0; n < N; n++)