Mojo struct

QRegisterBuffer

struct QRegisterBuffer[: Bool, dtype: DType, layout: Layout, address_space: AddressSpace, alignment: Int, origin: Origin[], masked: Bool, layout_int_type: DType, linear_idx_type: DType, //, mma_shape: IndexList[3], k_group_size: Int, WM: Int, WN: Int, BN: Int, BK: Int, depth: Int, thread_layout: Layout]

Fields

• ​gmem_tensor (LayoutTensor[dtype, layout, origin, address_space=address_space, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]):
• ​mma_tile (LayoutTensor[dtype, row_major((ceildiv[::CeilDivable](BK, (mma_shape.__getitem__[::Indexer](2) * k_group_size)) * ceildiv[::CeilDivable](WM, mma_shape.__getitem__[::Indexer](0)) * 0 if (BK == 0) else (div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">) + -1) if (((rem_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">) == 0) ^ True) & ((BK < 0) ^ (depth < 0))) else div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">)), simdwidthof[::DType,__mlir_type.!kgen.target]()), MutableAnyOrigin, address_space=AddressSpace(5)]):

Implemented traits

AnyType, UnknownDestructibility

Aliases

GlobalTensorType

alias GlobalTensorType = LayoutTensor[dtype, layout, origin, address_space=address_space, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]

MMA_K

alias MMA_K = mma_shape.__getitem__[::Indexer](2)

MMA_M

alias MMA_M = mma_shape.__getitem__[::Indexer](0)

num_k_tiles

alias num_k_tiles = ceildiv[::CeilDivable](BK, (mma_shape.__getitem__[::Indexer](2) * k_group_size))

num_mmas

alias num_mmas = ceildiv[::CeilDivable](WM, mma_shape.__getitem__[::Indexer](0))

num_tiles

alias num_tiles = 0 if (BK == 0) else (div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">) + -1) if (((rem_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">) == 0) ^ True) & ((BK < 0) ^ (depth < 0))) else div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">)

RegisterTileType

alias RegisterTileType = LayoutTensor[dtype, row_major((ceildiv[::CeilDivable](BK, (mma_shape.__getitem__[::Indexer](2) * k_group_size)) * ceildiv[::CeilDivable](WM, mma_shape.__getitem__[::Indexer](0)) * 0 if (BK == 0) else (div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">) + -1) if (((rem_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">) == 0) ^ True) & ((BK < 0) ^ (depth < 0))) else div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">)), simdwidthof[::DType,__mlir_type.!kgen.target]()), MutableAnyOrigin, address_space=AddressSpace(5)]

simd_width

alias simd_width = simdwidthof[::DType,__mlir_type.!kgen.target]()

TiledIteratorType

alias TiledIteratorType = LayoutTensorIter[dtype, _compute_tile_layout[*::Int]()[0], MutableAnyOrigin, address_space=AddressSpace(5), axis=OptionalReg[Int]({:@stdlib::@builtin::@int::@Int {0}, 0}), layout_int_type=_get_layout_type(row_major((ceildiv[::CeilDivable](BK, (mma_shape.__getitem__[::Indexer](2) * k_group_size)) * ceildiv[::CeilDivable](WM, mma_shape.__getitem__[::Indexer](0)) * 0 if (BK == 0) else (div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">) + -1) if (((rem_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">) == 0) ^ True) & ((BK < 0) ^ (depth < 0))) else div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">)), simdwidthof[::DType,__mlir_type.!kgen.target]()), AddressSpace(5)), linear_idx_type=_get_index_type(row_major((ceildiv[::CeilDivable](BK, (mma_shape.__getitem__[::Indexer](2) * k_group_size)) * ceildiv[::CeilDivable](WM, mma_shape.__getitem__[::Indexer](0)) * 0 if (BK == 0) else (div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">) + -1) if (((rem_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">) == 0) ^ True) & ((BK < 0) ^ (depth < 0))) else div_s(#lit.struct.extract<:@stdlib::@builtin::@int::@Int depth, "value">, #lit.struct.extract<:@stdlib::@builtin::@int::@Int cond(eq(#lit.struct.extract<:@stdlib::@builtin::@int::@Int BK, "value">, 0), {1}, BK), "value">)), simdwidthof[::DType,__mlir_type.!kgen.target]()), AddressSpace(5)), masked=_tile_is_masked[::Layout,*::Int]()]

Methods

__init__

__init__(out self, tensor: LayoutTensor[dtype, layout, origin, address_space=address_space, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment])

load_from_dram

load_from_dram(mut self)

get_mma_tile

get_mma_tile[tile_idx: Int, k_idx: Int](self) -> LayoutTensor[dtype, _compute_tile_layout[::Int,::Int]()[0], MutableAnyOrigin, address_space=AddressSpace(5)]

Returns:

LayoutTensor

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit