triton.extension.buffer.language.alloc

triton.extension.buffer.language.alloc(etype: dtype, shape: List[constexpr], _address_space: address_space = None, is_mem_unique: bool = False, _semantic=None) → buffer

Allocates a region of local memory with the specified shape and type.

参数:

• etype (tl.dtype) -- the element type of the buffer.

• shape (List[tl.constexpr]) -- A list of non-negative integers representing the shape of the buffer.

• _address_space (bl.address_space) -- (Optional) backend-specific local memory address space

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)

Special Restrictions

• DataType: Ascend supports int8, int16, int32, uint8, uint64, int64, fp32, bf16, bool. Does not support uint16, uint32, fp16.

• Shape: Each element must be a positive integer.

• Address space: must fit within the specified address space size limits.