[INFRA] Import NVIDIA/CCCL upstream as optimization reference library

CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

Competition optimization vectors mapped to CCCL:
- Output TPS (83% weight): warp_reduce, block_reduce, device_topk
- Input TPS (14% weight): device_scan, block_load, prefetch
- Cache TPS (3% weight): prefix caching strategy patterns
- Memory (0.9 util): pooled/cached/buddy allocators

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/check_preconditions.h

@@ -0,0 +1,80 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_CHECK_PRECONDITIONS_H
+#define _CUDA___MEMCPY_ASYNC_CHECK_PRECONDITIONS_H
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/__memory/aligned_size.h>
+#include <cuda/std/__algorithm/max.h>
+#include <cuda/std/__cstddef/types.h>
+#include <cuda/std/__memory/is_sufficiently_aligned.h>
+#include <cuda/std/cstdint>
+
+#include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+#ifndef _LIBCUDACXX_MEMCPY_ASYNC_PRE_TESTING
+#  define _LIBCUDACXX_MEMCPY_ASYNC_PRE_ASSERT(_Cond, _Msg) _CCCL_ASSERT(_Cond, _Msg)
+#else // ^^^ _LIBCUDACXX_MEMCPY_ASYNC_PRE_TESTING ^^^ / vvv !_LIBCUDACXX_MEMCPY_ASYNC_PRE_TESTING vvv
+#  define _LIBCUDACXX_MEMCPY_ASYNC_PRE_ASSERT(_Cond, _Msg) \
+    do                                                     \
+    {                                                      \
+      if (!(_Cond))                                        \
+      {                                                    \
+        return false;                                      \
+      }                                                    \
+    } while (false)
+#endif // _LIBCUDACXX_MEMCPY_ASYNC_PRE_TESTING
+
+// Check the memcpy_async preconditions, return value is intended for testing purposes exclusively
+template <class _Tp, class _Size>
+_CCCL_HOST_DEVICE_API inline bool __memcpy_async_check_pre(_Tp* __dst, const _Tp* __src, _Size __size)
+{
+  constexpr auto __align = ::cuda::std::max(alignof(_Tp), __get_size_align_v<_Size>);
+
+  const auto __dst_val = reinterpret_cast<uintptr_t>(__dst);
+  const auto __src_val = reinterpret_cast<uintptr_t>(__src);
+
+  // check src and dst alignment
+  _LIBCUDACXX_MEMCPY_ASYNC_PRE_ASSERT(::cuda::std::is_sufficiently_aligned<__align>(__dst),
+                                      "destination pointer must be aligned to the specified alignment");
+  _LIBCUDACXX_MEMCPY_ASYNC_PRE_ASSERT(
+    ::cuda::std::is_sufficiently_aligned<__align>(__src), "source pointer must be aligned to the specified alignment");
+
+  // check src and dst overlap
+  _LIBCUDACXX_MEMCPY_ASYNC_PRE_ASSERT(
+    !((__dst_val <= __src_val && __src_val < __dst_val + __size)
+      || (__src_val <= __dst_val && __dst_val < __src_val + __size)),
+    "destination and source buffers must not overlap");
+  return true;
+}
+
+template <class _Size>
+_CCCL_HOST_DEVICE_API inline bool __memcpy_async_check_pre(void* __dst, const void* __src, _Size __size)
+{
+  return ::cuda::__memcpy_async_check_pre(reinterpret_cast<char*>(__dst), reinterpret_cast<const char*>(__src), __size);
+}
+
+_CCCL_END_NAMESPACE_CUDA
+
+#include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CUDA___MEMCPY_ASYNC_CHECK_PRECONDITIONS_H

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/completion_mechanism.h

@@ -0,0 +1,47 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_COMPLETION_MECHANISM_H
+#define _CUDA___MEMCPY_ASYNC_COMPLETION_MECHANISM_H
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+//! @brief __completion_mechanism allows memcpy_async to report back what completion
+//! mechanism it used. This is necessary to determine in which way to synchronize
+//! the memcpy_async with a sync object (barrier or pipeline).
+//
+//! In addition, we use this enum to create bit flags so that calling functions
+//! can specify which completion mechanisms can be used (__sync is always
+//! allowed).
+enum class __completion_mechanism
+{
+  __sync                 = 0,
+  __mbarrier_complete_tx = 1 << 0, // Use powers of two here to support the
+  __async_group          = 1 << 1, // bit flag use case
+  __async_bulk_group     = 1 << 2,
+};
+
+_CCCL_END_NAMESPACE_CUDA
+
+#include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CUDA___MEMCPY_ASYNC_COMPLETION_MECHANISM_H

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/cp_async_bulk_shared_global.h

@@ -0,0 +1,71 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_CP_ASYNC_BULK_SHARED_GLOBAL_H_
+#define _CUDA___MEMCPY_ASYNC_CP_ASYNC_BULK_SHARED_GLOBAL_H_
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#if _CCCL_CUDA_COMPILATION()
+#  if __cccl_ptx_isa >= 800
+
+#    include <cuda/__memcpy_async/elect_one.h>
+#    include <cuda/__ptx/instructions/cp_async_bulk.h>
+#    include <cuda/__ptx/instructions/mbarrier_expect_tx.h>
+#    include <cuda/__ptx/ptx_dot_variants.h>
+#    include <cuda/__ptx/ptx_helper_functions.h>
+#    include <cuda/std/__type_traits/conditional.h>
+#    include <cuda/std/cstdint>
+
+#    include <nv/target>
+
+#    include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+extern "C" _CCCL_DEVICE void __cuda_ptx_cp_async_bulk_shared_global_is_not_supported_before_SM_90__();
+template <typename _Group>
+_CCCL_DEVICE_API inline void __cp_async_bulk_shared_global_and_expect_tx(
+  const _Group& __g, char* __dest, const char* __src, ::cuda::std::size_t __size, ::cuda::std::uint64_t* __bar_handle)
+{
+  // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk
+  NV_IF_ELSE_TARGET(
+    NV_PROVIDES_SM_90,
+    (if (::cuda::device::__group_elect_one(__g)) {
+      ::cuda::ptx::cp_async_bulk(
+        ::cuda::std::conditional_t<__cccl_ptx_isa >= 860, ::cuda::ptx::space_shared_t, ::cuda::ptx::space_cluster_t>{},
+        ::cuda::ptx::space_global,
+        __dest,
+        __src,
+        __size,
+        __bar_handle);
+      ::cuda::ptx::mbarrier_expect_tx(
+        ::cuda::ptx::sem_relaxed, ::cuda::ptx::scope_cta, ::cuda::ptx::space_shared, __bar_handle, __size);
+    }),
+    (::cuda::__cuda_ptx_cp_async_bulk_shared_global_is_not_supported_before_SM_90__();));
+}
+
+_CCCL_END_NAMESPACE_CUDA
+
+#    include <cuda/std/__cccl/epilogue.h>
+
+#  endif // __cccl_ptx_isa >= 800
+#endif // _CCCL_CUDA_COMPILATION()
+
+#endif // _CUDA___MEMCPY_ASYNC_CP_ASYNC_BULK_SHARED_GLOBAL_H_

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/cp_async_fallback.h

