WarpScan ExclusiveSum does not compile with vector types

[kaatish]

Trying to compile WarpScan::ExclusiveSum with a vector type like short2 results in a compilation error.

Code to reproduce error :

template <int LOGICAL_WARP_THREADS,
          class T,
          class ActionT>
__global__ void warp_scan_kernel(T *in, T *out, ActionT action)
{
  using warp_scan_t      = cub::WarpScan<T, LOGICAL_WARP_THREADS>;
  using storage_t         = typename warp_scan_t::TempStorage;

  __shared__ storage_t storage;

  const int tid = static_cast<int>(cub::RowMajorTid(blockDim.x, blockDim.y, blockDim.z));

  T thread_data = in[tid];
  warp_scan_t scan(storage);
  action(scan, thread_data);
  out[tid] = thread_data;
}

template <int LOGICAL_WARP_THREADS,
          class T,
          class ActionT>
void warp_scan(thrust::device_vector<T> &in, thrust::device_vector<T> &out, ActionT action)
{
  warp_scan_kernel<LOGICAL_WARP_THREADS, T, ActionT>
    <<<1, LOGICAL_WARP_THREADS>>>(thrust::raw_pointer_cast(in.data()),
                                              thrust::raw_pointer_cast(out.data()),
                                              action);
}

struct sum_op_t
{
  template <class WarpScanT, class T>
  __device__ void operator()(WarpScanT &scan, T &thread_data) const
  {
    scan.ExclusiveSum(thread_data, thread_data);
  }
};

static __host__ __device__ short2 operator+(const short2& lhs, const short2& rhs) {
  short2 result;
  result.x = lhs.x + rhs.x;
  result.y = lhs.y + rhs.y;
  return result;
}

int main (int argc, char* argv[]) {
  constexpr int len = 10;
  using type = short2;
  thrust::device_vector<type> in(len);
  thrust::device_vector<type> out(len);
  warp_scan<16>(in, out, sum_op_t{});
  return 0;
}

Error output :

$ nvcc --std=c++17 warp_scan.cu -o warp_scan
/usr/local/cuda-11.5/bin/../targets/x86_64-linux/include/cub/block/specializations/../../block/../block/specializations/../../warp/wa
rp_scan.cuh(348): error: no suitable constructor exists to convert from "int" to "short2"
          detected during:
            instantiation of "void cub::WarpScan<T, LOGICAL_WARP_THREADS, PTX_ARCH>::ExclusiveSum(T, T &) [with T=short2, LOGICAL_WAR
P_THREADS=16, PTX_ARCH=520]"
warp_scan.cu(40): here
            instantiation of "void sum_op_t::operator()(WarpScanT &, T &) const [with WarpScanT=cub::WarpScan<short2, 16, 520>, T=sho
rt2]"
warp_scan.cu(19): here
            instantiation of "void warp_scan_kernel<LOGICAL_WARP_THREADS,T,ActionT>(T *, T *, ActionT) [with LOGICAL_WARP_THREADS=16,
 T=short2, ActionT=sum_op_t]"
warp_scan.cu(31): here
            instantiation of "void warp_scan<LOGICAL_WARP_THREADS,T,ActionT>(thrust::device_vector<T, thrust::device_allocator<T>> &,
 thrust::device_vector<T, thrust::device_allocator<T>> &, ActionT) [with LOGICAL_WARP_THREADS=16, T=short2, ActionT=sum_op_t]"
warp_scan.cu(56): here