Skip to main content

Mojo function

col_major

col_major(var shape: Coord[shape.element_types]) -> Layout[shape.element_types, #kgen.variadic.reduce(shape.element_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[1] if (idx == 0)._mlir_value else RuntimeInt[VA[(add idx, -1)].DTYPE if VA[(add idx, -1)].is_static_value.__bool__().__invert__()._mlir_value else PrevV[(add idx, -1)].DTYPE] if VA[(add idx, -1)].is_static_value.__bool__().__invert__() if VA[(add idx, -1)].is_static_value.__bool__().__invert__()._mlir_value else PrevV[(add idx, -1)].is_static_value.__bool__().__invert__() else ComptimeInt[(VA[(add idx, -1)].static_value * PrevV[(add idx, -1)].static_value)]))]

Create a column-major layout from a shape.

Column-major means the first dimension has stride 1, and each subsequent dimension has stride equal to the product of all previous dimensions.

For shape (M, N, K):

  • row_major strides: (N*K, K, 1)
  • col_major strides: (1, M, M*N)

Args:

  • ​shape (Coord): The shape as a Coord.

Returns:

Layout: A Layout with column-major strides.

col_major[*idxs: Int]() -> Layout[#kgen.variadic.reduce(idxs, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[Int], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[VA[idx]])), #kgen.variadic.reduce(#kgen.variadic.reduce(idxs, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[Int], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[VA[idx]])), base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[1] if (idx == 0)._mlir_value else RuntimeInt[VA[(add idx, -1)].DTYPE if VA[(add idx, -1)].is_static_value.__bool__().__invert__()._mlir_value else PrevV[(add idx, -1)].DTYPE] if VA[(add idx, -1)].is_static_value.__bool__().__invert__() if VA[(add idx, -1)].is_static_value.__bool__().__invert__()._mlir_value else PrevV[(add idx, -1)].is_static_value.__bool__().__invert__() else ComptimeInt[(VA[(add idx, -1)].static_value * PrevV[(add idx, -1)].static_value)]))]

Create a column-major layout from compile-time shape dimensions.

Example:

from layout.tile_layout import col_major

var layout = col_major[3, 4]()
# shape: (3, 4), stride: (1, 3)

Parameters:

  • ​*idxs (Int): The shape dimensions as compile-time integers.

Returns:

Layout: A Layout with column-major strides.

col_major(idx: ComptimeInt[idx.val]) -> Layout[ComptimeInt[idx.val], ComptimeInt[1]]

Creates a 1D column-major layout from a compile-time dimension.

Args:

  • ​idx (ComptimeInt): The shape dimension as a ComptimeInt.

Returns:

Layout: A 1D Layout with stride 1.

col_major(idx: RuntimeInt[idx.dtype]) -> Layout[RuntimeInt[idx.dtype], ComptimeInt[1]]

Creates a 1D column-major layout from a runtime dimension.

Args:

  • ​idx (RuntimeInt): The shape dimension as a RuntimeInt.

Returns:

Layout: A 1D Layout with stride 1.

Was this page helpful?