Update

Author: Jokeren

bench/bench/bench_mlp.py

@@ -95,14 +95,15 @@ def bench_mlp(batch, dim1, dim2, n_expts_tot, n_expts_act, x_dtype, w_dtype,
 
     # -- analyze --
     gf, inclusive_metrics, exclusive_metrics, device_info = viewer.read(fpath)
-    matmuls = gf.dataframe[gf.dataframe["name"].str.contains("matmul")]
+    # Now the dataframe only contains leave nodes (i.e., kernels) that perform matmuls
+    matmuls = gf.dataframe[gf.dataframe["name"].str.contains("matmul") & gf.dataframe["device_id"].notna()]
     tot_bytes = matmuls["bytes"].sum()
     tot_flops = sum(matmuls[[c for c in ['flops8', 'flops16'] if c in matmuls.columns]].sum())
     tot_time = matmuls["time (ns)"].sum()
 
     # Calculate theoretical min time based on hardware limits
-    device_type = matmuls["device_type"].dropna().iloc[0]
-    device_id = matmuls["device_id"].dropna().iloc[0]
+    device_type = matmuls["device_type"].iloc[0]
+    device_id = matmuls["device_id"].iloc[0]
     info = device_info[device_type][device_id]
 
     min_time_flops_sec = sum(