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

Adjust ROCm runtime sandboxing to hook only the PJRT plugin and make hipblastlt bytecodes optional.

Browse Source
This commit is contained in:
Tarry Singh 2024-01-26 13:02:23 +00:00
parent 0ce36599da
commit edc2ac26f8
3 changed files with 26 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
runtimes/rocm/libpjrt_rocm.BUILD.bazel
View File

@ -60,6 +60,10 @@ cc_import(
":_hipblaslt": ["@hipblaslt//:runfiles"],
"//conditions:default": [],
}),
add_needed = ["libzmlxrocm.so.0"],
rename_dynamic_symbols = {
"dlopen": "zmlxrocm_dlopen",
},
shared_library = "libpjrt_rocm.so",
visibility = ["//visibility:public"],
deps = [

8
runtimes/rocm/rocm.bzl
View File

@ -119,10 +119,6 @@ load("@zml//runtimes/rocm:gfx.bzl", "bytecode_select")
cc_import(
name = "rocblas",
shared_library = "lib/librocblas.so.4",
add_needed = ["libzmlxrocm.so.0"],
rename_dynamic_symbols = {
"dlopen": "zmlxrocm_dlopen",
},
visibility = ["@libpjrt_rocm//:__subpackages__"],
)
@ -160,10 +156,6 @@ load("@zml//runtimes/rocm:gfx.bzl", "bytecode_select")
cc_import(
name = "hipblaslt",
shared_library = "lib/libhipblaslt.so.0",
add_needed = ["libzmlxrocm.so.0"],
rename_dynamic_symbols = {
"dlopen": "zmlxrocm_dlopen",
},
visibility = ["@libpjrt_rocm//:__subpackages__"],
)

38
runtimes/rocm/rocm.zig
View File

@ -8,6 +8,20 @@ const pjrt = @import("pjrt");
const runfiles = @import("runfiles");
const stdx = @import("stdx");
const ROCmEnvEntry = struct {
name: [:0]const u8,
rpath: []const u8,
dirname: bool,
mandatory: bool,
};
const rocm_env_entries: []const ROCmEnvEntry = &.{
.{ .name = "HIPBLASLT_EXT_OP_LIBRARY_PATH", .rpath = "hipblaslt/lib/hipblaslt/library/hipblasltExtOpLibrary.dat", .dirname = false, .mandatory = false },
.{ .name = "HIPBLASLT_TENSILE_LIBPATH", .rpath = "hipblaslt/lib/hipblaslt/library/TensileManifest.txt", .dirname = true, .mandatory = false },
.{ .name = "ROCBLAS_TENSILE_LIBPATH", .rpath = "rocblas/lib/rocblas/library/TensileManifest.txt", .dirname = true, .mandatory = true },
.{ .name = "ROCM_PATH", .rpath = "libpjrt_rocm/sandbox", .dirname = false, .mandatory = true },
};
pub fn isEnabled() bool {
return @hasDecl(c, "ZML_RUNTIME_ROCM");
}
@ -23,13 +37,6 @@ fn setupRocmEnv() !void {
var arena = std.heap.ArenaAllocator.init(std.heap.c_allocator);
defer arena.deinit();
const paths = .{
.{ "HIPBLASLT_EXT_OP_LIBRARY_PATH", "hipblaslt/lib/hipblaslt/library/hipblasltExtOpLibrary.dat", false },
.{ "HIPBLASLT_TENSILE_LIBPATH", "hipblaslt/lib/hipblaslt/library/TensileManifest.txt", true },
.{ "ROCBLAS_TENSILE_LIBPATH", "rocblas/lib/rocblas/library/TensileManifest.txt", true },
.{ "ROCM_PATH", "libpjrt_rocm/sandbox", false },
};
const r = blk: {
var r_ = try runfiles.Runfiles.create(.{ .allocator = arena.allocator() }) orelse {
stdx.debug.panic("Unable to find Runfiles directory", .{});
@ -38,22 +45,21 @@ fn setupRocmEnv() !void {
break :blk r_.withSourceRepo(source_repo);
};
inline for (paths) |path| {
const name = path[0];
const rpath = path[1];
const dirname = path[2];
var real_path = r.rlocationAlloc(arena.allocator(), rpath) catch null orelse {
stdx.debug.panic("Unable to find " ++ name ++ " in " ++ bazel_builtin.current_repository, .{});
for (rocm_env_entries) |entry| {
var real_path = r.rlocationAlloc(arena.allocator(), entry.rpath) catch null orelse {
if (entry.mandatory) {
stdx.debug.panic("Unable to find {s} in {s}", .{ entry.name, bazel_builtin.current_repository });
}
continue;
};
if (dirname) {
if (entry.dirname) {
real_path = std.fs.path.dirname(real_path) orelse {
stdx.debug.panic("Unable to dirname on {s}", .{real_path});
};
}
_ = c.setenv(name, try arena.allocator().dupeZ(u8, real_path), 1);
_ = c.setenv(entry.name, try arena.allocator().dupeZ(u8, real_path), 1);
}
}