Add TraceMe for Mosaic Kernel Injection Id

PiperOrigin-RevId: 803537235

Author: Google-ML-Automation

jax/_src/lib/__init__.py

@@ -100,7 +100,6 @@ def _parse_version(v: str) -> tuple[int, ...]:
 import jaxlib._jax as _jax  # noqa: F401
 
 
-
 import jaxlib.mlir._mlir_libs._jax_mlir_ext as jax_mlir_ext  # noqa: F401
 from jaxlib._jax import guard_lib as guard_lib  # noqa: F401
 from jaxlib._jax import jax_jit as jax_jit  # noqa: F401
@@ -110,6 +109,11 @@ def _parse_version(v: str) -> tuple[int, ...]:
 from jaxlib import _profiler as _profiler  # noqa: F401
 from jaxlib import _profile_data as _profile_data  # noqa: F401
 
+try:
+  from jaxlib import _gpu_ondevice_tracing as _gpu_ondevice_tracing  # noqa: F401
+except (ImportError, ModuleNotFoundError):
+  _gpu_ondevice_tracing = None
+
 from jaxlib._jax import ffi as ffi  # noqa: F401
 import jaxlib.cpu_sparse as cpu_sparse  # noqa: F401
 has_cpu_sparse = True

jax/_src/pallas/mosaic_gpu/core.py

@@ -121,9 +121,9 @@ def __post_init__(self):
       object.__setattr__(
           self, "dimension_semantics", tuple(self.dimension_semantics)
       )
-    if bool(self.profile_space) ^ bool(self.profile_dir):
+    if bool(self.profile_dir) and not bool(self.profile_space):
       raise ValueError(
-          "Either both profile_space and profile_dir must be set, or neither."
+          "If profile_dir is set, the profile_space must be set too."
       )
 
 

jax/_src/pallas/mosaic_gpu/pallas_call_registration.py

@@ -120,24 +120,33 @@ def zero_init_gmem_scratch():
   )
   if (prof_spec := lowering_result.profiler_spec) is not None:
     *outs, prof_buffer = outs
-    out_file = os.path.join(
-        prof_spec.dump_path,
-        f"{mlir.sanitize_name(debug_info.func_name)}-{time.time_ns()}-trace.json",
-    )
     def dump_profile(prof_buffer):
-      try:
-        with open(out_file, "x") as f:
-          prof_spec.dump(
-              prof_buffer,
-              f,
-              grid=lowering_result.grid,
-              block=lowering_result.block,
-          )
-      except FileExistsError:
-        warnings.warn(
-            f"Failed to dump profile for pallas_call {debug_info.func_src_info}, "
-            f"profile already exists at {out_file}"
+      if prof_spec.dump_path is None:
+        prof_spec.dump(
+            prof_buffer,
+            None,
+            grid=lowering_result.grid,
+            block=lowering_result.block,
         )
+      else:
+        out_file = os.path.join(
+            prof_spec.dump_path,
+            f"{mlir.sanitize_name(debug_info.func_name)}-{time.time_ns()}-trace.json",
+        )
+        try:
+          with open(out_file, "x") as f:
+            prof_spec.dump(
+                prof_buffer,
+                f,
+                grid=lowering_result.grid,
+                block=lowering_result.block,
+            )
+        except FileExistsError:
+          warnings.warn(
+              "Failed to dump profile for pallas_call"
+              f" {debug_info.func_src_info}, profile already exists at"
+              f" {out_file}"
+          )
     def do_callback(prof_buffer):
       jax.debug.callback(dump_profile, prof_buffer)
       return ()

jax/experimental/mosaic/gpu/core.py

@@ -925,10 +925,15 @@ def prof_kernel(*args):
       _check_args(*args)
       *results, prof_buffer = bind(*args)
       def dump_profile(prof_buffer):
-        out_file = os.path.join(prof_spec.dump_path, f"{time.time_ns()}-trace.json")
         try:
-          with open(out_file, "x") as f:
-            prof_spec.dump(prof_buffer, f, grid=grid, block=block)
+          if prof_spec.dump_path is not None:
+            out_file = os.path.join(
+                prof_spec.dump_path, f"{time.time_ns()}-trace.json"
+            )
+            with open(out_file, "x") as f:
+              prof_spec.dump(prof_buffer, f, grid=grid, block=block)
+          else:
+            prof_spec.dump(prof_buffer, None, grid=grid, block=block)
         except FileExistsError:
           pass  # TODO: Retry
       jax.debug.callback(dump_profile, prof_buffer)

jax/experimental/mosaic/gpu/profiler.py

@@ -26,6 +26,8 @@
 import jax
 from jax._src import stages
 from jax._src import util
+from jax._src.lib import _gpu_ondevice_tracing as gpu_ondevice_tracing
+from jax._src.lib import _profiler
 import jax.numpy as jnp
 from jaxlib.mlir import ir
 from jaxlib.mlir.dialects import arith
@@ -185,6 +187,20 @@ def __init__(self, entries_per_warpgroup: int, dump_path: str = "sponge"):
       )
     else:
       self.dump_path = dump_path
