Trait Triton 

pub trait Triton
where
    Self::I32Tensor: Add<i32, Output = Self::I32Tensor> + Sub<i32, Output = Self::I32Tensor> + Mul<i32, Output = Self::I32Tensor> + Div<i32, Output = Self::I32Tensor>,
    Self::BoolTensor: BitAnd<Output = Self::BoolTensor> + BitOr<Output = Self::BoolTensor>,
{
    type BoolTensor: Copy + Clone;
    type I32Tensor: Copy + Clone;
    type Tensor<D: Dtype>: Copy + Clone + Add<Self::Tensor<D>, Output = Self::Tensor<D>> + Sub<Self::Tensor<D>, Output = Self::Tensor<D>> + Mul<Self::Tensor<D>, Output = Self::Tensor<D>> + Div<Self::Tensor<D>, Output = Self::Tensor<D>> + Neg<Output = Self::Tensor<D>>;
    type Pointer<D: Dtype>: Copy + Clone + Dtype + Add<Self::Pointer<D>, Output = Self::Pointer<D>>;

    // Required methods
    fn program_id(axis: Axis) -> i32;
    fn num_programs(axis: Axis) -> i32;
    fn load_scalar_f32_as_i32(ptr: Self::Pointer<f32>, offset: i32) -> i32;
    fn arange(start: impl Into<i32>, end: impl Into<i32>) -> Self::I32Tensor;
    fn arange_f32(
        start: impl Into<i32>,
        end: impl Into<i32>,
    ) -> Self::Tensor<f32>;
    fn zeros<D: Dtype>(shape: &[i32]) -> Self::Tensor<D>;
    fn zeros_like<D: Dtype>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn full<D: Dtype>(shape: &[i32], value: D) -> Self::Tensor<D>;
    fn cast<Src: Dtype, Dst: Dtype>(
        x: Self::Tensor<Src>,
        fp_downcast_rounding: Option<FpDowncastRounding>,
        bitcast: bool,
    ) -> Self::Tensor<Dst>;
    fn cat<D: Dtype>(
        a: Self::Tensor<D>,
        b: Self::Tensor<D>,
        can_reorder: bool,
    ) -> Self::Tensor<D>;
    fn broadcast<D: Dtype>(
        a: Self::Tensor<D>,
        b: Self::Tensor<D>,
    ) -> (Self::Tensor<D>, Self::Tensor<D>);
    fn broadcast_to<D: Dtype>(
        x: Self::Tensor<D>,
        shape: &[i32],
    ) -> Self::Tensor<D>;
    fn expand_dims<D: Dtype>(x: Self::Tensor<D>, axis: i32) -> Self::Tensor<D>;
    fn permute<D: Dtype>(x: Self::Tensor<D>, dims: &[i32]) -> Self::Tensor<D>;
    fn reshape<D: Dtype>(
        x: Self::Tensor<D>,
        shape: &[i32],
        can_reorder: bool,
    ) -> Self::Tensor<D>;
    fn trans<D: Dtype>(x: Self::Tensor<D>, dims: &[i32]) -> Self::Tensor<D>;
    fn ravel<D: Dtype>(x: Self::Tensor<D>, can_reorder: bool) -> Self::Tensor<D>;
    fn view<D: Dtype>(x: Self::Tensor<D>, shape: &[i32]) -> Self::Tensor<D>;
    fn join<D: Dtype>(a: Self::Tensor<D>, b: Self::Tensor<D>) -> Self::Tensor<D>;
    fn interleave<D: Dtype>(
        a: Self::Tensor<D>,
        b: Self::Tensor<D>,
    ) -> Self::Tensor<D>;
    fn split<D: Dtype>(x: Self::Tensor<D>) -> (Self::Tensor<D>, Self::Tensor<D>);
    fn dot<D: Num, O: Num>(
        a: Self::Tensor<D>,
        b: Self::Tensor<D>,
        acc: Option<Self::Tensor<O>>,
        input_precision: InputPrecision,
        max_num_imprecise_acc: Option<i32>,
    ) -> Self::Tensor<O>;
    fn dot_scaled<D: Num, S: Num, O: Num>(
        lhs: Self::Tensor<D>,
        lhs_scale: Self::Tensor<S>,
        lhs_format: DotFormat,
        rhs: Self::Tensor<D>,
        rhs_scale: Self::Tensor<S>,
        rhs_format: DotFormat,
        acc: Option<Self::Tensor<O>>,
        fast_math: bool,
    ) -> Self::Tensor<O>;
    fn make_block_ptr<D: Dtype>(
        base: Self::Pointer<D>,
        shape: &[i32],
        strides: &[i32],
        offsets: &[i32],
        block_shape: &[i32],
        order: &[i32],
    ) -> Self::Pointer<D>;
    fn advance<D: Dtype>(
        ptr: Self::Pointer<D>,
        offsets: &[i32],
    ) -> Self::Pointer<D>;
    fn make_tensor_descriptor<D: Dtype>(
        base: Self::Pointer<D>,
        shape: &[i32],
        strides: &[i32],
        block_shape: &[i32],
        padding_option: Option<PaddingOption>,
    ) -> Self::Pointer<D>;
    fn load_tensor_descriptor<D: Dtype>(
        desc: Self::Pointer<D>,
        offsets: &[i32],
    ) -> Self::Tensor<D>;
    fn store_tensor_descriptor<D: Dtype>(
        desc: Self::Pointer<D>,
        offsets: &[i32],
        value: Self::Tensor<D>,
    );
    fn load<D: Dtype, const N: usize>(
        ptr: Self::Tensor<Self::Pointer<D>>,
        mask: Option<Self::BoolTensor>,
        other: Option<Self::Tensor<D>>,
        boundary_check: &[i32; N],
        padding_option: Option<PaddingOption>,
        cache_modifier: Option<CacheModifier>,
        eviction_policy: Option<EvictionPolicy>,
        volatile: bool,
    ) -> Self::Tensor<D>;
    fn store<D: Dtype, const N: usize>(
        dest: Self::Tensor<Self::Pointer<D>>,
        src: Self::Tensor<D>,
        mask: Option<Self::BoolTensor>,
        boundary_check: &[i32; N],
        cache_modifier: Option<CacheModifier>,
        eviction_policy: Option<EvictionPolicy>,
    );
    fn lt<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor;
    fn le<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor;
    fn gt<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor;
    fn ge<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor;
    fn eq<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor;
    fn ne<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor;
    fn lt_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor;
    fn le_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor;
    fn gt_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor;
    fn ge_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor;
    fn eq_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor;
    fn ne_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor;
    fn where_<D: Dtype>(
        cond: Self::BoolTensor,
        x: Self::Tensor<D>,
        y: Self::Tensor<D>,
    ) -> Self::Tensor<D>;
    fn flip<D: Dtype>(x: Self::Tensor<D>, dim: Option<i32>) -> Self::Tensor<D>;
    fn gather<D: Dtype>(
        src: Self::Tensor<D>,
        index: Self::I32Tensor,
        axis: i32,
    ) -> Self::Tensor<D>;
    fn abs<D: Dtype>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn ceil<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn floor<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn cos<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn sin<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn exp<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn exp2<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn log<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn log2<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn rsqrt<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn sigmoid<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn sqrt<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn sqrt_rn<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn erf<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn atan<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>;
    fn softmax<D: Float>(
        x: Self::Tensor<D>,
        dim: Option<i32>,
        keep_dims: bool,
        ieee_rounding: bool,
    ) -> Self::Tensor<D>;
    fn maximum<D: Num>(
        x: Self::Tensor<D>,
        y: Self::Tensor<D>,
    ) -> Self::Tensor<D>;
    fn minimum<D: Num>(
        x: Self::Tensor<D>,
        y: Self::Tensor<D>,
    ) -> Self::Tensor<D>;
    fn clamp<D: Num>(
        x: Self::Tensor<D>,
        lo: Self::Tensor<D>,
        hi: Self::Tensor<D>,
    ) -> Self::Tensor<D>;
    fn fma<D: Float>(
        x: Self::Tensor<D>,
        y: Self::Tensor<D>,
        z: Self::Tensor<D>,
    ) -> Self::Tensor<D>;
    fn fdiv<D: Float>(
        x: Self::Tensor<D>,
        y: Self::Tensor<D>,
        ieee_rounding: bool,
    ) -> Self::Tensor<D>;
    fn div_rn<D: Float>(
        x: Self::Tensor<D>,
        y: Self::Tensor<D>,
    ) -> Self::Tensor<D>;
    fn umulhi(x: Self::Tensor<u32>, y: Self::Tensor<u32>) -> Self::Tensor<u32>;
    fn cdiv(x: i32, div: i32) -> i32;
    fn swizzle2d(
        i: i32,
        j: i32,
        size_i: i32,
        size_j: i32,
        size_g: i32,
    ) -> (i32, i32);
    fn sum<D: Num>(
        x: Self::Tensor<D>,
        axis: Option<i32>,
        keep_dims: bool,
    ) -> Self::Tensor<D>;
    fn max<D: Num>(
        x: Self::Tensor<D>,
        axis: Option<i32>,
        keep_dims: bool,
    ) -> Self::Tensor<D>;
    fn max_with_indices<D: Num>(
        x: Self::Tensor<D>,
        axis: i32,
        tie_break_left: bool,
        keep_dims: bool,
    ) -> (Self::Tensor<D>, Self::I32Tensor);
    fn min<D: Num>(
        x: Self::Tensor<D>,
        axis: Option<i32>,
        keep_dims: bool,
    ) -> Self::Tensor<D>;
    fn min_with_indices<D: Num>(
        x: Self::Tensor<D>,
        axis: i32,
        tie_break_left: bool,
        keep_dims: bool,
    ) -> (Self::Tensor<D>, Self::I32Tensor);
    fn argmax<D: Num>(
        x: Self::Tensor<D>,
        axis: i32,
        tie_break_left: bool,
        keep_dims: bool,
    ) -> Self::I32Tensor;
    fn argmin<D: Num>(
        x: Self::Tensor<D>,
        axis: i32,
        tie_break_left: bool,
        keep_dims: bool,
    ) -> Self::I32Tensor;
    fn xor_sum<D: Int>(
        x: Self::Tensor<D>,
        axis: Option<i32>,
        keep_dims: bool,
    ) -> Self::Tensor<D>;
    fn cumsum<D: Num>(
        x: Self::Tensor<D>,
        axis: i32,
        reverse: bool,
    ) -> Self::Tensor<D>;
    fn cumprod<D: Num>(
        x: Self::Tensor<D>,
        axis: i32,
        reverse: bool,
    ) -> Self::Tensor<D>;
    fn sort<D: Num>(
        x: Self::Tensor<D>,
        dim: Option<i32>,
        descending: bool,
    ) -> Self::Tensor<D>;
    fn histogram(
        x: Self::I32Tensor,
        num_bins: i32,
        mask: Option<Self::BoolTensor>,
    ) -> Self::I32Tensor;
    fn reduce<D: Dtype, O: Dtype>(
        x: Self::Tensor<D>,
        axis: i32,
        combine_fn: fn(Self::Tensor<O>, Self::Tensor<O>) -> Self::Tensor<O>,
        keep_dims: bool,
    ) -> Self::Tensor<O>;
    fn associative_scan<D: Dtype>(
        x: Self::Tensor<D>,
        axis: i32,
        combine_fn: fn(Self::Tensor<D>, Self::Tensor<D>) -> Self::Tensor<D>,
        reverse: bool,
    ) -> Self::Tensor<D>;
    fn atomic_add<D: Num>(
        ptr: Self::Tensor<Self::Pointer<D>>,
        val: Self::Tensor<D>,
        mask: Option<Self::BoolTensor>,
        sem: Option<MemSem>,
        scope: Option<MemScope>,
    ) -> Self::Tensor<D>;
    fn atomic_and<D: Int>(
        ptr: Self::Tensor<Self::Pointer<D>>,
        val: Self::Tensor<D>,
        mask: Option<Self::BoolTensor>,
        sem: Option<MemSem>,
        scope: Option<MemScope>,
    ) -> Self::Tensor<D>;
    fn atomic_or<D: Int>(
        ptr: Self::Tensor<Self::Pointer<D>>,
        val: Self::Tensor<D>,
        mask: Option<Self::BoolTensor>,
        sem: Option<MemSem>,
        scope: Option<MemScope>,
    ) -> Self::Tensor<D>;
    fn atomic_xor<D: Int>(
        ptr: Self::Tensor<Self::Pointer<D>>,
        val: Self::Tensor<D>,
        mask: Option<Self::BoolTensor>,
        sem: Option<MemSem>,
        scope: Option<MemScope>,
    ) -> Self::Tensor<D>;
    fn atomic_max<D: Num>(
        ptr: Self::Tensor<Self::Pointer<D>>,
        val: Self::Tensor<D>,
        mask: Option<Self::BoolTensor>,
        sem: Option<MemSem>,
        scope: Option<MemScope>,
    ) -> Self::Tensor<D>;
    fn atomic_min<D: Num>(
        ptr: Self::Tensor<Self::Pointer<D>>,
        val: Self::Tensor<D>,
        mask: Option<Self::BoolTensor>,
        sem: Option<MemSem>,
        scope: Option<MemScope>,
    ) -> Self::Tensor<D>;
    fn atomic_xchg<D: Dtype>(
        ptr: Self::Tensor<Self::Pointer<D>>,
        val: Self::Tensor<D>,
        mask: Option<Self::BoolTensor>,
        sem: Option<MemSem>,
        scope: Option<MemScope>,
    ) -> Self::Tensor<D>;
    fn atomic_cas<D: Dtype>(
        ptr: Self::Tensor<Self::Pointer<D>>,
        cmp: Self::Tensor<D>,
        val: Self::Tensor<D>,
        sem: Option<MemSem>,
        scope: Option<MemScope>,
    ) -> Self::Tensor<D>;
    fn rand(
        seed: u32,
        offsets: Self::I32Tensor,
        n_rounds: i32,
    ) -> Self::Tensor<f32>;
    fn randn(
        seed: u32,
        offsets: Self::I32Tensor,
        n_rounds: i32,
    ) -> Self::Tensor<f32>;
    fn randint(
        seed: u32,
        offsets: Self::I32Tensor,
        n_rounds: i32,
    ) -> Self::I32Tensor;
    fn randint4x(
        seed: u32,
        offsets: Self::I32Tensor,
        n_rounds: i32,
    ) -> (Self::I32Tensor, Self::I32Tensor, Self::I32Tensor, Self::I32Tensor);
    fn inline_asm_elementwise<D: Dtype>(
        asm: &str,
        constraints: &str,
        is_pure: bool,
        pack: i32,
    ) -> Self::Tensor<D>;
    fn assume(cond: Self::BoolTensor);
    fn multiple_of<D: Dtype>(
        x: Self::Tensor<D>,
        values: &[i32],
    ) -> Self::Tensor<D>;
    fn max_contiguous<D: Dtype>(
        x: Self::Tensor<D>,
        values: &[i32],
    ) -> Self::Tensor<D>;
    fn max_constancy<D: Dtype>(
        x: Self::Tensor<D>,
        values: &[i32],
    ) -> Self::Tensor<D>;
    fn debug_barrier();
    fn device_assert(
        cond: Self::BoolTensor,
        msg: &str,
        mask: Option<Self::BoolTensor>,
    );
    fn device_print<D: Dtype>(prefix: &str, val: Self::Tensor<D>, hex: bool);
    fn static_assert(cond: bool, msg: &str);
    fn static_print(msg: &str);
}