@@ -0,0 +1,72 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_CP_ASYNC_FALLBACK_H_
+#define _CUDA___MEMCPY_ASYNC_CP_ASYNC_FALLBACK_H_
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/std/cstddef>
+
+#include <nv/target>
+
+#include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+template <::cuda::std::size_t _Copy_size>
+struct __copy_chunk
+{
+  _CCCL_ALIGNAS(_Copy_size) char data[_Copy_size];
+};
+
+template <::cuda::std::size_t _Alignment, typename _Group>
+_CCCL_HOST_DEVICE_API void
+__cp_async_fallback_mechanism(_Group __g, char* __dest, const char* __src, ::cuda::std::size_t __size)
+{
+  // Maximal copy size is 16 bytes
+  constexpr ::cuda::std::size_t __copy_size = (_Alignment > 16) ? 16 : _Alignment;
+
+  using __chunk_t = __copy_chunk<__copy_size>;
+
+  // "Group"-strided loop over memory
+  const ::cuda::std::size_t __stride = __g.size() * __copy_size;
+
+  // An unroll factor of 64 ought to be enough for anybody. This unroll pragma
+  // is mainly intended to place an upper bound on loop unrolling. The number
+  // is more than high enough for the intended use case: an unroll factor of
+  // 64 allows moving 4 * 64 * 256 = 64kb in one unrolled loop with 256
+  // threads (copying ints). On the other hand, in the unfortunate case that
+  // we have to move 1024 bytes / thread with char width, then we prevent
+  // fully unrolling the loop to 1024 copy instructions. This prevents the
+  // compile times from increasing unreasonably, and also has negligible
+  // impact on runtime performance.
+  _CCCL_PRAGMA_UNROLL(64)
+  for (::cuda::std::size_t __offset = __g.thread_rank() * __copy_size; __offset < __size; __offset += __stride)
+  {
+    __chunk_t tmp                                    = *reinterpret_cast<const __chunk_t*>(__src + __offset);
+    *reinterpret_cast<__chunk_t*>(__dest + __offset) = tmp;
+  }
+}
+
+_CCCL_END_NAMESPACE_CUDA
+
+#include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CUDA___MEMCPY_ASYNC_CP_ASYNC_FALLBACK_H_

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/cp_async_shared_global.h

@@ -0,0 +1,148 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_CP_ASYNC_SHARED_GLOBAL_H_
+#define _CUDA___MEMCPY_ASYNC_CP_ASYNC_SHARED_GLOBAL_H_
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#if _CCCL_CUDA_COMPILATION()
+
+#  include <cuda/__ptx/ptx_dot_variants.h>
+#  include <cuda/__ptx/ptx_helper_functions.h>
+#  include <cuda/std/cstdint>
+
+#  include <nv/target>
+
+#  include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+extern "C" _CCCL_DEVICE void __cuda_ptx_cp_async_shared_global_is_not_supported_before_SM_80__();
+
+#  if _CCCL_CUDA_COMPILER(NVCC, <, 12, 1) // WAR for compiler state space issues
+template <size_t _Copy_size>
+_CCCL_DEVICE_API void __cp_async_shared_global(char* __dest, const char* __src)
+{
+  // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async
+
+  // If `if constexpr` is not available, this function gets instantiated even
+  // if is not called. Do not static_assert in that case.
+  static_assert(_Copy_size == 4 || _Copy_size == 8 || _Copy_size == 16,
+                "cp.async.shared.global requires a copy size of 4, 8, or 16.");
+
+  NV_IF_ELSE_TARGET(
+    NV_PROVIDES_SM_80,
+    (asm volatile(R"XYZ(
+      {
+        .reg .b64 tmp;
+        .reg .b32 dst;
+
+        cvta.to.shared.u64 tmp, %0;
+        cvt.u32.u64 dst, tmp;
+        cvta.to.global.u64 tmp, %1;
+        cp.async.ca.shared.global [dst], [tmp], %2, %2;
+      }
+      )XYZ" : : "l"(__dest),
+                  "l"(__src),
+                  "n"(_Copy_size) : "memory");),
+    (::cuda::__cuda_ptx_cp_async_shared_global_is_not_supported_before_SM_80__();));
+}
+template <>
+_CCCL_DEVICE_API inline void __cp_async_shared_global<16>(char* __dest, const char* __src)
+{
+  // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async
+  // When copying 16 bytes, it is possible to skip L1 cache (.cg).
+  NV_IF_ELSE_TARGET(NV_PROVIDES_SM_80,
+                    (asm volatile(R"XYZ(
+      {
+        .reg .u64 tmp;
+        .reg .u32 dst;
+
+        cvta.to.shared.u64 tmp, %0;
+        cvt.u32.u64 dst, tmp;
+        cvta.to.global.u64 tmp, %1;
+        cp.async.cg.shared.global [dst], [tmp], 16, 16;
+      }
+      )XYZ" : : "l"(__dest),
+                                  "l"(__src) : "memory");),
+                    (::cuda::__cuda_ptx_cp_async_shared_global_is_not_supported_before_SM_80__();));
+}
+#  else // ^^^^ NVCC 12.0 / !NVCC 12.0 vvvvv WAR for compiler state space issues
+template <size_t _Copy_size>
+_CCCL_DEVICE_API void __cp_async_shared_global(char* __dest, const char* __src)
+{
+  // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async
+
+  // If `if constexpr` is not available, this function gets instantiated even
+  // if is not called. Do not static_assert in that case.
+  static_assert(_Copy_size == 4 || _Copy_size == 8 || _Copy_size == 16,
+                "cp.async.shared.global requires a copy size of 4, 8, or 16.");
+
+  NV_IF_ELSE_TARGET(
+    NV_PROVIDES_SM_80,
+    (asm volatile("cp.async.ca.shared.global [%0], [%1], %2, %2;" : : "r"(
+                    static_cast<::cuda::std::uint32_t>(::__cvta_generic_to_shared(__dest))),
+                  "l"(static_cast<::cuda::std::uint64_t>(::__cvta_generic_to_global(__src))),
+                  "n"(_Copy_size) : "memory");),
+    (::cuda::__cuda_ptx_cp_async_shared_global_is_not_supported_before_SM_80__();));
+}
+template <>
+_CCCL_DEVICE_API inline void __cp_async_shared_global<16>(char* __dest, const char* __src)
+{
+  // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async
+  // When copying 16 bytes, it is possible to skip L1 cache (.cg).
+  NV_IF_ELSE_TARGET(
+    NV_PROVIDES_SM_80,
+    (asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %2;" : : "r"(
+                    static_cast<::cuda::std::uint32_t>(::__cvta_generic_to_shared(__dest))),
+                  "l"(static_cast<::cuda::std::uint64_t>(::__cvta_generic_to_global(__src))),
+                  "n"(16) : "memory");),
+    (::cuda::__cuda_ptx_cp_async_shared_global_is_not_supported_before_SM_80__();));
+}
+#  endif // _CCCL_CUDA_COMPILER(NVCC, >=, 12, 1)
+
+template <size_t _Alignment, typename _Group>
+_CCCL_DEVICE_API void
+__cp_async_shared_global_mechanism(_Group __g, char* __dest, const char* __src, ::cuda::std::size_t __size)
+{
+  // If `if constexpr` is not available, this function gets instantiated even
+  // if is not called. Do not static_assert in that case.
+  static_assert(4 <= _Alignment, "cp.async requires at least 4-byte alignment");
+
+  // Maximal copy size is 16.
+  constexpr int __copy_size = (_Alignment > 16) ? 16 : _Alignment;
+  // We use an int offset here, because we are copying to shared memory,
+  // which is easily addressable using int.
+  const int __group_size = __g.size();
+  const int __group_rank = __g.thread_rank();
+  const int __stride     = __group_size * __copy_size;
+  for (int __offset = __group_rank * __copy_size; __offset < static_cast<int>(__size); __offset += __stride)
+  {
+    ::cuda::__cp_async_shared_global<__copy_size>(__dest + __offset, __src + __offset);
+  }
+}
+
+_CCCL_END_NAMESPACE_CUDA
+
+#  include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CCCL_CUDA_COMPILATION()
+
+#endif // _CUDA___MEMCPY_ASYNC_CP_ASYNC_SHARED_GLOBAL_H_

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/dispatch_memcpy_async.h

