Fix jagged_test_index_select_2d that hangs in OSS and revert skip tests (#2036)

Summary:
Pull Request resolved: #2036

Before c++20, `std::atomic_flag` is initialized to an unspecified state, hence the loop `while (lock.test_and_set(std::memory_order_acquire)` is never broken and causes the test to hang in OSS. This diff properly initializes the `std::atomic_flag`.

Reviewed By: q10, sryap

Differential Revision: D49528661

fbshipit-source-id: ba2213cb9bf8c0abbd1e169db03f0e32dd2a7ebb

Author: spcyppt

fbgemm_gpu/src/jagged_tensor_ops/jagged_tensor_ops_cpu.cpp

@@ -1151,6 +1151,13 @@ void jagged_index_add_2d_kernel(
   const auto num_cols = input.size(1);
   // Allocate one lock per row
   std::atomic_flag* locks = new std::atomic_flag[output.size(0)];
+  // Initialize all locks since before c++20 std::atomic_flag is initialized to
+  // an unspecified state.
+  // https://en.cppreference.com/w/cpp/atomic/atomic_flag/atomic_flag
+  for (auto i = 0; i < output.size(0); i++) {
+    locks[i].clear();
+  }
+
   at::parallel_for(0, num_dense_input_rows, 0, [&](int64_t start, int64_t end) {
     for (const auto dense_input_offset : c10::irange(start, end)) {
       int index_pos;

fbgemm_gpu/test/jagged_tensor_ops_test.py

@@ -27,7 +27,6 @@
         gpu_available,
         gpu_unavailable,
         on_arm_platform,
-        running_on_github,
         symint_vector_unsupported,
         TEST_WITH_ROCM,
     )
@@ -38,7 +37,6 @@
         gpu_available,
         gpu_unavailable,
         on_arm_platform,
-        running_on_github,
         symint_vector_unsupported,
         TEST_WITH_ROCM,
     )
@@ -1805,7 +1803,6 @@ def jagged_index_select_2d_ref(
         new_embeddings = torch.index_select(values, 0, all_indices)
         return new_embeddings
 
-    @unittest.skipIf(*running_on_github)
     @given(
         max_seq_length=st.integers(5, 10),
         batch_size=st.integers(1, 128),
@@ -1826,7 +1823,7 @@ def jagged_index_select_2d_ref(
         if (gpu_available and TEST_WITH_ROCM)
         else st.just(True),
     )
-    @settings(max_examples=20, deadline=None)
+    @settings(max_examples=20, deadline=None, verbosity=Verbosity.verbose)
     def test_jagged_index_select_2d(
         self,
         max_seq_length: int,
@@ -1899,7 +1896,6 @@ def test_jagged_index_select_2d(
             atol=1e-2 if jagged_tensor_dtype in [torch.half, torch.bfloat16] else None,
         )
 
-    @unittest.skipIf(*running_on_github)
     @given(
         max_seq_length=st.integers(5, 10),
         batch_size=st.integers(1, 128),