// 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/execution_engine.h"
#include <chrono>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <unordered_map>
#include "tenzr/common/error.h"
#include "tenzr/common/logging.h"
#include "tenzr/common/trace.h"
#include "tenzr/execution/ast_hash.h"
#include "tenzr/frontend/node_traversal.h"
#include "tenzr/runtime/pjrt_error.h"
#include "tenzr/runtime/pjrt_executable.h"
#include "tenzr/txpr/emitter.h"
#include "tenzr/txpr/node_to_txpr.h"
namespace tenzr {
namespace {
// Node traversal helpers live in tenzr/frontend/node_traversal.h.
} // namespace
ExecutionEngine::ExecutionEngine() : hw_(&HardwareManager::Instance()) {
if (!hw_->IsInitialized()) {
hw_->Initialize();
}
}
ExecutionEngine::ExecutionEngine(HardwareManager& hw_manager) : hw_(&hw_manager) {
if (!hw_->IsInitialized()) {
hw_->Initialize();
}
}
std::shared_ptr<CachedExecutable> ExecutionEngine::Compile(const Tensor& root) {
TENZR_CHECK(root.IsDefined(), "Cannot compile undefined tensor");
return CompileOptimized(root.node());
}
std::shared_ptr<CachedExecutable> ExecutionEngine::CompileOptimized(const std::shared_ptr<Node>& target_node) {
TENZR_TRACE_SCOPE_CAT("compiler", "ExecutionEngine::compile");
if (!hw_->IsInitialized()) {
hw_->Initialize();
}
uint64_t hash = ComputeAstHash(target_node);
const uint64_t cache_key = JitCache::DomainKey(JitCache::kEagerKeyDomain, hash);
auto cached = JitCache::Instance().Get(cache_key);
if (cached) {
TENZR_LOG_DEBUG << "JIT Cache HIT for AST hash " << hash << "\n";
return cached;
}
TENZR_LOG_DEBUG << "JIT Cache MISS for AST hash " << hash << ", lowering to StableHLO...\n";
txpr::Txpr tape = txpr::LowerNodeToTxpr(target_node);
txpr::TxprEmitter emitter;
txpr::TxprCompiledModule module = emitter.Emit(tape);
TENZR_LOG_DEBUG << "Emitted StableHLO MLIR:\n" << module.mlir_text << "\n";
const char* dump_env = std::getenv("TENZR_DUMP_MLIR");
const char* dump_dir_env = std::getenv("TENZR_DUMP_MLIR_DIR");
if (dump_env || dump_dir_env) {
std::string dir = dump_dir_env ? dump_dir_env : "/tmp/tenzr_mlir";
std::error_code ec;
std::filesystem::create_directories(dir, ec);
std::string file_path = dir + "/ast_" + std::to_string(hash) + ".mlir";
std::ofstream out(file_path);
if (out.is_open()) {
out << module.mlir_text;
out.close();
TENZR_LOG_INFO << "[MLIR Dump] Saved StableHLO MLIR to " << file_path << "\n";
}
}
const PJRT_Api* api = hw_->Api();
PJRT_Client* client = hw_->Client();
std::string format = "mlir";
PJRT_Program program;
std::memset(&program, 0, sizeof(program));
program.struct_size = PJRT_Program_STRUCT_SIZE;
program.code = module.mlir_text.data();
program.code_size = module.mlir_text.size();
program.format = format.c_str();
program.format_size = format.size();
// Protobuf wire format for CompileOptionsProto (1 replica, 1 partition, device 0):
const char opt_bytes[] = {0x1a, 0x06, 0x08, 0x00, 0x20, 0x01, 0x28, 0x01};
PJRT_Client_Compile_Args compile_args;
std::memset(&compile_args, 0, sizeof(compile_args));
compile_args.struct_size = PJRT_Client_Compile_Args_STRUCT_SIZE;
compile_args.client = client;
compile_args.program = &program;
compile_args.compile_options = opt_bytes;
compile_args.compile_options_size = sizeof(opt_bytes);
PJRT_Error* err = api->PJRT_Client_Compile(&compile_args);
if (err) {
std::string err_msg = PjrtErrorMessage(api, err);
PjrtErrorDestroy(api, err);
TENZR_THROW("PJRT compilation failed: " + err_msg);
}
auto cached_exec = std::make_shared<CachedExecutable>();
cached_exec->hash = hash;
cached_exec->executable = compile_args.executable;
const ExecutableInfo exec_info = PjrtExecutableInfo(api, compile_args.executable);
TENZR_CHECK(!exec_info.num_outputs_known || exec_info.num_outputs > 0, "PJRT compiled executable reports zero outputs");
TENZR_LOG_INFO << "PJRT executable: outputs=" << exec_info.num_outputs << " fingerprint=" << exec_info.fingerprint
<< " code_bytes=" << exec_info.generated_code_bytes << "\n";
cached_exec->in_shapes = module.in_shapes;
cached_exec->in_dtypes = module.in_dtypes;
cached_exec->out_shapes = module.out_shapes;
cached_exec->out_dtypes = module.out_dtypes;
TENZR_CHECK(!exec_info.num_outputs_known || exec_info.num_outputs == cached_exec->out_shapes.size(),
"PJRT executable output count mismatch (compiled " + std::to_string(exec_info.num_outputs) + ", descriptor " +
std::to_string(cached_exec->out_shapes.size()) + ")");
JitCache::Instance().Put(cache_key, cached_exec);
return cached_exec;
}
void ExecutionResult::EnsureHostBuffer() const {
TENZR_TRACE_SCOPE_CAT("transfer", "PJRT::D2H_Transfer");
if (output_buffer) {
return;
}
TENZR_CHECK(device_buffer != nullptr, "No buffer available in ExecutionResult");
const PJRT_Api* api_ptr = api ? api : HardwareManager::Instance().Api();
TENZR_CHECK(api_ptr != nullptr, "HardwareManager API uninitialized");
size_t out_num_elems = static_cast<size_t>(shape.NumElements());
size_t out_bytes = out_num_elems * SizeOf(dtype);
auto host_output = MemoryPool::Instance().Allocate(out_bytes);
PJRT_Buffer_ToHostBuffer_Args thb_args;
std::memset(&thb_args, 0, sizeof(thb_args));
thb_args.struct_size = PJRT_Buffer_ToHostBuffer_Args_STRUCT_SIZE;
thb_args.src = device_buffer->Handle();
thb_args.host_layout = nullptr;
thb_args.dst = host_output->Data();
thb_args.dst_size = out_bytes;
const auto transfer_start = std::chrono::steady_clock::now();
PJRT_Error* err = api_ptr->PJRT_Buffer_ToHostBuffer(&thb_args);
if (err) {
std::string err_msg = PjrtErrorMessage(api_ptr, err);
PjrtErrorDestroy(api_ptr, err);
TENZR_THROW("Buffer_ToHostBuffer failed: " + err_msg);
}
if (thb_args.event) {
PJRT_Event_Await_Args await_args{PJRT_Event_Await_Args_STRUCT_SIZE, nullptr, thb_args.event};
PJRT_Error* await_err = api_ptr->PJRT_Event_Await(&await_args);
if (await_err != nullptr) {
std::string err_msg = PjrtErrorMessage(api_ptr, await_err);
PjrtErrorDestroy(api_ptr, await_err);
TENZR_LOG_WARN << "PJRT_Event_Await failed: " << err_msg << "\n";
}
PJRT_Event_Destroy_Args d_args{PJRT_Event_Destroy_Args_STRUCT_SIZE, nullptr, thb_args.event};
(void)api_ptr->PJRT_Event_Destroy(&d_args);
}
const size_t device_bytes = device_buffer->DeviceSizeBytes();
const size_t moved_bytes = device_bytes > 0 ? device_bytes : out_bytes;
const auto transfer_us = std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - transfer_start).count();
profiler::RecordMemoryTransfer("PJRT::D2H", moved_bytes, false, static_cast<int64_t>(transfer_us));
output_buffer = std::move(host_output);
}
Array ExecutionResult::AsArray(std::string name) const {
auto var_node = std::make_shared<VariableNode>(shape, dtype, std::move(name));
if (device_buffer) {
var_node->SetDeviceData(device_buffer);
}
if (output_buffer) {
var_node->SetData(output_buffer);
}
return Array(std::move(var_node));
}
ExecutionResult ExecutionEngine::Execute(const Tensor& root) {
TENZR_TRACE_SCOPE_CAT("runtime", "ExecutionEngine::execute");
TENZR_CHECK(root.IsDefined(), "Cannot execute undefined tensor");
std::shared_ptr<Node> target_node = root.node();
auto cached = CompileOptimized(target_node);
const PJRT_Api* api = hw_->Api();
auto current_params = CollectParameters(target_node);
TENZR_CHECK(current_params.size() == cached->in_shapes.size(), "Parameter count mismatch between compiled module (" +
std::to_string(cached->in_shapes.size()) + ") and target graph (" +
std::to_string(current_params.size()) + ")");
for (size_t i = 0; i < current_params.size(); ++i) {
TENZR_CHECK(current_params[i]->shape() == cached->in_shapes[i] && current_params[i]->Dtype() == cached->in_dtypes[i],
"Parameter " + std::to_string(i) + " shape/dtype does not match the compiled executable");
}
// 1. Gather input device buffers (Zero-copy if already on device!)
std::vector<PJRT_Buffer*> in_device_buffers;
in_device_buffers.reserve(current_params.size());
for (const auto& param : current_params) {
TENZR_CHECK(param->HasData(), "Parameter '" + param->VariableName() + "' has no bound data");
// Ensure device residency (NO-OP if already on GPU)
param->SyncDevice();
in_device_buffers.push_back(param->DeviceData()->Handle());
}
// 2. Setup execution arguments (one PJRT output slot per descriptor output)
const size_t num_outputs = cached->out_shapes.empty() ? 1 : cached->out_shapes.size();
PJRT_Buffer* const* dev_arg_list[1] = {in_device_buffers.data()};
std::vector<PJRT_Buffer*> out_bufs(num_outputs, nullptr);
PJRT_Buffer** dev_out_list[1] = {out_bufs.data()};
PJRT_Event* comp_event[1] = {nullptr};
PJRT_ExecuteOptions exec_options;
std::memset(&exec_options, 0, sizeof(exec_options));
exec_options.struct_size = PJRT_ExecuteOptions_STRUCT_SIZE;
PJRT_LoadedExecutable_Execute_Args exec_args;
std::memset(&exec_args, 0, sizeof(exec_args));
exec_args.struct_size = PJRT_LoadedExecutable_Execute_Args_STRUCT_SIZE;
exec_args.executable = cached->executable;
exec_args.options = &exec_options;
exec_args.argument_lists = dev_arg_list;
exec_args.num_devices = 1;
exec_args.num_args = in_device_buffers.size();
exec_args.output_lists = dev_out_list;
exec_args.device_complete_events = comp_event;
exec_args.execute_device = nullptr;
{
TENZR_TRACE_SCOPE_CAT("runtime", "PJRT::LoadedExecutable_Execute");
PJRT_Error* err = api->PJRT_LoadedExecutable_Execute(&exec_args);
if (err) {
std::string err_msg = PjrtErrorMessage(api, err);
PjrtErrorDestroy(api, err);
TENZR_THROW("LoadedExecutable_Execute failed: " + err_msg);
}
}
if (comp_event[0]) {
PJRT_Event_Destroy_Args d_args{PJRT_Event_Destroy_Args_STRUCT_SIZE, nullptr, comp_event[0]};
(void)api->PJRT_Event_Destroy(&d_args);
}
// 3. Wrap every output buffer in an RAII DeviceBuffer (Resident on device!)
ExecutionResult result;
result.device_buffers.reserve(num_outputs);
for (size_t i = 0; i < num_outputs; ++i) {
result.device_buffers.push_back(out_bufs[i] ? std::make_shared<DeviceBuffer>(out_bufs[i], api) : nullptr);
}
result.device_buffer = result.device_buffers.empty() ? nullptr : result.device_buffers.front();
result.shape = cached->out_shapes.at(0);
result.dtype = cached->out_dtypes.at(0);
result.shapes = cached->out_shapes;
result.dtypes = cached->out_dtypes;
result.api = api;
return result;
}
} // namespace tenzr