Added parity waiting

Author: ogiroux

include/cuda/std/barrier

@@ -77,6 +77,31 @@ _LIBCUDACXX_END_NAMESPACE_CUDA_DEVICE
 
 _LIBCUDACXX_BEGIN_NAMESPACE_CUDA
 
+template<class __Barrier>
+struct __barrier_poll_tester_parity {
+    __Barrier const* __this;
+    bool __parity;
+
+    _LIBCUDACXX_INLINE_VISIBILITY
+    __barrier_poll_tester_parity(__Barrier const* __this_, bool __parity_)
+        : __this(__this_)
+        , __parity(__parity_)
+    {}
+
+    inline _LIBCUDACXX_INLINE_VISIBILITY
+    bool operator()() const
+    {
+        return __this->__try_wait_parity(__parity);
+    }
+};
+
+template<class __Barrier>
+inline _LIBCUDACXX_INLINE_VISIBILITY
+void barrier_wait_for_parity(__Barrier const* __self, bool __parity)
+{
+    _CUDA_VSTD::__libcpp_thread_poll_with_backoff(__barrier_poll_tester_parity<__Barrier>(__self, __parity));
+}
+
 template<>
 class barrier<thread_scope_block, std::__empty_completion> : public __block_scope_barrier_base {
     using __barrier_base = std::__barrier_base<std::__empty_completion, (int)thread_scope_block>;
@@ -88,24 +113,6 @@ class barrier<thread_scope_block, std::__empty_completion> : public __block_scop
 public:
     using arrival_token = typename __barrier_base::arrival_token;
 
-private:
-    struct __poll_tester {
-        barrier const* __this;
-        arrival_token __phase;
-
-        _LIBCUDACXX_INLINE_VISIBILITY
-        __poll_tester(barrier const* __this_, arrival_token&& __phase_)
-          : __this(__this_)
-          , __phase(_CUDA_VSTD::move(__phase_))
-        {}
-
-        inline _LIBCUDACXX_INLINE_VISIBILITY
-        bool operator()() const
-        {
-            return __this->__try_wait(__phase);
-        }
-    };
-
     _LIBCUDACXX_INLINE_VISIBILITY
     bool __try_wait(arrival_token __phase) const {
 #if __CUDA_ARCH__ >= 800
@@ -131,7 +138,28 @@ private:
     template<thread_scope>
     friend class pipeline;
 
-public:
+    _LIBCUDACXX_INLINE_VISIBILITY
+    bool __try_wait_parity(bool __parity) const {
+#if __CUDA_ARCH__ >= 800
+        if (__isShared(&__barrier)) {
+            int __ready = 0;
+            asm volatile ("{\n\t"
+                    ".reg .pred p;\n\t"
+                    "mbarrier.test_wait.parity.shared.b64 p, [%1], %2;\n\t"
+                    "selp.b32 %0, 1, 0, p;\n\t"
+                    "}"
+                : "=r"(__ready)
+                : "r"(static_cast<std::uint32_t>(__cvta_generic_to_shared(&__barrier))), "r"(static_cast<std::uint32_t>(__parity))
+                : "memory");
+            return bool(__ready);
+        }
+        else
+#endif
+        {
+            return __barrier.__try_wait_parity(__parity);
+        }
+    }
+
     barrier() = default;
 
     barrier(const barrier &) = delete;
@@ -216,7 +244,7 @@ public:
     _LIBCUDACXX_INLINE_VISIBILITY
     void wait(arrival_token && __phase) const
     {
-        _CUDA_VSTD::__libcpp_thread_poll_with_backoff(__poll_tester(this, _CUDA_VSTD::move(__phase)));
+        _CUDA_VSTD::__libcpp_thread_poll_with_backoff(std::__barrier_poll_tester<barrier>(this, _CUDA_VSTD::move(__phase)));
     }
 
     inline _LIBCUDACXX_INLINE_VISIBILITY

libcxx/include/barrier

@@ -270,6 +270,24 @@ public:
     }
 };
 
+template<class __Barrier>
+struct __barrier_poll_tester {
+    __Barrier const* __this;
+    typename __Barrier::arrival_token __phase;
+
+    _LIBCUDACXX_INLINE_VISIBILITY
+    __barrier_poll_tester(__Barrier const* __this_, typename __Barrier::arrival_token&& __phase_)
+        : __this(__this_)
+        , __phase(_CUDA_VSTD::move(__phase_))
+    {}
+
+    inline _LIBCUDACXX_INLINE_VISIBILITY
+    bool operator()() const
+    {
+        return __this->__try_wait(__phase);
+    }
+};
+
 template<int _Sco>
 class __barrier_base<__empty_completion, _Sco> {
 
@@ -285,23 +303,6 @@ public:
     using arrival_token = uint64_t;
 
 private:
-    struct __poll_tester {
-        __barrier_base const* __this;
-        arrival_token __phase;
-
-        _LIBCUDACXX_INLINE_VISIBILITY
-        __poll_tester(__barrier_base const* __this_, arrival_token&& __phase_)
-          : __this(__this_)
-          , __phase(_CUDA_VSTD::move(__phase_))
-        {}
-
-        inline _LIBCUDACXX_INLINE_VISIBILITY
-        bool operator()() const
-        {
-            return __this->__try_wait(__phase);
-        }
-    };
-
     static inline _LIBCUDACXX_INLINE_VISIBILITY _LIBCUDACXX_CONSTEXPR
     uint64_t __init(ptrdiff_t __count) _NOEXCEPT
     {
@@ -322,12 +323,22 @@ public:
     __barrier_base(__barrier_base const&) = delete;
     __barrier_base& operator=(__barrier_base const&) = delete;
 
-    _LIBCUDACXX_INLINE_VISIBILITY
-    bool __try_wait(arrival_token __phase) const
+    inline _LIBCUDACXX_INLINE_VISIBILITY
+    bool __try_wait_phase(uint64_t __phase) const
     {
         uint64_t const __current = __phase_arrived_expected.load(memory_order_acquire);
         return ((__current & __phase_bit) != __phase);
     }
+    inline _LIBCUDACXX_INLINE_VISIBILITY
+    bool __try_wait_parity(bool __parity) const
+    {
+        return __try_wait_phase(__parity ? __phase_bit : 0);
+    }
+    inline _LIBCUDACXX_INLINE_VISIBILITY
+    bool __try_wait(arrival_token __old) const
+    {
+        return __try_wait_phase(__old & __phase_bit);
+    }
 
     _LIBCUDACXX_NODISCARD_ATTRIBUTE inline _LIBCUDACXX_INLINE_VISIBILITY
     arrival_token arrive(ptrdiff_t __update = 1)
@@ -343,7 +354,7 @@ public:
     inline _LIBCUDACXX_INLINE_VISIBILITY
     void wait(arrival_token&& __phase) const
     {
-        __libcpp_thread_poll_with_backoff(__poll_tester(this, _CUDA_VSTD::move(__phase)));
+        __libcpp_thread_poll_with_backoff(__barrier_poll_tester<__barrier_base<__empty_completion, _Sco>>(this, _CUDA_VSTD::move(__phase)));
     }
     inline _LIBCUDACXX_INLINE_VISIBILITY
     void arrive_and_wait()