+    self.tracing_version = self.injection_id = 0
+    self.check_gpu_ondevice_tracing()
+
+  def check_gpu_ondevice_tracing(self):
+    if (
+        gpu_ondevice_tracing is not None
+        and self.tracing_version == 0
+        and self.injection_id == 0
+    ):
+      self.tracing_version = gpu_ondevice_tracing.active_version()
+      if self.tracing_version > 0:
+        self.injection_id = gpu_ondevice_tracing.start_injection_instance(
+            self.tracing_version
+        )
 
   def _num_warpgroups(
       self, grid: tuple[int, ...], block: tuple[int, ...]
@@ -294,7 +310,42 @@ def dump(self, buffer, f, grid: tuple[int, ...], block: tuple[int, ...]):
           events.append(block_events)
     events = sorted(events, key=lambda x: x[0]["ts"])
     flat_events = list(itertools.chain.from_iterable(events))
-    return json.dump({"displayTimeUnit": "ns", "traceEvents": flat_events}, f)
+
+    if f is not None:
+      json.dump({"displayTimeUnit": "ns", "traceEvents": flat_events}, f)
+
+    if (
+        gpu_ondevice_tracing is not None
+        and self.tracing_version > 0
+        and self.injection_id > 0
+    ):
+      with _profiler.TraceMe(
+          "MosaicGpuProfilerDump", inject_id=self.injection_id
+      ):
+        range_dict = {}
+        for event in flat_events:
+          range_key = (event["name"], event["pid"], event["tid"])
+          if event["ph"] == "B":
+            range_dict[range_key] = event["ts"]
+          elif event["ph"] == "E":
+            if range_key in range_dict:
+              begin_ts = range_dict[range_key]
+              range_dict.pop(range_key)
+              gpu_ondevice_tracing.inject(
+                  version=self.tracing_version,
+                  injection_instance_id=self.injection_id,
+                  tag_name=event["name"],
+                  tag_id=self.interned_names[event["name"]],
+                  pid=event["pid"],
+                  tid=event["tid"],
+                  start_time_ns=int(begin_ts * 1e3),
+                  duration_ps=int((event["ts"] - begin_ts) * 1e6),
+              )
+            else:
+              warnings.warn(
+                  "Event start time not found for ending event:"
+                  f" {event['name']}@{event['ts']}"
+              )
 
 
 @dataclasses.dataclass(frozen=True)

jaxlib/BUILD

@@ -92,6 +92,7 @@ pytype_strict_library(
         "//jaxlib/mosaic/python:gpu_dialect",
         "//jaxlib/mosaic/python:tpu_dialect",
         "//jaxlib/triton",
+        "@xla//xla/python:_gpu_ondevice_tracing",
         "@xla//xla/python:_profile_data",
         "@xla//xla/python:_profiler",
     ],
@@ -178,6 +179,7 @@ pywrap_library(
         "//jaxlib/mlir/_mlir_libs:_stablehlo",
         "//jaxlib/mlir/_mlir_libs:_tpu_ext",
         "//jaxlib/mlir/_mlir_libs:_triton_ext",
+        "@xla//xla/python:_gpu_ondevice_tracing",
         "@xla//xla/python:_profile_data",
         "@xla//xla/python:_profiler",
     ],

jaxlib/mosaic/gpu/custom_call.cc

@@ -662,6 +662,7 @@ void MosaicGPUCustomCall(void* stream, void** buffers, char* opaque,
     mosaic::gpu::NvshmemApi::Default().barrier_all_on_stream(
         reinterpret_cast<cudaStream_t>(stream));
   }
+  tsl::profiler::TraceMe trace("MosaicGpuLaunchKernel");
   std::get<1>(ctx_kernel_comm)(args);
 }
 
@@ -724,6 +725,7 @@ absl::Status MosaicGpuExecute(gpuStream_t stream, ffi::RemainingArgs inputs,
     mosaic::gpu::NvshmemApi::Default().barrier_all_on_stream(
         reinterpret_cast<cudaStream_t>(stream));
   }
+  tsl::profiler::TraceMe trace("MosaicGpuLaunchKernel");
   std::get<1>(ctx_kernel_comm)(args);
   return absl::OkStatus();
 }

jaxlib/tools/build_wheel.py

@@ -213,6 +213,7 @@ def prepare_wheel(wheel_sources_path: pathlib.Path, *, cpu, wheel_sources):
           f"{source_file_prefix}jaxlib/_jax.{pyext}",
           f"{source_file_prefix}jaxlib/_sdy_mpmd.{pyext}",
           f"{source_file_prefix}jaxlib/_pathways.{pyext}",
+          f"{source_file_prefix}jaxlib/_gpu_ondevice_tracing.{pyext}",
           f"{source_file_prefix}jaxlib/_profiler.{pyext}",
           f"{source_file_prefix}jaxlib/_profile_data.{pyext}",
       ],