Skip to content

cub::detail::WarpScanSmem could be optimized/cleaned up #5380

Description

@pauleonix

This is a rather low-priority issue given that this specialization of cub::WarpScan is only used for LOGICAL_WARP_THREADS is not a power of 2 which I assume is quite exotic. But while adding new member functions in #5379 I observed quite a few potential inefficiencies. Most of them come from the fact that when CUB was sprinkled with __syncwarp() for Volta/Independent Thread Scheduling in 83ab806, there was little regard for how else the code could be modified to fit the new paradigm.

  1. All loads and store from/to shared memory are done with ThreadLoad<LOAD_VOLATILE>/ThreadStore<STORE_VOLATILE> since long before the Volta changes. Doing volatile loads/stores should not be necessary in the presence of the explicit warp synchronization.
  2. Afaik there is at least one occurrence of redundant synchronization in line 417 here:
    ThreadStore<STORE_VOLATILE>(&temp_storage[HALF_WARP_THREADS + lane_id], (CellT) inclusive);
    __syncwarp(member_mask);
    warp_aggregate = (T) ThreadLoad<LOAD_VOLATILE>(&temp_storage[WARP_SMEM_ELEMENTS - 1]);
    __syncwarp(member_mask);
    // Update inclusive with initial value
    inclusive = scan_op(initial_value, inclusive);
    // Get exclusive from exclusive
    ThreadStore<STORE_VOLATILE>(&temp_storage[HALF_WARP_THREADS + lane_id - 1], (CellT) inclusive);

    The second store is explicitly done with an offset such that the aggregate is not overwritten.
  3. There is an explicit optimization for cuda::std::plus on primitive types based on knowing the identity value. That optimization could easily be done for other operators as well (cuda::std::multiplies, cuda::minimum, cuda::maximum, cuda::std::bit_or, cuda::std::bit_and, cuda::std::bit_xor) in
    // Update partial if addend is in range
    if (HAS_IDENTITY || (lane_id >= OFFSET))
    {
    T addend = (T) ThreadLoad<LOAD_VOLATILE>(&temp_storage[HALF_WARP_THREADS + lane_id - OFFSET]);
    partial = scan_op(addend, partial);
    }
    HAS_IDENTITY could be true for all of these operators. See also identity_v from CUB - Add internal thread and warp utils (Split WarpReduce PR) #5317
  4. Another optimization currently only done for cuda::std::plus on integers is calculating the exclusive scan from the inclusive scan as exclusive = inclusive - input instead of communicating through shared memory. The same could be done for cuda::std::bit_xor and in theory cuda::std::multiplies although integer division might not actually be faster than going through shared memory.
  5. Remove CellT and all the C-style casts associated with it.
    /// Storage cell type (workaround for SM1x compiler bugs with custom-ops like Max() on signed chars)
    using CellT = T;
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    • Status
      Todo

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions