File size: 2,585 Bytes
9dd3461 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
#pragma once
#include <ATen/cuda/Atomic.cuh>
namespace at {
namespace native {
__device__ __forceinline__ size_t
idx(const size_t nc,
const size_t height,
const size_t width,
const size_t h,
const size_t w) {
return (nc * height + h) * width + w;
}
// for channels-last
__device__ __forceinline__ size_t
idx_cl(
const size_t n, const size_t h, const size_t w, const size_t c,
const size_t height, const size_t width, const size_t channel
) {
return ((n * height + h) * width + w) * channel + c;
}
template <
typename scalar_t,
typename index_t,
typename std::enable_if<std::is_same<c10::Half, scalar_t>::value>::type* =
nullptr>
__device__ __forceinline__ void fastSpecializedAtomicAdd(
scalar_t* tensor,
index_t index,
const index_t numel,
scalar_t value) {
#if ( \
(defined(USE_ROCM)) || \
(defined(CUDA_VERSION) && (CUDA_VERSION < 10000)) || \
(defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700)))
gpuAtomicAddNoReturn(
reinterpret_cast<at::Half*>(tensor) + index,
static_cast<at::Half>(value));
#else
// Accounts for the chance tensor falls on an odd 16 bit alignment (ie, not 32 bit aligned)
__half* target_addr = reinterpret_cast<__half*>(tensor + index);
bool low_byte = (reinterpret_cast<std::uintptr_t>(target_addr) % sizeof(__half2) == 0);
if (low_byte && index < (numel - 1)) {
__half2 value2;
value2.x = value;
value2.y = __int2half_rz(0);
atomicAdd(reinterpret_cast<__half2*>(target_addr), value2);
} else if (!low_byte && index > 0) {
__half2 value2;
value2.x = __int2half_rz(0);
value2.y = value;
atomicAdd(reinterpret_cast<__half2*>(target_addr - 1), value2);
} else {
atomicAdd(
reinterpret_cast<__half*>(tensor) + index, static_cast<__half>(value));
}
#endif
}
template <
typename scalar_t,
typename index_t,
typename std::enable_if<!std::is_same<c10::Half, scalar_t>::value>::type* =
nullptr>
__device__ __forceinline__ void fastSpecializedAtomicAdd(
scalar_t* tensor,
index_t index,
const index_t numel,
scalar_t value) {
gpuAtomicAddNoReturn(tensor + index, value);
}
template <class scalar_t, class index_t>
__device__ __forceinline__ void fastAtomicAdd(
scalar_t* tensor,
index_t index,
const index_t numel,
scalar_t value,
bool fast_atomics) {
if (fast_atomics) {
fastSpecializedAtomicAdd(tensor, index, numel, value);
} else {
gpuAtomicAddNoReturn(tensor + index, value);
}
}
} // namespace native
} // namespace at
|