For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

Conv2dProblemShape

struct Conv2dProblemShape

Defines 2D convolution problem geometry.

Layouts:

• Activation: NHWC (batch, height, width, channels)
• Filter: KRSC (output_channels, filter_h, filter_s, input_channels)
• Output: NHWC (batch, out_height, out_width, output_channels)

For Fprop with stride=1, no dilation, this maps to GEMM as:

• M = N * H_out * W_out
• N = K (output channels)
• K = C * R * S (input channels * filter area)

Fields

• ​batch (Int):
• ​in_height (Int):
• ​in_width (Int):
• ​in_channels (Int):
• ​out_channels (Int):
• ​filter_h (Int):
• ​filter_w (Int):
• ​pad_h (Int):
• ​pad_w (Int):
• ​stride_h (Int):
• ​stride_w (Int):
• ​dilation_h (Int):
• ​dilation_w (Int):
• ​groups (Int):

Implemented traits

AnyType, Copyable, Deinitable, Movable

Methods

__init__

def __init__(out self, batch: Int, in_height: Int, in_width: Int, in_channels: Int, out_channels: Int, filter_h: Int, filter_w: Int, pad_h: Int = Int(0), pad_w: Int = Int(0), stride_h: Int = Int(1), stride_w: Int = Int(1), dilation_h: Int = Int(1), dilation_w: Int = Int(1), groups: Int = Int(1))

out_height

def out_height(self) -> Int

Compute output height.

Returns:

Int

out_width

def out_width(self) -> Int

Compute output width.

Returns:

Int

gemm_m

def gemm_m(self) -> Int

GEMM M dimension = batch * output_height * output_width.

Returns:

Int

gemm_n

def gemm_n(self) -> Int

GEMM N dimension = output_channels.

Returns:

Int

gemm_k

def gemm_k(self) -> Int

GEMM K dimension = input_channels * filter_height * filter_width.

Returns:

Int

num_m_tiles

def num_m_tiles(self, tile_m: Int) -> Int

Number of tiles in M dimension.

Args:

• ​tile_m (Int): Size of one tile along the M dimension, in elements. The M dimension equals batch * out_height * out_width.

Returns:

Int

num_n_tiles

def num_n_tiles(self, tile_n: Int) -> Int

Number of tiles in N dimension.

Args:

• ​tile_n (Int): Size of one tile along the N dimension, in elements. The N dimension equals out_channels.

Returns:

Int

num_k_tiles

def num_k_tiles(self, tile_k: Int) -> Int

Number of tiles in K dimension.

Args:

• ​tile_k (Int): Size of one tile along the K dimension, in elements. The K dimension equals in_channels * filter_h * filter_w.

Returns:

Int