Skip to main content

/

DocsCode

Mojo struct

VBufferTransposeLoads

struct VBufferTransposeLoads[dtype: DType, layout: Layout, address_space: AddressSpace, alignment: Int, origin: Origin[mut=origin.mut], masked: Bool, layout_int_type: DType, linear_idx_type: DType, //, tensor_core_mma: TiledTensorCore[tensor_core_mma.out_type, tensor_core_mma.in_type, tensor_core_mma.shape, tensor_core_mma.group_size, tensor_core_mma.transpose_b], BN: Int, BK: Int, depth: Int, num_threads: Int, num_stages: Int = 1]

Fields​

  • ​load_tile (VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].LoadTileType):
  • ​mma_tile (VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].MMATileType):
  • ​smem_ptr (UnsafePointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]):
  • ​global_iterator (VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].GlobalTiledIteratorType):
  • ​global_base_tile (VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].GlobalTensorType):
  • ​current_stage (Int):

Implemented traits​

AnyType, ImplicitlyDestructible, KVBuffer

comptime members​

base_layout​

comptime base_layout = Layout.row_major(VariadicList(VBufferTransposeLoads.pad[depth](), VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].simd_width))

depth_tile_size​

comptime depth_tile_size = min(depth, 128)

GlobalTensorType​

comptime 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]

GlobalTiledIteratorType​

comptime GlobalTiledIteratorType = LayoutTensorIter[dtype, LayoutTensor._compute_tile_layout[BK, depth]()[0], origin, address_space=address_space, axis=0, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked or _tile_is_masked[layout, BK, depth]()]

load_width​

comptime load_width = 4 if (depth == 64) else VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].simd_width

loads_per_thread_per_depth_tile​

comptime loads_per_thread_per_depth_tile = ((VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].depth_tile_size * BK) // (VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].load_width * num_threads))

LoadTileType​

comptime LoadTileType = LayoutTensor[dtype, Layout.row_major(VariadicList(((VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].loads_per_thread_per_depth_tile * (depth // VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].depth_tile_size)) * num_stages), VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].load_width)), MutAnyOrigin, address_space=AddressSpace.LOCAL]

MMA_K​

comptime MMA_K = shape.__getitem__[Int](2)

MMA_M​

comptime MMA_M = shape.__getitem__[Int](0)

mma_tile_layout​

comptime mma_tile_layout = Layout.row_major(VariadicList((depth // VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].MMA_M), VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].simd_width))

MMATileType​

comptime MMATileType = LayoutTensor[dtype, VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].mma_tile_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]

num_depth_tiles​

comptime num_depth_tiles = (depth // VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].MMA_M)

num_k_tiles​

comptime num_k_tiles = ceildiv(BK, (VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].MMA_K * group_size))

num_repeats​

comptime num_repeats = (BK // VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].simd_width)

SharedTileType​

comptime SharedTileType = LayoutTensor[dtype, VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].smem_layout, MutAnyOrigin, address_space=AddressSpace.SHARED]

simd_width​

comptime simd_width = simd_width_of[dtype]()

smem_layout​

comptime smem_layout = blocked_product(VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].base_layout, VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].tiler_layout, True)

tiler_layout​

comptime tiler_layout = Layout.row_major(VariadicList(1, VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].num_repeats))

Methods​

__init__​

__init__(out self, global_tile: LayoutTensor[dtype, layout, origin, address_space=address_space, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], shared_ptr: UnsafePointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED])

get_dtype​

static get_dtype() -> DType

Returns:

DType

pad​

static pad[dim: Int]() -> Int

Returns:

Int

load_from_dram​

load_from_dram(mut self)

get_mma_tile​

get_mma_tile(self) -> VBufferTransposeLoads[tensor_core_mma, BN, BK, depth, num_threads, num_stages].MMATileType

Returns:

VBufferTransposeLoads

copy_to_shared​

copy_to_shared[tile_id: Int = 0](self)

load_from_shared​

load_from_shared[k_mma: Int](self)

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact