Struct Barrier

pub struct Barrier<C: CubePrimitive> { /* private fields */ }

A mechanism for awaiting on asynchronous data transfers Behaviour is defined by its BarrierLevel.

Implementations

impl<C: CubePrimitive> Barrier<C>

pub fn tma_load_2d( &self, source: &TensorMap<C>, destination: &mut SliceMut<Line<C>>, y: i32, x: i32, )

Copy a tile from a global memory source to a shared memory destination, with the provided offsets.

pub fn __expand_tma_load_2d( scope: &mut Scope, expand: BarrierExpand<C>, source: ExpandElementTyped<TensorMap<C>>, destination: SliceExpand<Line<C>, ReadWrite>, y: ExpandElementTyped<i32>, x: ExpandElementTyped<i32>, )

impl<C: CubePrimitive> Barrier<C>

pub fn tma_load_3d( &self, source: &TensorMap<C>, destination: &mut SliceMut<Line<C>>, z: i32, y: i32, x: i32, )

Copy a tile from a global memory source to a shared memory destination, with the provided offsets.

pub fn __expand_tma_load_3d( scope: &mut Scope, expand: BarrierExpand<C>, source: ExpandElementTyped<TensorMap<C>>, destination: SliceExpand<Line<C>, ReadWrite>, z: ExpandElementTyped<i32>, y: ExpandElementTyped<i32>, x: ExpandElementTyped<i32>, )

impl<C: CubePrimitive> Barrier<C>

pub fn tma_load_4d( &self, source: &TensorMap<C>, destination: &mut SliceMut<Line<C>>, w: i32, z: i32, y: i32, x: i32, )

Copy a tile from a global memory source to a shared memory destination, with the provided offsets.

pub fn __expand_tma_load_4d( scope: &mut Scope, expand: BarrierExpand<C>, source: ExpandElementTyped<TensorMap<C>>, destination: SliceExpand<Line<C>, ReadWrite>, w: ExpandElementTyped<i32>, z: ExpandElementTyped<i32>, y: ExpandElementTyped<i32>, x: ExpandElementTyped<i32>, )

impl<C: CubePrimitive> Barrier<C>

pub fn tma_load_5d( &self, source: &TensorMap<C>, destination: &mut SliceMut<Line<C>>, v: i32, w: i32, z: i32, y: i32, x: i32, )

Copy a tile from a global memory source to a shared memory destination, with the provided offsets.

pub fn __expand_tma_load_5d( scope: &mut Scope, expand: BarrierExpand<C>, source: ExpandElementTyped<TensorMap<C>>, destination: SliceExpand<Line<C>, ReadWrite>, v: ExpandElementTyped<i32>, w: ExpandElementTyped<i32>, z: ExpandElementTyped<i32>, y: ExpandElementTyped<i32>, x: ExpandElementTyped<i32>, )

impl<C: CubePrimitive> Barrier<C>

pub fn tma_load_im2col_3d( &self, source: &TensorMap<C>, destination: &mut SliceMut<Line<C>>, n: i32, w: i32, c: i32, w_offset: u16, )

Copy a tile from a global memory source to a shared memory destination, with the provided offsets.

pub fn __expand_tma_load_im2col_3d( scope: &mut Scope, expand: BarrierExpand<C>, source: ExpandElementTyped<TensorMap<C>>, destination: SliceExpand<Line<C>, ReadWrite>, n: ExpandElementTyped<i32>, w: ExpandElementTyped<i32>, c: ExpandElementTyped<i32>, w_offset: ExpandElementTyped<u16>, )

impl<C: CubePrimitive> Barrier<C>

pub fn tma_load_im2col_4d( &self, source: &TensorMap<C>, destination: &mut SliceMut<Line<C>>, n: i32, h: i32, w: i32, c: i32, h_offset: u16, w_offset: u16, )

Copy a tile from a global memory source to a shared memory destination, with the provided offsets.

