Repository: cupy/cupy

Issue:
cupy.bincount can give incorrect results on CuPy 12.x with CUB enabled
### Description

I encountered some cases where bincount is giving incorrect results. It showed up on CI for cuCIM when using CUDA 12.0 and CuPy 12.3, but does not occur on our CUDA 11 CI runs.

I was also able to reproduce the issue locally from CuPy main and CUDA 12.2.

### To Reproduce

Here is an example taking bincount of `arange` so that all elements of the result should be 1. It can be seen that the output has zeros in some locations:

```py
import cupy as cp

cp.bincount(cp.arange(65536, dtype=cp.int32))
```
gives
```
array([1, 1, 1, ..., 0, 0, 0])
```
and `cp.bincount(cp.arange(65536, dtype=cp.int32)).sum()` is 32768 instead of the expected 65536.


Running the same thing in NumPy or in CuPy with `CUPY_ACCELERATORS=""` set gives the expected result that is all ones.


### Installation

None

### Environment

OS                           : Linux-5.15.0-76-generic-x86_64-with-glibc2.31
Python Version               : 3.10.12
CuPy Version                 : 13.0.0a1
CuPy Platform                : NVIDIA CUDA
NumPy Version                : 1.24.4
SciPy Version                : 1.11.1
Cython Build Version         : 0.29.30
Cython Runtime Version       : 0.29.35
CUDA Root                    : /usr/local/cuda
nvcc PATH                    : /usr/local/cuda/bin/nvcc
CUDA Build Version           : 12020
CUDA Driver Version          : 12020
CUDA Runtime Version         : 12020
cuBLAS Version               : (available)
cuFFT Version                : 11008
cuRAND Version               : 10303
cuSOLVER Version             : (11, 5, 0)
cuSPARSE Version             : (available)
NVRTC Version                : (12, 2)
Thrust Version               : 200101
CUB Build Version            : 200101
Jitify Build Version         : 4a37de0
cuDNN Build Version          : 8902
cuDNN Version                : 8902
NCCL Build Version           : 21803
NCCL Runtime Version         : 21803
cuTENSOR Version             : None
cuSPARSELt Build Version     : None
Device 0 Name                : NVIDIA RTX A6000
Device 0 Compute Capability  : 86
Device 0 PCI Bus ID          : 0000:17:00.0


### Additional Information

I will report the issue to the CUB devs as it seems it is likely an issue in CUB itself.


Discussion:
I forgot to mention, that the correct result is also given even with CUB enabled if the dtype is unsigned instead
```py
cp.bincount(cp.arange(65536, dtype=cp.uint32))
```
array([1, 1, 1, ..., 1, 1, 1])


but gives incorrect result with only 32768 ones again if `dtype=uint16`

