Kernel Reference#

Scalar Types#

class warp.int8[source]#

class warp.uint8[source]#

class warp.int16[source]#

class warp.uint16[source]#

class warp.int32[source]#

class warp.uint32[source]#

class warp.int64[source]#

class warp.uint64[source]#

class warp.float16[source]#

class warp.float32[source]#

class warp.float64[source]#

class warp.bool[source]#

Vector Types#

class warp.vec2b[source]#

class warp.vec2ub[source]#

class warp.vec2s[source]#

class warp.vec2us[source]#

class warp.vec2i[source]#

class warp.vec2ui[source]#

class warp.vec2l[source]#

class warp.vec2ul[source]#

class warp.vec2h[source]#

class warp.vec2f[source]#

class warp.vec2d[source]#

class warp.vec3b[source]#

class warp.vec3ub[source]#

class warp.vec3s[source]#

class warp.vec3us[source]#

class warp.vec3i[source]#

class warp.vec3ui[source]#

class warp.vec3l[source]#

class warp.vec3ul[source]#

class warp.vec3h[source]#

class warp.vec3f[source]#

class warp.vec3d[source]#

class warp.vec4b[source]#

class warp.vec4ub[source]#

class warp.vec4s[source]#

class warp.vec4us[source]#

class warp.vec4i[source]#

class warp.vec4ui[source]#

class warp.vec4l[source]#

class warp.vec4ul[source]#

class warp.vec4h[source]#

class warp.vec4f[source]#

class warp.vec4d[source]#

class warp.mat22h[source]#

class warp.mat22f[source]#

class warp.mat22d[source]#

class warp.mat33h[source]#

class warp.mat33f[source]#

class warp.mat33d[source]#

class warp.mat44h[source]#

class warp.mat44f[source]#

class warp.mat44d[source]#

class warp.quath[source]#

class warp.quatf[source]#

class warp.quatd[source]#

class warp.transformh[source]#

class warp.transformf[source]#

class warp.transformd[source]#

class warp.spatial_vectorh[source]#

class warp.spatial_vectorf[source]#

class warp.spatial_vectord[source]#

class warp.spatial_matrixh[source]#

class warp.spatial_matrixf[source]#

class warp.spatial_matrixd[source]#

Generic Types#

class warp.Int#

class warp.Float#

class warp.Scalar#

class warp.Vector#

class warp.Matrix#

class warp.Quaternion#

class warp.Transformation#

class warp.Array#

Scalar Math#

warp.min(a: Scalar, b: Scalar) → Scalar#: Return the minimum of two scalars.

warp.min( a: Vector[Any, Scalar], b: Vector[Any, Scalar], ) → Vector[Any, Scalar]: Return the element-wise minimum of two vectors.

warp.min(a: Vector[Any, Scalar]) → Scalar: Return the minimum element of a vector a.

warp.max(a: Scalar, b: Scalar) → Scalar#: Return the maximum of two scalars.

warp.max( a: Vector[Any, Scalar], b: Vector[Any, Scalar], ) → Vector[Any, Scalar]: Return the element-wise maximum of two vectors.

warp.max(a: Vector[Any, Scalar]) → Scalar: Return the maximum element of a vector a.

warp.clamp(x: Scalar, low: Scalar, high: Scalar) → Scalar#: Clamp the value of x to the range [low, high].

warp.abs(x: Scalar) → Scalar#: Return the absolute value of x.

warp.abs(x: Vector[Any, Scalar]) → Vector[Any, Scalar]: Return the absolute values of the elements of x.

warp.sign(x: Scalar) → Scalar#: Return -1 if x < 0, return 1 otherwise.

warp.sign(x: Vector[Any, Scalar]) → Scalar: Return -1 for the negative elements of x, and 1 otherwise.

warp.step(x: Scalar) → Scalar#: Return 1.0 if x < 0.0, return 0.0 otherwise.

warp.nonzero(x: Scalar) → Scalar#: Return 1.0 if x is not equal to zero, return 0.0 otherwise.

warp.sin(x: Float) → Float#: Return the sine of x in radians.

warp.cos(x: Float) → Float#: Return the cosine of x in radians.

warp.acos(x: Float) → Float#: Return arccos of x in radians. Inputs are automatically clamped to [-1.0, 1.0].

warp.asin(x: Float) → Float#: Return arcsin of x in radians. Inputs are automatically clamped to [-1.0, 1.0].

warp.sqrt(x: Float) → Float#: Return the square root of x, where x is positive.

warp.cbrt(x: Float) → Float#: Return the cube root of x.

warp.tan(x: Float) → Float#: Return the tangent of x in radians.

warp.atan(x: Float) → Float#: Return the arctangent of x in radians.

warp.atan2(y: Float, x: Float) → Float#: Return the 2-argument arctangent, atan2, of the point (x, y) in radians.

warp.sinh(x: Float) → Float#: Return the sinh of x.

warp.cosh(x: Float) → Float#: Return the cosh of x.

warp.tanh(x: Float) → Float#: Return the tanh of x.

warp.degrees(x: Float) → Float#: Convert x from radians into degrees.

warp.radians(x: Float) → Float#: Convert x from degrees into radians.

warp.log(x: Float) → Float#: Return the natural logarithm (base-e) of x, where x is positive.

warp.log2(x: Float) → Float#: Return the binary logarithm (base-2) of x, where x is positive.

warp.log10(x: Float) → Float#: Return the common logarithm (base-10) of x, where x is positive.

warp.exp(x: Float) → Float#: Return the value of the exponential function \(e^x\).

warp.pow(x: Float, y: Float) → Float#: Return the result of x raised to power of y.

warp.round(x: Float) → Float#

Return the nearest integer value to x, rounding halfway cases away from zero.

This is the most intuitive form of rounding in the colloquial sense, but can be slower than other options like warp.rint(). Differs from numpy.round(), which behaves the same way as numpy.rint().

warp.rint(x: Float) → Float#

Return the nearest integer value to x, rounding halfway cases to nearest even integer.

It is generally faster than warp.round(). Equivalent to numpy.rint().

warp.trunc(x: Float) → Float#

Return the nearest integer that is closer to zero than x.

In other words, it discards the fractional part of x. It is similar to casting float(int(a)), but preserves the negative sign when x is in the range [-0.0, -1.0). Equivalent to numpy.trunc() and numpy.fix().

warp.floor(x: Float) → Float#: Return the largest integer that is less than or equal to x.

warp.ceil(x: Float) → Float#: Return the smallest integer that is greater than or equal to x.

warp.frac(x: Float) → Float#

Retrieve the fractional part of x.

In other words, it discards the integer part of x and is equivalent to x - trunc(x).

warp.isfinite(a: Scalar) → bool#: Return True if a is a finite number, otherwise return False.

warp.isfinite(a: Vector[Any, Scalar]) → bool: Return True if all elements of the vector a are finite, otherwise return False.

warp.isfinite(a: Quaternion[Scalar]) → bool: Return True if all elements of the quaternion a are finite, otherwise return False.

warp.isfinite(a: Matrix[Any, Any, Scalar]) → bool: Return True if all elements of the matrix a are finite, otherwise return False.

warp.isnan(a: Scalar) → bool#: Return True if a is NaN, otherwise return False.

warp.isnan(a: Vector[Any, Scalar]) → bool: Return True if any element of the vector a is NaN, otherwise return False.

warp.isnan(a: Quaternion[Scalar]) → bool: Return True if any element of the quaternion a is NaN, otherwise return False.

warp.isnan(a: Matrix[Any, Any, Scalar]) → bool: Return True if any element of the matrix a is NaN, otherwise return False.

warp.isinf(a: Scalar) → bool#: Return True if a is positive or negative infinity, otherwise return False.

warp.isinf(a: Vector[Any, Scalar]) → bool: Return True if any element of the vector a is positive or negative infinity, otherwise return False.

warp.isinf(a: Quaternion[Scalar]) → bool: Return True if any element of the quaternion a is positive or negative infinity, otherwise return False.

warp.isinf(a: Matrix[Any, Any, Scalar]) → bool: Return True if any element of the matrix a is positive or negative infinity, otherwise return False.

Vector Math#

warp.dot(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) → Scalar#: Compute the dot product between two vectors.

warp.dot(a: Quaternion[Float], b: Quaternion[Float]) → Float: Compute the dot product between two quaternions.

warp.ddot(a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar]) → Scalar#: Compute the double dot product between two matrices.

warp.argmin(a: Vector[Any, Scalar]) → uint32#: Return the index of the minimum element of a vector a. [1]

warp.argmax(a: Vector[Any, Scalar]) → uint32#: Return the index of the maximum element of a vector a. [1]

warp.outer( a: Vector[Any, Scalar], b: Vector[Any, Scalar], ) → Matrix[Any, Any, Scalar]#: Compute the outer product a*b^T for two vectors.

warp.cross(a: Vector[3, Scalar], b: Vector[3, Scalar]) → Vector[3, Scalar]#: Compute the cross product of two 3D vectors.

warp.skew(vec: Vector[3, Scalar]) → Matrix[3, 3, Scalar]#: Compute the skew-symmetric 3x3 matrix for a 3D vector vec.

warp.length(a: Vector[Any, Float]) → Float#: Compute the length of a floating-point vector a.

warp.length(a: Quaternion[Float]) → Float: Compute the length of a quaternion a.

warp.length_sq(a: Vector[Any, Scalar]) → Scalar#: Compute the squared length of a vector a.

warp.length_sq(a: Quaternion[Scalar]) → Scalar: Compute the squared length of a quaternion a.

warp.normalize(a: Vector[Any, Float]) → Vector[Any, Float]#: Compute the normalized value of a. If length(a) is 0 then the zero vector is returned.

warp.normalize(a: Quaternion[Float]) → Quaternion[Float]: Compute the normalized value of a. If length(a) is 0, then the zero quaternion is returned.

warp.transpose(a: Matrix[Any, Any, Scalar]) → Matrix[Any, Any, Scalar]#: Return the transpose of the matrix a.

