IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

TMemEngine

struct TMemEngine

Implements TensorEngine over Tensor Memory via tcgen05.ld/st.

Element offsets are positions on the lane-by-column grid, so layouts over this engine use a lane stride of 1 and a column stride of TMEM_NUM_LANES; the engine encodes them into hardware addresses. It has no unsafe_ptr, so its only data path is copy_from in either direction, which moves consecutive columns of the lane the calling thread owns. A copy issues one tcgen05 instruction per power-of-two chunk of columns, largest first, none wider than 64 registers, and waits once per 64-column slice; the _async variants issue everything and leave the wait to the caller.

The tcgen05 access shape is not stored on the engine: its lane count is always 32, one thread per lane, and its bits per lane are the element type's width, so a float32 tile moves in the 32x32b shape. Only 4-byte element types are supported; two-byte types would need the pack::16b path.

Implemented traits​

AnyType, Deinitable, Movable, TensorEngine

comptime members​

element_size​

comptime element_size = 1

One scalar per logical element.

OffsetResultType​

comptime OffsetResultType[offset_types: TypeList[values]] = TMemEngine

Offsetting never changes the engine, so this is Self.

Parameters​

  • ​offset_types (TypeList[values]): The coordinate element types of the applied offset.

StorageType​

comptime StorageType[mut: Bool, //, dtype: DType, origin: Origin[mut=mut], address_space: AddressSpace] = TMemStorage[dtype, origin]

The encoded TMEM address handle.

The address_space is part of the TensorEngine interface but unused: TMEM has no pointer, so no AddressSpace describes it and the handle does not carry one.

Parameters​

  • ​mut (Bool): The mutability of the viewed storage, inferred from origin.
  • ​dtype (DType): The element data type of the viewed storage.
  • ​origin (Origin[mut=mut]): The origin tracking the lifetime of the allocation.
  • ​address_space (AddressSpace): Unused.

Methods​

write_type_name_to​

static def write_type_name_to(mut writer: T)

Writes the engine type name representation to the writer.

Args:

  • ​writer (T): The Writer to output to.

unsafe_cast​

static def unsafe_cast[to_mut: Bool, //, to_dtype: DType, to_origin: Origin[mut=to_mut], to_address_space: AddressSpace](storage: TMemStorage) -> TMemStorage[to_dtype, to_origin]

Reinterprets the handle with new type parameters.

The address is untyped, so this only re-labels the handle.

Parameters:

  • ​to_mut (Bool): The mutability of the new origin.
  • ​to_dtype (DType): The element data type to view the cells as.
  • ​to_origin (Origin[mut=to_mut]): The origin to reinterpret the storage as.
  • ​to_address_space (AddressSpace): The address space parameter to carry.

Args:

  • ​storage (TMemStorage): The storage to reinterpret.

Returns:

TMemStorage[to_dtype, to_origin]: A handle with the same address and the new parameters.

offset​

static def offset[offset_mut: Bool, offset_types: TypeList[offset_types.values], //, offset_dtype: DType, offset_origin: Origin[mut=offset_mut], offset_address_space: AddressSpace](var storage: TMemStorage[offset_dtype, offset_origin], var offset_coord: Coord[offset_types]) -> TMemStorage[offset_dtype, offset_origin]

Returns a handle advanced by a number of grid positions.

The offset is a flat index on the lane-by-column grid, as a layout with column stride TMEM_NUM_LANES produces; the engine re-encodes the resulting cell as a hardware address.

Parameters:

  • ​offset_mut (Bool): The mutability of the storage, inferred from offset_origin.
  • ​offset_types (TypeList[offset_types.values]): The coordinate element types of offset_coord.
  • ​offset_dtype (DType): The element data type of the storage.
  • ​offset_origin (Origin[mut=offset_mut]): The origin tracking the lifetime of the storage.
  • ​offset_address_space (AddressSpace): The address space parameter of the storage.

Args:

Returns:

TMemStorage[offset_dtype, offset_origin]: A handle at the cell offset positions after storage's.

distance​

static def distance[dtype: DType, address_space: AddressSpace, //](storage: TMemStorage[dtype], other: TMemStorage[dtype]) -> Int

Returns the number of grid positions from other to storage.

Parameters:

  • ​dtype (DType): The storages' DType.
  • ​address_space (AddressSpace): The storages' address space parameter.

Args:

Returns:

Int: The signed difference of the two cells' flat grid indices.

wait_store​

static def wait_store()

Waits for every tcgen05.st this thread has issued.

Call it after one or more copy_from_async before the source tiles are modified or the stored columns are read back.

wait_load​

static def wait_load()

Waits for every tcgen05.ld this thread has issued.

Call it after one or more copy_to_async before the destination tiles are read.

copy_from​

static def copy_from[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dst_dtype: DType, src_dtype: DType, OtherEngine: TensorEngine](storage: Tuple[TMemStorage[dst_dtype, self_origin], SelfLayoutType], other: Tuple[OtherEngine.StorageType[other_mut, origin_of(other_origin), src_dtype, other_origin, other_address_space], OtherLayoutType])

Copies the elements of other into this thread's lane, in place.

Reads the source through OtherEngine.unsafe_ptr, casting to dst_dtype, then issues one tcgen05.st per column chunk. A row of up to 64 columns is one slice with a single tcgen05.wait::st; a wider row moves in 64-column slices, each waited on, so at most 64 staging registers are live at once. storage must be a warp-base handle whose layout addresses consecutive columns of one lane, such as row 0 of a warp's (32, N) tile.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the destination storage.
  • ​self_origin (MutOrigin): The origin of the destination storage.
  • ​self_address_space (AddressSpace): The address space of the destination storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the source storage.
  • ​other_mut (Bool): The mutability of the source storage.
  • ​other_origin (Origin[mut=other_mut]): The origin of the source storage.
  • ​other_address_space (AddressSpace): The address space of the source storage.
  • ​dst_dtype (DType): The element data type of the destination storage.
  • ​src_dtype (DType): The element data type of the source storage.
  • ​OtherEngine (TensorEngine): The engine of the source. Must not be TMemEngine.

Args:

copy_from_async​

static def copy_from_async[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dst_dtype: DType, src_dtype: DType, OtherEngine: TensorEngine](storage: Tuple[TMemStorage[dst_dtype, self_origin], SelfLayoutType], other: Tuple[OtherEngine.StorageType[other_mut, origin_of(other_origin), src_dtype, other_origin, other_address_space], OtherLayoutType])

Issues the stores of copy_from without waiting for them.

The caller must call wait_store before the source tile is modified or the stored columns are read back; until then the registers the stores read from are in flight. Several copy_from_async calls can share one wait. The whole row is staged at once, so a row of N columns keeps N registers live until the wait.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the destination storage.
  • ​self_origin (MutOrigin): The origin of the destination storage.
  • ​self_address_space (AddressSpace): The address space of the destination storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the source storage.
  • ​other_mut (Bool): The mutability of the source storage.
  • ​other_origin (Origin[mut=other_mut]): The origin of the source storage.
  • ​other_address_space (AddressSpace): The address space of the source storage.
  • ​dst_dtype (DType): The element data type of the destination storage.
  • ​src_dtype (DType): The element data type of the source storage.
  • ​OtherEngine (TensorEngine): The engine of the source. Must not be TMemEngine.

Args:

copy_to​

static def copy_to[SelfLayoutType: TensorLayout, self_mut: Bool, self_origin: Origin[mut=self_mut], self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_origin: MutOrigin, other_address_space: AddressSpace, //, src_dtype: DType, dst_dtype: DType, OtherEngine: TensorEngine](storage: Tuple[TMemStorage[src_dtype, self_origin], SelfLayoutType], other: Tuple[OtherEngine.StorageType[True, origin_of(other_origin), dst_dtype, other_origin, other_address_space], OtherLayoutType])

Copies this thread's lane at storage into other, in place.

Issues one tcgen05.ld per column chunk, then writes the columns through OtherEngine.unsafe_ptr, casting to dst_dtype. A row of up to 64 columns is one slice with a single tcgen05.wait::ld; a wider row moves in 64-column slices, each loaded, waited on and written before the next, so at most 64 staging registers are live at once. storage must be a warp-base handle whose layout addresses consecutive columns of one lane, such as row 0 of a warp's (32, N) tile.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the source storage.
  • ​self_mut (Bool): The mutability of the source storage.
  • ​self_origin (Origin[mut=self_mut]): The origin of the source storage.
  • ​self_address_space (AddressSpace): The address space of the source storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the destination storage.
  • ​other_origin (MutOrigin): The origin of the destination storage.
  • ​other_address_space (AddressSpace): The address space of the destination storage.
  • ​src_dtype (DType): The element data type of the source storage.
  • ​dst_dtype (DType): The element data type of the destination storage.
  • ​OtherEngine (TensorEngine): The engine of the destination. Must not be TMemEngine.

Args:

copy_to_async​

static def copy_to_async[SelfLayoutType: TensorLayout, self_mut: Bool, self_origin: Origin[mut=self_mut], self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_origin: MutOrigin, other_address_space: AddressSpace, //, src_dtype: DType, dst_dtype: DType, OtherEngine: TensorEngine](storage: Tuple[TMemStorage[src_dtype, self_origin], SelfLayoutType], other: Tuple[OtherEngine.StorageType[True, origin_of(other_origin), dst_dtype, other_origin, other_address_space], OtherLayoutType])

Issues the loads of copy_to without waiting for them.

The destination must be a register tile: the loaded registers are only valid after wait_load, so a destination that reaches memory before the wait receives stale data. Several copy_to_async calls can share one wait, and the destinations must not be read until it returns. The whole row is in flight at once, so a row of N columns keeps N registers live until the wait.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the source storage.
  • ​self_mut (Bool): The mutability of the source storage.
  • ​self_origin (Origin[mut=self_mut]): The origin of the source storage.
  • ​self_address_space (AddressSpace): The address space of the source storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the destination storage.
  • ​other_origin (MutOrigin): The origin of the destination storage.
  • ​other_address_space (AddressSpace): The address space of the destination storage.
  • ​src_dtype (DType): The element data type of the source storage.
  • ​dst_dtype (DType): The element data type of the destination storage.
  • ​OtherEngine (TensorEngine): The engine of the destination. Must not be TMemEngine.

Args:

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community