Summary (on behalf of @grlee77 and @senior-zero, who have kindly triaged the issue quickly): It seems to happen with CUB shipped with CTK 12(.x?), but not with the upstream version. @senior-zero came up with a C++ reproducer: 
```c++
#include <iostream>
#include <thrust/device_vector.h>
#include <thrust/sequence.h>
#include <thrust/reduce.h>
#include <cub/device/device_histogram.cuh>

int main() {
  int num_samples = 65536;

  thrust::device_vector<int> samples(num_samples);
  thrust::sequence(samples.begin(), samples.end()); // np.arrange(65536) ?

  // size = int(cupy.max(x)) + 1
  int num_bins = num_samples;
  int num_levels = num_bins + 1;
  int lower_level = 0;
  int upper_level = num_bins;

  thrust::device_vector<unsigned long long int> bins(num_bins);

  std::uint8_t *d_temp_storage{};
  std::size_t temp_storage_bytes{};

  int *d_samples = thrust::raw_pointer_cast(samples.data());
  unsigned long long int *d_bins = thrust::raw_pointer_cast(bins.data());

  cub::DeviceHistogram::HistogramEven(d_temp_storage,
                                      temp_storage_bytes,
                                      d_samples,
                                      d_bins,
                                      num_levels,
                                      lower_level,
                                      upper_level,
                                      num_samples);
  
  thrust::device_vector<std::uint8_t> temp_storage(temp_storage_bytes);
  d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());

  cub::DeviceHistogram::HistogramEven(d_temp_storage,
                                      temp_storage_bytes,
                                      d_samples,
                                      d_bins,
                                      num_levels,
                                      lower_level,
                                      upper_level,
                                      num_samples);

  std::cout << thrust::reduce(bins.begin(), bins.end()) << std::endl;
}
```
This adds one more reason to moving to use upstream CCCL asap (https://github.com/cupy/cupy/issues/6700#issuecomment-1543507027). For now, given this and [the other bug](https://github.com/NVIDIA/cccl/issues/912) that we [worked around](https://github.com/cupy/cupy/pull/7316), I think we should just disable CUB histogram until we can pin at upstream instead of CTK.
> t disable CUB histogram until we can pin at upstream instead of CTK.

We think this is the best course of action for now!
cc @anaruse (for awareness)

Resolving patch:
diff --git a/cupy/_statistics/histogram.py b/cupy/_statistics/histogram.py
--- a/cupy/_statistics/histogram.py
+++ b/cupy/_statistics/histogram.py
@@ -5,10 +5,6 @@
 
 import cupy
 from cupy import _core
-from cupy._core import _accelerator
-from cupy.cuda import cub
-from cupy.cuda import common
-from cupy.cuda import runtime
 
 
 # rename builtin range for use in functions that take a range argument
@@ -220,43 +216,10 @@ def histogram(x, bins=10, range=None, weights=None, density=False):
 
     if weights is None:
         y = cupy.zeros(bin_edges.size - 1, dtype=cupy.int64)
-        for accelerator in _accelerator.get_routine_accelerators():
-            # CUB uses int for bin counts
-            # TODO(leofang): support >= 2^31 elements in x?
-            if (accelerator == _accelerator.ACCELERATOR_CUB
-                    and x.size <= 0x7fffffff and bin_edges.size <= 0x7fffffff):
-                # Need to ensure the dtype of bin_edges as it's needed for both
-                # the CUB call and the correction later
-                assert isinstance(bin_edges, cupy.ndarray)
-                if numpy.issubdtype(x.dtype, numpy.integer):
-                    bin_type = float
-                else:
-                    bin_type = numpy.result_type(bin_edges.dtype, x.dtype)
-                    if (bin_type == numpy.float16 and
-                            not common._is_fp16_supported()):
-                        bin_type = numpy.float32
-                    x = x.astype(bin_type, copy=False)
-                acc_bin_edge = bin_edges.astype(bin_type, copy=True)
-                # CUB's upper bin boundary is exclusive for all bins, including
-                # the last bin, so we must shift it to comply with NumPy
-                if x.dtype.kind in 'ui':
-                    acc_bin_edge[-1] += 1
-                elif x.dtype.kind == 'f':
-                    last = acc_bin_edge[-1]
-                    acc_bin_edge[-1] = cupy.nextafter(last, last + 1)
-                if runtime.is_hip:
-                    y = y.astype(cupy.uint64, copy=False)
-                out = cub.cub_histogram(x, y, acc_bin_edge)
-                if out is None:
-                    # fallback to CuPy impl
-                    continue
-                else:
-                    y = out
-                if runtime.is_hip:
-                    y = y.astype(cupy.int64, copy=False)
-                break
-        else:
-            _histogram_kernel(x, bin_edges, bin_edges.size, y)
+        # TODO(leofang): we temporarily remove CUB histogram support for now,
+        # see cupy/cupy#7698. When it's ready, revert the commit that checked
+        # in this comment to restore the support.
+        _histogram_kernel(x, bin_edges, bin_edges.size, y)
     else:
         simple_weights = (
             cupy.can_cast(weights.dtype, cupy.float64) or
@@ -556,21 +519,10 @@ def bincount(x, weights=None, minlength=None):
 
     if weights is None:
         b = cupy.zeros((size,), dtype=numpy.intp)
-
-        for accelerator in _accelerator.get_routine_accelerators():
-            # CUB uses int for bin counts
-            # TODO(leofang): support >= 2^31 elements in x?
-            if (not runtime.is_hip
-                    and accelerator == _accelerator.ACCELERATOR_CUB
-                    and x.size <= 0x7fffffff and size <= 0x7fffffff):
-                out = cub.cub_histogram(x, b, size+1)
-                if out is None:
-                    continue
-                else:
-                    b = out
-                    break
-        else:
-            _bincount_kernel(x, b)
+        # TODO(leofang): we temporarily remove CUB histogram support for now,
+        # see cupy/cupy#7698. When it's ready, revert the commit that checked
+        # in this comment to restore the support.
+        _bincount_kernel(x, b)
     else:
         b = cupy.zeros((size,), dtype=numpy.float64)
         _bincount_with_weight_kernel(x, weights, b)


Test patch:
diff --git a/tests/cupy_tests/statistics_tests/test_histogram.py b/tests/cupy_tests/statistics_tests/test_histogram.py
--- a/tests/cupy_tests/statistics_tests/test_histogram.py
+++ b/tests/cupy_tests/statistics_tests/test_histogram.py
@@ -6,7 +6,6 @@
 
 import cupy
 from cupy import testing
-from cupy._core import _accelerator
 
 
 # Note that numpy.bincount does not support uint64 on 64-bit environment
@@ -328,83 +327,9 @@ def test_bincount_too_small_minlength(self, dtype):
                 xp.bincount(x, minlength=-1)
 
 
-# This class compares CUB results against NumPy's
-@unittest.skipUnless(cupy.cuda.cub.available, 'The CUB routine is not enabled')
-class TestCubHistogram(unittest.TestCase):
-
-    def setUp(self):
-        self.old_accelerators = _accelerator.get_routine_accelerators()
-        _accelerator.set_routine_accelerators(['cub'])
-
-    def tearDown(self):
-        _accelerator.set_routine_accelerators(self.old_accelerators)
-
-    @testing.for_all_dtypes(no_bool=True, no_complex=True)
-    @testing.numpy_cupy_array_equal()
-    def test_histogram(self, xp, dtype):
-        x = testing.shaped_arange((10,), xp, dtype)
-
-        if xp is numpy:
-            return xp.histogram(x)
-
-        # xp is cupy, first ensure we really use CUB
-        cub_func = 'cupy._statistics.histogram.cub.device_histogram'
-        with testing.AssertFunctionIsCalled(cub_func):
-            xp.histogram(x)
-        # ...then perform the actual computation
-        return xp.histogram(x)
-
-    @testing.for_all_dtypes(no_bool=True, no_complex=True)
-    @testing.numpy_cupy_array_equal()
-    def test_histogram_range_float(self, xp, dtype):
-        a = testing.shaped_arange((10,), xp, dtype)
-        h, b = xp.histogram(a, testing.shaped_arange((10,), xp, numpy.float64))
-        assert int(h.sum()) == 10
-        return h, b
-
-    @testing.for_all_dtypes_combination(['dtype_a', 'dtype_b'],
-                                        no_bool=True, no_complex=True)
-    @testing.numpy_cupy_array_equal()
-    def test_histogram_with_bins(self, xp, dtype_a, dtype_b):
-        x = testing.shaped_arange((10,), xp, dtype_a)
-        bins = testing.shaped_arange((4,), xp, dtype_b)
-
-        if xp is numpy:
-            return xp.histogram(x, bins)[0]
-
-        # xp is cupy, first ensure we really use CUB
-        cub_func = 'cupy._statistics.histogram.cub.device_histogram'
-        with testing.AssertFunctionIsCalled(cub_func):
-            xp.histogram(x, bins)
-        # ...then perform the actual computation
-        return xp.histogram(x, bins)[0]
-
-    @testing.for_all_dtypes_combination(['dtype_a', 'dtype_b'],
-                                        no_bool=True, no_complex=True)
-    @testing.numpy_cupy_array_equal()
-    def test_histogram_with_bins2(self, xp, dtype_a, dtype_b):
-        x = testing.shaped_arange((10,), xp, dtype_a)
-        bins = testing.shaped_arange((4,), xp, dtype_b)
-
-        if xp is numpy:
-            return xp.histogram(x, bins)[1]
-
-        # xp is cupy, first ensure we really use CUB
-        cub_func = 'cupy._statistics.histogram.cub.device_histogram'
-        with testing.AssertFunctionIsCalled(cub_func):
-            xp.histogram(x, bins)
-        # ...then perform the actual computation
-        return xp.histogram(x, bins)[1]
-
-    @testing.slow
-    @testing.numpy_cupy_array_equal()
-    def test_no_oom(self, xp):
-        # ensure the workaround for NVIDIA/cub#613 kicks in
-        amax = 28854312
-        A = xp.linspace(0, amax, num=amax,
-                        endpoint=True, retstep=False, dtype=xp.int32)
-        out = xp.histogram(A, bins=amax, range=[0, amax])
-        return out
+# TODO(leofang): we temporarily remove CUB histogram support for now,
+# see cupy/cupy#7698. When it's ready, revert the commit that checked
+# in this comment to restore the support.
 
 
 @testing.parameterize(*testing.product(