This document is relevant for: Trn2, Trn3

Conv3D Backward Kernel API Reference#

3D Convolution backward pass.

Computes any subset of {dx, dw, db} for a 3D convolution forward pass. The planner selects a strategy and per-gradient compute method from the problem geometry.

Background#

The conv3d_bwd kernel computes any subset of the input, weight, and bias gradients (dx, dw, db) of a 3D convolution. A planner selects the strategy and per-gradient compute method from the problem geometry.

API Reference#

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

conv3d_bwd#

nkilib.experimental.conv.conv3d_bwd(dy: nl.NkiTensor, x_in: nl.NkiTensor, filters: nl.NkiTensor, 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), compute_dx: bool = True, compute_dw: bool = True, compute_db: bool = True) → tuple[nl.NkiTensor, ...]#

3D Convolution backward pass.

Parameters:

  • dy (nl.NkiTensor) – [B, C_out, D_out, H_out, W_out], Output gradient on HBM.

  • x_in (nl.NkiTensor) – [B, C_in, D, H, W], Forward input on HBM.

  • filters (nl.NkiTensor) – [K_d, K_h, K_w, C_in, C_out], Forward filters on HBM.

  • stride (tuple[int, int, int]) – (stride_d, stride_h, stride_w), Forward 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), Forward padding.

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

  • compute_dx (bool) – Whether to compute the input gradient dx.

  • compute_dw (bool) – Whether to compute the filter gradient dw.

  • compute_db (bool) – Whether to compute the bias gradient db.

Notes:

  • Supports the same 6-value (possibly asymmetric) padding as the forward conv3d kernel for all of dx, dw, and db.

  • The planner selects the strategy and per-gradient compute method from the problem geometry.

Dimensions:

  • B: Batch size

  • C_in: Number of input channels

This document is relevant for: Trn2, Trn3