include/tenzr/execution/execution_engine.h¶

// 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_EXECUTION_ENGINE_H_
#define TENZR_EXECUTION_EXECUTION_ENGINE_H_

#include <exception>
#include <future>
#include <memory>
#include <vector>

#include "tenzr/execution/jit_cache.h"
#include "tenzr/execution/stream.h"
#include "tenzr/frontend/tensor.h"
#include "tenzr/runtime/device_buffer.h"
#include "tenzr/runtime/hardware_manager.h"

namespace tenzr {

  struct ExecutionResult {
    std::shared_ptr<DeviceBuffer> device_buffer;

    mutable std::shared_ptr<PinnedBuffer> output_buffer;

    Shape shape;

    DType dtype;

    const PJRT_Api* api{nullptr};

    std::vector<std::shared_ptr<DeviceBuffer>> device_buffers;

    std::vector<Shape> shapes;

    std::vector<DType> dtypes;

    void EnsureHostBuffer() const;

    template <typename T> std::vector<T> ToVector() const {
      EnsureHostBuffer();
      TENZR_CHECK(output_buffer != nullptr, "Null output buffer");
      size_t count = static_cast<size_t>(shape.NumElements());
      if constexpr (std::is_same_v<T, bool>) {
        std::vector<bool> result(count);
        const uint8_t* src = reinterpret_cast<const uint8_t*>(output_buffer->Data());
        for (size_t i = 0; i < count; ++i) {
          result[i] = (src[i] != 0);
        }
        return result;
      } else {
        std::vector<T> result(count);
        std::memcpy(result.data(), output_buffer->Data(), count * sizeof(T));
        return result;
      }
    }

    Array AsArray(std::string name = "") const;

    Array AsTensor(std::string name = "") const { return AsArray(std::move(name)); }

    operator Array() const { return AsArray(); }
  };

  class ExecutionEngine {
   public:
    ExecutionEngine();

    explicit ExecutionEngine(HardwareManager& hw_manager);

    ~ExecutionEngine() = default;

    std::shared_ptr<CachedExecutable> Compile(const Tensor& root);

    std::shared_ptr<CachedExecutable> CompileOptimized(const std::shared_ptr<Node>& target_node);

    // Synchronous execution
    ExecutionResult Execute(const Tensor& root);

    std::future<ExecutionResult> ExecuteAsync(const Tensor& root) {
      auto prom = std::make_shared<std::promise<ExecutionResult>>();
      auto fut = prom->get_future();
      execution::ExecutionStream::DefaultStream().Submit([this, root, prom]() {
#if TENZR_EXCEPTIONS_ENABLED
        try {
          prom->set_value(this->Execute(root));
        } catch (...) {
          prom->set_exception(std::current_exception());
        }
#else
        prom->set_value(this->Execute(root));
#endif
      });
      return fut;
    }

   private:
    HardwareManager* hw_{nullptr};
  };

}  // namespace tenzr

#endif  // TENZR_EXECUTION_EXECUTION_ENGINE_H_