[PROTON] Set non-derivable metrics as nan (#4387)

Also improved the test cases

third_party/proton/proton/viewer.py

@@ -2,8 +2,8 @@
 from collections import namedtuple
 import json
 import pandas as pd
-from numpy import inf, nan
 import hatchet as ht
+import numpy as np
 from triton.profiler.hook import COMPUTE_METADATA_SCOPE_NAME, TritonHook
 
 
@@ -92,18 +92,20 @@ def get_min_time_bytes(df, device_info):
 def derive_metrics(gf, metrics, raw_metrics, device_info):
     derived_metrics = []
     original_metrics = []
-    time_metric_name = match_available_metrics([time_factor_dict.name], raw_metrics)[0]
-    time_unit = (time_factor_dict.name + "/" + time_metric_name.split("(")[1].split(")")[0])
+    internal_frame_indices = gf.dataframe["DeviceId"].isna()
 
     def get_time_seconds(df):
+        time_metric_name = match_available_metrics([time_factor_dict.name], raw_metrics)[0]
+        time_unit = (time_factor_dict.name + "/" + time_metric_name.split("(")[1].split(")")[0])
         return df[time_metric_name] * time_factor_dict.factor[time_unit]
 
     for metric in metrics:
         if metric == "util":  # Tensor core only
             min_time_bytes = get_min_time_bytes(gf.dataframe, device_info)
             min_time_flops = get_min_time_flops(gf.dataframe, device_info)
-            time_sec = get_time_seconds(gf.dataframe, time_metric_name, time_unit)
+            time_sec = get_time_seconds(gf.dataframe)
             gf.dataframe["util (inc)"] = min_time_flops["min_time"].combine(min_time_bytes["min_time"], max) / time_sec
+            gf.dataframe.loc[internal_frame_indices, "util (inc)"] = np.nan
             derived_metrics.append("util (inc)")
         elif metric in derivable_metrics:
             deriveable_metric = derivable_metrics[metric]
@@ -121,7 +123,8 @@ def get_time_seconds(df):
         elif metric in avg_time_factor_dict.factor:
             metric_time_unit = avg_time_factor_dict.name + "/" + metric.split("/")[1]
             gf.dataframe[f"{metric} (inc)"] = (get_time_seconds(gf.dataframe) / gf.dataframe['Count'] /
-                                               avg_time_factor_dict.factor[metric_time_unit]).replace([inf, -inf], nan)
+                                               avg_time_factor_dict.factor[metric_time_unit])
+            gf.dataframe.loc[internal_frame_indices, f"{metric} (inc)"] = np.nan
             derived_metrics.append(f"{metric} (inc)")
         else:
             original_metrics.append(metric)

third_party/proton/test/example_cuda.json

@@ -8,7 +8,7 @@
           "type": "function"
         },
         "metrics": {
-          "Count": 1,
+          "Count": 10,
           "DeviceId": "1",
           "DeviceType": "CUDA",
           "Time (ns)": 204800,

third_party/proton/test/test_viewer.py

@@ -2,7 +2,6 @@
 import subprocess
 from triton.profiler.viewer import get_min_time_flops, get_min_time_bytes, get_raw_metrics, format_frames, derive_metrics
 import numpy as np
-import pandas as pd
 
 file_path = __file__
 cuda_example_file = file_path.replace("test_viewer.py", "example_cuda.json")
@@ -78,13 +77,24 @@ def test_avg_time_derivation():
     metrics = ["avg_time/s", "avg_time/ms", "avg_time/us", "avg_time/ns"]
     with open(cuda_example_file, "r") as f:
         expected_data = {
-            'avg_time/s (inc)': [np.nan, 0.000205, 0.000205], 'avg_time/ms (inc)': [np.nan, 0.2048, 0.2048],
-            'avg_time/us (inc)': [np.nan, 204.8, 204.8], 'avg_time/ns (inc)': [np.nan, 204800.0, 204800.0]
+            'avg_time/s (inc)': [np.nan, 0.0000205, 0.000205], 'avg_time/ms (inc)': [np.nan, 0.02048, 0.2048],
+            'avg_time/us (inc)': [np.nan, 20.48, 204.8], 'avg_time/ns (inc)': [np.nan, 20480.0, 204800.0]
         }
         gf, raw_metrics, device_info = get_raw_metrics(f)
-        gf = format_frames(gf, format)
         assert len(raw_metrics) > 0, "No metrics found in the input file"
         gf.update_inclusive_columns()
-        d_metrics = derive_metrics(gf, metrics, raw_metrics, device_info)
-        for metric in d_metrics:
-            np.testing.assert_allclose(gf.dataframe[metric].to_numpy(), expected_data[metric], atol=1e-6)
+        derived_metrics = derive_metrics(gf, metrics, raw_metrics, device_info)
+        for derived_metric in derived_metrics:
+            np.testing.assert_allclose(gf.dataframe[derived_metric].to_numpy(), expected_data[derived_metric],
+                                       atol=1e-6)
+
+
+def test_util():
+    metrics = ["util"]
+    with open(cuda_example_file, "r") as f:
+        gf, raw_metrics, device_info = get_raw_metrics(f)
+        assert len(raw_metrics) > 0, "No metrics found in the input file"
+        gf.update_inclusive_columns()
+        derived_metrics = derive_metrics(gf, metrics, raw_metrics, device_info)
+        np.testing.assert_allclose(gf.dataframe[derived_metrics].to_numpy(), [[np.nan], [0.247044], [0.147830]],
+                                   atol=1e-6)