simd

Module

Implements SIMD struct.

These are Mojo built-ins, so you don’t need to import them.

Aliases:

  • Int8 = SIMD[si8, 1]: Represents an 8-bit signed scalar integer.
  • UInt8 = SIMD[ui8, 1]: Represents an 8-bit unsigned scalar integer.
  • Int16 = SIMD[si16, 1]: Represents a 16-bit signed scalar integer.
  • UInt16 = SIMD[ui16, 1]: Represents a 16-bit unsigned scalar integer.
  • Int32 = SIMD[si32, 1]: Represents a 32-bit signed scalar integer.
  • UInt32 = SIMD[ui32, 1]: Represents a 32-bit unsigned scalar integer.
  • Int64 = SIMD[si64, 1]: Represents a 64-bit signed scalar integer.
  • UInt64 = SIMD[ui64, 1]: Represents a 64-bit unsigned scalar integer.
  • Float16 = SIMD[f16, 1]: Represents a 16-bit floating point value.
  • Float32 = SIMD[f32, 1]: Represents a 32-bit floating point value.
  • Float64 = SIMD[f64, 1]: Represents a 64-bit floating point value.

SIMD

Represents a small vector that is backed by a hardware vector element.

SIMD allows a single instruction to be executed across the multiple data elements of the vector.

Constraints:

The size of the SIMD vector to be positive and a power of 2.

Parameters:

  • type (DType): The data type of SIMD vector elements.
  • size (Int): The size of the SIMD vector.

Aliases:

  • element_type = _65x13_type

Fields:

  • value (simd<#lit.struct.extract<:!kgen.declref<@"$builtin"::@"$int"::@Int> size, "value">, #lit.struct.extract<:!kgen.declref<@"$builtin"::@"$dtype"::@DType> type, "value">>): The underlying storage for the vector.

Functions:

__init__

__init__() -> Self

Default initializer of the SIMD vector.

By default the SIMD vectors are initialized to all zeros.

Returns:

SIMD vector whose elements are 0.

__init__(value: Int) -> Self

Initializes the SIMD vector with an integer.

The integer value is splatted across all the elements of the SIMD vector.

Args:

  • value (Int): The input value.

Returns:

SIMD vector whose elements have the specified value.

__init__(value: Bool) -> Self

Initializes the SIMD vector with a bool value.

The bool value is splatted across all elements of the SIMD vector.

Args:

  • value (Bool): The bool value.

Returns:

SIMD vector whose elements have the specified value.

__init__(value: simd<#lit.struct.extract<:!kgen.declref<_"$builtin"::_"$int"::_Int> size, "value">, #lit.struct.extract<:!kgen.declref<_"$builtin"::_"$dtype"::_DType> type, "value">>) -> Self

Initializes the SIMD vector with the underlying mlir value.

Args:

  • value (simd<#lit.struct.extract<:!kgen.declref<_"$builtin"::_"$int"::_Int> size, "value">, #lit.struct.extract<:!kgen.declref<_"$builtin"::_"$dtype"::_DType> type, "value">>): The input value.

Returns:

SIMD vector using the specified value.

__init__(*elems: SIMD[type, 1]) -> Self

Constructs a SIMD vector via a variadic list of elements.

If there is just one input value, then it is splatted to all elements of the SIMD vector. Otherwise, the input values are assigned to the corresponding elements of the SIMD vector.

Constraints:

The number of input values is 1 or equal to size of the SIMD vector.

Args:

  • elems (*SIMD[type, 1]): The variadic list of elements from which the SIMD vector is constructed.

Returns:

The constructed SIMD vector.

__init__(value: FloatLiteral) -> Self

Initializes the SIMD vector with a FP64 value.

The input value is splatted (broadcast) across all the elements of the SIMD vector.

Args:

  • value (FloatLiteral): The input value.

Returns:

A SIMD vector whose elements have the specified value.

__bool__

__bool__(self: Self) -> Bool

Converts the SIMD vector into a boolean scalar value.

Returns:

True if all the elements in the SIMD vector are non-zero and False otherwise.

__getitem__

__getitem__(self: Self, idx: Int) -> SIMD[type, 1]

Gets an element from the vector.

Args:

  • idx (Int): The element index.

Returns:

The value at position idx.

__setitem__

__setitem__(inout self: Self, idx: Int, val: SIMD[type, 1])

Sets an element in the vector.

Args:

  • idx (Int): The index to set.
  • val (SIMD[type, 1]): The value to set.

__setitem__(inout self: Self, idx: Int, val: scalar<#lit.struct.extract<:!kgen.declref<_"$builtin"::_"$dtype"::_DType> type, "value">>)

Sets an element in the vector.

Args:

  • idx (Int): The index to set.
  • val (scalar<#lit.struct.extract<:!kgen.declref<_"$builtin"::_"$dtype"::_DType> type, "value">>): The value to set.

__neg__

__neg__(self: Self) -> Self

Defines the unary - operation.

Returns:

The negation of this SIMD vector.

__pos__

__pos__(self: Self) -> Self

Defines the unary + operation.

Returns:

This SIMD vector.

__invert__

__invert__(self: Self) -> Self

Returns ~self.

Constraints:

The element type of the SIMD vector must be boolean or integral.

Returns:

The ~self value.

__lt__

__lt__(self: Self, rhs: Self) -> SIMD[bool, size]

Compares two SIMD vectors using less-than comparison.

Args:

  • rhs (Self): The rhs of the operation.

Returns:

A new bool SIMD vector of the same size whose element at position i is True or False depending on the expression self[i] < rhs[i].

__le__

__le__(self: Self, rhs: Self) -> SIMD[bool, size]

Compares two SIMD vectors using less-than-or-equal comparison.

Args:

  • rhs (Self): The rhs of the operation.

Returns:

A new bool SIMD vector of the same size whose element at position i is True or False depending on the expression self[i] <= rhs[i].

__eq__

__eq__(self: Self, rhs: Self) -> SIMD[bool, size]

Compares two SIMD vectors using equal-to comparison.

Args:

  • rhs (Self): The rhs of the operation.

Returns:

A new bool SIMD vector of the same size whose element at position i is True or False depending on the expression self[i] == rhs[i].

__ne__

__ne__(self: Self, rhs: Self) -> SIMD[bool, size]

Compares two SIMD vectors using not-equal comparison.

Args:

  • rhs (Self): The rhs of the operation.

Returns:

A new bool SIMD vector of the same size whose element at position i is True or False depending on the expression self[i] != rhs[i].

__gt__

__gt__(self: Self, rhs: Self) -> SIMD[bool, size]

Compares two SIMD vectors using greater-than comparison.

Args:

  • rhs (Self): The rhs of the operation.

Returns:

A new bool SIMD vector of the same size whose element at position i is True or False depending on the expression self[i] > rhs[i].

__ge__

__ge__(self: Self, rhs: Self) -> SIMD[bool, size]

Compares two SIMD vectors using greater-than-or-equal comparison.

Args:

  • rhs (Self): The rhs of the operation.

Returns:

A new bool SIMD vector of the same size whose element at position i is True or False depending on the expression self[i] >= rhs[i].

__add__

__add__(self: Self, rhs: Self) -> Self

Computes self + rhs.

Args:

  • rhs (Self): The rhs value.

Returns:

A new vector whose element at position i is computed as self[i] + rhs[i].

__sub__

__sub__(self: Self, rhs: Self) -> Self

Computes self - rhs.

Args:

  • rhs (Self): The rhs value.

Returns:

A new vector whose element at position i is computed as self[i] - rhs[i].

__mul__

__mul__(self: Self, rhs: Self) -> Self

Computes self * rhs.

Args:

  • rhs (Self): The rhs value.

Returns:

A new vector whose element at position i is computed as self[i] * rhs[i].

__truediv__

__truediv__(self: Self, rhs: Self) -> Self

Computes self / rhs.

Args:

  • rhs (Self): The rhs value.

Returns:

A new vector whose element at position i is computed as self[i] / rhs[i].

__floordiv__

__floordiv__(self: Self, rhs: Self) -> Self

Returns the division of self and rhs rounded down to the nearest integer.

Constraints:

The element type of the SIMD vector must be integral.

Args:

  • rhs (Self): The value to divide on.

Returns:

floor(self / rhs) value.

__mod__

__mod__(self: Self, rhs: Self) -> Self

Returns the remainder of self divided by rhs.

Args:

  • rhs (Self): The value to divide on.

Returns:

The remainder of dividing self by rhs.

__pow__

__pow__(self: Self, rhs: Int) -> Self

Computes the vector raised to the power of the input integer value.

Args:

  • rhs (Int): The exponential value.

Returns:

A SIMD vector where each element is raised to the power of the specified exponential value.

__pow__[rhs_type: DType](self: Self, rhs: SIMD[rhs_type, size]) -> Self

Computes the vector raised to the power of the input integer value.

Parameters:

  • rhs_type (DType): The dtype of the rhs SIMD vector.

Args:

  • rhs (SIMD[rhs_type, size]): The exponential value.

Returns:

A SIMD vector where each element is raised to the power of the specified exponential value.

__lshift__

__lshift__(self: Self, rhs: Self) -> Self

Returns self << rhs.

Constraints:

The element type of the SIMD vector must be integral.

Args:

  • rhs (Self): The RHS value.

Returns:

self << rhs.

__rshift__

__rshift__(self: Self, rhs: Self) -> Self

Returns self >> rhs.

Constraints:

The element type of the SIMD vector must be integral.

Args:

  • rhs (Self): The RHS value.

Returns:

self >> rhs.

__and__

__and__(self: Self, rhs: Self) -> Self

Returns self & rhs.

Constraints:

The element type of the SIMD vector must be bool or integral.

Args:

  • rhs (Self): The RHS value.

Returns:

self & rhs.

__or__

__or__(self: Self, rhs: Self) -> Self

Returns self | rhs.

Constraints:

The element type of the SIMD vector must be bool or integral.

Args:

  • rhs (Self): The RHS value.

Returns:

self | rhs.

__xor__

__xor__(self: Self, rhs: Self) -> Self

Returns self ^ rhs.

Constraints:

The element type of the SIMD vector must be bool or integral.

Args:

  • rhs (Self): The RHS value.

Returns:

self ^ rhs.

__radd__

__radd__(self: Self, value: Self) -> Self

Returns value + self.

Args:

  • value (Self): The other value.

Returns:

value + self.

__rsub__

__rsub__(self: Self, value: Self) -> Self

Returns value - self.

Args:

  • value (Self): The other value.

Returns:

value - self.

__rmul__

__rmul__(self: Self, value: Self) -> Self

Returns value * self.

Args:

  • value (Self): The other value.

Returns:

value * self.

__rtruediv__

__rtruediv__(self: Self, value: Self) -> Self

Returns value / self.

Args:

  • value (Self): The other value.

Returns:

value / self.

__rfloordiv__

__rfloordiv__(self: Self, value: Int) -> Int

Returns value // self.

Args:

  • value (Int): The other value.

Returns:

value // self.

__rmod__

__rmod__(self: Self, value: Int) -> Int

Returns value % self.

Args:

  • value (Int): The other value.

Returns:

value % self.

__rlshift__

__rlshift__(self: Self, value: Self) -> Self

Returns value << self.

Constraints:

The element type of the SIMD vector must be integral.

Args:

  • value (Self): The other value.

Returns:

value << self.

__rrshift__

__rrshift__(self: Self, value: Self) -> Self

Returns value >> self.

Constraints:

The element type of the SIMD vector must be integral.

Args:

  • value (Self): The other value.

Returns:

value >> self.

__rand__

__rand__(self: Self, value: Self) -> Self

Returns value & self.

Constraints:

The element type of the SIMD vector must be bool or integral.

Args:

  • value (Self): The other value.

Returns:

value & self.

__ror__

__ror__(self: Self, value: Self) -> Self

Returns value | self.

Constraints:

The element type of the SIMD vector must be bool or integral.

Args:

  • value (Self): The other value.

Returns:

value | self.

__rxor__

__rxor__(self: Self, value: Self) -> Self

Returns value ^ self.

Constraints:

The element type of the SIMD vector must be bool or integral.

Args:

  • value (Self): The other value.

Returns:

value ^ self.

__iadd__

__iadd__(inout self: Self, rhs: Self)

Performs in-place addition.

The vector is mutated where each element at position i is computed as self[i] + rhs[i].

Args:

  • rhs (Self): The rhs of the addition operation.

__isub__

__isub__(inout self: Self, rhs: Self)

Performs in-place subtraction.

The vector is mutated where each element at position i is computed as self[i] - rhs[i].

Args:

  • rhs (Self): The rhs of the operation.

__imul__

__imul__(inout self: Self, rhs: Self)

Performs in-place multiplication.

The vector is mutated where each element at position i is computed as self[i] * rhs[i].

Args:

  • rhs (Self): The rhs of the operation.

__itruediv__

__itruediv__(inout self: Self, rhs: Self)

In-place true divide operator.

The vector is mutated where each element at position i is computed as self[i] / rhs[i].

Args:

  • rhs (Self): The rhs of the operation.

__ifloordiv__

__ifloordiv__(inout self: Self, rhs: Self)

In-place flood div operator.

The vector is mutated where each element at position i is computed as self[i] // rhs[i].

Args:

  • rhs (Self): The rhs of the operation.

__imod__

__imod__(inout self: Self, rhs: Self)

In-place mod operator.

The vector is mutated where each element at position i is computed as self[i] % rhs[i].

Args:

  • rhs (Self): The rhs of the operation.

__ipow__

__ipow__(inout self: Self, rhs: Int)

In-place pow operator.

The vector is mutated where each element at position i is computed as pow(self[i], rhs).

Args:

  • rhs (Int): The rhs of the operation.

__ilshift__

__ilshift__(inout self: Self, rhs: Self)

Computes self << rhs and save the result in self.

Constraints:

The element type of the SIMD vector must be integral.

Args:

  • rhs (Self): The RHS value.

__irshift__

__irshift__(inout self: Self, rhs: Self)

Computes self >> rhs and save the result in self.

Constraints:

The element type of the SIMD vector must be integral.

Args:

  • rhs (Self): The RHS value.

__iand__

__iand__(inout self: Self, rhs: Self)

Computes self & rhs and save the result in self.

Constraints:

The element type of the SIMD vector must be bool or integral.

Args:

  • rhs (Self): The RHS value.

__ixor__

__ixor__(inout self: Self, rhs: Self)

Computes self ^ rhs and save the result in self.

Constraints:

The element type of the SIMD vector must be bool or integral.

Args:

  • rhs (Self): The RHS value.

__ior__

__ior__(inout self: Self, rhs: Self)

Computes self | rhs and save the result in self.

Constraints:

The element type of the SIMD vector must be bool or integral.

Args:

  • rhs (Self): The RHS value.

__len__

__len__(self: Self) -> Int

Gets the length of the SIMD vector.

Returns:

The length of the SIMD vector.

splat

splat(x: Bool) -> Self

Splats (broadcasts) the element onto the vector.

Args:

  • x (Bool): The input value.

Returns:

A new SIMD vector whose elements are the same as the input value.

splat(x: SIMD[type, 1]) -> Self

Splats (broadcasts) the element onto the vector.

Args:

  • x (SIMD[type, 1]): The input scalar value.

Returns:

A new SIMD vector whose elements are the same as the input value.

cast

cast[target: DType](self: Self) -> SIMD[target, size]

Casts the elements of the SIMD vector to the target element type.

Parameters:

  • target (DType): The target DType.

Returns:

A new SIMD vector whose elements have been casted to the target element type.

__int__

__int__(self: Self) -> Int

Casts to the value to an Int. If there is a fractional component, then the value is truncated towards zero.

Constraints:

The size of the SIMD vector must be 1.

Returns:

The value as an integer.

to_int

to_int(self: Self) -> Int

Casts to the value to an Int. If there is a fractional component, then the value is truncated towards zero.

Constraints:

The size of the SIMD vector must be 1.

Returns:

The value of the single integer element in the SIMD vector.

fma

fma(self: Self, multiplier: Self, accumulator: Self) -> Self

Performs a fused multiply-add operation, i.e. self*multiplier + accumulator.

Args:

  • multiplier (Self): The value to multiply.
  • accumulator (Self): The value to accumulate.

Returns:

A new vector whose element at position i is computed as self[i]*multiplier[i] + accumulator[i].

shuffle

shuffle[*mask: Int](self: Self) -> Self

Shuffles (also called blend) the values of the current vector with the other value using the specified mask (permutation).

Parameters:

  • mask (*Int): The permutation to use in the shuffle.

Returns:

A new vector of length len where the value at position i is (self)[permutation[i]].

shuffle[*mask: Int](self: Self, other: Self) -> Self

Shuffles (also called blend) the values of the current vector with the other value using the specified mask (permutation).

Parameters:

  • mask (*Int): The permutation to use in the shuffle.

Args:

  • other (Self): The other vector to shuffle with.

Returns:

A new vector of length len where the value at position i is (self+other)[permutation[i]].

slice

slice[output_width: Int](self: Self, offset: Int) -> SIMD[type, output_width]

Returns a slice of the vector of the specified width with the given offset.

Constraints:

output_width + offset must not exceed the size of this SIMD vector.

Parameters:

  • output_width (Int): The output SIMD vector size.

Args:

  • offset (Int): The given offset for the slice.

Returns:

A new vector whose elements map to self[offset:offset+output_width].

min

min(self: Self, other: Self) -> Self

Computes the elementwise minimum between the two vectors.

Args:

  • other (Self): The other SIMD vector.

Returns:

A new SIMD vector where each element at position i is min(self[i], other[i]).

max

max(self: Self, other: Self) -> Self

Computes the elementwise maximum between the two vectors.

Args:

  • other (Self): The other SIMD vector.

Returns:

A new SIMD vector where each element at position i is max(self[i], other[i]).

reduce

reduce[func: fn[DType, Int](SIMD[*(0,0), *(0,1)], SIMD[*(0,0), *(0,1)]) capturing -> SIMD[*(0,0), *(0,1)]](self: Self) -> SIMD[type, 1]

Reduces the vector using a provided reduce operator.

Parameters:

  • func (fn[DType, Int](SIMD[*(0,0), *(0,1)], SIMD[*(0,0), *(0,1)]) capturing -> SIMD[*(0,0), *(0,1)]): The reduce function to apply to elements in this SIMD.

Returns:

A new scalar which is the reduction of all vector elements.

reduce_max

reduce_max(self: Self) -> SIMD[type, 1]

Reduces the vector using the max operator.

Constraints:

The element type of the vector must be integer or FP.

Returns:

The maximum element of the vector.

reduce_min

reduce_min(self: Self) -> SIMD[type, 1]

Reduces the vector using the min operator.

Constraints:

The element type of the vector must be integer or FP.

Returns:

The minimum element of the vector.

reduce_add

reduce_add(self: Self) -> SIMD[type, 1]

Reduces the vector using the add operator.

Returns:

The sum of all vector elements.

reduce_mul

reduce_mul(self: Self) -> SIMD[type, 1]

Reduces the vector using the mul operator.

Constraints:

The element type of the vector must be integer or FP.

Returns:

The product of all vector elements.

reduce_and

reduce_and(self: Self) -> Bool

Reduces the boolean vector using the and operator.

Constraints:

The element type of the vector must be boolean.

Returns:

True if all element in the vector is True and False otherwise.

reduce_or

reduce_or(self: Self) -> Bool

Reduces the boolean vector using the or operator.

Constraints:

The element type of the vector must be boolean.

Returns:

True if any element in the vector is True and False otherwise.

select

select[result_type: DType](self: Self, true_case: SIMD[result_type, size], false_case: SIMD[result_type, size]) -> SIMD[result_type, size]

Selects the values of the true_case or the false_case based on the current boolean values of the SIMD vector.

Parameters:

  • result_type (DType): The element type of the input and output SIMD vectors.

Args:

  • true_case (SIMD[result_type, size]): The values selected if the positional value is True.
  • false_case (SIMD[result_type, size]): The values selected if the positional value is False.

Returns:

A new vector of the form [true_case[i] if elem else false_case[i] in enumerate(self)].

rotate_left

rotate_left[shift: Int](self: Self) -> Self

Shifts the elements of a SIMD vector to the left by shift elements (with wrap-around).

Constraints:

-size <= shift < size

Parameters:

  • shift (Int): The number of positions by which to rotate the elements of SIMD vector to the left (with wrap-around).

Returns:

The SIMD vector rotated to the left by shift elements (with wrap-around).

rotate_right

rotate_right[shift: Int](self: Self) -> Self

Shifts the elements of a SIMD vector to the right by shift elements (with wrap-around).

Constraints:

-size < shift <= size

Parameters:

  • shift (Int): The number of positions by which to rotate the elements of SIMD vector to the right (with wrap-around).

Returns:

The SIMD vector rotated to the right by shift elements (with wrap-around).

shift_left

shift_left[shift: Int](self: Self) -> Self

Shifts the elements of a SIMD vector to the left by shift elements (no wrap-around, fill with zero).

Constraints:

0 <= shift <= size

Parameters:

  • shift (Int): The number of positions by which to rotate the elements of SIMD vector to the left (no wrap-around, fill with zero).

Returns:

The SIMD vector rotated to the left by shift elements (no wrap-around, fill with zero).

shift_right

shift_right[shift: Int](self: Self) -> Self

Shifts the elements of a SIMD vector to the right by shift elements (no wrap-around, fill with zero).

Constraints:

0 <= shift <= size

Parameters:

  • shift (Int): The number of positions by which to rotate the elements of SIMD vector to the right (no wrap-around, fill with zero).

Returns:

The SIMD vector rotated to the right by shift elements (no wrap-around, fill with zero).