Program Listing for File gxf_cuda.hpp

            

            
/*
* SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

#ifndef HOLOSCAN_CORE_GXF_GXF_CUDA_HPP
#define HOLOSCAN_CORE_GXF_GXF_CUDA_HPP

#include <gxf/core/gxf.h>

#include <memory>
#include <string>
#include <unordered_map>
#include <vector>

#include "../operator_spec.hpp"
#include "../parameter.hpp"
#include "../resources/gxf/cuda_stream_pool.hpp"
#include "gxf/cuda/cuda_stream.hpp"
#include "gxf/cuda/cuda_stream_id.hpp"
#include "gxf/cuda/cuda_stream_pool.hpp"

namespace holoscan::gxf {

using CudaStreamId = nvidia::gxf::CudaStreamId;
using CudaStream = nvidia::gxf::CudaStream;
using CudaStreamHandle = nvidia::gxf::Handle<nvidia::gxf::CudaStream>;

class CudaObjectHandler {
 public:
  ~CudaObjectHandler();

  void init_from_operator(Operator* op);

  gxf_result_t add_stream(const CudaStreamHandle& stream_handle,
                          const std::string& output_port_name);

  gxf_result_t add_stream(const cudaStream_t stream, const std::string& output_port_name);

  expected<CudaStreamHandle, RuntimeError> get_cuda_stream_handle(
      gxf_context_t context, const std::string& input_port_name, bool allocate = true,
      bool sync_to_default = false);

  expected<std::vector<std::optional<CudaStreamHandle>>, RuntimeError> get_cuda_stream_handles(
      gxf_context_t context, const std::string& input_port_name);

  cudaStream_t get_cuda_stream(gxf_context_t context, const std::string& input_port_name,
                               bool allocate = false, bool sync_to_default = true);

  std::vector<std::optional<cudaStream_t>> get_cuda_streams(gxf_context_t context,
                                                            const std::string& input_port_name);

  gxf_result_t synchronize_streams(std::vector<std::optional<CudaStreamHandle>> stream_handles,
                                   CudaStreamHandle target_stream_handle,
                                   bool sync_to_default_stream = true);

  gxf_result_t synchronize_streams(std::vector<cudaStream_t> cuda_streams,
                                   cudaStream_t target_stream, bool sync_to_default_stream = true);

  cudaStream_t stream_from_stream_handle(CudaStreamHandle stream_handle);

  expected<CudaStreamHandle, RuntimeError> stream_handle_from_stream(cudaStream_t stream);

  expected<gxf_uid_t, ErrorCode> get_output_stream_cid(const std::string& output_port_name);

  gxf_result_t streams_from_message(gxf_context_t context, const nvidia::gxf::Entity& message,
                                    const std::string& input_name);

  expected<CudaStreamHandle, RuntimeError> allocate_internal_stream(gxf_context_t context,
                                                                    const std::string& stream_name);

  gxf_result_t release_internal_streams(gxf_context_t context);

  void clear_received_streams();

 private:
  expected<CudaStreamHandle, RuntimeError> allocate_cuda_stream(gxf_context_t context);

  gxf_result_t from_messages(gxf_context_t context, size_t message_count,
                             const nvidia::gxf::Entity* messages);

  Parameter<std::shared_ptr<CudaStreamPool>> cuda_stream_pool_;

  expected<nvidia::gxf::Handle<nvidia::gxf::CudaStreamPool>, RuntimeError> cuda_stream_pool_handle(
      gxf_context_t context);

  bool default_stream_warning_ = false;

  cudaEvent_t cuda_event_ = 0;

  bool event_created_ = false;

  std::unordered_map<std::string, std::vector<std::optional<CudaStreamId>>>
      received_cuda_stream_ids_{};
  std::unordered_map<std::string, std::vector<std::optional<CudaStreamHandle>>>
      received_cuda_stream_handles_{};

  std::unordered_map<std::string, CudaStreamHandle> allocated_cuda_stream_handles_{};

  std::unordered_map<std::string, gxf_uid_t> emitted_cuda_stream_cids_{};

  // If we want to allow emitting via a cudaStream_t, we need a way to get the handle it
  // corresponds to.
  std::unordered_map<cudaStream_t, CudaStreamHandle> stream_to_stream_handle_{};
};

}  // namespace holoscan::gxf

#endif/* HOLOSCAN_CORE_GXF_GXF_CUDA_HPP */