saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-05-26 22:35:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix TopK with NAN on Cuda (#10314)

Browse source 
* reset MIN for float/double

* better logics for float/double comparision for equals
This commit is contained in:
RandySheriffH 2022-01-27 16:19:55 -08:00 committed by GitHub
parent ff2057a817
commit ce081fe655
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
onnxruntime/core/providers/cuda/math/topk_impl.cu
View file

@ -197,13 +197,11 @@ __device__ __inline__ bool Equal(const T& t0, const T& t1) {
}
__device__ __inline__ bool Equal(const float& t0, const float& t1) {
auto t2 = t0 > t1 ? t0 - t1 : t1 - t0;
return t2 < std::numeric_limits<float>::epsilon();
return !(t0 > t1 || t1 > t0);
}
__device__ __inline__ bool Equal(const double& t0, const double& t1) {
auto t2 = t0 > t1 ? t0 - t1 : t1 - t0;
return t2 < std::numeric_limits<double>::epsilon();
return !(t0 > t1 || t1 > t0);
}
template<typename T>
@ -305,7 +303,7 @@ __global__ void RadixTopK(const T* X, T* V, int64_t* I, const TArray<int64_t> el
}
} else {
if (KK > negative) {
KK = dimension - KK + 1;
KK = dimension - KK + 1;
} else {
sign = (T)-1;
}