warp.inverse(a: Matrix[2, 2, Float]) → Matrix[Any, Any, Float]#: Return the inverse of a 2x2 matrix a.

warp.inverse(a: Matrix[3, 3, Float]) → Matrix[Any, Any, Float]: Return the inverse of a 3x3 matrix a.

warp.inverse(a: Matrix[4, 4, Float]) → Matrix[Any, Any, Float]: Return the inverse of a 4x4 matrix a.

warp.determinant(a: Matrix[2, 2, Float]) → Float#: Return the determinant of a 2x2 matrix a.

warp.determinant(a: Matrix[3, 3, Float]) → Float: Return the determinant of a 3x3 matrix a.

warp.determinant(a: Matrix[4, 4, Float]) → Float: Return the determinant of a 4x4 matrix a.

warp.trace(a: Matrix[Any, Any, Scalar]) → Scalar#: Return the trace of the matrix a.

warp.diag(vec: Vector[Any, Scalar]) → Matrix[Any, Any, Scalar]#: Returns a matrix with the components of the vector vec on the diagonal.

warp.get_diag(mat: Matrix[Any, Any, Scalar]) → Vector[Any, Scalar]#: Returns a vector containing the diagonal elements of the square matrix mat.

warp.cw_mul( a: Vector[Any, Scalar], b: Vector[Any, Scalar], ) → Vector[Any, Scalar]#: Component-wise multiplication of two vectors.

warp.cw_mul( a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar], ) → Matrix[Any, Any, Scalar]: Component-wise multiplication of two matrices.

warp.cw_div( a: Vector[Any, Scalar], b: Vector[Any, Scalar], ) → Vector[Any, Scalar]#: Component-wise division of two vectors.

warp.cw_div( a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar], ) → Matrix[Any, Any, Scalar]: Component-wise division of two matrices.

warp.vector( *args: Scalar, length: int32, dtype: Scalar, ) → Vector[Any, Scalar]#: Construct a vector of given length and dtype.

warp.matrix( pos: Vector[3, Float], rot: Quaternion[Float], scale: Vector[3, Float], dtype: Float, ) → Matrix[4, 4, Float]#: Construct a 4x4 transformation matrix that applies the transformations as Translation(pos)*Rotation(rot)*Scaling(scale) when applied to column vectors, i.e.: y = (TRS)*x

warp.matrix( *args: Scalar, shape: Tuple[int, int], dtype: Scalar, ) → Matrix[Any, Any, Scalar]: Construct a matrix. If the positional arg_types are not given, then matrix will be zero-initialized.

warp.identity(n: int32, dtype: Scalar) → Matrix[Any, Any, Scalar]#: Create an identity matrix with shape=(n,n) with the type given by dtype.

warp.svd3( A: Matrix[3, 3, Float], U: Matrix[3, 3, Float], sigma: Vector[3, Float], V: Matrix[3, 3, Scalar], ) → None#: Compute the SVD of a 3x3 matrix A. The singular values are returned in sigma, while the left and right basis vectors are returned in U and V.

warp.qr3( A: Matrix[3, 3, Float], Q: Matrix[3, 3, Float], R: Matrix[3, 3, Float], ) → None#: Compute the QR decomposition of a 3x3 matrix A. The orthogonal matrix is returned in Q, while the upper triangular matrix is returned in R.

warp.eig3( A: Matrix[3, 3, Float], Q: Matrix[3, 3, Float], d: Vector[3, Float], ) → None#: Compute the eigendecomposition of a 3x3 matrix A. The eigenvectors are returned as the columns of Q, while the corresponding eigenvalues are returned in d.

Quaternion Math#

warp.quaternion(dtype: Float) → Quaternion[Float]#: Construct a zero-initialized quaternion. Quaternions are laid out as [ix, iy, iz, r], where ix, iy, iz are the imaginary part, and r the real part.

warp.quaternion( x: Float, y: Float, z: Float, w: Float, ) → Quaternion[Float]: Create a quaternion using the supplied components (type inferred from component type).

warp.quaternion( ijk: Vector[3, Float], real: Float, dtype: Float, ) → Quaternion[Float]: Create a quaternion using the supplied vector/scalar (type inferred from scalar type).

warp.quaternion(quat: Quaternion[Float], dtype: Float) → Quaternion[Float]: Construct a quaternion of type dtype from another quaternion of a different dtype.

warp.quat_identity(dtype: Float) → quatf#: Construct an identity quaternion with zero imaginary part and real part of 1.0

warp.quat_from_axis_angle( axis: Vector[3, Float], angle: Float, ) → Quaternion[Float]#: Construct a quaternion representing a rotation of angle radians around the given axis.

warp.quat_to_axis_angle( quat: Quaternion[Float], axis: Vector[3, Float], angle: Float, ) → None#: Extract the rotation axis and angle radians a quaternion represents.

warp.quat_from_matrix(mat: Matrix[3, 3, Float]) → Quaternion[Float]#: Construct a quaternion from a 3x3 matrix.

