// 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/runtime/hardware_manager.h"
#include <dlfcn.h>
#include <sys/utsname.h>
#include <xla/pjrt/pjrt_loader.h>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <string>
#include <utility>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/common/logging.h"
#include "tenzr/runtime/pjrt_backends.h"
#include "tenzr/runtime/pjrt_error.h"
namespace tenzr {
namespace {
std::string HostArch() {
struct utsname info;
if (uname(&info) == 0) {
const std::string machine = info.machine;
if (machine == "amd64" || machine == "x86_64") {
return "x86_64";
}
if (machine == "arm64" || machine == "aarch64") {
return "aarch64";
}
}
return "x86_64";
}
std::string ModuleDirectory() {
static const char anchor = 0;
static const std::string directory = []() -> std::string {
Dl_info info{};
if (dladdr(&anchor, &info) != 0 && info.dli_fname != nullptr) {
return std::filesystem::path(info.dli_fname).parent_path().string();
}
return {};
}();
return directory;
}
} // namespace
HardwareManager& HardwareManager::Instance() {
static HardwareManager manager;
return manager;
}
HardwareManager::HardwareManager() = default;
HardwareManager::~HardwareManager() { Shutdown(); }
std::string HardwareManager::DiscoverPluginPath(const std::string& preferred_path) const {
if (!preferred_path.empty() && std::filesystem::exists(preferred_path)) {
return preferred_path;
}
const char* env_path = std::getenv("PJRT_PLUGIN_PATH");
if (env_path && std::filesystem::exists(env_path)) {
return std::string(env_path);
}
std::string device_pref;
if (const char* dev = std::getenv("TENZR_DEVICE")) {
device_pref = dev;
} else if (const char* backend = std::getenv("TENZR_BACKEND")) {
device_pref = backend;
}
for (char& ch : device_pref) {
ch = static_cast<char>(std::tolower(ch));
}
const std::string arch = HostArch();
// {device, path under a root}; accelerator-first, CPU last. The table is
// generated from cmake/packaging/pjrt-backends.cmake (the single backend manifest).
std::vector<std::pair<std::string, std::string>> plugins;
for (const auto& backend : detail::kPjrtBackends) {
std::string relative = "devices/" + std::string(backend.family) + "/";
if (backend.variant[0] != '\0') {
relative += std::string(backend.variant) + "/";
}
relative += arch + "/" + backend.library;
plugins.emplace_back(backend.device, relative);
}
auto matches = [&](const std::string& device) {
if (device_pref.empty()) {
return true;
}
if (device_pref == "cuda" || device_pref == "gpu" || device_pref == "nvidia") {
return device.rfind("cuda", 0) == 0;
}
if (device_pref == "rocm" || device_pref == "amd") {
return device.rfind("rocm", 0) == 0;
}
if (device_pref == "oneapi" || device_pref == "intel") {
return device == "xpu";
}
return device == device_pref;
};
std::vector<std::string> roots;
// Package-relative first: plugins bundled inside an installed wheel live
// next to the extension (tenzr/pjrt/<family>/...).
const std::string module_dir = ModuleDirectory();
if (!module_dir.empty()) {
roots.emplace_back(module_dir + "/pjrt");
}
#ifdef TENZR_PJRT_INSTALL_DIR
// Configured install root: covers DEB multiarch and RPM lib64.
roots.emplace_back(TENZR_PJRT_INSTALL_DIR);
#endif
roots.emplace_back("/usr/local/lib/tenzr/pjrt");
roots.emplace_back("/usr/lib/tenzr/pjrt");
#ifdef TENZR_DEPENDENCIES_DIR
roots.emplace_back(std::string(TENZR_DEPENDENCIES_DIR) + "/pjrt");
#endif
roots.emplace_back("dependencies/pjrt");
roots.emplace_back("../dependencies/pjrt");
roots.emplace_back("../../dependencies/pjrt");
for (const auto& plugin : plugins) {
if (!matches(plugin.first)) {
continue;
}
for (const auto& root : roots) {
const std::string candidate = root + "/" + plugin.second;
if (std::filesystem::exists(candidate)) {
return candidate;
}
}
}
TENZR_THROW("Could not find any PJRT plugin .so. Set PJRT_PLUGIN_PATH or run dependencies/pjrt/fetch.sh.");
}
void HardwareManager::EnsureCudaEnvironment() {
std::vector<std::string> nvidia_roots;
if (const char* tenzr_cuda = std::getenv("TENZR_CUDA_PATH")) {
nvidia_roots.emplace_back(tenzr_cuda);
}
if (const char* cuda_home = std::getenv("CUDA_HOME")) {
nvidia_roots.emplace_back(cuda_home);
}
if (const char* cuda_path = std::getenv("CUDA_PATH")) {
nvidia_roots.emplace_back(cuda_path);
}
if (const char* venv = std::getenv("VIRTUAL_ENV")) {
std::string lib_dir = std::string(venv) + "/lib";
if (std::filesystem::exists(lib_dir)) {
for (const auto& entry : std::filesystem::directory_iterator(lib_dir)) {
if (entry.is_directory() && entry.path().filename().string().rfind("python", 0) == 0) {
std::string nv = (entry.path() / "site-packages" / "nvidia").string();
if (std::filesystem::exists(nv)) {
nvidia_roots.push_back(nv);
}
}
}
}
}
nvidia_roots.emplace_back("/usr/local/cuda");
nvidia_roots.emplace_back("/usr/local/cuda-12");
for (const auto& root : nvidia_roots) {
if (!std::filesystem::exists(root)) {
continue;
}
std::string ptxas_dir = std::filesystem::exists(root + "/cuda_nvcc/bin") ? (root + "/cuda_nvcc/bin") : (root + "/bin");
if (std::filesystem::exists(ptxas_dir)) {
const char* current_path = std::getenv("PATH");
std::string new_path = ptxas_dir + ":" + (current_path ? current_path : "");
setenv("PATH", new_path.c_str(), 1);
}
std::string nvcc_dir = std::filesystem::exists(root + "/cuda_nvcc") ? (root + "/cuda_nvcc") : root;
if (std::filesystem::exists(nvcc_dir + "/nvvm/libdevice/libdevice.10.bc")) {
std::string data_flag = "--xla_gpu_cuda_data_dir=" + nvcc_dir;
const char* cur_flags = std::getenv("XLA_FLAGS");
if (!cur_flags || std::string(cur_flags).find("xla_gpu_cuda_data_dir") == std::string::npos) {
std::string updated_flags = data_flag + " " + (cur_flags ? std::string(cur_flags) : "");
setenv("XLA_FLAGS", updated_flags.c_str(), 1);
}
}
const char* current_flags = std::getenv("XLA_FLAGS");
if (!current_flags || std::string(current_flags).find("unsafe_fallback") == std::string::npos) {
std::string new_flags = "--xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found " + (current_flags ? std::string(current_flags) : "");
setenv("XLA_FLAGS", new_flags.c_str(), 1);
}
dlopen("libcuda.so.1", RTLD_NOW | RTLD_GLOBAL);
std::vector<std::string> companion_libs = {root + "/cuda_runtime/lib/libcudart.so.12",
root + "/cublas/lib/libcublas.so.12",
root + "/cudnn/lib/libcudnn.so.9",
root + "/lib64/libcudart.so.12",
root + "/lib64/libcublas.so.12",
root + "/lib64/libcudnn.so.9"};
for (const auto& lib : companion_libs) {
if (std::filesystem::exists(lib)) {
dlopen(lib.c_str(), RTLD_NOW | RTLD_GLOBAL);
}
}
break;
}
}
void HardwareManager::EnsureCpuEnvironment() {
#if defined(__x86_64__) || defined(_M_X64)
const char* mode = std::getenv("TENZR_CPU_TUNING");
if (mode == nullptr || std::string(mode) != "onednn") {
return;
}
static const std::pair<const char*, const char*> kDefaults[] = {
{"xla_cpu_use_onednn", "true"},
{"xla_cpu_experimental_onednn_fusion_type", "dot,eltwise"},
};
const char* current = std::getenv("XLA_FLAGS");
std::string flags = current != nullptr ? current : "";
bool changed = false;
for (const auto& [name, value] : kDefaults) {
const std::string key = std::string("--") + name + "=";
if (flags.find(key) == std::string::npos) {
if (!flags.empty() && flags.back() != ' ') {
flags += ' ';
}
flags += key + value;
changed = true;
}
}
if (changed) {
setenv("XLA_FLAGS", flags.c_str(), 1);
TENZR_LOG_INFO << "CPU backend tuning XLA_FLAGS: " << flags << "\n";
}
#endif
}
void HardwareManager::Initialize(const std::string& preferred_plugin_path) {
std::lock_guard<std::mutex> lock(mutex_);
if (client_ != nullptr) {
return; // Already initialized
}
extension_cache_.clear();
plugin_path_ = DiscoverPluginPath(preferred_plugin_path);
TENZR_LOG_INFO << "Loading PJRT plugin: " << plugin_path_ << "\n";
if (plugin_path_.find("cuda") != std::string::npos) {
EnsureCudaEnvironment();
} else if (plugin_path_.find("cpu") != std::string::npos) {
EnsureCpuEnvironment();
}
// Provide compiler runtime emulation routines (e.g. __extendhfsf2) in the global scope.
dlopen("libgcc_s.so.1", RTLD_NOW | RTLD_GLOBAL);
// Keep the handle for the process lifetime; dlclose faults inside some plugins' destructors.
dl_handle_ = dlopen(plugin_path_.c_str(), RTLD_LAZY | RTLD_LOCAL | RTLD_NODELETE);
TENZR_CHECK(dl_handle_ != nullptr, "dlopen failed for " + plugin_path_ + ": " + dlerror());
typedef const PJRT_Api* (*GetPjrtApiFn)();
GetPjrtApiFn get_api = reinterpret_cast<GetPjrtApiFn>(dlsym(dl_handle_, "GetPjrtApi"));
TENZR_CHECK(get_api != nullptr, "dlsym failed to locate 'GetPjrtApi' in " + plugin_path_ + ": " + dlerror());
api_ = get_api();
TENZR_CHECK(api_ != nullptr, "GetPjrtApi returned nullptr");
TENZR_LOG_INFO << "PJRT API Version: " << api_->pjrt_api_version.major_version << "." << api_->pjrt_api_version.minor_version << "\n";
// Every supported plugin must provide the client entry point; other calls are gated below.
TENZR_CHECK(PJRT_HAS(api_, PJRT_Client_Create), "PJRT plugin is missing PJRT_Client_Create");
auto error_message = [&](PJRT_Error* err) {
std::string message = "unknown PJRT error";
if (err != nullptr && PJRT_HAS(api_, PJRT_Error_Message) && api_->PJRT_Error_Message != nullptr) {
PJRT_Error_Message_Args msg_args;
std::memset(&msg_args, 0, sizeof(msg_args));
msg_args.struct_size = PJRT_Error_Message_Args_STRUCT_SIZE;
msg_args.error = err;
api_->PJRT_Error_Message(&msg_args);
message = std::string(msg_args.message, msg_args.message_size);
}
if (err != nullptr && PJRT_HAS(api_, PJRT_Error_Destroy) && api_->PJRT_Error_Destroy != nullptr) {
PJRT_Error_Destroy_Args destroy_err{PJRT_Error_Destroy_Args_STRUCT_SIZE, nullptr, err};
api_->PJRT_Error_Destroy(&destroy_err);
}
return message;
};
// 1. Initialize the plugin (older plugins omit this entry point).
if (PJRT_HAS(api_, PJRT_Plugin_Initialize) && api_->PJRT_Plugin_Initialize != nullptr) {
PJRT_Plugin_Initialize_Args init_args;
std::memset(&init_args, 0, sizeof(init_args));
init_args.struct_size = PJRT_Plugin_Initialize_Args_STRUCT_SIZE;
PJRT_Error* err = api_->PJRT_Plugin_Initialize(&init_args);
if (err != nullptr) {
TENZR_THROW("PJRT_Plugin_Initialize failed: " + error_message(err));
}
}
// 2. Create the client (required).
PJRT_Client_Create_Args create_args;
std::memset(&create_args, 0, sizeof(create_args));
create_args.struct_size = PJRT_Client_Create_Args_STRUCT_SIZE;
PJRT_Error* err = api_->PJRT_Client_Create(&create_args);
if (err != nullptr) {
TENZR_THROW("PJRT_Client_Create failed: " + error_message(err));
}
client_ = create_args.client;
TENZR_CHECK(client_ != nullptr, "PJRT_Client_Create returned null client");
// 3. Platform name and version (informational).
if (PJRT_HAS(api_, PJRT_Client_PlatformName) && api_->PJRT_Client_PlatformName != nullptr) {
PJRT_Client_PlatformName_Args name_args;
std::memset(&name_args, 0, sizeof(name_args));
name_args.struct_size = PJRT_Client_PlatformName_Args_STRUCT_SIZE;
name_args.client = client_;
(void)api_->PJRT_Client_PlatformName(&name_args);
platform_name_ = std::string(name_args.platform_name, name_args.platform_name_size);
}
if (PJRT_HAS(api_, PJRT_Client_PlatformVersion) && api_->PJRT_Client_PlatformVersion != nullptr) {
PJRT_Client_PlatformVersion_Args ver_args;
std::memset(&ver_args, 0, sizeof(ver_args));
ver_args.struct_size = PJRT_Client_PlatformVersion_Args_STRUCT_SIZE;
ver_args.client = client_;
(void)api_->PJRT_Client_PlatformVersion(&ver_args);
platform_version_ = std::string(ver_args.platform_version, ver_args.platform_version_size);
}
TENZR_LOG_INFO << "Hardware platform: " << platform_name_ << " (" << platform_version_ << ")\n";
// 4. Addressable devices.
if (PJRT_HAS(api_, PJRT_Client_AddressableDevices) && api_->PJRT_Client_AddressableDevices != nullptr) {
PJRT_Client_AddressableDevices_Args dev_args;
std::memset(&dev_args, 0, sizeof(dev_args));
dev_args.struct_size = PJRT_Client_AddressableDevices_Args_STRUCT_SIZE;
dev_args.client = client_;
err = api_->PJRT_Client_AddressableDevices(&dev_args);
if (err == nullptr) {
devices_.reserve(dev_args.num_addressable_devices);
for (size_t i = 0; i < dev_args.num_addressable_devices; ++i) {
PJRT_Device* dev = dev_args.addressable_devices[i];
DeviceInfo info;
info.handle = dev;
if (PJRT_HAS(api_, PJRT_Device_GetDescription) && api_->PJRT_Device_GetDescription != nullptr) {
PJRT_Device_GetDescription_Args desc_args;
std::memset(&desc_args, 0, sizeof(desc_args));
desc_args.struct_size = PJRT_Device_GetDescription_Args_STRUCT_SIZE;
desc_args.device = dev;
(void)api_->PJRT_Device_GetDescription(&desc_args);
if (desc_args.device_description != nullptr) {
if (PJRT_HAS(api_, PJRT_DeviceDescription_Id) && api_->PJRT_DeviceDescription_Id != nullptr) {
PJRT_DeviceDescription_Id_Args id_args;
std::memset(&id_args, 0, sizeof(id_args));
id_args.struct_size = PJRT_DeviceDescription_Id_Args_STRUCT_SIZE;
id_args.device_description = desc_args.device_description;
(void)api_->PJRT_DeviceDescription_Id(&id_args);
info.id = id_args.id;
}
if (PJRT_HAS(api_, PJRT_DeviceDescription_ProcessIndex) && api_->PJRT_DeviceDescription_ProcessIndex != nullptr) {
PJRT_DeviceDescription_ProcessIndex_Args proc_args;
std::memset(&proc_args, 0, sizeof(proc_args));
proc_args.struct_size = PJRT_DeviceDescription_ProcessIndex_Args_STRUCT_SIZE;
proc_args.device_description = desc_args.device_description;
(void)api_->PJRT_DeviceDescription_ProcessIndex(&proc_args);
info.process_index = proc_args.process_index;
}
if (PJRT_HAS(api_, PJRT_DeviceDescription_Kind) && api_->PJRT_DeviceDescription_Kind != nullptr) {
PJRT_DeviceDescription_Kind_Args kind_args;
std::memset(&kind_args, 0, sizeof(kind_args));
kind_args.struct_size = PJRT_DeviceDescription_Kind_Args_STRUCT_SIZE;
kind_args.device_description = desc_args.device_description;
(void)api_->PJRT_DeviceDescription_Kind(&kind_args);
info.kind = std::string(kind_args.device_kind, kind_args.device_kind_size);
}
if (PJRT_HAS(api_, PJRT_DeviceDescription_DebugString) && api_->PJRT_DeviceDescription_DebugString != nullptr) {
PJRT_DeviceDescription_DebugString_Args dbg_args;
std::memset(&dbg_args, 0, sizeof(dbg_args));
dbg_args.struct_size = PJRT_DeviceDescription_DebugString_Args_STRUCT_SIZE;
dbg_args.device_description = desc_args.device_description;
(void)api_->PJRT_DeviceDescription_DebugString(&dbg_args);
info.debug_string = std::string(dbg_args.debug_string, dbg_args.debug_string_size);
}
}
}
devices_.push_back(std::move(info));
}
}
}
TENZR_LOG_INFO << "Discovered " << devices_.size() << " addressable device(s)\n";
}
PJRT_Device* HardwareManager::DefaultDevice() const {
TENZR_CHECK(!devices_.empty(), "No addressable devices discovered");
return devices_[0].handle;
}
const PJRT_Extension_Base* HardwareManager::Extension(PJRT_Extension_Type type) const {
std::lock_guard<std::mutex> lock(mutex_);
auto cached = extension_cache_.find(type);
if (cached != extension_cache_.end()) {
return cached->second;
}
const PJRT_Extension_Base* ext = pjrt::find_extension(api_, type);
extension_cache_.emplace(type, ext);
return ext;
}
HardwareManager::DeviceMemoryStats HardwareManager::MemoryStats(PJRT_Device* device) const {
DeviceMemoryStats stats;
const PJRT_Api* api = api_;
if (api == nullptr || !PJRT_HAS(api, PJRT_Device_MemoryStats)) {
return stats;
}
if (device == nullptr) {
device = DefaultDevice();
}
if (device == nullptr) {
return stats;
}
PJRT_Device_MemoryStats_Args args;
std::memset(&args, 0, sizeof(args));
args.struct_size = PJRT_Device_MemoryStats_Args_STRUCT_SIZE;
args.device = device;
PJRT_Error* err = api->PJRT_Device_MemoryStats(&args);
if (err != nullptr) {
PjrtErrorDestroy(api, err);
return stats;
}
stats.available = true;
stats.bytes_in_use = args.bytes_in_use;
stats.peak_bytes_in_use = args.peak_bytes_in_use_is_set ? args.peak_bytes_in_use : args.bytes_in_use;
stats.bytes_limit = args.bytes_limit_is_set ? args.bytes_limit : 0;
return stats;
}
void HardwareManager::Shutdown() {
std::lock_guard<std::mutex> lock(mutex_);
if (client_ != nullptr && api_ != nullptr && PJRT_HAS(api_, PJRT_Client_Destroy)) {
PJRT_Client_Destroy_Args destroy_args;
std::memset(&destroy_args, 0, sizeof(destroy_args));
destroy_args.struct_size = PJRT_Client_Destroy_Args_STRUCT_SIZE;
destroy_args.client = client_;
PJRT_Error* err = api_->PJRT_Client_Destroy(&destroy_args);
PjrtErrorDestroy(api_, err);
}
client_ = nullptr;
devices_.clear();
extension_cache_.clear();
api_ = nullptr;
}
} // namespace tenzr