Wire reduce_sum op to WIP kernel with safe fallback

Author: debashishc

forge_cute_py/kernels/reduce_sum.py

@@ -20,7 +20,7 @@ def block_reduce(
     """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])
+    num_warps = cute.size(reduction_buffer)
 
     if lane_idx == 0:
         reduction_buffer[warp_idx] = val
@@ -56,7 +56,7 @@ def row_reduce(
         warp_op,
         threads_in_group=min(threads_per_row, cute.arch.WARP_SIZE),
     )
-    if const_expr(cute.size(reduction_buffer.shape[0]) > 1):
+    if const_expr(cute.size(reduction_buffer) > 1):
         val = block_reduce(val, warp_op, reduction_buffer, init_val=init_val)
     return val
 
@@ -132,12 +132,7 @@ def kernel(
             tXgX = thr_copy.partition_S(gX)
             tXsX = thr_copy.partition_D(sX)
             tXrX = cute.make_fragment_like(tXgX)
-            tXpX = cute.make_fragment_like(tXgX, cutlass.Boolean)
-            base = cutlass.Int32(tile) * tile_n + cutlass.Int32(tidx) * vec_size
-            limit = cutlass.Int32(n_cols)
-            for v in cutlass.range_constexpr(tXpX.shape[0]):
-                tXpX[v] = cute.elem_less(base + v, limit)
-            cute.copy(copy_atom, tXgX, tXsX, pred=tXpX)
+            cute.copy(copy_atom, tXgX, tXsX)
             cute.arch.cp_async_commit_group()
             cute.arch.cp_async_wait_group(0)
             cute.arch.barrier()

forge_cute_py/ops/reduce_sum.py

@@ -1,60 +1,113 @@
+import cutlass.cute as cute
 import torch
+from cutlass import BFloat16, Float16, Float32
+from cutlass.cute.runtime import from_dlpack
+
+from forge_cute_py.kernels.reduce_sum import ReduceSumRow
+
+
+def _reduce_sum_ref_fallback(x: torch.Tensor, out: torch.Tensor, dim: int) -> None:
+    from forge_cute_py.ref import reduce_sum as reduce_sum_ref
+
+    out.copy_(reduce_sum_ref(x, dim=dim))
 
 
 @torch.library.custom_op("forge_cute_py::_reduce_sum", mutates_args={"out"})
 def _reduce_sum(x: torch.Tensor, out: torch.Tensor, dim: int = -1) -> None:
-    """Row/column sum reduction (reference implementation stub).
+    """Row-wise sum reduction using CuTe DSL.
 
     Args:
         x: Input tensor of shape (M, N)
         out: Output tensor (mutated in-place)
-        dim: Dimension to reduce over (-1, 0, or 1)
+        dim: Dimension to reduce over (-1 or 1)
     """
     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 dim in (-1, 1), f"reduce_sum expects dim in {{-1, 1}} (row-wise), got {dim}"
     assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], (
         f"Unsupported dtype: {x.dtype}"
     )
 
-    # Normalize dim to positive index
     dim = dim if dim >= 0 else x.ndim + dim
     if dim != 1:
         raise ValueError(f"reduce_sum supports dim in {{-1, 1}} (row-wise), got {dim}")
 
-    # 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
-
-    result = reduce_sum_ref(x, dim=dim)
-    out.copy_(result)
+    dtype_map = {
+        torch.float16: Float16,
+        torch.float32: Float32,
+        torch.bfloat16: BFloat16,
+    }
+    cute_dtype = dtype_map[x.dtype]
+    block_size = 256
+    elem_bytes = x.element_size()
+    vec_size = 16 // elem_bytes
+    tile_n = block_size * vec_size
+
+    if x.shape[1] < vec_size:
+        raise ValueError(
+            f"reduce_sum requires N >= {vec_size} for 128-bit vectorized loads. Got N={x.shape[1]}."
+        )
+    if x.data_ptr() % 16 != 0 or (x.stride(0) * elem_bytes) % 16 != 0:
+        raise ValueError(
+            "reduce_sum requires 16-byte aligned rows for vectorized loads. "
+            f"Got data_ptr alignment={x.data_ptr() % 16} and row_stride_bytes={x.stride(0) * elem_bytes}."
+        )
+    if x.shape[1] % tile_n != 0:
+        _reduce_sum_ref_fallback(x, out, dim)
+        return
+
+    compile_key = (cute_dtype, block_size)
+    try:
+        if compile_key not in _reduce_sum.compile_cache:
+            m = cute.sym_int()
+            n = cute.sym_int()
+            input_cute = cute.runtime.make_fake_compact_tensor(
+                cute_dtype, (m, n), stride_order=(1, 0), assumed_align=16
+            )
+            output_cute = cute.runtime.make_fake_compact_tensor(
+                cute_dtype, (m,), stride_order=(0,), assumed_align=16
+            )
+            _reduce_sum.compile_cache[compile_key] = cute.compile(
+                ReduceSumRow(cute_dtype, block_size=block_size),
+                input_cute,
+                output_cute,
+                options="--enable-tvm-ffi",
+            )
+
+        x_cute = from_dlpack(x, assumed_align=16, enable_tvm_ffi=True)
+        out_cute = from_dlpack(out, assumed_align=16, enable_tvm_ffi=True)
+        _reduce_sum.compile_cache[compile_key](x_cute, out_cute)
+    except Exception:
+        _reduce_sum_ref_fallback(x, out, dim)
 
 
 _reduce_sum.compile_cache = {}
 
 
 def reduce_sum(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
-    """Row/column sum reduction.
+    """Row-wise sum reduction.
 
     Args:
         x: Input tensor of shape (M, N)
         dim: Dimension to reduce over (-1 for last dim, or 1)
 
     Returns:
-        Reduced tensor of shape (M,) if dim=1 or (N,) if dim=0
+        Reduced tensor of shape (M,)
 
     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
 
     if dim != 1:
         raise ValueError(f"Invalid dim={dim} for row-wise reduce_sum")
 
+    if not x.is_contiguous():
+        x = x.contiguous()
+
     out_shape = (x.shape[0],)
 
     out = torch.empty(out_shape, dtype=x.dtype, device=x.device)