diff --git a/bench/benchmark_reduce_sum.py b/bench/benchmark_reduce_sum.py new file mode 100644 index 0000000..f44ad40 --- /dev/null +++ b/bench/benchmark_reduce_sum.py @@ -0,0 +1,47 @@ +import argparse + +import torch + +from forge_cute_py.ops import reduce_sum +from forge_cute_py.ref import reduce_sum as ref_reduce_sum +from forge_cute_py.util.bench import do_bench, estimate_bandwidth, summarize_times + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark copy/transpose") + parser.add_argument("--m", type=int, default=1024) + parser.add_argument("--n", type=int, default=1024) + parser.add_argument("--dtype", choices=["float16", "bfloat16", "float32"], default="float16") + parser.add_argument("--dim", type=int, default=-1) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--compile-ref", action="store_true") + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA required for benchmarking") + + dtype = getattr(torch, args.dtype) + x = torch.randn(args.m, args.n, device="cuda", dtype=dtype) + dim = args.dim + + def fn(): + return reduce_sum(x, dim=args.dim) + + times = do_bench(fn, warmup=args.warmup, rep=args.iterations) + stats = summarize_times(times) + bytes_moved = (x.numel() + x.numel() / x.shape[dim]) * x.element_size() + bw = estimate_bandwidth(bytes_moved, stats["p50_ms"]) + print(f"copy_transpose p50: {stats['p50_ms']:.4f} ms, BW: {bw:.2f} GB/s") + + ref = lambda: ref_reduce_sum(x, dim=dim) + if args.compile_ref and hasattr(torch, "compile"): + ref = torch.compile(ref, fullgraph=True) + ref_times = do_bench(ref, warmup=args.warmup, rep=args.iterations) + ref_stats = summarize_times(ref_times) + ref_bw = estimate_bandwidth(bytes_moved, ref_stats["p50_ms"]) + print(f"reference p50: {ref_stats['p50_ms']:.4f} ms, BW: {ref_bw:.2f} GB/s") + + +if __name__ == "__main__": + main() diff --git a/forge_cute_py/kernels/reduce_sum.py b/forge_cute_py/kernels/reduce_sum.py new file mode 100644 index 0000000..83d8f57 --- /dev/null +++ b/forge_cute_py/kernels/reduce_sum.py @@ -0,0 +1,289 @@ +""" +Reduction kernel using CuTe DSL. + +Implements reduction over a specified dimension: + dim=1/-1 (row reduction): mO[row] = sum_{j} mX[row, j] + dim=0 (column reduction): mO[col] = sum_{i} mX[i, col] + +Variants (for row reduction): +- naive: vecsize=1 (no vectorization), 32 threads/row (single warp) +- improved: vecsize=128/dtype.width (vectorized loads), 32 threads/row (single warp) +- shfl: vecsize=128/dtype.width, up to 128 threads/row (multi-warp with shared memory) + +Note: Column reduction (dim=0) uses a simpler per-column accumulation strategy +since column elements are strided in memory (not amenable to vectorized loads). +""" + +from typing import Literal, Type + +import cutlass +import cutlass.cute as cute + +ReductionOp = Literal["sum", "amax", "amin", "prod"] +Variant = Literal["naive", "improved", "shfl"] + +# Warp size is always 32 on NVIDIA GPUs +WARP_SIZE = 32 + + +class Reduction: + _NUM_THREADS = 128 + _VEC_LOAD_BITS = 128 + + def __init__( + self, + dtype: Type[cutlass.Numeric], + N: int, + M: int, + reduction_dtype: Type[cutlass.Numeric] | None = cutlass.Float32, + reduction_op: ReductionOp = "sum", + dim: int = -1, + variant: Variant = "shfl", + ): + self.dtype = dtype + self.N = int(N) + self.M = int(M) + self.reduction_dtype = dtype if reduction_dtype is None else reduction_dtype + self.reduction_op = reduction_op + self.dim = dim + self.variant = variant + + self._validate_config() + + def _validate_config(self) -> None: + if self.dim not in (-1, 0, 1): + raise ValueError(f"dim must be -1, 0, or 1. Got: {self.dim}") + + if self.variant not in ("naive", "improved", "shfl"): + raise ValueError(f"Unknown variant={self.variant}") + + if self.reduction_op != "sum": + raise NotImplementedError(f"Only support reduction_op=sum, got {self.reduction_op}") + + def _threads_per_row(self) -> int: + """ + Threads per row based on variant. + + - naive/improved: 32 threads (one warp) + - shfl: up to 128 threads (4 warps) for large N + """ + if self.variant == "shfl" and self.N >= 512: + return min(self._NUM_THREADS, 128) + return WARP_SIZE + + def _pick_vecsize(self) -> int: + """ + Number of elements per vector load. + + - naive: vecsize=1 + - improved/shfl: target 128-bit loads, reduced if N not divisible + """ + if self.variant == "naive": + return 1 + + elems_per_128b = self._VEC_LOAD_BITS // self.dtype.width + vecsize = max(1, elems_per_128b) + + while vecsize > 1 and (self.N % vecsize) != 0: + vecsize //= 2 + + return vecsize + + def _adjust_n_blocks(self, n_blocks: int) -> int: + """Adjust n_blocks to avoid power-of-2 values that can cause codegen issues.""" + is_pow2 = n_blocks.bit_count() == 1 + return n_blocks + 1 if n_blocks >= 8 and is_pow2 else n_blocks + + def _get_tiled_copy(self, vecsize: int): + """ + Build tile shape (tileM, tileN) and tiled copy operator. + """ + threads_per_row = self._threads_per_row() + tile_m = self._NUM_THREADS // threads_per_row + + # Cover N in blocks of threads_per_row; tileN becomes multiple of (threads_per_row * vecsize). + n_vec_elems = (self.N + vecsize - 1) // vecsize + n_blocks = (n_vec_elems + threads_per_row - 1) // threads_per_row + n_blocks = self._adjust_n_blocks(n_blocks) + + tile_n = vecsize * n_blocks * threads_per_row + tiler_mn = (tile_m, tile_n) + + num_copy_bits = vecsize * self.dtype.width + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dtype, + num_bits_per_copy=num_copy_bits, + ) + + thr_layout = cute.make_ordered_layout( + (tile_m, threads_per_row), + order=(1, 0), + ) + val_layout = cute.make_layout((1, vecsize)) + tiled_copy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout) + + return tiler_mn, tiled_copy, threads_per_row + + @cute.jit + def __call__(self, mX: cute.Tensor, mO: cute.Tensor, stream=None): + # Dispatch to row or column reduction kernel + if self.dim == 0: + self._call_col_reduce(mX, mO, stream) + else: + self._call_row_reduce(mX, mO, stream) + + def _call_row_reduce(self, mX: cute.Tensor, mO: cute.Tensor, stream=None): + """Row reduction: out[m] = sum_n x[m, n]""" + vecsize = self._pick_vecsize() + tiler_mn, tiled_copy, threads_per_row = self._get_tiled_copy(vecsize=vecsize) + + num_threads = tiled_copy.size + warps_per_row = threads_per_row // WARP_SIZE + + self.kernel_row_reduce( + mX, + mO, + tiler_mn, + tiled_copy, + threads_per_row, + warps_per_row, + self.M, + self.N, + ).launch( + grid=[cute.ceil_div(self.M, tiler_mn[0]), 1, 1], + block=[num_threads, 1, 1], + stream=stream, + ) + + def _call_col_reduce(self, mX: cute.Tensor, mO: cute.Tensor, stream=None): + """Column reduction: out[n] = sum_m x[m, n]""" + num_threads = self._NUM_THREADS + + self.kernel_col_reduce( + mX, + mO, + self.M, + self.N, + ).launch( + grid=[cute.ceil_div(self.N, num_threads), 1, 1], + block=[num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel_row_reduce( + self, + mX: cute.Tensor, + mO: cute.Tensor, + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + threads_per_row: cutlass.Constexpr[int], + warps_per_row: cutlass.Constexpr[int], + M: int, + N: int, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + tile_m = tiler_mn[0] + tile_n = tiler_mn[1] + num_threads = threads_per_row * tile_m + + # Tile of X: (tileM, tileN) + gX = cute.local_tile(mX, tiler_mn, (bidx, 0)) + + thr_copy = tiled_copy.get_slice(tidx) + tXgX = thr_copy.partition_S(gX) + tXrX = cute.make_rmem_tensor_like(tXgX) + + # Predicate tensor for bounds checking + tXcX = thr_copy.partition_S(cute.make_identity_tensor((tile_m, tile_n))) + + # Two-phase predicated copy: + # - Use element 0 to construct a "zero" of the right type + first_val = tXgX[0] + zero_val = first_val - first_val + + for i in range(cute.size(tXrX)): + coord = tXcX[i] + row = coord[0] + (bidx * tile_m) + col = coord[1] + + if row < M and col < N: + tXrX[i] = tXgX[i] + else: + tXrX[i] = zero_val + + # Register accumulation then warp shuffle reduction + x = tXrX.load().to(self.reduction_dtype) + val = x.reduce(cute.ReductionOp.ADD, init_val=0.0, reduction_profile=0) + val = cute.arch.warp_reduction_sum(val) + + lane_id = cute.arch.lane_idx() + warp_id = cute.arch.warp_idx() + + if warps_per_row == 1: + # One warp per row: lane 0 writes output + if lane_id == 0: + global_row = warp_id + (tile_m * bidx) + if global_row < M: + mO[global_row] = val.to(self.dtype) + else: + # Multi-warp per row: shared memory for inter-warp reduction + smem = cutlass.utils.SmemAllocator() + num_warps = num_threads // cute.arch.WARP_SIZE + smem_layout = cute.make_layout((num_warps,), stride=(1,)) + partials = smem.allocate_tensor( + self.reduction_dtype, + smem_layout, + byte_alignment=16, + ) + + if lane_id == 0: + partials[warp_id] = val + + cute.arch.sync_threads() + + # Warp 0..(tile_m-1) each reduce their row's partials using lanes [0..warps_per_row-1] + if warp_id < tile_m and lane_id < warps_per_row: + row_warp_base = warp_id * warps_per_row + partial_val = partials[row_warp_base + lane_id] + final_sum = cute.arch.warp_reduction_sum(partial_val) + + if lane_id == 0: + row_global = warp_id + (tile_m * bidx) + if row_global < M: + mO[row_global] = final_sum.to(self.dtype) + + @cute.kernel + def kernel_col_reduce( + self, + mX: cute.Tensor, + mO: cute.Tensor, + M: int, + N: int, + ): + """ + Column reduction kernel: out[n] = sum_m x[m, n] + + Each thread handles one column, iterating over all rows. + This avoids the transpose + contiguous copy at the cost of strided memory access. + """ + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + num_threads = self._NUM_THREADS + col = tidx + bidx * num_threads + + if col < N: + # Initialize accumulator from first element, converted to reduction dtype + first_val = mX[0, col] + acc = first_val.to(self.reduction_dtype) - first_val.to(self.reduction_dtype) + + # Accumulate all rows for this column + for row in range(M): + val = mX[row, col] + acc = acc + val.to(self.reduction_dtype) + + mO[col] = acc.to(self.dtype) diff --git a/forge_cute_py/ops/reduce_sum.py b/forge_cute_py/ops/reduce_sum.py index d501eca..15ca02f 100644 --- a/forge_cute_py/ops/reduce_sum.py +++ b/forge_cute_py/ops/reduce_sum.py @@ -1,35 +1,138 @@ +from __future__ import annotations + +from typing import Callable, Type + +import cutlass +import cutlass.cute as cute import torch +from cutlass import BFloat16, Float16, Float32 + +from forge_cute_py.kernels.reduce_sum import Reduction, Variant + +_DTYPE_TO_CUTLASS: dict[torch.dtype, Type[cutlass.Numeric]] = { + torch.float16: Float16, + torch.float32: Float32, + torch.bfloat16: BFloat16, +} + +_SUPPORTED_VARIANTS: tuple[str, ...] = ("naive", "improved", "shfl") +_SUPPORTED_DTYPE: tuple[torch.dtype, ...] = tuple(_DTYPE_TO_CUTLASS.keys()) + +# Compile cache keyed by: (cute_dtype, dim, variant, M, N) +_COMPILE_CACHE: dict[tuple[Type[cutlass.Numeric], int, str, int, int], Callable] = {} + + +def _normalize_dim(ndim: int, dim: int) -> int: + if dim < 0: + return ndim + dim + return dim + + +def _get_cute_dtype(dtype: torch.dtype) -> Type[cutlass.Numeric]: + if dtype not in _DTYPE_TO_CUTLASS: + raise AssertionError(f"Unsupported dtype: {dtype}") + return _DTYPE_TO_CUTLASS[dtype] + + +def _compile_reduction_kernel( + cute_dtype: Type[cutlass.Numeric], + M: int, + N: int, + dim: int = 1, + variant: Variant = "shfl", +) -> Callable: + m = cute.sym_int() + n = cute.sym_int() + input_cute = cute.runtime.make_fake_compact_tensor( + cute_dtype, + (m, n), + stride_order=(1, 0), + ) + # Output shape depends on dim: (M,) for row reduction, (N,) for column reduction + out_sym = m if dim == 1 else n + output_cute = cute.runtime.make_fake_compact_tensor(cute_dtype, (out_sym,)) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + return cute.compile( + Reduction( + cute_dtype, + N=N, + M=M, + dim=dim, + variant=variant, + ), + input_cute, + output_cute, + fake_stream, + options="--enable-tvm-ffi", + ) + + +def _validate_inputs(x: torch.Tensor, out: torch.Tensor | None, dim: int, variant: str) -> None: + assert x.dim() == 2, "reduce_sum expects a 2D tensor" + assert x.is_cuda, f"reduce_sum is CUDA-only, got device={x.device}" + assert dim in (-1, 0, 1), f"reduce_sum expects dim in {{-1, 0, 1}}, got {dim}" + assert variant in _SUPPORTED_VARIANTS, f"Unsupported variant: {variant}" + assert x.dtype in _SUPPORTED_DTYPE, f"Unsupported dtype: {x.dtype}" + + if out is None: + return + + assert out.is_cuda, "out must be CUDA tensor" + assert out.is_contiguous(), "out must be contiguous" + assert out.dtype == x.dtype, "out dtype must match x dtype" + + +def _get_or_compile_kernel( + cute_dtype: Type[cutlass.Numeric], + dim: int, + variant: str, + M: int, + N: int, +) -> Callable: + key = (cute_dtype, dim, variant, M, N) + kernel = _COMPILE_CACHE.get(key) + if kernel is None: + kernel = _compile_reduction_kernel(cute_dtype, M, N, dim=dim, variant=variant) + _COMPILE_CACHE[key] = kernel + return kernel @torch.library.custom_op("forge_cute_py::_reduce_sum", mutates_args={"out"}) -def _reduce_sum(x: torch.Tensor, out: torch.Tensor, dim: int = -1, variant: str = "shfl") -> None: - """Row/column sum reduction (reference implementation stub). +def _reduce_sum( + x: torch.Tensor, + out: torch.Tensor, + dim: int = -1, + variant: str = "shfl", +) -> None: + """Row/column sum reduction using CuTe kernel. Args: x: Input tensor of shape (M, N) out: Output tensor (mutated in-place) dim: Dimension to reduce over (-1, 0, or 1) - variant: Reduction variant (naive, improved, shfl) - currently unused + variant: Reduction variant ("naive", "improved", "shfl") """ - assert x.dim() == 2, "reduce_sum expects a 2D tensor" - assert x.is_cuda, f"reduce_sum is CUDA-only, got device={x.device}" - assert dim in (-1, 0, 1), f"reduce_sum expects dim in {{-1, 0, 1}}, got {dim}" - assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], ( - f"Unsupported dtype: {x.dtype}" - ) + _validate_inputs(x, out, dim=dim, variant=variant) - # Normalize dim to positive index - dim = dim if dim >= 0 else x.ndim + dim + normalized_dim = _normalize_dim(x.ndim, dim) + M, N = x.shape + cute_dtype = _get_cute_dtype(x.dtype) - # For now, use reference implementation - # Future: call kernel implementation based on variant when available - from forge_cute_py.ref import reduce_sum as reduce_sum_ref + if normalized_dim == 1: + assert out.shape == (M,), f"out must be shape {(M,)}, got {tuple(out.shape)}" + kernel = _get_or_compile_kernel(cute_dtype, dim=1, variant=variant, M=M, N=N) + kernel(x, out) + return - result = reduce_sum_ref(x, dim=dim) - out.copy_(result) + if normalized_dim == 0: + assert out.shape == (N,), f"out must be shape {(N,)}, got {tuple(out.shape)}" + # Direct column reduction: out[n] = sum_m x[m, n] + kernel = _get_or_compile_kernel(cute_dtype, dim=0, variant=variant, M=M, N=N) + kernel(x, out) + return - -_reduce_sum.compile_cache = {} + raise ValueError(f"Invalid dim={normalized_dim} for 2D tensor") def reduce_sum(x: torch.Tensor, dim: int = -1, variant: str = "shfl") -> torch.Tensor: @@ -37,29 +140,15 @@ def reduce_sum(x: torch.Tensor, dim: int = -1, variant: str = "shfl") -> torch.T Args: x: Input tensor of shape (M, N) - dim: Dimension to reduce over (-1 for last dim, 0 or 1) - variant: Reduction variant (naive, improved, shfl) - currently unused + dim: -1 / 0 / 1 + variant: "naive" / "improved" / "shfl" + """ + _validate_inputs(x, out=None, dim=dim, variant=variant) - Returns: - Reduced tensor of shape (M,) if dim=1 or (N,) if dim=0 + normalized_dim = _normalize_dim(x.ndim, dim) + M, N = x.shape + out_shape = (M,) if normalized_dim == 1 else (N,) + out = torch.empty(out_shape, device=x.device, dtype=x.dtype) - Examples: - >>> x = torch.randn(32, 128, device='cuda', dtype=torch.float16) - >>> y = reduce_sum(x, dim=-1) # Sum over columns, result shape: (32,) - >>> y.shape - torch.Size([32]) - """ - # Normalize dim to positive index - dim = dim if dim >= 0 else x.ndim + dim - - # Determine output shape - if dim == 0: - out_shape = (x.shape[1],) - elif dim == 1: - out_shape = (x.shape[0],) - else: - raise ValueError(f"Invalid dim={dim} for 2D tensor") - - out = torch.empty(out_shape, dtype=x.dtype, device=x.device) - _reduce_sum(x, out, dim, variant) + _reduce_sum(x, out, dim=normalized_dim, variant=variant) return out diff --git a/tests/test_reduce_sum.py b/tests/test_reduce_sum.py index dd773b5..27d5b10 100644 --- a/tests/test_reduce_sum.py +++ b/tests/test_reduce_sum.py @@ -8,7 +8,19 @@ @pytest.mark.parametrize( "shape, dim", [ + # Small shapes ((4, 8), -1), + ((8, 4), -1), + ((3, 7), -1), # non-power-of-2 + ((1, 16), -1), # single row + ((7, 1), -1), # single column + # Medium shapes + ((32, 64), -1), + ((64, 128), -1), + # Large shapes + ((4096, 1024), 1), + ((1024, 4096), -1), + # Edge cases with dim=0 (handled via transpose) ((8, 4), 0), ], ) @@ -25,8 +37,8 @@ def test_reduce_sum_correctness(shape, dim, dtype, atol, rtol, variant): x = torch.randn(*shape, device="cuda", dtype=dtype) try: y = reduce_sum(x, dim=dim, variant=variant) - except NotImplementedError: - pytest.skip(f"reduce_sum variant {variant} not implemented") + except NotImplementedError as e: + pytest.skip(f"reduce_sum variant {variant} not implemented\n{e}") y_ref = ref_reduce_sum(x, dim=dim) torch.testing.assert_close(y, y_ref, atol=atol, rtol=rtol)