src/runtime/pjrt_profiler.cc¶

// 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/pjrt_profiler.h"

#include <cstring>
#include <fstream>

#include "tenzr/common/logging.h"
#include "tenzr/runtime/hardware_manager.h"

namespace tenzr {
  namespace {
    std::string ProfilerErrorMessage(const PLUGIN_Profiler_Api* api, PLUGIN_Profiler_Error* err) {
      if (err == nullptr || api == nullptr || api->error_message == nullptr) {
        return "unknown profiler error";
      }
      PLUGIN_Profiler_Error_Message_Args args;
      std::memset(&args, 0, sizeof(args));
      args.struct_size = PLUGIN_Profiler_Error_Message_Args_STRUCT_SIZE;
      args.error = err;
      api->error_message(&args);
      return args.message != nullptr ? std::string(args.message, args.message_size) : "unknown profiler error";
    }

    void ProfilerErrorDestroy(const PLUGIN_Profiler_Api* api, PLUGIN_Profiler_Error* err) {
      if (err != nullptr && api != nullptr && api->error_destroy != nullptr) {
        PLUGIN_Profiler_Error_Destroy_Args args;
        std::memset(&args, 0, sizeof(args));
        args.struct_size = PLUGIN_Profiler_Error_Destroy_Args_STRUCT_SIZE;
        args.error = err;
        api->error_destroy(&args);
      }
    }

    void ProfilerReportError(const PLUGIN_Profiler_Api* api, PLUGIN_Profiler_Error* err, const char* context) {
      TENZR_LOG_WARN << context << ": " << ProfilerErrorMessage(api, err) << "\n";
      ProfilerErrorDestroy(api, err);
    }
  }  // namespace

  PjrtProfiler::PjrtProfiler() {
    auto& hw = HardwareManager::Instance();
    if (!hw.IsInitialized()) {
#if TENZR_EXCEPTIONS_ENABLED
      try {
        hw.Initialize();
      } catch (...) {
        return;
      }
#else
      hw.Initialize();
#endif
    }
    const auto* ext = reinterpret_cast<const PJRT_Profiler_Extension*>(hw.Extension(PJRT_Extension_Type_Profiler));
    if (ext == nullptr || ext->profiler_api == nullptr) {
      return;
    }
    api_ = ext->profiler_api;
    if (api_->create == nullptr) {
      api_ = nullptr;
      return;
    }

    PLUGIN_Profiler_Create_Args args;
    std::memset(&args, 0, sizeof(args));
    args.struct_size = PLUGIN_Profiler_Create_Args_STRUCT_SIZE;
    PLUGIN_Profiler_Error* err = api_->create(&args);
    if (err != nullptr) {
      ProfilerReportError(api_, err, "profiler create failed");
      api_ = nullptr;
      return;
    }
    profiler_ = args.profiler;
  }

  PjrtProfiler::~PjrtProfiler() {
    if (profiler_ != nullptr && api_ != nullptr && api_->destroy != nullptr) {
      PLUGIN_Profiler_Destroy_Args args;
      std::memset(&args, 0, sizeof(args));
      args.struct_size = PLUGIN_Profiler_Destroy_Args_STRUCT_SIZE;
      args.profiler = profiler_;
      PLUGIN_Profiler_Error* err = api_->destroy(&args);
      ProfilerReportError(api_, err, "profiler destroy failed");
    }
  }

  bool PjrtProfiler::Supported() const { return api_ != nullptr && profiler_ != nullptr; }

  bool PjrtProfiler::Start() {
    if (!Supported() || started_ || api_->start == nullptr) {
      return false;
    }
    PLUGIN_Profiler_Start_Args args;
    std::memset(&args, 0, sizeof(args));
    args.struct_size = PLUGIN_Profiler_Start_Args_STRUCT_SIZE;
    args.profiler = profiler_;
    PLUGIN_Profiler_Error* err = api_->start(&args);
    if (err != nullptr) {
      ProfilerReportError(api_, err, "profiler start failed");
      return false;
    }
    started_ = true;
    return true;
  }

  bool PjrtProfiler::Stop() {
    if (!Supported() || !started_ || api_->stop == nullptr) {
      return false;
    }
    PLUGIN_Profiler_Stop_Args args;
    std::memset(&args, 0, sizeof(args));
    args.struct_size = PLUGIN_Profiler_Stop_Args_STRUCT_SIZE;
    args.profiler = profiler_;
    PLUGIN_Profiler_Error* err = api_->stop(&args);
    started_ = false;
    if (err != nullptr) {
      ProfilerReportError(api_, err, "profiler stop failed");
      return false;
    }
    return true;
  }

  std::vector<uint8_t> PjrtProfiler::CollectData() const {
    std::vector<uint8_t> data;
    if (!Supported() || api_->collect_data == nullptr) {
      return data;
    }

    PLUGIN_Profiler_CollectData_Args args;
    std::memset(&args, 0, sizeof(args));
    args.struct_size = PLUGIN_Profiler_CollectData_Args_STRUCT_SIZE;
    args.profiler = profiler_;
    args.buffer = nullptr;
    PLUGIN_Profiler_Error* err = api_->collect_data(&args);
    if (err != nullptr) {
      ProfilerReportError(api_, err, "profiler collect_data failed");
      return data;
    }
    if (args.buffer_size_in_bytes == 0) {
      return data;
    }

    data.resize(args.buffer_size_in_bytes);
    std::memset(&args, 0, sizeof(args));
    args.struct_size = PLUGIN_Profiler_CollectData_Args_STRUCT_SIZE;
    args.profiler = profiler_;
    args.buffer = data.data();
    args.buffer_size_in_bytes = data.size();
    err = api_->collect_data(&args);
    if (err != nullptr) {
      ProfilerReportError(api_, err, "profiler collect_data failed");
      data.clear();
      return data;
    }
    data.resize(args.buffer_size_in_bytes);
    return data;
  }

  bool PjrtProfiler::Dump(const std::string& path) const {
    const std::vector<uint8_t> data = CollectData();
    if (data.empty()) {
      return false;
    }
    std::ofstream out(path, std::ios::binary);
    if (!out) {
      return false;
    }
    out.write(reinterpret_cast<const char*>(data.data()), static_cast<std::streamsize>(data.size()));
    return out.good();
  }
}  // namespace tenzr