start work on auto opt (#2034)

* start work on auto opt

* lin failure

* not beating hcopt

* greedy

* timing is fast

* codegen.search

* greedy search in handcode_opt

* track running gflops

* clean up those files

* no failure

examples/handcode_resnet50_opt.py

@@ -2,11 +2,14 @@ from typing import List
 from models.resnet import ResNet50
 from tinygrad.tensor import Tensor
 from tinygrad.ops import LoadOps, Device, Compiled
-from tinygrad.codegen.kernel import LinearizerOptions
 from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.helpers import ansilen, DEBUG
+from tinygrad.codegen.search import bufs_from_lin, time_linearizer, get_linearizer_actions
+from tinygrad.helpers import ansilen, DEBUG, getenv
 from tinygrad.graph import print_tree
 
+import shelve
+global_db = shelve.open("/tmp/greedy_cache")
+
 if __name__ == "__main__":
   mdl = ResNet50()
   seen = set()
@@ -26,50 +29,55 @@ if __name__ == "__main__":
 
   # work with the schedule
   total_tm = 0
+  running_gflops = 0
   for i,si in enumerate(sched):
     if DEBUG >= 2: print_tree(si.ast)
 
-    # enable only one kernel to focus on it
-    #if i != 1: continue
+    # create output/input buffers (NOTE: bufs_from_lin is slower, so we don't use it. TODO: fix)
+    rawbufs = [device.buffer(si.out.st.size(), si.out.dtype)] + [device.buffer(x.st.size(), x.dtype) for x in si.inputs]
+    #rawbufs = bufs_from_lin(lin)
 
     # "linearize" the op into uops in different ways
     lins:List[Linearizer] = []
 
-    if Device.DEFAULT == "METAL" and i == 1:
-      # through careful work, we discovered 1,8,0
-      for big_chomp in [1,2]: #[1,2,4,8,16]:
-        for lil_chomp in [2,4,7,8,14]:
-          for upcasted in [0,1,2]:
-            lin = Linearizer(si.ast, device.linearizer_opts)
-            lin.reshape_and_permute(lambda x: (4096//big_chomp,big_chomp,56//lil_chomp,lil_chomp,56//lil_chomp,lil_chomp)+x[-2:], [0,2,4,1,3,5,6,7])
-            lin.upcasted += upcasted
-            lin.local_dims += 3
-            lins.append(lin)
-    else:
-      # try with and without tensor cores
-      for tc in [0,1]:
-        lin = Linearizer(si.ast, device.linearizer_opts)
-        lin.hand_coded_optimizations(use_tensor_cores=tc)
-        lins.append(lin)
+    # always try hand coded opt
+    lin = Linearizer(si.ast, device.linearizer_opts)
+    lin.hand_coded_optimizations()
+    lins.append(lin)
 
-    # create output/input buffers
-    rawbufs = [device.buffer(si.out.st.size(), si.out.dtype)] + [device.buffer(x.st.size(), x.dtype) for x in si.inputs]
+    # maybe try tensor cores
+    lin = Linearizer(si.ast, device.linearizer_opts)
+    if lin.apply_tensor_cores():
+      lins.append(lin)
+
+    # try a greedy search
+    if getenv("GREEDY"):
+      lin = Linearizer(si.ast, device.linearizer_opts)
+      if str(lin.ast) in global_db:
+        for ao in global_db[str(lin.ast)]:
+          lin.apply_opt(ao)
+      else:
+        while 1:
+          acted_lins = get_linearizer_actions(lin)
+          tm, gflops = time_linearizer(lin, rawbufs)
+          timed_lins = {k:time_linearizer(v, rawbufs)[0] for k,v in acted_lins.items()}
+          opts = sorted(timed_lins.items(), key=lambda x: x[1])
+          if len(opts) == 0 or opts[0][1] >= tm: break   # we are done
+          lin = acted_lins[opts[0][0]]
+          if DEBUG >= 1: print(f"{opts[0][1]*1e3:10.2f} ms from {len(opts):3d} actions", lin.colored_shape())
+        global_db[str(lin.ast)] = lin.applied_opts
+      lins.append(lin)
 
     # benchmark the programs
     choices = []
     for lin in lins:
-      prg = device.to_program(lin)
-
-      # benchmark it by running 10 times
-      try:
-        tm = min([prg(rawbufs, force_wait=True) for _ in range(10)])
-        choices.append((tm, lin))
-      except AssertionError:
-        tm = float('inf')
+      tm, gflops = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, should_copy=False)
+      choices.append((tm, gflops, lin))
 
       # print all kernels
-      if DEBUG >= 1: print(f"                 kernel {i:2d} {lin.display_name+' '*(37-ansilen(lin.display_name))} {str(lin.global_size):18s} {str(lin.local_size):12s} takes {tm*1000:7.2f} ms, {lin.info.flops*1e-9/tm:6.0f} GFLOPS")
-    tm, lin = sorted(choices, key=lambda x: x[0])[0]
-    print(f"*** {total_tm*1000:7.2f} ms : kernel {i:2d} {lin.display_name+' '*(37-ansilen(lin.display_name))} {str(lin.global_size):18s} {str(lin.local_size):12s} takes {tm*1000:7.2f} ms, {lin.info.flops*1e-9/tm:6.0f} GFLOPS")
+      if DEBUG >= 1: print(f"                 kernel {i:2d} {lin.display_name+' '*(37-ansilen(lin.display_name))} {str(lin.global_size):18s} {str(lin.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS")
+    tm, gflops, lin = sorted(choices, key=lambda x: x[0])[0]
+    print(f"*** {total_tm*1000:7.2f} ms : kernel {i:2d} {lin.display_name+' '*(37-ansilen(lin.display_name))} {str(lin.global_size):18s} {str(lin.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS")
     total_tm += tm
-  print(f"******* total {total_tm*1000:.2f} ms")
+    running_gflops += gflops * tm
+  print(f"******* total {total_tm*1000:.2f} ms, {running_gflops/total_tm:6.0f} GFLOPS")

extra/optimization/go.py

@@ -1,99 +0,0 @@
-import sys
-import random
-from collections import defaultdict
-from tqdm import tqdm
-from tinygrad.helpers import dedup, ImageDType, getenv, ansilen
-from tinygrad.graph import print_tree
-from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.lazy import vars_from_ast
-from tinygrad.shape.symbolic import sym_infer
-from tinygrad.ops import Device, Compiled, MemBuffer
-
-# stuff needed to unpack a kernel
-from tinygrad.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
-from tinygrad.helpers import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.shape.symbolic import Variable
-inf, nan = float('inf'), float('nan')
-
-if __name__ == "__main__":
-  ast_strs = dedup(open(sys.argv[1]).read().strip().split("\n"))
-
-  # reduce kernels only, no ImageDType
-  ast_strs = [x for x in ast_strs if "dtypes.image" not in x]
-  #ast_strs = [x for x in ast_strs if "ReduceOps" in x and "dtypes.image" not in x]
-
-  # the device we are optimizing for
-  device: Compiled = Device[Device.DEFAULT]
-  print(f"optimizing for {Device.DEFAULT}")
-
-  # random first kernels
-  random.seed(1337)
-  random.shuffle(ast_strs)
-  #ast_strs = ast_strs[:5000]
-
-  print(f"loaded {len(ast_strs)} kernels")
-
-  atm = []
-  agflops = []
-  for ast_str in tqdm(ast_strs):
-    ast = eval(ast_str)
-
-    # linearize
-    lin = Linearizer(ast)
-    preopt = lin.colored_shape()
-    if not lin.apply_tensor_cores(getenv("TC", 1)): lin.hand_coded_optimizations()
-    postopt = lin.colored_shape()
-    #print(lin.applied_opts)
-
-    """
-    # linearize_alt
-    lin2 = Linearizer(ast)
-    lin2.hand_coded_optimizations_old()
-    postopt_alt = lin2.colored_shape()
-
-    print(preopt+' '*(37-ansilen(preopt)), "->", postopt)
-    assert postopt == postopt_alt, f"{postopt} != {postopt_alt}"
-    for s1,s2 in zip(lin.sts, lin2.sts): assert s1 == s2
-    continue
-    """
-
-    lin.linearize()
-
-    # create output/input buffers
-    bufsts = defaultdict(list)
-    for x in lin.bufs:
-      if isinstance(x, MemBuffer):
-        bufsts[x.idx].append(x)
-    buffer_count = len(bufsts)
-    rawbufs = [None]*buffer_count
-    for k,x in bufsts.items():
-      rawbufs[k] = device.buffer(max(y.st.size() for y in x), x[0].dtype)
-    assert all(x is not None for x in rawbufs)
-
-    # example var vals
-    var_vals = {k:k.min for k in vars_from_ast(ast)}
-
-    # time
-    prg = device.to_program(lin)
-    tm = min([prg(rawbufs, var_vals, force_wait=True) for _ in range(10)])
-    atm.append(tm)
-
-    # print
-    #print_tree(ast)
-    #for u in lin.uops: print(u)
-    gflops = sym_infer(lin.info.flops, var_vals)*1e-9/tm
-    agflops.append(gflops)
-    if tm*1e6 > 100:
-      print(f"{len(lin.uops):4d} uops, {str(lin.global_size):18s} {str(lin.local_size):12s} {tm*1e6:8.2f} us {gflops:7.2f} GFLOPS", preopt+' '*(37-ansilen(preopt)), "->", postopt)
-
-  print(f"all kernels ran in {sum(atm)*1e3:.2f} ms")
-
-  if getenv("SHOW"):
-    import matplotlib.pyplot as plt
-    #plt.hist(agflops, bins=100)
-    #plt.yscale('log')
-    plt.scatter(atm, agflops)
-    plt.xscale('log')
-    plt.show()

extra/optimization/greedy.py

@@ -0,0 +1,43 @@
+import numpy as np
+import random
+from copy import deepcopy
+from tqdm import tqdm
+from tinygrad.helpers import dedup, ansilen
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
+from tinygrad.codegen.linearizer import Linearizer
+from tinygrad.tensor import Tensor
+
+from tinygrad.codegen.search import get_linearizer_actions, time_linearizer, bufs_from_lin, actions
+
+#from extra.optimization.pretrain import PolicyNet
+from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats
+
+if __name__ == "__main__":
+  # load worlds
+  ast_strs = load_worlds()
+  for ep_num,ast_str in enumerate(ast_strs):
+    print("\nEPISODE", ep_num)
+    lin = ast_str_to_lin(ast_str)
+
+    linhc = deepcopy(lin)
+    linhc.hand_coded_optimizations()
+    if not all(x in actions for x in linhc.applied_opts):
+      print("skipping", linhc.colored_shape())
+      continue
+
+    rawbufs = bufs_from_lin(lin)
+    tm1, gf1 = time_linearizer(linhc, rawbufs)
+    print(f"{tm1:10.2f}", linhc.colored_shape(), f"with {len(linhc.applied_opts)} actions from {len(actions)} action space")
+
+    while 1:
+      tm, gflops = time_linearizer(lin, rawbufs)
+      print(f"{tm:10.2f}", lin.colored_shape())
+      acted_lins = get_linearizer_actions(lin)
+      if len(acted_lins) == 0: break
+
+      best_tm, best_lin = tm, lin
+      for l in list(acted_lins.values()):
+        tm, gflops = time_linearizer(l, rawbufs)
+        if tm < best_tm: best_tm, best_lin = tm, l
+      if lin == best_lin: break
+      lin = best_lin

extra/optimization/helpers.py

@@ -0,0 +1,59 @@
+# stuff needed to unpack a kernel
+from tinygrad.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
+from tinygrad.helpers import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+from tinygrad.shape.symbolic import Variable
+inf, nan = float('inf'), float('nan')
+
+# kernel unpacker
+from tinygrad.codegen.linearizer import Linearizer
+def ast_str_to_lin(ast_str): return Linearizer(eval(ast_str))
+
+# load worlds
+import random
+from tinygrad.helpers import dedup
+def load_worlds():
+  ast_strs = dedup(open("/tmp/sops").read().strip().split("\n"))
+  ast_strs = [x for x in ast_strs if "ReduceOps" in x and "dtypes.image" not in x and "Variable" not in x]
+  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
+from tinygrad.shape.symbolic import Node
+
+MAX_DIMS = 16
+def lin_to_feats(lin):
+  all_colors = ["blue", "cyan", "white", "green", "red", "magenta", "yellow"]
+  lc = [all_colors.index(x) for x in lin.colors()]
+  #my_sts = dedup([(x.shape == lin.full_shape, x.real_strides()) for x in lin.sts[1:]])
+
+  # first, the full shape, including the colors
+  ret = []
+  for s,c in zip(lin.full_shape,lc):
+    if isinstance(s, Node):
+      ret.append(False)
+      ret += [0]*7
+    else:
+      ret.append(True)
+      ret.append(math.log2(s))
+      ret.append(min(33, s))
+      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] * (15*(MAX_DIMS-len(lin.full_shape)))
+  ret = [float(x) for x in ret]
+
+  assert len(ret) == 240, f"wrong len {len(ret)}"
+  return ret

extra/optimization/pretrain.py

@@ -0,0 +1,64 @@
+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.search import actions
+from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
+
+INNER = 256
+class PolicyNet:
+  def __init__(self):
+    self.l1 = Linear(240,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()
+    return self.l3(x).log_softmax()
+
+if __name__ == "__main__":
+  ast_strs = load_worlds()
+
+  net = PolicyNet()
+  load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
+  optim = Adam(get_parameters(net))
+
+  X,Y = [], []
+  steps = 0
+
+  for ast_str in ast_strs:
+    lin = ast_str_to_lin(ast_str)
+    linhc = deepcopy(lin)
+    linhc.hand_coded_optimizations()
+    print(lin.colored_shape(50), "->", linhc.colored_shape())
+
+    lin2 = deepcopy(lin)
+    for o in linhc.applied_opts:
+      X.append(lin_to_feats(lin2))
+      Y.append(actions.index(o)+1)
+      lin2.apply_opt(o)
+    X.append(lin_to_feats(lin2))
+    Y.append(0)
+    assert_same_lin(linhc, lin2)
+
+    BS = 64
+    if len(X) >= BS:
+      Tensor.no_grad, Tensor.training = False, True
+      x,y = Tensor(X[:BS]), Tensor(Y[:BS])
+      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()
+      print(loss.numpy(), accuracy.numpy())
+
+      X = X[BS:]
+      Y = Y[BS:]
+
+      if steps%10 == 0:
+        safe_save(get_state_dict(net), "/tmp/policynet.safetensors")
+        print("saved model")
+      steps += 1

extra/optimization/testpolicynet.py

@@ -0,0 +1,37 @@
+import numpy as np
+np.set_printoptions(suppress=True)
+from copy import deepcopy
+from tinygrad.tensor import Tensor
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
+from tinygrad.codegen.search import bufs_from_lin, time_linearizer, actions
+from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats
+from extra.optimization.pretrain import PolicyNet
+
+if __name__ == "__main__":
+  net = PolicyNet()
+  load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
+
+  ast_strs = load_worlds()
+
+  for ep_num,ast_str in enumerate(ast_strs):
+    print("\nEPISODE", ep_num)
+    lin = ast_str_to_lin(ast_str)
+    rawbufs = bufs_from_lin(lin)
+
+    linhc = deepcopy(lin)
+    linhc.hand_coded_optimizations()
+    tm, gflops = time_linearizer(linhc, rawbufs)
+    print(f"{tm:10.2f}", linhc.colored_shape())
+
+    while 1:
+      probs = net(Tensor([lin_to_feats(lin)]))
+      dist = probs.exp().numpy()
+      act = dist.argmax()
+      if act == 0: break
+      try:
+        lin.apply_opt(actions[act-1])
+      except Exception:
+        print("FAILED")
+        break
+      tm, gflops = time_linearizer(lin, rawbufs)
+      print(f"{tm:10.2f}", lin.colored_shape())

test/test_linearizer_failures.py

@@ -0,0 +1,21 @@
+import unittest
+from tinygrad.codegen.linearizer import Linearizer
+from tinygrad.ops import Device
+
+# stuff needed to unpack a kernel
+from tinygrad.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
+from tinygrad.helpers import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+from tinygrad.shape.symbolic import Variable
+inf, nan = float('inf'), float('nan')
+
+class TestLinearizerFailures(unittest.TestCase):
+  @unittest.skip("this is currently failing")
+  def test_failure_1(self):
+    ast = LazyOp(op=BinaryOps.ADD, src=(LazyOp(op=BinaryOps.ADD, src=(LazyOp(op=ReduceOps.SUM, src=(LazyOp(op=BufferOps.MEM, src=(), arg=MemBuffer(idx=1, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(32, 16, 16), strides=(16, 1, 0), offset=0, mask=None, contiguous=False),)))),), arg=(32, 16, 1)), LazyOp(op=BufferOps.MEM, src=(), arg=MemBuffer(idx=2, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(32, 16, 1), strides=(0, 1, 0), offset=0, mask=None, contiguous=False),))))), arg=None), LazyOp(op=BufferOps.MEM, src=(), arg=MemBuffer(idx=1, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(32, 16, 1), strides=(16, 1, 0), offset=0, mask=None, contiguous=True),))))), arg=None)
+    lin = Linearizer(ast)
+    prg = Device[Device.DEFAULT].to_program(lin)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad/codegen/kernel.py

