Issue 20 reduce sum

DEVELOPMENT.md

@@ -7,7 +7,7 @@ For quickstart usage, see `README.md`.
 
 v0.1 targets KernelHeim Weeks 0-2:
 - Week 0: copy/transpose
-- Week 1: reductions (sum) variants
+- Week 1: reductions (row-wise sum)
 - Week 2: single-pass online softmax
 
 Out of scope for v0: FlashAttention kernels, KV-cache decode, FP8, NCCL, C++

README.md

@@ -106,7 +106,7 @@ ncu --set full -o profiles/copy_transpose uv run python bench/benchmark_copy_tra
 | Op | Status | Variants | Notes |
 | --- | --- | --- | --- |
 | copy_transpose | Implemented | tile_size=16/32 | CuTe DSL kernel with tiled shared memory |
-| reduce_sum | Stub (ref) | naive/improved/shfl | Uses PyTorch reference; kernel to be implemented |
+| reduce_sum | Stub (ref) | row-wise | Uses PyTorch reference; kernel to be implemented |
 | softmax_online | Stub (ref) | single-pass | Uses PyTorch reference with autograd support; kernel to be implemented |
 
 ---
@@ -174,7 +174,7 @@ For kernel development workflow and architecture details, see:
 ## Roadmap (v0.1 completion)
 
 * [x] Week 0 copy/transpose: end-to-end correctness + benchmark + profile scripts
-* [ ] Week 1 reductions: multiple variants, correctness + benchmark coverage
+* [ ] Week 1 reductions: row-wise sum, correctness + benchmark coverage
 * [ ] Week 2 online softmax: correctness + benchmark coverage + profiling notes
 * [ ] CI: run correctness on supported GPU runners; optional perf smoke checks
 

ROADMAP.md

@@ -39,32 +39,22 @@ Target: Harness infrastructure + Week 0-2 kernel implementations aligned to Kern
 
 ---
 
-## Week 1: Reductions (Sum)
+## Week 1: Reductions (Row-wise Sum)
 
-**Goal**: Multiple variants (naive → improved → shuffle) with correctness + benchmark coverage
+**Goal**: Row-wise sum kernel with correctness + benchmark coverage
 
 ### Infrastructure
 - [x] Test infrastructure (using PyTorch reference)
 - [x] Reference implementation (`forge_cute_py/ref/reduce_sum.py`)
 - [x] Ops registration and API design
 - [x] Benchmark integration
 
-### Kernel Implementations
-- [ ] **Naive variant**: Simple reduction without optimizations
+### Kernel Implementation
+- [ ] **Row-wise sum kernel**
   - [ ] CuTe DSL kernel implementation
   - [ ] Correctness tests vs PyTorch reference
   - [ ] Benchmark baseline
 
-- [ ] **Improved variant**: Optimized reduction with shared memory
-  - [ ] CuTe DSL kernel implementation
-  - [ ] Correctness tests vs PyTorch reference
-  - [ ] Benchmark comparison vs naive
-
-- [ ] **Shuffle variant**: Warp-level shuffle reduction
-  - [ ] CuTe DSL kernel implementation
-  - [ ] Correctness tests vs PyTorch reference
-  - [ ] Benchmark comparison vs improved
-
 - [ ] **Documentation**: Profiling notes and performance analysis
 
 **Status**: ⏳ Test infrastructure ready, kernels pending

bench/benchmark_reduce_sum.py

