tvm_tirx.backend.cuda.lang.SMEMPool

class tvm_tirx.backend.cuda.lang.SMEMPool(*args, **kwargs)

Bases: object

Bump allocator over a contiguous shared memory region.

Parameters:

ptr (Var or None, optional) – If omitted, an alloc_tensor(ty_args=[Tensor([0], "uint8", scope="shared.dyn")]) is created automatically and commit() must be called after all allocations to emit the size annotation. If a Var is provided, the caller manages the backing buffer and commit() is a no-op.

alloc(shape, dtype='float32', scope='shared.dyn', align=0, layout='default')

alloc_tcgen05_mma_AB(shape, dtype='float16', swizzle_mode='auto', align=1024)

Allocate MMA-compatible shared memory with an inferred swizzle layout.

move_base_to(offset)

commit(size=None)

Record the pool resource requirement on the enclosing kernel entry.

Must be called after all alloc() / move_base_to() calls.

Parameters:

size (int, optional) – Explicit shared memory size in bytes. When None (the default), the high-water mark max_offset tracked by the allocator is used.