Skip to main content

/

Docs

Nightly

v25.1

Log in

Mojo struct

LayoutTensor

@register_passable(trivial) struct LayoutTensor[dtype: DType, layout: Layout, /, *, mut: Bool = True, origin: Origin[mut] = SomeAnyOrigin, address_space: AddressSpace = AddressSpace(0), element_layout: Layout = __init__[::Origin[::Bool(IntTuple(1), IntTuple(1)), layout_bitwidth: Int = bitwidthof[::DType,__mlir_type.!kgen.target](), masked: Bool = False, alignment: Int = alignof[::DType,__mlir_type.!kgen.target]()]

This is a Tensor type that has a specified memory layout and rank. The following example demonstrate a LayoutTensor of float32 with a row major layout of shape (5, 4).

alias f32 = DType.float32
var tensor_5x4 = LayoutTensor[f32, Layout.row_major(5,4)].stack_allocation()
alias f32 = DType.float32
var tensor_5x4 = LayoutTensor[f32, Layout.row_major(5,4)].stack_allocation()

Parameters​

  • ​dtype (DType): The data type of the underlying pointer.
  • ​layout (Layout): The memory layout of the Tensor.
  • ​mut (Bool): The mutability of the underlying pointer.
  • ​origin (Origin[mut]): The origin of the underlying pointer.
  • ​address_space (AddressSpace): The address space of the underlying pointer.
  • ​element_layout (Layout): The memory layout of each element in the Tensor.
  • ​layout_bitwidth (Int): The bitwidth of each dimension of runtime layout.
  • ​masked (Bool): If true the tensor is masked and runtime layouts determine the shape.
  • ​alignment (Int): Alignment of the data pointer.

Aliases​

  • rank = layout.rank():
  • index_type = _get_index_type(layout, address_space):
  • uint_type = SIMD[_get_unsigned_type(layout, address_space), 1]:
  • element_size = element_layout.size():
  • element_type = SIMD[dtype, element_layout.size()]:

Fields​

  • ​ptr (UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]):
  • ​runtime_layout (RuntimeLayout[layout, bitwidth=layout_bitwidth]):
  • ​runtime_element_layout (RuntimeLayout[element_layout]):

Implemented traits​

AnyType, CollectionElement, CollectionElementNew, Copyable, ExplicitlyCopyable, Movable, Stringable, UnknownDestructibility, Writable

Methods​

__init__​

@implicit __init__(ptr: UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]) -> Self

Create a LayoutTensor with an UnsafePointer. Expect layout to be fully static.

Args:

  • ​ptr (UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]): The UnsafePointer pointing to the underlying data.

__init__(ptr: UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin], runtime_layout: RuntimeLayout[layout, bitwidth=bitwidth]) -> Self

Create a LayoutTensor with an UnsafePointer. Expect element layout to be fully static.

Args:

  • ​ptr (UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]): The UnsafePointer pointing to the underlying data.
  • ​runtime_layout (RuntimeLayout[layout, bitwidth=bitwidth]): The runtime layout of the LayoutTensor.

__init__(ptr: UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin], runtime_layout: RuntimeLayout[layout, bitwidth=layout_bitwidth], element_runtime_layout: RuntimeLayout[element_layout]) -> Self

Create a LayoutTensor with an UnsafePointer, a runtime layout of the Tensor, the runtime layout of each element.