@@ -0,0 +1,98 @@
+"""Benchmark reduce_sum kernel variants against torch.sum."""
+
+import argparse
+
+import torch
+
+from forge_cute_py.ops.reduce_sum import reduce_sum
+from forge_cute_py.util.bench import do_bench, estimate_bandwidth, summarize_times
+
+DEFAULT_SIZES = [1024, 2048, 4096, 8192]
+DEFAULT_DTYPES = ["float16", "bfloat16", "float32"]
+DEFAULT_DIMS = [1]
+
+
+def parse_int_list(s: str) -> list[int]:
+    return [int(x.strip()) for x in s.split(",")]
+
+
+def parse_str_list(s: str) -> list[str]:
+    return [x.strip() for x in s.split(",")]
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Benchmark reduce_sum variants")
+    parser.add_argument("--sizes", type=parse_int_list, default=DEFAULT_SIZES)
+    parser.add_argument("--dtypes", type=parse_str_list, default=DEFAULT_DTYPES)
+    parser.add_argument("--dims", type=parse_int_list, default=DEFAULT_DIMS)
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iterations", type=int, default=100)
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA required for benchmarking")
+
+    gpu_name = torch.cuda.get_device_name(0)
+    print(f"reduce_sum benchmarks ({gpu_name})")
+    print()
+
+    header = f"{'Size':>6}  {'Dtype':<10} {'Dim':>3}  {'Op':<12} {'p50 (ms)':>10} {'BW (GB/s)':>10} {'vs torch':>10}"
+    print(header)
+    print("-" * len(header))
+
+    for size in args.sizes:
+        for dtype_str in args.dtypes:
+            dtype = getattr(torch, dtype_str)
+            for dim in args.dims:
+                x = torch.randn(size, size, device="cuda", dtype=dtype)
+
+                # Output shape determines output bytes
+                if dim == 1:
+                    out_numel = size
+                else:
+                    out_numel = size
+                input_bytes = x.numel() * x.element_size()
+                output_bytes = out_numel * x.element_size()
+                total_bytes = input_bytes + output_bytes
+
+                # Benchmark torch.sum as reference
+                torch_fn = lambda: torch.sum(x, dim=dim)
+                torch_times = do_bench(torch_fn, warmup=args.warmup, rep=args.iterations)
+                torch_stats = summarize_times(torch_times)
+                torch_p50 = torch_stats["p50_ms"]
+                torch_bw = estimate_bandwidth(total_bytes, torch_p50)
+                print(
+                    f"{size:>6}  {dtype_str:<10} {dim:>3}  {'torch.sum':<14} "
+                    f"{torch_p50:>10.4f} {torch_bw:>10.2f} {1.0:>10.2f}x"
+                )
+
+                if dim not in (-1, 1):
+                    raise ValueError(f"Unsupported dim={dim}; reduce_sum supports only -1 or 1")
+
+                # Warm the JIT cache
+                try:
+                    reduce_sum(x, dim=dim)
+                except Exception as e:
+                    print(
+                        f"{size:>6}  {dtype_str:<10} {dim:>3}  {'reduce_sum':<12} "
+                        f"{'ERROR':>10} {'':>10} {'':>10}  {e}"
+                    )
+                    print()
+                    continue
+
+                fn = lambda d=dim: reduce_sum(x, dim=d)
+                times = do_bench(fn, warmup=args.warmup, rep=args.iterations)
+                stats = summarize_times(times)
+                p50 = stats["p50_ms"]
+                bw = estimate_bandwidth(total_bytes, p50)
+                ratio = p50 / torch_p50 if torch_p50 > 0 else float("inf")
+                print(
+                    f"{size:>6}  {dtype_str:<10} {dim:>3}  {'reduce_sum':<12} "
+                    f"{p50:>10.4f} {bw:>10.2f} {ratio:>10.2f}x"
+                )
+
+                print()
+
+
+if __name__ == "__main__":
+    main()

bench/run.py

@@ -84,16 +84,12 @@ def fn():
         if not hasattr(ops, "reduce_sum"):
             return {"status": "skipped", "reason": "reduce_sum not available"}
         dim = case.get("dim", -1)
-        variant = case.get("variant", "shfl")
         x = torch.randn(*shape, device="cuda", dtype=dtype)
 
         def fn():
-            return ops.reduce_sum(x, dim=dim, variant=variant)
+            return ops.reduce_sum(x, dim=dim)
 
-        try:
-            times = do_bench(fn, warmup=warmup, rep=iterations)
-        except NotImplementedError:
-            return {"status": "skipped", "reason": f"variant {variant} not implemented"}
+        times = do_bench(fn, warmup=warmup, rep=iterations)
         stats = summarize_times(times)
         bytes_moved = _estimate_bytes(op_name, shape, dtype, dim=dim)
         bw = estimate_bandwidth(bytes_moved, stats["p50_ms"])
@@ -103,7 +99,6 @@ def fn():
             "shape": shape,
             "dtype": str(dtype).replace("torch.", ""),
             "dim": dim,
-            "variant": variant,
             "times_ms": stats,
             "bandwidth_gbps": bw,
         }

bench/suites.yaml

@@ -11,7 +11,6 @@ suites:
         shape: [1024, 1024]
         dtype: float16
         dim: -1
-        variant: shfl
       - op: softmax_online
         shape: [256, 2048]
         dtype: float16

forge_cute_py/kernels/reduce_sum.py

