include/tenzr/runtime/pjrt_executable.h¶

// 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_PJRT_EXECUTABLE_H_
#define TENZR_RUNTIME_PJRT_EXECUTABLE_H_

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

#include <cstring>
#include <string>

#include "tenzr/runtime/pjrt_error.h"

namespace tenzr {
  struct ExecutableInfo {
    bool num_outputs_known{false};
    size_t num_outputs{0};
    std::string fingerprint;
    bool generated_code_known{false};
    int64_t generated_code_bytes{0};
  };

  [[nodiscard]] inline PJRT_Executable* PjrtLoadedExecutableGet(const PJRT_Api* api, PJRT_LoadedExecutable* loaded) {
    if (api == nullptr || loaded == nullptr || !PJRT_HAS(api, PJRT_LoadedExecutable_GetExecutable)) {
      return nullptr;
    }
    PJRT_LoadedExecutable_GetExecutable_Args args;
    std::memset(&args, 0, sizeof(args));
    args.struct_size = PJRT_LoadedExecutable_GetExecutable_Args_STRUCT_SIZE;
    args.loaded_executable = loaded;
    PJRT_Error* err = api->PJRT_LoadedExecutable_GetExecutable(&args);
    if (err != nullptr) {
      PjrtErrorDestroy(api, err);
      return nullptr;
    }
    return args.executable;
  }

  [[nodiscard]] inline ExecutableInfo PjrtExecutableInfo(const PJRT_Api* api, PJRT_LoadedExecutable* loaded) {
    ExecutableInfo info;
    PJRT_Executable* exec = PjrtLoadedExecutableGet(api, loaded);
    if (exec == nullptr) {
      return info;
    }

    if (PJRT_HAS(api, PJRT_Executable_NumOutputs)) {
      PJRT_Executable_NumOutputs_Args args;
      std::memset(&args, 0, sizeof(args));
      args.struct_size = PJRT_Executable_NumOutputs_Args_STRUCT_SIZE;
      args.executable = exec;
      PJRT_Error* err = api->PJRT_Executable_NumOutputs(&args);
      if (err == nullptr) {
        info.num_outputs_known = true;
        info.num_outputs = args.num_outputs;
      } else {
        PjrtErrorDestroy(api, err);
      }
    }

    if (PJRT_HAS(api, PJRT_Executable_Fingerprint)) {
      PJRT_Executable_Fingerprint_Args args;
      std::memset(&args, 0, sizeof(args));
      args.struct_size = PJRT_Executable_Fingerprint_Args_STRUCT_SIZE;
      args.executable = exec;
      PJRT_Error* err = api->PJRT_Executable_Fingerprint(&args);
      if (err == nullptr && args.executable_fingerprint != nullptr) {
        info.fingerprint.assign(args.executable_fingerprint, args.executable_fingerprint_size);
      } else if (err != nullptr) {
        PjrtErrorDestroy(api, err);
      }
    }

    if (PJRT_HAS(api, PJRT_Executable_SizeOfGeneratedCodeInBytes)) {
      PJRT_Executable_SizeOfGeneratedCodeInBytes_Args args;
      std::memset(&args, 0, sizeof(args));
      args.struct_size = PJRT_Executable_SizeOfGeneratedCodeInBytes_Args_STRUCT_SIZE;
      args.executable = exec;
      PJRT_Error* err = api->PJRT_Executable_SizeOfGeneratedCodeInBytes(&args);
      if (err == nullptr) {
        info.generated_code_known = true;
        info.generated_code_bytes = args.size_in_bytes;
      } else {
        PjrtErrorDestroy(api, err);
      }
    }

    return info;
  }
}  // namespace tenzr

#endif  // TENZR_RUNTIME_PJRT_EXECUTABLE_H_