diff --git a/src/cunumeric/stat/bincount.cu b/src/cunumeric/stat/bincount.cu index 2e56d1d690..921288fb16 100644 --- a/src/cunumeric/stat/bincount.cu +++ b/src/cunumeric/stat/bincount.cu @@ -98,6 +98,19 @@ static __global__ void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM) } } +template +static __global__ void bincount_kernel_rd_global(AccessorRD, false, 1> lhs, + AccessorRO rhs, + const size_t volume, + Point<1> origin) +{ + // Just blast out the atomic writes into global memory. + auto idx = global_tid_1d(); + if (idx >= volume) return; + auto bin = rhs[idx + origin[0]]; + lhs[bin] <<= 1; +} + template static __global__ void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM) weighted_bincount_kernel_rd(AccessorRD, false, 1> lhs, @@ -117,6 +130,21 @@ static __global__ void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM) } } +template +static __global__ void weighted_bincount_kernel_rd_global( + AccessorRD, false, 1> lhs, + AccessorRO rhs, + AccessorRO weights, + const size_t volume, + Point<1> origin) +{ + // Just blast out the atomic writes into global memory. + auto idx = global_tid_1d(); + if (idx >= volume) return; + auto bin = rhs[idx + origin[0]]; + lhs[bin] <<= weights[idx + origin[0]]; +} + template struct BincountImplBody { using VAL = legate_type_of; @@ -129,15 +157,24 @@ struct BincountImplBody { const auto volume = rect.volume(); const auto num_bins = lhs_rect.volume(); const auto bin_size = num_bins * sizeof(int32_t); + auto stream = get_cached_stream(); int32_t num_ctas = 0; cudaOccupancyMaxActiveBlocksPerMultiprocessor( &num_ctas, bincount_kernel_rd, THREADS_PER_BLOCK, bin_size); - assert(num_ctas > 0); - // Launch a kernel with this number of CTAs - auto stream = get_cached_stream(); - bincount_kernel_rd - <<>>(lhs, rhs, volume, num_bins, rect.lo); + // If the number of bins is relatively low, attempt to use an algorithm that + // buffers bincounts local to each SM in shared memory. If there are too many + // bins to fit in shared memory, fall back to an approach that just blasts + // updates out to global memory. + if (num_ctas > 0) { + // Launch a kernel with this number of CTAs + bincount_kernel_rd + <<>>(lhs, rhs, volume, num_bins, rect.lo); + } else { + auto blocks = (volume + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + bincount_kernel_rd_global + <<>>(lhs, rhs, volume, rect.lo); + } CHECK_CUDA_STREAM(stream); } @@ -150,15 +187,23 @@ struct BincountImplBody { const auto volume = rect.volume(); const auto num_bins = lhs_rect.volume(); const auto bin_size = num_bins * sizeof(double); + auto stream = get_cached_stream(); int32_t num_ctas = 0; cudaOccupancyMaxActiveBlocksPerMultiprocessor( &num_ctas, weighted_bincount_kernel_rd, THREADS_PER_BLOCK, bin_size); - assert(num_ctas > 0); - // Launch a kernel with this number of CTAs - auto stream = get_cached_stream(); - weighted_bincount_kernel_rd<<>>( - lhs, rhs, weights, volume, num_bins, rect.lo); + // If the number of bins is relatively low, attempt to use an algorithm that + // buffers bincounts local to each SM in shared memory. If there are too many + // bins to fit in shared memory, fall back to an approach that just blasts + // updates out to global memory. + if (num_ctas > 0) { + weighted_bincount_kernel_rd<<>>( + lhs, rhs, weights, volume, num_bins, rect.lo); + } else { + auto blocks = (volume + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + weighted_bincount_kernel_rd_global + <<>>(lhs, rhs, weights, volume, rect.lo); + } CHECK_CUDA_STREAM(stream); } }; diff --git a/tests/integration/test_bincount.py b/tests/integration/test_bincount.py index 4c3d9e1910..b1d9fd4c5b 100644 --- a/tests/integration/test_bincount.py +++ b/tests/integration/test_bincount.py @@ -21,6 +21,8 @@ N = 8000 MAX_VAL = 9 +LARGE_NUM_BINS = 20000 + DTYPES = [np.int64, np.int32, np.int16] MINLENGTHS = [0, 5, 15] @@ -29,11 +31,10 @@ @pytest.mark.parametrize("minlength", MINLENGTHS) def test_bincount_basic(dtype, minlength): v_num = num.random.randint(0, MAX_VAL, size=N, dtype=dtype) + out_num = num.bincount(v_num, minlength=minlength) v_np = v_num.__array__() - out_np = np.bincount(v_np, minlength=minlength) - out_num = num.bincount(v_num, minlength=minlength) assert num.array_equal(out_np, out_num) @@ -41,12 +42,36 @@ def test_bincount_basic(dtype, minlength): def test_bincount_weights(dtype): v_num = num.random.randint(0, MAX_VAL, size=N, dtype=dtype) w_num = num.random.randn(N) + out_num = num.bincount(v_num, weights=w_num) v_np = v_num.__array__() w_np = w_num.__array__() - out_np = np.bincount(v_np, weights=w_np) + + assert allclose(out_np, out_num) + + +@pytest.mark.parametrize("dtype", DTYPES) +def test_bincount_high_bins(dtype): + v_num = num.array([0, LARGE_NUM_BINS], dtype=dtype) + out_num = num.bincount(v_num) + + v_np = v_num.__array__() + out_np = np.bincount(v_np) + + assert num.array_equal(out_np, out_num) + + +@pytest.mark.parametrize("dtype", DTYPES) +def test_bincount_weights_high_bins(dtype): + v_num = num.array([0, LARGE_NUM_BINS], dtype=dtype) + w_num = num.random.randn(2) out_num = num.bincount(v_num, weights=w_num) + + v_np = v_num.__array__() + w_np = w_num.__array__() + out_np = np.bincount(v_np, weights=w_np) + assert allclose(out_np, out_num)