Function compile_net

extra/export_model.py:20–45

(linear:UOp, output_bufs:List[Buffer])

  return (u for u in linear.toposort(gate=lambda x: x.op not in _KERNEL_ASTS) if u.op is Ops.CALL and u.src[0].op in _KERNEL_ASTS)

def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], List, Dict[str,Tuple[int,DType,int]], Dict[str,Buffer]]:
  output_name = {id(b): f"output{i}" for i, b in enumerate(output_bufs)}
  functions, bufs, bufs_to_save, statements, n = {}, {}, {}, [], 0

  def name_of(bu:UOp, is_out:bool) -> str:
    nonlocal n
    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg), f"input{bu.arg}", prod(bu.shape)*bu.dtype.itemsize
    else:
      b = bu.buffer
      key, size = (id(b.base), b.offset, b.size, b.dtype), b.size*b.dtype.itemsize
      if key in bufs: return bufs[key][0]
      if (name:=output_name.get(id(b))) is None:
        name, n = f"buf_{n}", n+1
        if not is_out: bufs_to_save[name] = b
    bufs[key] = (name, size, bu.dtype, key)
    return name

  for call in iter_kernel_calls(linear):
    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
    info = prg.arg
    functions[info.function_name] = prg.src[3].arg
    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + [v for v in info.vars if v.op is Ops.DEFINE_VAR]
    statements.append((info.function_name, cargs, info.global_size, info.local_size))

  return functions, statements, {name:(size, dtype, key) for name, size, dtype, key in bufs.values()}, bufs_to_save

def jit_model(model, *args) -> Tuple[UOp, List[Buffer]]:
  assert hasattr(model, "forward") or callable(model), "model needs a forward function"