[PROTON] Change the output format of pc sampling lines (triton-lang#5711)

This way users can directly open the file using IDEs like vscode and
jump to the corresponding lines, by holding the `ctrl` key and click the
line on the terminal.

Also, this PR emits an error instead of using workarounds for CUPTI
compatibility problems and adds more instructions for using PC sampling
in the tutorial.


![image](https://github.com/user-attachments/assets/9ecb90c3-0953-43bd-8db4-605dc13c38a2)

Author: Jokeren

third_party/proton/csrc/lib/Profiler/Cupti/CuptiPCSampling.cpp

@@ -135,15 +135,23 @@ CUpti_PCSamplingData allocPCSamplingData(size_t collectNumPCs,
   uint32_t libVersion = 0;
   cupti::getVersion<true>(&libVersion);
   size_t pcDataSize = sizeof(CUpti_PCSamplingPCData);
-  // Check cupti api version < 12.4 but cupti header version >= 12.4
-  // If so, we subtract 4 bytes from the size of CUpti_PCSamplingPCData
-  // because it introduces a new field (i.e., correlationId) at the end of the
-  // struct, which is not compatible with the previous versions.
-  if (libVersion < CUPTI_CUDA12_4_VERSION &&
-      CUPTI_API_VERSION >= CUPTI_CUDA12_4_VERSION)
-    pcDataSize -= CUPTI_CUDA12_4_PC_DATA_PADDING_SIZE;
+  // Since CUPTI 12.4, a new field (i.e., correlationId) is added to
+  // CUpti_PCSamplingPCData, which breaks the ABI compatibility.
+  // Instead of using workarounds, we emit an error message and exit the
+  // application.
+  if ((libVersion < CUPTI_CUDA12_4_VERSION &&
+       CUPTI_API_VERSION >= CUPTI_CUDA12_4_VERSION) ||
+      (libVersion >= CUPTI_CUDA12_4_VERSION &&
+       CUPTI_API_VERSION < CUPTI_CUDA12_4_VERSION)) {
+    throw std::runtime_error(
+        "[PROTON] CUPTI API version: " + std::to_string(CUPTI_API_VERSION) +
+        " and CUPTI driver version: " + std::to_string(libVersion) +
+        " are not compatible. Please set the environment variable "
+        " TRITON_CUPTI_INCLUDE_PATH and TRITON_CUPTI_LIB_PATH to resolve the "
+        "problem.");
+  }
   CUpti_PCSamplingData pcSamplingData{
-      /*size=*/pcDataSize,
+      /*size=*/sizeof(CUpti_PCSamplingData),
       /*collectNumPcs=*/collectNumPCs,
       /*totalSamples=*/0,
       /*droppedSamples=*/0,
@@ -372,16 +380,16 @@ void CuptiPCSampling::processPCSamplingData(ConfigureData *configureData,
         auto *stallReason = &pcData->stallReason[j];
         if (!configureData->stallReasonIndexToMetricIndex.count(
                 stallReason->pcSamplingStallReasonIndex))
-          throw std::runtime_error("Invalid stall reason index");
+          throw std::runtime_error("[PROTON] Invalid stall reason index");
         for (auto *data : dataSet) {
           auto scopeId = externId;
           if (isAPI)
             scopeId = data->addOp(externId, lineInfo.functionName);
           if (lineInfo.fileName.size())
-            scopeId = data->addOp(scopeId,
-                                  lineInfo.dirName + "/" + lineInfo.fileName +
-                                      ":" + lineInfo.functionName + "@" +
-                                      std::to_string(lineInfo.lineNumber));
+            scopeId = data->addOp(
+                scopeId, lineInfo.dirName + "/" + lineInfo.fileName + ":" +
+                             std::to_string(lineInfo.lineNumber) + "@" +
+                             lineInfo.functionName);
           auto metricKind = static_cast<PCSamplingMetric::PCSamplingMetricKind>(
               configureData->stallReasonIndexToMetricIndex
                   [stallReason->pcSamplingStallReasonIndex]);

third_party/proton/csrc/lib/Profiler/Cupti/CuptiProfiler.cpp

@@ -227,7 +227,7 @@ void CuptiProfiler::CuptiProfilerPimpl::allocBuffer(uint8_t **buffer,
                                                     size_t *maxNumRecords) {
   *buffer = static_cast<uint8_t *>(aligned_alloc(AlignSize, BufferSize));
   if (*buffer == nullptr) {
-    throw std::runtime_error("aligned_alloc failed");
+    throw std::runtime_error("[PROTON] aligned_alloc failed");
   }
   *bufferSize = BufferSize;
   *maxNumRecords = 0;
@@ -253,7 +253,7 @@ void CuptiProfiler::CuptiProfilerPimpl::completeBuffer(CUcontext ctx,
     } else if (status == CUPTI_ERROR_MAX_LIMIT_REACHED) {
       break;
     } else {
-      throw std::runtime_error("cupti::activityGetNextRecord failed");
+      throw std::runtime_error("[PROTON] cupti::activityGetNextRecord failed");
     }
   } while (true);
 

third_party/proton/proton/viewer.py

@@ -214,7 +214,8 @@ def format_frames(gf, format):
     elif format == "function_line":
         gf.dataframe["name"] = gf.dataframe["name"].apply(lambda x: x.split(":")[-1])
     elif format == "file_function":
-        gf.dataframe["name"] = gf.dataframe["name"].apply(lambda x: x.split("/")[-1].split("@")[0])
+        gf.dataframe["name"] = gf.dataframe["name"].apply(
+            lambda x: f"{x.split('/')[-1].split(':')[0]}@{x.split('@')[-1].split(':')[0]}")
     return gf
 
 

third_party/proton/test/examples/frame.json

@@ -6,7 +6,7 @@
           {
             "children": [],
             "frame": {
-              "name": "/home/user/projects/example.py/test.py:foo@1",
+              "name": "/home/user/projects/example.py/test.py:1@foo",
               "type": "function"
             },
             "metrics": {

third_party/proton/test/test_viewer.py

@@ -50,13 +50,13 @@ def test_format_frames(option):
         gf, _, _, _ = get_raw_metrics(f)
         gf = format_frames(gf, option)
         if option == "full":
-            idx = gf.dataframe["name"] == "/home/user/projects/example.py/test.py:foo@1"
+            idx = gf.dataframe["name"] == "/home/user/projects/example.py/test.py:1@foo"
         elif option == "file_function_line":
-            idx = gf.dataframe["name"] == "test.py:foo@1"
+            idx = gf.dataframe["name"] == "test.py:1@foo"
         elif option == "function_line":
-            idx = gf.dataframe["name"] == "foo@1"
+            idx = gf.dataframe["name"] == "1@foo"
         elif option == "file_function":
-            idx = gf.dataframe["name"] == "test.py:foo"
+            idx = gf.dataframe["name"] == "test.py@foo"
         assert idx.sum() == 1
 
 

third_party/proton/tutorials/matmul.py

@@ -238,6 +238,7 @@ def grid(META):
 
 argparser = argparse.ArgumentParser()
 argparser.add_argument("--profile", action="store_true")
+argparser.add_argument("--pcsampling", action="store_true", default=False)
 argparser.add_argument("--cudagraph", action="store_true", default=False)
 args = argparser.parse_args()
 
@@ -305,9 +306,13 @@ def perf(ms):
 
 
 if args.profile:
-    proton.start("matmul", hook="triton")
+    if args.pcsampling:
+        # proton-viewer -m num_samples/%,time/s ./matmul.hatchet
+        proton.start("matmul", hook="triton", backend="cupti_pcsampling")
+    else:
+        # proton-viewer -m tflop/s,time/s ./matmul.hatchet
+        proton.start("matmul", hook="triton")
     benchmark.run(show_plots=True, print_data=True)
     proton.finalize()
-    # proton-viewer -m tflop/s,time/s ./matmul.hatchet
 else:
     benchmark.run(show_plots=True, print_data=True)