GDA 64 bit atomics APIs (#254)

feat(gda): add support for
  * Standard atomics
    - atomic_CAS
    - atomic_fetch_CAS
  * Extended atomics
    - atomic set
    - atomic swap
  * Bitwise atomics
    - atomic_fetch_and
    - atomic_and
    - atomic_fetch_or
    - atomic_or
    - atomic_fetch_xor
    - atomic_xor

src/gda/context_gda_tmpl_device.hpp

@@ -108,70 +108,141 @@ __device__ void GDAContext::amo_add(void *dst, T value, int pe) {
 
 template <typename T>
 __device__ void GDAContext::amo_set(void *dst, T value, int pe) {
-  if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_set not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t
-  uint64_t L_offset = reinterpret_cast<char *>(dst) - base_heap[my_pe];
-  T ret_val;
-  T cond = 0;
-  for (int i = 0; i < WF_SIZE; i++) { //TODO: this looks wrong
-    while ((ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, value, cond, pe))) {
-      if (ret_val == cond) { break; }
-      cond = ret_val;
-    }
-  }
+  amo_swap(dst, value, pe);
 }
 
 template <typename T>
 __device__ T GDAContext::amo_swap(void *dst, T value, int pe) {
-  printf("rocshmem::gda:amo_swap not implemented\n");
-  abort();
-  return 0;
+  if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_set not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t
+  uint64_t L_offset = reinterpret_cast<char *>(dst) - base_heap[my_pe];
+  bool need_turn {true};
+  uint64_t turns = __ballot(need_turn);
+  T ret_val;
+  T cond = 0;
+  while (turns) {
+    uint8_t lane = __ffsll((unsigned long long)turns) - 1;
+    int pe_turn = __shfl(pe, lane);
+    if (pe_turn == pe) {
+      /**
+       * Guess that the remote memory is zero by setting condition to zero.
+       * The compare-and-swap loop will execute at least twice if wrong.
+       * It may run additional times if contention on memory location.
+       */
+      while ((ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, value,
+                         cond, pe)) != cond) {
+        cond = ret_val;
+      }
+      need_turn = false;
+    }
+    turns = __ballot(need_turn);
+  }
+  return ret_val;
 }
 
 template <typename T>
 __device__ T GDAContext::amo_fetch_and(void *dst, T value, int pe) {
-  printf("rocshmem::gda:amo_fetch_and not implemented\n");
-  abort();
-  return 0;
+  if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_fetch_and not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t
+  uint64_t L_offset = reinterpret_cast<char *>(dst) - base_heap[my_pe];
+  bool need_turn {true};
+  uint64_t turns = __ballot(need_turn);
+  T ret_val;
+  T cond = 0;
+  T desired_val = cond & value;
+  while (turns) {
+    uint8_t lane = __ffsll((unsigned long long)turns) - 1;
+    int pe_turn = __shfl(pe, lane);
+    if (pe_turn == pe) {
+      while ((ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset,
+                         desired_val, cond, pe)) != cond) {
+        cond = ret_val;
+        desired_val = ret_val & value;
+      }
+      need_turn = false;
+    }
+    turns = __ballot(need_turn);
+  }
+  return ret_val;
 }
 
 template <typename T>
 __device__ void GDAContext::amo_and(void *dst, T value, int pe) {
-  printf("rocshmem::gda:amo_and not implemented\n");
-  abort();
+  amo_fetch_and(dst, value, pe);
 }
 
 template <typename T>
 __device__ T GDAContext::amo_fetch_or(void *dst, T value, int pe) {
-  printf("rocshmem::gda:amo_fetch_or not implemented\n");
-  abort();
-  return 0;
+  if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_fetch_or not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t
+  uint64_t L_offset = reinterpret_cast<char *>(dst) - base_heap[my_pe];
+  bool need_turn {true};
+  uint64_t turns = __ballot(need_turn);
+  T ret_val;
+  T cond = 0;
+  T desired_val = cond | value;
+  while (turns) {
+    uint8_t lane = __ffsll((unsigned long long)turns) - 1;
+    int pe_turn = __shfl(pe, lane);
+    if (pe_turn == pe) {
+      while ((ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset,
+                         desired_val, cond, pe)) != cond) {
+        cond = ret_val;
+        desired_val = ret_val | value;
+      }
+      need_turn = false;
+    }
+    turns = __ballot(need_turn);
+  }
+  return ret_val;
 }
 
 template <typename T>
 __device__ void GDAContext::amo_or(void *dst, T value, int pe) {
-  printf("rocshmem::gda:amo_or not implemented\n");
-  abort();
+  amo_fetch_or(dst, value, pe);
 }
 
 template <typename T>
 __device__ T GDAContext::amo_fetch_xor(void *dst, T value, int pe) {
-  printf("rocshmem::gda:amo_fetch_xor not implemented\n");
-  abort();
-  return 0;
+  if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_fetch_xor not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t
+  uint64_t L_offset = reinterpret_cast<char *>(dst) - base_heap[my_pe];
+  bool need_turn {true};
+  uint64_t turns = __ballot(need_turn);
+  T ret_val;
+  T cond = 0;
+  T desired_val = cond ^ value;
+  while (turns) {
+    uint8_t lane = __ffsll((unsigned long long)turns) - 1;
+    int pe_turn = __shfl(pe, lane);
+    if (pe_turn == pe) {
+      while ((ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset,
+                         desired_val, cond, pe)) != cond) {
+        cond = ret_val;
+        desired_val = ret_val ^ value;
+      }
+      need_turn = false;
+    }
+    turns = __ballot(need_turn);
+  }
+  return ret_val;
 }
 
 template <typename T>
 __device__ void GDAContext::amo_xor(void *dst, T value, int pe) {
-  printf("rocshmem::gda:amo_xor not implemented\n");
-  abort();
+  amo_fetch_xor(dst, value, pe);
 }
 
 template <typename T>
 __device__ void GDAContext::amo_cas(void *dst, T value, T cond, int pe) {
   if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_cas not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t
   uint64_t L_offset = reinterpret_cast<char *>(dst) - base_heap[my_pe];
-  for (int i = 0; i < WF_SIZE; i++) { //TODO: this looks wrong
-    qps[pe].atomic_cas_nofetch(base_heap[pe] + L_offset, value, cond, pe);
+  bool need_turn {true};
+  uint64_t turns = __ballot(need_turn);
+  while (turns) {
+    uint8_t lane = __ffsll((unsigned long long)turns) - 1;
+    int pe_turn = __shfl(pe, lane);
+    if (pe_turn == pe) {
+      qps[pe].atomic_cas_nofetch(base_heap[pe] + L_offset, value, cond, pe);
+      need_turn = false;
+    }
+    turns = __ballot(need_turn);
   }
 }
 
@@ -198,9 +269,17 @@ template <typename T>
 __device__ T GDAContext::amo_fetch_cas(void *dst, T value, T cond, int pe) {
   if constexpr (sizeof(T) != 8) { printf("rocshmem::gda:amo_fcas not implemented for non-64bit types.\n"); abort(); }//TODO:support for non-uint64t
   uint64_t L_offset = reinterpret_cast<char *>(dst) - base_heap[my_pe];
+  bool need_turn {true};
+  uint64_t turns = __ballot(need_turn);
   T ret_val;
-  for (int i = 0; i < WF_SIZE; i++) {
-    ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, value, cond, pe);
+  while (turns) {
+    uint8_t lane = __ffsll((unsigned long long)turns) - 1;
+    int pe_turn = __shfl(pe, lane);
+    if (pe_turn == pe) {
+      ret_val = qps[pe].atomic_cas(base_heap[pe] + L_offset, value, cond, pe);
+      need_turn = false;
+    }
+    turns = __ballot(need_turn);
   }
   return ret_val;
 }