Test Mark changes

Signed-off-by: Anatoly Myachev <anatoly.myachev@intel.com>

python/triton/testing.py

@@ -377,12 +377,14 @@ def _run(self, bench: Benchmark, save_path: str, show_plots: bool, print_data: b
         y_min = [f'{x}-min' for x in bench.line_names]
         y_max = [f'{x}-max' for x in bench.line_names]
         x_names = list(bench.x_names)
-        df = pd.DataFrame(columns=x_names + y_mean + y_min + y_max)
-        for x in bench.x_vals:
+        if len(x_names) == 1:
+            index = pd.Index(bench.x_vals, name=x_names[0])
+        else:
             # x can be a single value or a sequence of values.
-            if not isinstance(x, (list, tuple)):
-                x = [x for _ in x_names]
-
+            x_vals = [tuple(x for _ in x_names) for x in bench.x_vals if not isinstance(x, (list, tuple))]
+            index = pd.MultiIndex.from_tuples(x_vals, names=x_names)
+        df = pd.DataFrame(index=index, columns=y_mean + y_min + y_max)
+        for x in df.index:
             if len(x) != len(x_names):
                 raise ValueError(f"Expected {len(x_names)} values, got {x}")
             x_args = dict(zip(x_names, x))
@@ -397,24 +399,28 @@ def _run(self, bench: Benchmark, save_path: str, show_plots: bool, print_data: b
                 row_mean += [y_mean]
                 row_min += [y_min]
                 row_max += [y_max]
-            df.loc[len(df)] = list(x) + row_mean + row_min + row_max
+            df.loc[x] = row_mean + row_min + row_max
 
         if bench.plot_name:
             plt.figure()
             ax = plt.subplot()
             # Plot first x value on x axis if there are multiple.
-            first_x = x_names[0]
+            index_name = x_names[0]
+            if len(x_names) == 1:
+                index = df.index
+            else:
+                index = df.index.get_level_values(0)
             for i, y in enumerate(bench.line_names):
                 y_min, y_max = df[y + '-min'], df[y + '-max']
                 col = bench.styles[i][0] if bench.styles else None
                 sty = bench.styles[i][1] if bench.styles else None
-                ax.plot(df[first_x], df[y], label=y, color=col, ls=sty)
+                ax.plot(index, df[y], label=y, color=col, ls=sty)
                 if not y_min.isnull().all() and not y_max.isnull().all():
                     y_min = y_min.astype(float)
                     y_max = y_max.astype(float)
-                    ax.fill_between(df[first_x], y_min, y_max, alpha=0.15, color=col)
+                    ax.fill_between(index, y_min, y_max, alpha=0.15, color=col)
             ax.legend()
-            ax.set_xlabel(bench.xlabel or first_x)
+            ax.set_xlabel(bench.xlabel or index_name)
             ax.set_ylabel(bench.ylabel)
             # ax.set_title(bench.plot_name)
             ax.set_xscale("log" if bench.x_log else "linear")
@@ -423,7 +429,7 @@ def _run(self, bench: Benchmark, save_path: str, show_plots: bool, print_data: b
                 plt.show()
             if save_path:
                 plt.savefig(os.path.join(save_path, f"{bench.plot_name}.png"))
-        df = df[x_names + bench.line_names]
+        df = df[bench.line_names]
         if diff_col and df.shape[1] == 2:
             col0, col1 = df.columns.tolist()
             df['Diff'] = df[col1] - df[col0]
@@ -432,8 +438,7 @@ def _run(self, bench: Benchmark, save_path: str, show_plots: bool, print_data: b
             print(bench.plot_name + ':')
             print(df.to_string())
         if save_path:
-            df.to_csv(os.path.join(save_path, f"{bench.plot_name}.csv"), float_format=f"%.{save_precision}f",
-                      index=False)
+            df.to_csv(os.path.join(save_path, f"{bench.plot_name}.csv"), float_format=f"%.{save_precision}f")
         return df
 
     def run(self, show_plots=False, print_data=False, save_path='', return_df=False, **kwargs):