src/execution/stream.cc¶

// Copyright 2026 The Tenzr Project Authors
//
// 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
//
//      https://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.

#include "tenzr/execution/stream.h"

#include <mutex>
#include <unordered_map>
#include <utility>

namespace tenzr::execution {

  void StreamEventState::SetReady() {
    std::vector<std::function<void()>> cbs;
    {
      std::lock_guard<std::mutex> lock(mutex);
      ready.store(true, std::memory_order_release);
      cbs = std::move(callbacks);
      callbacks.clear();
    }
    cv.notify_all();
    for (auto& cb : cbs) {
      if (cb) {
        cb();
      }
    }
  }

  void StreamEventState::Wait() {
    if (ready.load(std::memory_order_acquire)) {
      return;
    }
    std::unique_lock<std::mutex> lock(mutex);
    cv.wait(lock, [this]() { return ready.load(std::memory_order_acquire); });
  }

  bool StreamEventState::IsReady() const { return ready.load(std::memory_order_acquire); }

  void StreamEventState::AddCallback(std::function<void()> cb) {
    bool call_now = false;
    {
      std::lock_guard<std::mutex> lock(mutex);
      if (ready.load(std::memory_order_acquire)) {
        call_now = true;
      } else {
        callbacks.push_back(std::move(cb));
      }
    }
    if (call_now && cb) {
      cb();
    }
  }

  // StreamEvent
  StreamEvent::StreamEvent() : state_(std::make_shared<StreamEventState>()) {}

  StreamEvent::StreamEvent(PJRT_Event* pjrt_event) : state_(std::make_shared<StreamEventState>()) { state_->pjrt_event = pjrt_event; }

  void StreamEvent::Record(ExecutionStream& stream) { stream.RecordEvent(*this); }

  bool StreamEvent::Query() const {
    TENZR_CHECK(state_ != nullptr, "Null StreamEvent state");
    return state_->IsReady();
  }

  void StreamEvent::Synchronize() {
    TENZR_CHECK(state_ != nullptr, "Null StreamEvent state");
    state_->Wait();
  }

  // ExecutionStream
  ExecutionStream::ExecutionStream(int stream_id, int device_id) : stream_id_(stream_id), device_id_(device_id) {
    worker_ = std::thread(&ExecutionStream::WorkerLoop, this);
  }

  ExecutionStream::~ExecutionStream() {
    {
      std::lock_guard<std::mutex> lock(queue_mutex_);
      stop_.store(true, std::memory_order_release);
    }
    queue_cv_.notify_all();
    if (worker_.joinable()) {
      worker_.join();
    }
  }

  void ExecutionStream::RecordEvent(StreamEvent& event) {
    auto state = event.State();
    TENZR_CHECK(state != nullptr, "Cannot record event with null state");
    Submit([state]() { state->SetReady(); });
  }

  void ExecutionStream::WaitEvent(const StreamEvent& event) {
    auto state = event.State();
    TENZR_CHECK(state != nullptr, "Cannot wait on event with null state");
    Submit([state]() { state->Wait(); });
  }

  void ExecutionStream::Synchronize() {
    std::unique_lock<std::mutex> lock(queue_mutex_);
    empty_cv_.wait(lock, [this]() { return tasks_.empty() && active_tasks_ == 0; });
  }

  void ExecutionStream::Submit(std::function<void()> task) {
    TENZR_CHECK(task != nullptr, "Cannot submit null task to stream");
    {
      std::lock_guard<std::mutex> lock(queue_mutex_);
      tasks_.push(std::move(task));
    }
    queue_cv_.notify_one();
  }

  void ExecutionStream::WorkerLoop() {
    while (true) {
      std::function<void()> task;
      {
        std::unique_lock<std::mutex> lock(queue_mutex_);
        queue_cv_.wait(lock, [this]() { return stop_.load(std::memory_order_acquire) || !tasks_.empty(); });

        if (stop_.load(std::memory_order_acquire) && tasks_.empty()) {
          break;
        }

        task = std::move(tasks_.front());
        tasks_.pop();
        ++active_tasks_;
      }

#if TENZR_EXCEPTIONS_ENABLED
      try {
        task();
      } catch (...) {
        // Preserve stream worker liveness on task exception
      }
#else
      task();
#endif

      {
        std::lock_guard<std::mutex> lock(queue_mutex_);
        --active_tasks_;
        if (tasks_.empty() && active_tasks_ == 0) {
          empty_cv_.notify_all();
        }
      }
    }
  }

  ExecutionStream& ExecutionStream::DefaultStream(int device_id) {
    static std::mutex s_mutex;
    static std::unordered_map<int, std::shared_ptr<ExecutionStream>> streams;
    std::lock_guard<std::mutex> lock(s_mutex);
    auto it = streams.find(device_id);
    if (it == streams.end()) {
      it = streams.emplace(device_id, std::make_shared<ExecutionStream>(0, device_id)).first;
    }
    return *it->second;
  }

  std::shared_ptr<ExecutionStream> ExecutionStream::Create(int device_id) {
    static std::atomic<int> next_stream_id{1};
    return std::make_shared<ExecutionStream>(next_stream_id.fetch_add(1), device_id);
  }

}  // namespace tenzr::execution