@@ -0,0 +1,165 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_DISPATCH_MEMCPY_ASYNC_H_
+#define _CUDA___MEMCPY_ASYNC_DISPATCH_MEMCPY_ASYNC_H_
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/__memcpy_async/completion_mechanism.h>
+#include <cuda/__memcpy_async/cp_async_bulk_shared_global.h>
+#include <cuda/__memcpy_async/cp_async_fallback.h>
+#include <cuda/__memcpy_async/cp_async_shared_global.h>
+#include <cuda/__memory/address_space.h>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/cstring>
+
+#include <nv/target>
+
+#include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+/***********************************************************************
+ * cuda::memcpy_async dispatch
+ *
+ * The dispatch mechanism takes all the arguments and dispatches to the
+ * fastest asynchronous copy mechanism available.
+ *
+ * It returns a __completion_mechanism that indicates which completion mechanism
+ * was used by the copy mechanism. This value can be used by the sync object to
+ * further synchronize if necessary.
+ *
+ ***********************************************************************/
+
+template <::cuda::std::size_t _Align, typename _Group>
+[[nodiscard]] _CCCL_DEVICE_API __completion_mechanism __dispatch_memcpy_async_any_to_any(
+  _Group const& __group,
+  char* __dest_char,
+  char const* __src_char,
+  ::cuda::std::size_t __size,
+  [[maybe_unused]] ::cuda::std::uint32_t __allowed_completions,
+  [[maybe_unused]] ::cuda::std::uint64_t* __bar_handle)
+{
+  ::cuda::__cp_async_fallback_mechanism<_Align>(__group, __dest_char, __src_char, __size);
+  return __completion_mechanism::__sync;
+}
+
+template <::cuda::std::size_t _Align, typename _Group>
+[[nodiscard]] _CCCL_DEVICE_API __completion_mechanism __dispatch_memcpy_async_global_to_shared(
+  _Group const& __group,
+  char* __dest_char,
+  char const* __src_char,
+  ::cuda::std::size_t __size,
+  [[maybe_unused]] ::cuda::std::uint32_t __allowed_completions,
+  [[maybe_unused]] ::cuda::std::uint64_t* __bar_handle)
+{
+#if __cccl_ptx_isa >= 800
+  NV_IF_TARGET(
+    NV_PROVIDES_SM_90,
+    ([[maybe_unused]] const bool __can_use_complete_tx =
+       __allowed_completions & uint32_t(__completion_mechanism::__mbarrier_complete_tx);
+     _CCCL_ASSERT(__can_use_complete_tx == (nullptr != __bar_handle),
+                  "Pass non-null bar_handle if and only if can_use_complete_tx.");
+     if constexpr (_Align >= 16) {
+       if (__can_use_complete_tx
+           && ::cuda::device::is_address_from(__bar_handle, ::cuda::device::address_space::shared))
+       {
+         ::cuda::__cp_async_bulk_shared_global_and_expect_tx(__group, __dest_char, __src_char, __size, __bar_handle);
+         return __completion_mechanism::__mbarrier_complete_tx;
+       }
+     }
+     // Fallthrough to SM 80..
+     ));
+#endif // __cccl_ptx_isa >= 800
+
+  NV_IF_TARGET(
+    NV_PROVIDES_SM_80,
+    (if constexpr (_Align >= 4) {
+      const bool __can_use_async_group = __allowed_completions & uint32_t(__completion_mechanism::__async_group);
+      if (__can_use_async_group)
+      {
+        ::cuda::__cp_async_shared_global_mechanism<_Align>(__group, __dest_char, __src_char, __size);
+        return __completion_mechanism::__async_group;
+      }
+    }
+     // Fallthrough..
+     ));
+
+  ::cuda::__cp_async_fallback_mechanism<_Align>(__group, __dest_char, __src_char, __size);
+  return __completion_mechanism::__sync;
+}
+
+// __dispatch_memcpy_async is the internal entry point for dispatching to the correct memcpy_async implementation.
+template <::cuda::std::size_t _Align, typename _Group>
+[[nodiscard]] _CCCL_HOST_DEVICE_API __completion_mechanism __dispatch_memcpy_async(
+  _Group const& __group,
+  char* __dest_char,
+  char const* __src_char,
+  ::cuda::std::size_t __size,
+  [[maybe_unused]] ::cuda::std::uint32_t __allowed_completions,
+  [[maybe_unused]] ::cuda::std::uint64_t* __bar_handle)
+{
+  NV_IF_ELSE_TARGET(
+    NV_IS_DEVICE,
+    (
+      // Dispatch based on direction of the copy: global to shared, shared to
+      // global, etc.
+
+      // CUDA compilers <= 12.2 may not propagate assumptions about the state space
+      // of pointers correctly. Therefore, we
+      // 1) put the code for each copy direction in a separate function, and
+      // 2) make sure none of the code paths can reach each other by "falling through".
+      //
+      // See nvbug 4074679 and also PR #478.
+      if (::cuda::device::is_address_from(__src_char, ::cuda::device::address_space::global)
+          && ::cuda::device::is_address_from(__dest_char, ::cuda::device::address_space::shared)) {
+        return ::cuda::__dispatch_memcpy_async_global_to_shared<_Align>(
+          __group, __dest_char, __src_char, __size, __allowed_completions, __bar_handle);
+      } else {
+        return ::cuda::__dispatch_memcpy_async_any_to_any<_Align>(
+          __group, __dest_char, __src_char, __size, __allowed_completions, __bar_handle);
+      }),
+    (
+      // Host code path:
+      if (__group.thread_rank() == 0) {
+        ::cuda::std::memcpy(__dest_char, __src_char, __size);
+      } return __completion_mechanism::__sync;));
+}
+
+template <::cuda::std::size_t _Align, typename _Group>
+[[nodiscard]] _CCCL_HOST_DEVICE_API __completion_mechanism __dispatch_memcpy_async(
+  _Group const& __group,
+  char* __dest_char,
+  char const* __src_char,
+  [[maybe_unused]] ::cuda::std::size_t __size,
+  [[maybe_unused]] ::cuda::std::uint32_t __allowed_completions)
+{
+  _CCCL_ASSERT(!(__allowed_completions & uint32_t(__completion_mechanism::__mbarrier_complete_tx)),
+               "Cannot allow mbarrier_complete_tx completion mechanism when not passing a barrier. ");
+  return ::cuda::__dispatch_memcpy_async<_Align>(
+    __group, __dest_char, __src_char, __size, __allowed_completions, nullptr);
+}
+
+_CCCL_END_NAMESPACE_CUDA
+
+#include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CUDA___MEMCPY_ASYNC_DISPATCH_MEMCPY_ASYNC_H_

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/elect_one.h

