Support reduce and scan for more than 2^31 items

[milthorpe]

cub::DeviceReduce and cub::DeviceScan specify the input size as int num_items, which limits reductions/scans to 2^31-1 items. Other CUB functions have the type of the number of items as a template parameter e.g. cub::DeviceRadixSort accepts NumItemsT num_items.

The num_items parameter should be replaced with a templated parameter NumItemsT num_items for both cub::DeviceReduce and cub::DeviceScan, to allow reductions and scans of more than 2^31 items.

A simple test code for DeviceReduce with size_t num_items would be as follows:

#include <cub/cub.cuh>

using namespace cub;

template <typename T> void cubSum(const T *d_in, T *d_out, size_t num_items) {
  void *d_temp_storage = NULL;
  size_t temp_storage_bytes = 0;
  CubDebugExit(DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items));
  cudaMalloc(&d_temp_storage, temp_storage_bytes);

  CubDebugExit(DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items));

  cudaFree(d_temp_storage);
}

int main(int argc, char* argv[]) {
    size_t num_items = 2l<<30;
    double *h_in = new double[num_items];
    for (size_t i=0; i<num_items; i++) h_in[i] = i;
    double result = 0;

    CachingDeviceAllocator  g_allocator(true);

    double *d_in = NULL;
    CubDebugExit(g_allocator.DeviceAllocate((void**)&d_in, sizeof(double) * num_items));

    // Initialize device input
    CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(double) * num_items, cudaMemcpyHostToDevice));

    double *d_out = NULL;
    CubDebugExit(g_allocator.DeviceAllocate((void**)&d_out, sizeof(double)));

    cubSum(d_in, d_out, num_items);

    CubDebugExit(cudaMemcpy(&result, d_out, sizeof(double), cudaMemcpyDeviceToHost));
    printf("result: %e\n", result);

    if (h_in) delete[] h_in;
    if (d_in) CubDebugExit(g_allocator.DeviceFree(d_in));
    if (d_out) CubDebugExit(g_allocator.DeviceFree(d_out));
}

[milthorpe]

I'd be happy to work on a patch for this if the CUB developers think the proposed change would be acceptable.

[gevtushenko]

Hello, @milthorpe!

Thank you for offering help!
I've just filed a PR that illustrates what we expect from patches that introduce > 2^31 items support into our device algorithms. If you'd like, you can start looking into scan as soon as the PR is merged. Alternatively, you can wait till we prioritize this work, it's definitely on our list.

[wmaxey]

Does anyone have plans on tackling this change?

[gevtushenko]

@wmaxey unless @milthorpe would like to contribute 64-bit support for scan, we'll create a separate issue for that.

[milthorpe]

I don't actually have a current use-case for scan (only reduce, which @senior-zero kindly fixed already), which is why I've been so slow to respond - sorry! I think it is better to create a separate issue.

[gevtushenko]

@milthorpe I'm closing this issue then