// 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_DEVICE_BUFFER_H_
#define TENZR_RUNTIME_DEVICE_BUFFER_H_
#include <xla/pjrt/c/pjrt_c_api.h>
#include <cstring>
#include <memory>
#include <vector>
#include "tenzr/common/error.h"
namespace tenzr::runtime {
class DeviceBuffer {
public:
DeviceBuffer(PJRT_Buffer* buffer, const PJRT_Api* api) : buffer_(buffer), api_(api) {}
~DeviceBuffer() {
if (buffer_ && api_) {
if (api_->PJRT_Buffer_ReadyEvent && api_->PJRT_Event_Await && api_->PJRT_Event_Destroy) {
PJRT_Buffer_ReadyEvent_Args ready_args;
std::memset(&ready_args, 0, sizeof(ready_args));
ready_args.struct_size = PJRT_Buffer_ReadyEvent_Args_STRUCT_SIZE;
ready_args.buffer = buffer_;
PJRT_Error* err = api_->PJRT_Buffer_ReadyEvent(&ready_args);
if (!err && ready_args.event) {
PJRT_Event_Await_Args await_args{PJRT_Event_Await_Args_STRUCT_SIZE, nullptr, ready_args.event};
(void)api_->PJRT_Event_Await(&await_args);
PJRT_Event_Destroy_Args d_args{PJRT_Event_Destroy_Args_STRUCT_SIZE, nullptr, ready_args.event};
(void)api_->PJRT_Event_Destroy(&d_args);
} else if (err && api_->PJRT_Error_Destroy) {
PJRT_Error_Destroy_Args d_args{PJRT_Error_Destroy_Args_STRUCT_SIZE, nullptr, err};
api_->PJRT_Error_Destroy(&d_args);
}
}
if (api_->PJRT_Buffer_Destroy) {
PJRT_Buffer_Destroy_Args args;
std::memset(&args, 0, sizeof(args));
args.struct_size = PJRT_Buffer_Destroy_Args_STRUCT_SIZE;
args.buffer = buffer_;
(void)api_->PJRT_Buffer_Destroy(&args);
buffer_ = nullptr;
}
}
}
DeviceBuffer(const DeviceBuffer&) = delete;
DeviceBuffer& operator=(const DeviceBuffer&) = delete;
DeviceBuffer(DeviceBuffer&& other) noexcept : buffer_(other.buffer_), api_(other.api_) { other.buffer_ = nullptr; }
DeviceBuffer& operator=(DeviceBuffer&& other) noexcept {
if (this != &other) {
if (buffer_ && api_ && api_->PJRT_Buffer_Destroy) {
PJRT_Buffer_Destroy_Args args;
std::memset(&args, 0, sizeof(args));
args.struct_size = PJRT_Buffer_Destroy_Args_STRUCT_SIZE;
args.buffer = buffer_;
(void)api_->PJRT_Buffer_Destroy(&args);
}
buffer_ = other.buffer_;
api_ = other.api_;
other.buffer_ = nullptr;
}
return *this;
}
PJRT_Buffer* Handle() const { return buffer_; }
const PJRT_Api* Api() const { return api_; }
void* UnsafePointer() const {
TENZR_CHECK(buffer_ != nullptr && api_ != nullptr && api_->PJRT_Buffer_UnsafePointer != nullptr,
"DeviceBuffer::unsafe_pointer is unavailable for this PJRT plugin");
PJRT_Buffer_UnsafePointer_Args args;
std::memset(&args, 0, sizeof(args));
args.struct_size = PJRT_Buffer_UnsafePointer_Args_STRUCT_SIZE;
args.buffer = buffer_;
PJRT_Error* err = api_->PJRT_Buffer_UnsafePointer(&args);
if (err != nullptr) {
DestroyError(err);
TENZR_THROW("PJRT_Buffer_UnsafePointer failed");
}
return reinterpret_cast<void*>(args.buffer_pointer);
}
size_t DeviceSizeBytes() const {
if (buffer_ == nullptr || api_ == nullptr || api_->PJRT_Buffer_OnDeviceSizeInBytes == nullptr) {
return 0;
}
PJRT_Buffer_OnDeviceSizeInBytes_Args args;
std::memset(&args, 0, sizeof(args));
args.struct_size = PJRT_Buffer_OnDeviceSizeInBytes_Args_STRUCT_SIZE;
args.buffer = buffer_;
PJRT_Error* err = api_->PJRT_Buffer_OnDeviceSizeInBytes(&args);
if (err != nullptr) {
DestroyError(err);
return 0;
}
return args.on_device_size_in_bytes;
}
void AwaitReady() const {
if (buffer_ == nullptr || api_ == nullptr || api_->PJRT_Buffer_ReadyEvent == nullptr || api_->PJRT_Event_Await == nullptr) {
return;
}
PJRT_Buffer_ReadyEvent_Args ready_args;
std::memset(&ready_args, 0, sizeof(ready_args));
ready_args.struct_size = PJRT_Buffer_ReadyEvent_Args_STRUCT_SIZE;
ready_args.buffer = buffer_;
PJRT_Error* err = api_->PJRT_Buffer_ReadyEvent(&ready_args);
if (err != nullptr) {
DestroyError(err);
return;
}
if (ready_args.event != nullptr) {
PJRT_Event_Await_Args await_args{PJRT_Event_Await_Args_STRUCT_SIZE, nullptr, ready_args.event};
(void)api_->PJRT_Event_Await(&await_args);
if (api_->PJRT_Event_Destroy != nullptr) {
PJRT_Event_Destroy_Args d_args{PJRT_Event_Destroy_Args_STRUCT_SIZE, nullptr, ready_args.event};
(void)api_->PJRT_Event_Destroy(&d_args);
}
}
}
int DeviceId() const {
if (buffer_ == nullptr || api_ == nullptr || api_->PJRT_Buffer_Device == nullptr || api_->PJRT_Device_LocalHardwareId == nullptr) {
return 0;
}
PJRT_Buffer_Device_Args dev_args;
std::memset(&dev_args, 0, sizeof(dev_args));
dev_args.struct_size = PJRT_Buffer_Device_Args_STRUCT_SIZE;
dev_args.buffer = buffer_;
PJRT_Error* err = api_->PJRT_Buffer_Device(&dev_args);
if (err != nullptr) {
DestroyError(err);
return 0;
}
if (dev_args.device == nullptr) {
return 0;
}
PJRT_Device_LocalHardwareId_Args id_args;
std::memset(&id_args, 0, sizeof(id_args));
id_args.struct_size = PJRT_Device_LocalHardwareId_Args_STRUCT_SIZE;
id_args.device = dev_args.device;
err = api_->PJRT_Device_LocalHardwareId(&id_args);
if (err != nullptr) {
DestroyError(err);
return 0;
}
return static_cast<int>(id_args.local_hardware_id);
}
static std::shared_ptr<DeviceBuffer> CreateView(const PJRT_Api* api, PJRT_Client* client, PJRT_Device* device, void* device_buffer_ptr,
const std::vector<int64_t>& dims, PJRT_Buffer_Type element_type, intptr_t stream,
void (*on_delete)(void*, void*), void* on_delete_arg) {
TENZR_CHECK(api != nullptr && client != nullptr && api->PJRT_Client_CreateViewOfDeviceBuffer != nullptr,
"PJRT_Client_CreateViewOfDeviceBuffer is unavailable for this PJRT plugin");
PJRT_Client_CreateViewOfDeviceBuffer_Args args;
std::memset(&args, 0, sizeof(args));
args.struct_size = PJRT_Client_CreateViewOfDeviceBuffer_Args_STRUCT_SIZE;
args.client = client;
args.device_buffer_ptr = device_buffer_ptr;
args.dims = dims.data();
args.num_dims = dims.size();
args.element_type = element_type;
args.layout = nullptr;
args.device = device;
args.on_delete_callback = on_delete;
args.on_delete_callback_arg = on_delete_arg;
args.stream = stream;
args.memory = nullptr;
PJRT_Error* err = api->PJRT_Client_CreateViewOfDeviceBuffer(&args);
if (err != nullptr) {
std::string msg = "PJRT_Client_CreateViewOfDeviceBuffer failed";
if (api->PJRT_Error_Message != nullptr) {
PJRT_Error_Message_Args m_args;
std::memset(&m_args, 0, sizeof(m_args));
m_args.struct_size = PJRT_Error_Message_Args_STRUCT_SIZE;
m_args.error = err;
api->PJRT_Error_Message(&m_args);
if (m_args.message != nullptr) {
msg += ": ";
msg.append(m_args.message, m_args.message_size);
}
}
if (api->PJRT_Error_Destroy != nullptr) {
PJRT_Error_Destroy_Args d_args{PJRT_Error_Destroy_Args_STRUCT_SIZE, nullptr, err};
api->PJRT_Error_Destroy(&d_args);
}
TENZR_THROW(msg);
}
TENZR_CHECK(args.buffer != nullptr, "PJRT_Client_CreateViewOfDeviceBuffer returned a null buffer");
return std::make_shared<DeviceBuffer>(args.buffer, api);
}
void IncreaseExternalReferenceCount() const {
if (buffer_ == nullptr || api_ == nullptr || api_->PJRT_Buffer_IncreaseExternalReferenceCount == nullptr) {
return;
}
PJRT_Buffer_IncreaseExternalReferenceCount_Args args;
std::memset(&args, 0, sizeof(args));
args.struct_size = PJRT_Buffer_IncreaseExternalReferenceCount_Args_STRUCT_SIZE;
args.buffer = buffer_;
PJRT_Error* err = api_->PJRT_Buffer_IncreaseExternalReferenceCount(&args);
if (err != nullptr) {
DestroyError(err);
}
}
void DecreaseExternalReferenceCount() const {
if (buffer_ == nullptr || api_ == nullptr || api_->PJRT_Buffer_DecreaseExternalReferenceCount == nullptr) {
return;
}
PJRT_Buffer_DecreaseExternalReferenceCount_Args args;
std::memset(&args, 0, sizeof(args));
args.struct_size = PJRT_Buffer_DecreaseExternalReferenceCount_Args_STRUCT_SIZE;
args.buffer = buffer_;
PJRT_Error* err = api_->PJRT_Buffer_DecreaseExternalReferenceCount(&args);
if (err != nullptr) {
DestroyError(err);
}
}
private:
void DestroyError(PJRT_Error* err) const {
if (err != nullptr && api_ != nullptr && api_->PJRT_Error_Destroy != nullptr) {
PJRT_Error_Destroy_Args d_args{PJRT_Error_Destroy_Args_STRUCT_SIZE, nullptr, err};
api_->PJRT_Error_Destroy(&d_args);
}
}
PJRT_Buffer* buffer_{nullptr};
const PJRT_Api* api_{nullptr};
};
} // namespace tenzr::runtime
namespace tenzr {
using runtime::DeviceBuffer;
}
#endif // TENZR_RUNTIME_DEVICE_BUFFER_H_