triton.language.make_tensor_descriptor

triton.language.make_tensor_descriptor(base: tensor, shape: List[tensor], strides: List[tensor], block_shape: List[constexpr], _semantic=None) → tensor_descriptor

创建一个张量描述符对象

参数:

• base – 张量的基指针，必须是16字节对齐的

• shape – 一个表示张量形状的非负整数列表

• strides – 张量步幅列表。前导维度必须是16字节步幅的倍数，并且最后一个维度必须是连续的。

• block_shape – 从全局内存加载/存储的块的形状

备注

在支持TMA的NVIDIA GPU上，这将生成一个TMA描述符对象，从该描述符进行的加载和存储将由TMA硬件支持。

目前仅支持2-5维张量。

示例

@triton.jit
def inplace_abs(in_out_ptr, M, N, M_BLOCK: tl.constexpr, N_BLOCK: tl.constexpr):
    desc = tl.make_tensor_descriptor(
        in_out_ptr,
        shape=[M, N],
        strides=[N, 1],
        block_shape=[M_BLOCK, N_BLOCK],
    )

    moffset = tl.program_id(0) * M_BLOCK
    noffset = tl.program_id(1) * N_BLOCK

    value = desc.load([moffset, noffset])
    desc.store([moffset, noffset], tl.abs(value))

# TMA descriptors require a global memory allocation
def alloc_fn(size: int, alignment: int, stream: Optional[int]):
    return torch.empty(size, device="cuda", dtype=torch.int8)

triton.set_allocator(alloc_fn)

M, N = 256, 256
x = torch.randn(M, N, device="cuda")
M_BLOCK, N_BLOCK = 32, 32
grid = (M / M_BLOCK, N / N_BLOCK)
inplace_abs[grid](x, M, N, M_BLOCK, N_BLOCK)