codekingpro/portable-devtools

1//===----------------------------------------------------------------------===//2//3// Part of libcu++, the C++ Standard Library for your entire system,4// under the Apache License v2.0 with LLVM Exceptions.5// See https://llvm.org/LICENSE.txt for license information.6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception7// SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES.8//9//===----------------------------------------------------------------------===//10 11#ifndef _CUDA_STREAM_REF12#define _CUDA_STREAM_REF13 14// clang-format off15/*16    stream_ref synopsis17namespace cuda {18class stream_ref {19    using value_type = cudaStream_t;20 21    stream_ref() = default;22    stream_ref(cudaStream_t stream_) noexcept : stream(stream_) {}23 24    stream_ref(int) = delete;25    stream_ref(nullptr_t) = delete;26 27    [[nodiscard]] value_type get() const noexcept;28 29    void wait() const;30 31    [[nodiscard]] bool ready() const;32 33    [[nodiscard]] friend bool operator==(stream_ref, stream_ref);34    [[nodiscard]] friend bool operator!=(stream_ref, stream_ref);35 36private:37  cudaStream_t stream = 0; // exposition only38};39}  // cuda40*/41 42#ifdef LIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE43 44#include <cuda_runtime_api.h> // cuda_runtime_api needs to come first45// clang-format on46 47#include <cuda/std/detail/__config>48 49#include <cuda/std/detail/__pragma_push>50 51#include <cuda/std/array>52#include <cuda/std/type_traits>53 54_LIBCUDACXX_BEGIN_NAMESPACE_CUDA55 56/**57 * \brief A non-owning wrapper for a `cudaStream_t`.58 *59 * `stream_view` is a non-owning "view" type similar to `std::span` or60 * `std::string_view`. \see https://en.cppreference.com/w/cpp/container/span and61 * \see https://en.cppreference.com/w/cpp/string/basic_string_view62 *63 */64class stream_ref65{66private:67  ::cudaStream_t __stream{0};68 69public:70  using value_type = ::cudaStream_t;71 72  /**73   * \brief Constructs a `stream_view` of the "default" CUDA stream.74   *75   * For behavior of the default stream,76   * \see77   * https://docs.nvidia.com/cuda/cuda-runtime-api/stream-sync-behavior.html78   *79   */80  stream_ref() = default;81 82  /**83   * \brief Constructs a `stream_view` from a `cudaStream_t` handle.84   *85   * This constructor provides implicit conversion from `cudaStream_t`.86   *87   * \note: It is the callers responsibilty to ensure the `stream_view` does not88   * outlive the stream identified by the `cudaStream_t` handle.89   *90   */91  constexpr stream_ref(value_type __stream_) noexcept92      : __stream{__stream_}93  {}94 95  /// Disallow construction from an `int`, e.g., `0`.96  stream_ref(int) = delete;97 98  /// Disallow construction from `nullptr`.99  stream_ref(_CUDA_VSTD::nullptr_t) = delete;100 101  /**102   * \brief Compares two `stream_view`s for equality103   *104   * \note Allows comparison with `cudaStream_t` due to implicit conversion to105   * `stream_view`.106   *107   * \param lhs The first `stream_view` to compare108   * \param rhs The second `stream_view` to compare109   * \return true if equal, false if unequal110   */111  _LIBCUDACXX_NODISCARD_FRIEND constexpr bool operator==(const stream_ref& __lhs, const stream_ref& __rhs) noexcept112  {113    return __lhs.__stream == __rhs.__stream;114  }115 116  /**117   * \brief Compares two `stream_view`s for inequality118   *119   * \note Allows comparison with `cudaStream_t` due to implicit conversion to120   * `stream_view`.121   *122   * \param lhs The first `stream_view` to compare123   * \param rhs The second `stream_view` to compare124   * \return true if unequal, false if equal125   */126  _LIBCUDACXX_NODISCARD_FRIEND constexpr bool operator!=(const stream_ref& __lhs, const stream_ref& __rhs) noexcept127  {128    return __lhs.__stream != __rhs.__stream;129  }130 131  /// Returns the wrapped `cudaStream_t` handle.132  _LIBCUDACXX_NODISCARD_ATTRIBUTE constexpr value_type get() const noexcept { return __stream; }133 134  /**135   * \brief Synchronizes the wrapped stream.136   *137   * \throws cuda::cuda_error if synchronization fails.138   *139   */140  void wait() const141  {142    const auto __result = ::cudaStreamSynchronize(get());143    switch (__result)144    {145      case ::cudaSuccess:146        return;147      default:148        ::cudaGetLastError(); // Clear CUDA error state149#ifndef _LIBCUDACXX_NO_EXCEPTIONS150        throw cuda::cuda_error{__result, "Failed to synchronize stream."};151#else152        _LIBCUDACXX_UNREACHABLE();153#endif154    }155  }156 157  /**158   * \brief Queries if all operations on the wrapped stream have completed.159   *160   * \throws cuda::cuda_error if the query fails.161   *162   * \return `true` if all operations have completed, or `false` if not.163   */164  _LIBCUDACXX_NODISCARD_ATTRIBUTE bool ready() const165  {166    const auto __result = ::cudaStreamQuery(get());167    switch (__result)168    {169      case ::cudaSuccess:170        return true;171      case ::cudaErrorNotReady:172        return false;173      default:174        ::cudaGetLastError(); // Clear CUDA error state175#ifndef _LIBCUDACXX_NO_EXCEPTIONS176        throw cuda::cuda_error{__result, ""};177#else178        // _LIBCUDACXX_UNREACHABLE();179#endif180        return false;181    }182  }183};184 185_LIBCUDACXX_END_NAMESPACE_CUDA186 187#include <cuda/std/detail/__pragma_pop>188 189#endif // LIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE190 191#endif //_CUDA_STREAM_REF192