The Triton-like kernel DSL: tensor/pointer types and the operations (creation, shape manipulation, linear algebra, memory, math, reduction, scan/sort, atomics, RNG) available inside a #[kernel]-annotated function. See the module docs for how this compiles.

Required Associated Types

type BoolTensor: Copy + Clone

A tensor of bool, produced by comparisons and used as a mask.

type I32Tensor: Copy + Clone

A tensor of i32, e.g. produced by Triton::arange.

type Tensor<D: Dtype>: Copy + Clone + Add<Self::Tensor<D>, Output = Self::Tensor<D>> + Sub<Self::Tensor<D>, Output = Self::Tensor<D>> + Mul<Self::Tensor<D>, Output = Self::Tensor<D>> + Div<Self::Tensor<D>, Output = Self::Tensor<D>> + Neg<Output = Self::Tensor<D>>

A tensor of dtype D.

type Pointer<D: Dtype>: Copy + Clone + Dtype + Add<Self::Pointer<D>, Output = Self::Pointer<D>>

A device pointer to elements of dtype D.

Required Methods

fn program_id(axis: Axis) -> i32

The current program’s index along axis within the launch grid.

fn num_programs(axis: Axis) -> i32

The total number of programs launched along axis.

fn load_scalar_f32_as_i32(ptr: Self::Pointer<f32>, offset: i32) -> i32