@@ -0,0 +1,73 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_ELECT_ONE_H_
+#define _CUDA___MEMCPY_ASYNC_ELECT_ONE_H_
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/__memcpy_async/group_traits.h>
+#include <cuda/__ptx/instructions/elect_sync.h>
+
+#include <cuda/std/__cccl/prologue.h>
+
+#if _CCCL_CUDA_COMPILATION()
+
+_CCCL_BEGIN_NAMESPACE_CUDA_DEVICE
+
+//! Elects a single leader thread from a one dimensional thread block. For SM90+ will use ptx::elect_sync() etc.,
+//! otherwise just selects the thread with ID 0. If the returned value is used as condition for an if statement, the
+//! compiler will emit a uniform data path for the branch.
+[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE bool __block_elect_one() noexcept
+{
+  _CCCL_ASSERT(blockDim.y == 1 && blockDim.z == 1, "The block must by one dimensional");
+
+  NV_IF_ELSE_TARGET(
+    NV_PROVIDES_SM_90,
+    (const auto tid             = threadIdx.x; //
+     const auto warp_id         = tid / 32;
+     const auto uniform_warp_id = ::__shfl_sync(~0, warp_id, 0); // broadcast from lane 0
+     return uniform_warp_id == 0 && ::cuda::ptx::elect_sync(~0); // elect a leader thread among warp 0
+     ),
+    (return threadIdx.x == 0;));
+}
+
+template <typename _Group>
+[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE bool __group_elect_one(const _Group& __g) noexcept
+{
+  NV_IF_TARGET(NV_PROVIDES_SM_90,
+               (
+                 if constexpr (is_thread_block_group_v<_Group>) {
+                   // cooperative groups maps a multidimensional thread id into the thread rank the same way as warps do
+                   const unsigned __tid             = __g.thread_rank();
+                   const unsigned __warp_id         = __tid / 32;
+                   const unsigned __uniform_warp_id = ::__shfl_sync(~0, __warp_id, 0); // broadcast from lane 0
+                   return __uniform_warp_id == 0 && ::cuda::ptx::elect_sync(~0); // elect a leader thread among warp 0
+                 } else if constexpr (is_warp_group_v<_Group>) { return ::cuda::ptx::elect_sync(~0); }));
+
+  return __g.thread_rank() == 0;
+}
+
+#endif // _CCCL_CUDA_COMPILATION()
+
+_CCCL_END_NAMESPACE_CUDA_DEVICE
+
+#include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CUDA___MEMCPY_ASYNC_ELECT_ONE_H_

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/group_traits.h

@@ -0,0 +1,61 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_GROUP_TRAITS_H_
+#define _CUDA___MEMCPY_ASYNC_GROUP_TRAITS_H_
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/std/__cccl/prologue.h>
+
+// forward declare cooperative groups types. we cannot include <cooperative_groups.h> since it does not work with NVHPC
+namespace cooperative_groups
+{
+namespace __v1
+{
+class thread_block;
+
+template <unsigned int Size, typename ParentT>
+class thread_block_tile;
+} // namespace __v1
+using namespace __v1;
+} // namespace cooperative_groups
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+//! Trait to detect whether a group represents a CUDA thread block, for example: ``cooperative_groups::thread_block``.
+template <typename _Group>
+inline constexpr bool is_thread_block_group_v = false;
+
+template <>
+inline constexpr bool is_thread_block_group_v<::cooperative_groups::thread_block> = true;
+
+//! Trait to detect whether a group represents a CUDA warp, for example:
+//! ``cooperative_groups::thread_block_tile<32, ...>``.
+template <typename _Group>
+inline constexpr bool is_warp_group_v = false;
+
+template <typename _Parent>
+inline constexpr bool is_warp_group_v<::cooperative_groups::thread_block_tile<32, _Parent>> = true;
+
+_CCCL_END_NAMESPACE_CUDA
+
+#include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CUDA___MEMCPY_ASYNC_GROUP_TRAITS_H_

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/is_local_smem_barrier.h

@@ -0,0 +1,53 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_IS_LOCAL_SMEM_BARRIER_H
+#define _CUDA___MEMCPY_ASYNC_IS_LOCAL_SMEM_BARRIER_H
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/__barrier/barrier.h>
+#include <cuda/__memory/address_space.h>
+#include <cuda/std/__atomic/scopes.h>
+#include <cuda/std/__barrier/empty_completion.h>
+#include <cuda/std/__type_traits/is_same.h>
+
+#include <nv/target>
+
+#include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+//! @brief __is_local_smem_barrier returns true if barrier is (1) block-scoped and (2) located in shared memory.
+template <thread_scope _Sco,
+          typename _CompF,
+          bool _Is_mbarrier = (_Sco == thread_scope_block)
+                           && ::cuda::std::is_same_v<_CompF, ::cuda::std::__empty_completion>>
+_CCCL_HOST_DEVICE_API inline bool __is_local_smem_barrier([[maybe_unused]] barrier<_Sco, _CompF>& __barrier)
+{
+  NV_IF_ELSE_TARGET(
+    NV_IS_DEVICE,
+    (return _Is_mbarrier && ::cuda::device::is_object_from(__barrier, ::cuda::device::address_space::shared);),
+    (return false;));
+}
+
+_CCCL_END_NAMESPACE_CUDA
+
+#include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CUDA___MEMCPY_ASYNC_IS_LOCAL_SMEM_BARRIER_H

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/memcpy_async.h

@@ -0,0 +1,179 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_MEMCPY_ASYNC_H_
+#define _CUDA___MEMCPY_ASYNC_MEMCPY_ASYNC_H_
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#if _CCCL_CUDA_COMPILATION()
+
+#  include <cuda/__barrier/async_contract_fulfillment.h>
+#  include <cuda/__barrier/barrier.h>
+#  include <cuda/__barrier/barrier_block_scope.h>
+#  include <cuda/__barrier/barrier_thread_scope.h>
+#  include <cuda/__memcpy_async/check_preconditions.h>
+#  include <cuda/__memcpy_async/memcpy_async_barrier.h>
+#  include <cuda/__memory/aligned_size.h>
+#  include <cuda/std/__atomic/scopes.h>
+#  include <cuda/std/__type_traits/void_t.h>
+#  include <cuda/std/cstddef>
+#  include <cuda/std/cstdint>
+
+#  include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+/***********************************************************************
+ * memcpy_async code:
+ *
+ * A call to cuda::memcpy_async(dest, src, size, barrier) can dispatch to any of
+ * these PTX instructions:
+ *
+ * 1. normal synchronous copy (fallback)
+ * 2. cp.async:      shared  <- global
+ * 3. cp.async.bulk: shared  <- global
+ * 4. TODO: cp.async.bulk: global  <- shared
+ * 5. TODO: cp.async.bulk: cluster <- shared
+ *
+ * Which of these options is chosen, depends on:
+ *
+ * 1. The alignment of dest, src, and size;
+ * 2. The direction of the copy
+ * 3. The current compute capability
+ * 4. The requested completion mechanism
+ *
+ * PTX has 3 asynchronous completion mechanisms:
+ *
+ * 1. Async group           - local to a thread. Used by cp.async
+ * 2. Bulk async group      - local to a thread. Used by cp.async.bulk (shared -> global)
+ * 3. mbarrier::complete_tx - shared memory barier. Used by cp.async.bulk (other directions)
+ *
+ * The code is organized as follows:
+ *
+ * 1. Asynchronous copy mechanisms that wrap the PTX instructions
+ *
+ * 2. Device memcpy_async implementation per copy direction (global to shared,
+ *    shared to global, etc). Dispatches to fastest mechanism based on requested
+ *    completion mechanism(s), pointer alignment, and architecture.
+ *
+ * 3. Host and device memcpy_async implementations. Host implementation is
+ *    basically a memcpy wrapper; device implementation dispatches based on the
+ *    direction of the copy.
+ *
+ * 4. __memcpy_async_barrier:
+ *    a) Sets the allowed completion mechanisms based on the barrier location
+ *    b) Calls the host or device memcpy_async implementation
+ *    c) If necessary, synchronizes with the barrier based on the returned
+ *    completion mechanism.
+ *
+ * 5. The public memcpy_async function overloads. Call into
+ *    __memcpy_async_barrier.
+ *
+ ***********************************************************************/
+
+/***********************************************************************
+ * Asynchronous copy mechanisms:
+ *
+ * 1. cp.async.bulk: shared  <- global
+ * 2. TODO: cp.async.bulk: cluster <- shared
+ * 3. TODO: cp.async.bulk: global  <- shared
+ * 4. cp.async:      shared  <- global
+ * 5. normal synchronous copy (fallback)
+ ***********************************************************************/
+
+template <typename _Group, class _Tp, ::cuda::std::size_t _Alignment, thread_scope _Sco, typename _CompF>
+_CCCL_HOST_DEVICE_API inline async_contract_fulfillment memcpy_async(
+  _Group const& __group,
+  _Tp* __destination,
+  _Tp const* __source,
+  aligned_size_t<_Alignment> __size,
+  barrier<_Sco, _CompF>& __barrier)
+{
+  static_assert(_Alignment >= alignof(_Tp), "alignment must be at least the alignof(T)");
+  _CCCL_ASSERT(::cuda::__memcpy_async_check_pre(__destination, __source, __size), "memcpy_async preconditions unmet");
+  return ::cuda::__memcpy_async_barrier(__group, __destination, __source, __size, __barrier);
+}
+
+template <class _Tp, typename _Size, thread_scope _Sco, typename _CompF>
+_CCCL_HOST_DEVICE_API inline async_contract_fulfillment
+memcpy_async(_Tp* __destination, _Tp const* __source, _Size __size, barrier<_Sco, _CompF>& __barrier)
+{
+  _CCCL_ASSERT(::cuda::__memcpy_async_check_pre(__destination, __source, __size), "memcpy_async preconditions unmet");
+  return ::cuda::__memcpy_async_barrier(__single_thread_group{}, __destination, __source, __size, __barrier);
+}
+
+template <typename _Group, class _Tp, thread_scope _Sco, typename _CompF>
+_CCCL_HOST_DEVICE_API inline async_contract_fulfillment memcpy_async(
+  _Group const& __group,
+  _Tp* __destination,
+  _Tp const* __source,
+  ::cuda::std::size_t __size,
+  barrier<_Sco, _CompF>& __barrier)
+{
+  _CCCL_ASSERT(::cuda::__memcpy_async_check_pre(__destination, __source, __size), "memcpy_async preconditions unmet");
+  return ::cuda::__memcpy_async_barrier(__group, __destination, __source, __size, __barrier);
+}
+
+template <typename _Group, thread_scope _Sco, typename _CompF>
+_CCCL_HOST_DEVICE_API inline async_contract_fulfillment memcpy_async(
+  _Group const& __group,
+  void* __destination,
+  void const* __source,
+  ::cuda::std::size_t __size,
+  barrier<_Sco, _CompF>& __barrier)
+{
+  _CCCL_ASSERT(::cuda::__memcpy_async_check_pre(__destination, __source, __size), "memcpy_async preconditions unmet");
+  return ::cuda::__memcpy_async_barrier(
+    __group, reinterpret_cast<char*>(__destination), reinterpret_cast<char const*>(__source), __size, __barrier);
+}
+
+template <typename _Group, ::cuda::std::size_t _Alignment, thread_scope _Sco, typename _CompF>
+_CCCL_HOST_DEVICE_API inline async_contract_fulfillment memcpy_async(
+  _Group const& __group,
+  void* __destination,
+  void const* __source,
+  aligned_size_t<_Alignment> __size,
+  barrier<_Sco, _CompF>& __barrier)
+{
+  _CCCL_ASSERT(::cuda::__memcpy_async_check_pre(__destination, __source, __size), "memcpy_async preconditions unmet");
+  return ::cuda::__memcpy_async_barrier(
+    __group, reinterpret_cast<char*>(__destination), reinterpret_cast<char const*>(__source), __size, __barrier);
+}
+
+template <typename _Size, thread_scope _Sco, typename _CompF>
+_CCCL_HOST_DEVICE_API inline async_contract_fulfillment
+memcpy_async(void* __destination, void const* __source, _Size __size, barrier<_Sco, _CompF>& __barrier)
+{
+  _CCCL_ASSERT(::cuda::__memcpy_async_check_pre(__destination, __source, __size), "memcpy_async preconditions unmet");
+  return ::cuda::__memcpy_async_barrier(
+    __single_thread_group{},
+    reinterpret_cast<char*>(__destination),
+    reinterpret_cast<char const*>(__source),
+    __size,
+    __barrier);
+}
+
+_CCCL_END_NAMESPACE_CUDA
+
+#  include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CCCL_CUDA_COMPILATION()
+
+#endif // _CUDA___MEMCPY_ASYNC_MEMCPY_ASYNC_H_

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/memcpy_async_barrier.h

