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

NhwcCoordsCompilationArg

NhwcCoordsExpand

NhwcCoordsLaunch

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

TmaIm2colLayout

Im2col layout, producing both the position and offset

TmaIm2colLayoutCompilationArg

TmaIm2colLayoutExpand

TmaIm2colLayoutLaunch

TmaOutGradLayout

Weight backwards needs a consolidated layout to work properly across the combined k dimension. Padding to an even tile shape on width isn’t valid, because im2col doesn’t do this. Wouldn’t be necessary with im2colWide, should investigate at some point.

TmaOutGradLayoutCompilationArg

TmaOutGradLayoutExpand

TmaOutGradLayoutLaunch

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

Enums

NhwcCheck