pjrt: bind PJRT_Device_AddressableMemories.

This commit is contained in:
Tarry Singh 2024-09-26 12:02:04 +00:00
parent 2863c1f5e0
commit 8b53ebbd9b

View File

@ -502,6 +502,13 @@ pub const Device = opaque {
}) catch unreachable; }) catch unreachable;
return @intCast(ret.local_hardware_id); return @intCast(ret.local_hardware_id);
} }
pub fn addressableMemories(self: *const Device, api: *const Api) ApiError![]const *Memory {
const ret = try api.call(.PJRT_Device_AddressableMemories, .{
.device = self.inner(),
});
return @ptrCast(ret.memories[0..ret.num_memories]);
}
}; };
pub const DeviceDescription = opaque { pub const DeviceDescription = opaque {