Improvemenet to uniform floating point random.

More expected range, slight increase to generation cost.
FLAMEGPU · Oct 23, 2020 · ef6ada9 · ef6ada9
1 parent ea0a348
commit ef6ada9
Showing 1 changed file with 15 additions and 7 deletions.
diff --git a/include/flamegpu/runtime/utility/AgentRandom.cuh b/include/flamegpu/runtime/utility/AgentRandom.cuh
@@ -2,7 +2,8 @@
 #define INCLUDE_FLAMEGPU_RUNTIME_UTILITY_AGENTRANDOM_CUH_
 
 #include <curand_kernel.h>
-#include <cassert>
+#include <limits>
+#include <cfloat>
 
 #include "flamegpu/exception/FGPUStaticAssert.h"
 #include "flamegpu/exception/FGPUDeviceException.h"
@@ -23,7 +24,7 @@ class AgentRandom {
     __forceinline__ __device__ AgentRandom(curandState *d_rng);
     /**
      * Returns a float uniformly distributed between 0.0 and 1.0. 
-     * @note It may return from 0.0 to 1.0, where 1.0 is included and 0.0 is excluded.
+     * @note It may return from 0.0 to 1.0, where 0.0 is included and 1.0 is excluded.
      * @note Available as float or double
      */
     template<typename T>
@@ -65,11 +66,18 @@ __forceinline__ __device__ AgentRandom::AgentRandom(curandState *d_rng) : d_rand
  */
 template<>
 __forceinline__ __device__ float AgentRandom::uniform() const {
-    return curand_uniform(d_random_state);
+    // curand naturally generates the range (0, 1], we want [0, 1)
+    // https://github.com/pytorch/pytorch/blob/059aa34b124916dfd761f3cbdb5fa97d7a01fc93/aten/src/ATen/native/cuda/Distributions.cu#L71-L77
+    uint32_t val = curand(d_random_state);  // need just bits
+    constexpr auto MASK = static_cast<uint32_t>((static_cast<uint64_t>(1) << std::numeric_limits<float>::digits) - 1);
+    constexpr auto DIVISOR = static_cast<float>(1) / (static_cast<uint32_t>(1) << std::numeric_limits<float>::digits);
+    return (val & MASK) * DIVISOR;
 }
 template<>
 __forceinline__ __device__ double AgentRandom::uniform() const {
-    return curand_uniform_double(d_random_state);
+    // curand naturally generates the range (0, 1], we want [0, 1)
+    double d = curand_uniform_double(d_random_state) - DBL_EPSILON;
+    return d < 0.0 ? 0.0 : d;
 }
 
 /**
@@ -105,7 +113,7 @@ __forceinline__ __device__ T AgentRandom::uniform(const T& min, const T& max) co
         DTHROW("Invalid arguments passed to AgentRandom::uniform(), %lld > %lld\n", static_cast<int64_t>(min), static_cast<int64_t>(max));
     }
 #endif
-    return static_cast<T>(min + (1 + max - min) * (1.0 - uniform<float>()));
+    return static_cast<T>(min + (1 + max - min) * uniform<float>());
 }
 template<>
 __forceinline__ __device__ int64_t AgentRandom::uniform(const int64_t& min, const int64_t& max) const {
@@ -114,7 +122,7 @@ __forceinline__ __device__ int64_t AgentRandom::uniform(const int64_t& min, cons
         DTHROW("Invalid arguments passed to AgentRandom::uniform(), %lld > %lld\n", static_cast<int64_t>(min), static_cast<int64_t>(max));
     }
 #endif
-    return static_cast<int64_t>(min + (1 + max - min) * (1.0 - uniform<double>()));
+    return static_cast<int64_t>(min + (1 + max - min) * uniform<double>());
 }
 template<>
 __forceinline__ __device__ uint64_t AgentRandom::uniform(const uint64_t& min, const uint64_t& max) const {
@@ -123,6 +131,6 @@ __forceinline__ __device__ uint64_t AgentRandom::uniform(const uint64_t& min, co
         DTHROW("Invalid arguments passed to AgentRandom::uniform(), %lld > %lld\n", static_cast<int64_t>(min), static_cast<int64_t>(max));
     }
 #endif
-    return static_cast<uint64_t>(min + (1 + max - min) * (1.0 - uniform<double>()));
+    return static_cast<uint64_t>(min + (1 + max - min) * uniform<double>());
 }
 #endif  // INCLUDE_FLAMEGPU_RUNTIME_UTILITY_AGENTRANDOM_CUH_