@@ -1,7 +1,7 @@
 from typing import NamedTuple, Optional, List, Tuple, cast, Dict
 import itertools
 from tinygrad.ops import LazyOp, FlopCounter, get_lazyop_info, ReduceOps, MemBuffer, BufferOps, Device, Compiled
-from tinygrad.helpers import dedup, dtypes, colored, ImageDType, DType, all_int
+from tinygrad.helpers import dedup, dtypes, colored, ImageDType, DType, all_int, ansilen
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.symbolic import sint
 from tinygrad.shape.view import strides_for_shape
@@ -83,6 +83,9 @@ class Kernel:
     self.global_size: Optional[List[int]] = None
     self.local_size: Optional[List[int]] = None
 
+  @property
+  def membufs(self) -> List[MemBuffer]: return [x for x in self.bufs if isinstance(x, MemBuffer)]
+
   def has_variable_shape(self) -> bool:
     for b in self.bufs:
       if not all_int(b.st.views[-1].shape): return True
@@ -152,7 +155,10 @@ class Kernel:
     assert len(colors) == self.shape_len, "colors size mismatch"
     return colors
 
-  def colored_shape(self) -> str: return ' '.join(colored(s, color) for s,color in zip([f"{s:4d}" if isinstance(s, int) else s for s in self.full_shape], self.colors()))
+  def colored_shape(self, pad=None) -> str:
+    ret = ' '.join(colored(s, color) for s,color in zip([f"{s:4d}" if isinstance(s, int) else s for s in self.full_shape], self.colors()))
+    if pad: ret += ' '*(pad-ansilen(ret))
+    return ret
   def printbufs(self, prefix=""):
     for i,st in enumerate(self.sts):
       print(prefix, f"{i:3d} {str(self.bufs[i]):47s}", st.views)

