[Question]: optimization for block_radix_rank

[hlyix]

Is this a duplicate?

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

Area

Thrust

Is your feature request related to a problem? Please describe.

ccl/cub/cub/block/block_radix_rank.cuh

I assume that we can reduce some shared memory access in function ScanCounters.

Will it result in some unexpected behavior?

Describe the solution you'd like

change

    __device__ __forceinline__ void ScanCounters()
    {
        // Upsweep scan
        PackedCounter raking_partial = Upsweep();

        // Compute exclusive sum
        PackedCounter exclusive_partial;
        PrefixCallBack prefix_call_back;
        BlockScan(temp_storage.block_scan).ExclusiveSum(raking_partial, exclusive_partial, prefix_call_back);

        // Downsweep scan with exclusive partial
        ExclusiveDownsweep(exclusive_partial);
    }

to

    __device__ __forceinline__ void ScanCounters()
    {
        // Upsweep scan
        PackedCounter raking_partial = Upsweep();

        // Compute exclusive sum
        PackedCounter exclusive_partial;
        PrefixCallBack prefix_call_back;
        BlockScan(temp_storage.block_scan).ExclusiveSum(raking_partial, exclusive_partial, prefix_call_back);
        if (raking_partial  == 0) return; // if raking_partial is zero, it's unnecessary to update. 
        // Downsweep scan with exclusive partial
        ExclusiveDownsweep(exclusive_partial);
    }

Describe alternatives you've considered

No response

Additional context

No response

[elstehle]

Thank you for your suggestion, @hlyix!

While an early return might seem appealing for optimization, warp divergence becomes a concern in this context. If threads within the same warp take divergent execution paths (some returning early while others proceed), this serializes instruction execution across the warp, potentially negating any performance gains and even causing regressions.

More importantly, the proposed change would compromise correctness: The ExclusiveDownsweep operation isn't just a conditional scan, but it also integrates the exclusive_partial value into the thread's items.