feat: add bincount operator implementation, tests and benchmark

Author: factnn

benchmark/test_special_perf.py

@@ -1535,3 +1535,81 @@ def unique_consecutive_input_fn(shape, dtype, device):
         dtypes=INT_DTYPES,
     )
     bench.run()
+
+
+class BincountBenchmark(Benchmark):
+    """Benchmark for bincount operation."""
+
+    def __init__(self, op_name, torch_op, dtypes):
+        super().__init__(op_name=op_name, torch_op=torch_op, dtypes=dtypes)
+
+    def set_shapes(self, shape_file_path=None):
+        bincount_configs = [
+            (1000, 100),
+            (10000, 100),
+            (10000, 1000),
+            (100000, 100),
+            (100000, 1000),
+            (1000000, 100),
+            (1000000, 1000),
+            (1000000, 10000),
+        ]
+        self.shapes = bincount_configs
+
+    def get_input_iter(self, cur_dtype):
+        for config in self.shapes:
+            yield from self.bincount_input_fn(config, cur_dtype, self.device)
+
+    def bincount_input_fn(self, config, dtype, device):
+        input_size, max_val = config
+        inp = torch.randint(0, max_val, (input_size,), dtype=torch.int64, device=device)
+        yield inp,
+
+
+class BincountWeightsBenchmark(Benchmark):
+    """Benchmark for bincount operation with weights."""
+
+    def __init__(self, op_name, torch_op, dtypes):
+        super().__init__(op_name=op_name, torch_op=torch_op, dtypes=dtypes)
+
+    def set_shapes(self, shape_file_path=None):
+        bincount_configs = [
+            (1000, 100),
+            (10000, 100),
+            (10000, 1000),
+            (100000, 100),
+            (100000, 1000),
+            (1000000, 100),
+            (1000000, 1000),
+        ]
+        self.shapes = bincount_configs
+
+    def get_input_iter(self, cur_dtype):
+        for config in self.shapes:
+            yield from self.bincount_weights_input_fn(config, cur_dtype, self.device)
+
+    def bincount_weights_input_fn(self, config, dtype, device):
+        input_size, max_val = config
+        inp = torch.randint(0, max_val, (input_size,), dtype=torch.int64, device=device)
+        weights = torch.randn(input_size, dtype=dtype, device=device)
+        yield inp, {"weights": weights}
+
+
+@pytest.mark.bincount
+def test_perf_bincount():
+    bench = BincountBenchmark(
+        op_name="bincount",
+        torch_op=torch.bincount,
+        dtypes=[torch.int64],
+    )
+    bench.run()
+
+
+@pytest.mark.bincount
+def test_perf_bincount_with_weights():
+    bench = BincountWeightsBenchmark(
+        op_name="bincount_with_weights",
+        torch_op=torch.bincount,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()

src/flag_gems/__init__.py

@@ -115,6 +115,7 @@ def torch_ge(v):
     ("bitwise_or_.Scalar", bitwise_or_scalar_),
     ("bitwise_or_.Tensor", bitwise_or_tensor_),
     ("bitwise_right_shift", bitwise_right_shift),
+    ("bincount", bincount),
     ("bmm", bmm),
     ("bmm.out", bmm_out),
     ("cat", cat),

src/flag_gems/ops/__init__.py

@@ -42,6 +42,7 @@
 from flag_gems.ops.baddbmm import baddbmm
 from flag_gems.ops.batch_norm import batch_norm, batch_norm_backward
 from flag_gems.ops.bernoulli_ import bernoulli_
+from flag_gems.ops.bincount import bincount
 from flag_gems.ops.bitwise_and import (
     bitwise_and_scalar,
     bitwise_and_scalar_,
@@ -396,6 +397,7 @@
     "bitwise_or_tensor",
     "bitwise_or_tensor_",
     "bitwise_right_shift",
+    "bincount",
     "bmm",
     "bmm_out",
     "cat",

src/flag_gems/ops/bincount.py

@@ -0,0 +1,138 @@
+import logging
+
+import torch
+import triton
+import triton.language as tl
+
+from ..utils import libentry
+
+logger = logging.getLogger(__name__)
+
+
+@libentry()
+@triton.jit
+def bincount_kernel(
+    inp_ptr,
+    out_ptr,
+    N,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Kernel for bincount without weights."""
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < N
+
+    # Load input values (indices)
+    indices = tl.load(inp_ptr + offsets, mask=mask, other=0)
+
+    # Atomic add 1 to the output at each index
+    # Use int64 for the atomic add
+    ones = tl.full((BLOCK_SIZE,), 1, dtype=tl.int64)
+    tl.atomic_add(out_ptr + indices, ones, mask=mask, sem="relaxed")
+
+
+@libentry()
+@triton.jit
+def bincount_weights_kernel(
+    inp_ptr,
+    weights_ptr,
+    out_ptr,
+    N,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Kernel for bincount with weights."""
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < N
+
+    # Load input values (indices) and weights
+    indices = tl.load(inp_ptr + offsets, mask=mask, other=0)
+    weights = tl.load(weights_ptr + offsets, mask=mask, other=0.0)
+
+    # Atomic add weights to the output at each index
+    tl.atomic_add(out_ptr + indices, weights, mask=mask, sem="relaxed")
+
+
+def bincount(inp, weights=None, minlength=0):
+    """
+    Count the frequency of each value in an array of non-negative ints.
+
+    Args:
+        inp: 1-d int tensor of non-negative integers
+        weights: optional weights tensor of same size as inp
+        minlength: optional minimum number of bins
+
+    Returns:
+        Tensor of shape (max(inp) + 1,) or (minlength,) if minlength is larger
+    """
+    logger.debug("GEMS BINCOUNT")
+
+    # Input validation
+    assert inp.ndim == 1, "bincount only supports 1-d tensors"
+    assert inp.dtype in (
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.uint8,
+    ), "bincount only supports integer tensors"
+
+    N = inp.numel()
+
+    # Handle empty input
+    if N == 0:
+        if weights is not None:
+            return torch.zeros(minlength, dtype=weights.dtype, device=inp.device)
+        return torch.zeros(minlength, dtype=torch.int64, device=inp.device)
+
+    # Compute output size
+    max_val = int(inp.max().item())
+    output_size = max(max_val + 1, minlength)
+
+    # Ensure input is contiguous
+    inp = inp.contiguous()
+
+    if weights is not None:
+        assert weights.shape == inp.shape, "weights must have same shape as input"
+        weights = weights.contiguous()
+
+        # Output dtype matches weights dtype
+        # For atomic_add compatibility, convert to float32 if float16/bfloat16
+        weights_dtype = weights.dtype
+        if weights_dtype in (torch.float16, torch.bfloat16):
+            weights = weights.to(torch.float32)
+            out = torch.zeros(output_size, dtype=torch.float32, device=inp.device)
+        else:
+            out = torch.zeros(output_size, dtype=weights.dtype, device=inp.device)
+
+        BLOCK_SIZE = 1024
+        grid = (triton.cdiv(N, BLOCK_SIZE),)
+
+        bincount_weights_kernel[grid](
+            inp,
+            weights,
+            out,
+            N,
+            BLOCK_SIZE=BLOCK_SIZE,
+        )
+
+        # Convert back if needed
+        if weights_dtype in (torch.float16, torch.bfloat16):
+            out = out.to(weights_dtype)
+
+        return out
+    else:
+        # No weights: count occurrences
+        out = torch.zeros(output_size, dtype=torch.int64, device=inp.device)
+
+        BLOCK_SIZE = 1024
+        grid = (triton.cdiv(N, BLOCK_SIZE),)
+
+        bincount_kernel[grid](
+            inp,
+            out,
+            N,
+            BLOCK_SIZE=BLOCK_SIZE,
+        )
+
+        return out

tests/test_bincount.py

@@ -0,0 +1,116 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+from .conftest import QUICK_MODE
+
+BINCOUNT_SIZES = [16, 100, 1024, 10000] if not QUICK_MODE else [100, 1024]
+BINCOUNT_MAXVALS = [10, 100, 1000] if not QUICK_MODE else [100]
+
+
+@pytest.mark.bincount
+@pytest.mark.parametrize("size", BINCOUNT_SIZES)
+@pytest.mark.parametrize("max_val", BINCOUNT_MAXVALS)
+def test_accuracy_bincount(size, max_val):
+    """Test bincount without weights."""
+    inp = torch.randint(0, max_val, (size,), dtype=torch.int64, device=flag_gems.device)
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.bincount(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.bincount(inp)
+
+    utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.bincount
+@pytest.mark.parametrize("size", BINCOUNT_SIZES)
+@pytest.mark.parametrize("max_val", BINCOUNT_MAXVALS)
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+def test_accuracy_bincount_with_weights(size, max_val, dtype):
+    """Test bincount with weights."""
+    inp = torch.randint(0, max_val, (size,), dtype=torch.int64, device=flag_gems.device)
+    weights = torch.randn(size, dtype=dtype, device=flag_gems.device)
+    ref_inp = utils.to_reference(inp)
+    ref_weights = utils.to_reference(weights)
+
+    ref_out = torch.bincount(ref_inp, weights=ref_weights)
+    with flag_gems.use_gems():
+        res_out = torch.bincount(inp, weights=weights)
+
+    utils.gems_assert_close(res_out, ref_out, dtype)
+
+
+@pytest.mark.bincount
+@pytest.mark.parametrize("size", BINCOUNT_SIZES)
+@pytest.mark.parametrize("max_val", BINCOUNT_MAXVALS)
+@pytest.mark.parametrize("minlength", [0, 50, 2000])
+def test_accuracy_bincount_with_minlength(size, max_val, minlength):
+    """Test bincount with minlength parameter."""
+    inp = torch.randint(0, max_val, (size,), dtype=torch.int64, device=flag_gems.device)
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.bincount(ref_inp, minlength=minlength)
+    with flag_gems.use_gems():
+        res_out = torch.bincount(inp, minlength=minlength)
+
+    utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.bincount
+def test_accuracy_bincount_empty():
+    """Test bincount with empty input."""
+    inp = torch.tensor([], dtype=torch.int64, device=flag_gems.device)
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.bincount(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.bincount(inp)
+
+    utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.bincount
+def test_accuracy_bincount_single():
+    """Test bincount with single element."""
+    inp = torch.tensor([5], dtype=torch.int64, device=flag_gems.device)
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.bincount(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.bincount(inp)
+
+    utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.bincount
+def test_accuracy_bincount_all_zeros():
+    """Test bincount with all zeros."""
+    inp = torch.zeros(100, dtype=torch.int64, device=flag_gems.device)
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.bincount(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.bincount(inp)
+
+    utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.bincount
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+def test_accuracy_bincount_weights_edge_cases(dtype):
+    """Test bincount with edge case weights."""
+    inp = torch.tensor([0, 1, 2, 1, 0], dtype=torch.int64, device=flag_gems.device)
+    weights = torch.tensor(
+        [1.0, 2.0, 3.0, 4.0, 5.0], dtype=dtype, device=flag_gems.device
+    )
+    ref_inp = utils.to_reference(inp)
+    ref_weights = utils.to_reference(weights)
+
+    ref_out = torch.bincount(ref_inp, weights=ref_weights)
+    with flag_gems.use_gems():
+        res_out = torch.bincount(inp, weights=weights)
+
+    utils.gems_assert_close(res_out, ref_out, dtype)