// 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.
// 2. C system headers
#include "tenzr/txpr/jit.h"
#include <cstring>
#include <memory>
#include <string>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/common/logging.h"
#include "tenzr/common/trace.h"
#include "tenzr/execution/jit_cache.h"
#include "tenzr/runtime/device_buffer.h"
#include "tenzr/runtime/hardware_manager.h"
#include "tenzr/runtime/pjrt_error.h"
#include "tenzr/runtime/pjrt_executable.h"
#include "tenzr/txpr/emitter.h"
#include "tenzr/txpr/txpr.h"
namespace tenzr::txpr {
std::shared_ptr<CachedExecutable> CompileTxprModule(uint64_t hash, const Txpr& txpr) {
TENZR_TRACE_SCOPE_CAT("compiler", "compile_txpr_module");
auto& hw = HardwareManager::Instance();
if (!hw.IsInitialized()) {
hw.Initialize();
}
TxprEmitter emitter;
TxprCompiledModule module = emitter.Emit(txpr);
const PJRT_Api* api = hw.Api();
PJRT_Client* client = hw.Client();
TENZR_CHECK(api != nullptr && client != nullptr, "HardwareManager not initialized");
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();
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 of traced Txpr failed: " + err_msg + "\nMLIR:\n" + module.mlir_text);
}
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";
for (VarId v : txpr.InBinders()) {
cached_exec->in_shapes.push_back(txpr.Variable(v).shape);
cached_exec->in_dtypes.push_back(txpr.Variable(v).dtype);
}
for (size_t i = 0; i < module.out_shapes.size(); ++i) {
cached_exec->out_shapes.push_back(module.out_shapes[i]);
cached_exec->out_dtypes.push_back(module.out_dtypes[i]);
}
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(JitCache::DomainKey(JitCache::kJitKeyDomain, hash), cached_exec);
return cached_exec;
}
ExecutionResult ExecuteJittedExecutable(const std::shared_ptr<CachedExecutable>& cached, const std::vector<Array>& inputs) {
TENZR_TRACE_SCOPE_CAT("runtime", "execute_jitted_executable");
TENZR_CHECK(cached != nullptr && cached->executable != nullptr, "Null cached executable");
TENZR_CHECK(inputs.size() == cached->in_shapes.size(),
"JIT input count mismatch (executable expects " + std::to_string(cached->in_shapes.size()) + ", got " + std::to_string(inputs.size()) + ")");
auto& hw = HardwareManager::Instance();
if (!hw.IsInitialized()) {
hw.Initialize();
}
const PJRT_Api* api = hw.Api();
// 1. Gather input device buffers (Zero-copy if already resident on GPU!)
std::vector<PJRT_Buffer*> in_device_buffers;
in_device_buffers.reserve(inputs.size());
{
TENZR_TRACE_SCOPE_CAT("transfer", "jit_input_sync_device");
for (const auto& in_arr : inputs) {
TENZR_CHECK(in_arr.IsDefined(), "Input Array to JIT execution is undefined");
in_arr.SyncDevice();
TENZR_CHECK(in_arr.HasDeviceData(), "Input Array failed to resident on device");
in_device_buffers.push_back(in_arr.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 in JIT: " + 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!)
TENZR_TRACE_SCOPE_CAT("runtime", "jit_output_wrapping");
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::txpr