TensorStorage

Defines a non-owning interface for accessing tensor storage.

A conforming type describes how to access storage that is owned elsewhere. It provides a concrete StorageType handle along with static operations to load from, store to, offset into, and reinterpret values of that handle. The trait never owns the underlying memory; the handle's origin parameter tracks the lifetime and mutability of the borrowed storage.

Implemented traits

AnyType

comptime members

StorageType

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

The concrete, register-passable handle to the borrowed storage.

Every operation in this trait acts on values of this type. It is parameterized on the element dtype, the origin that tracks the lifetime and mutability of the borrowed storage, and the address_space the storage resides in, so a single conforming type describes a whole family of handles.

Parameters

• ​mut (Bool): The mutability of the borrowed storage, inferred from origin.
• ​dtype (DType): The element data type of the borrowed storage.
• ​origin (Origin[mut=mut]): The origin tracking the lifetime of the borrowed storage.
• ​address_space (AddressSpace): The address space the borrowed storage resides in.

Required methods

unsafe_cast

static def unsafe_cast[to_mut: Bool, //, to_dtype: DType, to_origin: Origin[mut=to_mut], to_address_space: AddressSpace](storage: _Self.StorageType[origin.mut, origin._mlir_origin, dtype, origin, address_space]) -> _Self.StorageType[to_mut, to_origin._mlir_origin, to_dtype, to_origin, to_address_space]

Reinterprets a storage handle with new type parameters.

This performs an unchecked reinterpretation of the underlying reference; no conversion of the stored elements takes place. The caller is responsible for ensuring the new dtype, origin, and address_space are valid for the referenced storage.

Parameters:

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

Args:

• ​storage (_Self.StorageType[origin.mut, origin._mlir_origin, dtype, origin, address_space]): The storage to reinterpret.

Returns:

_Self.StorageType[to_mut, to_origin._mlir_origin, to_dtype, to_origin, to_address_space]: A handle referring to the same storage, viewed with the new type parameters.

load

static def load[dtype: DType, //, width: SIMDSize, alignment: Int](storage: _Self.StorageType[False, origin._mlir_origin, dtype, origin, address_space]) -> SIMD[dtype, width]

Loads a SIMD value from the storage.

Parameters:

• ​width (SIMDSize): The number of elements to load.
• ​alignment (Int): The alignment guarantee for the load.

Args:

• ​storage (_Self.StorageType[False, origin._mlir_origin, dtype, origin, address_space]): The storage to load from.

Returns:

SIMD[dtype, width]: The loaded SIMD value.

static def load[dtype: DType, //, width: SIMDSize, alignment: Int](storage: _Self.StorageType[False, origin._mlir_origin, dtype, origin, address_space], offset: T) -> SIMD[dtype, width]

Loads a SIMD value at an element offset from the storage.

Parameters:

• ​width (SIMDSize): The number of elements to load.
• ​alignment (Int): The alignment guarantee for the load.

Args:

• ​storage (_Self.StorageType[False, origin._mlir_origin, dtype, origin, address_space]): The storage to load from.
• ​offset (T): The element offset to load at.

Returns:

SIMD[dtype, width]: The loaded SIMD value.

store

static def store[dtype: DType, //, alignment: Int](storage: _Self.StorageType[True, origin._mlir_origin, dtype, origin, address_space], value: SIMD[dtype])

Stores a SIMD value into the storage.

Parameters:

• ​alignment (Int): The alignment guarantee for the store.

Args:

• ​storage (_Self.StorageType[True, origin._mlir_origin, dtype, origin, address_space]): The storage to store into.
• ​value (SIMD[dtype]): The SIMD value to store.

static def store[dtype: DType, //, alignment: Int](storage: _Self.StorageType[True, origin._mlir_origin, dtype, origin, address_space], offset: T, value: SIMD[dtype])

Stores a SIMD value at an element offset in the storage.

Parameters:

• ​alignment (Int): The alignment guarantee for the store.

Args:

• ​storage (_Self.StorageType[True, origin._mlir_origin, dtype, origin, address_space]): The storage to store into.
• ​offset (T): The element offset to store at.
• ​value (SIMD[dtype]): The SIMD value to store.

offset

static def offset(storage: _Self.StorageType[origin.mut, origin._mlir_origin, dtype, origin, address_space], offset: T) -> _Self.StorageType[origin.mut, origin._mlir_origin, dtype, origin, address_space]

Returns a storage handle offset by a number of elements.

The returned handle refers to the same externally owned storage, advanced by offset elements.

Args:

• ​storage (_Self.StorageType[origin.mut, origin._mlir_origin, dtype, origin, address_space]): The storage to offset from.
• ​offset (T): The number of elements to advance the handle by.

Returns:

_Self.StorageType[origin.mut, origin._mlir_origin, dtype, origin, address_space]: A handle of the same type starting offset elements into the referenced storage.

distance

static def distance[dtype: DType, address_space: AddressSpace, //](storage: _Self.StorageType[False, origin._mlir_origin, dtype, origin, address_space], other: _Self.StorageType[False, origin._mlir_origin, dtype, origin, address_space]) -> Int

Returns the element distance from other to storage.

Parameters:

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

Args:

• ​storage (_Self.StorageType[False, origin._mlir_origin, dtype, origin, address_space]): The storage to measure the distance to.
• ​other (_Self.StorageType[False, origin._mlir_origin, dtype, origin, address_space]): The storage to measure the distance from.

Returns:

Int: The number of elements separating the two handles. The value is positive when storage is ahead of other and negative when it precedes other.