warp.quat_rpy(roll: Float, pitch: Float, yaw: Float) → Quaternion[Float]#: Construct a quaternion representing a combined roll (z), pitch (x), yaw rotations (y) in radians.

warp.quat_inverse(quat: Quaternion[Float]) → Quaternion[Float]#: Compute quaternion conjugate.

warp.quat_rotate( quat: Quaternion[Float], vec: Vector[3, Float], ) → Vector[3, Float]#: Rotate a vector by a quaternion.

warp.quat_rotate_inv( quat: Quaternion[Float], vec: Vector[3, Float], ) → Vector[3, Float]#: Rotate a vector by the inverse of a quaternion.

warp.quat_slerp( a: Quaternion[Float], b: Quaternion[Float], t: Float, ) → Quaternion[Float]#: Linearly interpolate between two quaternions.

warp.quat_to_matrix(quat: Quaternion[Float]) → Matrix[3, 3, Float]#: Convert a quaternion to a 3x3 rotation matrix.

Transformations#

warp.transformation( pos: Vector[3, Float], rot: Quaternion[Float], dtype: Float, ) → Transformation[Float]#: Construct a rigid-body transformation with translation part pos and rotation rot.

warp.transform_identity(dtype: Float) → transformf#: Construct an identity transform with zero translation and identity rotation.

warp.transform_get_translation( xform: Transformation[Float], ) → Vector[3, Float]#: Return the translational part of a transform xform.

warp.transform_get_rotation( xform: Transformation[Float], ) → Quaternion[Float]#: Return the rotational part of a transform xform.

warp.transform_multiply( a: Transformation[Float], b: Transformation[Float], ) → Transformation[Float]#: Multiply two rigid body transformations together.

warp.transform_point( xform: Transformation[Float], point: Vector[3, Float], ) → Vector[3, Float]#: Apply the transform to a point point treating the homogeneous coordinate as w=1 (translation and rotation).

warp.transform_point( mat: Matrix[4, 4, Float], point: Vector[3, Float], ) → Vector[3, Float]

Apply the transform to a point point treating the homogeneous coordinate as w=1.

The transformation is applied treating point as a column vector, e.g.: y = mat*point.

This is in contrast to some libraries, notably USD, which applies transforms to row vectors, y^T = point^T*mat^T. If the transform is coming from a library that uses row-vectors, then users should transpose the transformation matrix before calling this method.

warp.transform_vector( xform: Transformation[Float], vec: Vector[3, Float], ) → Vector[3, Float]#: Apply the transform to a vector vec treating the homogeneous coordinate as w=0 (rotation only).

warp.transform_vector( mat: Matrix[4, 4, Float], vec: Vector[3, Float], ) → Vector[3, Float]

Apply the transform to a vector vec treating the homogeneous coordinate as w=0.

The transformation is applied treating vec as a column vector, e.g.: y = mat*vec.

This is in contrast to some libraries, notably USD, which applies transforms to row vectors, y^T = vec^T*mat^T. If the transform is coming from a library that uses row-vectors, then users should transpose the transformation matrix before calling this method.

warp.transform_inverse( xform: Transformation[Float], ) → Transformation[Float]#: Compute the inverse of the transformation xform.

Spatial Math#

warp.spatial_vector(dtype: Float) → Vector[6, Float][source]#: Zero-initialize a 6D screw vector.

warp.spatial_vector( w: Vector[3, Float], v: Vector[3, Float], dtype: Float, ) → Vector[6, Float][source]: Construct a 6D screw vector from two 3D vectors.

warp.spatial_vector( wx: Float, wy: Float, wz: Float, vx: Float, vy: Float, vz: Float, dtype: Float, ) → Vector[6, Float][source]: Construct a 6D screw vector from six values.

warp.spatial_adjoint( r: Matrix[3, 3, Float], s: Matrix[3, 3, Float], ) → Matrix[6, 6, Float]#: Construct a 6x6 spatial inertial matrix from two 3x3 diagonal blocks.

warp.spatial_dot(a: Vector[6, Float], b: Vector[6, Float]) → Float#: Compute the dot product of two 6D screw vectors.

warp.spatial_cross( a: Vector[6, Float], b: Vector[6, Float], ) → Vector[6, Float]#: Compute the cross product of two 6D screw vectors.

warp.spatial_cross_dual( a: Vector[6, Float], b: Vector[6, Float], ) → Vector[6, Float]#: Compute the dual cross product of two 6D screw vectors.

warp.spatial_top(svec: Vector[6, Float]) → Vector[3, Float]#: Return the top (first) part of a 6D screw vector.

warp.spatial_bottom(svec: Vector[6, Float]) → Vector[3, Float]#: Return the bottom (second) part of a 6D screw vector.

warp.spatial_jacobian( S: Array[Vector[6, Float]], joint_parents: Array[int32], joint_qd_start: Array[int32], joint_start: int32, joint_count: int32, J_start: int32, J_out: Array[Float], ) → None#

warp.spatial_mass( I_s: Array[Matrix[6, 6, Float]], joint_start: int32, joint_count: int32, M_start: int32, M: Array[Float], ) → None#