@@ -0,0 +1,99 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_MEMCPY_ASYNC_BARRIER_H_
+#define _CUDA___MEMCPY_ASYNC_MEMCPY_ASYNC_BARRIER_H_
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/__barrier/barrier.h>
+#include <cuda/__barrier/barrier_block_scope.h>
+#include <cuda/__barrier/barrier_thread_scope.h>
+#include <cuda/__memcpy_async/completion_mechanism.h>
+#include <cuda/__memcpy_async/dispatch_memcpy_async.h>
+#include <cuda/__memcpy_async/is_local_smem_barrier.h>
+#include <cuda/__memcpy_async/memcpy_completion.h>
+#include <cuda/__memcpy_async/try_get_barrier_handle.h>
+#include <cuda/__memory/aligned_size.h>
+#include <cuda/__type_traits/is_trivially_copyable.h>
+#include <cuda/std/__algorithm/max.h>
+#include <cuda/std/__atomic/scopes.h>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+
+#include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+struct __single_thread_group
+{
+  _CCCL_HOST_DEVICE_API inline void sync() const {}
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::size_t size() const
+  {
+    return 1;
+  };
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::size_t thread_rank() const
+  {
+    return 0;
+  };
+};
+
+template <typename _Group, class _Tp, typename _Size, thread_scope _Sco, typename _CompF>
+_CCCL_HOST_DEVICE_API inline async_contract_fulfillment __memcpy_async_barrier(
+  _Group const& __group, _Tp* __destination, _Tp const* __source, _Size __size, barrier<_Sco, _CompF>& __barrier)
+{
+  static_assert(::cuda::is_trivially_copyable_v<_Tp>, "memcpy_async requires a trivially copyable type");
+
+  // 1. Determine which completion mechanisms can be used with the current
+  // barrier. A local shared memory barrier, i.e., block-scope barrier in local
+  // shared memory, supports the mbarrier_complete_tx mechanism in addition to
+  // the async group mechanism.
+  ::cuda::std::uint32_t __allowed_completions =
+    ::cuda::__is_local_smem_barrier(__barrier)
+      ? (::cuda::std::uint32_t(__completion_mechanism::__async_group)
+         | ::cuda::std::uint32_t(__completion_mechanism::__mbarrier_complete_tx))
+      : ::cuda::std::uint32_t(__completion_mechanism::__async_group);
+
+  // Alignment: Use the maximum of the alignment of _Tp and that of a possible cuda::aligned_size_t.
+  constexpr auto __align = ::cuda::std::max(alignof(_Tp), __get_size_align_v<_Size>);
+  // Cast to char pointers. We don't need the type for alignment anymore and
+  // erasing the types reduces the number of instantiations of down-stream
+  // functions.
+  char* __dest_char      = reinterpret_cast<char*>(__destination);
+  char const* __src_char = reinterpret_cast<char const*>(__source);
+
+  // 2. Issue actual copy instructions.
+  ::cuda::std::uint64_t* __bh = nullptr;
+#if __cccl_ptx_isa >= 800
+  NV_IF_TARGET(
+    NV_PROVIDES_SM_90,
+    (__bh = ::cuda::__is_local_smem_barrier(__barrier) ? ::cuda::__try_get_barrier_handle(__barrier) : nullptr;))
+#endif // __cccl_ptx_isa >= 800
+  auto __cm =
+    ::cuda::__dispatch_memcpy_async<__align>(__group, __dest_char, __src_char, __size, __allowed_completions, __bh);
+
+  // 3. Synchronize barrier with copy instructions.
+  return __memcpy_completion_impl::__defer(__cm, __group, __size, __barrier);
+}
+
+_CCCL_END_NAMESPACE_CUDA
+
+#include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CUDA___MEMCPY_ASYNC_MEMCPY_ASYNC_BARRIER_H_

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/memcpy_async_tx.h

