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 function
col_major
def col_major[*element_types: CoordLike](var *elements: *element_types.values) -> Layout[element_types, TypeList[#kgen.param_list.reduce(element_types.values[0]._ParamListType if element_types.values[0].is_tuple if identical(len(element_types.values), 1) else identical(len(element_types.values), 1) else element_types.values, base=, reducer=[PrevV: KGENParamList[CoordLike], VA: KGENParamList[CoordLike], idx: __mlir_type.index] #kgen.param_list.concat(PrevV, ComptimeInt[Int(1)] if identical(idx, 0) else Scalar[element_types.values[0]._ParamListType if element_types.values[0].is_tuple if identical(len(element_types.values), 1) else identical(len(element_types.values), 1) else element_types.values[(add idx, -1)].DTYPE if (xor element_types.values[0]._ParamListType if element_types.values[0].is_tuple if identical(len(element_types.values), 1) else identical(len(element_types.values), 1) else element_types.values[(add idx, -1)].is_static_value, True) else PrevV[(add idx, -1)].DTYPE] if (xor element_types.values[0]._ParamListType if element_types.values[0].is_tuple if identical(len(element_types.values), 1) else identical(len(element_types.values), 1) else element_types.values[(add idx, -1)].is_static_value, True) if (xor element_types.values[0]._ParamListType if element_types.values[0].is_tuple if identical(len(element_types.values), 1) else identical(len(element_types.values), 1) else element_types.values[(add idx, -1)].is_static_value, True) else (xor PrevV[(add idx, -1)].is_static_value, True) else ComptimeInt[Int((mul element_types.values[0]._ParamListType if element_types.values[0].is_tuple if identical(len(element_types.values), 1) else identical(len(element_types.values), 1) else element_types.values[(add idx, -1)].static_value, PrevV[(add idx, -1)].static_value))]))]()]
Create a column-major layout from variadic arguments.
Column-major means the first dimension has stride 1, and each subsequent dimension has stride equal to the product of all previous dimensions.
Parameters:
- *element_types (
CoordLike): The variadic pack of element types that implementCoordLike.
Args:
- *elements (
*element_types.values): The shape dimensions.
Returns:
def col_major(var shape: Coord) -> Layout[shape.element_types, TypeList[#kgen.param_list.reduce(shape.element_types.values[0]._ParamListType if shape.element_types.values[0].is_tuple if identical(len(shape.element_types.values), 1) else identical(len(shape.element_types.values), 1) else shape.element_types.values, base=, reducer=[PrevV: KGENParamList[CoordLike], VA: KGENParamList[CoordLike], idx: __mlir_type.index] #kgen.param_list.concat(PrevV, ComptimeInt[Int(1)] if identical(idx, 0) else Scalar[shape.element_types.values[0]._ParamListType if shape.element_types.values[0].is_tuple if identical(len(shape.element_types.values), 1) else identical(len(shape.element_types.values), 1) else shape.element_types.values[(add idx, -1)].DTYPE if (xor shape.element_types.values[0]._ParamListType if shape.element_types.values[0].is_tuple if identical(len(shape.element_types.values), 1) else identical(len(shape.element_types.values), 1) else shape.element_types.values[(add idx, -1)].is_static_value, True) else PrevV[(add idx, -1)].DTYPE] if (xor shape.element_types.values[0]._ParamListType if shape.element_types.values[0].is_tuple if identical(len(shape.element_types.values), 1) else identical(len(shape.element_types.values), 1) else shape.element_types.values[(add idx, -1)].is_static_value, True) if (xor shape.element_types.values[0]._ParamListType if shape.element_types.values[0].is_tuple if identical(len(shape.element_types.values), 1) else identical(len(shape.element_types.values), 1) else shape.element_types.values[(add idx, -1)].is_static_value, True) else (xor PrevV[(add idx, -1)].is_static_value, True) else ComptimeInt[Int((mul shape.element_types.values[0]._ParamListType if shape.element_types.values[0].is_tuple if identical(len(shape.element_types.values), 1) else identical(len(shape.element_types.values), 1) else shape.element_types.values[(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:
def col_major[*idxs: Int]() -> Layout[TypeList[#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])](), TypeList[#kgen.param_list.reduce(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0]._ParamListType if #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0].is_tuple if identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]]), base=, reducer=[PrevV: KGENParamList[CoordLike], VA: KGENParamList[CoordLike], idx: __mlir_type.index] #kgen.param_list.concat(PrevV, ComptimeInt[Int(1)] if identical(idx, 0) else Scalar[#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0]._ParamListType if #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0].is_tuple if identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[(add idx, -1)].DTYPE if (xor #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0]._ParamListType if #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0].is_tuple if identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[(add idx, -1)].is_static_value, True) else PrevV[(add idx, -1)].DTYPE] if (xor #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0]._ParamListType if #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0].is_tuple if identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[(add idx, -1)].is_static_value, True) if (xor #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0]._ParamListType if #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0].is_tuple if identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[(add idx, -1)].is_static_value, True) else (xor PrevV[(add idx, -1)].is_static_value, True) else ComptimeInt[Int((mul #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0]._ParamListType if #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[0].is_tuple if identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else identical(len(#kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])), 1) else #kgen.param_list.tabulate(len(idxs.values), [idx: __mlir_type.index] ComptimeInt[idxs.values[idx]])[(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:
def col_major(idx: ComptimeInt) -> Layout[TypeList[ComptimeInt[idx.val]](), TypeList[ComptimeInt[Int(1)]]()]
Creates a 1D column-major layout from a compile-time dimension.
Args:
- idx (
ComptimeInt): The shape dimension as aComptimeInt.
Returns:
Layout[TypeList[ComptimeInt[idx.val]](), TypeList[ComptimeInt[Int(1)]]()]: A 1D Layout with stride 1.
def col_major(idx: Scalar) -> Layout[TypeList[Scalar[idx.dtype]](), TypeList[ComptimeInt[Int(1)]]()]
Creates a 1D column-major layout from a runtime dimension.
Args:
- idx (
Scalar): The shape dimension as aScalar.
Returns:
Layout[TypeList[Scalar[idx.dtype]](), TypeList[ComptimeInt[Int(1)]]()]: A 1D Layout with stride 1.