Args:

  • ​ptr (UnsafePointer[SIMD[dtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin]): The UnsafePointer pointing to the underlying data.
  • ​runtime_layout (RuntimeLayout[layout, bitwidth=layout_bitwidth]): The runtime layout of the LayoutTensor.
  • ​element_runtime_layout (RuntimeLayout[element_layout]): The runtime layout of each element.

@implicit __init__(device_buffer: DeviceBuffer[dtype, address_space, mut, origin]) -> Self

Create a LayoutTensor from a DeviceBuffer. The layout must have statically known dimensions.

from gpu.host import DeviceContext, DeviceBuffer
from layout import Layout, LayoutTensor

alias dtype = DType.float32

var ctx = DeviceContext()
var dev_buf = ctx.enqueue_create_buffer[dtype](8)

alias layout = Layout.row_major(4, 4)
var tensor = LayoutTensor[dtype, layout](dev_buf)
from gpu.host import DeviceContext, DeviceBuffer
from layout import Layout, LayoutTensor

alias dtype = DType.float32

var ctx = DeviceContext()
var dev_buf = ctx.enqueue_create_buffer[dtype](8)

alias layout = Layout.row_major(4, 4)
var tensor = LayoutTensor[dtype, layout](dev_buf)

Args:

  • ​device_buffer (DeviceBuffer[dtype, address_space, mut, origin]): Contains the underlying data to point to.

__init__(device_buffer: DeviceBuffer[dtype, address_space, mut, origin], runtime_layout: RuntimeLayout[layout, bitwidth=bitwidth]) -> Self

Create a LayoutTensor from a DeviceBuffer. The layout must have statically known dimensions.

Args:

  • ​device_buffer (DeviceBuffer[dtype, address_space, mut, origin]): The DeviceBuffer containing to the underlying data.
  • ​runtime_layout (RuntimeLayout[layout, bitwidth=bitwidth]): The runtime layout of the LayoutTensor.

__init__(device_buffer: DeviceBuffer[dtype, address_space, mut, origin], runtime_layout: RuntimeLayout[layout, bitwidth=layout_bitwidth], element_runtime_layout: RuntimeLayout[element_layout]) -> Self

Create a LayoutTensor from a DeviceBuffer, a runtime layout of the Tensor, and the runtime layout of each element.

Args:

  • ​device_buffer (DeviceBuffer[dtype, address_space, mut, origin]): The DeviceBuffer containing to the underlying data.
  • ​runtime_layout (RuntimeLayout[layout, bitwidth=layout_bitwidth]): The runtime layout of the LayoutTensor.
  • ​element_runtime_layout (RuntimeLayout[element_layout]): The runtime layout of each element.

__getitem__​

__getitem__(self, *dims: Int) -> SIMD[dtype, element_layout.size()]

Get the element of the tensor with a specified index. Note that the size of index has to match the rank of the tensor.

Args:

  • ​*dims (Int): The indexes that specify which element to retrieve.

__setitem__​

__setitem__(self, d0: Int, val: SIMD[dtype, element_layout.size()])

Set the element of the tensor with a specified index and value.

Args:

  • ​d0 (Int): The first dimensional index.
  • ​val (SIMD[dtype, element_layout.size()]): The value writing to the tensor.

__setitem__(self, d0: Int, d1: Int, val: SIMD[dtype, element_layout.size()])

Set the element of the tensor with a specified index and value.

Args:

  • ​d0 (Int): The first dimensional index.
  • ​d1 (Int): The second dimensional index.
  • ​val (SIMD[dtype, element_layout.size()]): The value writing to the tensor.

__setitem__(self, d0: Int, d1: Int, d2: Int, val: SIMD[dtype, element_layout.size()])

Set the element of the tensor with a specified index and value.

Args:

  • ​d0 (Int): The first dimensional index.
  • ​d1 (Int): The second dimensional index.
  • ​d2 (Int): The third dimensional index.
  • ​val (SIMD[dtype, element_layout.size()]): The value writing to the tensor.

__setitem__(self, d0: Int, d1: Int, d2: Int, d3: Int, val: SIMD[dtype, element_layout.size()])

Set the element of the tensor with a specified index and value.

Args:

  • ​d0 (Int): The first dimensional index.
  • ​d1 (Int): The second dimensional index.
  • ​d2 (Int): The third dimensional index.
  • ​d3 (Int): The fourth dimensional index.
  • ​val (SIMD[dtype, element_layout.size()]): The value writing to the tensor.

__add__​

__add__(self, other: SIMD[dtype, 1]) -> LayoutTensor[dtype, layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

Add the LayoutTensor with a scalar value. The scalar value will be broadcasted to the entire tensor.

Args:

  • ​other (SIMD[dtype, 1]): The scalar value.

__add__[other_layout: Layout](self, other: LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]) -> LayoutTensor[dtype, layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

Do an addition with another LayoutTensor and return the added tensor. Currently only support tensors of the same shape if the rank is the same and also tensors of rank-2.

Parameters:

  • ​other_layout (Layout): The layout of the other tensor.

Args:

  • ​other (LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]): The other tensor to be added to.

__sub__​

__sub__(self, other: SIMD[dtype, 1]) -> LayoutTensor[dtype, layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

Subtract the LayoutTensor with a scalar value. The scalar value will be broadcasted to the entire tensor.

Args:

  • ​other (SIMD[dtype, 1]): The scalar value.

__sub__[other_layout: Layout](self, other: LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]) -> LayoutTensor[dtype, layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

Do an subtraction with another LayoutTensor and return the subtracted tensor. Currently only support tensors of the same shape if the rank is the same and also tensors of rank-2.

Parameters:

  • ​other_layout (Layout): The layout of the other tensor.

Args:

  • ​other (LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]): The other tensor to be subtract from.

__mul__​

__mul__(self, other: SIMD[dtype, 1]) -> LayoutTensor[dtype, layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

Multiply the LayoutTensor with a scalar value. The scalar value will be broadcasted to the entire tensor.

Args:

  • ​other (SIMD[dtype, 1]): The scalar value.

__mul__[other_layout: Layout](self, other: LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]) -> LayoutTensor[dtype, layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

Perform a multiplication with another LayoutTensor and return the resulting tensor.

Currently, only tensors of the same shape are supported if the ranks are the same. Additionally, tensors of rank-2 are supported.

Parameters:

  • ​other_layout (Layout): The layout of the other tensor.

Args:

  • ​other (LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]): The other tensor to be multiplied with.

Returns:

The resulting tensor after multiplication.

__truediv__​

__truediv__(self, other: SIMD[dtype, 1]) -> LayoutTensor[dtype, layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

Truediv the LayoutTensor with a scalar value. The scalar value will be broadcasted to the entire tensor.

Args:

  • ​other (SIMD[dtype, 1]): The scalar value.

__truediv__[other_layout: Layout](self, other: LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]) -> LayoutTensor[dtype, layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

Do an truediv with another LayoutTensor and return the divided tensor. Currently only support tensors of the same shape if the rank is the same and also tensors of rank-2.

Parameters:

  • ​other_layout (Layout): The layout of the other tensor.

Args:

  • ​other (LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]): The other tensor to be subtract from.

__iadd__​

__iadd__(self, other: SIMD[dtype, 1])

Adds scalar value to the LayoutTensor. The scalar value will be broadcasted to the entire tensor.

Args:

  • ​other (SIMD[dtype, 1]): The scalar value.

__iadd__[other_layout: Layout](self, other: LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth])

