triton.extension.buffer.language.to_buffer

triton.extension.buffer.language.to_buffer(tensor: tensor, space: address_space = None, bind_buffer: buffer = None, _semantic=None) → buffer

Convert a tensor to a buffer.

参数:

• tensor (tl.tensor) -- the tensor to convert.

• space (address_space) -- the address space for the buffer (optional).

Example

@triton.jit
def to_buffer_kernel():
    a = tl.full((32, 2, 4), 0, dtype=tl.int64)
    a_buf = bl.to_buffer(a)
    b = tl.full((32, 2, 4), 0, dtype=tl.int64)
    b_buf = bl.to_buffer(b, al.ascend_address_space.UB)
    c = tl.full((32, 2, 4), 0, dtype=tl.int64)
    c_buf = bl.to_buffer(c, al.ascend_address_space.L1)

Special Restrictions

• Address space must be one of UB, L1, L0A, L0B, L0C.

• When using bind_buffer, tensor and bind_buffer must have identical shapes and element types.

• A tensor cannot be bound to multiple buffers.