execution_context.hpp

/*
 * SPDX-FileCopyrightText: Copyright (c) 2026 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.
 */

#pragma once

#include <common.hpp>
#include <nve_types.hpp>
#include <cuda_support.hpp>
#include <unordered_map>
#include <memory>
#include <mutex>
#include <vector>

namespace nve {

class ResizeableBuffer;

// Internal class to share functionality across different execution contexts
// This is not part of the external API, applications should use Layer/Table::create_context() instead
class ExecutionContext {
 public:
  NVE_PREVENT_COPY_AND_MOVE_(ExecutionContext);
  virtual ~ExecutionContext();

  // Getters
  inline cudaStream_t get_lookup_stream() const { return lookup_stream_; };
  inline cudaStream_t get_modify_stream() const { return modify_stream_; };
  inline thread_pool_ptr_t get_thread_pool() const { return thread_pool_; }
  inline allocator_ptr_t get_allocator() const { return allocator_; }

  // Get a temporary buffer. Storage lookup and resizing are synchronized, but the caller
  // is responsible to not call get_buffer again with the same name and a larger size
  // (potentially triggering a realloc) before the work on the buffer is done.
  void* get_buffer(const std::string& name, size_t size, bool host_alloc);

  // get aux streams
  virtual std::vector<cudaStream_t> get_aux_streams(const std::string& name, size_t num_streams);

  // Wait until pending work is complete
  // Note that this may include additional tasks offloaded to other threads
  //
  // CUDA stream sync is gated on driver_available_: on a system with no CUDA
  // driver (host-only inference) the runtime calls would fail, so we skip them.
  // The flag is read directly off a base member rather than via a virtual, so it
  // stays correct even when wait() is invoked from ~ExecutionContext() (where
  // virtual dispatch resolves to this base implementation, not a derived override).
  virtual void wait();

 protected:
  // using nullptr for threadpool/allocator implies use the default one.
  ExecutionContext(
    cudaStream_t lookup_stream,
    cudaStream_t modify_stream,
    thread_pool_ptr_t thread_pool,
    allocator_ptr_t allocator);

  cudaStream_t lookup_stream_;
  cudaStream_t modify_stream_;
  thread_pool_ptr_t thread_pool_;
  allocator_ptr_t allocator_;
  // Whether a usable CUDA driver is present in this process (detected via cuInit).
  // When false, the context must not enter the CUDA runtime during wait()/teardown.
  const bool driver_available_;
  std::mutex buffer_storage_mutex_;
  std::unordered_map<std::string, std::shared_ptr<ResizeableBuffer>> buffer_storage_;
  std::mutex aux_streams_mutex_;
  std::unordered_map<std::string, std::vector<cudaStream_t>> aux_streams_storage_;
  static std::string internal_name(const std::string& name, bool host_alloc);
  std::vector<cudaStream_t> snapshot_aux_streams();
};

}  // namespace nve