Do an addition with another LayoutTensor in place. Currently only support tensors of the same shape if the rank is the same and also tensors of rank-2.

Parameters:

  • ​other_layout (Layout): The layout of the other tensor.

Args:

  • ​other (LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]): The other tensor to be added to.

__isub__​

__isub__(self, other: SIMD[dtype, 1])

Subtract scalar value from the LayoutTensor. The scalar value will be broadcasted to the entire tensor.

Args:

  • ​other (SIMD[dtype, 1]): The scalar value.

__isub__[other_layout: Layout](self, other: LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth])

Subtracts other from the LayoutTensor. Currently only support tensors of the same shape if the rank is the same and also tensors of rank-2.

Parameters:

  • ​other_layout (Layout): The layout of the other tensor.

Args:

  • ​other (LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]): The other tensor to be subtract from.

__imul__​

__imul__(self, other: SIMD[dtype, 1])

Multiply the LayoutTensor with a scalar value inplace. The scalar value will be broadcasted to the entire tensor.

Args:

  • ​other (SIMD[dtype, 1]): The scalar value.

__imul__[other_layout: Layout](self, other: LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth])

Do a multiplication with another LayoutTensor in place. Currently only support tensors of the same shape if the rank is the same and also tensors of rank-2.

Parameters:

  • ​other_layout (Layout): The layout of the other tensor.

