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

pjrt: bind PJRT_Device_AddressableMemories.

Browse Source
This commit is contained in:
Tarry Singh 2024-09-26 12:02:04 +00:00
parent 2863c1f5e0
commit 8b53ebbd9b
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
pjrt/pjrt.zig
View File

@ -502,6 +502,13 @@ pub const Device = opaque {
}) catch unreachable;
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 {