Implement long / long long shuffles (#1829)

Implement additional data-types for shuffles (long and long long).
Based upon the double implementation.
This commit is contained in:
Nick Curtis
2020-02-14 22:21:09 -06:00
committed by GitHub
parent 13052f27c2
commit 5b0f34c59b
2 changed files with 244 additions and 0 deletions
@@ -319,6 +319,36 @@ double __shfl(double var, int src_lane, int width = warpSize) {
double tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
long __shfl(long var, int src_lane, int width = warpSize)
{
static_assert(sizeof(long) == 2 * sizeof(int), "");
static_assert(sizeof(long) == sizeof(uint64_t), "");
int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
tmp[0] = __shfl(tmp[0], src_lane, width);
tmp[1] = __shfl(tmp[1], src_lane, width);
uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
long tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
long long __shfl(long long var, int src_lane, int width = warpSize)
{
static_assert(sizeof(long long) == 2 * sizeof(int), "");
static_assert(sizeof(long long) == sizeof(uint64_t), "");
int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
tmp[0] = __shfl(tmp[0], src_lane, width);
tmp[1] = __shfl(tmp[1], src_lane, width);
uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
long long tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
@@ -356,6 +386,34 @@ double __shfl_up(double var, unsigned int lane_delta, int width = warpSize) {
double tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
long __shfl_up(long var, unsigned int lane_delta, int width = warpSize)
{
static_assert(sizeof(long) == 2 * sizeof(int), "");
static_assert(sizeof(long) == sizeof(uint64_t), "");
int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
tmp[0] = __shfl_up(tmp[0], lane_delta, width);
tmp[1] = __shfl_up(tmp[1], lane_delta, width);
uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
long tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
long long __shfl_up(long long var, unsigned int lane_delta, int width = warpSize)
{
static_assert(sizeof(long long) == 2 * sizeof(int), "");
static_assert(sizeof(long long) == sizeof(uint64_t), "");
int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
tmp[0] = __shfl_up(tmp[0], lane_delta, width);
tmp[1] = __shfl_up(tmp[1], lane_delta, width);
uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
long long tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
@@ -393,6 +451,34 @@ double __shfl_down(double var, unsigned int lane_delta, int width = warpSize) {
double tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
long __shfl_down(long var, unsigned int lane_delta, int width = warpSize)
{
static_assert(sizeof(long) == 2 * sizeof(int), "");
static_assert(sizeof(long) == sizeof(uint64_t), "");
int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
tmp[0] = __shfl_down(tmp[0], lane_delta, width);
tmp[1] = __shfl_down(tmp[1], lane_delta, width);
uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
long tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
long long __shfl_down(long long var, unsigned int lane_delta, int width = warpSize)
{
static_assert(sizeof(long long) == 2 * sizeof(int), "");
static_assert(sizeof(long long) == sizeof(uint64_t), "");
int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
tmp[0] = __shfl_down(tmp[0], lane_delta, width);
tmp[1] = __shfl_down(tmp[1], lane_delta, width);
uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
long long tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
@@ -430,6 +516,34 @@ double __shfl_xor(double var, int lane_mask, int width = warpSize) {
double tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
long __shfl_xor(long var, int lane_mask, int width = warpSize)
{
static_assert(sizeof(long) == 2 * sizeof(int), "");
static_assert(sizeof(long) == sizeof(uint64_t), "");
int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
tmp[0] = __shfl_xor(tmp[0], lane_mask, width);
tmp[1] = __shfl_xor(tmp[1], lane_mask, width);
uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
long tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
__device__
inline
long long __shfl_xor(long long var, int lane_mask, int width = warpSize)
{
static_assert(sizeof(long long) == 2 * sizeof(int), "");
static_assert(sizeof(long long) == sizeof(uint64_t), "");
int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
tmp[0] = __shfl_xor(tmp[0], lane_mask, width);
tmp[1] = __shfl_xor(tmp[1], lane_mask, width);
uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
long long tmp1; __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
return tmp1;
}
#define MASK1 0x00ff00ff
#define MASK2 0xff00ff00