Numerical kernels
Raw-pointer APIs, the dispatch table, and structural operations.
01Low-level contract#
nexus_model::kernels operates on row-major float32 pointers without Tensor ownership or lifetime management. Callers allocate correctly sized output, workspace, and gradient buffers. Dimensions mostly use int or size_t; callers must prevent oversized dimension overflows.
02KernelTable#
cpu() returns a const reference to the selected function table; wrappers dispatch through it. set_force_scalar rebuilds the table. This raw API is close to implementation details; prefer Module for general usage.
| Family | Operations |
|---|---|
| Vector | fill, add, mul, scale, axpy, dot |
| Activation | relu, leaky_relu, sigmoid, tanh_fwd, gelu, gelu_tanh, silu, elu, mish |
| Activation backward | corresponding *_bwd functions |
| Row reductions | softmax, log_softmax, softmax_bwd, log_softmax_bwd |
| Dense | linear_forward, linear_backward |
| Normalization | layernorm, layernorm_bwd |
03Accumulation and caches#
linear_backward adds into dX, dW, and db without clearing. LayerNorm backward uses forward mean and reciprocal-std buffers. Sigmoid/Tanh derivatives use output y; ReLU uses input x. Passing the wrong cache gives wrong results even when shapes match.
04Structural kernels#
ConvDesc contains n/c/k, three-element spatial_in/out, kernel, stride, pad, dilation arrays, groups=1, and rank=2. conv_col_size determines im2col workspace size. Conv may use SIMD GEMM internally; structural loops are separate kernels.
| Family | Operations |
|---|---|
| Normalization | batch_norm / batch_norm_bwd / group_norm / group_norm_bwd |
| Convolution | conv_out_size / conv_col_size / conv_forward / conv_backward |
| Max pooling | max_pool2d / max_pool2d_bwd; saved int indices |
| Average pooling | avg_pool2d / avg_pool2d_bwd; count_include_pad |
| Embedding | embedding_forward / embedding_backward; float indices |
| Dropout | dropout_mask(mask,n,p,seed) |