Safe bet for temp memory requirements - device wide ReduceByKey/Radix SortPairs

[stndbye]

Hello,

I'd gladly sacrifice some memory to be able to avoid mem operations as much as possible.
Is there a safe overestimate for reduce and radix sort temp memory requirements?
I'd allocate a temp buffer once and reuse it for multiple CUB calls if possible.

[pauleonix]

You could use a memory resource that is able to reuse allocations. See thrust::mr::disjoint_unsynchronized_pool_resource. CUDA actually has something like this built-in at the runtime level. See Using the NVIDIA CUDA Stream-Ordered Memory Allocator, Part 1. The most complete C++ API for this (i.e. including streams) is probably in RMM, but this adds a dependency.

[stndbye]

Thanks, sounds nice! Maybe i can call my kernels in an order that hides the memory operations.

Can i expect the temp buffer requirement to be lower or equal to the the problem size in worst case scenario?

[pauleonix]

@stndbye
I'm not a CUB developer, but I think these buffer size calculations are handled as implementation details. They might still be able to give you some formulas, but you probably should still write your code in a way that still works if those heuristics fail (i.e. due to a future algorithmic change).

I think I have read issues before where people asked which parameters need to be available to calculate the buffer size, i.e. you might get away with passing nullptr instead of some actual parameters. Depending on the circumstances this could be used to query the max. temp buffer size of all used algorithms before allocating the first buffer.

[pauleonix]

See #847

[elstehle]

As @pauleonix correctly stated, temporary storage size calculations are handled as implementation details. Usually, it's a function of various types and the number of items. If you're planning to invoke the same algorithm for the same types repeatedly but for varying number of items, you can query temporary storage requirements for the maximum number of items that you are going to invoke that algorithm for and take that as an upper bound in the current implementation.