Fix a memory bug in the `free_state` function of random pools. #6290

Shihab-Shahriar · 2023-07-17T19:46:13Z

The free_state function was missing a memory fence between two of it's global memory operations.

  KOKKOS_INLINE_FUNCTION
  void free_state(const Random_XorShift64<DeviceType>& state) const {
    state_(state.state_idx_, 0) = state.state_;
    locks_(state.state_idx_, 0) = 0;
  }

Since not all devices guarantee order in global memory operations, without a fence, the lock could get released before the state_ array has been updated in global memory.

…tate update and lock release operations.

dalg24-jenkins · 2023-07-17T19:46:15Z

Can one of the admins verify this patch?

masterleinad · 2023-07-17T20:54:46Z

Is there any chance that we can get a unit test for this?

Shihab-Shahriar · 2023-07-18T07:27:25Z

I added an unit test. In summary, this is how it works:

Since this bug gets triggered only when states get reused by next thread with similar ID, the number of streams (n_streams) should be significantly larger than number of states (which equals execution space's concurrency()). I used n_streams = ExecutionSpace{}.concurrency() * 4.
Generate 8 samples for each stream on device, copy back to a (n_streams ,8) shaped 2d array on host.
To quickly find the duplicates- sort all the streams based on first number, if that's equal then second number and so on. (there's some room for optimization here- we don't need to do entire sort if a duplicate is found. But I opted for simplicity.)
Check each pair of neighbors in the sorted order for duplicates.

I have been able to reliably reproduce the bug for 64 bit random pool in my device (Nvidia V100), but not for 1024 bit one. I think this is because as the pool has 16 states that get updated in free_state, the probability of at least one of them reaching global memory before before lock gets released is much higher.

I also want to note that the test has a narrow focus of reproducing the issue reported in #6140 i.e. checking for exact duplicates. Statistical tests that check for independence of the streams [1] can potentially identify the issue with 1024 bit pool without memory fences.

[1] Salmon, John K., Mark A. Moraes, Ron O. Dror, and David E. Shaw. "Parallel random numbers: as easy as 1, 2, 3." In Proceedings of 2011 international conference for high performance computing, networking, storage and analysis, pp. 1-12. 2011.

masterleinad · 2023-07-18T12:32:22Z

OK to test.

masterleinad · 2023-07-18T12:53:51Z

PLease fix the indentation (according to clang-format-8):

+ ./scripts/docker/check_format_cpp.sh
diff --git a/algorithms/src/Kokkos_Random.hpp b/algorithms/src/Kokkos_Random.hpp
index aada5fe27..2d7d236d2 100644
--- a/algorithms/src/Kokkos_Random.hpp
+++ b/algorithms/src/Kokkos_Random.hpp
@@ -1210,7 +1210,7 @@ class Random_XorShift1024_Pool {
   KOKKOS_INLINE_FUNCTION
   void free_state(const Random_XorShift1024<DeviceType>& state) const {
     for (int i = 0; i < 16; i++) state_(state.state_idx_, i) = state.state_[i];
-    p_(state.state_idx_, 0)     = state.p_;
+    p_(state.state_idx_, 0) = state.p_;
     // Release the lock only after the state has been updated in memory
     Kokkos::memory_fence();
     locks_(state.state_idx_, 0) = 0;
diff --git a/algorithms/unit_tests/TestRandom.hpp b/algorithms/unit_tests/TestRandom.hpp
index 7a31411b9..fae2f029d 100644
--- a/algorithms/unit_tests/TestRandom.hpp
+++ b/algorithms/unit_tests/TestRandom.hpp
@@ -480,8 +480,7 @@ struct generate_random_stream {
   GeneratorPool rand_pool;
   int samples;
 
-  generate_random_stream(ViewType vals_, GeneratorPool rand_pool_,
-                  int samples_)
+  generate_random_stream(ViewType vals_, GeneratorPool rand_pool_, int samples_)
       : vals(vals_), rand_pool(rand_pool_), samples(samples_) {}
 
   KOKKOS_INLINE_FUNCTION
@@ -494,50 +493,48 @@ struct generate_random_stream {
   }
 };
 
-
 // NOTE: this doesn't test the statistical independence of multiple streams
-// generated by a Random pool, it only tests for complete duplicates. 
+// generated by a Random pool, it only tests for complete duplicates.
 template <class ExecutionSpace, class Pool>
-void test_duplicate_stream(){
-  using ViewType = Kokkos::View<uint64_t**, ExecutionSpace>; 
+void test_duplicate_stream() {
+  using ViewType = Kokkos::View<uint64_t**, ExecutionSpace>;
 
   // Heuristic to create a "large enough" number of streams.
   int n_streams = ExecutionSpace{}.concurrency() * 4;
-  int samples = 8;
+  int samples   = 8;
 
   Pool rand_pool(42);
   ViewType vals_d("Vals", n_streams, samples);
   typename ViewType::HostMirror vals_h = Kokkos::create_mirror_view(vals_d);
 
-  Kokkos::parallel_for(n_streams,generate_random_stream<ExecutionSpace, Pool>(
-                          vals_d, rand_pool, samples));
+  Kokkos::parallel_for(n_streams, generate_random_stream<ExecutionSpace, Pool>(
+                                      vals_d, rand_pool, samples));
 
   Kokkos::fence();
   Kokkos::deep_copy(vals_h, vals_d);
 
   /*
   To quickly find streams that are identical, we sort them by the first number,
-  if that's equal then the second and so on. We then test each neighbor pair 
+  if that's equal then the second and so on. We then test each neighbor pair
   for duplicates.
-  */ 
+  */
   std::vector<size_t> indices(n_streams);
   std::iota(indices.begin(), indices.end(), 0);
 
-  auto comparator = [&](int i, int j){
-                  for (int k=0; k<samples; k++){
-                    if(vals_h(i, k)!=vals_h(j, k)) 
-                      return vals_h(i, k) < vals_h(j, k);
-                  }
-                  return false;
-              };
+  auto comparator = [&](int i, int j) {
+    for (int k = 0; k < samples; k++) {
+      if (vals_h(i, k) != vals_h(j, k)) return vals_h(i, k) < vals_h(j, k);
+    }
+    return false;
+  };
   std::sort(indices.begin(), indices.end(), comparator);
 
-  for(int i=0; i< n_streams-1; i++){
-    int idx1 = indices[i], idx2 = indices[i+1];
+  for (int i = 0; i < n_streams - 1; i++) {
+    int idx1 = indices[i], idx2 = indices[i + 1];
 
     int k = 0;
-    while(k < samples && vals_h(idx1, k)==vals_h(idx2,k)) k++;
-    ASSERT_LT(k, samples)  << "Duplicate streams found";
+    while (k < samples && vals_h(idx1, k) == vals_h(idx2, k)) k++;
+    ASSERT_LT(k, samples) << "Duplicate streams found";
   }
 }
 
@@ -583,8 +580,8 @@ TEST(TEST_CATEGORY, Random_XorShift1024_0) {
 
 TEST(TEST_CATEGORY, Multi_streams) {
   using ExecutionSpace = TEST_EXECSPACE;
-  using Pool64 = Kokkos::Random_XorShift64_Pool<ExecutionSpace>;
-  using Pool1024 = Kokkos::Random_XorShift1024_Pool<ExecutionSpace>;
+  using Pool64         = Kokkos::Random_XorShift64_Pool<ExecutionSpace>;
+  using Pool1024       = Kokkos::Random_XorShift1024_Pool<ExecutionSpace>;
 
   AlgoRandomImpl::test_duplicate_stream<ExecutionSpace, Pool64>();
   AlgoRandomImpl::test_duplicate_stream<ExecutionSpace, Pool1024>();

algorithms/unit_tests/TestRandom.hpp

dalg24 · 2023-07-19T12:40:25Z

SYCL failure is unrelated but OpenMPTarget has an issue

29: [ RUN      ] openmptarget.Multi_streams
29: CUDA error: an illegal memory access was encountered 
29: Libomptarget error: Call to targetDataEnd failed, abort target.
29: Libomptarget error: Failed to process data after launching the kernel.
29: Libomptarget error: Consult https://openmp.llvm.org/design/Runtimes.html for debugging options.
29: Kokkos_OpenMPTarget_ParallelFor_Range.hpp:54:1: Libomptarget fatal error 1: failure of target construct while offloading is mandatory

we could disable if we can't figure it out. @rgayatri23 please advise.

rgayatri23 · 2023-07-19T15:16:03Z

SYCL failure is unrelated but OpenMPTarget has an issue

29: [ RUN      ] openmptarget.Multi_streams
29: CUDA error: an illegal memory access was encountered 
29: Libomptarget error: Call to targetDataEnd failed, abort target.
29: Libomptarget error: Failed to process data after launching the kernel.
29: Libomptarget error: Consult https://openmp.llvm.org/design/Runtimes.html for debugging options.
29: Kokkos_OpenMPTarget_ParallelFor_Range.hpp:54:1: Libomptarget fatal error 1: failure of target construct while offloading is mandatory

we could disable if we can't figure it out. @rgayatri23 please advise.

Yeah this is very likely a compiler bug. Can we disable it for OpenMPTarget for now? Once it is merged, I will go line-by-line to see what triggers it and take the necessary action.

algorithms/unit_tests/TestRandom.hpp

dalg24 · 2023-07-20T03:57:01Z

Uggg of course now SYCL fails

39: [ RUN      ] sycl.Multi_streams
39: /var/jenkins/workspace/Kokkos/algorithms/unit_tests/TestRandom.hpp:538: Failure
39: Expected: (k) < (samples), actual: 8 vs 8
39: Duplicate streams found

@masterleinad please advise

masterleinad · 2023-07-20T12:41:06Z

@masterleinad please advise

It works on Intel GPUs so we might just disable the test for SYCL with a FIXME if KOKKOS_ARCH_INTEL_GPU is not defined.

dalg24 · 2023-07-20T12:45:19Z

@masterleinad please advise

It works on Intel GPUs so we might just disable the test for SYCL with a FIXME if KOKKOS_ARCH_INTEL_GPU is not defined.

What about

#if defined(KOKKOS_ENABLE_SYCL) && defined(KOKKOS_IMPL_ARCH_NVIDIA_GPU)
  if constexpr (std::is_same_v<ExecutionSpace,
                               Kokkos::Experimental::SYCL>) {
    GTEST_SKIP() << "Failing on NVIDIA GPUs";  // FIXME_SYCL
  }
#endif

masterleinad · 2023-07-20T12:52:52Z

Sure, we can also just disable for SYCL+CUDA. It boils down to the same restriction anyway for now.

dalg24 · 2023-07-21T19:22:17Z

Ignoring HIP builds that timed out.

Thank you for fixing this!

Insert memory fences in random pools' free_state() function between s…

8dc2ecc

…tate update and lock release operations.

added unit test to check for duplicate streams generated by random pools

06d6e79

apply clang-format-8 on source files

a7865af

masterleinad approved these changes Jul 18, 2023

View reviewed changes

dalg24 reviewed Jul 18, 2023

View reviewed changes

algorithms/unit_tests/TestRandom.hpp Outdated Show resolved Hide resolved

algorithms/unit_tests/TestRandom.hpp Outdated Show resolved Hide resolved

algorithms/unit_tests/TestRandom.hpp Outdated Show resolved Hide resolved

Addressed reviewer comments

306d54d

dalg24 approved these changes Jul 19, 2023

View reviewed changes

algorithms/unit_tests/TestRandom.hpp Outdated Show resolved Hide resolved

little refactor

98951e4

dalg24 reviewed Jul 19, 2023

View reviewed changes

algorithms/unit_tests/TestRandom.hpp Show resolved Hide resolved

skip openmp target for now

ec12805

Skip the test for SYCL+CUDA platform

4f651f5

dalg24 merged commit 4d1c6c3 into kokkos:develop Jul 21, 2023
27 of 28 checks passed

dalg24 mentioned this pull request Jul 14, 2023

CHANGELOG: 4.2.0 #6197

Closed

ajpowelsnl mentioned this pull request Aug 21, 2023

Kokkos_Random: Random number sequence seems to get repeated for different threads on CUDA #6140

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix a memory bug in the `free_state` function of random pools. #6290

Fix a memory bug in the `free_state` function of random pools. #6290

Shihab-Shahriar commented Jul 17, 2023

dalg24-jenkins commented Jul 17, 2023

masterleinad commented Jul 17, 2023

Shihab-Shahriar commented Jul 18, 2023

masterleinad commented Jul 18, 2023

masterleinad commented Jul 18, 2023

dalg24 commented Jul 19, 2023

rgayatri23 commented Jul 19, 2023

dalg24 commented Jul 20, 2023

masterleinad commented Jul 20, 2023

dalg24 commented Jul 20, 2023 •

edited

masterleinad commented Jul 20, 2023

dalg24 commented Jul 21, 2023

Fix a memory bug in the free_state function of random pools. #6290

Fix a memory bug in the free_state function of random pools. #6290

Conversation

Shihab-Shahriar commented Jul 17, 2023

dalg24-jenkins commented Jul 17, 2023

masterleinad commented Jul 17, 2023

Shihab-Shahriar commented Jul 18, 2023

masterleinad commented Jul 18, 2023

masterleinad commented Jul 18, 2023

dalg24 commented Jul 19, 2023

rgayatri23 commented Jul 19, 2023

dalg24 commented Jul 20, 2023

masterleinad commented Jul 20, 2023

dalg24 commented Jul 20, 2023 • edited

masterleinad commented Jul 20, 2023

dalg24 commented Jul 21, 2023

Fix a memory bug in the `free_state` function of random pools. #6290

Fix a memory bug in the `free_state` function of random pools. #6290

dalg24 commented Jul 20, 2023 •

edited