/**
* \file dnn/src/cuda/atomic.cuh
* MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
*
* Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
* implied.
*/
#pragma once
#include <cuda_runtime_api.h>
#include <stdint.h>
#include "cuda.h"
#include "include/megdnn/dtype.h"
namespace megdnn {
namespace cuda {
#if MEGDNN_CC_CUDA
template <typename T>
static inline MEGDNN_DEVICE void atomic_add(T* address, T val);
template <>
MEGDNN_DEVICE void atomic_add<dt_float32>(dt_float32* address, dt_float32 val) {
::atomicAdd(reinterpret_cast<float*>(address), static_cast<float>(val));
}
// overload atomicAdd for half precision
// Taken from:
// https://github.com/torch/cutorch/blob/master/lib/THC/THCAtomic.cuh
template <>
MEGDNN_DEVICE void atomic_add(dt_float16* address, dt_float16 val) {
#if (__CUDA_ARCH__ < 700 || __CUDACC_VER_MAJOR__ <= 9)
unsigned int* address_as_ui = reinterpret_cast<unsigned int*>(
reinterpret_cast<char*>(address) - (reinterpret_cast<size_t>(address) & 2));
unsigned int old = *address_as_ui;
unsigned int assumed;
do {
assumed = old;
unsigned short data =
reinterpret_cast<size_t>(address) & 2 ? (old >> 16) : (old & 0xffff);
dt_float16 hsum = *reinterpret_cast<dt_float16*>(&data);
hsum += val;
data = *reinterpret_cast<unsigned short*>(&hsum);
old = reinterpret_cast<size_t>(address) & 2 ? (old & 0xffff) | (data << 16)
: (old & 0xffff0000) | data;
old = ::atomicCAS(address_as_ui, assumed, old);
} while (assumed != old);
#else
::atomicAdd(reinterpret_cast<__half*>(address), static_cast<__half>(val));
#endif
}
template <>
MEGDNN_DEVICE void atomic_add(dt_bfloat16* address, dt_bfloat16 val) {
unsigned int* address_as_ui = reinterpret_cast<unsigned int*>(
reinterpret_cast<char*>(address) - (reinterpret_cast<size_t>(address) & 2));
unsigned int old = *address_as_ui;
unsigned int assumed;
do {
assumed = old;
unsigned short data =
reinterpret_cast<size_t>(address) & 2 ? (old >> 16) : (old & 0xffff);
dt_bfloat16 hsum = *reinterpret_cast<dt_bfloat16*>(&data);
hsum += val;
data = *reinterpret_cast<unsigned short*>(&hsum);
old = reinterpret_cast<size_t>(address) & 2 ? (old & 0xffff) | (data << 16)
: (old & 0xffff0000) | data;
old = ::atomicCAS(address_as_ui, assumed, old);
} while (assumed != old);
}
template <typename T, size_t n>
struct AtomicAddIntegerImpl;
template <typename T>
struct AtomicAddIntegerImpl<T, 1> {
inline __device__ void operator()(T* address, T val) {
size_t offset = (size_t)address & 3;
uint32_t* address_as_ui = (uint32_t*)((char*)address - offset);
uint32_t old = *address_as_ui;
uint32_t shift = offset * 8;
uint32_t old_byte;
uint32_t newval;
uint32_t assumed;
do {
assumed = old;
old_byte = (old >> shift) & 0xff;
// preserve size in initial cast. Casting directly to uint32_t pads
// negative signed values with 1's (e.g. signed -1 = unsigned ~0).
newval = static_cast<uint8_t>(
static_cast<T>(val) + static_cast<T>(old_byte));
// newval = static_cast<uint8_t>(THCNumerics<T>::add(val,
// old_byte));
newval = (old & ~(0x000000ff << shift)) | (newval << shift);
old = atomicCAS(address_as_ui, assumed, newval);
} while (assumed != old);
}
};
template <typename T>
struct AtomicAddIntegerImpl<T, 2> {
inline __device__ void operator()(T* address, T val) {
size_t offset = (size_t)address & 2;
uint32_t* address_as_ui = (uint32_t*)((char*)address - offset);
bool is_32_align = offset;
uint32_t old = *address_as_ui;
uint32_t old_bytes;
uint32_t newval;
uint32_t assumed;
do {
assumed = old;
old_bytes = is_32_align ? old >> 16 : old & 0xffff;
// preserve size in initial cast. Casting directly to uint32_t pads
// negative signed values with 1's (e.g. signed -1 = unsigned ~0).
newval = static_cast<uint16_t>(
static_cast<T>(val) + static_cast<T>(old_bytes));
// newval = static_cast<uint16_t>(THCNumerics<T>::add(val,
// old_bytes));
newval = is_32_align ? (old & 0xffff) | (newval << 16)
: (old & 0xffff0000) | newval;
old = atomicCAS(address_as_ui, assumed, newval);
} while (assumed != old);
}
};
template <>
MEGDNN_DEVICE void atomic_add(dt_int32* address, dt_int32 val) {
::atomicAdd(reinterpret_cast<int*>(address), static_cast<int>(val));
}
// we assume quantized int in the same tensor with same scale
template <>
MEGDNN_DEVICE void atomic_add(dt_qint32* address, dt_qint32 val) {
::atomicAdd(reinterpret_cast<int*>(address), val.as_int32());
}
template <>
MEGDNN_DEVICE void atomic_add(dt_int16* address, dt_int16 val) {
AtomicAddIntegerImpl<dt_int16, sizeof(dt_int16)>()(address, val);
}
template <>
MEGDNN_DEVICE void atomic_add(dt_uint16* address, dt_uint16 val) {
AtomicAddIntegerImpl<dt_uint16, sizeof(dt_uint16)>()(address, val);
}
// we assume quantized int in the same tensor with same scale
template <>
MEGDNN_DEVICE void atomic_add(dt_qint16* address, dt_qint16 val) {
AtomicAddIntegerImpl<dt_int16, sizeof(dt_qint16)>()(
reinterpret_cast<dt_int16*>(address), val.as_int16());
}
// be careful! may case over flow
#if 0
template <>
MEGDNN_DEVICE void atomic_add(dt_int8* address, dt_int8 val) {
AtomicAddIntegerImpl<dt_int8, sizeof(dt_int8)>()(address, val);
}
template <>
MEGDNN_DEVICE void atomic_add(dt_uint8* address, dt_uint8 val) {
AtomicAddIntegerImpl<dt_uint8, sizeof(dt_uint8)>()(address, val);
}
// we assume quantized int in the same tensor with same scale
template <>
MEGDNN_DEVICE void atomic_add(dt_quint8* address, dt_quint8 val) {
AtomicAddIntegerImpl<dt_uint8, sizeof(dt_quint8)>()(reinterpret_cast<dt_uint8*>(address), val.as_uint8());
}
// we assume quantized int in the same tensor with same scale
template <>
MEGDNN_DEVICE void atomic_add(dt_qint8* address, dt_qint8 val) {
AtomicAddIntegerImpl<dt_int8, sizeof(dt_qint8)>()(reinterpret_cast<dt_int8*>(address), val.as_int8());
}
#endif
#endif
} // namespace cuda
} // namespace megdnn