hardware_manager.h¶

Namespaces¶

Name
tenzr

Classes¶

	Name
struct	tenzr::DeviceInfo Description of one PJRT device: its ids, platform kind, debug string, and handle.
class	tenzr::HardwareManager Owns the loaded PJRT plugin and client, exposing device enumeration and platform metadata.
struct	tenzr::HardwareManager::DeviceMemoryStats Device allocator statistics (from PJRT_Device_MemoryStats when the plugin supports it).

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_RUNTIME_HARDWARE_MANAGER_H_
#define TENZR_RUNTIME_HARDWARE_MANAGER_H_

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

#include <cstdint>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>

namespace tenzr {

  struct DeviceInfo {
    int id{0};

    int process_index{0};

    std::string kind;

    std::string debug_string;

    PJRT_Device* handle{nullptr};
  };

  class HardwareManager {
   public:
    static HardwareManager& Instance();

    HardwareManager();

    ~HardwareManager();

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

    void Initialize(const std::string& preferred_plugin_path = "");

    void Shutdown();

    bool IsInitialized() const { return client_ != nullptr; }

    const PJRT_Api* Api() const { return api_; }

    PJRT_Client* Client() const { return client_; }

    const std::string& PlatformName() const { return platform_name_; }

    const std::string& PlatformVersion() const { return platform_version_; }

    const std::string& LoadedPluginPath() const { return plugin_path_; }

    const std::vector<DeviceInfo>& Devices() const { return devices_; }

    PJRT_Device* DefaultDevice() const;

    const PJRT_Extension_Base* Extension(PJRT_Extension_Type type) const;

    struct DeviceMemoryStats {
      bool available{false};
      int64_t bytes_in_use{0};
      int64_t peak_bytes_in_use{0};
      int64_t bytes_limit{0};
    };

    DeviceMemoryStats MemoryStats(PJRT_Device* device = nullptr) const;

   private:
    std::string DiscoverPluginPath(const std::string& preferred_path) const;
    void EnsureCudaEnvironment();
    void EnsureCpuEnvironment();

    void* dl_handle_{nullptr};
    const PJRT_Api* api_{nullptr};
    PJRT_Client* client_{nullptr};

    std::string plugin_path_;
    std::string platform_name_;
    std::string platform_version_;
    std::vector<DeviceInfo> devices_;
    mutable std::unordered_map<PJRT_Extension_Type, const PJRT_Extension_Base*> extension_cache_;
    mutable std::mutex mutex_;
  };

}  // namespace tenzr

#endif  // TENZR_RUNTIME_HARDWARE_MANAGER_H_