tinygrad/codegen/optimizer.py

@@ -15,6 +15,7 @@ class Opt:
   op: OptOps
   axis: int
   amt: int
+  def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, amt={self.amt})"
 
 class OptimizedKernel(Kernel):
   def __init__(self, ast:LazyOp, opts:Optional[LinearizerOptions]=None):
@@ -155,7 +156,7 @@ class OptimizedKernel(Kernel):
   # ******************** high level optimizers ********************
 
   # TODO: unify this
-  def apply_tensor_cores(self, use_tensor_cores):
+  def apply_tensor_cores(self, use_tensor_cores=1):
     # should use HIP tensor cores?
     if use_tensor_cores != 0 and self.opts.device == "HIP" and self.reduceop and self.reduceop.op == ReduceOps.SUM and \
         isinstance(self.reduceop.src[0], LazyOp) and self.reduceop.src[0].op == UnaryOps.CAST and \
@@ -320,7 +321,7 @@ class OptimizedKernel(Kernel):
     self.applied_opts.append(opt)
     assert self.full_shape[opt.axis] % opt.amt == 0, "no longer valid shift"
     if opt.op == OptOps.LOCAL:        # cyan
-      assert opt.axis < self.first_reduce, "can't local reduce axis"
+      assert opt.axis < (self.first_reduce-self.local_dims), "can't local a local or reduce"
       self.shift_to(opt.axis, opt.amt, insert_before=self.first_reduce)
       self.local_dims += 1
     elif opt.op == OptOps.GROUP:      # green
