From 98a81b36e1ca82e5b6dae4d24ba95f21e9026d06 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Fri, 27 Sep 2024 10:26:46 +0800
Subject: [PATCH] viz table view (#6743)

* fix matcher with ctx

* current_kernel fix

* add table

* make the right things clickable

* some more init work

* add kernel resizer

* Revert "add kernel resizer"

This reverts commit 035eef37039aa1e848a766a29e3c4e81bbff2bab.

* allow scroll
---
 tinygrad/codegen/kernel.py | 2 +-
 viz/index.html             | 7 ++++---
 viz/serve.py               | 3 ++-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py
index bd332384..7925a960 100644
--- a/tinygrad/codegen/kernel.py
+++ b/tinygrad/codegen/kernel.py
@@ -704,6 +704,7 @@ class Kernel:
   # **** this is the lowerer ****
 
   def linearize(self) -> Kernel:
+    if TRACK_MATCH_STATS >= 2: _CURRENT_KERNEL.set(self.name)
     modified_ast = self.get_optimized_ast()
 
     if DEBUG >= 3:
@@ -713,7 +714,6 @@ class Kernel:
       print(self.applied_opts)
     verify_ast(modified_ast)
 
-    if TRACK_MATCH_STATS >= 2: _CURRENT_KERNEL.set(self.name)
     self.uops:List[UOp] = linearize_uop(full_graph_rewrite(ast_to_uop(modified_ast, self.opts), self.opts))
     if TRACK_MATCH_STATS >= 2: _CURRENT_KERNEL.set(None)
     if DEBUG >= 5: print_uops(self.uops)
diff --git a/viz/index.html b/viz/index.html
index a118ecc0..7d0d0b6a 100644
--- a/viz/index.html
+++ b/viz/index.html
@@ -42,6 +42,7 @@
   ul {
     padding: 0;
     color: #7c7d85;
+    white-space: nowrap;
     cursor: pointer;
   }
   ul.active {
@@ -242,11 +243,11 @@
     const kernelList = document.querySelector(".container.kernel-list");
     kernelList.innerHTML = "";
     kernels.forEach((k, i) => {
-      kernelUl = Object.assign(document.createElement("ul"), { key: `kernel-${i}`, className: i === currentKernel ? "active" : "" });
-      const p = Object.assign(document.createElement("p"), {id: `kernel-${k.name}`, innerText: k.name})
+      kernelUl = Object.assign(document.createElement("ul"), { key: `kernel-${i}`, className: i === currentKernel ? "active" : "", style: "overflow-x: auto; cursor: initial;" });
+      const p = Object.assign(document.createElement("p"), {id: `kernel-${k.name}`, innerText: k.name, style: "cursor: pointer;"});
       kernelUl.appendChild(p)
       k.ctxs.forEach((u, j) => {
-        const rwUl = Object.assign(document.createElement("ul"), { innerText: `${u.filename} - ${u.match_count}`, key: `uop-rewrite-${j}`, className: (j === currentUOp && i == currentKernel) ? "active" : "" })
+        const rwUl = Object.assign(document.createElement("ul"), { innerText: `${u.filename} - ${u.match_count} - ${u.matcher_name}`, key: `uop-rewrite-${j}`, className: (j === currentUOp && i == currentKernel) ? "active" : "" })
         rwUl.style.display = i === currentKernel && expandKernel ? "block" : "none";
         rwUl.onclick = (e) => {
           e.stopPropagation();
diff --git a/viz/serve.py b/viz/serve.py
index 84d40136..d2d2c456 100755
--- a/viz/serve.py
+++ b/viz/serve.py
@@ -30,7 +30,8 @@ class RewriteLocation:
     fp, lineno = ctx.loc
     p = r"graph_rewrite\([^,]+,\s*([^>]+)\)"
     match = re.search(p, code:=lines(fp)[lineno-1].strip())
-    return RewriteLocation(f"{fp.split('/')[-1]}:{lineno}", code, match.group(1) if match is not None else None, len(graph_rewrites(ctx)))
+    return RewriteLocation(f"{fp.split('/')[-1]}:{lineno}", code, match.group(1).split(",")[0] if match is not None else None,
+                           len(graph_rewrites(ctx)))
   def to_json(self): return asdict(self)
 
 @dataclass(frozen=True)