Module layout 

Structs

BiasLayout

BiasLayoutCompilationArg

BiasLayoutExpand

BiasLayoutLaunch

Im2colLayout

Maps a 4D NHWC tensor to a 2D column matrix using the im2col transformation It first decomposes the (m, k) matrix into ((n, out_h, out_w), (k_h, k_w, c)), then applies the convolution parameters to calculate the position in the input tensor for that kernel element.

Im2colLayoutCompilationArg

Im2colLayoutExpand

Im2colLayoutLaunch

NhwcCoords

NhwcCoordsExpand

NhwcLayout

Layout for a spatial (i.e. NHWC) tensor. Bounds check only applies to spatial dimensions, not channel or batch (because these are implicitly checked in the layouts used with spatial tensors).

NhwcLayoutCompilationArg

NhwcLayoutExpand

NhwcLayoutLaunch

OutLayout

Maps a 4D NHWC out tensor of shape ((n, h, w), c) to a col-major 2D matmul tile with shape (m, n)

OutLayoutCompilationArg

OutLayoutExpand

OutLayoutLaunch

WeightLayout

Maps a 4D weight tensor of shape (out_c, (k_h, k_w, in_c)) to a col-major 2D matmul tile with shape (n, k)

WeightLayoutCompilationArg

WeightLayoutExpand

WeightLayoutLaunch