From 54ecee04f4e6248ad4cd0cf778e4e1d6c06b27ea Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Thu, 11 Aug 2022 11:51:26 -0700 Subject: [PATCH 1/4] src/cunumeric: handle high number of bins in GPU bincount The existing bincount implementation on GPUs attempts to allocate a workspace for all bins within the shared memory available on each SM. This commit updates the implementation to fall back to a slower kernel that reduces to global memory when there are too many bins to fit into shared memory. Fixes #503. --- src/cunumeric/stat/bincount.cu | 133 ++++++++++++++++++++++++----- tests/integration/test_bincount.py | 26 ++++++ 2 files changed, 139 insertions(+), 20 deletions(-) diff --git a/src/cunumeric/stat/bincount.cu b/src/cunumeric/stat/bincount.cu index ea52a679cc..b4133409f5 100644 --- a/src/cunumeric/stat/bincount.cu +++ b/src/cunumeric/stat/bincount.cu @@ -98,6 +98,21 @@ 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, + const size_t num_bins, + 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) bincount_kernel_rw(AccessorRW lhs, @@ -116,6 +131,21 @@ static __global__ void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM) } } +template +static __global__ void + bincount_kernel_rw_global(AccessorRW lhs, + AccessorRO rhs, + const size_t volume, + const size_t num_bins, + 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]]; + SumReduction::fold(lhs[bin], 1); +} + template static __global__ void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM) weighted_bincount_kernel_rd(AccessorRD, false, 1> lhs, @@ -135,6 +165,22 @@ 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, + const size_t num_bins, + 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 static __global__ void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM) weighted_bincount_kernel_rw(AccessorRW lhs, @@ -154,6 +200,22 @@ static __global__ void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM) } } +template +static __global__ void + weighted_bincount_kernel_rw_global(AccessorRW lhs, + AccessorRO rhs, + AccessorRO weights, + const size_t volume, + const size_t num_bins, + 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]]; + SumReduction::fold(lhs[bin], weights[idx + origin[0]]); +} + template struct BincountImplBody { using VAL = legate_type_of; @@ -166,15 +228,23 @@ 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, num_bins, rect.lo); + } CHECK_CUDA_STREAM(stream); } @@ -186,15 +256,22 @@ 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_rw, THREADS_PER_BLOCK, bin_size); - assert(num_ctas > 0); - // Launch a kernel with this number of CTAs - auto stream = get_cached_stream(); - bincount_kernel_rw - <<>>(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) { + bincount_kernel_rw + <<>>(lhs, rhs, volume, num_bins, rect.lo); + } else { + auto blocks = (volume + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + bincount_kernel_rw_global<<>>(lhs, rhs, volume, num_bins, rect.lo); + } CHECK_CUDA_STREAM(stream); } @@ -207,15 +284,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, num_bins, rect.lo); + } CHECK_CUDA_STREAM(stream); } @@ -228,15 +313,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_rw, 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_rw<<>>( - 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_rw<<>>( + lhs, rhs, weights, volume, num_bins, rect.lo); + } else { + auto blocks = (volume + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + weighted_bincount_kernel_rw_global<<>>( + lhs, rhs, weights, volume, num_bins, rect.lo); + } CHECK_CUDA_STREAM(stream); } }; diff --git a/tests/integration/test_bincount.py b/tests/integration/test_bincount.py index 4c3d9e1910..bde9c42846 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] @@ -50,6 +52,30 @@ def test_bincount_weights(dtype): 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) + + v_np = v_num.__array__() + + out_np = np.bincount(v_np) + out_num = num.bincount(v_num) + 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) + + v_np = v_num.__array__() + w_np = w_num.__array__() + + out_np = np.bincount(v_np, weights=w_np) + out_num = num.bincount(v_num, weights=w_num) + assert allclose(out_np, out_num) + + if __name__ == "__main__": import sys From 1606af10c93907a7940a1d175df399327e0b05c8 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Thu, 11 Aug 2022 13:47:11 -0700 Subject: [PATCH 2/4] src/cunumeric/stat: fix launch parameters in bincount kernels --- src/cunumeric/stat/bincount.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/cunumeric/stat/bincount.cu b/src/cunumeric/stat/bincount.cu index b4133409f5..b4ce4129d8 100644 --- a/src/cunumeric/stat/bincount.cu +++ b/src/cunumeric/stat/bincount.cu @@ -298,7 +298,7 @@ struct BincountImplBody { lhs, rhs, weights, volume, num_bins, rect.lo); } else { auto blocks = (volume + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; - weighted_bincount_kernel_rd_global<<>>( + weighted_bincount_kernel_rd_global<<>>( lhs, rhs, weights, volume, num_bins, rect.lo); } CHECK_CUDA_STREAM(stream); @@ -327,7 +327,7 @@ struct BincountImplBody { lhs, rhs, weights, volume, num_bins, rect.lo); } else { auto blocks = (volume + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; - weighted_bincount_kernel_rw_global<<>>( + weighted_bincount_kernel_rw_global<<>>( lhs, rhs, weights, volume, num_bins, rect.lo); } CHECK_CUDA_STREAM(stream); From ac88141d0b0a154c5d9fb2ddf39dcf2de05f32c5 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Thu, 11 Aug 2022 13:48:42 -0700 Subject: [PATCH 3/4] tests/integration: refactor test_bincount.py for better readability --- tests/integration/test_bincount.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/tests/integration/test_bincount.py b/tests/integration/test_bincount.py index bde9c42846..b1d9fd4c5b 100644 --- a/tests/integration/test_bincount.py +++ b/tests/integration/test_bincount.py @@ -31,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) @@ -43,23 +42,23 @@ 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) - out_num = num.bincount(v_num, weights=w_num) + 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) - out_num = num.bincount(v_num) + assert num.array_equal(out_np, out_num) @@ -67,12 +66,12 @@ def test_bincount_high_bins(dtype): 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) - out_num = num.bincount(v_num, weights=w_num) + assert allclose(out_np, out_num) From d40b3e190268eed6dcd176ad57c51a5a1a006220 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Aug 2022 18:54:40 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/cunumeric/stat/bincount.cu | 71 +++++++++++++++++----------------- 1 file changed, 35 insertions(+), 36 deletions(-) diff --git a/src/cunumeric/stat/bincount.cu b/src/cunumeric/stat/bincount.cu index b4ce4129d8..f4596b0b15 100644 --- a/src/cunumeric/stat/bincount.cu +++ b/src/cunumeric/stat/bincount.cu @@ -99,12 +99,11 @@ 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, - const size_t num_bins, - Point<1> origin) +static __global__ void bincount_kernel_rd_global(AccessorRD, false, 1> lhs, + AccessorRO rhs, + const size_t volume, + const size_t num_bins, + Point<1> origin) { // Just blast out the atomic writes into global memory. auto idx = global_tid_1d(); @@ -132,12 +131,11 @@ static __global__ void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM) } template -static __global__ void - bincount_kernel_rw_global(AccessorRW lhs, - AccessorRO rhs, - const size_t volume, - const size_t num_bins, - Point<1> origin) +static __global__ void bincount_kernel_rw_global(AccessorRW lhs, + AccessorRO rhs, + const size_t volume, + const size_t num_bins, + Point<1> origin) { // Just blast out the atomic writes into global memory. auto idx = global_tid_1d(); @@ -166,13 +164,13 @@ 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, - const size_t num_bins, - Point<1> origin) +static __global__ void weighted_bincount_kernel_rd_global( + AccessorRD, false, 1> lhs, + AccessorRO rhs, + AccessorRO weights, + const size_t volume, + const size_t num_bins, + Point<1> origin) { // Just blast out the atomic writes into global memory. auto idx = global_tid_1d(); @@ -201,13 +199,12 @@ static __global__ void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM) } template -static __global__ void - weighted_bincount_kernel_rw_global(AccessorRW lhs, - AccessorRO rhs, - AccessorRO weights, - const size_t volume, - const size_t num_bins, - Point<1> origin) +static __global__ void weighted_bincount_kernel_rw_global(AccessorRW lhs, + AccessorRO rhs, + AccessorRO weights, + const size_t volume, + const size_t num_bins, + Point<1> origin) { // Just blast out the atomic writes into global memory. auto idx = global_tid_1d(); @@ -228,7 +225,7 @@ 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(); + auto stream = get_cached_stream(); int32_t num_ctas = 0; cudaOccupancyMaxActiveBlocksPerMultiprocessor( @@ -243,7 +240,8 @@ struct BincountImplBody { <<>>(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, num_bins, rect.lo); + bincount_kernel_rd_global + <<>>(lhs, rhs, volume, num_bins, rect.lo); } CHECK_CUDA_STREAM(stream); } @@ -256,7 +254,7 @@ 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(); + auto stream = get_cached_stream(); int32_t num_ctas = 0; cudaOccupancyMaxActiveBlocksPerMultiprocessor( @@ -270,7 +268,8 @@ struct BincountImplBody { <<>>(lhs, rhs, volume, num_bins, rect.lo); } else { auto blocks = (volume + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; - bincount_kernel_rw_global<<>>(lhs, rhs, volume, num_bins, rect.lo); + bincount_kernel_rw_global + <<>>(lhs, rhs, volume, num_bins, rect.lo); } CHECK_CUDA_STREAM(stream); } @@ -284,7 +283,7 @@ 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(); + auto stream = get_cached_stream(); int32_t num_ctas = 0; cudaOccupancyMaxActiveBlocksPerMultiprocessor( @@ -298,8 +297,8 @@ struct BincountImplBody { 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, num_bins, rect.lo); + weighted_bincount_kernel_rd_global + <<>>(lhs, rhs, weights, volume, num_bins, rect.lo); } CHECK_CUDA_STREAM(stream); } @@ -313,7 +312,7 @@ 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(); + auto stream = get_cached_stream(); int32_t num_ctas = 0; cudaOccupancyMaxActiveBlocksPerMultiprocessor( @@ -327,8 +326,8 @@ struct BincountImplBody { lhs, rhs, weights, volume, num_bins, rect.lo); } else { auto blocks = (volume + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; - weighted_bincount_kernel_rw_global<<>>( - lhs, rhs, weights, volume, num_bins, rect.lo); + weighted_bincount_kernel_rw_global + <<>>(lhs, rhs, weights, volume, num_bins, rect.lo); } CHECK_CUDA_STREAM(stream); }