Use CUB's ScanByKey implementation

[alliepiper]

I will update thrust_benchmark to test this for perf regressions before merging. (Hence the blocked tag)

cc: @zasdfgbnm for awareness (don't want to duplicate effort)

run tests

[alliepiper]

run tests

[alliepiper]

DVS CL: 30580794

[zasdfgbnm]

There is one thing that might impact the performance on thrust and libcxx pointers. In cub, https://github.com/NVIDIA/cub/blob/main/cub/agent/agent_scan_by_key.cuh#L124-L129

    using WrappedKeysInputIteratorT = typename If<IsPointer<KeysInputIteratorT>::VALUE,
        CacheModifiedInputIterator<AgentScanByKeyPolicyT::LOAD_MODIFIER, KeyT, OffsetT>,   // Wrap the native input pointer with CacheModifiedInputIterator
        KeysInputIteratorT>::Type;
    using WrappedValuesInputIteratorT = typename If<IsPointer<ValuesInputIteratorT>::VALUE,
        CacheModifiedInputIterator<AgentScanByKeyPolicyT::LOAD_MODIFIER, InputT, OffsetT>,   // Wrap the native input pointer with CacheModifiedInputIterator
        ValuesInputIteratorT>::Type;

These two lines are copied from agent_scan.cuh. But for the implementation of scan-by-key in thrust, the condition for using CacheModifiedInputIterator is is_contiguous_iterator, which is

template <typename Iterator>
struct is_contiguous_iterator_impl
  : integral_constant<
      bool
    ,    is_pointer<Iterator>::value
      || is_thrust_pointer<Iterator>::value
      || is_libcxx_wrap_iter<Iterator>::value
      || is_libstdcxx_normal_iterator<Iterator>::value
      || is_msvc_contiguous_iterator<Iterator>::value
      || proclaim_contiguous_iterator<Iterator>::value
    >
{};

[gevtushenko]

It might be worth addressing @zasdfgbnm comment by implementing a make_input_iterator version without caching policy:

  template <class It>
  auto __device__ __forceinline__
  make_load_iterator_impl(It it, thrust::detail::true_type /* is_trivial */)
  {
    return raw_pointer_cast(&*it);
  }

  template <class It>
  It __device__ __forceinline__
  make_load_iterator_impl(It it, thrust::detail::false_type /* is_trivial */)
  {
    return it;
  }

  template <class It>
  typename LoadIterator<It>::type __device__ __forceinline__
  make_load_iterator(It it)
  {
    return make_load_iterator_impl(
        it, typename is_contiguous_iterator<It>::type());
  }

This version would be used before dispatching into CUB code. I think will need this kind of facility later in any case.

[alliepiper]

It might be worth addressing @zasdfgbnm comment by implementing a make_input_iterator version without caching policy:
[snip]
This version would be used before dispatching into CUB code. I think will need this kind of facility later in any case.

Agreed -- I'll add something like this to the contiguous iterator implementation.

[alliepiper]

run tests

[alliepiper]

• Added thrust::detail::try_unwrap_contiguous_iterator(Iterator it) that returns a raw pointer for contiguous iterators and passes-through other iterators.
• Use above in new scan_by_key implementation.
• Restored the short-circuit num_items==0 that avoids temp mem allocs.

Still need to write/verify benchmarks, and if we're happy with the new contiguous iterator unwrapping API I'll update the other CUB-backed algorithms in a followup.

[alliepiper]

run tests

[alliepiper]

run tests

[alliepiper]

run tests

[alliepiper]

run tests

[alliepiper]

run tests

[alliepiper]

Split the scan_by_key test into separate exclusive and inclusive variants to reduce memory usage during compilation.

run tests

[alliepiper]

run tests

[gevtushenko]

A few comments, but nothing critical. Thank you for separating the test!

[gevtushenko]

Is there a reason to test for check_unwrapped_iterator<T, T, true> + check_unwrapped_iterator<T, T, false> that I'm missing?

[alliepiper]

Since the unwrapped raw pointer is the same type as the pass-through, both the true and false cases should pass here. Testing both just ensures that everything is working as expected.

[alliepiper]

Since the unwrapped raw pointer is the same type as the pass-through, both the true and false cases should pass here. Testing both just ensures that everything is working as expected.

[gevtushenko]

Optional: this function is a duplicate on inclusive version. We might always have InitValueT and pass cub::NullType{} to indicate exclusiveness. The error messages will get back to scan_by_key, but it was like that before anyway.

[alliepiper]

Good catch, I'll clean this up to reduce duplication.

[alliepiper]

On second thought, let's keep these separate -- otherwise we'll just have to re-split them when we address NVIDIA/cub#384 and add InitialValue support to inclusive scans.

[alliepiper]

run tests