Tile Primitives#

warp.tile_zeros(m: int32, n: int32, dtype: Any, storage: str) → Tile#

Allocates a tile of zero-initialized items.

Parameters:

m – Size of the first dimension of the output tile
n – Size of the second dimension of the output tile
dtype – Datatype of output tile’s elements
storage – The storage location for the tile: "register" for registers (default) or "shared" for shared memory.

Returns:

A zero-initialized tile with shape=(m,n) and the specified datatype [1]

warp.tile_ones(m: int32, n: int32, dtype: Any, storage: str) → Tile#

Allocates a tile of one-initialized items.

Parameters:

m – Size of the first dimension of the output tile
n – Size of the second dimension of the output tile
dtype – Datatype of output tile’s elements
storage – The storage location for the tile: "register" for registers (default) or "shared" for shared memory.

Returns:

A one-initialized tile with shape=(m,n) and the specified dtype [1]

warp.tile_arange(*args: Scalar, dtype: Any, storage: str) → Tile#

Generates a tile of linearly spaced elements.

Parameters:

args –
Variable-length positional arguments, interpreted as:
- (stop,): Generates values from 0 to stop - 1
- (start, stop): Generates values from start to stop - 1
- (start, stop, step): Generates values from start to stop - 1 with a step size
dtype – Datatype of output tile’s elements (optional, default: int)
storage – The storage location for the tile: "register" for registers (default) or "shared" for shared memory.

Returns:

A tile with shape=(1,n) with linearly spaced elements of specified dtype [1]

warp.tile_load(a: Array[Any], i: int32, n: int32, storage: str) → Tile#

Loads a 1D tile from a global memory array.

This method will cooperatively load a tile from global memory using all threads in the block.

Parameters:

a – The source array in global memory
i – Offset in the source array measured in multiples of n, i.e.: offset=i*n
n – The number of elements in the tile
storage – The storage location for the tile: "register" for registers (default) or "shared" for shared memory.

Returns:

A tile with shape=(1,n) and dtype the same as the source array

warp.tile_load( a: Array[Any], i: int32, j: int32, m: int32, n: int32, storage: str, ) → Tile

Loads a 2D tile from a global memory array.

This method will cooperatively load a tile from global memory using all threads in the block.

Parameters:

a – The source array in global memory
i – Offset in the source array measured in multiples of m, i.e.: row=i*m
j – Offset in the source array measured in multiples of n, i.e.; col=j*n
m – The size of the tile’s first dimension
n – The size of the tile’s second dimension
storage – The storage location for the tile: "register" for registers (default) or "shared" for shared memory.

Returns:

A tile with shape=(m,n) and dtype the same as the source array

warp.tile_store(a: Array[Any], i: int32, t: Any) → None#

Stores a 1D tile to a global memory array.

This method will cooperatively store a tile to global memory using all threads in the block.

Parameters:

a – The destination array in global memory
i – Offset in the destination array measured in multiples of n, i.e.: offset=i*n
t – The source tile to store data from, must have the same dtype as the destination array

warp.tile_store(a: Array[Any], i: int32, j: int32, t: Any) → None

Stores a tile to a global memory array.

This method will cooperatively store a tile to global memory using all threads in the block.

Parameters:

a – The destination array in global memory
i – Offset in the destination array measured in multiples of m, i.e.: row=i*m
j – Offset in the destination array measured in multiples of n, i.e.; col=j*n
t – The source tile to store data from, must have the same dtype as the destination array

warp.tile_atomic_add(a: Array[Any], x: int32, y: int32, t: Any) → Tile#

Atomically add a tile to the array a, each element will be updated atomically.

Parameters:

a – Array in global memory, should have the same dtype as the input tile
x – Offset in the destination array measured in multiples of m, i.e.: i=x*M where M is the first tile dimension
y – Offset in the destination array measured in multiples of n, i.e.: j=y*N where N is the second tile dimension
t – Source tile to add to the destination array

Returns:

A tile with the same dimensions and type as the source tile, holding the original value of the destination elements

warp.tile_view(t: Tile, i: int32, j: int32, m: int32, n: int32) → Tile#

Return a subrange of a given tile from coordinates (i,j) to (i+m, j+n).

Parameters:

t – Input tile to extract a subrange from
i – Offset in the source tile along the first dimension
j – Offset in the source tile along the second dimensions
m – Size of the subrange to return along the first dimension
n – Size of the subrange to return along the second dimension

Returns:

A tile with dimensions (m,n) and the same datatype as the input tile

warp.tile_assign(dst: Tile, i: int32, j: int32, src: Tile) → None#

Assign a tile to a subrange of a destination tile at coordinates (i,j).

Parameters:

t – The destination tile to assign to
i – Offset in the source tile along the first dimension
j – Offset in the source tile along the second dimensions
src – The source tile to read values from

warp.tile(x: Any) → Tile#

Constructs a new Tile from per-thread kernel values.

This function converts values computed using scalar kernel code to a tile representation for input into collective operations.

