update script

Author: xadupre

_doc/technical/plot_gemm_or_matmul_add.py

@@ -145,7 +145,7 @@ def matrix_diff(tensors):
 # a lot higher.
 
 B = (torch.arange(512, dtype=torch.float32) + 1) / 512 * 16384
-labels = ["linear", *[o.name for o in model.graph.output], "a @ x + b"]
+labels = ["F.linear", *[o.name for o in model.graph.output], "a @ x + b"]
 all_results = {}
 
 for itype, dtype, device in [
@@ -187,28 +187,58 @@ def matrix_diff(tensors):
 # bias value vs discrepancies
 # ===========================
 #
-# Let's compare GemmOnly (so bias is included) and Gemm+Add.
-
-i, j = 1, -1
-labs = labels[i], labels[j]
-
-fig, ax = plt.subplots(len(all_results), 2, figsize=(8, 2.5 * len(results)))
-for pos, ((device, dtype), results) in enumerate(all_results.items()):
-    m1, m2 = results[i], results[j]
-    diff = torch.abs(m1.to(torch.float32) - m2.to(torch.float32)).max(dim=0)[0]
-    print(f"labels={labs}, {device}/{dtype}: max(diff)={diff.max()}")
-    expand = 0.5 if diff.max() >= 1 else diff.max().detach().cpu() / 2
-    ax[pos, 0].plot(B.tolist(), (diff.detach().cpu() + torch.rand(512) * expand).tolist(), ".")
-    ax[pos, 0].set_title(f"{labs[0]}-{labs[1]} {device}/{dtype}")
-
-    corr = matrix_diff(results)
-    ax[pos, 1].imshow(corr, cmap="Blues", vmin=0, vmax=corr.max())
-    # ax[pos,1].colorbar(label=f'Discrepancies {device}/{dtype}')
-    ax[pos, 1].set_xticks(range(len(labels)), labels, rotation=45)
-    ax[pos, 1].set_yticks(range(len(labels)), labels)
-    ax[pos, 1].set_title(f"max={diff.max()}")
+# Let's compare torch linear with GemmOnly.
+
+
+def make_figure_axis(all_results, i, j):
+    labs = labels[i], labels[j]
+    fig, ax = plt.subplots(len(all_results), 2, figsize=(12, 4 * len(all_results)))
+    for pos, ((device, dtype), results) in enumerate(all_results.items()):
+        m1, m2 = results[i], results[j]
+        diff = torch.abs(m1.to(torch.float32) - m2.to(torch.float32)).max(dim=0)[0]
+        print(f"labels={labs}, {device}/{dtype}: max(diff)={diff.max()}")
+        expand = 0.5 if diff.max() >= 1 else diff.max().detach().cpu() / 2
+        ax[pos, 0].plot(
+            B.tolist(), (diff.detach().cpu() + torch.rand(512) * expand).tolist(), "."
+        )
+        ax[pos, 0].set_title(f"{labs[0]}-{labs[1]} {device}/{dtype}", fontsize=10)
+
+        corr = matrix_diff(results)
+        ax[pos, 1].imshow(corr, cmap="Wistia", vmin=0, vmax=corr.max())
+        # ax[pos,1].colorbar(label=f'Discrepancies {device}/{dtype}')
+        ax[pos, 1].set_xticks(range(len(labels)), labels, rotation=45, ha="right", fontsize=10)
+        ax[pos, 1].set_yticks(range(len(labels)), labels, fontsize=10)
+        ax[pos, 1].set_title(f"max={diff.max():1.2g}", fontsize=10)
+        for _i in range(corr.shape[0]):
+            for _j in range(corr.shape[1]):
+                ax[pos, 1].text(
+                    _j,
+                    _i,
+                    f"{corr[_i, _j]:1.1g}",
+                    ha="center",
+                    va="center",
+                    color="black",
+                    fontsize=8,
+                )
+    fig.suptitle(
+        f"Left column: discrepancies {labs[0]} VS {labs[1]}\n"
+        f"Right column: max absolute error, accross all configuration\n"
+        f"white is good, orange is not"
+    )
+    return fig, ax
+
+
+fig, ax = make_figure_axis(all_results, 0, 1)
+fig.tight_layout()
+fig.savefig("plot_gemm_or_matmul_add1.png")
+
+# %%
+# Let's compare with ``a @ x + b``.
+
+fig, ax = make_figure_axis(all_results, -1, 1)
 fig.tight_layout()
-fig.savefig("plot_gemm_or_matmul_add.png")
+fig.savefig("plot_gemm_or_matmul_add2.png")
+
 
 # %%
 # Discrepancies do not happen all the time but it is very likely to happen.