pub fn __expand_tma_load_im2col_4d( scope: &mut Scope, expand: BarrierExpand<C>, source: ExpandElementTyped<TensorMap<C>>, destination: SliceExpand<Line<C>, ReadWrite>, n: ExpandElementTyped<i32>, h: ExpandElementTyped<i32>, w: ExpandElementTyped<i32>, c: ExpandElementTyped<i32>, h_offset: ExpandElementTyped<u16>, w_offset: ExpandElementTyped<u16>, )

impl<C: CubePrimitive> Barrier<C>

pub fn tma_load_im2col_5d( &self, source: &TensorMap<C>, destination: &mut SliceMut<Line<C>>, n: i32, d: i32, h: i32, w: i32, c: i32, d_offset: u16, h_offset: u16, w_offset: u16, )

Copy a tile from a global memory source to a shared memory destination, with the provided offsets.

pub fn __expand_tma_load_im2col_5d( scope: &mut Scope, expand: BarrierExpand<C>, source: ExpandElementTyped<TensorMap<C>>, destination: SliceExpand<Line<C>, ReadWrite>, n: ExpandElementTyped<i32>, d: ExpandElementTyped<i32>, h: ExpandElementTyped<i32>, w: ExpandElementTyped<i32>, c: ExpandElementTyped<i32>, d_offset: ExpandElementTyped<u16>, h_offset: ExpandElementTyped<u16>, w_offset: ExpandElementTyped<u16>, )

impl<C: CubePrimitive> Barrier<C>

pub fn new(_level: BarrierLevel) -> Self

Creates a barrier using a user defined comptime barrier level

pub fn new_with_tma_proxy(_level: BarrierLevel) -> Self

Creates a new barrier for use with TMA instructions. Adds a shared memory proxy barrier to the initialization.

pub fn memcpy_async( &self, _source: &Slice<Line<C>>, _destination: &mut SliceMut<Line<C>>, )

Copy the source slice to destination

Safety

This will try to copy the whole source slice, so make sure source length <= destination length

pub fn arrive(&self)

Arrive at the barrier, decrementing arrival count

pub fn arrive_tx(&self, _arrival_count: u32, _transaction_count: u32)

Arrive at the barrier, decrementing arrival count. Additionally increments expected count.

pub fn expect_tx(&self, _expected_count: u32)

Increments the expected count of the barrier.

pub fn wait(&self)

Wait at the barrier until all arrivals are done

pub fn arrive_and_wait(&self)

Wait until all data is loaded

pub fn __expand_new(scope: &mut Scope, level: BarrierLevel) -> BarrierExpand<C>

pub fn __expand_new_with_tma_proxy( scope: &mut Scope, level: BarrierLevel, ) -> BarrierExpand<C>

pub fn __expand_memcpy_async( scope: &mut Scope, expand: BarrierExpand<C>, source: SliceExpand<Line<C>, ReadOnly>, destination: SliceExpand<Line<C>, ReadWrite>, )

pub fn __expand_arrive(scope: &mut Scope, expand: BarrierExpand<C>)

pub fn __expand_arrive_tx( scope: &mut Scope, expand: BarrierExpand<C>, arrival_count: ExpandElementTyped<u32>, transaction_count: ExpandElementTyped<u32>, )

pub fn __expand_expect_tx( scope: &mut Scope, expand: BarrierExpand<C>, expected_count: ExpandElementTyped<u32>, )

pub fn __expand_wait(scope: &mut Scope, expand: BarrierExpand<C>)

pub fn __expand_arrive_and_wait(scope: &mut Scope, expand: BarrierExpand<C>)

Trait Implementations

impl<C: Clone + CubePrimitive> Clone for Barrier<C>

fn clone(&self) -> Barrier<C>

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl<C: CubePrimitive> CubeType for Barrier<C>

type ExpandType = BarrierExpand<C>

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

Wrapper around the init method, necessary to type inference.

impl<C: Copy + CubePrimitive> Copy for Barrier<C>