use higher threshold

ngimel · ngimel · commit 64642dbdc123 · 2025-01-07T18:19:30.000-08:00
diff --git a/aten/src/ATen/core/DistributionsHelper.h b/aten/src/ATen/core/DistributionsHelper.h
@@ -40,7 +40,7 @@ struct uniform_int_from_to_distribution {
 
   template <typename RNG>
   C10_HOST_DEVICE inline T operator()(RNG generator) {
-    if (range_ >= 1ULL << 25) // allow approx 1% skew in uniform int generation using %
+    if (range_ >= 1ULL << 28) // allow approx 5% skew in uniform int generation using %
     {
       return transformation::uniform_int_from_to<T>(generator->random64(), range_, base_);
     } else {
diff --git a/aten/src/ATen/native/cuda/DistributionTemplates.h b/aten/src/ATen/native/cuda/DistributionTemplates.h
@@ -280,7 +280,7 @@ namespace cuda {
 template<typename RNG>
 void random_from_to_kernel(TensorIteratorBase& iter, uint64_t range, int64_t base, RNG gen) {
   AT_DISPATCH_V2(iter.dtype(), "random_from_to_kernel_cuda", AT_WRAP([&] {
-    if (range >= 1ULL << 25) // allow approx 1% skew in uniform int generation using %
+    if (range >= 1ULL << 28) // allow approx 5% skew in uniform int generation using %
     {
       // define lambda to mod with range and add base
       auto random_func = [range, base] __device__ (uint64_t rand) {
diff --git a/aten/src/ATen/test/rng_test.h b/aten/src/ATen/test/rng_test.h
@@ -138,8 +138,8 @@ void test_random_from_to(const at::Device& device) {
               from_case_covered = true;
             }
             // this is leaking details of implementation into test
-            // we are starting to use random64() at 2^25 to minimize skew due to %
-            if (range < (1ULL << 25)) {
+            // we are starting to use random64() at 2^28 to minimize skew due to %
+            if (range < (1ULL << 28)) {
               exp = static_cast<T>(static_cast<int64_t>((static_cast<uint32_t>(val) % range + from)));
             } else {
               exp = static_cast<T>(static_cast<int64_t>((val % range + from)));

Original file line number	Diff line number	Diff line change
`@@ -40,7 +40,7 @@ struct uniform_int_from_to_distribution {`
`40`	`40`
`41`	`41`	`template <typename RNG>`
`42`	`42`	`C10_HOST_DEVICE inline T operator()(RNG generator) {`
`43`		`- if (range_ >= 1ULL << 25) // allow approx 1% skew in uniform int generation using %`
	`43`	`+ if (range_ >= 1ULL << 28) // allow approx 5% skew in uniform int generation using %`
`44`	`44`	`{`
`45`	`45`	`return transformation::uniform_int_from_to<T>(generator->random64(), range_, base_);`
`46`	`46`	`} else {`
Original file line number	Diff line number	Diff line change
`@@ -280,7 +280,7 @@ namespace cuda {`
`280`	`280`	`template<typename RNG>`
`281`	`281`	`void random_from_to_kernel(TensorIteratorBase& iter, uint64_t range, int64_t base, RNG gen) {`
`282`	`282`	`AT_DISPATCH_V2(iter.dtype(), "random_from_to_kernel_cuda", AT_WRAP([&] {`
`283`		`- if (range >= 1ULL << 25) // allow approx 1% skew in uniform int generation using %`
	`283`	`+ if (range >= 1ULL << 28) // allow approx 5% skew in uniform int generation using %`
`284`	`284`	`{`
`285`	`285`	`// define lambda to mod with range and add base`
`286`	`286`	`auto random_func = [range, base] __device__ (uint64_t rand) {`