stream.h¶

Namespaces¶

Name
tenzr
tenzr::execution

Classes¶

	Name
struct	tenzr::execution::StreamEventState Shared readiness state for a stream event, holding the ready flag, the PJRT event, and completion callbacks.
class	tenzr::execution::StreamEvent Handle to a recorded point in a stream's work, used to query completion or to make later work wait.
class	tenzr::execution::ExecutionStream Serialized work queue backed by a worker thread, used to submit and synchronize device work.

Source code¶

// 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.

#ifndef TENZR_EXECUTION_STREAM_H_
#define TENZR_EXECUTION_STREAM_H_

#include <xla/pjrt/c/pjrt_c_api.h>

#include <atomic>
#include <condition_variable>
#include <cstdint>
#include <exception>
#include <functional>
#include <future>
#include <memory>
#include <mutex>
#include <queue>
#include <thread>
#include <type_traits>
#include <vector>

#include "tenzr/common/error.h"

namespace tenzr::execution {

  class ExecutionStream;

  struct StreamEventState {
    std::atomic<bool> ready{false};

    std::mutex mutex;

    std::condition_variable cv;

    std::vector<std::function<void()>> callbacks;

    PJRT_Event* pjrt_event{nullptr};

    void SetReady();

    void Wait();

    bool IsReady() const;

    void AddCallback(std::function<void()> cb);
  };

  class StreamEvent {
   public:
    StreamEvent();

    explicit StreamEvent(PJRT_Event* pjrt_event);

    ~StreamEvent() = default;

    StreamEvent(const StreamEvent&) = default;
    StreamEvent& operator=(const StreamEvent&) = default;
    StreamEvent(StreamEvent&&) noexcept = default;
    StreamEvent& operator=(StreamEvent&&) noexcept = default;

    void Record(ExecutionStream& stream);

    bool Query() const;

    void Synchronize();

    std::shared_ptr<StreamEventState> State() const { return state_; }

   private:
    std::shared_ptr<StreamEventState> state_;
  };

  class ExecutionStream : public std::enable_shared_from_this<ExecutionStream> {
   public:
    explicit ExecutionStream(int stream_id = 0, int device_id = 0);

    ~ExecutionStream();

    ExecutionStream(const ExecutionStream&) = delete;
    ExecutionStream& operator=(const ExecutionStream&) = delete;

    int StreamId() const { return stream_id_; }

    int DeviceId() const { return device_id_; }

    void RecordEvent(StreamEvent& event);

    void WaitEvent(const StreamEvent& event);

    void Synchronize();

    void Submit(std::function<void()> task);

    template <typename F> auto AsyncSubmit(F&& task) {
      using Ret = std::invoke_result_t<F>;
      auto prom = std::make_shared<std::promise<Ret>>();
      auto fut = prom->get_future();

      Submit([prom, captured_task = std::forward<F>(task)]() mutable {
#if TENZR_EXCEPTIONS_ENABLED
        try {
          if constexpr (std::is_void_v<Ret>) {
            captured_task();
            prom->set_value();
          } else {
            prom->set_value(captured_task());
          }
        } catch (...) {
          prom->set_exception(std::current_exception());
        }
#else
        if constexpr (std::is_void_v<Ret>) {
          captured_task();
          prom->set_value();
        } else {
          prom->set_value(captured_task());
        }
#endif
      });

      return fut;
    }

    static ExecutionStream& DefaultStream(int device_id = 0);

    static std::shared_ptr<ExecutionStream> Create(int device_id = 0);

   private:
    void WorkerLoop();

    int stream_id_{0};
    int device_id_{0};
    std::atomic<bool> stop_{false};
    std::thread worker_;
    std::mutex queue_mutex_;
    std::condition_variable queue_cv_;
    std::condition_variable empty_cv_;
    std::queue<std::function<void()>> tasks_;
    size_t active_tasks_{0};
  };

}  // namespace tenzr::execution

#endif  // TENZR_EXECUTION_STREAM_H_