Scalar gather: load the f32 at ptr + offset, truncate to i32, and return it as a plain Rust i32 usable in scalar arithmetic (e.g. as an addend to arange results via I32Tensor + i32).

Used when integer indices are stored as f32 (the graph’s default dtype).

fn arange(start: impl Into<i32>, end: impl Into<i32>) -> Self::I32Tensor

Create a 1-D i32 tensor with values [start, start+1, ..., end-1].

fn arange_f32(start: impl Into<i32>, end: impl Into<i32>) -> Self::Tensor<f32>

Create a 1-D f32 tensor with values [start as f32, start+1, ..., end-1].

Equivalent to casting arange(start, end) to f32, but avoids the intermediate I32Tensor copy that some backends cannot handle.

fn zeros<D: Dtype>(shape: &[i32]) -> Self::Tensor<D>

Create a tensor of the given shape filled with zeros.

fn zeros_like<D: Dtype>(x: Self::Tensor<D>) -> Self::Tensor<D>

Create a zero-filled tensor with the same shape/dtype as x.

fn full<D: Dtype>(shape: &[i32], value: D) -> Self::Tensor<D>

Create a tensor of the given shape filled with value.

fn cast<Src: Dtype, Dst: Dtype>( x: Self::Tensor<Src>, fp_downcast_rounding: Option<FpDowncastRounding>, bitcast: bool, ) -> Self::Tensor<Dst>

