This document is relevant for: Trn2, Trn3

Conv3D Temporal Unroll Kernel API Reference#

3D convolution with temporal unrolling and column tiling for small C_out.

The companion should_use_temporal_unroll advisory check reports whether a given problem shape benefits from temporal unrolling before you call the kernel. It returns True when: D_out temporal positions fit in a single PSUM bank column-tiled, C_in is large enough to amortize filter caching, and W_out exceeds F_MAX so multiple W tiles are needed (where the baseline is slow).

Background#

The conv3d_temporal_unroll kernel performs 3D convolution with temporal unrolling and column tiling for small C_out.

API Reference#

Source code for this kernel API can be found at: conv3d_temporal_unroll.py

should_use_temporal_unroll#

nkilib.experimental.conv.should_use_temporal_unroll(C_out: int, D_out: int, C_in: int, K_d: int, W_out: int) → bool#

Advisory check: whether conv3d_temporal_unroll is applicable for this shape.

Parameters:

  • C_out (int) – Number of output channels.

  • D_out (int) – Number of output depth positions.

  • C_in (int) – Number of input channels.

  • K_d (int) – Filter depth dimension.

  • W_out (int) – Output width.

Returns:

True if temporal unroll should be used.

Return type:

nl.ndarray

conv3d_temporal_unroll#

nkilib.experimental.conv.conv3d_temporal_unroll(x_in: nl.ndarray, filters: nl.ndarray, bias: Optional[nl.ndarray] = None, stride: tuple[int, int, int] = (1, 1, 1), padding: tuple[int, int, int, int, int, int] = (0, 0, 0, 0, 0, 0), dilation: tuple[int, int, int] = (1, 1, 1), activation_fn: Optional[ActFnType] = None, lnc_shard: bool = False) → nl.ndarray#

3D convolution with temporal unrolling and column tiling for small C_out.

Parameters:

  • x_in (nl.ndarray) – [B, C_in, D, H, W], Input tensor on HBM.

  • filters (nl.ndarray) – [K_d, K_h, K_w, C_in, C_out], Filter weights on HBM.

  • bias (Optional[nl.ndarray]) – [C_out], Optional bias tensor on HBM.

  • stride (tuple[int, int, int]) – (stride_d, stride_h, stride_w), Convolution strides.

  • padding (tuple[int, int, int, int, int, int]) – (pad_d_left, pad_d_right, pad_h_top, pad_h_bottom, pad_w_left, pad_w_right), Padding for each spatial dimension.

  • dilation (tuple[int, int, int]) – (dilation_d, dilation_h, dilation_w), Dilation factors.

  • activation_fn (Optional[ActFnType]) – Optional activation function to apply after conv.

  • lnc_shard (bool) – Enable LNC sharding across neuron cores (shards on H_out).

Returns:

[B, C_out, D_out, H_out, W_out], Output tensor on HBM.

Return type:

nl.ndarray

Dimensions:

  • B: Batch size

  • C_in: Number of input channels

  • C_out: Number of output channels

  • D: Input depth

  • H: Input height

  • W: Input width

  • K_d: Filter depth

  • K_h: Filter height

  • K_w: Filter width

  • D_out: Output depth = (D + pad_d_left + pad_d_right - dilation_d * (K_d - 1) - 1) // stride_d + 1

  • H_out: Output height = (H + pad_h_top + pad_h_bottom - dilation_h * (K_h - 1) - 1) // stride_h + 1

This document is relevant for: Trn2, Trn3