zml/pjrtx: execute bufferFromHostBuffer on the thread pool to avoid blocking and improve weight loading performance.

This commit is contained in:
Tarry Singh 2023-08-29 10:28:51 +00:00
parent c081cb9ad6
commit aa7fae449e

View File

@ -64,7 +64,7 @@ pub const Client = opaque {
pub const BufferFromHostBufferArgs = pjrt.Client.BufferFromHostBufferArgs; pub const BufferFromHostBufferArgs = pjrt.Client.BufferFromHostBufferArgs;
pub fn bufferFromHostBuffer(self: *const Client, api: *const Api, args: BufferFromHostBufferArgs) !*Buffer { pub fn bufferFromHostBuffer(self: *const Client, api: *const Api, args: BufferFromHostBufferArgs) !*Buffer {
const buffer, const event_ = try self.inner().bufferFromHostBuffer(api, args); const buffer, const event_ = try asynk.callBlocking(pjrt.Client.bufferFromHostBuffer, .{ self.inner(), api, args });
if (event_) |event__| { if (event_) |event__| {
const event: *Event = @ptrCast(event__); const event: *Event = @ptrCast(event__);
try event.await_(api); try event.await_(api);