[Cherry-pick] [proton] Add init and final timestamps to chrome trace (#8870) (#719)

Summary:
Cherry-picked from upstream OAI repository.

Original Commit: ecbb77c
Original Author: Srivatsan Ramesh
Original Date: 2025-12-02 12:32:13 -0500

Original commit message:
```
[proton] Add init and final timestamps to chrome trace (#8870)

Currently the trace timing information are all obained from GPU clock
and are scaled assuming a frequency of 1GHz but if the GPU operates at a
different frequency the timing information becomes a bit misleading. In
this PR, enhanced the Proton instrumentation trace output to include
block-level timing metadata in each trace event, enabling better
temporal analysis.
```

This PR was automatically cherry-picked from the upstream triton-lang/triton repository.

Pull Request resolved: #719

Reviewed By: fywkevin

Differential Revision: D88415845

Pulled By: srivatsan-ramesh

fbshipit-source-id: 3167a0e90fa06be42bd83ba121d46aad58add751

Author: srivatsan-ramesh

third_party/proton/common/lib/TraceDataIO/TraceWriter.cpp

@@ -232,6 +232,8 @@ void StreamChromeTraceWriter::writeKernel(json &object,
           element["ts"] = static_cast<double>(ts) / freq;
           element["dur"] = static_cast<double>(dur) / freq;
           json args;
+          args["Init Time (ns)"] = bt->initTime;
+          args["Post Final Time (ns)"] = bt->postFinalTime;
           args["Finalization Time (ns)"] = bt->postFinalTime - bt->preFinalTime;
           args["Frequency (MHz)"] = freq;
           element["args"] = args;

third_party/proton/test/test_instrumentation.py

@@ -701,6 +701,63 @@ def add_kernel(
         assert len(warp1_events) == 2
 
 
+def test_event_args(tmp_path: pathlib.Path):
+
+    @triton.jit
+    def add_kernel(
+        x_ptr,
+        y_ptr,
+        output_ptr,
+        n_elements,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        with pl.scope("kernel"):
+            pid = tl.program_id(axis=0)
+            block_start = pid * BLOCK_SIZE
+            offsets = block_start + tl.arange(0, BLOCK_SIZE)
+            mask = offsets < n_elements
+            x = tl.load(x_ptr + offsets, mask=mask)
+            y = tl.load(y_ptr + offsets, mask=mask)
+            output = x + y
+            tl.store(output_ptr + offsets, output, mask=mask)
+
+    size = 256
+    x = torch.rand(size, device="cuda")
+    y = torch.rand(size, device="cuda")
+    temp_file = tmp_path / "test_block_metadata.chrome_trace"
+    output = torch.empty_like(x)
+    n_elements = output.numel()
+    grid = (1, 1, 1)
+    proton.start(str(temp_file.with_suffix("")), backend="instrumentation", data="trace")
+    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024, num_warps=2)
+    proton.finalize()
+
+    with open(temp_file, "rb") as f:
+        data = json.load(f)
+        events = data["traceEvents"]
+
+        # Verify we have events
+        assert len(events) > 0
+
+        # Verify each event has the required metadata in args
+        for event in events:
+            assert "args" in event
+            args = event["args"]
+
+            assert "Init Time (ns)" in args
+            assert "Post Final Time (ns)" in args
+            assert "Finalization Time (ns)" in args
+
+            # Verify timing values are reasonable
+            init_time = args["Init Time (ns)"]
+            post_final_time = args["Post Final Time (ns)"]
+            finalization_time = args["Finalization Time (ns)"]
+
+            assert init_time >= 0
+            assert post_final_time >= 0
+            assert finalization_time >= 0
+
+
 def test_threaded_kernel_call(tmp_path: pathlib.Path):
 
     import threading