Update Bazel build configurations (zig.bzl, BUILD files) for MLIR, PJRT, Neuron, ROCm, tokenizer, and tools, fixing broken dependencies.

bazel/zig.bzl

@@ -7,6 +7,7 @@ def zig_cc_binary(
         env = None,
         data = [],
         deps = [],
+        tags = [],
         visibility = None,
         **kwargs):
     zig_binary(
@@ -24,6 +25,7 @@ def zig_cc_binary(
         env = env,
         data = data,
         deps = [":{}_lib".format(name)],
+        tags = tags,
         visibility = visibility,
     )
 
@@ -34,12 +36,14 @@ def zig_cc_test(
         data = [],
         deps = [],
         test_runner = None,
+        tags = [],
         visibility = None,
         **kwargs):
     zig_binary(
         name = "{}_test_lib".format(name),
         kind = BINARY_KIND.test_lib,
         test_runner = test_runner,
+        tags = tags,
         copts = copts + ["-lc", "-fcompiler-rt"],
         deps = deps + [
             "@rules_zig//zig/lib:libc",
@@ -51,6 +55,7 @@ def zig_cc_test(
         env = env,
         data = data,
         deps = [":{}_test_lib".format(name)],
+        tags = tags,
         visibility = visibility,
         linkstatic = True,
     )

mlir/BUILD.bazel

@@ -25,6 +25,7 @@ cc_library(
 
 zig_library(
     name = "mlir",
+    copts = ["-lc"],
     main = "mlir.zig",
     visibility = ["//visibility:public"],
     deps = [

pjrt/BUILD.bazel

@@ -2,10 +2,12 @@ load("@rules_zig//zig:defs.bzl", "zig_library")
 load("@zml//bazel:zig.bzl", "zig_cc_binary")
 load("//bazel:zig_proto_library.bzl", "zig_proto_library")
 
-
 cc_library(
     name = "dlfcn",
     hdrs = ["dlfcn.h"],
+    target_compatible_with = [
+        "@platforms//os:linux",
+    ],
 )
 
 zig_library(
@@ -40,10 +42,10 @@ zig_proto_library(
 )
 
 zig_proto_library(
-     name = "trace_events_proto",
-     import_name = "//tsl:trace_events_proto",
-     deps = ["@tsl//tsl/profiler/protobuf:trace_events_proto"],
- )
+    name = "trace_events_proto",
+    import_name = "//tsl:trace_events_proto",
+    deps = ["@tsl//tsl/profiler/protobuf:trace_events_proto"],
+)
 
 zig_cc_binary(
     name = "xspace_to_json",
@@ -53,6 +55,6 @@ zig_cc_binary(
     deps = [
         ":trace_events_proto",
         ":xplane_proto",
-        "//examples/third_party/tigerbeetle:flags",
+        "//stdx",
     ],
 )

pjrt/xspace_to_json.zig

@@ -1,5 +1,6 @@
-const flags = @import("tigerbeetle/flags");
 const std = @import("std");
+const stdx = @import("stdx");
+const flags = stdx.flags;
 
 const TraceContainer = @import("convert/trace_container.zig").TraceContainer;
 
@@ -26,8 +27,9 @@ pub fn main() !void {
     defer allocator.free(pb_buffer);
     if (pb_buffer.len == 0) return error.EmptyBuffer;
 
-    var converter = try TraceContainer.init(allocator, pb_buffer, cli_args.max_events);
+    var converter = TraceContainer.init(allocator);
     defer converter.deinit();
+    try converter.parseXSpaceBytes(pb_buffer, cli_args.max_events);
 
     var path_buffer: [1028]u8 = undefined;
 

runtimes/neuron/BUILD.bazel

@@ -64,7 +64,7 @@ compile_pip_requirements(
     name = "requirements",
     src = "requirements.in",
     py_binary = py_binary_with_script,
-    requirements_txt = "requirements_lock.txt",
+    requirements_txt = "requirements.lock.txt",
 )
 
 runfiles_to_default(

runtimes/rocm/libpjrt_rocm.BUILD.bazel

@@ -5,15 +5,13 @@ load("@zml//bazel:cc_import.bzl", "cc_import")
 string_list_flag(
     name = "gfx",
     build_setting_default = ["all"],
-    visibility = [
-        "@rocblas//:__subpackages__",
-        "@hipblaslt//:__subpackages__",
-    ],
+    visibility = ["//visibility:public"],
 )
 
 bool_flag(
     name = "hipblaslt",
     build_setting_default = True,
+    visibility = ["//visibility:public"],
 )
 
 config_setting(

zml/tokenizer/BUILD.bazel

@@ -1,22 +1,11 @@
 load("@rules_zig//zig:defs.bzl", "zig_library")
 load("@zml//bazel:zig.bzl", "zig_cc_binary")
-load("//bazel:swig.bzl", "swig_cc_library")
-
-swig_cc_library(
-    name = "sentencepiece_swig",
-    interface = "sentencepiece.i",
-    module = "sentencepiece",
-    deps = [
-        "//ffi:cc",
-        "@com_google_sentencepiece//:sentencepiece_processor",
-    ],
-)
 
 zig_library(
     name = "tokenizer",
+    srcs = ["homemade.zig"],
     import_name = "zml/tokenizer",
     main = "tokenizer.zig",
-    srcs = ["homemade.zig"],
     visibility = ["//visibility:public"],
     deps = [
         "//async",
@@ -31,8 +20,8 @@ zig_cc_binary(
     main = "main.zig",
     visibility = ["//visibility:public"],
     deps = [
-        "//stdx",
-        "//async",
         ":tokenizer",
+        "//async",
+        "//stdx",
     ],
 )

zml/tokenizer/hftokenizers/BUILD.bazel

@@ -1,6 +1,5 @@
 load("@rules_rust//rust:defs.bzl", "rust_static_library")
 load("@rules_zig//zig:defs.bzl", "zig_library")
-load("@zml//bazel:zig.bzl", "zig_cc_binary")
 
 rust_static_library(
     name = "hftokenizers_rs",
@@ -29,10 +28,3 @@ zig_library(
         "//ffi:zig",
     ],
 )
-
-zig_cc_binary(
-    name = "main",
-    main = "main.zig",
-    visibility = ["//visibility:public"],
-    deps = [":hftokenizers"],
-)

zml/tokenizer/sentencepiece/BUILD.bazel

@@ -1,5 +1,4 @@
 load("@rules_zig//zig:defs.bzl", "zig_library")
-load("@zml//bazel:zig.bzl", "zig_cc_binary")
 load("//bazel:swig.bzl", "swig_cc_library")
 
 swig_cc_library(
@@ -22,14 +21,3 @@ zig_library(
         "//ffi:zig",
     ],
 )
-
-zig_cc_binary(
-    name = "main",
-    srcs = ["sentencepiece.zig"],
-    main = "main.zig",
-    visibility = ["//visibility:public"],
-    deps = [
-        ":sentencepiece_swig",
-        "//ffi:zig",
-    ],
-)

zml/tools/BUILD.bazel

@@ -4,6 +4,9 @@ cc_library(
     name = "macos_c",
     srcs = ["macos.c"],
     hdrs = ["macos.h"],
+    target_compatible_with = [
+        "@platforms//os:macos",
+    ],
     visibility = ["//zml/tools:__subpackages__"],
 )