nx_f32_conv2d_grouped.nx source

1// nx_f32_conv2d_grouped.nx -- software-f32 GROUPED 2D convolution (NCHW), generalizing nx_f32_conv2d with a 2// `groups` parameter. groups=1 is the dense conv; groups=C_in (with C_out a multiple) is DEPTHWISE -- the op a 3// MobileNet/EfficientNet-class pose backbone is almost entirely built from, and the single biggest gap the 4// reconnaissance found (nx_f32_conv2d has no groups; the graph-IR has a DEPTHWISE kind with no kernel behind it). 5// Input C_in and output C_out are each split into `groups` contiguous groups; output channel co (in group 6// g=co/(C_out/groups)) convolves ONLY the input channels of group g. Weight layout matches PyTorch grouped conv: 7// [C_out, C_in/groups, KH, KW]. Composes ONLY nx_f32_mul / nx_f32_add. license_tier: ORIGINAL 8import "nx_syscalls.nx" 9import "nx_f32.nx" 10 11const NX_GCV_OK: i64 = 0 12const NX_GCV_ERR_ARGS: i64 = 4 13 14func nx_f32_conv2d_grouped(input: *i64, N: i64, C_in: i64, H: i64, W: i64, 15 weight: *i64, C_out: i64, KH: i64, KW: i64, 16 stride: i64, pad: i64, groups: i64, 17 bias: *i64, out: *i64) -> i64 { 18 if stride <= 0 { return NX_GCV_ERR_ARGS } 19 if groups <= 0 { return NX_GCV_ERR_ARGS } 20 if KH <= 0 { return NX_GCV_ERR_ARGS } 21 if KW <= 0 { return NX_GCV_ERR_ARGS } 22 if C_in - (C_in / groups) * groups != 0 { return NX_GCV_ERR_ARGS } // C_in % groups != 0 23 if C_out - (C_out / groups) * groups != 0 { return NX_GCV_ERR_ARGS } // C_out % groups != 0 24 let OH: i64 = (H + 2 * pad - KH) / stride + 1 25 let OW: i64 = (W + 2 * pad - KW) / stride + 1 26 if OH <= 0 { return NX_GCV_ERR_ARGS } 27 if OW <= 0 { return NX_GCV_ERR_ARGS } 28 let cig: i64 = C_in / groups // input channels per group 29 let cog: i64 = C_out / groups // output channels per group 30 let in_cs: i64 = H * W 31 let in_bs: i64 = C_in * H * W 32 let wt_cs: i64 = KH * KW 33 let wt_os: i64 = cig * KH * KW 34 let out_cs: i64 = OH * OW 35 let out_bs: i64 = C_out * OH * OW 36 var n: i64 = 0 37 while n < N { 38 var co: i64 = 0 39 while co < C_out { 40 let g: i64 = co / cog 41 let ci_base: i64 = g * cig 42 var bv: i64 = 0 43 if (bias as i64) != 0 { bv = bias[co] } 44 var oh: i64 = 0 45 while oh < OH { 46 var ow: i64 = 0 47 while ow < OW { 48 var acc: i64 = bv 49 var cl: i64 = 0 50 while cl < cig { 51 let ci: i64 = ci_base + cl 52 var kh: i64 = 0 53 while kh < KH { 54 let ih: i64 = oh * stride + kh - pad 55 if ih >= 0 { if ih < H { 56 var kw: i64 = 0 57 while kw < KW { 58 let iw: i64 = ow * stride + kw - pad 59 if iw >= 0 { if iw < W { 60 let in_idx: i64 = n * in_bs + ci * in_cs + ih * W + iw 61 let wt_idx: i64 = co * wt_os + cl * wt_cs + kh * KW + kw 62 acc = nx_f32_add(acc, nx_f32_mul(input[in_idx], weight[wt_idx])) 63 } } 64 kw = kw + 1 65 } 66 } } 67 kh = kh + 1 68 } 69 cl = cl + 1 70 } 71 out[n * out_bs + co * out_cs + oh * OW + ow] = acc 72 ow = ow + 1 73 } 74 oh = oh + 1 75 } 76 co = co + 1 77 } 78 n = n + 1 79 } 80 return NX_GCV_OK 81}