Mojo struct

ManagedTensorSlice

struct ManagedTensorSlice[mut: Bool, input: IO, dtype: DType, rank: Int, InFusion: InputFusion, OutFusion: OutputFusion, ComputeFusion: ComputeOutputFusion, //, io_spec: IOSpec[mut, input], *, static_spec: StaticTensorSpec[dtype, rank, static_spec.static_layout, InFusion, OutFusion, ComputeFusion]]

A view of a tensor that does not own the underlying allocated pointer. When the object lifetime ends it does not free the underlying pointer. Conversely, if a ManagedTensorSlice is created, it will not extend the life of the underlying pointer.

Therefore, the user must take care to keep the pointer alive until the last use of a ManagedTensorSlice instance. This class is useful for writing custom operations where memory is managed by an external runtime like in MAX's inference stack.

Fields

in_fusion (InFusion):
out_fusion (OutFusion):
compute_fusion (ComputeFusion):

Implemented traits

AnyType, Copyable, DevicePassable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, RegisterPassable, TrivialRegisterPassable, Writable

`comptime` members

`address_space`

comptime address_space = static_spec.address_space

`alignment`

comptime alignment = static_spec.alignment

`device_type`

comptime device_type = LayoutTensor[dtype, static_spec.to_layout(), MutAnyOrigin]

`exclusive`

comptime exclusive = static_spec.exclusive

Methods

`init`

__init__(ptr: UnsafePointer[Scalar[dtype]], slices: InlineArray[Slice, rank], slicer_spec: RuntimeTensorSpec[dtype, rank]) -> Self

Initializes a ManagedTensorSlice from a pointer, array of slices and tensor spec.

In general, custom operations should not create ManagedTensorSlice instances, but instead use the ones provided by the MAX inference engine.

__init__(ptr: UnsafePointer[Scalar[dtype], MutAnyOrigin], spec: RuntimeTensorSpec[dtype, rank], strides: IndexList[rank]) -> Self

Initializes a ManagedTensorSlice from a pointer, runtime tensor spec, and strides.

In general, custom operations should not create ManagedTensorSlice instances, but instead use the ones provided by the MAX inference engine.

__init__(ptr: UnsafePointer[Scalar[dtype], MutAnyOrigin], shape: IndexList[rank]) -> Self

Initializes a ManagedTensorSlice from a pointer and shape.

In general, custom operations should not create ManagedTensorSlice instances, but instead use the ones provided by the MAX inference engine.

__init__(ptr: UnsafePointer[Scalar[dtype], MutAnyOrigin], shape: IndexList[rank], strides: IndexList[rank]) -> Self

Initializes a ManagedTensorSlice from a pointer, shape, and strides.

In general, custom operations should not create ManagedTensorSlice instances, but instead use the ones provided by the MAX inference engine.

`getitem`

__getitem__(self, indices: IndexList[rank]) -> Scalar[dtype]

Gets the value at the specified indices.

Args:

indices (IndexList[rank]): The indices of the value to retrieve.

Returns:

Scalar[dtype]: The value at the specified indices.

__getitem__(self, *indices: Int) -> Scalar[dtype]

Gets the value at the specified indices.

Args:

*indices (Int): The indices of the value to retrieve.

Returns:

Scalar[dtype]: The value at the specified indices.

`setitem`

__setitem__(self, *indices: Int, *, val: Scalar[dtype])

Stores the value at the specified indices.

Args:

*indices (Int): The indices of the value to store.
val (Scalar[dtype]): The value to store.

__setitem__(self, indices: IndexList[rank], val: Scalar[dtype])

Stores the value at the specified indices.

Args:

indices (IndexList[rank]): The indices of the value to store.
val (Scalar[dtype]): The value to store.

`get_type_name`

static get_type_name() -> String

Returns:

String

`spec`

spec(self) -> RuntimeTensorSpec[dtype, rank]

Gets the TensorSpec of this tensor slice, which provides meta-data about the tensor slice.

Returns:

RuntimeTensorSpec[dtype, rank]: The static TensorSpec for this tensor slice.

`shape`

shape(self) -> IndexList[rank]

Gets the shape of this tensor slice, as an IndexList.

Returns:

IndexList[rank]: The shape of this tensor slice.

`dim_size`

dim_size(self, index: Int) -> Int

Gets the size of a given dimension of this tensor slice using a run time value.

Args:

index (Int): The zero-based index of the dimension.

Returns:

Int: The size of the tensor slice in the given dimension.

dim_size[index: Int](self) -> Int

Gets the size of a given dimension of this tensor slice using a compile time value.

Parameters:

index (Int): The zero-based index of the dimension.

Returns:

Int: The size of the tensor slice in the given dimension.

`strides`

strides(self) -> IndexList[rank]

Gets the strides of this tensor slice, as an IndexList.

Returns:

IndexList[rank]: The strides of this tensor slice.

`stride_length`

stride_length(self, index: Int) -> Int

Gets the length of the stride of a given dimension of this tensor slice using a run time value.

Args:

index (Int): The zero-based index of the dimension.

Returns:

Int: The size of the tensor slice in the given dimension.

stride_length[index: Int](self) -> Int

Gets the length of the stride of a given dimension of this tensor slice using a compile time value.

Parameters:

index (Int): The zero-based index of the dimension.

Returns:

Int: The size of the tensor slice in the given dimension.

`size`

size(self) -> Int

Computes the tensor slice's number of elements.

Returns:

Int: The total number of elements in the tensor slice.

`bytecount`

bytecount(self) -> Int

Returns the size of the tensor slice in bytes.

Returns:

Int: The total number of bytes in the tensor slice.

`unsafe_ptr`

unsafe_ptr[_dtype: DType = dtype](self) -> UnsafePointer[Scalar[_dtype], MutAnyOrigin]

Get the pointer stored in this tensor slice.

Since this method obtains the pointer stored in this tensor slice, it can modify the invariants of this tensor slice and lead to unexpected behavior. It should be used with caution.

Parameters:

_dtype (DType): The type of the UnsafePointer in this tensor slice.

Returns:

UnsafePointer[Scalar[_dtype], MutAnyOrigin]: The UnsafePointer which contains the data for this tensor slice.

`to_device_buffer`

to_device_buffer(self, ctx: DeviceContext) -> DeviceBuffer[dtype]

Returns:

DeviceBuffer[dtype]

`load`

load[width: Int, _rank: Int, element_alignment: Int = 1](self, index: IndexList[_rank]) -> SIMD[dtype, width]

Gets data from this tensor slice as a SIMD.

Parameters:

width (Int): The width of the SIMD value. This must be large enough to contain the data from this tensor slice.
_rank (Int): The rank of the tensor slice.
element_alignment (Int): Indicate the alignment of the pointer stored to memory. This is needed to issue vector load for GPUs with strict alignment requirements.

Args:

index (IndexList[_rank]): An IndexList of size _rank to indicate the dimension of the tensor slice to obtain data from.

Returns:

SIMD[dtype, width]: Data from this tensor slice at dimension index.

`store`

store[width: Int, _rank: Int, element_alignment: Int = 1](self: ManagedTensorSlice[static_spec=static_spec], index: IndexList[_rank], val: SIMD[dtype, width])

Sets data in this tensor slice from a SIMD.

Parameters:

width (Int): The width of the SIMD value.
_rank (Int): The rank of the tensor slice.
element_alignment (Int): Indicate the alignment of the pointer stored to memory. This is needed to issue vector store for GPUs with strict alignment requirements.

Args:

index (IndexList[_rank]): An IndexList of size _rank to indicate the dimension of the tensor slice to set data in.
val (SIMD[dtype, width]): The data to set into this tensor slice.

`with_tile_layout`

with_tile_layout[new_layout: TensorLayout](self, new_runtime_shape: IndexList[new_layout.rank], new_runtime_strides: IndexList[new_layout.rank], offset_ptr: Optional[UnsafePointer[Scalar[dtype], MutAnyOrigin]] = None) -> ManagedTensorSlice[io_spec, static_spec=static_spec.with_tile_layout[new_layout]()]

Returns:

ManagedTensorSlice[io_spec, static_spec=static_spec.with_tile_layout[new_layout]()]

`to_layout_tensor`

to_layout_tensor(self) -> LayoutTensor[dtype, static_spec.to_layout(), MutAnyOrigin]

Returns:

LayoutTensor[dtype, static_spec.to_layout(), MutAnyOrigin]

`to_tile_tensor`

to_tile_tensor[coord_dtype: DType = DType.int64](self) -> TileTensor[dtype, Layout[static_spec.static_layout._shape_types, static_spec.static_layout._stride_types], MutExternalOrigin]

Returns:

TileTensor[dtype, Layout[static_spec.static_layout._shape_types, static_spec.static_layout._stride_types], MutExternalOrigin]

`write_to`

write_to(self, mut writer: T)

Formats this buffer to the provided Writer.

Args:

writer (T): The object to write to.

`write_repr_to`

write_repr_to(self, mut writer: T)

Formats this buffer to the provided Writer.

Args:

writer (T): The object to write to.

Fields​

Implemented traits​

comptime members​

address_space​

alignment​

device_type​

exclusive​

Methods​

__init__​

__getitem__​

__setitem__​

get_type_name​

spec​

shape​

dim_size​

strides​

stride_length​

size​

bytecount​

unsafe_ptr​

to_device_buffer​

load​

store​

with_tile_layout​

to_layout_tensor​

to_tile_tensor​

write_to​

write_repr_to​

Fields

Implemented traits

`comptime` members

`address_space`

`alignment`

`device_type`

`exclusive`

Methods

`init`

`getitem`

`setitem`

`get_type_name`

`spec`

`shape`

`dim_size`

`strides`

`stride_length`

`size`

`bytecount`

`unsafe_ptr`

`to_device_buffer`

`load`

`store`

`with_tile_layout`

`to_layout_tensor`

`to_tile_tensor`

`write_to`

`write_repr_to`