[BUG]: Unique by key doesn't use allocated vsmem

[gevtushenko]

Is this a duplicate?

• I confirmed there appear to be no duplicate issues for this bug and that I agree to the Code of Conduct

Type of Bug

Runtime Error

Component

CUB

Describe the bug

Unique by key implementation pretends to take care of vsmem requesting temporary storage for it:

std::size_t vshmem_size = detail::VshmemSize(max_shmem, sizeof(typename UniqueByKeyAgentT::TempStorage), num_tiles);

// Specify temporary storage allocation requirements
size_t allocation_sizes[2] = {0, vshmem_size};
if (CubDebug(error = AliasTemporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) break;

But the allocated memory is never used. After progress on #548, this should be addressed as well.

How to Reproduce

Pass a large type that would lead to exceeding shared memory limitations.

Expected behavior

The algorithm should work with large types.

Reproduction link

No response

Operating System

No response

nvidia-smi output

No response

NVCC version

No response

[gevtushenko]

@elstehle this might be the next step in porting Thrust kernels to CUB. CUB already have unique by key but thrust doesn't use it. If we introduce your vsmem abstraction into unique by key, we should be able to remove thrust implementation. This would make recent performance improvements in CUB accessible to Thrust.