341 lines
11 KiB
Python
341 lines
11 KiB
Python
# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
|
|
|
|
import math
|
|
from typing import Optional, Type, Callable
|
|
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
from cutlass import Float32, Int32, const_expr
|
|
from cutlass.cute.nvgpu import cpasync
|
|
import cutlass.utils.blackwell_helpers as sm100_utils
|
|
from cutlass.cutlass_dsl import T, dsl_user_op
|
|
from cutlass._mlir.dialects import llvm
|
|
import cutlass.pipeline
|
|
|
|
|
|
@dsl_user_op
|
|
def cvt_copy(
|
|
atom: cute.CopyAtom,
|
|
src: cute.Tensor,
|
|
dst: cute.Tensor,
|
|
*,
|
|
pred: Optional[cute.Tensor] = None,
|
|
loc=None,
|
|
ip=None,
|
|
**kwargs,
|
|
) -> None:
|
|
assert isinstance(src.iterator, cute.Pointer) and src.memspace == cute.AddressSpace.rmem
|
|
if const_expr(src.element_type != dst.element_type):
|
|
src_cvt = cute.make_fragment_like(src, dst.element_type, loc=loc, ip=ip)
|
|
src_cvt.store(src.load().to(dst.element_type))
|
|
src = src_cvt
|
|
cute.copy(atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
|
|
|
|
|
|
@dsl_user_op
|
|
def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
|
|
dst = cute.make_fragment_like(src, src.element_type, loc=loc, ip=ip)
|
|
cute.autovec_copy(src, dst, loc=loc, ip=ip)
|
|
return dst
|
|
|
|
|
|
@dsl_user_op
|
|
def get_copy_atom(
|
|
dtype: Type[cutlass.Numeric], num_copy_elems: int, is_async: bool = False, *, loc=None, ip=None
|
|
) -> cute.CopyAtom:
|
|
num_copy_bits = const_expr(min(128, num_copy_elems * dtype.width))
|
|
copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
|
|
return cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
|
|
|
|
|
|
@dsl_user_op
|
|
def make_tmem_copy(
|
|
tmem_copy_atom: cute.CopyAtom, num_wg: int = 1, *, loc=None, ip=None
|
|
) -> cute.CopyAtom:
|
|
num_dp, num_bits, num_rep, _ = sm100_utils.get_tmem_copy_properties(tmem_copy_atom)
|
|
assert num_dp == 32
|
|
assert num_bits == 32
|
|
tiler_mn = (cute.make_layout((128 * num_rep * num_wg // 32, 32), stride=(32, 1)),)
|
|
layout_tv = cute.make_layout(
|
|
((32, 4, num_wg), (num_rep, 32)), stride=((0, 1, 4 * num_rep), (4, 4 * num_rep * num_wg))
|
|
)
|
|
return cute.make_tiled_copy(tmem_copy_atom, layout_tv, tiler_mn)
|
|
|
|
|
|
@dsl_user_op
|
|
def copy(
|
|
src: cute.Tensor,
|
|
dst: cute.Tensor,
|
|
*,
|
|
pred: Optional[cute.Tensor] = None,
|
|
num_copy_elems: int = 1,
|
|
is_async: bool = False,
|
|
loc=None,
|
|
ip=None,
|
|
**kwargs,
|
|
) -> None:
|
|
copy_atom = get_copy_atom(src.element_type, num_copy_elems, is_async)
|
|
cute.copy(copy_atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
|
|
|
|
|
|
def tiled_copy_1d(
|
|
dtype: Type[cutlass.Numeric], num_threads: int, num_copy_elems: int = 1, is_async: bool = False
|
|
) -> cute.TiledCopy:
|
|
num_copy_bits = num_copy_elems * dtype.width
|
|
copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
|
|
copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
|
|
thr_layout = cute.make_layout(num_threads)
|
|
val_layout = cute.make_layout(num_copy_elems)
|
|
return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
|
|
|
|
|
|
def tiled_copy_2d(
|
|
dtype: Type[cutlass.Numeric], major_mode_size: int, num_threads: int, is_async: bool = False
|
|
) -> cute.TiledCopy:
|
|
num_copy_bits = math.gcd(major_mode_size, 128 // dtype.width) * dtype.width
|
|
copy_elems = num_copy_bits // dtype.width
|
|
copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
|
|
copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
|
|
gmem_threads_per_row = major_mode_size // copy_elems
|
|
assert num_threads % gmem_threads_per_row == 0
|
|
thr_layout = cute.make_ordered_layout(
|
|
(num_threads // gmem_threads_per_row, gmem_threads_per_row),
|
|
order=(1, 0),
|
|
)
|
|
val_layout = cute.make_layout((1, copy_elems))
|
|
return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
|
|
|
|
|
|
@dsl_user_op
|
|
def atomic_add_fp32x4(
|
|
a: Float32, b: Float32, c: Float32, d: Float32, gmem_ptr: cute.Pointer, *, loc=None, ip=None
|
|
) -> None:
|
|
gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value()
|
|
# cache_hint = cutlass.Int64(0x12F0000000000000)
|
|
llvm.inline_asm(
|
|
None,
|
|
[
|
|
gmem_ptr_i64,
|
|
Float32(a).ir_value(loc=loc, ip=ip),
|
|
Float32(b).ir_value(loc=loc, ip=ip),
|
|
Float32(c).ir_value(loc=loc, ip=ip),
|
|
Float32(d).ir_value(loc=loc, ip=ip),
|
|
],
|
|
# [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip), cache_hint.ir_value()],
|
|
"{\n\t"
|
|
# ".reg .b128 abcd;\n\t"
|
|
# "mov.b128 abcd, {$1, $2, $3, $4};\n\t"
|
|
".reg .v4 .f32 abcd;\n\t"
|
|
# "mov.b128 abcd, {$1, $2, $3, $4};\n\t"
|
|
"mov.f32 abcd.x, $1;\n\t"
|
|
"mov.f32 abcd.y, $2;\n\t"
|
|
"mov.f32 abcd.z, $3;\n\t"
|
|
"mov.f32 abcd.w, $4;\n\t"
|
|
"red.global.add.v4.f32 [$0], abcd;\n\t"
|
|
# "red.global.add.L2::cache_hint.v4.f32 [$0], abcd, 0x14F0000000000000;\n\t"
|
|
"}\n",
|
|
# "red.global.add.L2::cache_hint.f32 [$0], $1, 0x12F0000000000000;",
|
|
# "red.global.add.L2::cache_hint.f32 [$0], $1, $2;",
|
|
"l,f,f,f,f",
|
|
# "l,f,l",
|
|
has_side_effects=True,
|
|
is_align_stack=False,
|
|
asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def set_block_rank(
|
|
smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None
|
|
) -> Int32:
|
|
"""Map the given smem pointer to the address at another CTA rank in the cluster."""
|
|
smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
|
|
return Int32(
|
|
llvm.inline_asm(
|
|
T.i32(),
|
|
[smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()],
|
|
"mapa.shared::cluster.u32 $0, $1, $2;",
|
|
"=r,r,r",
|
|
has_side_effects=False,
|
|
is_align_stack=False,
|
|
asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
)
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def store_shared_remote_fp32x4(
|
|
a: Float32,
|
|
b: Float32,
|
|
c: Float32,
|
|
d: Float32,
|
|
smem_ptr: cute.Pointer,
|
|
mbar_ptr: cute.Pointer,
|
|
peer_cta_rank_in_cluster: Int32,
|
|
*,
|
|
loc=None,
|
|
ip=None,
|
|
) -> None:
|
|
remote_smem_ptr_i32 = set_block_rank(
|
|
smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
|
|
).ir_value()
|
|
remote_mbar_ptr_i32 = set_block_rank(
|
|
mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
|
|
).ir_value()
|
|
llvm.inline_asm(
|
|
None,
|
|
[
|
|
remote_smem_ptr_i32,
|
|
remote_mbar_ptr_i32,
|
|
Float32(a).ir_value(loc=loc, ip=ip),
|
|
Float32(b).ir_value(loc=loc, ip=ip),
|
|
Float32(c).ir_value(loc=loc, ip=ip),
|
|
Float32(d).ir_value(loc=loc, ip=ip),
|
|
],
|
|
"{\n\t"
|
|
".reg .v4 .f32 abcd;\n\t"
|
|
"mov.f32 abcd.x, $2;\n\t"
|
|
"mov.f32 abcd.y, $3;\n\t"
|
|
"mov.f32 abcd.z, $4;\n\t"
|
|
"mov.f32 abcd.w, $5;\n\t"
|
|
"st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.f32 [$0], abcd, [$1];\n\t"
|
|
"}\n",
|
|
"r,r,f,f,f,f",
|
|
has_side_effects=True,
|
|
is_align_stack=False,
|
|
asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def cpasync_bulk_g2s(
|
|
gmem_ptr: cute.Pointer,
|
|
smem_ptr: cute.Pointer,
|
|
tma_bar_ptr: cute.Pointer,
|
|
size: int | Int32,
|
|
*,
|
|
loc=None,
|
|
ip=None,
|
|
):
|
|
gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value()
|
|
smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
|
|
mbar_ptr_i32 = tma_bar_ptr.toint(loc=loc, ip=ip).ir_value()
|
|
llvm.inline_asm(
|
|
None,
|
|
[gmem_ptr_i64, smem_ptr_i32, mbar_ptr_i32, Int32(size).ir_value()],
|
|
"cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [$1], [$0], $3, [$2];",
|
|
"l,r,r,r",
|
|
has_side_effects=True,
|
|
is_align_stack=False,
|
|
asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
)
|
|
|
|
|
|
@dsl_user_op
|
|
def cpasync_reduce_bulk_add_f32(
|
|
smem_ptr: cute.Pointer,
|
|
gmem_ptr: cute.Pointer,
|
|
store_bytes: int | Int32,
|
|
*,
|
|
loc=None,
|
|
ip=None,
|
|
):
|
|
smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
|
|
# cache_hint = cutlass.Int64(0x14F0000000000000) # EVICT_LAST
|
|
llvm.inline_asm(
|
|
None,
|
|
[gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value()],
|
|
"cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;",
|
|
"l,r,r",
|
|
# [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value(), cache_hint.ir_value()],
|
|
# "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.f32 [$0], [$1], $2, $3;",
|
|
# "l,r,r,l",
|
|
has_side_effects=True,
|
|
is_align_stack=False,
|
|
asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
)
|
|
|
|
|
|
def cpasync_bulk_get_copy_fn(
|
|
src_tensor: cute.Tensor,
|
|
dst_tensor: cute.Tensor,
|
|
single_stage: bool = False,
|
|
**kwargs,
|
|
) -> Callable:
|
|
# src_is_smem = const_expr(
|
|
# isinstance(src_tensor.iterator, cute.Pointer)
|
|
# and src_tensor.memspace == cute.AddressSpace.smem
|
|
# )
|
|
group_rank_src = const_expr(cute.rank(src_tensor) - (1 if not single_stage else 0))
|
|
group_rank_dst = const_expr(cute.rank(dst_tensor) - (1 if not single_stage else 0))
|
|
# ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK)
|
|
src = cute.group_modes(src_tensor, 0, group_rank_src)
|
|
dst = cute.group_modes(dst_tensor, 0, group_rank_dst)
|
|
|
|
def copy_bulk(src_idx, dst_idx, **new_kwargs):
|
|
size = const_expr(cute.size(src.shape[:-1]) * src.element_type.width // 8)
|
|
cpasync_bulk_g2s(
|
|
src[None, src_idx].iterator,
|
|
dst[None, dst_idx].iterator,
|
|
size=size,
|
|
**new_kwargs,
|
|
**kwargs,
|
|
)
|
|
|
|
def copy_bulk_single_stage(**new_kwargs):
|
|
size = const_expr(cute.size(src.shape) * src.element_type.width // 8)
|
|
cpasync_bulk_g2s(src.iterator, dst.iterator, size=size, **new_kwargs, **kwargs)
|
|
|
|
return copy_bulk if const_expr(not single_stage) else copy_bulk_single_stage
|
|
|
|
|
|
def tma_get_copy_fn(
|
|
atom: cute.CopyAtom,
|
|
cta_coord: cute.Coord,
|
|
cta_layout: cute.Layout,
|
|
src_tensor: cute.Tensor,
|
|
dst_tensor: cute.Tensor,
|
|
filter_zeros: bool = False,
|
|
single_stage: bool = False,
|
|
**kwargs,
|
|
) -> Callable:
|
|
src_is_smem = const_expr(
|
|
isinstance(src_tensor.iterator, cute.Pointer)
|
|
and src_tensor.memspace == cute.AddressSpace.smem
|
|
)
|
|
smem_tensor, gmem_tensor = (src_tensor, dst_tensor) if src_is_smem else (dst_tensor, src_tensor)
|
|
group_rank_smem = const_expr(cute.rank(smem_tensor) - (1 if not single_stage else 0))
|
|
group_rank_gmem = const_expr(cute.rank(gmem_tensor) - (1 if not single_stage else 0))
|
|
# ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK)
|
|
s, g = cpasync.tma_partition(
|
|
atom,
|
|
cta_coord,
|
|
cta_layout,
|
|
cute.group_modes(smem_tensor, 0, group_rank_smem),
|
|
cute.group_modes(gmem_tensor, 0, group_rank_gmem),
|
|
)
|
|
if const_expr(filter_zeros):
|
|
s = cute.filter_zeros(s)
|
|
g = cute.filter_zeros(g)
|
|
src, dst = (s, g) if src_is_smem else (g, s)
|
|
|
|
def copy_tma(src_idx, dst_idx, **new_kwargs):
|
|
cute.copy(atom, src[None, src_idx], dst[None, dst_idx], **new_kwargs, **kwargs)
|
|
|
|
def copy_tma_single_stage(**new_kwargs):
|
|
cute.copy(atom, src, dst, **new_kwargs, **kwargs)
|
|
|
|
return (copy_tma if const_expr(not single_stage) else copy_tma_single_stage), s, g
|
|
|
|
|
|
def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsync):
|
|
def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs):
|
|
copy(
|
|
src_idx=src_idx,
|
|
dst_idx=producer_state.index,
|
|
tma_bar_ptr=pipeline.producer_get_barrier(producer_state),
|
|
**new_kwargs,
|
|
)
|
|
|
|
return copy_fn
|