If the input value is a scalar, then the resulting tile has shape=(1, block_dim)
If the input value is a vector, then the resulting tile has shape=(length(vector), block_dim)

Parameters:: x – A per-thread local value, e.g.: scalar, vector, or matrix.
Returns:: A tile with first dimension according to the value type length and a second dimension equal to block_dim

This example shows how to create a linear sequence from thread variables:

@wp.kernel
def compute():
    i = wp.tid()
    t = wp.tile(i*2)
    print(t)

wp.launch(compute, dim=16, inputs=[], block_dim=16)

Prints:

tile(m=1, n=16, storage=register) = [[0 2 4 6 8 ...]]

warp.untile(a: Any) → Scalar#

Convert a Tile back to per-thread values.

This function converts a block-wide tile back to per-thread values.

If the input tile is 1-dimensional then the resulting value will be a per-thread scalar
If the input tile is 2-dimensional then the resulting value will be a per-thread vector of length M

Parameters:: a – A tile with dimensions shape=(M, block_dim)
Returns:: A single value per-thread with the same dtype as the tile

This example shows how to create a linear sequence from thread variables:

@wp.kernel
def compute():
    i = wp.tid()

    # create block-wide tile
    t = wp.tile(i)*2

    # convert back to per-thread values
    s = wp.untile()

    print(s)

wp.launch(compute, dim=16, inputs=[], block_dim=16)

Prints:

0
2
4
6
8
...

warp.tile_extract(a: Tile, i: int32, j: int32) → Scalar#

Extracts a single element from the tile and returns it as a scalar type.

This function will extract an element from the tile and broadcast its value to all threads in the block.

Note that this may incur additional synchronization if the source tile is a register tile.

Parameters:

a – Tile to extract the element from
i – Coordinate of element on first dimension
j – Coordinate of element on the second dimension

Returns:

The value of the element at the specified tile location, with the same type as the input tile’s per-element dtype

warp.tile_transpose(a: Tile) → Tile#

Transpose a tile.

For shared memory tiles this operation will alias the input tile, register tiles will first be transferred to shared memory before transposition.

Parameters:: a – Tile to transpose with shape=(M,N)
Returns:: Tile with shape=(N,M)

warp.tile_broadcast(a: Tile, m: int32, n: int32) → Tile#

Broadcast a tile.

This method will attempt to broadcast the input tile a to the destination shape (m, n), broadcasting follows NumPy broadcast rules.

Parameters:: a – Tile to broadcast
Returns:: Tile with broadcast shape=(m, n)

warp.tile_sum(a: Tile) → Tile#

Cooperatively compute the sum of the tile elements using all threads in the block.

Parameters:: a – The tile to compute the sum of
Returns:: A single-element tile with dimensions of (1,1) holding the sum

Example:

@wp.kernel
def compute():

    t = wp.tile_ones(dtype=float, m=16, n=16)
    s = wp.tile_sum(t)

    print(t)

wp.launch(compute, dim=[64], inputs=[])

Prints:

tile(m=1, n=1, storage=register) = [[256]]

warp.tile_min(a: Tile) → Tile#

Cooperatively compute the minimum of the tile elements using all threads in the block.

Parameters:: a – The tile to compute the minimum of
Returns:: A single-element tile with dimensions of (1,1) holding the minimum value

Example:

@wp.kernel
def compute():

    t = wp.tile_arange(start=--10, stop=10, dtype=float)
    s = wp.tile_min(t)

    print(t)

wp.launch(compute, dim=[64], inputs=[])

Prints:

tile(m=1, n=1, storage=register) = [[-10]]

warp.tile_max(a: Tile) → Tile#

Cooperatively compute the maximum of the tile elements using all threads in the block.

Parameters:: a – The tile to compute the maximum from
Returns:: A single-element tile with dimensions of (1,1) holding the maximum value

Example:

@wp.kernel
def compute():

    t = wp.tile_arange(start=--10, stop=10, dtype=float)
    s = wp.tile_min(t)

    print(t)

wp.launch(compute, dim=[64], inputs=[])

Prints:

tile(m=1, n=1, storage=register) = [[10]]

warp.tile_reduce(op: Callable, a: Any) → Tile#

Apply a custom reduction operator across the tile.

This function cooperatively performs a reduction using the provided operator across the tile.

Parameters:

op – A callable function that accepts two arguments and returns one argument, may be a user function or builtin
a – The input tile, the operator (or one of its overloads) must be able to accept the tile’s dtype

Returns:

A single-element tile with shape=(1,1) with the same datatype as the input tile.

Example:

@wp.kernel
def factorial():

    t = wp.tile_arange(1, 10, dtype=int)
    s = wp.tile_reduce(wp.mul, t)

    print(s)

wp.launch(factorial, dim=[16], inputs=[], block_dim=16)

Prints:

tile(m=1, n=1, storage=register) = [[362880]]

warp.tile_map(op: Callable, a: Any) → Tile#

Apply a unary function onto the tile.

This function cooperatively applies a unary function to each element of the tile using all threads in the block.

Parameters:

op – A callable function that accepts one argument and returns one argument, may be a user function or builtin
a – The input tile, the operator (or one of its overloads) must be able to accept the tile’s dtype

Returns:

A tile with the same dimensions and datatype as the input tile.

Example:

@wp.kernel
def compute():

    t = wp.tile_arange(0.0, 1.0, 0.1, dtype=float)
    s = wp.tile_map(wp.sin, t)

    print(s)

wp.launch(compute, dim=[16], inputs=[])

Prints:

tile(m=1, n=10, storage=register) = [[0 0.0998334 0.198669 0.29552 ...]]

warp.tile_map(op: Callable, a: Any, b: Any) → Tile

Apply a binary function onto the tile.

This function cooperatively applies a binary function to each element of the tiles using all threads in the block. Both input tiles must have the same dimensions and datatype.

Parameters:

op – A callable function that accepts two arguments and returns one argument, all of the same type, may be a user function or builtin
a – The first input tile, the operator (or one of its overloads) must be able to accept the tile’s dtype
b – The second input tile, the operator (or one of its overloads) must be able to accept the tile’s dtype

Returns:

A tile with the same dimensions and datatype as the input tiles.

Example:

@wp.kernel
def compute():

    a = wp.tile_arange(0.0, 1.0, 0.1, dtype=float)
    b = wp.tile_ones(m=1, n=10, dtype=float)

    s = wp.tile_map(wp.add, a, b)

    print(s)

wp.launch(compute, dim=[16], inputs=[])

Prints:

tile(m=1, n=10, storage=register) = [[1 1.1 1.2 1.3 1.4 1.5 1.6 1.7 1.8 1.9]]

warp.tile_matmul(a: Tile, b: Tile, out: Tile) → Tile#

Computes the matrix product and accumulates out += a*b.

Supported datatypes are:

fp16, fp32, fp64 (real)
vec2h, vec2f, vec2d (complex)

All input and output tiles must have the same datatype. Tile data will be automatically be migrated to shared memory if necessary and will use TensorCore operations when available.

Parameters:

a – A tile with shape=(M, K)
b – A tile with shape=(K, N)
out – A tile with shape=(M, N)

warp.tile_matmul(a: Tile, b: Tile) → Tile

Computes the matrix product out = a*b.

Supported datatypes are:

fp16, fp32, fp64 (real)
vec2h, vec2f, vec2d (complex)

Both input tiles must have the same datatype. Tile data will be automatically be migrated to shared memory if necessary and will use TensorCore operations when available.

Parameters:

a – A tile with shape=(M, K)
b – A tile with shape=(K, N)

Returns:

A tile with shape=(M, N)

warp.tile_fft(inout: Tile) → Tile#

Compute the forward FFT along the second dimension of a 2D tile of data.

This function cooperatively computes the forward FFT on a tile of data inplace, treating each row individually.

Supported datatypes are:

vec2f, vec2d

Parameters:: inout – The input/output tile

warp.tile_ifft(inout: Tile) → Tile#

Compute the inverse FFT along the second dimension of a 2D tile of data.

This function cooperatively computes the inverse FFT on a tile of data inplace, treating each row individually.

Supported datatypes are:

vec2f, vec2d

Parameters:: inout – The input/output tile

Utility#

warp.mlp( weights: Array[float32], bias: Array[float32], activation: Callable, index: int32, x: Array[float32], out: Array[float32], ) → None#

Evaluate a multi-layer perceptron (MLP) layer in the form: out = act(weights*x + bias).

Parameters:

weights – A layer’s network weights with dimensions (m, n).
bias – An array with dimensions (n).
activation – A wp.func function that takes a single scalar float as input and returns a scalar float as output
index – The batch item to process, typically each thread will process one item in the batch, in which case index should be wp.tid()
x – The feature matrix with dimensions (n, b)
out – The network output with dimensions (m, b)

Note:

Feature and output matrices are transposed compared to some other frameworks such as PyTorch. All matrices are assumed to be stored in flattened row-major memory layout (NumPy default).

warp.reversed(range: range_t) → range_t#: Returns the range in reversed order.

warp.printf(fmt: str, *args: Any) → None#: Allows printing formatted strings using C-style format specifiers.

warp.print(value: Any) → None#: Print variable to stdout

warp.breakpoint() → None#: Debugger breakpoint

warp.tid() → int#

Return the current thread index for a 1D kernel launch.

Note that this is the global index of the thread in the range [0, dim) where dim is the parameter passed to kernel launch.

This function may not be called from user-defined Warp functions.

warp.tid() → Tuple[int, int]

Return the current thread indices for a 2D kernel launch.

Use i,j = wp.tid() syntax to retrieve the coordinates inside the kernel thread grid.

This function may not be called from user-defined Warp functions.

warp.tid() → Tuple[int, int, int]

Return the current thread indices for a 3D kernel launch.

Use i,j,k = wp.tid() syntax to retrieve the coordinates inside the kernel thread grid.

This function may not be called from user-defined Warp functions.

warp.tid() → Tuple[int, int, int, int]

Return the current thread indices for a 4D kernel launch.