Cast a tensor to a different dtype.

• fp_downcast_rounding: rounding mode when narrowing float types (default None = unspecified).
• bitcast: reinterpret bits without conversion (default false).

fn cat<D: Dtype>( a: Self::Tensor<D>, b: Self::Tensor<D>, can_reorder: bool, ) -> Self::Tensor<D>

Concatenate two tensors.

• can_reorder: allow the compiler to reorder elements (default false).

fn broadcast<D: Dtype>( a: Self::Tensor<D>, b: Self::Tensor<D>, ) -> (Self::Tensor<D>, Self::Tensor<D>)

Broadcast two tensors to a common compatible shape.

fn broadcast_to<D: Dtype>(x: Self::Tensor<D>, shape: &[i32]) -> Self::Tensor<D>

Broadcast x to shape.

fn expand_dims<D: Dtype>(x: Self::Tensor<D>, axis: i32) -> Self::Tensor<D>

Insert a size-1 dimension at axis.

fn permute<D: Dtype>(x: Self::Tensor<D>, dims: &[i32]) -> Self::Tensor<D>

Permute x’s dimensions according to dims.

fn reshape<D: Dtype>( x: Self::Tensor<D>, shape: &[i32], can_reorder: bool, ) -> Self::Tensor<D>

Reshape a tensor.