@@ -330,6 +331,7 @@ class OptimizedKernel(Kernel):
       self.shift_to(opt.axis, opt.amt, top=True, insert_before=self.first_reduce + len(self.group_for_reduce))
       self.group_for_reduce.append(opt.amt)
     elif opt.op == OptOps.UPCAST:     # yellow (or purple if it's a reduce axis)
+      assert opt.axis < self.shape_len-self.upcasted, "can't upcasted already upcasted"
       self.shift_to(opt.axis, opt.amt, insert_before=None if opt.axis < self.first_reduce else len(self.full_unupcasted_shape))
       self.upcast()
     self.simplify_ones()
@@ -427,9 +429,10 @@ class OptimizedKernel(Kernel):
     # if last dim is small(ish) and it's a reduce dim, upcast the reduce (loop unrolling). no simplify needed since it's just an upcast. NOTE: careful, this has broken VALIDHACKS
     if self.first_reduce < (self.shape_len-self.upcasted) and (len(list(self.shape_offsets(self.full_buf_index))) <= 4 or not any(r for _,_,r in self.upcasted_axis(self.full_buf_index))):
       if (s:=self.full_unupcasted_shape[-1]) <= 32 and isinstance(s, int):  # NOTE: cannot loop unroll symbolic axis
-        self.upcast()
+        self.apply_opt(Opt(OptOps.UPCAST, len(self.full_unupcasted_shape)-1, s))
         # if it's small, upcast a second reduce dimension too
