delete files that import ShapeTracker (#13805)

extra/optimization/extract_policynet.py

@@ -1,114 +0,0 @@
-import os, sys, sqlite3, pickle, random
-from tqdm import tqdm, trange
-from copy import deepcopy
-from tinygrad.nn import Linear
-from tinygrad.tensor import Tensor
-from tinygrad.nn.optim import Adam
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.codegen.opt.search import actions
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.helpers import getenv
-
-# stuff needed to unpack a kernel
-from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
-from tinygrad.dtype import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.uop.ops import Variable
-inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.opt.kernel import Opt, OptOps
-
-INNER = 256
-class PolicyNet:
-  def __init__(self):
-    self.l1 = Linear(1021,INNER)
-    self.l2 = Linear(INNER,INNER)
-    self.l3 = Linear(INNER,1+len(actions))
-  def __call__(self, x):
-    x = self.l1(x).relu()
-    x = self.l2(x).relu().dropout(0.9)
-    return self.l3(x).log_softmax()
-
-def dataset_from_cache(fn):
-  conn = sqlite3.connect(fn)
-  cur = conn.cursor()
-  cur.execute("SELECT * FROM beam_search")
-  X,A = [], []
-  for f in tqdm(cur.fetchall()):
-    Xs,As = [], []
-    try:
-      lin = Kernel(eval(f[0]))
-      opts = pickle.loads(f[-1])
-      for o in opts:
-        Xs.append(lin_to_feats(lin, use_sts=True))
-        As.append(actions.index(o))
-        lin.apply_opt(o)
-      Xs.append(lin_to_feats(lin, use_sts=True))
-      As.append(0)
-    except Exception:
-      pass
-    X += Xs
-    A += As
-  return X,A
-
-if __name__ == "__main__":
-  if getenv("REGEN"):
-    X,V = dataset_from_cache(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tinygrad_cache")
-    safe_save({"X": Tensor(X), "V": Tensor(V)}, "/tmp/dataset_policy")
-  else:
-    ld = safe_load("/tmp/dataset_policy")
-    X,V = ld['X'].numpy(), ld['V'].numpy()
-
-  print(X.shape, V.shape)
-  order = list(range(X.shape[0]))
-  random.shuffle(order)
-  X, V = X[order], V[order]
-
-  ratio = -256
-  X_test, V_test = Tensor(X[ratio:]), Tensor(V[ratio:])
-  X,V = X[:ratio], V[:ratio]
-  print(X.shape, V.shape)
-
-  net = PolicyNet()
-  #if os.path.isfile("/tmp/policynet.safetensors"): load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
-  optim = Adam(get_parameters(net))
-
-  def get_minibatch(X,Y,bs):
-    xs, ys = [], []
-    for _ in range(bs):
-      sel = random.randint(0, len(X)-1)
-      xs.append(X[sel])
-      ys.append(Y[sel])
-    return Tensor(xs), Tensor(ys)
-
-  Tensor.training = True
-  losses = []
-  test_losses = []
-  test_accuracy = 0
-  test_loss = float('inf')
-  for i in (t:=trange(500)):
-    x,y = get_minibatch(X,V,bs=256)
-    out = net(x)
-    loss = out.sparse_categorical_crossentropy(y)
-    optim.zero_grad()
-    loss.backward()
-    optim.step()
-    cat = out.argmax(axis=-1)
-    accuracy = (cat == y).mean()
-    t.set_description(f"loss {loss.numpy():7.2f} accuracy {accuracy.numpy()*100:7.2f}%, test loss {test_loss:7.2f} test accuracy {test_accuracy*100:7.2f}%")
-
-    losses.append(loss.numpy().item())
-    test_losses.append(test_loss)
-    if i % 10:
-      out = net(X_test)
-      test_loss = out.sparse_categorical_crossentropy(V_test).square().mean().numpy().item()
-      cat = out.argmax(axis=-1)
-      test_accuracy = (cat == y).mean().numpy()
-
-  safe_save(get_state_dict(net), "/tmp/policynet.safetensors")
-
-  import matplotlib.pyplot as plt
-  plt.plot(losses[10:])
-  plt.plot(test_losses[10:])
-  plt.show()

extra/optimization/extract_sa_pairs.py

@@ -1,129 +0,0 @@
-import sys, sqlite3, pickle, math
-from collections import defaultdict
-from tqdm import tqdm, trange
-import numpy as np
-
-# stuff needed to unpack a kernel
-from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
-from tinygrad.dtype import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.uop.ops import Variable
-inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.opt.kernel import Opt, OptOps
-
-# more stuff
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.search import actions
-from extra.optimization.helpers import lin_to_feats
-from extra.optimization.pretrain_valuenet import ValueNet
-from tinygrad.nn.optim import Adam
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-import random
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv
-
-def dataset_from_cache(fn):
-  conn = sqlite3.connect(fn)
-  cur = conn.cursor()
-  cur.execute("SELECT * FROM time_linearizer")
-  grouped = defaultdict(dict)
-  for f in tqdm(cur.fetchall()): grouped[f[0]][f[1:-1]] = pickle.loads(f[-1])
-
-  opts_to_outcome = {}
-
-  for ast,sk in grouped.items():
-    cnts = defaultdict(int)
-    for sks,tm in sk.items():
-      if sks[1] != 1: continue
-      opts = eval(sks[0])
-      cnts[(len(opts), sks[1])] += 1
-      opts_to_outcome[(ast, tuple(opts))] = tm
-    #print(cnts)
-
-  S,A,V = [], [], []
-  for ast,k in tqdm(opts_to_outcome):
-    if len(k) == 0: continue
-    old_tm = min(opts_to_outcome[(ast,k[:-1])])
-    new_tm = min(opts_to_outcome[(ast,k)])
-    if math.isinf(old_tm) or math.isinf(new_tm) or old_tm < 1e-9 or new_tm < 1e-9: continue
-    try:
-      lin = Kernel(eval(ast))
-    except Exception:
-      continue
-    lin.apply_opts(k[:-1])
-    act = k[-1]
-    log_ratio = math.log(old_tm/new_tm)
-    #print(f"ratio: {old_tm/new_tm:6.2f}x (log {log_ratio:5.2f}) from {str(act):50s} on {lin.colored_shape()}")
-    S.append(lin_to_feats(lin, use_sts=True))
-    A.append(actions.index(act))
-    V.append([log_ratio])  # NOTE: i have written the bug many times with this having the wrong dim
-
-  S, A, V = np.array(S), np.array(A), np.array(V, dtype=np.float32)
-  X = np.zeros((S.shape[0], S.shape[1]+len(actions)), dtype=np.float32)
-  X[:, :S.shape[1]] = S
-  X[range(S.shape[0]), S.shape[1]+A] = 1.0
-  return X, V
-
-def log_likelihood(x:Tensor, mu:Tensor, log_sigma:Tensor):
-  #print(x.shape, mu.shape, log_sigma.shape)
-  #return (x-mu).abs() * (-log_sigma).exp() + log_sigma
-  return (x-mu).square() * (-2*log_sigma).exp() / 2 + log_sigma
-
-if __name__ == "__main__":
-  if getenv("REGEN"):
-    X,V = dataset_from_cache(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tinygrad_cache")
-    safe_save({"X": Tensor(X), "V": Tensor(V)}, "/tmp/dataset")
-  else:
-    ld = safe_load("/tmp/dataset")
-    X,V = ld['X'].numpy(), ld['V'].numpy()
-
-  print(X.shape, V.shape)
-  order = list(range(X.shape[0]))
-  random.shuffle(order)
-  X, V = X[order], V[order]
-
-  ratio = -512
-  X_test, V_test = Tensor(X[ratio:]), Tensor(V[ratio:])
-  X,V = X[:ratio], V[:ratio]
-  print(X.shape, V.shape)
-
-  #print(X[0], V[0])
-  #print(X[-1], V[-1])
-  print(X.shape)
-
-  net = ValueNet(X.shape[1], 2)
-  optim = Adam(get_parameters(net))
-
-  def get_minibatch(X,Y,bs):
-    xs, ys = [], []
-    #random.seed(1337)
-    for _ in range(bs):
-      sel = random.randint(0, len(X)-1)
-      xs.append(X[sel])
-      ys.append(Y[sel])
-    return Tensor(xs), Tensor(ys)
-
-  Tensor.training = True
-  losses = []
-  test_losses = []
-  test_loss = float('inf')
-  for i in (t:=trange(2000)):
-    x,y = get_minibatch(X,V,bs=256)
-    out = net(x)
-    #loss = (out-y).square().mean()
-    loss = log_likelihood(y, out[:, 0:1], out[:, 1:2]).mean()
-    optim.zero_grad()
-    loss.backward()
-    optim.step()
-    t.set_description(f"loss {loss.numpy():7.2f}, test loss {test_loss:7.2f}")
-    losses.append(loss.numpy().item())
-    test_losses.append(test_loss)
-    if i % 10: test_loss = (net(X_test)[:, 0:1]-V_test).square().mean().numpy().item()
-
-  safe_save(get_state_dict(net), "/tmp/qnet.safetensors")
-
-  import matplotlib.pyplot as plt
-  plt.plot(losses[20:])
-  plt.plot(test_losses[20:])
-  plt.show()

extra/optimization/helpers.py

@@ -1,124 +0,0 @@
-# stuff needed to unpack a kernel
-from tinygrad import Variable
-from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.dtype import dtypes, PtrDType
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.helpers import getenv
-from tinygrad.engine.realize import get_program
-inf, nan = float('inf'), float('nan')
-UOps = Ops
-
-# kernel unpacker
-from tinygrad.codegen.opt.kernel import Kernel
-def ast_str_to_ast(ast_str:str) -> UOp: return eval(ast_str)
-def ast_str_to_lin(ast_str:str, opts=None): return Kernel(ast_str_to_ast(ast_str), opts=opts)
-def kern_str_to_lin(kern_str:str, opts=None):
-  (ast, applied_opts,) = eval(kern_str)
-  k = Kernel(ast, opts=opts)
-  k.apply_opts(applied_opts)
-  return k
-
-# load worlds, a dataset of about 12k kernels
-import gzip
-from pathlib import Path
-import random
-from tinygrad.helpers import dedup, DEBUG
-def load_worlds(filter_reduce=True, filter_noimage=True, filter_novariable=True):
-  fn = Path(__file__).parent.parent / "datasets/sops.gz"
-  ast_strs = dedup(gzip.open(fn).read().decode('utf-8').strip().split("\n"))
-  assert len(ast_strs) >= getenv("MIN_ASTS", 1000), f"dataset size = {len(ast_strs)} is too small"
-  if DEBUG >= 1: print(f"loaded {len(ast_strs)=} before filters")
-  if filter_reduce: ast_strs = [x for x in ast_strs if "REDUCE_AXIS" in x]
-  if filter_noimage: ast_strs = [x for x in ast_strs if "dtypes.image" not in x]
-  if filter_novariable: ast_strs = [x for x in ast_strs if "DEFINE_VAR" not in x]
-  if DEBUG >= 1: print(f"loaded {len(ast_strs)=} after filters {filter_reduce=}, {filter_noimage=}, {filter_novariable=}")
-  random.seed(1337)
-  random.shuffle(ast_strs)
-  return ast_strs
-
-def assert_same_lin(l1, l2):
-  assert l1.colored_shape() == l2.colored_shape()
-  assert all(x==y for x,y in zip(l1.sts, l2.sts))
-
-# get features
-import math
-
-MAX_DIMS = 16
-MAX_BUFS = 9
-def lin_to_feats(lin:Kernel, use_sts=True):
-  assert lin.shape_len < MAX_DIMS, "too many dims"
-
-  all_colors = ["blue", "cyan", "white", "green", "red", "magenta", "yellow"]
-  lc = [all_colors.index(x) for x in lin.colors()]
-
-  ret = []
-  # before, some generic linearizer stuff
-  ret.append(lin.upcasted)
-  ret.append(lin.local_dims)
-
-  # first, the full shape, including the colors
-  for s,os,c in zip(lin.full_shape,lin.output_shape,lc):
-    if isinstance(s, UOp):
-      ret.append(False)
-      ret += [0]*9
-    else:
-      ret.append(True)
-      ret.append(math.log2(s))
-      ret.append(min(33, s))
-      ret.append(math.log2(os))
-      ret.append(min(33, os))
-      ret.append(s%2 == 0)
-      ret.append(s%3 == 0)
-      ret.append(s%4 == 0)
-      ret.append(s%8 == 0)
-      ret.append(s%16 == 0)
-    cc = [0]*7
-    cc[c] = 1
-    ret += cc
-  ret += [0] * (17*(MAX_DIMS-len(lin.full_shape)))
-  ret = [float(x) for x in ret]
-
-  if use_sts:
-    my_sts = dedup([(x.shape == lin.full_shape, x.is_expanded(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
-    assert len(my_sts) < MAX_BUFS
-    sts_len = 3 + 5*MAX_DIMS
-    for s in my_sts:
-      ret.append(s[0])  # reduce
-      ret.append(s[2])  # has mask
-      ret.append(s[3])  # len views
-      for d in s[1]:
-        ret.append(d is None)
-        ret.append(d == 0)
-        ret.append(d == 1)
-        ret.append(min(33, d) if d is not None else -1)
-        if d is not None and d >= 1: ret.append(math.log2(d))
-        else: ret.append(-1)
-      ret += [0] * (5*(MAX_DIMS - len(s[1])))
-    ret += [0] * (sts_len*(MAX_BUFS - len(my_sts)))
-    assert len(ret) == 1021, f"wrong len {len(ret)}"
-  else:
-    assert len(ret) == 274, f"wrong len {len(ret)}"
-  return ret
-
-from tinygrad.device import Device, Buffer
-from tinygrad.codegen.opt.search import _ensure_buffer_alloc, _time_program
-from tinygrad.helpers import to_function_name, CACHELEVEL, diskcache_get, diskcache_put
-
-def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float:  # noqa: E501
-  key = {"ast": lin.ast.key, "opts": str(lin.applied_opts), "allow_test_size": allow_test_size,
-         "max_global_size": max_global_size, "clear_l2": clear_l2, "device": lin.opts.device, "suffix": lin.opts.suffix}
-  if not disable_cache and CACHELEVEL >= 2 and (val:=diskcache_get("time_linearizer", key)) is not None: return min(val)
-
-  dev = Device[lin.opts.device]
-  assert dev.compiler is not None
-
-  rawbufs = _ensure_buffer_alloc(rawbufs)
-  var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
-  p = get_program(lin.get_optimized_ast(), lin.opts)
-  tms = _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
-                      max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))
-
-  if CACHELEVEL >= 2: diskcache_put("time_linearizer", key, tms)
-  return min(tms)

extra/optimization/pretrain_valuenet.py

@@ -1,88 +0,0 @@
-from tinygrad.codegen.opt.kernel import Kernel
-from tqdm import tqdm, trange
-import math
-import random
-from tinygrad.tensor import Tensor
-from tinygrad.nn import Linear
-from tinygrad.nn.optim import Adam
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-
-# stuff needed to unpack a kernel
-from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
-from tinygrad.dtype import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.uop.ops import Variable
-inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.opt.kernel import Opt, OptOps
-
-from extra.optimization.helpers import lin_to_feats, MAX_DIMS
-
-# NOTE: this is not real value of the state, it's just a prediction of the runtime
-INNER = 512
-class ValueNet:
-  def __init__(self, feats=240, out=1):
-    self.l1 = Linear(feats,INNER)
-    self.l2 = Linear(INNER,INNER)
-    self.l3 = Linear(INNER,INNER)
-    self.l4 = Linear(INNER,out)
-  def __call__(self, x):
-    x = self.l1(x).relu()
-    x = self.l2(x).relu()
-    x = self.l3(x).relu().dropout(0.8)
-    return self.l4(x)
-
-if __name__ == "__main__":
-  net = ValueNet()
-  optim = Adam(get_parameters(net))
-
-  TEST_SIZE = 256
-
-  dset = open("/tmp/logtm").read().strip().split("\n")
-  random.seed(1337)
-  random.shuffle(dset)
-
-  X,Y = [], []
-  for i,x in enumerate(tqdm(dset)):
-    ast, opts, tms = eval(x)
-    lin = Kernel(ast)
-    for o in opts: lin.apply_opt(o)
-    if lin.shape_len >= MAX_DIMS: continue
-    if min(tms) == float('inf'): continue
-    X.append(lin_to_feats(lin))
-    Y.append([math.log(min(tms))])
-  print(f"got {len(X)} samples")
-
-  X_test,Y_test = Tensor(X[-TEST_SIZE:]), Tensor(Y[-TEST_SIZE:])
-  X,Y = X[:-TEST_SIZE], Y[:-TEST_SIZE]
-
-  def get_minibatch(X,Y,bs):
-    xs, ys = [], []
-    for _ in range(bs):
-      sel = random.randint(0, len(X)-1)
-      xs.append(X[sel])
-      ys.append(Y[sel])
-    return Tensor(xs), Tensor(ys)
-
-  Tensor.training = True
-  losses = []
-  test_losses = []
-  test_loss = float('inf')
-  for i in (t:=trange(2000)):
-    x,y = get_minibatch(X,Y,bs=256)
-    out = net(x)
-    loss = (out-y).square().mean()
-    optim.zero_grad()
-    loss.backward()
-    optim.step()
-    t.set_description(f"loss {loss.numpy():7.2f}, test loss {test_loss:7.2f}")
-    losses.append(loss.numpy().item())
-    test_losses.append(test_loss)
-    if i % 10: test_loss = (net(X_test)-Y_test).square().mean().numpy().item()
-
-  safe_save(get_state_dict(net), "/tmp/valuenet.safetensors")
-
-  import matplotlib.pyplot as plt
-  plt.plot(losses[200:])
-  plt.plot(test_losses[200:])
-  plt.show()

extra/test_pyrender.py

@@ -1,20 +0,0 @@
-from extra.optimization.helpers import load_worlds, ast_str_to_ast
-from tinygrad.helpers import tqdm
-from tinygrad.uop.ops import pyrender, UOp, Ops
-from tinygrad import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker, View
-inf, nan = float('inf'), float('nan')
-
-if __name__ == "__main__":
-  ast_strs = load_worlds()
-  for i, ast_str in enumerate(tqdm(ast_strs)):
-    good_ast = ast_str_to_ast(ast_str)
-    code = '\n'.join(pyrender(good_ast))
-    print("\n***************\n\n"+code)
-    exec(code)
-    if str(good_ast) != str(ast):
-      print(code)
-      print("MISMATCH")
-      print(good_ast)
-      print(ast)
-      break