ba9f6d8f26
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
24 lines
638 B
Plaintext
24 lines
638 B
Plaintext
#pragma once
|
|
#include <sgl_kernel/utils.cuh>
|
|
|
|
namespace device::atomic {
|
|
|
|
SGL_DEVICE float max(float* addr, float value) {
|
|
#ifndef USE_ROCM
|
|
float old;
|
|
old = (value >= 0) ? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
|
|
: __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value)));
|
|
return old;
|
|
#else
|
|
int* addr_as_i = (int*)addr;
|
|
int old = *addr_as_i, assumed;
|
|
do {
|
|
assumed = old;
|
|
old = atomicCAS(addr_as_i, assumed, __float_as_int(fmaxf(value, __int_as_float(assumed))));
|
|
} while (assumed != old);
|
|
return __int_as_float(old);
|
|
#endif
|
|
}
|
|
|
|
} // namespace device::atomic
|