triton.language.extra.cann.extension.parallel

class triton.language.extra.cann.extension.parallel(arg1, arg2=None, step=None, num_stages=None, loop_unroll_factor=None, bind_sub_block: bool = False)

Iterator that counts upward forever, with parallel execution semantics.

This is a special iterator used to implement similar semantics to Python's range in the context of triton.jit functions. In addition, it allows the user to pass extra attributes to the compiler.

参数:

• arg1 (int) -- If arg2 is not given, the end of the range; otherwise the start of the range.

• arg2 (int) -- The end of the range, if given.

• step (int) -- The step size of the iteration. Defaults to 1.

• num_stages (int) -- Number of pipeline stages for the loop.

• loop_unroll_factor (int) -- The unroll factor applied to the loop body.

• bind_sub_block (bool) -- Tells the compiler whether multiple vector cores participate in the loop. This is used in the mixed cube-vector kernel on 910B. The number of vector cores is determined by the number of iterations in this loop. Currently on 910B, at most 2 vector cores can be used.

Example

import torch
import torch_npu
from torch.testing import assert_close

import triton
import triton.language as tl
import triton.language.extra.cann.extension as extension


@triton.jit
def parallel_kernel(x_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr):
    # Load the full [M, N] block
    offs_m = tl.arange(0, M)
    offs_n = tl.arange(0, N)
    block = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :])

    # Split M across 2 vector sub-blocks; iterations are independent (no cross-iteration dependency)
    SUB_M: tl.constexpr = M // 2
    for s in extension.parallel(0, 2, bind_sub_block=True):
        sub = extension.extract_slice(block, (s * SUB_M, 0), (SUB_M, N), (1, 1))
        sub = sub * 2.0  # a compute instruction on the sub-block
        offs_sub_m = s * SUB_M + tl.arange(0, SUB_M)
        out_ptrs = out_ptr + offs_sub_m[:, None] * N + offs_n[None, :]
        tl.store(out_ptrs, sub)


def test_parallel():
    M = 128
    N = 128
    x = torch.randn((M, N), device="npu", dtype=torch.float32)
    out = torch.empty((M, N), device="npu", dtype=torch.float32)
    parallel_kernel[(1, )](x, out, M=M, N=N)
    torch.npu.synchronize()
    assert_close(out, x * 2, rtol=1e-3, atol=1e-3)


if __name__ == "__main__":
    test_parallel()
    print("test_parallel PASSED!")

Special Restrictions

• DataType: Ascend A2/A3 does not support uint16/uint32/uint64/fp64, Ascend 950 does not support fp64 (hardware limitation).

• When using parallel as an iterator, the operations within the loop body must have no dependencies and no conflicts between them:

• Memory access (load/store) operations can be executed in parallel.

• Compute operations are allowed, but there must be no more than one compute operation. Multiple compute operations would produce intermediate UB buffers, which cannot be accessed in parallel.

__init__(arg1, arg2=None, step=None, num_stages=None, loop_unroll_factor=None, bind_sub_block: bool = False)

Methods

__init__(arg1[, arg2, step, num_stages, ...])

Attributes

type