-        if self.first_reduce < (self.shape_len-self.upcasted) and s <= 3 and self.full_unupcasted_shape[-1] <= 3: self.upcast()
+        if self.first_reduce < (self.shape_len-self.upcasted) and s <= 3 and (s2:=self.full_unupcasted_shape[-1]) <= 3 and isinstance(s2, int):
+          self.apply_opt(Opt(OptOps.UPCAST, len(self.full_unupcasted_shape)-1, s2))
       else:
         for splits in [4]:
           if self.full_unupcasted_shape[-1]%splits == 0:

tinygrad/nn/__init__.py

@@ -74,7 +74,7 @@ class Linear:
     bound = 1 / math.sqrt(self.weight.shape[1])
     self.bias = Tensor.uniform(out_features, low=-bound, high=bound) if bias else None
 
-  def __call__(self, x):
+  def __call__(self, x:Tensor):
     return x.linear(self.weight.transpose(), self.bias)
 
 class GroupNorm:

tinygrad/tensor.py

@@ -425,7 +425,7 @@ class Tensor:
 
   # ***** reduce ops *****
 
-  def _reduce(self, fxn:Type[Function], axis:Optional[Union[int, Tuple[int, ...]]]=None, keepdim=False):
+  def _reduce(self, fxn:Type[Function], axis:Optional[Union[int, Tuple[int, ...]]]=None, keepdim=False) -> Tensor:
     axis_: List[int] = list(range(len(self.shape))) if axis is None else ([axis] if axis.__class__ is int else list(axis)) # type: ignore
     axis_ = [x if x >= 0 else x+len(self.shape) for x in axis_]
     shape = [s for i,s in enumerate(self.shape) if i not in axis_]
@@ -439,11 +439,11 @@ class Tensor:
   def mean(self, axis=None, keepdim=False):
     assert all_int(self.shape), "does not support symbolic shape"
     out = self.sum(axis=axis, keepdim=keepdim)
-    return out * (prod(out.shape)/prod(self.shape))
+    return out.mul(prod(out.shape)/prod(self.shape))
   def std(self, axis=None, keepdim=False, correction=1):
     assert all_int(self.shape), "does not support symbolic shape"
     square_sum = ((self - self.mean(axis=axis, keepdim=True)).square()).sum(axis=axis, keepdim=keepdim)
-    return (square_sum / (prod(self.shape)/prod(square_sum.shape)-correction)).sqrt()
+    return square_sum.div(prod(self.shape)/prod(square_sum.shape)-correction).sqrt()
   def _softmax(self, axis):
     m = self - self.max(axis=axis, keepdim=True)
     e = m.exp()