40 lines
1.6 KiB
Zig
40 lines
1.6 KiB
Zig
const std = @import("std");
|
|
const runfiles = @import("runfiles");
|
|
const bazel_builtin = @import("bazel_builtin");
|
|
const stdx = @import("stdx");
|
|
|
|
const log = std.log.scoped(.@"zml/runtimes/neuron");
|
|
|
|
pub export fn GetPjrtApi() *anyopaque {
|
|
var arena = std.heap.ArenaAllocator.init(std.heap.c_allocator);
|
|
defer arena.deinit();
|
|
|
|
var r_ = runfiles.Runfiles.create(.{ .allocator = arena.allocator() }) catch |err| {
|
|
stdx.debug.panic("Unable to find runfiles: {}", .{err});
|
|
} orelse stdx.debug.panic("Runfiles not availeabwewefle", .{});
|
|
|
|
const source_repo = bazel_builtin.current_repository;
|
|
const r = r_.withSourceRepo(source_repo);
|
|
|
|
var path_buf: [std.fs.max_path_bytes]u8 = undefined;
|
|
const sandbox_path = r.rlocation("libpjrt_neuron/sandbox", &path_buf) catch |err| {
|
|
stdx.debug.panic("Failed to find sandbox path for NEURON runtime: {}", .{err});
|
|
} orelse stdx.debug.panic("No NEURON sandbox path found", .{});
|
|
|
|
var lib_path_buf: [std.fs.max_path_bytes]u8 = undefined;
|
|
const library = stdx.fs.path.bufJoinZ(&lib_path_buf, &.{ sandbox_path, "lib", "libneuronpjrt.so" }) catch unreachable;
|
|
|
|
var lib: std.DynLib = blk: {
|
|
const handle = std.c.dlopen(library, .{ .LAZY = true, .GLOBAL = true, .NODELETE = true }) orelse {
|
|
stdx.debug.panic("Unable to dlopen plugin: {s}", .{library});
|
|
};
|
|
break :blk .{ .inner = .{ .handle = handle } };
|
|
};
|
|
|
|
const sym = lib.lookup(*const fn () callconv(.C) *anyopaque, "GetPjrtApi") orelse {
|
|
stdx.debug.panic("Unable to find symbol GetPjrtApi in plugin: {s}", .{library});
|
|
};
|
|
|
|
return sym();
|
|
}
|