From b4ea45b4a656290b5e3848cfc923255c5825efc3 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 28 Mar 2025 15:24:53 +0800 Subject: [PATCH 1/5] fix viz recenter + worker cleanup (#9607) --- tinygrad/viz/index.html | 5 +++-- tinygrad/viz/lib/graph.js | 26 +++++++++++--------------- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html index 8f33163715..606f820c60 100644 --- a/tinygrad/viz/index.html +++ b/tinygrad/viz/index.html @@ -262,7 +262,7 @@ const x0 = document.querySelector(".kernel-list-parent").getBoundingClientRect().right; const x1 = document.querySelector(".metadata").getBoundingClientRect().left; const pad = 16; - const R = { x: x0+pad, y: mainRect.top+pad, width: x1-x0-2*pad, height: mainRect.height-2*pad }; + const R = { x: x0+pad, y: mainRect.top+pad, width: (x1>0 ? x1-x0 : mainRect.width)-2*pad, height: mainRect.height-2*pad }; const r = document.querySelector("#render").getBoundingClientRect(); if (r.width === 0) return; const scale = Math.min(R.width/r.width, R.height/r.height); @@ -416,7 +416,8 @@ document.querySelector(".collapse-btn").addEventListener("click", (e) => { isCollapsed = !isCollapsed; mainContainer.classList.toggle("collapsed", isCollapsed); - e.target.style.transform = isCollapsed ? "rotate(180deg)" : "rotate(0deg)"; + e.currentTarget.blur(); + e.currentTarget.style.transform = isCollapsed ? "rotate(180deg)" : "rotate(0deg)"; }); // **** resizer function appendResizer(element, { minWidth, maxWidth }, left=false) { diff --git a/tinygrad/viz/lib/graph.js b/tinygrad/viz/lib/graph.js index c3c08a69dd..e8e09b11b5 100644 --- a/tinygrad/viz/lib/graph.js +++ b/tinygrad/viz/lib/graph.js @@ -8,31 +8,27 @@ function intersectRect(r1, r2) { return {x:r1.x+dx*scale, y:r1.y+dy*scale}; } -const allWorkers = []; -let workerUrl = null; +let [workerUrl, worker, timeout] = [null, null, null]; window.renderGraph = async function(graph, additions, name) { - if (workerUrl == null) { - const resp = await Promise.all(["/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js","/lib/worker.js"].map(u => fetch(u))); - workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" })); - } - while (allWorkers.length) { - const { worker, timeout } = allWorkers.pop(); - worker.terminate(); - clearTimeout(timeout); - } - if (name === "View Memory Graph") { return renderMemoryGraph(graph); } d3.select("#bars").html(""); // ** start calculating the new layout (non-blocking) - worker = new Worker(workerUrl); + if (worker == null) { + const resp = await Promise.all(["/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js","/lib/worker.js"].map(u => fetch(u))); + workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" })); + worker = new Worker(workerUrl); + } else { + worker.terminate(); + worker = new Worker(workerUrl); + } + if (timeout != null) clearTimeout(timeout); const progressMessage = document.querySelector(".progress-message"); - const timeout = setTimeout(() => { + timeout = setTimeout(() => { progressMessage.style.display = "block"; }, 2000); - allWorkers.push({worker, timeout}); worker.postMessage({graph, additions}); worker.onmessage = (e) => { From d1e8598c81a1c5c3e887a4c577f0ed9e6695e43c Mon Sep 17 00:00:00 2001 From: Harsh Natuskar <74592384+hurrrsh@users.noreply.github.com> Date: Fri, 28 Mar 2025 14:22:21 +0530 Subject: [PATCH 2/5] add copy button in VIZ code-block (#9605) * works * only second block has copy * better function * better * ... * smol function * update copy-btn css * updates --- tinygrad/viz/index.html | 41 +++++++++++++++++++++++++++++++---------- 1 file changed, 31 insertions(+), 10 deletions(-) diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html index 606f820c60..76fdf3ca41 100644 --- a/tinygrad/viz/index.html +++ b/tinygrad/viz/index.html @@ -131,7 +131,7 @@ flex-direction: row; gap: 8px; } - .btn { + .btn, .copy-btn { outline: none; background-color: #1a1b26; border: 1px solid #4a4b56; @@ -149,6 +149,14 @@ background-color: #2a2b36; border-color: #5a5b66; } + .copy-btn { + position: absolute; + top: 5px; + right: 10px; + } + .code-container { + position: relative; + } .collapsed .container { display: none; } @@ -232,14 +240,28 @@ const toPath = ([fp, lineno]) => `${fp.replaceAll("\\", "/").split("/").pop()}:${lineno}`; const vsCodeOpener = (parts) => Object.assign(document.createElement("a"), { textContent: parts[parts.length-1]+"\n\n", href: "vscode://file"+parts.join("/"), style: "font-family: monospace; margin: 4px 0;" }); - const highlightedCodeBlock = (code, lang, wrap) => { - const pre = Object.assign(document.createElement("pre"), {className: wrap ? "wrap" : ""}); - // NOTE: since code is in textContent, we don't need DOMPurify - const codeEl = Object.assign(document.createElement("code"), { className: `language-${lang} code-block`, textContent: code}); - pre.appendChild(codeEl); + const highlightedCodeBlock = (code, lang, wrap, addCopy = false) => { + const container = document.createElement("div"); + container.className = "code-container"; + const pre = document.createElement("pre"); + pre.className = wrap ? "wrap" : ""; + const codeEl = pre.appendChild(document.createElement("code")); + codeEl.className = `language-${lang} code-block`; + codeEl.textContent = code; + container.appendChild(pre); + if (addCopy) { + const btn = container.appendChild(document.createElement("button")); + btn.className = "copy-btn"; + btn.textContent = "Copy"; + btn.onclick = () => navigator.clipboard.writeText(code).then(() => { + btn.textContent = "Copied!"; + setTimeout(() => btn.textContent = "Copy", 1500); + }); + } hljs.highlightElement(codeEl); - return pre; + return container; }; + const coloredToHTML = (str) => { const colors = ['gray','red','green','yellow','blue','magenta','cyan','white']; return str.replace(/\u001b\[(\d+)m(.*?)\u001b\[0m/g, (_, code, st) => { @@ -362,7 +384,7 @@ const metadata = document.querySelector(".container.metadata"); metadata.innerHTML = ""; metadata.appendChild(vsCodeOpener(kernel.loc.join(":").split("/"))); - metadata.appendChild(highlightedCodeBlock(kernel.code_line, "python", true)); + metadata.appendChild(highlightedCodeBlock(kernel.code_line, "python", true, false)); appendResizer(metadata, { minWidth: 20, maxWidth: 50 }); // ** code blocks let code = ret[currentRewrite].uop; @@ -371,8 +393,7 @@ code = kernel.kernel_code; lang = "cpp"; } - const codeBlock = highlightedCodeBlock(code, lang, false); - metadata.appendChild(codeBlock); + metadata.appendChild(highlightedCodeBlock(code, lang, false, true)); // ** rewrite list if (kernel.match_count >= 1) { const rewriteList = Object.assign(document.createElement("div"), { className: "rewrite-list" }) From 392a3113123d2585818a76d0cf59d1974b415446 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 28 Mar 2025 17:05:44 +0800 Subject: [PATCH 3/5] Revert "add copy button in VIZ code-block (#9605)" (#9610) This reverts commit d1e8598c81a1c5c3e887a4c577f0ed9e6695e43c. --- tinygrad/viz/index.html | 41 ++++++++++------------------------------- 1 file changed, 10 insertions(+), 31 deletions(-) diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html index 76fdf3ca41..606f820c60 100644 --- a/tinygrad/viz/index.html +++ b/tinygrad/viz/index.html @@ -131,7 +131,7 @@ flex-direction: row; gap: 8px; } - .btn, .copy-btn { + .btn { outline: none; background-color: #1a1b26; border: 1px solid #4a4b56; @@ -149,14 +149,6 @@ background-color: #2a2b36; border-color: #5a5b66; } - .copy-btn { - position: absolute; - top: 5px; - right: 10px; - } - .code-container { - position: relative; - } .collapsed .container { display: none; } @@ -240,28 +232,14 @@ const toPath = ([fp, lineno]) => `${fp.replaceAll("\\", "/").split("/").pop()}:${lineno}`; const vsCodeOpener = (parts) => Object.assign(document.createElement("a"), { textContent: parts[parts.length-1]+"\n\n", href: "vscode://file"+parts.join("/"), style: "font-family: monospace; margin: 4px 0;" }); - const highlightedCodeBlock = (code, lang, wrap, addCopy = false) => { - const container = document.createElement("div"); - container.className = "code-container"; - const pre = document.createElement("pre"); - pre.className = wrap ? "wrap" : ""; - const codeEl = pre.appendChild(document.createElement("code")); - codeEl.className = `language-${lang} code-block`; - codeEl.textContent = code; - container.appendChild(pre); - if (addCopy) { - const btn = container.appendChild(document.createElement("button")); - btn.className = "copy-btn"; - btn.textContent = "Copy"; - btn.onclick = () => navigator.clipboard.writeText(code).then(() => { - btn.textContent = "Copied!"; - setTimeout(() => btn.textContent = "Copy", 1500); - }); - } + const highlightedCodeBlock = (code, lang, wrap) => { + const pre = Object.assign(document.createElement("pre"), {className: wrap ? "wrap" : ""}); + // NOTE: since code is in textContent, we don't need DOMPurify + const codeEl = Object.assign(document.createElement("code"), { className: `language-${lang} code-block`, textContent: code}); + pre.appendChild(codeEl); hljs.highlightElement(codeEl); - return container; + return pre; }; - const coloredToHTML = (str) => { const colors = ['gray','red','green','yellow','blue','magenta','cyan','white']; return str.replace(/\u001b\[(\d+)m(.*?)\u001b\[0m/g, (_, code, st) => { @@ -384,7 +362,7 @@ const metadata = document.querySelector(".container.metadata"); metadata.innerHTML = ""; metadata.appendChild(vsCodeOpener(kernel.loc.join(":").split("/"))); - metadata.appendChild(highlightedCodeBlock(kernel.code_line, "python", true, false)); + metadata.appendChild(highlightedCodeBlock(kernel.code_line, "python", true)); appendResizer(metadata, { minWidth: 20, maxWidth: 50 }); // ** code blocks let code = ret[currentRewrite].uop; @@ -393,7 +371,8 @@ code = kernel.kernel_code; lang = "cpp"; } - metadata.appendChild(highlightedCodeBlock(code, lang, false, true)); + const codeBlock = highlightedCodeBlock(code, lang, false); + metadata.appendChild(codeBlock); // ** rewrite list if (kernel.match_count >= 1) { const rewriteList = Object.assign(document.createElement("div"), { className: "rewrite-list" }) From fa0ebbd2371f741c7a851c96423b1be1ba76a225 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Fri, 28 Mar 2025 19:06:09 +0700 Subject: [PATCH 4/5] jit: optimize before pickle (#9611) * jit: optimize before pickle * optimize weights * fix * mypy * mypy2 --- test/test_jit.py | 19 +++++++++++++++++++ tinygrad/engine/jit.py | 9 +++++++++ tinygrad/engine/memory.py | 7 ++++--- 3 files changed, 32 insertions(+), 3 deletions(-) diff --git a/test/test_jit.py b/test/test_jit.py index d1fee92976..4a146f5cf9 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -605,5 +605,24 @@ class TestJitFree(unittest.TestCase): fxn(Tensor([2])) self.assertEqual(x.item(), 8) + def test_optimize_weights(self): + if not hasattr(Device[Device.DEFAULT].allocator, '_offset'): raise unittest.SkipTest("optimize_weights useless") + + ext_tensor = Tensor([1,24,23,45,1]) + ext_tensor_2 = Tensor([2,2,2,2,2]) + @TinyJit + def fxn(x:Tensor): + out = (x*ext_tensor_2+ext_tensor).reshape(5,1).expand(5, 100).contiguous() + return out.sum() + for i in range(5): + out = fxn(Tensor([i,1,2,3,4])) + self.assertEqual(out.item(), 11400+200*i) + assert len(set([b.base for item in fxn.captured.jit_cache for b in item.bufs if b is not None])) == 4 + fxn.captured.optimize_weights() + assert len(set([b.base for item in fxn.captured.jit_cache for b in item.bufs if b is not None])) == 2 + + out = fxn(Tensor([11,1,2,3,4])) + self.assertEqual(out.item(), 13600) + if __name__ == '__main__': unittest.main() diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 125a549ccf..dfd5556456 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -150,6 +150,7 @@ class CapturedJit(Generic[ReturnType]): def __reduce__(self): # TODO: free_intermediates here? + self.optimize_weights() return self.__class__, (self.ret, self.jit_cache, self.input_replace, self.extra_view_inputs, self.expected_names, self.expected_st_vars_dtype_device) @@ -171,6 +172,14 @@ class CapturedJit(Generic[ReturnType]): if b._base is not None and b._base.allocated_views == 0: b._base.deallocate() self.__post_init__() # reset the graph state + def optimize_weights(self): + blacklist = [t.lazydata.buffer for t in get_parameters(self.ret)] + asgn = _internal_memory_planner([[b for item in self.jit_cache for b in item.bufs if b is not None and b not in blacklist]], ignore_checks=True) + self.jit_cache = [ExecItem(item.prg, [asgn.get(b,b) if b is not None else None for b in item.bufs]) for item in self.jit_cache] + for old, new in asgn.items(): + if old.is_allocated(): new.ensure_allocated().copyin(old.as_buffer()) + self.__post_init__() + # jit exec def __call__(self, input_buffers:list[Buffer], var_vals:dict[Variable, int]) -> ReturnType: # assign inputs diff --git a/tinygrad/engine/memory.py b/tinygrad/engine/memory.py index d53ae565ea..5314b0af79 100644 --- a/tinygrad/engine/memory.py +++ b/tinygrad/engine/memory.py @@ -9,12 +9,13 @@ from tinygrad.runtime.support.allocator import TLSFAllocator # **************** memory planning **************** -def _internal_memory_planner(buffers:list[list[Buffer]|tuple[Buffer, ...]], noopt_buffers=None, debug_prefix="") -> dict[Buffer, Buffer]: +def _internal_memory_planner(buffers:list[list[Buffer]], noopt_buffers=None, ignore_checks=False, debug_prefix="") -> dict[Buffer, Buffer]: if NO_MEMORY_PLANNER: return {} first_appearance, last_appearance, buf_to_opt = {}, {}, set() for i,u in enumerate(buffers): for buf in u: - if buf.is_allocated() or buf.base.is_allocated() or buf.lb_refcount > 0 or (noopt_buffers is not None and buf.base in noopt_buffers): continue + should_skip = buf.is_allocated() or buf.base.is_allocated() or buf.lb_refcount > 0 or (noopt_buffers is not None and buf.base in noopt_buffers) + if not ignore_checks and should_skip: continue if buf.base not in first_appearance: first_appearance[buf.base] = i last_appearance[buf.base] = i buf_to_opt.add(buf) @@ -63,6 +64,6 @@ def _internal_memory_planner(buffers:list[list[Buffer]|tuple[Buffer, ...]], noop def memory_planner(schedule:list[ScheduleItem]) -> list[ScheduleItem]: # Exclude buffers involved in load ops (e.g transfers) to preserve parallelism in graphs. - assigned = _internal_memory_planner([si.bufs for si in schedule], + assigned = _internal_memory_planner([list(si.bufs) for si in schedule], noopt_buffers={b for si in schedule if si.ast.op is not Ops.SINK for b in si.bufs}) return [ScheduleItem(si.ast, tuple(assigned.get(x, x) for x in si.bufs), si.metadata) for si in schedule] From a8ff85369ebe0771613d1af803ea123a51068e8c Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Fri, 28 Mar 2025 19:06:31 +0700 Subject: [PATCH 5/5] cpugraph for dsp (#9601) * cpugraph init * fixes * no cpu for now * mypy * fix --- tinygrad/runtime/graph/cpu.py | 44 +++++++++++++++++++++++++++++++++++ tinygrad/runtime/ops_dsp.py | 5 +++- 2 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 tinygrad/runtime/graph/cpu.py diff --git a/tinygrad/runtime/graph/cpu.py b/tinygrad/runtime/graph/cpu.py new file mode 100644 index 0000000000..454ce5f411 --- /dev/null +++ b/tinygrad/runtime/graph/cpu.py @@ -0,0 +1,44 @@ +from typing import cast +import itertools +from tinygrad.helpers import dedup, DEBUG, to_function_name +from tinygrad.engine.jit import GraphRunner, GraphException +from tinygrad.device import Buffer +from tinygrad.engine.realize import ExecItem, CompiledRunner +from tinygrad.ops import Variable +from tinygrad.dtype import dtypes +from tinygrad.renderer.cstyle import ClangRenderer + +class CPUGraph(GraphRunner): + def __init__(self, device, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[Variable, int]): + if not issubclass(type(device.renderer), ClangRenderer) and not isinstance(device.renderer, ClangRenderer): raise GraphException + super().__init__(jit_cache, input_rawbuffers, var_vals) + + self.base_bufs = dedup(b.base for ji in jit_cache for b in ji.bufs if b is not None and b not in input_rawbuffers) + self.base_rawbufs = [b._buf for b in self.base_bufs] + + targs = [(f"arg{i}", (x.dtype.ptr(), False)) for i,x in enumerate(input_rawbuffers)] + \ + [(f"cbuf{i}", (dtypes.char.ptr(), False)) for i in range(len(self.base_bufs))] + \ + sorted([(f"{v.expr}", (dtypes.int, False)) for v in var_vals]) + + def render_arg(buf): + if buf in input_rawbuffers: return f"arg{input_rawbuffers.index(buf)}" + return f"({device.renderer.render_dtype(buf.dtype)}*)(cbuf{self.base_bufs.index(buf.base)} + {buf.offset})" + + batched = ["void batched("+','.join([f"{device.renderer.render_dtype(x[1][0])} {x[0]}" for x in targs])+") {"] + for i, ji in enumerate(jit_cache): + args = [render_arg(buf) for buf in ji.bufs] + [x.expr for x in cast(CompiledRunner, ji.prg).p.vars] + batched.append(f" {to_function_name(cast(CompiledRunner, ji.prg).p.name)}({','.join(args)});") + batched.append("}") + + prep = [device.renderer._render(cast(CompiledRunner, ji.prg).p.uops) for i,ji in enumerate(jit_cache)] + funcs = dedup(device.renderer._render_body(prep[i][0], *prep[i][1:], cast(CompiledRunner, ji.prg).p.uops) for i,ji in enumerate(jit_cache)) + + defines = '\n'.join(set(itertools.chain.from_iterable(device.renderer._render_defines(cast(CompiledRunner, ji.prg).p.uops) for ji in jit_cache))) + entry = device.renderer._render_entry("batched", targs) + code = defines + '\n' + '\n'.join([''.join(f) for f in funcs]) + '\n'.join(batched) + '\n' + entry + + if DEBUG >= 4: print(code) + self.clprg = device.runtime("batched", device.compiler.compile_cached(code)) + + def __call__(self, rawbufs: list[Buffer], var_vals: dict[Variable, int], wait=False): + return self.clprg(*[x._buf for x in rawbufs], *self.base_rawbufs, *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0].expr)], wait=wait) diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 7a2102d043..e31d4db8b1 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -142,8 +142,11 @@ class DSPDevice(Compiled): with tempfile.NamedTemporaryFile(delete=False) as self.link_ld: self.link_ld.write(f"SECTIONS {{ . = 0x0; {sections_link}\n /DISCARD/ : {{ *(.note .note.* .gnu.hash .comment) }} }}".encode()) self.link_ld.flush() + + from tinygrad.runtime.graph.cpu import CPUGraph super().__init__(device, DSPAllocator(self), DSPRenderer(), - ClangCompiler("compile_dsp", ["-shared"] + compiler_args + [f"-T{self.link_ld.name}"], 'llvm-objdump'), functools.partial(DSPProgram, self)) + ClangCompiler("compile_dsp", ["-shared"] + compiler_args + [f"-T{self.link_ld.name}"], 'llvm-objdump'), functools.partial(DSPProgram, self), + functools.partial(CPUGraph, self)) fastrpc_shell = memoryview(bytearray(pathlib.Path('/dsp/cdsp/fastrpc_shell_3').read_bytes())) self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True)) ctypes.memmove(self.shell_buf.va_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes)