nx_winograd_conv.nx
buildroot/runtime/nx_winograd_conv.nx
about
nx_winograd_conv.nx -- 2D convolution via Winograd F(2x2, 3x3).
Algorithm-led mult-reduction kernel. For every 3x3 conv kernel
producing a 2x2 output tile:
Direct convolution: 2 * 2 * 3 * 3 = 36 multiplies
Winograd F(2x2, 3x3): 4 * 4 = 16 multiplies
Reduction: 36 / 16 = 2.25x (Lavin & Gray 2016)
More additions (transform matrices) -- but adds are 1-cycle ALU
on every CPU including MCUs; mults take 3-5. Net win on every
hardware floor we target.
Per the min-hardware-floor + algo-led cardinal: this brick beats
direct conv on a $5 microcontroller as much as on a 5090, without
any SIMD / GPU dependency. When SIMD lands (compiler I3.x) the
16 mults become 4 SIMD ops, compounding the win.
Reference output composes via nx_compute_node as NX_CN_K_CONV2D
kernel kind.
Winograd transform matrices (F(2x2, 3x3), Lavin & Gray 2016):
B^T = [[1, 0, -1, 0],
[0, 1, 1, 0],
[0, -1, 1, 0],
[0, 1, 0, -1]]
G = [[1, 0, 0],
[1/2, 1/2, 1/2],
[1/2, -1/2, 1/2],
[0, 0, 1]]
A^T = [[1, 1, 1, 0],
[0, 1, -1, -1]]
In Q-format on i64 substrate:
* Input tile (4x4) is just i64 values
* G filter transform: G * g * G^T -> 4x4 transformed kernel
* B input transform: B^T * d * B -> 4x4 transformed input
dependencies 3 imports · 1 importers
imports: nx_syscalls.nxnx_tier.nxnx_tensor.nx
imported by: nx_winograd_conv_test.nx
structs
| none |
consts
| 61 | const NX_WG_Q10: nx_int = 1024 |
| 65 | const NX_WG_OK: nx_int = 0 |
| 66 | const NX_WG_ERR_BAD_DTYPE: nx_int = 1 |
| 67 | const NX_WG_ERR_BAD_KERNEL_SIZE: nx_int = 2 // not 3x3 (v1 only supports F(2x2,3x3)) |
| 68 | const NX_WG_ERR_BAD_INPUT_SHAPE: nx_int = 3 // input dims not 4x4 per tile |
| 69 | const NX_WG_ERR_NOT_CONTIGUOUS: nx_int = 4 |
| 70 | const NX_WG_N_VERDICTS: nx_int = 5 |
| 304 | const NX_WG_DIRECT_MULTS: nx_int = 36 |
| 305 | const NX_WG_WINOGRAD_MULTS: nx_int = 16 |
functions
| 72 | func nx_wg_verdict_is_valid(v: nx_int) -> nx_int called by 1: main |
| 105 | func nx_wg_filter_transform(g: *i64, U: *i64) -> nx_int |
| 158 | func nx_wg_input_transform(d: *i64, V: *i64) -> nx_int |
| 194 | func nx_wg_hadamard_4x4(U: *i64, V: *i64, M: *i64) -> nx_int called by 1: nx_wg_conv_tile |
| 216 | func nx_wg_output_transform(M: *i64, Y: *i64) -> nx_int |
| 248 | func nx_wg_conv_tile(filter: *i64, input_tile: *i64, output_tile: *i64) -> nx_int |
| 274 | func nx_wg_direct_conv_reference(filter: *i64, input_tile: *i64, called by 1: main |
| 307 | func nx_wg_mult_reduction_q10() -> nx_int called by 1: main |