@@ -0,0 +1,157 @@
+"""
+Row-wise sum reduction kernel using CuTe DSL (WIP).
+"""
+
+import operator
+from typing import Callable
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import const_expr
+
+
+@cute.jit
+def block_reduce(
+    val: cute.Numeric,
+    op: Callable,
+    reduction_buffer: cute.Tensor,
+    init_val: cute.Numeric,
+) -> cute.Numeric:
+    """Block-wide reduction using warp shuffles + shared memory cross-warp step."""
+    lane_idx = cute.arch.lane_idx()
+    warp_idx = cute.arch.warp_idx()
+    num_warps = cute.size(reduction_buffer.shape[0])
+
+    if lane_idx == 0:
+        reduction_buffer[warp_idx] = val
+    cute.arch.barrier()
+
+    block_reduce_val = init_val
+    if lane_idx < num_warps:
+        block_reduce_val = reduction_buffer[lane_idx]
+    return cute.arch.warp_reduction(block_reduce_val, op)
+
+
+@cute.jit
+def row_reduce(
+    x: cute.TensorSSA | cute.Numeric,
+    op: cute.ReductionOp,
+    threads_per_row: cutlass.Constexpr[int],
+    reduction_buffer: cute.Tensor,
+    init_val: cute.Numeric = 0.0,
+) -> cute.Numeric:
+    """Thread + warp + block reduction for a single row."""
+    if const_expr(isinstance(x, cute.TensorSSA)):
+        val = x.reduce(op, init_val=init_val, reduction_profile=0)
+    else:
+        val = x
+    warp_op = {
+        cute.ReductionOp.ADD: operator.add,
+        cute.ReductionOp.MAX: cute.arch.fmax,
+        cute.ReductionOp.MIN: min,
+        cute.ReductionOp.MUL: operator.mul,
+    }[op]
+    val = cute.arch.warp_reduction(
+        val,
+        warp_op,
+        threads_in_group=min(threads_per_row, cute.arch.WARP_SIZE),
+    )
+    if const_expr(cute.size(reduction_buffer.shape[0]) > 1):
+        val = block_reduce(val, warp_op, reduction_buffer, init_val=init_val)
+    return val
+
+
+class ReduceSumRow:
+    """Row-wise sum reduction (one block per row)."""
+
+    def __init__(self, dtype: type, block_size: int = 256):
+        self.dtype = dtype
+        self.block_size = block_size
+
+    @cute.jit
+    def __call__(
+        self,
+        x: cute.Tensor,
+        out: cute.Tensor,
+    ):
+        """
+        Launch row-wise reduction.
+
+        Args:
+            x: Input tensor of shape (M, N)
+            out: Output tensor of shape (M,)
+            stream: CUDA stream
+        """
+        M, _ = x.shape
+        block = const_expr(self.block_size)
+        self.kernel(x, out).launch(
+            grid=[M, 1, 1],
+            block=[block, 1, 1],
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        x: cute.Tensor,
+        out: cute.Tensor,
+    ):
+        """Row-wise sum reduction with vectorized loads and block reduction."""
+        tidx, _, _ = cute.arch.thread_idx()
+        row_idx, _, _ = cute.arch.block_idx()
+
+        threads_per_block = const_expr(self.block_size)
+        num_warps = threads_per_block // 32
+
+        smem = cutlass.utils.SmemAllocator()
+        warp_sums = smem.allocate_tensor(
+            cutlass.Float32, layout=cute.make_layout(num_warps), byte_alignment=16
+        )
+
+        row = x[(row_idx, None)]
+        n_cols = row.shape[0]
+        acc = cutlass.Float32(0.0)
+
+        vec_size = const_expr(128 // x.element_type.width)
+        tile_n = threads_per_block * vec_size
+        num_full_tiles = n_cols // tile_n
+
+        sX = smem.allocate_tensor(
+            x.element_type, layout=cute.make_layout(tile_n), byte_alignment=16
+        )
+        copy_atom = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(),
+            x.element_type,
+            num_bits_per_copy=128,
+        )
+        thr_layout = cute.make_layout(threads_per_block)
+        val_layout = cute.make_layout(vec_size)
+        tiled_copy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
+        thr_copy = tiled_copy.get_slice(tidx)
+
+        for tile in cutlass.range(0, num_full_tiles, 1):
+            gX = cute.local_tile(row, (tile_n,), (tile,))
+            tXgX = thr_copy.partition_S(gX)
+            tXsX = thr_copy.partition_D(sX)
+            tXrX = cute.make_fragment_like(tXgX)
+            cute.copy(copy_atom, tXgX, tXsX)
+            cute.arch.cp_async_commit_group()
+            cute.arch.cp_async_wait_group(0)
+            cute.arch.barrier()
+
+            cute.autovec_copy(tXsX, tXrX)
+            x_vec = tXrX.load().to(cutlass.Float32)
+            acc += x_vec.reduce(cute.ReductionOp.ADD, init_val=0.0, reduction_profile=0)
+
+        tail_start = num_full_tiles * tile_n + tidx
+        for col in cutlass.range(tail_start, n_cols, threads_per_block):
+            acc += row[col].to(cutlass.Float32)
+
+        block_val = row_reduce(
+            acc,
+            cute.ReductionOp.ADD,
+            threads_per_row=threads_per_block,
+            reduction_buffer=warp_sums,
+            init_val=cutlass.Float32(0.0),
+        )
+        if tidx == 0:
+            out[row_idx] = block_val.to(out.element_type)