tarry.singh/Radix
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

pjrt: emit warning instead of panic when FFI Extension is missing (e.g., on TPU).

Browse Source
This commit is contained in:
Tarry Singh 2025-03-24 09:40:44 +00:00
parent dc121fce4f
commit a5420068b1
1 changed files with 5 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
zml/context.zig
View File

@ -218,13 +218,11 @@ pub const Context = struct {
const CustomCall = struct { const CustomCall = struct {
pub fn registerZmlCustomCalls(platform: Platform) !void { pub fn registerZmlCustomCalls(platform: Platform) !void {
const maybe_ffi = platform.pjrt_api.ffi(); const ffi = platform.pjrt_api.ffi() orelse {
log.warn("Registering custom calls failed: No FFI Extension found in {s} PJRT Plugin.", .{@tagName(platform.target)});
if (maybe_ffi) |ffi| { return;
};
try ffi.register(platform.pjrt_api, "zmlHostBufferCallback", @tagName(platform.target), &hostBufferCallback, .{}); try ffi.register(platform.pjrt_api, "zmlHostBufferCallback", @tagName(platform.target), &hostBufferCallback, .{});
} else {
stdx.debug.panic("Registering custom calls failed", .{});
}
} }
fn hostBufferCallback(call_frame: *pjrt.ffi.CallFrame) callconv(.C) ?*pjrt.ffi.Error { fn hostBufferCallback(call_frame: *pjrt.ffi.CallFrame) callconv(.C) ?*pjrt.ffi.Error {