@@ -0,0 +1,102 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_MEMCPY_ASYNC_TX_H_
+#define _CUDA___MEMCPY_ASYNC_MEMCPY_ASYNC_TX_H_
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#if _CCCL_CUDA_COMPILATION()
+#  if __cccl_ptx_isa >= 800
+
+#    include <cuda/__barrier/async_contract_fulfillment.h>
+#    include <cuda/__barrier/barrier_block_scope.h>
+#    include <cuda/__memcpy_async/check_preconditions.h>
+#    include <cuda/__memory/address_space.h>
+#    include <cuda/__memory/aligned_size.h>
+#    include <cuda/__ptx/instructions/cp_async_bulk.h>
+#    include <cuda/__ptx/ptx_dot_variants.h>
+#    include <cuda/__ptx/ptx_helper_functions.h>
+#    include <cuda/__type_traits/is_trivially_copyable.h>
+#    include <cuda/std/__atomic/scopes.h>
+#    include <cuda/std/__type_traits/conditional.h>
+#    include <cuda/std/cstdint>
+
+#    include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA_DEVICE
+
+extern "C" _CCCL_DEVICE void __cuda_ptx_memcpy_async_tx_is_not_supported_before_SM_90__();
+template <typename _Tp, ::cuda::std::size_t _Alignment>
+_CCCL_DEVICE_API async_contract_fulfillment memcpy_async_tx(
+  _Tp* __dest,
+  const _Tp* __src,
+  ::cuda::aligned_size_t<_Alignment> __size,
+  ::cuda::barrier<::cuda::thread_scope_block>& __b)
+{
+  // When compiling with NVCC and GCC 4.8, certain user defined types that _are_ trivially copyable are
+  // incorrectly classified as not trivially copyable. Remove this assertion to allow for their usage with
+  // memcpy_async when compiling with GCC 4.8.
+  // FIXME: remove the #if once GCC 4.8 is no longer supported.
+#    if !_CCCL_COMPILER(GCC) || _CCCL_COMPILER(GCC, >, 4, 8)
+  static_assert(::cuda::is_trivially_copyable_v<_Tp>, "memcpy_async_tx requires a trivially copyable type");
+#    endif
+  static_assert(16 <= _Alignment, "mempcy_async_tx expects arguments to be at least 16 byte aligned.");
+  static_assert(_Alignment >= alignof(_Tp), "alignment must be at least the alignof(T)");
+
+  _CCCL_ASSERT(::cuda::__memcpy_async_check_pre(__dest, __src, __size), "memcpy_async_tx preconditions unmet");
+
+  _CCCL_ASSERT(
+    ::cuda::device::is_address_from(::cuda::device::barrier_native_handle(__b), ::cuda::device::address_space::shared),
+    "Barrier must be located in local shared memory.");
+  _CCCL_ASSERT(::cuda::device::is_address_from(__dest, ::cuda::device::address_space::shared),
+               "dest must point to shared memory.");
+  _CCCL_ASSERT(::cuda::device::is_address_from(__src, ::cuda::device::address_space::global),
+               "src must point to global memory.");
+
+  NV_IF_ELSE_TARGET(
+    NV_PROVIDES_SM_90,
+    (
+      if (::cuda::device::is_address_from(__dest, ::cuda::device::address_space::shared)
+          && ::cuda::device::is_address_from(__src, ::cuda::device::address_space::global)) {
+        ::cuda::ptx::cp_async_bulk(
+          ::cuda::std::conditional_t<__cccl_ptx_isa >= 860, ::cuda::ptx::space_shared_t, ::cuda::ptx::space_cluster_t>{},
+          ::cuda::ptx::space_global,
+          __dest,
+          __src,
+          static_cast<uint32_t>(__size),
+          ::cuda::device::barrier_native_handle(__b));
+      } else {
+        _CCCL_VERIFY(false,
+                     "memcpy_async_tx only supports copying from global to shared or from shared to remote cluster "
+                     "dsmem. To copy to remote is not yet implemented.");
+      }),
+    (::cuda::device::__cuda_ptx_memcpy_async_tx_is_not_supported_before_SM_90__();));
+
+  return async_contract_fulfillment::async;
+}
+
+_CCCL_END_NAMESPACE_CUDA_DEVICE
+
+#    include <cuda/std/__cccl/epilogue.h>
+
+#  endif // __cccl_ptx_isa >= 800
+#endif // _CCCL_CUDA_COMPILATION()
+
+#endif // _CUDA___MEMCPY_ASYNC_MEMCPY_ASYNC_TX_H_

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/memcpy_completion.h