Args:

  • ​other (LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]): The other tensor to be added to.

__itruediv__​

__itruediv__(self, other: SIMD[dtype, 1])

Truediv the LayoutTensor with a scalar value. The scalar value will be broadcasted to the entire tensor.

Args:

  • ​other (SIMD[dtype, 1]): The scalar value.

__itruediv__[other_layout: Layout](self, other: LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth])

Do an truediv with another LayoutTensor and return the divided tensor. Currently only support tensors of the same shape if the rank is the same and also tensors of rank-2.

Parameters:

  • ​other_layout (Layout): The layout of the other tensor.

Args:

  • ​other (LayoutTensor[dtype, other_layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth]): The other tensor to be subtract from.

copy​

copy(self) -> Self

Explicitly copy the other LayoutTensor.

Returns:

A copy of the value.

bitcast​

bitcast[new_type: DType, /, address_space: AddressSpace = address_space, element_layout: Layout = element_layout](self) -> LayoutTensor[new_type, layout, mut=mut, origin=origin, address_space=address_space, element_layout=element_layout, masked=masked]

Bitcast the underlying pointer to a new data type.

Parameters:

  • ​new_type (DType): The new data type it is casting to.
  • ​address_space (AddressSpace): The address space of the returned LayoutTensor.
  • ​element_layout (Layout): The element layout of the returned LayoutTensor.

origin_cast​

