Add declarations for atomic add for fp16 and unsigned short. (#2483)

* add template for fp16 atomic add

* add template for unsigned short atomic add

* use atomicCAS in atomic add for fp16 and unsigned short

* revrt back to atomic add using casting
This commit is contained in:
Illia Silin
2025-07-10 07:18:56 -07:00
committed by GitHub
parent d9b37c7121
commit 1b66f3f4a3
2 changed files with 22 additions and 22 deletions

View File

@@ -32,6 +32,22 @@ __device__ float atomic_add<float>(float* p_dst, const float& x)
return atomicAdd(p_dst, x);
}
template <>
__device__ unsigned short atomic_add<unsigned short>(unsigned short* p_dst, const unsigned short& x)
{
// Use atomicAdd with unsigned int
return static_cast<unsigned short>(
atomicAdd(reinterpret_cast<unsigned int*>(p_dst), static_cast<unsigned int>(x)));
}
template <>
__device__ _Float16 atomic_add<_Float16>(_Float16* p_dst, const _Float16& x)
{
// Use atomicAdd with unsigned int
return static_cast<_Float16>(
atomicAdd(reinterpret_cast<unsigned int*>(p_dst), static_cast<unsigned int>(x)));
}
template <>
__device__ double atomic_add<double>(double* p_dst, const double& x)
{