Function name_of

extra/export_model.py:24–35

(bu:UOp, is_out:bool)

  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]