• can_reorder: allow element reordering during reshape (default false).

fn trans<D: Dtype>(x: Self::Tensor<D>, dims: &[i32]) -> Self::Tensor<D>

Permute dimensions. Alias for permute.

fn ravel<D: Dtype>(x: Self::Tensor<D>, can_reorder: bool) -> Self::Tensor<D>

Flatten to 1-D.

• can_reorder: allow element reordering (default false).

fn view<D: Dtype>(x: Self::Tensor<D>, shape: &[i32]) -> Self::Tensor<D>

View with a new shape (order not preserved).

fn join<D: Dtype>(a: Self::Tensor<D>, b: Self::Tensor<D>) -> Self::Tensor<D>

Join two tensors along a new minor dimension.

fn interleave<D: Dtype>( a: Self::Tensor<D>, b: Self::Tensor<D>, ) -> Self::Tensor<D>

Interleave two tensors along their last dimension.

fn split<D: Dtype>(x: Self::Tensor<D>) -> (Self::Tensor<D>, Self::Tensor<D>)

Split a tensor in two along its last dimension (which must have size 2).

fn dot<D: Num, O: Num>( a: Self::Tensor<D>, b: Self::Tensor<D>, acc: Option<Self::Tensor<O>>, input_precision: InputPrecision, max_num_imprecise_acc: Option<i32>, ) -> Self::Tensor<O>

Matrix (or batched matrix) multiply.

• acc: optional accumulator tensor added to the result.
• input_precision: Tensor Core precision for f32 × f32. No default — every call site must pick one explicitly. IEEE matches tt.dot’s own default (full-precision, no tensor cores); TF32 routes through tensor cores at reduced mantissa precision. A silent default here is exactly what caused kernels to unknowingly run at IEEE instead of the intended TF32.
• max_num_imprecise_acc: limit on imprecise accumulations (default None).

fn dot_scaled<D: Num, S: Num, O: Num>( lhs: Self::Tensor<D>, lhs_scale: Self::Tensor<S>, lhs_format: DotFormat, rhs: Self::Tensor<D>, rhs_scale: Self::Tensor<S>, rhs_format: DotFormat, acc: Option<Self::Tensor<O>>, fast_math: bool, ) -> Self::Tensor<O>

Scaled mixed-precision matrix multiply (FP8 / narrow formats).

• acc: optional accumulator (default None).
• fast_math: allow reduced precision accumulation (default false).

fn make_block_ptr<D: Dtype>( base: Self::Pointer<D>, shape: &[i32], strides: &[i32], offsets: &[i32], block_shape: &[i32], order: &[i32], ) -> Self::Pointer<D>

Create a block pointer encoding shape, strides, offsets, and tile shape.

fn advance<D: Dtype>(ptr: Self::Pointer<D>, offsets: &[i32]) -> Self::Pointer<D>

Advance a block pointer by the given per-dimension offsets.

fn make_tensor_descriptor<D: Dtype>( base: Self::Pointer<D>, shape: &[i32], strides: &[i32], block_shape: &[i32], padding_option: Option<PaddingOption>, ) -> Self::Pointer<D>

Create a tensor descriptor for TMA (Tensor Memory Accelerator) operations.

• padding_option: out-of-bounds padding behaviour (default PaddingOption::Zero).

fn load_tensor_descriptor<D: Dtype>( desc: Self::Pointer<D>, offsets: &[i32], ) -> Self::Tensor<D>

Load a tile from memory using a tensor descriptor and per-dimension offsets.

fn store_tensor_descriptor<D: Dtype>( desc: Self::Pointer<D>, offsets: &[i32], value: Self::Tensor<D>, )

Store a tile to memory using a tensor descriptor and per-dimension offsets.

fn load<D: Dtype, const N: usize>( ptr: Self::Tensor<Self::Pointer<D>>, mask: Option<Self::BoolTensor>, other: Option<Self::Tensor<D>>, boundary_check: &[i32; N], padding_option: Option<PaddingOption>, cache_modifier: Option<CacheModifier>, eviction_policy: Option<EvictionPolicy>, volatile: bool, ) -> Self::Tensor<D>

Load a tensor from memory.

