338 lines
13 KiB
Python
338 lines
13 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
|
#
|
|
# Use of this software is governed by the terms and conditions of the
|
|
# NVIDIA End User License Agreement (EULA), available at:
|
|
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
|
|
#
|
|
# Any use, reproduction, disclosure, or distribution of this software
|
|
# and related documentation outside the scope permitted by the EULA
|
|
# is strictly prohibited.
|
|
|
|
from typing import Optional, Type
|
|
import inspect
|
|
|
|
from cutlass import const_expr
|
|
from cutlass.cutlass_dsl import (
|
|
Numeric,
|
|
Float32,
|
|
extract_mlir_values,
|
|
new_from_mlir_values,
|
|
dsl_user_op,
|
|
)
|
|
import cutlass.pipeline as pipeline
|
|
import cutlass.cute as cute
|
|
from cutlass._mlir import ir
|
|
|
|
|
|
class TmemAllocator:
|
|
"""A class for managing tensor memory allocation on Blackwell GPU.
|
|
|
|
This class manages allocation/deallocation of tensor memory, including the mbarrier
|
|
synchronization for two cta use case.
|
|
|
|
:ivar _alloc_result_dst_smem_ptr: The smem pointer that holds the base address of allocated tensor memory.
|
|
:type _alloc_result_dst_smem_ptr: cute.Pointer
|
|
:ivar _barrier_for_retrieve: The barrier for retrieving tensor memory ptr.
|
|
:type _barrier_for_retrieve: pipeline.NamedBarrier
|
|
:ivar _allocator_warp_id: The warp id of the allocator warp.
|
|
:type _allocator_warp_id: int
|
|
:ivar _is_two_cta: Whether the allocator is for two cta.
|
|
:type _is_two_cta: bool
|
|
:ivar _num_allocated_columns: The number of columns allocated in the tensor memory.
|
|
:type _num_allocated_columns: int
|
|
:ivar _two_cta_tmem_dealloc_mbar_ptr: The mbarrier pointer required when deallocating tensor memory for two cta.
|
|
:type _two_cta_tmem_dealloc_mbar_ptr: cute.Pointer
|
|
"""
|
|
|
|
@dsl_user_op
|
|
@cute.jit
|
|
def _init_dealloc_mbarrier(self, *, loc=None, ip=None):
|
|
assert self._two_cta_tmem_dealloc_mbar_ptr is not None, (
|
|
"two_cta_tmem_dealloc_mbar_ptr is required for two cta"
|
|
)
|
|
warp_idx = cute.arch.warp_idx(loc=loc, ip=ip)
|
|
warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip)
|
|
_is_allocator_warp = warp_idx == self._allocator_warp_id
|
|
if _is_allocator_warp:
|
|
num_tmem_dealloc_threads = 32
|
|
with cute.arch.elect_one(loc=loc, ip=ip):
|
|
cute.arch.mbarrier_init(
|
|
self._two_cta_tmem_dealloc_mbar_ptr,
|
|
num_tmem_dealloc_threads,
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
cute.arch.mbarrier_init_fence()
|
|
|
|
@dsl_user_op
|
|
def __init__(
|
|
self,
|
|
alloc_result_dst_smem_ptr: cute.Pointer,
|
|
barrier_for_retrieve: pipeline.NamedBarrier,
|
|
allocator_warp_id: int = 0,
|
|
is_two_cta: bool = False,
|
|
num_allocated_columns: int = 0,
|
|
two_cta_tmem_dealloc_mbar_ptr: Optional[cute.Pointer] = None,
|
|
*,
|
|
loc=None,
|
|
ip=None,
|
|
):
|
|
"""
|
|
Initialize a TmemAllocator instance for managing tensor memory on Blackwell GPUs.
|
|
|
|
This initializer sets up the allocator's state, including the shared memory (smem) pointer
|
|
holding the base address of the allocated tensor memory, barrier synchronization for
|
|
retrieving the tensor memory pointer, allocator warp ID, whether the allocator is being used
|
|
for a 2-SM configuration, number of allocated columns in tensor
|
|
memory, and the optional mbarrier pointer for deallocation in the 2-SM case.
|
|
|
|
If `is_two_cta` is set to True, this will initialize the mbarrier pointer required for tensor
|
|
memory deallocation across two CTAs.
|
|
|
|
:param alloc_result_dst_smem_ptr: The shared memory pointer that holds the base address of allocated tensor memory.
|
|
:type alloc_result_dst_smem_ptr: cute.Pointer
|
|
:param barrier_for_retrieve: The named barrier for retrieving the tensor memory pointer.
|
|
:type barrier_for_retrieve: pipeline.NamedBarrier
|
|
:param allocator_warp_id: The warp ID of the allocator warp, defaults to 0.
|
|
:type allocator_warp_id: int, optional
|
|
:param is_two_cta: Whether the allocator should coordinate two CTAs, defaults to False.
|
|
:type is_two_cta: bool, optional
|
|
:param num_allocated_columns: The number of columns allocated in tensor memory, defaults to 0.
|
|
:type num_allocated_columns: int, optional
|
|
:param two_cta_tmem_dealloc_mbar_ptr: The mbarrier pointer required for two-CTA tensor memory deallocation, optional.
|
|
:type two_cta_tmem_dealloc_mbar_ptr: cute.Pointer, optional
|
|
:param loc: Optional codegen location for debugging and error reporting.
|
|
:type loc: Any, optional
|
|
:param ip: Optional insertion point for codegen.
|
|
:type ip: Any, optional
|
|
|
|
:raises AssertionError: If two_cta_tmem_dealloc_mbar_ptr is None while is_two_cta is True.
|
|
"""
|
|
# TODO: automatically maintain a smem address
|
|
self._alloc_result_dst_smem_ptr = alloc_result_dst_smem_ptr
|
|
self._allocator_warp_id = allocator_warp_id
|
|
self._is_two_cta = is_two_cta
|
|
self._num_allocated_columns = num_allocated_columns
|
|
self._two_cta_tmem_dealloc_mbar_ptr = two_cta_tmem_dealloc_mbar_ptr
|
|
self._barrier_for_retrieve = barrier_for_retrieve
|
|
|
|
# Init tmem dealloc mbarrier if two cta
|
|
if const_expr(self._is_two_cta):
|
|
self._init_dealloc_mbarrier(loc=loc, ip=ip)
|
|
|
|
def __extract_mlir_values__(self) -> list[ir.Value]:
|
|
values = extract_mlir_values(self._alloc_result_dst_smem_ptr)
|
|
if self._is_two_cta:
|
|
assert self._two_cta_tmem_dealloc_mbar_ptr is not None, (
|
|
"2CTA mode requires the dealloc mbarrier"
|
|
)
|
|
values.extend(extract_mlir_values(self._two_cta_tmem_dealloc_mbar_ptr))
|
|
return values
|
|
|
|
def __new_from_mlir_values__(self, values: list[ir.Value]) -> "TmemAllocator":
|
|
assert len(values) == 2 if self._is_two_cta else 1
|
|
new_alloc_result_dst_smem_ptr = new_from_mlir_values(
|
|
self._alloc_result_dst_smem_ptr, [values[0]]
|
|
)
|
|
new_two_cta_tmem_dealloc_mbar_ptr = (
|
|
new_from_mlir_values(self._two_cta_tmem_dealloc_mbar_ptr, [values[1]])
|
|
if self._is_two_cta
|
|
else None
|
|
)
|
|
return TmemAllocator(
|
|
new_alloc_result_dst_smem_ptr,
|
|
pipeline.NamedBarrier(
|
|
barrier_id=self._barrier_for_retrieve.barrier_id,
|
|
num_threads=self._barrier_for_retrieve.num_threads,
|
|
),
|
|
self._allocator_warp_id,
|
|
self._is_two_cta,
|
|
self._num_allocated_columns,
|
|
new_two_cta_tmem_dealloc_mbar_ptr,
|
|
)
|
|
|
|
@cute.jit
|
|
def check_valid_num_columns(self, num_columns: int):
|
|
"""Check if the number of columns is valid.
|
|
|
|
This method checks if the number of columns is valid.
|
|
It checks if the number of columns is larger than 0, smaller than 512, a multiple of 32, and a power of two.
|
|
"""
|
|
# larger than 0
|
|
if const_expr(num_columns < 0):
|
|
return False
|
|
# smaller than 512
|
|
if const_expr(num_columns > 512):
|
|
return False
|
|
# multiple of 32
|
|
if const_expr(num_columns % 32 != 0):
|
|
return False
|
|
# power of two
|
|
if const_expr(num_columns & (num_columns - 1) != 0):
|
|
return False
|
|
return True
|
|
|
|
@dsl_user_op
|
|
@cute.jit
|
|
def allocate(self, num_columns: int, *, loc=None, ip=None):
|
|
"""Allocate a block of tensor memory.
|
|
|
|
This method allocates a block of tensor memory from allocator warp and returns a handle to retrieve
|
|
the allocated tensor memory address.
|
|
"""
|
|
|
|
assert self.check_valid_num_columns(num_columns), (
|
|
"num_columns must be multiple of 32 and power of two, and between 0 and 512"
|
|
)
|
|
assert self._num_allocated_columns + num_columns <= 512, (
|
|
"total allocated columns must be less than or equal to 512"
|
|
)
|
|
|
|
warp_idx = cute.arch.warp_idx(loc=loc, ip=ip)
|
|
warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip)
|
|
_is_allocator_warp = warp_idx == self._allocator_warp_id
|
|
if _is_allocator_warp:
|
|
cute.arch.alloc_tmem(
|
|
num_columns,
|
|
self._alloc_result_dst_smem_ptr,
|
|
is_two_cta=self._is_two_cta,
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
self._num_allocated_columns += num_columns
|
|
|
|
@dsl_user_op
|
|
def wait_for_alloc(self, *, loc=None, ip=None):
|
|
"""Wait for the allocator warp to finish allocation.
|
|
|
|
This method is used to synchronize the allocator warp with the other warps before retrieving tmem ptr.
|
|
"""
|
|
self._barrier_for_retrieve.arrive_and_wait(loc=loc, ip=ip)
|
|
|
|
@dsl_user_op
|
|
def retrieve_ptr(
|
|
self,
|
|
dtype: Type[Numeric] = Float32,
|
|
*,
|
|
loc=None,
|
|
ip=None,
|
|
) -> cute.Pointer:
|
|
"""Retrieve the pointer to the allocated tensor memory.
|
|
|
|
This method can be called by all warps after allocation has been performed
|
|
by the allocator warp.
|
|
"""
|
|
return cute.arch.retrieve_tmem_ptr(
|
|
dtype,
|
|
alignment=16,
|
|
ptr_to_buffer_holding_addr=self._alloc_result_dst_smem_ptr,
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
|
|
@dsl_user_op
|
|
@cute.jit
|
|
def relinquish_alloc_permit(self, *, loc=None, ip=None):
|
|
"""Relinquish the tensor memory allocation permit.
|
|
|
|
This method relinquishes the tensor memory allocation permit for the allocator warp, promising
|
|
the allocator warp will not allocate any more tensor memory.
|
|
"""
|
|
warp_idx = cute.arch.warp_idx(loc=loc, ip=ip)
|
|
warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip)
|
|
_is_allocator_warp = warp_idx == self._allocator_warp_id
|
|
if _is_allocator_warp:
|
|
cute.arch.relinquish_tmem_alloc_permit(
|
|
is_two_cta=self._is_two_cta, loc=loc, ip=ip
|
|
)
|
|
|
|
@dsl_user_op
|
|
@cute.jit
|
|
def free(self, tmem_ptr: cute.Pointer, num_columns: int = 0, *, loc=None, ip=None):
|
|
"""Deallocate the tensor memory.
|
|
|
|
This method sync on mbarrier (for two cta use case) and deallocates the tensor memory from the allocator warp.
|
|
User can optionally specify the number of columns to deallocate. If not specified, all allocated columns will be deallocated.
|
|
"""
|
|
warp_idx = cute.arch.warp_idx(loc=loc, ip=ip)
|
|
warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip)
|
|
_is_allocator_warp = warp_idx == self._allocator_warp_id
|
|
|
|
assert num_columns <= self._num_allocated_columns, (
|
|
"num_columns must be less than or equal to num_allocated_columns"
|
|
)
|
|
if const_expr(num_columns != 0):
|
|
assert self.check_valid_num_columns(num_columns), "num_columns is invalid"
|
|
|
|
num_deallocate_columns = (
|
|
self._num_allocated_columns if num_columns == 0 else num_columns
|
|
) # if num_columns is 0, deallocate all allocated columns
|
|
self._num_allocated_columns -= num_deallocate_columns
|
|
if _is_allocator_warp:
|
|
if const_expr(self._is_two_cta):
|
|
bid_in_cluster = cute.arch.block_idx_in_cluster(loc=loc, ip=ip)
|
|
_cta_rank_in_cluster = cute.arch.make_warp_uniform(
|
|
bid_in_cluster, loc=loc, ip=ip
|
|
)
|
|
# Arrive and wait for dealloc signal from peer cta
|
|
cute.arch.mbarrier_arrive(
|
|
self._two_cta_tmem_dealloc_mbar_ptr,
|
|
_cta_rank_in_cluster ^ 1,
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
cute.arch.mbarrier_wait(
|
|
self._two_cta_tmem_dealloc_mbar_ptr, 0, loc=loc, ip=ip
|
|
)
|
|
# Deallocate tmem
|
|
cute.arch.dealloc_tmem(
|
|
tmem_ptr,
|
|
num_deallocate_columns,
|
|
is_two_cta=self._is_two_cta,
|
|
loc=loc,
|
|
ip=ip,
|
|
)
|
|
|
|
|
|
# Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator
|
|
TmemAllocator.__init__.__signature__ = inspect.Signature(
|
|
[
|
|
inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
|
|
inspect.Parameter(
|
|
"alloc_result_dst_smem_ptr",
|
|
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
|
annotation=cute.Pointer,
|
|
),
|
|
inspect.Parameter(
|
|
"barrier_for_retrieve",
|
|
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
|
annotation=pipeline.NamedBarrier,
|
|
),
|
|
inspect.Parameter(
|
|
"allocator_warp_id",
|
|
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
|
default=0,
|
|
annotation=int,
|
|
),
|
|
inspect.Parameter(
|
|
"is_two_cta",
|
|
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
|
default=False,
|
|
annotation=bool,
|
|
),
|
|
inspect.Parameter(
|
|
"num_allocated_columns",
|
|
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
|
default=0,
|
|
annotation=int,
|
|
),
|
|
inspect.Parameter(
|
|
"two_cta_tmem_dealloc_mbar_ptr",
|
|
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
|
default=None,
|
|
annotation=Optional[cute.Pointer],
|
|
),
|
|
]
|
|
)
|