856 lines
30 KiB
Zig
856 lines
30 KiB
Zig
const builtin = @import("builtin");
|
|
const std = @import("std");
|
|
|
|
const c = @import("c");
|
|
|
|
const log = std.log.scoped(.pjrt);
|
|
|
|
pub const Profiler = @import("profiler.zig").Profiler;
|
|
|
|
test {
|
|
std.testing.refAllDecls(@This());
|
|
}
|
|
|
|
// We could calculate it like PJRT does, but it turns out that some of those
|
|
// were wrong in PJRT itself [1], which gets propagated to binary plugins. In
|
|
// order to mirror that, we just the value as computed by PJRT itself, through
|
|
// comptime reflection. We could make the argument to remove that one day since
|
|
// [1] has been fixed. The problem is that this problem could happen again in
|
|
// as the way PJRT does it is not very robust.
|
|
//
|
|
// 1. https://github.com/openxla/xla/issues/10032
|
|
fn pjrtStructSize(comptime T: type) usize {
|
|
// unsafe on purpose, we want this to fail if that ever changes
|
|
const typedef_name = comptime blk: {
|
|
const needle = ".struct_";
|
|
const idx = std.mem.indexOf(u8, @typeName(T), needle).?;
|
|
break :blk @typeName(T)[idx + needle.len ..];
|
|
};
|
|
return @field(c, typedef_name ++ "_STRUCT_SIZE");
|
|
}
|
|
|
|
inline fn pjrtStruct(v: anytype) @TypeOf(v) {
|
|
var ret = v;
|
|
ret.struct_size = pjrtStructSize(@TypeOf(v));
|
|
return ret;
|
|
}
|
|
|
|
pub const ApiError = error{
|
|
Cancelled,
|
|
Unknown,
|
|
InvalidArgument,
|
|
DeadlineExceeded,
|
|
NotFound,
|
|
AlreadyExists,
|
|
PermissionDenied,
|
|
ResourceExhausted,
|
|
FailedPrecondition,
|
|
Aborted,
|
|
OutOfRange,
|
|
Unimplemented,
|
|
Internal,
|
|
Unavailable,
|
|
DataLoss,
|
|
Unauthenticated,
|
|
};
|
|
|
|
fn InnerMixin(comptime innerT: type) type {
|
|
return struct {
|
|
fn inner(self: anytype) *innerT {
|
|
return @ptrCast(@constCast(@alignCast(self)));
|
|
}
|
|
};
|
|
}
|
|
|
|
pub const Api = struct {
|
|
pub const Version = struct {
|
|
major: i64,
|
|
minor: i64,
|
|
};
|
|
|
|
const Funcs = std.meta.FieldEnum(c.PJRT_Api);
|
|
|
|
inner: c.PJRT_Api,
|
|
|
|
pub fn loadFrom(library: []const u8) !*const Api {
|
|
var lib: std.DynLib = switch (builtin.os.tag) {
|
|
.linux => blk: {
|
|
const library_c = try std.posix.toPosixPath(library);
|
|
break :blk .{
|
|
.inner = .{
|
|
.handle = c.dlopen(&library_c, c.RTLD_LAZY | c.RTLD_LOCAL | c.RTLD_NODELETE) orelse {
|
|
return error.FileNotFound;
|
|
},
|
|
},
|
|
};
|
|
},
|
|
else => try std.DynLib.open(library),
|
|
};
|
|
const DynGetPjrtApi = lib.lookup(*const fn () callconv(.C) *const Api, "GetPjrtApi") orelse {
|
|
std.debug.panic("Unable to find GetPjrtApi symbol in library: {s}", .{library});
|
|
};
|
|
|
|
const api = DynGetPjrtApi();
|
|
log.info("Loaded library: {s}", .{library});
|
|
_ = api.call(.PJRT_Plugin_Initialize, .{}) catch unreachable;
|
|
|
|
return api;
|
|
}
|
|
|
|
fn CallFnArgType(comptime func: Funcs) type {
|
|
const fti = @typeInfo(std.meta.FieldType(c.PJRT_Api, func));
|
|
const fn_ptr = @typeInfo(fti.Optional.child);
|
|
const fn_type_info = @typeInfo(fn_ptr.Pointer.child);
|
|
const arg_array_type_info = @typeInfo(fn_type_info.Fn.params[0].type.?);
|
|
return arg_array_type_info.Pointer.child;
|
|
}
|
|
|
|
inline fn call(self: *const Api, comptime method: Funcs, arg: CallFnArgType(method)) ApiError!@TypeOf(arg) {
|
|
var ret = pjrtStruct(arg);
|
|
const fn_ptr = @field(&self.inner, @tagName(method)).?;
|
|
const result = fn_ptr(&ret);
|
|
if (@TypeOf(result) == void) {
|
|
return ret;
|
|
}
|
|
|
|
if (result) |pjrt_c_error| {
|
|
const pjrt_error: *Error = @ptrCast(pjrt_c_error);
|
|
log.err("[{s}] {s}", .{ @tagName(method), pjrt_error.getMessage(self) });
|
|
return pjrt_error.getCode(self).toApiError();
|
|
}
|
|
|
|
return ret;
|
|
}
|
|
|
|
pub fn lookupExtension(self: *const Api, comptime ExtensionT: type, ext_id: c_int) ?*const ExtensionT {
|
|
var cur: [*c]const c.PJRT_Extension_Base = @alignCast(@ptrCast(self.inner.extension_start));
|
|
while (cur != null) : (cur = cur.*.next) {
|
|
if (cur.*.type == ext_id) {
|
|
return @alignCast(@ptrCast(cur));
|
|
}
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
pub inline fn version(self: *const Api) Version {
|
|
return .{
|
|
.major = @intCast(self.inner.pjrt_api_version.major_version),
|
|
.minor = @intCast(self.inner.pjrt_api_version.minor_version),
|
|
};
|
|
}
|
|
};
|
|
|
|
pub const ErrorCode = enum(c.PJRT_Error_Code) {
|
|
cancelled = c.PJRT_Error_Code_CANCELLED,
|
|
unknown = c.PJRT_Error_Code_UNKNOWN,
|
|
invalid_argument = c.PJRT_Error_Code_INVALID_ARGUMENT,
|
|
deadline_exceeded = c.PJRT_Error_Code_DEADLINE_EXCEEDED,
|
|
not_found = c.PJRT_Error_Code_NOT_FOUND,
|
|
already_exists = c.PJRT_Error_Code_ALREADY_EXISTS,
|
|
permission_denied = c.PJRT_Error_Code_PERMISSION_DENIED,
|
|
resource_exhausted = c.PJRT_Error_Code_RESOURCE_EXHAUSTED,
|
|
failed_precondition = c.PJRT_Error_Code_FAILED_PRECONDITION,
|
|
aborted = c.PJRT_Error_Code_ABORTED,
|
|
out_of_range = c.PJRT_Error_Code_OUT_OF_RANGE,
|
|
unimplemented = c.PJRT_Error_Code_UNIMPLEMENTED,
|
|
internal = c.PJRT_Error_Code_INTERNAL,
|
|
unavailable = c.PJRT_Error_Code_UNAVAILABLE,
|
|
data_loss = c.PJRT_Error_Code_DATA_LOSS,
|
|
unauthenticated = c.PJRT_Error_Code_UNAUTHENTICATED,
|
|
|
|
pub fn toApiError(code: ErrorCode) ApiError {
|
|
return switch (code) {
|
|
.cancelled => error.Cancelled,
|
|
.unknown => error.Unknown,
|
|
.invalid_argument => error.InvalidArgument,
|
|
.deadline_exceeded => error.DeadlineExceeded,
|
|
.not_found => error.NotFound,
|
|
.already_exists => error.AlreadyExists,
|
|
.permission_denied => error.PermissionDenied,
|
|
.resource_exhausted => error.ResourceExhausted,
|
|
.failed_precondition => error.FailedPrecondition,
|
|
.aborted => error.Aborted,
|
|
.out_of_range => error.OutOfRange,
|
|
.unimplemented => error.Unimplemented,
|
|
.internal => error.Internal,
|
|
.unavailable => error.Unavailable,
|
|
.data_loss => error.DataLoss,
|
|
.unauthenticated => error.Unauthenticated,
|
|
};
|
|
}
|
|
};
|
|
|
|
pub const Error = opaque {
|
|
pub fn deinit(self: *Error, api: *const Api) void {
|
|
_ = api.call(.PJRT_Error_Destroy, .{
|
|
.@"error" = @ptrCast(self),
|
|
}) catch unreachable;
|
|
}
|
|
|
|
pub fn getCode(self: *Error, api: *const Api) ErrorCode {
|
|
const ret = api.call(.PJRT_Error_GetCode, .{
|
|
.@"error" = @ptrCast(self),
|
|
}) catch unreachable;
|
|
return @enumFromInt(ret.code);
|
|
}
|
|
|
|
pub fn getMessage(self: *Error, api: *const Api) []const u8 {
|
|
const ret = api.call(.PJRT_Error_Message, .{
|
|
.@"error" = @ptrCast(self),
|
|
}) catch unreachable;
|
|
return ret.message[0..ret.message_size];
|
|
}
|
|
};
|
|
|
|
pub const ClientInitError = error{LoadingFailed} || ApiError;
|
|
|
|
pub const Client = opaque {
|
|
const inner = InnerMixin(c.PJRT_Client).inner;
|
|
|
|
pub const ProgramFormat = enum {
|
|
hlo,
|
|
mlir,
|
|
};
|
|
|
|
pub fn init(api: *const Api, create_options: []const NamedValue) ClientInitError!*Client {
|
|
// log.info("Loaded PJRT runtime plugin: {s}", .{api.Platform});
|
|
const ret = try api.call(.PJRT_Client_Create, .{
|
|
.create_options = @ptrCast(create_options.ptr),
|
|
.num_options = create_options.len,
|
|
.kv_get_callback = null,
|
|
.kv_put_callback = null,
|
|
.kv_put_user_arg = null,
|
|
.kv_get_user_arg = null,
|
|
});
|
|
return @ptrCast(ret.client.?);
|
|
}
|
|
|
|
pub fn deinit(self: *Client, api: *const Api) void {
|
|
_ = api.call(.PJRT_Client_Destroy, .{
|
|
.client = self.inner(),
|
|
}) catch {};
|
|
}
|
|
|
|
pub fn getPlatformName(self: *const Client, api: *const Api) []const u8 {
|
|
const ret = api.call(.PJRT_Client_PlatformName, .{
|
|
.client = self.inner(),
|
|
}) catch unreachable;
|
|
return ret.platform_name[0..ret.platform_name_size];
|
|
}
|
|
|
|
pub fn getDevices(self: *const Client, api: *const Api) []const *Device {
|
|
const ret = api.call(.PJRT_Client_Devices, .{
|
|
.client = self.inner(),
|
|
}) catch unreachable;
|
|
return @ptrCast(ret.devices[0..ret.num_devices]);
|
|
}
|
|
|
|
pub fn getAddressableDevices(self: *const Client, api: *const Api) []const *Device {
|
|
const ret = api.call(.PJRT_Client_AddressableDevices, .{
|
|
.client = self.inner(),
|
|
}) catch unreachable;
|
|
return @ptrCast(ret.addressable_devices[0..ret.num_addressable_devices]);
|
|
}
|
|
|
|
pub const CompileArgs = struct {
|
|
bytecode: []const u8,
|
|
bytecode_format: ProgramFormat,
|
|
compile_options_pb: []const u8,
|
|
};
|
|
|
|
pub fn compile(self: *const Client, api: *const Api, args: CompileArgs) ApiError!*LoadedExecutable {
|
|
const bytecode_format_ = @tagName(args.bytecode_format);
|
|
const ret = try api.call(.PJRT_Client_Compile, .{
|
|
.program = &pjrtStruct(c.PJRT_Program{
|
|
.code = @ptrCast(@constCast(args.bytecode.ptr)),
|
|
.code_size = args.bytecode.len,
|
|
.format = @ptrCast(@constCast(bytecode_format_.ptr)),
|
|
.format_size = bytecode_format_.len,
|
|
}),
|
|
.compile_options = @ptrCast(@constCast(args.compile_options_pb.ptr)),
|
|
.compile_options_size = args.compile_options_pb.len,
|
|
.client = self.inner(),
|
|
});
|
|
return @ptrCast(ret.executable.?);
|
|
}
|
|
|
|
pub const BufferFromHostBufferArgs = struct {
|
|
data: []const u8,
|
|
buffer_type: BufferType,
|
|
dims: []const i64,
|
|
byte_strides: ?[]const i64,
|
|
device: *const Device,
|
|
host_buffer_semantics: HostBufferSemantics,
|
|
};
|
|
|
|
pub fn bufferFromHostBuffer(self: *const Client, api: *const Api, args: BufferFromHostBufferArgs) ApiError!struct { *Buffer, *Event } {
|
|
const ret = try api.call(.PJRT_Client_BufferFromHostBuffer, .{
|
|
.client = self.inner(),
|
|
.data = @ptrCast(@constCast(args.data.ptr)),
|
|
.type = @intFromEnum(args.buffer_type),
|
|
.dims = @ptrCast(@constCast(args.dims.ptr)),
|
|
.num_dims = args.dims.len,
|
|
.byte_strides = if (args.byte_strides) |bs| @ptrCast(@constCast(bs.ptr)) else null,
|
|
.num_byte_strides = if (args.byte_strides) |bs| bs.len else 0,
|
|
.host_buffer_semantics = @intFromEnum(args.host_buffer_semantics),
|
|
.device = @ptrCast(@constCast(args.device)),
|
|
.memory = null, // TODO
|
|
.device_layout = null, // TODO
|
|
.done_with_host_buffer = null,
|
|
.buffer = null,
|
|
});
|
|
return .{
|
|
@ptrCast(ret.buffer.?),
|
|
@ptrCast(ret.done_with_host_buffer.?),
|
|
};
|
|
}
|
|
|
|
/// Returns the Profiler for this API.
|
|
/// Not all platform have a profiling api, for those the profiler object will do nothing.
|
|
/// Platforms with known profiler extensions: cuda, xpu
|
|
pub fn getProfiler(self: *const Client, api: *const Api, options: Profiler.Options) Profiler {
|
|
if (api.version().minor >= 45) {
|
|
if (api.lookupExtension(c.PJRT_Profiler_Extension, c.PJRT_Extension_Type_Profiler)) |ext| {
|
|
return Profiler.init(ext.profiler_api.*, options);
|
|
}
|
|
}
|
|
log.warn("No profiler found for platform: {}", .{self});
|
|
return Profiler.init(null, options);
|
|
}
|
|
|
|
pub fn deserializeAndLoad(self: *const Client, api: *const Api, bytes: []const u8) ApiError!*LoadedExecutable {
|
|
const ret = try api.call(.PJRT_Executable_DeserializeAndLoad, .{
|
|
.client = self.inner(),
|
|
.serialized_executable = bytes.ptr,
|
|
.serialized_executable_size = bytes.len,
|
|
});
|
|
return @ptrCast(ret.loaded_executable.?);
|
|
}
|
|
|
|
pub const CreateViewOfDeviceBufferArgs = struct {
|
|
data: []const u8,
|
|
dims: []const i64,
|
|
element_type: BufferType,
|
|
layout: MemoryLayout,
|
|
device: *const Device,
|
|
on_delete_callback: ?*const fn (device_buffer_ptr: ?*anyopaque, ctx: ?*anyopaque) callconv(.C) void = null,
|
|
on_delete_callback_arg: ?*anyopaque = null,
|
|
stream: ?isize = null,
|
|
};
|
|
|
|
pub fn createViewOfDeviceBuffer(self: *const Client, api: *const Api, args: CreateViewOfDeviceBufferArgs) ApiError!*Buffer {
|
|
const layout = args.layout.toCStruct();
|
|
const ret = try api.call(.PJRT_Client_CreateViewOfDeviceBuffer, .{
|
|
.client = self.inner(),
|
|
.device_buffer_ptr = @ptrCast(@constCast(args.data.ptr)),
|
|
.dims = args.dims.ptr,
|
|
.num_dims = args.dims.len,
|
|
.element_type = @intFromEnum(args.element_type),
|
|
.layout = @ptrCast(@constCast(&layout)),
|
|
.device = @ptrCast(@constCast(args.device)),
|
|
.on_delete_callback = args.on_delete_callback,
|
|
.on_delete_callback_arg = args.on_delete_callback_arg,
|
|
.stream = if (args.stream) |stream| stream else 0,
|
|
});
|
|
return @ptrCast(ret.buffer.?);
|
|
}
|
|
};
|
|
|
|
pub const Device = opaque {
|
|
const inner = InnerMixin(c.PJRT_Device).inner;
|
|
|
|
pub fn getDescription(self: *const Device, api: *const Api) *const DeviceDescription {
|
|
const ret = api.call(.PJRT_Device_GetDescription, .{
|
|
.device = self.inner(),
|
|
}) catch unreachable;
|
|
return @ptrCast(ret.device_description.?);
|
|
}
|
|
|
|
pub fn isAddressable(self: *const Device, api: *const Api) bool {
|
|
const ret = api.call(.PJRT_Device_IsAddressable, .{
|
|
.device = self.inner(),
|
|
}) catch unreachable;
|
|
return ret.is_addressable;
|
|
}
|
|
|
|
pub fn getLocalHardwareId(self: *const Device, api: *const Api) usize {
|
|
const ret = api.call(.PJRT_Device_LocalHardwareId, .{
|
|
.device = self.inner(),
|
|
}) catch unreachable;
|
|
return @intCast(ret.local_hardware_id);
|
|
}
|
|
};
|
|
|
|
pub const DeviceDescription = opaque {
|
|
const inner = InnerMixin(c.PJRT_DeviceDescription).inner;
|
|
|
|
pub fn getId(self: *const DeviceDescription, api: *const Api) usize {
|
|
const ret = api.call(.PJRT_DeviceDescription_Id, .{
|
|
.device_description = self.inner(),
|
|
}) catch unreachable;
|
|
return @intCast(ret.id);
|
|
}
|
|
|
|
pub fn getProcessIndex(self: *const DeviceDescription, api: *const Api) usize {
|
|
const ret = api.call(.PJRT_DeviceDescription_ProcessIndex, .{
|
|
.device_description = self.inner(),
|
|
}) catch unreachable;
|
|
return @intCast(ret.process_index);
|
|
}
|
|
|
|
pub fn getKind(self: *const DeviceDescription, api: *const Api) []const u8 {
|
|
const ret = api.call(.PJRT_DeviceDescription_Kind, .{
|
|
.device_description = self.inner(),
|
|
}) catch unreachable;
|
|
return ret.device_kind[0..ret.device_kind_size];
|
|
}
|
|
|
|
pub fn debugString(self: *const DeviceDescription, api: *const Api) []const u8 {
|
|
const ret = api.call(.PJRT_DeviceDescription_DebugString, .{
|
|
.device_description = self.inner(),
|
|
}) catch unreachable;
|
|
return ret.debug_string[0..ret.debug_string_size];
|
|
}
|
|
|
|
pub fn toString(self: *const DeviceDescription, api: *const Api) []const u8 {
|
|
const ret = api.call(.PJRT_DeviceDescription_ToString, .{
|
|
.device_description = self.inner(),
|
|
}) catch unreachable;
|
|
return ret.to_string[0..ret.to_string_size];
|
|
}
|
|
};
|
|
|
|
pub const GetCostAnalysisError = std.mem.Allocator.Error || ApiError;
|
|
|
|
pub const SerializeResult = struct {
|
|
bytes: []const u8,
|
|
handle: *anyopaque,
|
|
deleter: *const fn (?*anyopaque) callconv(.C) void,
|
|
|
|
pub fn deinit(self: *SerializeResult) void {
|
|
self.deleter(self.handle);
|
|
self.bytes = &.{};
|
|
self.* = undefined;
|
|
}
|
|
};
|
|
|
|
pub const Executable = opaque {
|
|
const inner = InnerMixin(c.PJRT_Executable).inner;
|
|
|
|
pub fn deinit(self: *Executable, api: *const Api) void {
|
|
_ = api.call(.PJRT_Executable_Destroy, .{
|
|
.executable = self.inner(),
|
|
}) catch unreachable;
|
|
}
|
|
|
|
pub fn getCostAnalysis(self: *const Executable, api: *const Api) GetCostAnalysisError![]*const NamedValue {
|
|
const ret = try api.call(.PJRT_Executable_GetCostAnalysis, .{
|
|
.executable = self.inner(),
|
|
});
|
|
const values: [*]*const NamedValue = @ptrCast(ret.properties);
|
|
return values[0..ret.num_properties];
|
|
}
|
|
|
|
pub fn serialize(self: *const Executable, api: *const Api) ApiError!SerializeResult {
|
|
const ret = try api.call(.PJRT_Executable_Serialize, .{
|
|
.executable = self.inner(),
|
|
});
|
|
|
|
return .{
|
|
.bytes = ret.serialized_bytes[0..ret.serialized_bytes_size],
|
|
.handle = ret.serialized_executable.?,
|
|
.deleter = @ptrCast(ret.serialized_executable_deleter.?),
|
|
};
|
|
}
|
|
};
|
|
|
|
pub const LoadedExecutable = opaque {
|
|
const inner = InnerMixin(c.PJRT_LoadedExecutable).inner;
|
|
|
|
pub fn deinit(self: *LoadedExecutable, api: *const Api) void {
|
|
_ = api.call(.PJRT_LoadedExecutable_Destroy, .{
|
|
.executable = self.inner(),
|
|
}) catch {};
|
|
self.* = undefined;
|
|
}
|
|
|
|
pub fn delete(self: *LoadedExecutable, api: *const Api) void {
|
|
_ = api.call(.PJRT_LoadedExecutable_Delete, .{
|
|
.executable = self.inner(),
|
|
}) catch unreachable;
|
|
}
|
|
|
|
pub fn isDeleted(self: *const LoadedExecutable, api: *const Api) bool {
|
|
const ret = api.call(.PJRT_LoadedExecutable_IsDeleted, .{
|
|
.executable = self.inner(),
|
|
}) catch unreachable;
|
|
return ret.is_deleted;
|
|
}
|
|
|
|
pub fn getAddressableDevices(self: *const LoadedExecutable, api: *const Api) []Device {
|
|
const ret = api.call(.PJRT_LoadedExecutable_AddressableDevices, .{
|
|
.executable = self.inner(),
|
|
}) catch unreachable;
|
|
return @ptrCast(ret.addressable_devices);
|
|
}
|
|
|
|
pub fn execute(self: *const LoadedExecutable, api: *const Api, args: struct {
|
|
num_args: usize,
|
|
arguments: []const [*]const *const Buffer,
|
|
results: []const [*]*Buffer,
|
|
events: []*Event,
|
|
non_donatable_input_indices: []const i64 = &.{},
|
|
}) ApiError!void {
|
|
var options = pjrtStruct(c.PJRT_ExecuteOptions{
|
|
.send_callbacks = null,
|
|
.recv_callbacks = null,
|
|
.num_send_ops = 0,
|
|
.num_recv_ops = 0,
|
|
.launch_id = 0,
|
|
.non_donatable_input_indices = @ptrCast(args.non_donatable_input_indices.ptr),
|
|
.num_non_donatable_input_indices = args.non_donatable_input_indices.len,
|
|
});
|
|
_ = try api.call(.PJRT_LoadedExecutable_Execute, .{
|
|
.executable = self.inner(),
|
|
.options = @ptrCast(&options),
|
|
.argument_lists = @ptrCast(args.arguments.ptr),
|
|
.num_devices = @intCast(args.arguments.len),
|
|
.num_args = args.num_args,
|
|
.output_lists = @ptrCast(args.results.ptr),
|
|
.device_complete_events = @ptrCast(args.events.ptr),
|
|
.execute_device = null,
|
|
});
|
|
}
|
|
|
|
pub fn getExecutable(self: *LoadedExecutable, api: *const Api) ApiError!*Executable {
|
|
const ret = try api.call(.PJRT_LoadedExecutable_GetExecutable, .{
|
|
.loaded_executable = self.inner(),
|
|
});
|
|
return @ptrCast(ret.executable.?);
|
|
}
|
|
};
|
|
|
|
pub const BufferType = enum(c.PJRT_Buffer_Type) {
|
|
INVALID = c.PJRT_Buffer_Type_INVALID,
|
|
PRED = c.PJRT_Buffer_Type_PRED,
|
|
S8 = c.PJRT_Buffer_Type_S8,
|
|
S16 = c.PJRT_Buffer_Type_S16,
|
|
S32 = c.PJRT_Buffer_Type_S32,
|
|
S64 = c.PJRT_Buffer_Type_S64,
|
|
U8 = c.PJRT_Buffer_Type_U8,
|
|
U16 = c.PJRT_Buffer_Type_U16,
|
|
U32 = c.PJRT_Buffer_Type_U32,
|
|
U64 = c.PJRT_Buffer_Type_U64,
|
|
F16 = c.PJRT_Buffer_Type_F16,
|
|
F32 = c.PJRT_Buffer_Type_F32,
|
|
F64 = c.PJRT_Buffer_Type_F64,
|
|
BF16 = c.PJRT_Buffer_Type_BF16,
|
|
C64 = c.PJRT_Buffer_Type_C64,
|
|
C128 = c.PJRT_Buffer_Type_C128,
|
|
F8E5M2 = c.PJRT_Buffer_Type_F8E5M2,
|
|
F8E4M3FN = c.PJRT_Buffer_Type_F8E4M3FN,
|
|
F8E4M3B11FNUZ = c.PJRT_Buffer_Type_F8E4M3B11FNUZ,
|
|
F8E5M2FNUZ = c.PJRT_Buffer_Type_F8E5M2FNUZ,
|
|
F8E4M3FNUZ = c.PJRT_Buffer_Type_F8E4M3FNUZ,
|
|
S4 = c.PJRT_Buffer_Type_S4,
|
|
U4 = c.PJRT_Buffer_Type_U4,
|
|
};
|
|
|
|
pub const MemoryLayoutType = enum(c.PJRT_Buffer_MemoryLayout_Type) {
|
|
Tiled = c.PJRT_Buffer_MemoryLayout_Type_Tiled,
|
|
Strides = c.PJRT_Buffer_MemoryLayout_Type_Strides,
|
|
};
|
|
|
|
pub const MemoryLayout = union(MemoryLayoutType) {
|
|
pub const Type = MemoryLayoutType;
|
|
|
|
pub const Tiled = struct {
|
|
minor_to_major: []const i64,
|
|
tile_dims: []const i64,
|
|
tile_dims_sizes: []const usize,
|
|
};
|
|
|
|
pub const Strides = struct {
|
|
byte_strides: []const i64,
|
|
};
|
|
|
|
Tiled: Tiled,
|
|
Strides: Strides,
|
|
|
|
fn toCStruct(self: MemoryLayout) c.PJRT_Buffer_MemoryLayout {
|
|
return pjrtStruct(switch (self) {
|
|
.Tiled => |v| c.PJRT_Buffer_MemoryLayout{
|
|
.type = c.PJRT_Buffer_MemoryLayout_Type_Tiled,
|
|
.unnamed_0 = .{
|
|
.tiled = c.PJRT_Buffer_MemoryLayout_Tiled{
|
|
.minor_to_major = v.minor_to_major.ptr,
|
|
.minor_to_major_size = v.minor_to_major.len,
|
|
.tile_dims = v.tile_dims.ptr,
|
|
.tile_dim_sizes = v.tile_dims_sizes.ptr,
|
|
.num_tiles = v.tile_dims_sizes.len,
|
|
},
|
|
},
|
|
},
|
|
.Strides => |v| c.PJRT_Buffer_MemoryLayout{
|
|
.type = c.PJRT_Buffer_MemoryLayout_Type_Strides,
|
|
.unnamed_0 = .{
|
|
.strides = c.PJRT_Buffer_MemoryLayout_Strides{
|
|
.byte_strides = v.byte_strides.ptr,
|
|
.num_byte_strides = v.byte_strides.len,
|
|
},
|
|
},
|
|
},
|
|
});
|
|
}
|
|
};
|
|
|
|
pub const HostBufferSemantics = enum(c.PJRT_HostBufferSemantics) {
|
|
ImmutableOnlyDuringCall = c.PJRT_HostBufferSemantics_kImmutableOnlyDuringCall,
|
|
ImmutableUntilTransferCompletes = c.PJRT_HostBufferSemantics_kImmutableUntilTransferCompletes,
|
|
ImmutableZeroCopy = c.PJRT_HostBufferSemantics_kImmutableZeroCopy,
|
|
MutableZeroCopy = c.PJRT_HostBufferSemantics_kMutableZeroCopy,
|
|
};
|
|
|
|
pub const Buffer = opaque {
|
|
const inner = InnerMixin(c.PJRT_Buffer).inner;
|
|
|
|
pub fn deinit(self: *Buffer, api: *const Api) void {
|
|
_ = api.call(.PJRT_Buffer_Destroy, .{
|
|
.buffer = self.inner(),
|
|
}) catch unreachable;
|
|
}
|
|
|
|
pub fn getDevice(self: *const Buffer, api: *const Api) ApiError!*Device {
|
|
const ret = try api.call(.PJRT_Buffer_Device, .{
|
|
.buffer = self.inner(),
|
|
});
|
|
return @ptrCast(ret.device.?);
|
|
}
|
|
|
|
pub fn delete(self: *Buffer, api: *const Api) void {
|
|
_ = api.call(.PJRT_Buffer_Delete, .{
|
|
.buffer = self.inner(),
|
|
}) catch unreachable;
|
|
}
|
|
|
|
pub fn isDeleted(self: *const Buffer, api: *const Api) bool {
|
|
const ret = api.call(.PJRT_Buffer_IsDeleted, .{
|
|
.buffer = self.inner(),
|
|
}) catch unreachable;
|
|
return ret.is_deleted;
|
|
}
|
|
|
|
pub fn isOnCpu(self: *const Buffer, api: *const Api) bool {
|
|
const ret = api.call(.PJRT_Buffer_IsOnCpu, .{
|
|
.buffer = self.inner(),
|
|
}) catch unreachable;
|
|
return ret.is_on_cpu;
|
|
}
|
|
|
|
pub fn toHostBuffer(self: *const Buffer, api: *const Api, dst: []u8) ApiError!*Event {
|
|
const ret = try api.call(.PJRT_Buffer_ToHostBuffer, .{
|
|
.src = self.inner(),
|
|
.dst = @ptrCast(dst.ptr),
|
|
.dst_size = dst.len,
|
|
});
|
|
return @ptrCast(ret.event.?);
|
|
}
|
|
|
|
pub fn getElementType(self: *const Buffer, api: *const Api) BufferType {
|
|
const ret = api.call(.PJRT_Buffer_ElementType, .{
|
|
.buffer = self.inner(),
|
|
}) catch unreachable;
|
|
return @enumFromInt(ret.type);
|
|
}
|
|
|
|
pub fn getDimensions(self: *const Buffer, api: *const Api) []const i64 {
|
|
const ret = api.call(.PJRT_Buffer_Dimensions, .{
|
|
.buffer = self.inner(),
|
|
}) catch unreachable;
|
|
if (ret.num_dims == 0) {
|
|
return &.{};
|
|
}
|
|
return ret.dims[0..ret.num_dims];
|
|
}
|
|
|
|
pub fn getUnpaddedDimensions(self: *const Buffer, api: *const Api) ApiError![]const i64 {
|
|
const ret = try api.call(.PJRT_Buffer_UnpaddedDimensions, .{
|
|
.buffer = self.inner(),
|
|
});
|
|
return ret.dims[0..ret.num_dims];
|
|
}
|
|
|
|
pub fn getOnDeviceSizeInBytes(self: *const Buffer, api: *const Api) ApiError!usize {
|
|
const ret = try api.call(.PJRT_Buffer_OnDeviceSizeInBytes, .{
|
|
.buffer = self.inner(),
|
|
});
|
|
return ret.on_device_size_in_bytes;
|
|
}
|
|
|
|
pub fn copyToDevice(self: *const Buffer, api: *const Api, device: Device) ApiError!Buffer {
|
|
const ret = try api.call(.PJRT_Buffer_CopyToDevice, .{
|
|
.buffer = self.inner(),
|
|
.dst_device = device.inner,
|
|
});
|
|
return @ptrCast(ret.dst_buffer.?);
|
|
}
|
|
|
|
pub fn getReadyEvent(self: *const Buffer, api: *const Api) *Event {
|
|
const ret = api.call(.PJRT_Buffer_ReadyEvent, .{
|
|
.buffer = self.inner(),
|
|
}) catch unreachable;
|
|
return @ptrCast(ret.event.?);
|
|
}
|
|
|
|
pub fn getOpaqueDeviceMemoryDataPointer(self: *const Buffer, api: *const Api) ApiError!*anyopaque {
|
|
const ret = try api.call(.PJRT_Buffer_OpaqueDeviceMemoryDataPointer, .{
|
|
.buffer = self.inner(),
|
|
});
|
|
return ret.device_memory_ptr.?;
|
|
}
|
|
};
|
|
|
|
pub const Event = opaque {
|
|
const inner = InnerMixin(c.PJRT_Event).inner;
|
|
|
|
pub fn deinit(self: *Event, api: *const Api) void {
|
|
_ = api.call(.PJRT_Event_Destroy, .{
|
|
.event = self.inner(),
|
|
}) catch unreachable;
|
|
}
|
|
|
|
pub fn isReady(self: *const Event, api: *const Api) bool {
|
|
const ret = api.call(.PJRT_Event_IsReady, .{
|
|
.event = self.inner(),
|
|
}) catch unreachable;
|
|
return ret.is_ready;
|
|
}
|
|
|
|
pub fn getEventError(self: *const Event, api: *const Api) ApiError!?*Error {
|
|
const ret = try api.call(.PJRT_Event_Error, .{
|
|
.event = self.inner(),
|
|
});
|
|
return @ptrCast(ret);
|
|
}
|
|
|
|
pub fn await_(self: *const Event, api: *const Api) ApiError!void {
|
|
_ = try api.call(.PJRT_Event_Await, .{
|
|
.event = self.inner(),
|
|
});
|
|
}
|
|
|
|
pub fn onReady(self: *Event, api: *const Api, func: *const fn (err: ?*Error, user_arg: ?*anyopaque) callconv(.C) void, user_arg: ?*anyopaque) ApiError!void {
|
|
_ = try api.call(.PJRT_Event_OnReady, .{
|
|
.event = self.inner(),
|
|
.callback = @ptrCast(func),
|
|
.user_arg = user_arg,
|
|
});
|
|
}
|
|
};
|
|
|
|
pub const NamedValue = extern struct {
|
|
comptime {
|
|
std.debug.assert(@sizeOf(NamedValue) == @sizeOf(c.PJRT_NamedValue));
|
|
}
|
|
|
|
inner: c.PJRT_NamedValue,
|
|
|
|
pub const Kind = enum(c.PJRT_NamedValue_Type) {
|
|
string = c.PJRT_NamedValue_kString,
|
|
int64 = c.PJRT_NamedValue_kInt64,
|
|
int64list = c.PJRT_NamedValue_kInt64List,
|
|
float = c.PJRT_NamedValue_kFloat,
|
|
bool = c.PJRT_NamedValue_kBool,
|
|
};
|
|
|
|
pub fn kind(self: NamedValue) Kind {
|
|
return @enumFromInt(self.inner.type);
|
|
}
|
|
|
|
pub fn name(self: NamedValue) []const u8 {
|
|
return self.inner.name[0..self.inner.name_size];
|
|
}
|
|
|
|
pub fn from(name_: []const u8, value: anytype) NamedValue {
|
|
return switch (@TypeOf(value)) {
|
|
[]u8, []const u8 => fromString(name_, value),
|
|
i64 => fromInt64(name_, value),
|
|
[]i64, []const i64 => fromInt64List(name_, value),
|
|
f32 => fromFloat(name_, value),
|
|
bool => fromBool(name_, value),
|
|
else => unreachable,
|
|
};
|
|
}
|
|
|
|
pub fn fromString(name_: []const u8, value: []const u8) NamedValue {
|
|
return .{ .inner = pjrtStruct(c.PJRT_NamedValue{
|
|
.name = @ptrCast(@constCast(name_.ptr)),
|
|
.name_size = name_.len,
|
|
.type = c.PJRT_NamedValue_kString,
|
|
.unnamed_0 = .{ .string_value = @ptrCast(@constCast(value.ptr)) },
|
|
.value_size = value.len,
|
|
}) };
|
|
}
|
|
|
|
pub fn fromInt64(name_: []const u8, value: i64) NamedValue {
|
|
return .{ .inner = pjrtStruct(c.PJRT_NamedValue{
|
|
.name = @ptrCast(@constCast(name_.ptr)),
|
|
.name_size = name_.len,
|
|
.type = c.PJRT_NamedValue_kInt64,
|
|
.unnamed_0 = .{ .int64_value = value },
|
|
.value_size = 1,
|
|
}) };
|
|
}
|
|
|
|
pub fn fromInt64List(name_: []const u8, value: []const i64) NamedValue {
|
|
return .{ .inner = pjrtStruct(c.PJRT_NamedValue{
|
|
.name = @ptrCast(@constCast(name_.ptr)),
|
|
.name_size = name_.len,
|
|
.type = c.PJRT_NamedValue_kInt64List,
|
|
.unnamed_0 = .{ .int64_array_value = @ptrCast(@constCast(value.ptr)) },
|
|
.value_size = value.len,
|
|
}) };
|
|
}
|
|
|
|
pub fn fromFloat(name_: []const u8, value: f32) NamedValue {
|
|
return .{ .inner = pjrtStruct(c.PJRT_NamedValue{
|
|
.name = @ptrCast(@constCast(name_.ptr)),
|
|
.name_size = name_.len,
|
|
.type = c.PJRT_NamedValue_kFloat,
|
|
.unnamed_0 = .{ .float_value = value },
|
|
.value_size = 1,
|
|
}) };
|
|
}
|
|
|
|
pub fn fromBool(name_: []const u8, value: bool) NamedValue {
|
|
return .{ .inner = pjrtStruct(c.PJRT_NamedValue{
|
|
.name = @ptrCast(@constCast(name_.ptr)),
|
|
.name_size = name_.len,
|
|
.type = c.PJRT_NamedValue_kBool,
|
|
.unnamed_0 = .{ .bool_value = value },
|
|
.value_size = 1,
|
|
}) };
|
|
}
|
|
|
|
pub fn format(
|
|
self: NamedValue,
|
|
comptime fmt: []const u8,
|
|
options: std.fmt.FormatOptions,
|
|
writer: anytype,
|
|
) !void {
|
|
_ = fmt;
|
|
_ = options;
|
|
try writer.print("{s}{{ .name = {s},", .{ @typeName(NamedValue), self.inner.name[0..self.inner.name_size] });
|
|
const u = self.inner.unnamed_0;
|
|
switch (self.kind()) {
|
|
.string => try writer.print(" .string = {s} ", .{u.string_value[0..self.inner.value_size]}),
|
|
.int64 => try writer.print(" .int64 = {d} ", .{u.int64_value}),
|
|
.int64list => try writer.print(" .int64list = {d} ", .{u.int64_array_value[0..self.inner.value_size]}),
|
|
.float => try writer.print(" .float = {d} ", .{u.float_value}),
|
|
.bool => try writer.print(" .bool = {} ", .{u.bool_value}),
|
|
}
|
|
try writer.writeAll("}");
|
|
}
|
|
};
|