Use i,j,k,l = wp.tid() syntax to retrieve the coordinates inside the kernel thread grid.

This function may not be called from user-defined Warp functions.

warp.select(cond: bool, value_if_false: Any, value_if_true: Any) → Any#: Select between two arguments, if cond is False then return value_if_false, otherwise return value_if_true

warp.select(cond: int8, value_if_false: Any, value_if_true: Any) → Any: Select between two arguments, if cond is False then return value_if_false, otherwise return value_if_true

warp.select(cond: uint8, value_if_false: Any, value_if_true: Any) → Any: Select between two arguments, if cond is False then return value_if_false, otherwise return value_if_true

warp.select(cond: int16, value_if_false: Any, value_if_true: Any) → Any: Select between two arguments, if cond is False then return value_if_false, otherwise return value_if_true

warp.select(cond: uint16, value_if_false: Any, value_if_true: Any) → Any: Select between two arguments, if cond is False then return value_if_false, otherwise return value_if_true

warp.select(cond: int32, value_if_false: Any, value_if_true: Any) → Any: Select between two arguments, if cond is False then return value_if_false, otherwise return value_if_true

warp.select(cond: uint32, value_if_false: Any, value_if_true: Any) → Any: Select between two arguments, if cond is False then return value_if_false, otherwise return value_if_true

warp.select(cond: int64, value_if_false: Any, value_if_true: Any) → Any: Select between two arguments, if cond is False then return value_if_false, otherwise return value_if_true

warp.select(cond: uint64, value_if_false: Any, value_if_true: Any) → Any: Select between two arguments, if cond is False then return value_if_false, otherwise return value_if_true

warp.select( arr: Array[Any], value_if_false: Any, value_if_true: Any, ) → Any: Select between two arguments, if arr is null then return value_if_false, otherwise return value_if_true

warp.atomic_add(arr: Array[Any], i: Int, value: Any) → Any#: Atomically add value onto arr[i] and return the old value.

warp.atomic_add(arr: Array[Any], i: Int, j: Int, value: Any) → Any: Atomically add value onto arr[i,j] and return the old value.

warp.atomic_add(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) → Any: Atomically add value onto arr[i,j,k] and return the old value.

warp.atomic_add( arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any, ) → Any: Atomically add value onto arr[i,j,k,l] and return the old value.

warp.atomic_add(arr: FabricArray[Any], i: Int, value: Any) → Any: Atomically add value onto arr[i] and return the old value.

warp.atomic_add(arr: FabricArray[Any], i: Int, j: Int, value: Any) → Any: Atomically add value onto arr[i,j] and return the old value.

warp.atomic_add( arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any, ) → Any: Atomically add value onto arr[i,j,k] and return the old value.

warp.atomic_add( arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any, ) → Any: Atomically add value onto arr[i,j,k,l] and return the old value.

warp.atomic_add(arr: IndexedFabricArray[Any], i: Int, value: Any) → Any: Atomically add value onto arr[i] and return the old value.

warp.atomic_add( arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any, ) → Any: Atomically add value onto arr[i,j] and return the old value.

warp.atomic_add( arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any, ) → Any: Atomically add value onto arr[i,j,k] and return the old value.

warp.atomic_add( arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any, ) → Any: Atomically add value onto arr[i,j,k,l] and return the old value.

warp.atomic_sub(arr: Array[Any], i: Int, value: Any) → Any#: Atomically subtract value onto arr[i] and return the old value.

warp.atomic_sub(arr: Array[Any], i: Int, j: Int, value: Any) → Any: Atomically subtract value onto arr[i,j] and return the old value.

warp.atomic_sub(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) → Any: Atomically subtract value onto arr[i,j,k] and return the old value.

warp.atomic_sub( arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any, ) → Any: Atomically subtract value onto arr[i,j,k,l] and return the old value.

warp.atomic_sub(arr: FabricArray[Any], i: Int, value: Any) → Any: Atomically subtract value onto arr[i] and return the old value.

warp.atomic_sub(arr: FabricArray[Any], i: Int, j: Int, value: Any) → Any: Atomically subtract value onto arr[i,j] and return the old value.

warp.atomic_sub( arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any, ) → Any: Atomically subtract value onto arr[i,j,k] and return the old value.

warp.atomic_sub( arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any, ) → Any: Atomically subtract value onto arr[i,j,k,l] and return the old value.

warp.atomic_sub(arr: IndexedFabricArray[Any], i: Int, value: Any) → Any: Atomically subtract value onto arr[i] and return the old value.

warp.atomic_sub( arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any, ) → Any: Atomically subtract value onto arr[i,j] and return the old value.

warp.atomic_sub( arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any, ) → Any: Atomically subtract value onto arr[i,j,k] and return the old value.

warp.atomic_sub( arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any, ) → Any: Atomically subtract value onto arr[i,j,k,l] and return the old value.

warp.atomic_min(arr: Array[Any], i: Int, value: Any) → Any#

Compute the minimum of value and arr[i], atomically update the array, and return the old value.