Kernel API

LocalTensor	A local tensor is a multi-dimensional array of values in local memory (Unified Buffer, L1 Cache, etc.)
GlobalTensor	A global tensor is a contiguous ND-array of values in Global Memory.

Programming model

global_tensor	Define a new tensor descriptor for accessing data in global memory.
block_idx	Returns the current block (NPU core) index.
block_num	Returns the total number of blocks (NPU cores) allocated for the kernel.

Iterators

range	A range loop construct for use in JIT functions.
static_range	A static range loop construct for use in JIT functions.

Creation operations

cast	Cast a tensor or scalar value to a different data type.
concat	Concatenate tensors along the first dimension.
full	Create a tensor filled with a scalar value.
full_like	Create a tensor filled with a scalar value, with the same shape and dtype as the input tensor.
zeros	Create a tensor filled with zeros.
zeros_acc	Create a zero-initialized accumulator tensor in L0C memory for matrix multiplication.
zeros_like	Create a tensor filled with zeros, with the same shape and dtype as the input tensor.

Memory operations

copy	Copy a local tensor to a new local tensor, optionally reshaping and relocating.
copy_in	Copy data from a global tensor into a local tensor or scalar value.
copy_out	Copy data from a local tensor or scalar value to a global tensor.

Arithmetic operations

add	Computes the element-wise addition of input and other.
div	Computes the element-wise division of input and other.
equal	Computes the element-wise 'equality' comparison of input and other.
greater	Computes the element-wise 'greater' comparison of input and other.
greater_equal	Computes the element-wise 'greater or equal' comparison of input and other.
left_shift	Computes the element-wise left shift (bitwise) of input and other.
less	Computes the element-wise 'less' comparison of input and other.
less_equal	Computes the element-wise 'less or equal' comparison of input and other.
maximum	Computes the element-wise maximum of input and other.
minimum	Computes the element-wise minimum of input and other.
mul	Computes the element-wise multiplication of input and other.
negative	Computes the element-wise negation of input.
not_equal	Computes the element-wise 'inequality' comparison of input and other.
right_shift	Computes the element-wise right shift (bitwise) of input and other.
sub	Computes the element-wise subtraction of input and other.

Math operations

abs	Computes the element-wise absolute value of input.
ceil	Computes the element-wise ceil rounding of input.
cos	Computes the element-wise cosine of input.
cosh	Computes the element-wise hyperbolic cosine of input.
erf	Computes the element-wise error function of input.
exp	Computes the element-wise exponential of input.
exp2	Computes the element-wise exponential (base 2) of input.
floor	Computes the element-wise floor rounding of input.
log	Computes the element-wise natural logarithm of input.
log2	Computes the element-wise logarithm (base 2) of input.
relu	Computes the element-wise ReLU value of input.
rms_norm	Computes Root Mean Square Layer Normalization of input.
rsqrt	Computes the element-wise inverse square root of input.
sin	Computes the element-wise sine of input.
sinh	Computes the element-wise hyperbolic sine of input.
softmax	Computes the row-wise softmax of input.
sqrt	Computes the element-wise square root of input.
tan	Computes the element-wise tangent of input.
tanh	Computes the element-wise hyperbolic tangent of input.

Matrix multiplication operations

matmul	Computes the matrix multiplication of input and other with optional bias.
matmul_acc	Computes the matrix multiplication of input and other and accumulates the result into acc.

Indexing operations

mask	[Experimental] A context manager for masked operations on tensors.
where	Select elements from two sources based on a mask.

Reduction operations

reduce_min	Returns the minimum value of each row of the input tensor in the given dimensions dims.
reduce_max	Returns the maximum value of each row of the input tensor in the given dimensions dims.
reduce_sum	Returns the sum of each row of the input tensor in the given dimensions dims.
reduce_prod	Returns the product of each row of the input tensor in the given dimensions dims.

Shape manipulation operations

broadcast_to	Creates new tensor of a given shape broadcasting data from the input tensor.
expand_dims	Insert new dimensions of size 1 at the specified positions.
ravel	Flatten a tensor into a 1D tensor.
reshape	Reshape a tensor to a new shape without changing its data.
squeeze	Remove dimensions of size 1 from the tensor.
transpose	Rearrange tensor dimensions in specific order.

Atomic operations

atomic_add	Atomically add local tensor elements to a global tensor at specified offsets.
atomic_max	Atomically compute the maximum between local tensor elements and global tensor elements at specified offsets.
atomic_min	Atomically compute the minimum between local tensor elements and global tensor elements at specified offsets.

Debug operations

inline	Inject raw Ascend C++ code into the generated kernel.
inline_vf	Embed Ascend C VF (vector function) code within a kernel.

Utility functions

ceildiv

Compute ceiling division of two integers.