origin_cast[mut: Bool = mut, origin: Origin[$0] = (mutcast origin._mlir_origin)](self) -> LayoutTensor[dtype, layout, mut=mut, origin=origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

Changes the origin or mutability of a pointer.

Parameters:

  • ​mut (Bool): Whether the origin is mutable.
  • ​origin (Origin[$0]): Origin of the destination pointer.

Returns:

A new UnsafePointer object with the same type and the same address, as the original UnsafePointer and the new specified mutability and origin.

load​

load[width: Int](self, m: Int, n: Int) -> SIMD[dtype, width]

Load a value from a specified location.

Parameters:

  • ​width (Int): The simd width of the returned value.

Args:

  • ​m (Int): The m dimension of the value.
  • ​n (Int): The n dimension of the value.

prefetch​

prefetch(self, m: Int, n: Int)

Do software prefetching of a value from a specified location.

Args:

  • ​m (Int): The m dimension of the value.
  • ​n (Int): The n dimension of the value.

aligned_load​

aligned_load[width: Int](self, m: Int, n: Int) -> SIMD[dtype, width]

Do a load with a specified alignment base on the dtype and simd width.

Parameters:

  • ​width (Int): The simd width if the returned value.

Args:

  • ​m (Int): The m dimension of the value.
  • ​n (Int): The n dimension of the value.

store​

store[width: Int](self, m: Int, n: Int, val: SIMD[dtype, width])

Store a value to a specified location.

Parameters:

  • ​width (Int): The simd width of the stored value.

Args:

  • ​m (Int): The m dimensional index to the tensor.
  • ​n (Int): The n dimensional index to the tensor.
  • ​val (SIMD[dtype, width]): The value to be stored.

aligned_store​

aligned_store[width: Int](self, m: Int, n: Int, val: SIMD[dtype, width])

Do a store with a specified alignment base on the dtype and simd width.

Parameters:

  • ​width (Int): The simd width if the stored value.

Args:

  • ​m (Int): The m dimensional index to the tensor.
  • ​n (Int): The n dimensional index to the tensor.
  • ​val (SIMD[dtype, width]): The value to be stored.

stack_allocation​

static stack_allocation[*, alignment: Int = alignment]() -> LayoutTensor[dtype, layout, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

Allocates stack memory for a LayoutTensor. Expects layout to be fully static.

Parameters:

  • ​alignment (Int): Alignment of the allocation. It must be multiple of the tensor's alignment, which is the minimum required by arch, instruction, etc.

shape​

static shape[idx: Int]() -> Int

Returns the shape of the tensor given a index.

Parameters:

  • ​idx (Int): The index to the shape of the tensor.

stride​

static stride[idx: Int]() -> Int

Returns the stride of the tensor given a index.

Parameters:

  • ​idx (Int): The index to the stride of the tensor.

dim​

dim(self, idx: Int) -> Int

Returns the dimension of the tensor given a index.

Arguments: idx: The index to the dimension of the tensor.

coalesce​

coalesce(self) -> LayoutTensor[dtype, coalesce(layout, False), mut=mut, origin=origin, address_space=address_space, element_layout=element_layout]

Returns a LayoutTensor with a coalesced Layout.

tile​

tile[*tile_sizes: Int](self, *tile_coords: Int) -> LayoutTensor[dtype, _compute_tile_layout[*::Int]().__getitem__(0), mut=mut, origin=origin, address_space=address_space, element_layout=element_layout, masked=masked if masked else _tile_is_masked[layout::layout::Layout,*::Int]()]

Tiles the layout and returns a tensor tile with the specified tile_sizes at specific tile coordinates.

Example:

Memory Layout of
                [1 2 3 4]
                [2 3 4 5]
                [5 4 3 2]
                [1 1 1 1]

tile[2, 2](1, 0) will give you
                [5 4]
                [1 1]
Memory Layout of
                [1 2 3 4]
                [2 3 4 5]
                [5 4 3 2]
                [1 1 1 1]

tile[2, 2](1, 0) will give you
                [5 4]
                [1 1]

Parameters:

  • ​*tile_sizes (Int): The tile sizes of the returned LayoutTensor.

Args:

  • ​*tile_coords (Int): The tile coordinate. This refer to the coordinate of the tile after the tiled layout. Consider the following example.

tiled_iterator​

tiled_iterator[*tile_sizes: Int, *, axis: Int = 0](self, *tile_coords: Int) -> LayoutTensorIter[dtype, _compute_tile_layout[*::Int]().__getitem__(0), mut=mut, origin=origin, address_space=address_space, axis=OptionalReg[Int]({:_stdlib::_builtin::_int::_Int axis, 0}), layout_bitwidth=layout_bitwidth, masked=masked if masked else _tile_is_masked[layout::layout::Layout,*::Int]()]

Returns the tiled iterator of the LayoutTensor.

Parameters:

  • ​*tile_sizes (Int): Tile sizes of each tile the iterator will iterate through.
  • ​axis (Int): Axis of the LayoutTensor the iterator will iterate through.

Args:

  • ​*tile_coords (Int): The tile coordinate that the iterator will point to.

split​

split[count: Int, axis: Int = 0](self) -> StaticTuple[LayoutTensor[dtype, _compute_tile_layout[::Int,::Int]().__getitem__(0), mut=mut, origin=origin, address_space=address_space, element_layout=element_layout, alignment=alignment], count]

Split the LayoutTensor along a axis and return an InlineArray of LayoutTensor.

Parameters:

  • ​count (Int): Number of portion to split.
  • ​axis (Int): The axis where the split is applied to.

split[axis: Int = 0, alignment: Int = 1](self, count: Int, idx: Int) -> LayoutTensor[dtype, layout.make_shape_unknown[::Int](), mut=mut, origin=origin, address_space=address_space, element_layout=element_layout]

distribute​

distribute[threads_layout: Layout, axis: OptionalReg[Int] = OptionalReg[Int]({:i1 0, 1}), swizzle: OptionalReg[Swizzle] = OptionalReg[Swizzle]({:i1 0, 1}), submode_axis: OptionalReg[Int] = OptionalReg[Int]({:i1 0, 1})](self, thread_id: UInt) -> LayoutTensor[dtype, _compute_distribute_layout[layout::layout::Layout,layout::layout::Layout,stdlib::collections::optional::OptionalReg[::Int]]().__getitem__(1), mut=mut, origin=origin, address_space=address_space, element_layout=element_layout, masked=masked if masked else _distribute_is_masked[layout::layout::Layout,layout::layout::Layout,stdlib::collections::optional::OptionalReg[::Int]]()]

Distribute tiled workload to threads.

If the axis is given, for example, using axis = 0 for 4 threads: TH_0 TH_2 TH_1 TH_3 This means the tensor is only distributed to threads in axis = 0, i.e., threads 0 and 1. Threads 2 and 3 gets the same tile as 0 and 1, respectively. This is useful when threads load same vectors from a row in A matrix and some threads share the same vector.

vectorize​

vectorize[*vector_shape: Int](self) -> LayoutTensor[dtype, coalesce(_compute_tile_layout[*::Int]().__getitem__(1), True), mut=mut, origin=origin, address_space=address_space, element_layout=_divide_tiles[*::Int]().__getitem__(0), masked=masked]

slice​

slice[d0_slice: Slice, d1_slice: Slice](self) -> LayoutTensor[dtype, _compute_slice_layout(d0_slice, d1_slice), mut=mut, origin=origin, address_space=address_space, element_layout=element_layout]

slice[d0_slice: Slice, d1_slice: Slice, slice_indices: Index[2], __offset_dims: Int = (layout.rank() + -2)](self, offsets: Index[__offset_dims]) -> LayoutTensor[dtype, _compute_slice_layout(d0_slice, d1_slice, slice_indices.__getitem__[::Indexer](0), slice_indices.__getitem__[::Indexer](1)), mut=mut, origin=origin, address_space=address_space, element_layout=element_layout]

slice_1d​

slice_1d[d0_slice: Slice, slice_indices: Index[1], __offset_dims: Int = (layout.rank() + -1)](self, offsets: Index[__offset_dims]) -> LayoutTensor[dtype, _compute_slice_layout(d0_slice, slice_indices.__getitem__[::Indexer](0)), mut=mut, origin=origin, address_space=address_space, element_layout=element_layout]

transpose​

transpose[M: Int = shape[::Int](), N: Int = shape[::Int]()](self) -> LayoutTensor[dtype, composition(layout, __init__[::Origin[::Bool(IntTuple(N, M), IntTuple(M, 1))), mut=mut, origin=origin, address_space=address_space, element_layout=element_layout]

reshape​

reshape[dst_layout: Layout](self) -> LayoutTensor[dtype, dst_layout, mut=mut, origin=origin, address_space=address_space, element_layout=element_layout, masked=masked]

composition​

composition[rhs_layout: Layout, dst_layout: Layout = composition(layout, $0)](self) -> LayoutTensor[dtype, dst_layout, mut=mut, origin=origin, address_space=address_space, element_layout=element_layout]

distance​

distance[_uint_dtype: DType = uint32 if (address_space == AddressSpace(3)) else uint64](self, addr: UnsafePointer[SIMD[dtype, 1], address_space=address_space]) -> SIMD[_uint_dtype, 1]

Returns the distance from the input address.

distance[_layout: Layout, _uint_dtype: DType = _get_unsigned_type($0, address_space)](self, src: LayoutTensor[dtype, _layout, address_space=address_space]) -> SIMD[_uint_dtype, 1]

Returns the distance from the input address.

copy_from​

copy_from(self, other: LayoutTensor[dtype, layout, mut=mut, origin=origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment])

copy_from_async​

copy_from_async[is_masked: Bool = False, swizzle: OptionalReg[Swizzle] = OptionalReg[Swizzle]({:i1 0, 1}), fill: Fill = Fill(0), eviction_policy: CacheEviction = CacheEviction(0)](self, src: LayoutTensor[dtype, layout, mut=mut, origin=origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment], src_idx_bound: SIMD[_get_index_type(layout, address_space), 1] = __init__[__mlir_type.!kgen.int_literal](0), base_offset: SIMD[_get_unsigned_type(layout, address_space), 1] = __init__[__mlir_type.!kgen.int_literal](0))

fill​

fill(self: LayoutTensor[dtype, layout, origin=origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment], val: SIMD[dtype, 1]) -> LayoutTensor[dtype, layout, origin=origin, address_space=address_space, element_layout=element_layout, layout_bitwidth=layout_bitwidth, masked=masked, alignment=alignment]

__str__​

__str__(self) -> String

write_to​

write_to[W: Writer](self, mut writer: W)

Format 2D tensor in 2D, otherwise print all values in column major coordinate order.