• mask: when Some, lanes where mask is false are not loaded (default None = unconditional).
• other: fill value for masked-off lanes (default None = undefined).
• boundary_check: dimensions to check for out-of-bounds (block-pointer mode only, default &[]).
• padding_option: fill for out-of-bounds lanes in block-pointer mode (default None).
• cache_modifier: L1/L2 cache behaviour (default None).
• eviction_policy: eviction priority hint (default None).
• volatile: always fetch fresh from memory (default false).

fn store<D: Dtype, const N: usize>( dest: Self::Tensor<Self::Pointer<D>>, src: Self::Tensor<D>, mask: Option<Self::BoolTensor>, boundary_check: &[i32; N], cache_modifier: Option<CacheModifier>, eviction_policy: Option<EvictionPolicy>, )

Store a tensor to memory.

• mask: when Some, lanes where mask is false are not stored (default None = unconditional).
• boundary_check: dimensions to check for out-of-bounds (block-pointer mode only, default &[]).
• cache_modifier: L1/L2 cache behaviour (default None).
• eviction_policy: eviction priority hint (default None).

fn lt<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor

Element-wise less-than between two tensors.

fn le<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor

Element-wise less-than-or-equal between two tensors.

fn gt<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor

Element-wise greater-than between two tensors.

fn ge<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor

Element-wise greater-than-or-equal between two tensors.

fn eq<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor

Element-wise equality between two tensors.

fn ne<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::BoolTensor

Element-wise inequality between two tensors.

fn lt_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor

Element-wise less-than against a scalar.

fn le_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor

Element-wise less-than-or-equal against a scalar.

fn gt_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor

Element-wise greater-than against a scalar.

fn ge_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor

Element-wise greater-than-or-equal against a scalar.

fn eq_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor

Element-wise equality against a scalar.

fn ne_scalar<D: Num>(x: Self::Tensor<D>, y: D) -> Self::BoolTensor

Element-wise inequality against a scalar.

fn where_<D: Dtype>( cond: Self::BoolTensor, x: Self::Tensor<D>, y: Self::Tensor<D>, ) -> Self::Tensor<D>

Conditional element selection — corresponds to tl.where. Named where_ to avoid collision with the Rust keyword where.

fn flip<D: Dtype>(x: Self::Tensor<D>, dim: Option<i32>) -> Self::Tensor<D>

Reverse a tensor along dim. None reverses all dimensions.

fn gather<D: Dtype>( src: Self::Tensor<D>, index: Self::I32Tensor, axis: i32, ) -> Self::Tensor<D>

Gather elements from src along axis using index.

fn abs<D: Dtype>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise absolute value.

fn ceil<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise ceiling.

fn floor<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise floor.

fn cos<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise cosine.

fn sin<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise sine.

fn exp<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise natural exponential (e^x).

fn exp2<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise base-2 exponential (2^x).

fn log<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise natural logarithm.

fn log2<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise base-2 logarithm.

fn rsqrt<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise reciprocal square root (1/sqrt(x)).

fn sigmoid<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise sigmoid (1/(1+e^-x)).

fn sqrt<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise square root.

fn sqrt_rn<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise square root, round-to-nearest.

fn erf<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise error function.

