warp.h

Go to the documentation of this file.

#pragma once
 
#include "backend/api.h"
 
#include <cstdint>
 
namespace fz {
namespace backend {
 
#if defined(FZGMOD_BACKEND_HIP)
using warp_mask_t = unsigned long long;
inline constexpr warp_mask_t kFullMask = ~0ULL;
#else
using warp_mask_t = unsigned;
inline constexpr warp_mask_t kFullMask = 0xffffffffu;
#endif
 
template <typename T>
__device__ inline T shflUp(T val, unsigned delta, int width) {
    return __shfl_up_sync(kFullMask, val, delta, width);
}
 
template <typename T>
__device__ inline T shflDown(T val, unsigned delta, int width) {
    return __shfl_down_sync(kFullMask, val, delta, width);
}
 
template <typename T>
__device__ inline T shflXor(T val, int laneMask, int width) {
    return __shfl_xor_sync(kFullMask, val, laneMask, width);
}
 
template <typename T>
__device__ inline T shfl(T val, int srcLane, int width) {
    return __shfl_sync(kFullMask, val, srcLane, width);
}
 
__device__ inline uint32_t ballotSync32(int pred) {
#if defined(FZGMOD_BACKEND_HIP)
    bool upper = (__lane_id() & 32u) != 0u;
    warp_mask_t mask = upper ? 0xFFFFFFFF00000000ULL : 0x00000000FFFFFFFFULL;
    unsigned long long b = __ballot_sync(mask, pred);
    return static_cast<uint32_t>(upper ? (b >> 32) : b);
#else
    return __ballot_sync(kFullMask, pred);
#endif
}
 
__device__ inline bool anySync32(int pred) {
#if defined(FZGMOD_BACKEND_HIP)
    bool upper = (__lane_id() & 32u) != 0u;
    warp_mask_t mask = upper ? 0xFFFFFFFF00000000ULL : 0x00000000FFFFFFFFULL;
    return __any_sync(mask, pred) != 0;
#else
    return __any_sync(kFullMask, pred) != 0;
#endif
}
 
} // namespace backend
} // namespace fz