Method test_link_sched_codegen

test/null/test_viz.py:338–357

(self)

  # schedule graph CALL nodes have a link to jump to codegen
  def test_link_sched_codegen(self):
    with save_viz() as viz:
      c1 = Tensor.empty(4, device="NULL").add(1)
      c2 = Tensor.empty(8, device="NULL").add(1)
      with Context(SCACHE=0):
        sched = c1.schedule_linear(c2)
      from tinygrad.engine.realize import compile_linear
      sched = compile_linear(sched)
      with Context(NO_COLOR=0):
        prgs = [to_program(si.src[0], Device[c1.device].renderer).arg.name for si in sched.src]
    lst = viz.list_items()
    sched_idx = next(i for i,l in enumerate(lst) if l["name"].startswith("Schedule"))
    viz_kernel = next(i for i,s in enumerate(lst[sched_idx]["steps"]) if s["name"] == "View Kernel Graph")
    with Context(NO_COLOR=1):
      graph = next(viz.get_details(sched_idx, viz_kernel))["graph"]
    call_nodes = [n for n in graph.values() if n["label"].startswith("CALL")]
    for i,n in enumerate(call_nodes):
      assert n["ref"] is not None
      self.assertEqual(lst[n["ref"]]["name"], prgs[i])
      assert ansistrip(prgs[i]) in n["label"], f"CALL must contain kernel name, got {n['label']}"

  def test_link_sched_codegen_beam(self):
    with Context(BEAM=2):