triton.language.extra.cann.extension.sub_vec_id

triton.language.extra.cann.extension.sub_vec_id(_semantic=None) tensor

Get the Vector Core index on the AI Core.

Example

@triton.jit
def verify_sub_vec_id_kernel(out_ptr, N: tl.constexpr):
    with al.scope(core_mode="vector"):
        sub_id = al.sub_vec_id()
        offs = sub_id * N + tl.arange(0, N)
        out_ptrs = out_ptr + offs
        tl.store(out_ptrs, sub_id.to(tl.int32))

Special Restrictions

  • Only valid in mixed AIC+AIV scenarios (Cube + Vector cores).

  • Using it in pure-Cube or pure-Vector kernels will cause a compilation error.