From ab8c5d04ab36ff1e2992290e565ee3d42dec2eb4 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Tue, 24 Jun 2025 13:59:37 +0300
Subject: [PATCH] viz: convert to function_name in server [pr] (#10951)

* viz: convert to function_name in server [pr]

* it exists
---
 tinygrad/viz/js/index.js  | 2 +-
 tinygrad/viz/js/worker.js | 2 +-
 tinygrad/viz/serve.py     | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index f79d407d27..a9a4e5ac5d 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -252,7 +252,7 @@ async function renderProfiler() {
     }
   }
   const kernelMap = new Map();
-  for (const [i, c] of ctxs.entries()) kernelMap.set(c.name.replace(/\x1b\[\d+m(.*?)\x1b\[0m/g, "$1"), { name:c.name, i });
+  for (const [i, c] of ctxs.entries()) kernelMap.set(c.function_name, { name:c.name, i });
   // place devices on the y axis and set vertical positions
   const [tickSize, padding] = [10, 8];
   const deviceList = document.getElementById("device-list");
diff --git a/tinygrad/viz/js/worker.js b/tinygrad/viz/js/worker.js
index 5ff131e6c7..94fdbb8f81 100644
--- a/tinygrad/viz/js/worker.js
+++ b/tinygrad/viz/js/worker.js
@@ -12,7 +12,7 @@ onmessage = (e) => {
   for (let [k, {label, src, ref, ...rest }] of Object.entries(graph)) {
     const idx = ref ? ctxs.findIndex(k => k.ref === ref) : -1;
     // replace colors in label
-    if (idx != -1) label += `\ncodegen@${ctxs[idx].name.replace(/\x1b\[\d+m(.*?)\x1b\[0m/g, "$1")}`;
+    if (idx != -1) label += `\ncodegen@${ctxs[idx].function_name}`;
     // adjust node dims by label size + add padding
     let [width, height] = [0, 0];
     for (line of label.split("\n")) {
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index 1b58120975..46674b9690 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -25,7 +25,7 @@ def get_metadata(keys:list[Any], contexts:list[list[TrackedGraphRewrite]]) -> li
   ret = []
   for k,v in zip(keys, contexts):
     steps = [{"name":s.name, "loc":s.loc, "depth":s.depth, "match_count":len(s.matches), "code_line":lines(s.loc[0])[s.loc[1]-1].strip()} for s in v]
-    if isinstance(k, ProgramSpec): ret.append({"name":k.name, "kernel_code":k.src, "ref":id(k.ast), "steps":steps})
+    if isinstance(k, ProgramSpec): ret.append({"name":k.name, "kernel_code":k.src, "ref":id(k.ast), "function_name":k.function_name, "steps":steps})
     else: ret.append({"name":str(k), "steps":steps})
   return ret