Struct TensorMap

pub struct TensorMap<E: CubePrimitive> { /* private fields */ }

A CUDA CUtensorMap object. Represents a tensor encoded with a lot of metadata, and is an opaque packed object at runtime. Does not support retrieving any shapes or strides, nor does it give access to the pointer. So these need to be passed separately in an aliased Tensor if needed.

Also see [cubecl_common::tma].

Implementations

impl<T: CubePrimitive> TensorMap<T>

pub fn stride<C: Index>(&self, _dim: C) -> u32

Obtain the stride of input at dimension dim

pub fn shape<C: Index>(&self, _dim: C) -> u32

Obtain the shape of input at dimension dim

pub fn coordinate<I: Index, D: Index>(&self, _index: I, _dim: D) -> u32

Obtain the coordinate corresponding to the given index of the tensor at dimension dim.

A coordinate is a list of indices corresponding to the multi-dimensional position of an element in the tensor. The dim element in a coordinate is the position along the dim dimension of the tensor.

pub fn len(&self) -> u32

The number of vectorized elements in the tensor.

Warning

The length will be affected by the vectorization factor. To obtain the number of elements, you should multiply the length by the vectorization factor.

pub fn buffer_len(&self) -> u32

The length of the buffer representing the tensor in terms of vectorized elements.

Warning

The buffer length will be affected by the vectorization factor. To obtain the number of elements, you should multiply the length by the vectorization factor.

pub fn rank(&self) -> u32

Returns the rank of the tensor.

pub fn try_cast_unchecked<E: CubePrimitive>(&self) -> TensorMap<E>

Try to cast the tensormap to the given type and panic if the type isn’t the same.

This function should only be used to satisfy the Rust type system, when two generic types are supposed to be the same.

pub fn __expand_stride<C: Index>( scope: &mut Scope, expand: ExpandElementTyped<Tensor<T>>, dim: ExpandElementTyped<u32>, ) -> ExpandElementTyped<u32>

pub fn __expand_shape<C: Index>( scope: &mut Scope, expand: ExpandElementTyped<Tensor<T>>, dim: ExpandElementTyped<u32>, ) -> ExpandElementTyped<u32>

pub fn __expand_coordinate<I: Index, D: Index>( scope: &mut Scope, expand: ExpandElementTyped<Tensor<T>>, index: ExpandElementTyped<u32>, dim: ExpandElementTyped<u32>, ) -> ExpandElementTyped<u32>

pub fn __expand_len<C: Index>( scope: &mut Scope, expand: ExpandElementTyped<Tensor<T>>, ) -> ExpandElementTyped<u32>

pub fn __expand_buffer_len<C: Index>( scope: &mut Scope, expand: ExpandElementTyped<Tensor<T>>, ) -> ExpandElementTyped<u32>

pub fn __expand_rank<C: Index>( scope: &mut Scope, expand: ExpandElementTyped<Tensor<T>>, ) -> ExpandElementTyped<u32>

impl<E: CubePrimitive> TensorMap<E>

pub fn dummy() -> Self

Create a dummy tensor map to satisfy the type checker. Not actually valid, so the code should panic if this is ever reached.

pub fn __expand_dummy(_scope: &mut Scope) -> ExpandElementTyped<Self>

Trait Implementations

impl<E: Clone + CubePrimitive> Clone for TensorMap<E>

fn clone(&self) -> TensorMap<E>

Returns a copy of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<E: CubePrimitive> CubeType for TensorMap<E>

type ExpandType = ExpandElementTyped<TensorMap<E>>

fn init(scope: &mut Scope, expand: Self::ExpandType) -> Self::ExpandType

Wrapper around the init method, necessary to type inference.

impl<E: CubePrimitive> ExpandElementBaseInit for TensorMap<E>

fn init_elem(_scope: &mut Scope, elem: ExpandElement) -> ExpandElement

impl<E: CubePrimitive> LaunchArg for TensorMap<E>

type RuntimeArg<'a, R: Runtime> = TensorMapArg<'a, R>

The runtime argument for the kernel.

fn compilation_arg<R: Runtime>( _runtime_arg: &Self::RuntimeArg<'_, R>, ) -> Self::CompilationArg

impl<E: CubePrimitive> LaunchArgExpand for TensorMap<E>

type CompilationArg = TensorMapCompilationArg

Compilation argument.

fn expand( _arg: &Self::CompilationArg, builder: &mut KernelBuilder, ) -> ExpandElementTyped<TensorMap<E>>

Register an input variable during compilation that fill the KernelBuilder.

fn expand_output( _arg: &Self::CompilationArg, builder: &mut KernelBuilder, ) -> ExpandElementTyped<TensorMap<E>>

Register an output variable during compilation that fill the KernelBuilder.

impl<E: CubePrimitive> Copy for TensorMap<E>