[PROTON] Add kernel call stack to intra kernel events (#8071)

<img width="2000" height="243" alt="image"
src="https://github.com/user-attachments/assets/d577d056-8b03-4528-904c-2638ff294ddc"
/>

Author: Jokeren

third_party/proton/common/include/TraceDataIO/TraceWriter.h

@@ -19,6 +19,7 @@ const uint64_t kKernelTimeGap = 10000000;
 struct KernelMetadata {
   std::map<int, std::string> scopeName;
   std::string kernelName;
+  std::vector<std::string> callStack;
 };
 
 using KernelTrace = std::pair<std::shared_ptr<CircularLayoutParserResult>,

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

@@ -151,6 +151,11 @@ void StreamChromeTraceWriter::writeKernel(json &object,
   auto result = kernelTrace.first;
   auto metadata = kernelTrace.second;
 
+  json callStack = json::array();
+  for (auto const &frame : metadata->callStack) {
+    callStack.push_back(frame);
+  }
+
   int curColorIndex = 0;
   // scope id -> color index in chrome color
   std::map<int, int> scopeColor;
@@ -209,6 +214,8 @@ void StreamChromeTraceWriter::writeKernel(json &object,
           args["Unit"] = "GPU cycle";
           args["Kernel Gap"] = std::to_string(kKernelTimeGap) + "cycle(ns)";
           element["args"] = args;
+          element["args"]["call_stack"] = callStack;
+
           object["traceEvents"].push_back(element);
 
           eventIdx++;

third_party/proton/csrc/lib/Data/TraceData.cpp

@@ -385,9 +385,19 @@ convertToTimelineTrace(TraceData::Trace *trace,
       }
       parserResult->blockTraces.push_back(std::move(blockTrace));
     }
+    std::vector<std::string> callStack;
+    if (!sortedEvents.empty()) {
+      auto contexts = trace->getContexts(kernelEvent.contextId);
+      if (!contexts.empty()) {
+        callStack.resize(contexts.size() - 1);
+        std::transform(contexts.begin(), contexts.end() - 1, callStack.begin(),
+                       [](const Context &c) { return c.name; });
+      }
+    }
     metadata->kernelName =
         getStringValue(kernelEvent.cycleMetric, CycleMetric::KernelName);
     metadata->scopeName = scopeIdToName;
+    metadata->callStack = std::move(callStack);
     if (timeShiftCost > 0)
       timeShift(timeShiftCost, parserResult);
     results.emplace_back(parserResult, metadata);

third_party/proton/test/test_instrumentation.py

@@ -609,3 +609,5 @@ def foo(x, y, size: tl.constexpr):
         trace_events = data["traceEvents"]
         assert len(trace_events) == 12
         assert trace_events[-1]["tid"][0:4] == "warp"
+        assert trace_events[-1]["args"]["call_stack"][-1] == "foo"
+        assert trace_events[-1]["args"]["call_stack"][-2] == "test"