triton.language.extra.cann.extension.ascend_address_space

triton.language.extra.cann.extension.ascend_address_space = <triton.language.extra.cann.extension.core.ascend_address_space_group object>

Ascend hardware address space constants for buffer allocation.

Provides named address space specifiers that map to Ascend NPU memory regions:

• UB — Unified Buffer (on-chip shared memory)

• L1 — L1 cache buffer

• L0A — L0 buffer A (Cube unit input)

• L0B — L0 buffer B (Cube unit input)

• L0C — L0 buffer C (Cube unit output)

Usage with bl.alloc():

buf = bl.alloc(dtype, shape, al.ascend_address_space.UB)

Example

@triton.jit
def allocate_local_buffer(XBLOCK: tl.constexpr):
    bl.alloc(tl.float32, [XBLOCK])
    bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.UB)
    bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L1)
    bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L0A)
    bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L0B)
    bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L0C)
    bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.UB, is_mem_unique=True)

Special Restrictions

• Provides UB, L1, L0A, L0B, L0C address space constants for use with bl.alloc.