Can cub::DeviceSegmentedReduce::Reduce support self-defined functor for struct variable instead of just integer?

[zlwu92]

Hi,

I want to know if the cub::DeviceSegmentedReduce::Reduce support for struct variable, the source code example in comments I found is for integer.

Suppose we have two integer in the struct, one is id, one is value, firstly we compare value, choose the min, if two values are equal then we choose the minimal id, that's basically two levels of comparison in my self-defined functor.

Btw, if the warp level scan can pass a self-defined customized functor into it? And if this can support arbitrary length or just 32 which is equal to warpsize.

[gevtushenko]

Hello @zlwu92!

We support arbitrary operations in cub::DeviceSegmentedReduce:

#include <cub/cub.cuh>
#include <thrust/device_vector.h>

struct pair_t {
  int id;
  int val;
};

struct operation_t {
  __device__ pair_t operator()(const pair_t &a, const pair_t &b) {
    if (a.val == b.val) {
      return a.id < b.id ? a : b;
    }

    return a.val < b.val ? a : b;
  }
};

int main() {
  thrust::device_vector<pair_t> in = {
    pair_t{0, 1},
    pair_t{1, 2},
    pair_t{2, 2},
    pair_t{3, 2},
  };

  const int num_segments = 2;
  thrust::device_vector<pair_t> out(num_segments);

  thrust::device_vector<int> offsets = {
    0, 2, 4
  };

  const pair_t init{42, 42};

  std::size_t temp_storage_bytes = 0;
  std::uint8_t *d_temp_storage = nullptr;
  cub::DeviceSegmentedReduce::Reduce(
    d_temp_storage,
    temp_storage_bytes,
    in.begin(),
    out.begin(),
    num_segments,
    offsets.begin(),
    offsets.begin() + 1,
    operation_t(),
    init);
  
  thrust::device_vector<std::uint8_t> temp_storage(temp_storage_bytes);
  d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());

  cub::DeviceSegmentedReduce::Reduce(
    d_temp_storage,
    temp_storage_bytes,
    in.begin(),
    out.begin(),
    num_segments,
    offsets.begin(),
    offsets.begin() + 1,
    operation_t(),
    init);
  cudaDeviceSynchronize();

  for (pair_t pair: out) {
    std::cout << pair.id << " " << pair.val << std::endl;
  }

  // 0 1
  // 2 2
}

Btw, if the warp level scan can pass a self-defined customized functor into it?

Our cub::WarpScan also supports user-defined operators.

And if this can support arbitrary length or just 32 which is equal to warpsize.

Maximal size is currently equal to the virtual warp size (up to 32).

[zlwu92]

I see. Okay. Thank you very much! @senior-zero