Merge branch 'master' into dsp_search

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()

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

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]

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)

tinygrad/runtime/ops_dsp.py

@@ -428,8 +428,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)

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) {

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) => {