[FEA]: Possible improvement of deviceHistogram::HistogramEven

### Is this a duplicate?

- [x] I confirmed there appear to be no [duplicate issues](https://github.com/NVIDIA/cccl/issues) for this request and that I agree to the [Code of Conduct](CODE_OF_CONDUCT.md)

### Area

CUB

### Is your feature request related to a problem? Please describe.

For my application I had to create a custom histogram since I could not use CUB's, because I had to do other things apart from histogramming. Out of curiosity, I benchmarked it against the CUB implementation, and it seems to be considerably faster. Here's a plot of the results for an RTX 2080Ti:


My use case is creating the histogram of a text composed of an alphabet. For the tests I used uniformly randomly distributed texts, from an alphabet ´[0, alphabet_size)´. Also, for the alphabet size of 64'000 and the character 33'555, the CUB histogram entry is 0, which shouldn't be, and for the alphabet size of 100'000 and the character 0, the results are different. I've tested my implementation throughly with many random inputs, so am relatively sure the CUB result is wrong.

Here's the script I used for benchmarking:

```
#include

#include
#include
#include

typedef unsigned long long cu_size_t;

static uint32_t kMaxTPB = 0;
static cudaDeviceProp prop;

#define WS 32

#if defined(__CUDA_ARCH__)
#if __CUDA_ARCH__ > 800 && __CUDA_ARCH__ < 900
#define MAX_TPB 768
#define MIN_BPM 2
#elif __CUDA_ARCH__ == 750
#define MAX_TPB 1024
#define MIN_BPM 1
#else
#define MAX_TPB 1024
#define MIN_BPM 2
#endif

#define LB(x, y) __launch_bounds__(x, y)
#else
#define LB(x, y)
#endif

#define gpuErrchk(ans) \
{ \
gpuAssert((ans), __FILE__, __LINE__); \
}
__host__ __device__ inline void gpuAssert(cudaError_t code, const char* file,
int line, bool abort = true) {
if (code != cudaSuccess) {
printf("GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
if (abort) {
#ifdef __CUDA_ARCH__
asm("trap;");
#else
exit(EXIT_FAILURE);
#endif
}
}
}

__host__ cudaDeviceProp& getDeviceProperties() {
assert(prop.totalGlobalMem != 0);
return prop;
}

__host__ void checkWarpSize(uint32_t const GPU_index) {
if (prop.totalGlobalMem == 0) {
gpuErrchk(cudaSetDevice(GPU_index));
cudaGetDeviceProperties(&prop, GPU_index);
auto const threads_per_sm = prop.maxThreadsPerMultiProcessor;
kMaxTPB = prop.maxThreadsPerBlock;
// find max block size that can still fully load an SM
while (threads_per_sm % kMaxTPB != 0) {
kMaxTPB -= WS;
}
assert(kMaxTPB > WS);
}
if (prop.warpSize != WS) {
fprintf(stderr, "Warp size must be 32, but is %d\n", prop.warpSize);
exit(EXIT_FAILURE);
}
}

template
std::vector generateRandomData(std::vector const& alphabet,
size_t const data_size) {
std::vector data(data_size);
#pragma omp parallel
{
// Create a thread-local random number generator
std::random_device rd;
std::mt19937 gen(rd() + omp_get_thread_num()); // Add thread number to seed
// for better randomness
std::uniform_int_distribution dis(0, alphabet.size() - 1);

#pragma omp for
for (size_t i = 0; i < data_size; i++) {
data[i] = alphabet[dis(gen)];
}
}

return data;
}

template
__host__ T findLargestDivisor(T const n, T const divisor) {
if (divisor == 0) return 1;
return divisor - n % divisor;
}

template
__global__ LB(MAX_TPB, MIN_BPM) void computeGlobalHistogramKernel(
T* data, size_t const data_size, size_t* counts, T* const alphabet,
size_t const alphabet_size, uint16_t const hists_per_block) {
assert(blockDim.x % WS == 0);
extern __shared__ size_t shared_hist[];
size_t offset;
if constexpr (UseShmem) {
offset = (threadIdx.x % hists_per_block) * alphabet_size;
for (size_t i = threadIdx.x; i < alphabet_size * hists_per_block;
i += blockDim.x) {
shared_hist[i] = 0;
}
__syncthreads();
}

size_t const total_threads = blockDim.x * gridDim.x;
size_t const global_t_id = blockIdx.x * blockDim.x + threadIdx.x;
T char_data;
for (size_t i = global_t_id; i < data_size; i += total_threads) {
char_data = data[i];
if constexpr (UseShmem) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ <= 520
atomicAdd((cu_size_t*)&shared_hist[offset + char_data], size_t(1));
#else
atomicAdd_block((cu_size_t*)&shared_hist[offset + char_data], size_t(1));
#endif
} else {
atomicAdd((cu_size_t*)&counts[char_data], size_t(1));
}
}

if constexpr (UseShmem) {
__syncthreads();
// Reduce shared histograms to first one
for (size_t i = threadIdx.x; i < alphabet_size; i += blockDim.x) {
size_t sum = shared_hist[i];
for (size_t j = 1; j < hists_per_block; ++j) {
sum += shared_hist[j * alphabet_size + i];
}
atomicAdd((cu_size_t*)&counts[i], sum);
}
}
}

template
__host__ void computeGlobalHistogram(size_t const data_size, T* d_data,
T* d_alphabet, size_t* d_histogram,
size_t const alphabet_size) {
struct cudaFuncAttributes funcAttrib;
gpuErrchk(cudaFuncGetAttributes(&funcAttrib,
computeGlobalHistogramKernel));

struct cudaDeviceProp& prop = getDeviceProperties();

uint32_t max_TPB =
std::min(kMaxTPB, static_cast(funcAttrib.maxThreadsPerBlock));

max_TPB = findLargestDivisor(kMaxTPB, max_TPB);

auto const max_shmem_per_SM = prop.sharedMemPerMultiprocessor;
auto const max_threads_per_SM = prop.maxThreadsPerMultiProcessor;
size_t const hist_size = sizeof(size_t) * alphabet_size;

// Compute global_histogram and change text to min_alphabet
size_t const total_threads = std::min(
(data_size / WS) * WS,
static_cast(max_threads_per_SM * prop.multiProcessorCount));

auto const threads_per_block = max_TPB;
auto const num_blocks = total_threads / threads_per_block;

uint16_t const blocks_per_SM = max_threads_per_SM / threads_per_block;

size_t const used_shmem =
std::min(max_shmem_per_SM / blocks_per_SM, prop.sharedMemPerBlock);

uint16_t const hists_per_block =
std::min(static_cast(threads_per_block), used_shmem / hist_size);
if (hists_per_block > 0) {
computeGlobalHistogramKernel
<<>>(
d_data, data_size, d_histogram, d_alphabet, alphabet_size,
hists_per_block);
} else {
computeGlobalHistogramKernel<<>>(
d_data, data_size, d_histogram, d_alphabet, alphabet_size,
hists_per_block);
}
}

template
void BM_HistComputation(size_t const data_size, size_t const alphabet_size,
size_t const num_iters, std::string const& output) {
auto alphabet = std::vector(alphabet_size);
std::iota(alphabet.begin(), alphabet.end(), 0ULL);
auto data = generateRandomData(alphabet, data_size);

T* d_data;
T* d_alphabet;
size_t* d_histogram;
gpuErrchk(cudaMalloc(&d_data, data_size * sizeof(T)));
gpuErrchk(cudaMemcpy(d_data, data.data(), data_size * sizeof(T),
cudaMemcpyHostToDevice));
gpuErrchk(cudaMalloc(&d_alphabet, alphabet_size * sizeof(T)));
gpuErrchk(cudaMemcpy(d_alphabet, alphabet.data(), alphabet_size * sizeof(T),
cudaMemcpyHostToDevice));
gpuErrchk(cudaMalloc(&d_histogram, alphabet_size * sizeof(size_t)));

cudaEvent_t start, stop;
gpuErrchk(cudaEventCreate(&start));
gpuErrchk(cudaEventCreate(&stop));
std::vector times(num_iters);

void* d_temp_storage = nullptr;
size_t temp_storage_bytes = 0;
cub::DeviceHistogram::HistogramEven(
d_temp_storage, temp_storage_bytes, d_data, (cu_size_t*)d_histogram,
alphabet_size + 1, T(0), T(alphabet_size), data_size);
gpuErrchk(cudaMalloc(&d_temp_storage, temp_storage_bytes));

// Warmup
for (int i = 0; i < 5; ++i) {
gpuErrchk(cudaMemset(d_histogram, 0, alphabet_size * sizeof(size_t)));
cub::DeviceHistogram::HistogramEven(
d_temp_storage, temp_storage_bytes, d_data, (cu_size_t*)d_histogram,
alphabet_size + 1, T(0), T(alphabet_size), data_size);
}

float median_cub = 0;
std::vector cub_hist(alphabet_size);
for (size_t i = 0; i < num_iters; ++i) {
gpuErrchk(cudaMemset(d_histogram, 0, alphabet_size * sizeof(size_t)));
gpuErrchk(cudaEventRecord(start));
cub::DeviceHistogram::HistogramEven(
d_temp_storage, temp_storage_bytes, d_data, (cu_size_t*)d_histogram,
alphabet_size + 1, T(0), T(alphabet_size), data_size);
gpuErrchk(cudaEventRecord(stop));
gpuErrchk(cudaEventSynchronize(stop));
float milliseconds = 0;
gpuErrchk(cudaEventElapsedTime(&milliseconds, start, stop));
times[i] = milliseconds;
}
// Get median time
std::nth_element(times.begin(), times.begin() + times.size() / 2,
times.end());
median_cub = times[times.size() / 2];
gpuErrchk(cudaFree(d_temp_storage));
gpuErrchk(cudaMemcpy(cub_hist.data(), d_histogram,
alphabet_size * sizeof(size_t), cudaMemcpyDeviceToHost));

// Warmup
for (int i = 0; i < 5; ++i) {
gpuErrchk(cudaMemset(d_histogram, 0, alphabet_size * sizeof(size_t)));
computeGlobalHistogram(data_size, d_data, d_alphabet, d_histogram,
alphabet_size);
}

float median_custom = 0;
std::vector custom_hist(alphabet_size);
for (size_t i = 0; i < num_iters; ++i) {
gpuErrchk(cudaMemset(d_histogram, 0, alphabet_size * sizeof(size_t)));
gpuErrchk(cudaEventRecord(start));
computeGlobalHistogram(data_size, d_data, d_alphabet, d_histogram,
alphabet_size);
gpuErrchk(cudaEventRecord(stop));
gpuErrchk(cudaEventSynchronize(stop));
float milliseconds = 0;
gpuErrchk(cudaEventElapsedTime(&milliseconds, start, stop));
times[i] = milliseconds;
}
// Get median time
std::nth_element(times.begin(), times.begin() + times.size() / 2,
times.end());
median_custom = times[times.size() / 2];
std::ofstream out(output, std::ios::app);
out << data_size << "," << alphabet_size << "," << median_cub << ","
<< median_custom << std::endl;
out.close();

gpuErrchk(cudaMemcpy(custom_hist.data(), d_histogram,
alphabet_size * sizeof(size_t), cudaMemcpyDeviceToHost));
for (size_t i = 0; i < alphabet_size; ++i) {
if (cub_hist[i] != custom_hist[i]) {
std::cerr << "Mismatch at index " << i << " cub_hist: " << cub_hist[i]
<< " custom_hist: " << custom_hist[i] << std::endl;
break;
}
}
gpuErrchk(cudaFree(d_data));
gpuErrchk(cudaFree(d_alphabet));
gpuErrchk(cudaFree(d_histogram));
}

int main(int argc, char** argv) {
if (argc != 4) {
std::cerr << "Usage: " << argv[0] << " "
<< std::endl;
return EXIT_FAILURE;
}

uint32_t const GPU_index = std::stoi(argv[1]);
uint32_t const num_iters = std::stoi(argv[2]);
std::string const output = argv[3];

std::ofstream out(output);
out << "data_size,alphabet_size,median_cub(ms),median_custom(ms)"
<< std::endl;
out.close();

checkWarpSize(GPU_index);
std::vector const data_sizes = {500'000'000, 1'000'000'000,
2'000'000'000};
std::vector const alphabet_sizes = {
4, 8, 16, 32, 64, 128, 256, 512,
1024, 2048, 4096, 8192, 16384, 32768, 64000, 100000};

for (auto const data_size : data_sizes) {
for (auto const alphabet_size : alphabet_sizes) {
if (alphabet_size >= std::numeric_limits::max()) {
BM_HistComputation(data_size, alphabet_size, num_iters,
output);
} else if (alphabet_size >= std::numeric_limits::max()) {
BM_HistComputation(data_size, alphabet_size, num_iters,
output);
} else {
BM_HistComputation(data_size, alphabet_size, num_iters,
output);
}
}
}
return EXIT_SUCCESS;
}
```

I used `nvcc -O3 -Xcompiler -fopenmp -arch=sm_75 hist_comp.cu -o hist_comp` and ran it for 20 iterations, using version 12.0.

My implementation packs as many histograms as possible into shared memory, and assigns a local histogram to each thread in the block, in a round robin fashion 2 minimize how many threads within a warp share the same local histogram. If a histogram doesnt fit in shared memory, it just performs it atomically to global memory.

Here is also the CSV of the results, with more data sizes: [2080Ti.csv](https://github.com/user-attachments/files/19372362/2080Ti.csv).

Is your implementation optimized for other specific use cases? Just leaving this here in case you're interested in improving the performance for this use case.

### Describe the solution you'd like

Improve the performance of HistogramEven for creating histograms of texts.

### Describe alternatives you've considered

_No response_

### Additional context

_No response_