fn atan<D: Float>(x: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise arctangent.

fn softmax<D: Float>( x: Self::Tensor<D>, dim: Option<i32>, keep_dims: bool, ieee_rounding: bool, ) -> Self::Tensor<D>

Numerically-stable softmax along dim. dim = None defaults to the last dimension.

• keep_dims: retain the reduced dimension with length 1 (default false).
• ieee_rounding: use IEEE-754 rounding (default false).

fn maximum<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise maximum of two tensors.

fn minimum<D: Num>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise minimum of two tensors.

fn clamp<D: Num>( x: Self::Tensor<D>, lo: Self::Tensor<D>, hi: Self::Tensor<D>, ) -> Self::Tensor<D>

Element-wise clamp of x to [lo, hi].

fn fma<D: Float>( x: Self::Tensor<D>, y: Self::Tensor<D>, z: Self::Tensor<D>, ) -> Self::Tensor<D>

Element-wise fused multiply-add: x * y + z.

fn fdiv<D: Float>( x: Self::Tensor<D>, y: Self::Tensor<D>, ieee_rounding: bool, ) -> Self::Tensor<D>

Element-wise floating-point division.

• ieee_rounding: use IEEE-754 rounding (default false).

fn div_rn<D: Float>(x: Self::Tensor<D>, y: Self::Tensor<D>) -> Self::Tensor<D>

Element-wise division, round-to-nearest.

fn umulhi(x: Self::Tensor<u32>, y: Self::Tensor<u32>) -> Self::Tensor<u32>

Element-wise high 32 bits of an unsigned 32×32→64-bit multiply.

fn cdiv(x: i32, div: i32) -> i32

Ceiling integer division: ceil(x / div).

fn swizzle2d( i: i32, j: i32, size_i: i32, size_j: i32, size_g: i32, ) -> (i32, i32)

Swizzle 2-D indices for shared-memory bank-conflict avoidance. Returns the remapped (i, j) indices.

fn sum<D: Num>( x: Self::Tensor<D>, axis: Option<i32>, keep_dims: bool, ) -> Self::Tensor<D>

Sum all elements along axis. axis = None reduces all dimensions.

fn max<D: Num>( x: Self::Tensor<D>, axis: Option<i32>, keep_dims: bool, ) -> Self::Tensor<D>

Maximum along axis. axis = None reduces all dimensions.

fn max_with_indices<D: Num>( x: Self::Tensor<D>, axis: i32, tie_break_left: bool, keep_dims: bool, ) -> (Self::Tensor<D>, Self::I32Tensor)

Maximum along axis, also returning the index of the maximum.

• tie_break_left: when true, the leftmost index wins on ties (default true).

fn min<D: Num>( x: Self::Tensor<D>, axis: Option<i32>, keep_dims: bool, ) -> Self::Tensor<D>

Minimum along axis. axis = None reduces all dimensions.

fn min_with_indices<D: Num>( x: Self::Tensor<D>, axis: i32, tie_break_left: bool, keep_dims: bool, ) -> (Self::Tensor<D>, Self::I32Tensor)

Minimum along axis, also returning the index of the minimum.

• tie_break_left: when true, the leftmost index wins on ties (default true).

fn argmax<D: Num>( x: Self::Tensor<D>, axis: i32, tie_break_left: bool, keep_dims: bool, ) -> Self::I32Tensor

Index of the maximum along axis.

• tie_break_left: when true, the leftmost index wins on ties (default true).

fn argmin<D: Num>( x: Self::Tensor<D>, axis: i32, tie_break_left: bool, keep_dims: bool, ) -> Self::I32Tensor

Index of the minimum along axis.

• tie_break_left: when true, the leftmost index wins on ties (default true).

fn xor_sum<D: Int>( x: Self::Tensor<D>, axis: Option<i32>, keep_dims: bool, ) -> Self::Tensor<D>

XOR-reduction along axis. axis = None reduces all dimensions.

fn cumsum<D: Num>( x: Self::Tensor<D>, axis: i32, reverse: bool, ) -> Self::Tensor<D>

Cumulative sum along axis.

fn cumprod<D: Num>( x: Self::Tensor<D>, axis: i32, reverse: bool, ) -> Self::Tensor<D>

Cumulative product along axis.

fn sort<D: Num>( x: Self::Tensor<D>, dim: Option<i32>, descending: bool, ) -> Self::Tensor<D>

Sort along dim. dim = None sorts along the last dimension.

fn histogram( x: Self::I32Tensor, num_bins: i32, mask: Option<Self::BoolTensor>, ) -> Self::I32Tensor

Compute a histogram with num_bins bins (width 1, starting at 0).

• mask: when Some, masked-off elements are excluded (default None).

fn reduce<D: Dtype, O: Dtype>( x: Self::Tensor<D>, axis: i32, combine_fn: fn(Self::Tensor<O>, Self::Tensor<O>) -> Self::Tensor<O>, keep_dims: bool, ) -> Self::Tensor<O>

Generic reduction along axis using a user-supplied combine function.

combine_fn must be a statically-known function pointer (corresponds to a @triton.jit-decorated helper in Python Triton).

fn associative_scan<D: Dtype>( x: Self::Tensor<D>, axis: i32, combine_fn: fn(Self::Tensor<D>, Self::Tensor<D>) -> Self::Tensor<D>, reverse: bool, ) -> Self::Tensor<D>

Generic prefix-scan along axis using a user-supplied combine function.

• reverse: scan in the reverse direction (default false).

fn atomic_add<D: Num>( ptr: Self::Tensor<Self::Pointer<D>>, val: Self::Tensor<D>, mask: Option<Self::BoolTensor>, sem: Option<MemSem>, scope: Option<MemScope>, ) -> Self::Tensor<D>

Atomic add. Returns the previous value.

• mask: when Some, only masked lanes perform the operation (default None).
• sem: memory ordering semantics (default None = AcqRel).
• scope: synchronization scope (default None = Gpu).

fn atomic_and<D: Int>( ptr: Self::Tensor<Self::Pointer<D>>, val: Self::Tensor<D>, mask: Option<Self::BoolTensor>, sem: Option<MemSem>, scope: Option<MemScope>, ) -> Self::Tensor<D>

Atomic bitwise AND. Returns the previous value. See Triton::atomic_add for parameters.

fn atomic_or<D: Int>( ptr: Self::Tensor<Self::Pointer<D>>, val: Self::Tensor<D>, mask: Option<Self::BoolTensor>, sem: Option<MemSem>, scope: Option<MemScope>, ) -> Self::Tensor<D>

Atomic bitwise OR. Returns the previous value. See Triton::atomic_add for parameters.

fn atomic_xor<D: Int>( ptr: Self::Tensor<Self::Pointer<D>>, val: Self::Tensor<D>, mask: Option<Self::BoolTensor>, sem: Option<MemSem>, scope: Option<MemScope>, ) -> Self::Tensor<D>

Atomic bitwise XOR. Returns the previous value. See Triton::atomic_add for parameters.

fn atomic_max<D: Num>( ptr: Self::Tensor<Self::Pointer<D>>, val: Self::Tensor<D>, mask: Option<Self::BoolTensor>, sem: Option<MemSem>, scope: Option<MemScope>, ) -> Self::Tensor<D>

Atomic maximum. Returns the previous value. See Triton::atomic_add for parameters.

fn atomic_min<D: Num>( ptr: Self::Tensor<Self::Pointer<D>>, val: Self::Tensor<D>, mask: Option<Self::BoolTensor>, sem: Option<MemSem>, scope: Option<MemScope>, ) -> Self::Tensor<D>

Atomic minimum. Returns the previous value. See Triton::atomic_add for parameters.

fn atomic_xchg<D: Dtype>( ptr: Self::Tensor<Self::Pointer<D>>, val: Self::Tensor<D>, mask: Option<Self::BoolTensor>, sem: Option<MemSem>, scope: Option<MemScope>, ) -> Self::Tensor<D>

Atomic exchange. Returns the previous value. See Triton::atomic_add for parameters.

fn atomic_cas<D: Dtype>( ptr: Self::Tensor<Self::Pointer<D>>, cmp: Self::Tensor<D>, val: Self::Tensor<D>, sem: Option<MemSem>, scope: Option<MemScope>, ) -> Self::Tensor<D>

Atomic compare-and-swap. Returns the previous value.

fn rand(seed: u32, offsets: Self::I32Tensor, n_rounds: i32) -> Self::Tensor<f32>

Uniform random f32 in [0, 1).

• n_rounds: number of Philox rounds (default 10).

fn randn( seed: u32, offsets: Self::I32Tensor, n_rounds: i32, ) -> Self::Tensor<f32>

Standard-normal random f32.

• n_rounds: number of Philox rounds (default 10).

fn randint( seed: u32, offsets: Self::I32Tensor, n_rounds: i32, ) -> Self::I32Tensor

Random i32.

• n_rounds: number of Philox rounds (default 10).

fn randint4x( seed: u32, offsets: Self::I32Tensor, n_rounds: i32, ) -> (Self::I32Tensor, Self::I32Tensor, Self::I32Tensor, Self::I32Tensor)

Four random i32 streams (maximally efficient Philox entry point).

• n_rounds: number of Philox rounds (default 10).

fn inline_asm_elementwise<D: Dtype>( asm: &str, constraints: &str, is_pure: bool, pack: i32, ) -> Self::Tensor<D>

Emit inline PTX/assembly applied element-wise across a tensor.

• asm: the assembly template string.
• constraints: register constraint string.
• is_pure: whether the assembly has no side-effects (may be CSE’d).
• pack: number of elements packed into each register.

fn assume(cond: Self::BoolTensor)

Assert that cond is always true, allowing the compiler to assume so.

fn multiple_of<D: Dtype>(x: Self::Tensor<D>, values: &[i32]) -> Self::Tensor<D>

Hint that values of x are always multiples of the given constants.

fn max_contiguous<D: Dtype>( x: Self::Tensor<D>, values: &[i32], ) -> Self::Tensor<D>

Hint that x has values[i] contiguous elements along dimension i.

fn max_constancy<D: Dtype>( x: Self::Tensor<D>, values: &[i32], ) -> Self::Tensor<D>

Hint that x has values[i] constant elements along dimension i.

fn debug_barrier()

Insert a memory barrier for debugging purposes.

fn device_assert( cond: Self::BoolTensor, msg: &str, mask: Option<Self::BoolTensor>, )

Emit a runtime assertion on the device. No-op when cond is true.

• msg: message shown on assertion failure (default "").
• mask: when Some, only lanes where mask is true check the assertion.

fn device_print<D: Dtype>(prefix: &str, val: Self::Tensor<D>, hex: bool)

Print a tensor value from device code for debugging.

• hex: print values in hexadecimal (default false).

fn static_assert(cond: bool, msg: &str)

Compile-time assertion (evaluated before kernel launch).

fn static_print(msg: &str)

Compile-time print (evaluated before kernel launch).

Dyn Compatibility

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors

impl Triton for LlvmTriton

type BoolTensor = LlvmTensor<bool>

type I32Tensor = LlvmTensor<i32>

type Tensor<D: Dtype> = LlvmTensor<D>

type Pointer<D: Dtype> = LlvmPointer<D>