@@ -0,0 +1,164 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_MEMCPY_COMPLETION_H
+#define _CUDA___MEMCPY_ASYNC_MEMCPY_COMPLETION_H
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/__barrier/async_contract_fulfillment.h>
+#include <cuda/__barrier/barrier_block_scope.h>
+#include <cuda/__barrier/barrier_expect_tx.h>
+#include <cuda/__fwd/pipeline.h>
+#include <cuda/__memcpy_async/completion_mechanism.h>
+#include <cuda/__memcpy_async/is_local_smem_barrier.h>
+#include <cuda/__memcpy_async/try_get_barrier_handle.h>
+#include <cuda/std/__atomic/scopes.h>
+#include <cuda/std/cstdint>
+
+#if _CCCL_CUDA_COMPILATION()
+#  include <cuda/__ptx/ptx_dot_variants.h>
+#  include <cuda/__ptx/ptx_helper_functions.h>
+#endif // _CCCL_CUDA_COMPILATION()
+
+#include <nv/target>
+
+#include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+// This struct contains functions to defer the completion of a barrier phase
+// or pipeline stage until a specific memcpy_async operation *initiated by
+// this thread* has completed.
+
+// The user is still responsible for arriving and waiting on (or otherwise
+// synchronizing with) the barrier or pipeline barrier to see the results of
+// copies from other threads participating in the synchronization object.
+struct __memcpy_completion_impl
+{
+  template <typename _Group>
+  [[nodiscard]] _CCCL_HOST_DEVICE_API inline static async_contract_fulfillment
+  __defer(__completion_mechanism __cm,
+          _Group const& __group,
+          ::cuda::std::size_t __size,
+          barrier<::cuda::thread_scope_block>& __barrier)
+  {
+    // In principle, this is the overload for shared memory barriers. However, a
+    // block-scope barrier may also be located in global memory. Therefore, we
+    // check if the barrier is a non-smem barrier and handle that separately.
+    if (!::cuda::__is_local_smem_barrier(__barrier))
+    {
+      return __defer_non_smem_barrier(__cm, __group, __size, __barrier);
+    }
+
+    switch (__cm)
+    {
+      case __completion_mechanism::__async_group:
+        // Pre-SM80, the async_group mechanism is not available.
+        NV_IF_TARGET(
+          NV_PROVIDES_SM_80,
+          (
+            // Non-Blocking: unbalance barrier by 1, barrier will be
+            // rebalanced when all thread-local cp.async instructions
+            // have completed writing to shared memory.
+            ::cuda::std::uint64_t* __bh = ::cuda::__try_get_barrier_handle(__barrier);
+
+            asm volatile("cp.async.mbarrier.arrive.shared.b64 [%0];" ::"r"(
+              static_cast<::cuda::std::uint32_t>(::__cvta_generic_to_shared(__bh))) : "memory");));
+        return async_contract_fulfillment::async;
+      case __completion_mechanism::__async_bulk_group:
+        // This completion mechanism should not be used with a shared
+        // memory barrier. Or at least, we do not currently envision
+        // bulk group to be used with shared memory barriers.
+        _CCCL_UNREACHABLE();
+      case __completion_mechanism::__mbarrier_complete_tx:
+        // we already updated the mbarrier's tx count when we issued the bulk copy
+        return async_contract_fulfillment::async;
+      case __completion_mechanism::__sync:
+        // sync: In this case, we do not need to do anything. The user will have
+        // to issue `bar.arrive_wait();` to see the effect of the transaction.
+        return async_contract_fulfillment::none;
+      default:
+        // Get rid of "control reaches end of non-void function":
+        _CCCL_UNREACHABLE();
+    }
+  }
+
+  template <typename _Group, thread_scope _Sco, typename _CompF>
+  [[nodiscard]] _CCCL_HOST_DEVICE_API inline static async_contract_fulfillment __defer(
+    __completion_mechanism __cm, _Group const& __group, ::cuda::std::size_t __size, barrier<_Sco, _CompF>& __barrier)
+  {
+    return __defer_non_smem_barrier(__cm, __group, __size, __barrier);
+  }
+
+  template <typename _Group, thread_scope _Sco, typename _CompF>
+  [[nodiscard]] _CCCL_HOST_DEVICE_API inline static async_contract_fulfillment
+  __defer_non_smem_barrier(__completion_mechanism __cm, _Group const&, ::cuda::std::size_t, barrier<_Sco, _CompF>&)
+  {
+    // Overload for non-smem barriers.
+    switch (__cm)
+    {
+      case __completion_mechanism::__async_group:
+        // Pre-SM80, the async_group mechanism is not available.
+        NV_IF_TARGET(NV_PROVIDES_SM_80,
+                     (
+                       // Blocking: wait for all thread-local cp.async instructions to have
+                       // completed writing to shared memory.
+                       asm volatile("cp.async.wait_all;" :: : "memory");));
+        return async_contract_fulfillment::async;
+      case __completion_mechanism::__async_bulk_group:
+        [[fallthrough]];
+        // This completion mechanism is currently not expected to be used with barriers.
+      case __completion_mechanism::__mbarrier_complete_tx:
+        // Non-smem barriers do not have an mbarrier_complete_tx mechanism.
+        _CCCL_UNREACHABLE();
+      case __completion_mechanism::__sync:
+        // sync: In this case, we do not need to do anything.
+        return async_contract_fulfillment::none;
+      default:
+        // Get rid of "control reaches end of non-void function":
+        _CCCL_UNREACHABLE();
+    }
+  }
+
+  template <typename _Group, thread_scope _Sco>
+  [[nodiscard]] _CCCL_HOST_DEVICE_API inline static async_contract_fulfillment
+  __defer(__completion_mechanism __cm, _Group const&, ::cuda::std::size_t, pipeline<_Sco>&)
+  {
+    switch (__cm)
+    {
+      case __completion_mechanism::__async_group:
+        [[fallthrough]];
+      case __completion_mechanism::__async_bulk_group:
+        [[fallthrough]];
+      case __completion_mechanism::__mbarrier_complete_tx:
+        return async_contract_fulfillment::async;
+      case __completion_mechanism::__sync:
+        return async_contract_fulfillment::none;
+      default:
+        // Get rid of "control reaches end of non-void function":
+        _CCCL_UNREACHABLE();
+    }
+  }
+};
+
+_CCCL_END_NAMESPACE_CUDA
+
+#include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CUDA___MEMCPY_ASYNC_MEMCPY_COMPLETION_H

cccl_upstream/libcudacxx/include/cuda/__memcpy_async/try_get_barrier_handle.h

@@ -0,0 +1,57 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of libcu++, the C++ Standard Library for your entire system,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _CUDA___MEMCPY_ASYNC_TRY_GET_BARRIER_HANDLE_H
+#define _CUDA___MEMCPY_ASYNC_TRY_GET_BARRIER_HANDLE_H
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/__barrier/barrier_block_scope.h>
+#include <cuda/std/__atomic/scopes.h>
+#include <cuda/std/__barrier/barrier.h>
+#include <cuda/std/__barrier/empty_completion.h>
+#include <cuda/std/__type_traits/is_same.h>
+#include <cuda/std/cstdint>
+
+#include <nv/target>
+
+#include <cuda/std/__cccl/prologue.h>
+
+_CCCL_BEGIN_NAMESPACE_CUDA
+
+//! @brief __try_get_barrier_handle returns barrier handle of block-scoped barriers and a nullptr otherwise.
+template <thread_scope _Sco, typename _CompF>
+_CCCL_HOST_DEVICE_API inline ::cuda::std::uint64_t* __try_get_barrier_handle(barrier<_Sco, _CompF>&)
+{
+  return nullptr;
+}
+
+template <>
+_CCCL_HOST_DEVICE_API inline ::cuda::std::uint64_t*
+__try_get_barrier_handle<::cuda::thread_scope_block, ::cuda::std::__empty_completion>(
+  [[maybe_unused]] barrier<thread_scope_block>& __barrier)
+{
+  NV_DISPATCH_TARGET(
+    NV_IS_DEVICE, (return ::cuda::device::barrier_native_handle(__barrier);), NV_ANY_TARGET, (return nullptr;));
+}
+
+_CCCL_END_NAMESPACE_CUDA
+
+#include <cuda/std/__cccl/epilogue.h>
+
+#endif // _CUDA___MEMCPY_ASYNC_TRY_GET_BARRIER_HANDLE_H