version bump

* improve factorization

* cleanups

.github/workflows/test.yml

@@ -36,7 +36,7 @@ jobs:
     - name: Lint with pylint
       run: python -m pylint --disable=all -e W0311 --jobs=0 --indent-string='  ' **/*.py
     - name: Lint with flake8
-      run: flake8 tinygrad/ --indent-size=2 --select=F,E112,E113,E304,E502,E702,E703,E71,E72,E731,W191,W6 --statistics -j4
+      run: flake8 tinygrad/ --indent-size=2 --select=F,E112,E113,E203,E304,E502,E702,E703,E71,E72,E731,W191,W6 --statistics -j4
     - name: Lint tinygrad with pylint
       run: pylint tinygrad/
     - name: Run mypy
@@ -72,14 +72,14 @@ jobs:
     - name: Install Dependencies
       run: pip install -e .
     - name: Compile EfficientNet to C
-      run: CLANG=1 python3 examples/compile_efficientnet.py > recognize.c
+      run: PYTHONPATH="." CLANG=1 python3 examples/compile_efficientnet.py > recognize.c
     - name: Compile C to native
       run: clang -O2 recognize.c -lm -o recognize
     - name: Test EfficientNet
       run: curl https://media.istockphoto.com/photos/hen-picture-id831791190 | ./recognize | grep hen
 
   testllvm:
-    name: LLVM Tests
+    name: LLVM Tests (w method cache)
     runs-on: ubuntu-latest
 
     steps:
@@ -130,11 +130,11 @@ jobs:
       with:
         python-version: 3.8
     - name: Install Dependencies
-      run: pip install -e '.[gpu,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+      run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
     - name: Run Optimizer Test (OPT 2 and 3)
       run: |
-        OPT=2 GPU=1 python test/external/external_test_opt.py
-        OPT=3 GPU=1 python test/external/external_test_opt.py
+        PYTHONPATH="." OPT=2 GPU=1 python test/external/external_test_opt.py
+        PYTHONPATH="." OPT=3 GPU=1 python test/external/external_test_opt.py
     - name: Run Pytest (default)
       run: GPU=1 python -m pytest -s -v -n=auto
 
@@ -158,12 +158,15 @@ jobs:
       with:
         python-version: 3.8
     - name: Install Dependencies
-      run: pip install -e '.[gpu,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+      run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
     - name: Test GPU IMAGE ops
-      run: GPU=1 IMAGE=2 python3 test/test_ops.py
+      run: |
+        GPU=1 IMAGE=1 python3 test/test_ops.py
+        FORWARD_ONLY=1 GPU=1 IMAGE=2 python3 test/test_ops.py
     - name: Test openpilot model
       run: |
-        ALLOWED_KERNEL_COUNT=197 FLOAT16=1 VALIDHACKS=1 DEBUGCL=1 GPU=1 IMAGE=2 python3 openpilot/compile.py
+        ALLOWED_KERNEL_COUNT=199 FLOAT16=1 VALIDHACKS=1 DEBUGCL=1 GPU=1 IMAGE=2 python3 openpilot/compile.py
+        python3 -c 'import os; assert os.path.getsize("/tmp/output.thneed") < 100_000_000'
         DEBUGCL=1 GPU=1 IMAGE=2 python3 openpilot/compile.py
         VALIDHACKS=1 DEBUGCL=1 GPU=1 IMAGE=2 python3 openpilot/compile.py
 

.gitignore

@@ -1,4 +1,6 @@
 __pycache__
+.venv/
+.vscode
 notebooks
 .*.swp
 .*.swo
@@ -16,3 +18,8 @@ vertex.bin
 recognize*
 .idea
 disassemblers/applegpu
+disassemblers/cuda_ioctl_sniffer
+*.prof
+datasets/cifar-10-python.tar.gz
+datasets/librispeech/
+datasets/imagenet/

.pre-commit-config.yaml

@@ -1,21 +1,33 @@
 repos:
   - repo: local
     hooks:
+      - id: docs
+        name: docs
+        entry: python3 docs/abstractions.py
+        language: system
+        always_run: true
+        pass_filenames: false
+      - id: flake8
+        name: flake8
+        entry: flake8 tinygrad/ --indent-size=2 --select=F,E112,E113,E203,E304,E502,E702,E703,E71,E72,E731,W191,W6 --statistics -j4
+        language: system
+        always_run: true
+        pass_filenames: false
+      - id: mypy
+        name: mypy
+        entry: mypy tinygrad/ extra/helpers.py --check-untyped-defs --explicit-package-bases --warn-unreachable # --warn-return-any
+        language: system
+        always_run: true
+        pass_filenames: false
+      - id: tests
+        name: subset of (CPU) tests
+        entry: env CPU=1 EXCLUDE_DEVICES=GPU pytest test/unit/ test/test_ops.py
+        language: system
+        always_run: true
+        pass_filenames: false
       - id: pylint
         name: pylint
         entry: pylint tinygrad/
         language: system
         always_run: true
         pass_filenames: false
-      - id: flake8
-        name: flake8
-        entry: flake8 tinygrad/ --indent-size=2 --select=F,E112,E113,E304,E502,E702,E703,E71,E72,E731,W191,W6 --statistics -j4
-        language: system
-        always_run: true
-        pass_filenames: false
-      - id: mypy
-        name: mypy
-        entry: mypy tinygrad/ --check-untyped-defs --explicit-package-bases --warn-unreachable # --warn-return-any
-        language: system
-        always_run: true
-        pass_filenames: false

.pylintrc

@@ -469,4 +469,4 @@ check-str-concat-over-line-jumps=yes
 
 # Exceptions that will emit a warning when being caught. Defaults to
 # "Exception"
-overgeneral-exceptions=Exception
+overgeneral-exceptions=builtins.Exception

LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2020 George Hotz
+Copyright (c) 2023 George Hotz
 
 Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 

README.md

@@ -23,6 +23,9 @@ This project is maintained by [tiny corp](https://tinygrad.org/).
 ### Installation
 
 ```bash
+python3 -m pip install git+https://[email protected]/geohot/tinygrad.git
+
+# or
 git clone https://github.com/geohot/tinygrad.git
 cd tinygrad
 python3 -m pip install -e .
@@ -71,7 +74,7 @@ print(y.grad)  # dz/dy
 Try a matmul. See how, despite the style, it is fused into one kernel with the power of laziness.
 
 ```python
-DEBUG=3 OPTLOCAL=1 GPU=1 python3 -c "from tinygrad.tensor import Tensor;
+DEBUG=3 OPTLOCAL=1 python3 -c "from tinygrad.tensor import Tensor;
 N = 1024; a, b = Tensor.randn(N, N), Tensor.randn(N, N);
 c = (a.reshape(N, 1, N) * b.permute(1,0).reshape(1, N, N)).sum(axis=2);
 print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
@@ -81,9 +84,9 @@ Change to `DEBUG=4` to see the generated code.
 
 ## Neural networks?
 
-It turns out, a decent autograd tensor library is 90% of what you need for neural networks. Add an optimizer (SGD, RMSprop, and Adam implemented) from tinygrad.nn.optim, write some boilerplate minibatching code, and you have all you need.
+It turns out, a decent autograd tensor library is 90% of what you need for neural networks. Add an optimizer (SGD, Adam, AdamW implemented) from tinygrad.nn.optim, write some boilerplate minibatching code, and you have all you need.
 
-### Neural network example (from test/test_mnist.py)
+### Neural network example (from test/models/test_mnist.py)
 
 ```python
 from tinygrad.tensor import Tensor
@@ -145,7 +148,7 @@ hlops are syntactic sugar around mlops. They support most things torch does.
 mlops are mid level ops. They understand derivatives. They are very simple.
 
 ```
-Log, Exp                                       # unary ops
+Relu, Log, Exp                                 # unary ops
 Sum, Max                                       # reduce ops (with axis argument)
 Maximum, Add, Sub, Mul, Pow, Div, Equal        # binary ops (no broadcasting, use expand)
 Expand, Reshape, Permute, Pad, Shrink, Flip    # movement ops
@@ -158,12 +161,12 @@ You no longer need to write mlops for a new accelerator
 The autodiff stuff is all in mlops now so you can focus on the raw operations
 
 ```
-Buffer                                                     # class of memory on this device
-unary_op  (NOOP, NEG, NOT, EXP, LOG)                       # A -> A
-reduce_op (SUM, MAX)                                       # A -> B (smaller size, B has 1 in shape)
-binary_op (ADD, SUB, MUL, DIV, POW, CMPEQ, MAX)            # A + A -> A (all the same size)
-movement_op (EXPAND, RESHAPE, PERMUTE, PAD, SHRINK, FLIP)  # A -> B (different size)
-fused_op [[optional]] (MULACC)                             # A * A -> B
+Buffer                                                       # class of memory on this device
+unary_op  (NOOP, EXP, LOG, CAST)                             # A -> A
+reduce_op (SUM, MAX)                                         # A -> B (smaller size, B has 1 in shape)
+binary_op (ADD, SUB, MUL, DIV, POW, CMPEQ, MAX)              # A + A -> A (all the same size)
+movement_op (EXPAND, RESHAPE, PERMUTE, PAD, SHRINK, STRIDE)  # A -> B (different size)
+fused_op [[optional]] (MULACC)                               # A * A -> B
 ```
 
 ## ImageNet inference
@@ -180,15 +183,13 @@ Or, if you have a webcam and cv2 installed
 ipython3 examples/efficientnet.py webcam
 ```
 
-PROTIP: Set "GPU=1" environment variable if you want this to go faster.
-
-PROPROTIP: Set "DEBUG=1" environment variable if you want to see why it's slow.
+PROTIP: Set "DEBUG=1" environment variable if you want to see why it's slow.
 
 ### tinygrad supports Stable Diffusion!
 
 You might need to download the [weight](https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt) of Stable Diffusion and put it into weights/
 
-Run `GPU=1 python3 examples/stable_diffusion.py`
+Run `python3 examples/stable_diffusion.py`
 
 <p align="center">
   <img src="https://raw.githubusercontent.com/geohot/tinygrad/master/docs/stable_diffusion_by_tinygrad.jpg">
@@ -227,7 +228,7 @@ tinygrad will always be below 1000 lines. If it isn't, we will revert commits un
 * Purple edge is intermediates created in the forward
 
 ```bash
-GRAPH=1 python3 test/test_mnist.py TestMNIST.test_sgd_onestep
+GRAPH=1 python3 test/models/test_mnist.py TestMNIST.test_sgd_onestep
 # requires dot, outputs /tmp/net.svg
 ```
 

accel/triton/ops_triton.py

@@ -11,7 +11,7 @@ import triton.language as tl  # type: ignore # noqa: F401
 
 from typing import Union, Tuple, Optional, Dict
 from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, LazyOp, Op, ExplicitExecAST, GlobalCounters
-from tinygrad.shape import ShapeTracker
+from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.helpers import prod, DEBUG
 from tinygrad.runtime.cuda import CLBuffer
 from tinygrad.compiler.ast import ASTKernel
@@ -57,7 +57,7 @@ class TritonASTKernel(ASTKernel):
     self.kernel = ["@triton.jit"]
     self.kernel.append("def fxn("+','.join(f"data{i}" for i in range(len(self.bufs)))+"):")
 
-    self.output_shape = list(self.sts[0].shape[:self.first_reduce]) 
+    self.output_shape = list(self.sts[0].shape[:self.first_reduce])
 
     # copied from ops_gpu
     # TODO CUDA only supports a grid of (2^31-1, 65535, 65535), that results in invalid kernel launches for some shapes, so flattern the grid for now.
@@ -70,7 +70,7 @@ class TritonASTKernel(ASTKernel):
       self.output_shape = [prod(self.output_shape[0:final_dimension+1])] + list(self.output_shape[final_dimension+1:])
       if DEBUG >= 3: print(f"replaced output shape with {self.output_shape}")
     elif len(self.output_shape) == 0: self.output_shape = [1]
-    
+
     if self.reduceop:
       full_shape = [st.shape for st in self.sts if st.shape != self.sts[0].shape]
       full_shape = self.sts[0].shape if len(full_shape) == 0 else full_shape[0]

compile.sh

@@ -1,6 +1,8 @@
 #!/bin/bash
-mypyc --explicit-package-bases \
-  tinygrad/llops/ops_gpu.py tinygrad/shape/__init__.py tinygrad/ops.py tinygrad/ast.py \
-  tinygrad/helpers.py tinygrad/mlops.py tinygrad/nn/__init__.py tinygrad/graph.py tinygrad/lazy.py \
-  tinygrad/tensor.py tinygrad/llops/ops_cpu.py tinygrad/llops/ops_torch.py tinygrad/nn/optim.py
-
+# note: if we compile tinygrad/nn/__init__.py __dict__ no longer works, and optimizers will silently fail
+mypyc --check-untyped-defs --explicit-package-bases --warn-unreachable tinygrad/shape/shapetracker.py tinygrad/shape/symbolic.py \
+  tinygrad/helpers.py tinygrad/mlops.py tinygrad/tensor.py tinygrad/graph.py \
+  #tinygrad/codegen/gpu.py tinygrad/runtime/ops_metal.py
+  #tinygrad/codegen/ast.py
+  #tinygrad/nn/__init__.py
+  #tinygrad/ops.py tinygrad/runtime/ops_metal.py tinygrad/runtime/ops_gpu.py tinygrad/runtime/ops_cpu.py tinygrad/lazy.py

datasets/__init__.py

@@ -1,10 +1,9 @@
 import os
-import io
 import numpy as np
 import gzip
 import tarfile
 import pickle
-from extra.utils import fetch
+from extra.utils import download_file
 
 def fetch_mnist():
   parse = lambda file: np.frombuffer(gzip.open(file).read(), dtype=np.uint8).copy()
@@ -17,9 +16,10 @@ def fetch_mnist():
 def fetch_cifar(train=True):
   cifar10_mean = np.array([0.4913997551666284, 0.48215855929893703, 0.4465309133731618], dtype=np.float32).reshape(1,3,1,1)
   cifar10_std = np.array([0.24703225141799082, 0.24348516474564, 0.26158783926049628], dtype=np.float32).reshape(1,3,1,1)
-  tt = tarfile.open(fileobj=io.BytesIO(fetch('https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz')), mode='r:gz')
+  fn = os.path.dirname(__file__)+"/cifar-10-python.tar.gz"
+  download_file('https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz', fn)
+  tt = tarfile.open(fn, mode='r:gz')
   if train:
-    # TODO: data_batch 2-5
     db = [pickle.load(tt.extractfile(f'cifar-10-batches-py/data_batch_{i}'), encoding="bytes") for i in range(1,6)]
   else:
     db = [pickle.load(tt.extractfile('cifar-10-batches-py/test_batch'), encoding="bytes")]

datasets/imagenet.py

@@ -1,28 +1,49 @@
-import os
+# for imagenet download prepare.sh and run it
+import glob, random
 import json
 import numpy as np
 from PIL import Image
-import torchvision.transforms as transforms
+import functools, pathlib
 
-BASEDIR = "/Users/kafka/fun/imagenet"
-train_files = open(os.path.join(BASEDIR, "train_files")).read().strip().split("\n")
-val_files = open(os.path.join(BASEDIR, "val_files")).read().strip().split("\n")
-ci = json.load(open(os.path.join(BASEDIR, "imagenet_class_index.json")))
+BASEDIR = pathlib.Path(__file__).parent.parent / "datasets/imagenet"
+ci = json.load(open(BASEDIR / "imagenet_class_index.json"))
 cir = {v[0]: int(k) for k,v in ci.items()}
 
-rrc = transforms.RandomResizedCrop(224)
+@functools.lru_cache(None)
+def get_train_files():
+  train_files = open(BASEDIR / "train_files").read().strip().split("\n")
+  return [(BASEDIR / "train" / x) for x in train_files]
+
+@functools.lru_cache(None)
+def get_val_files():
+  val_files = glob.glob(str(BASEDIR / "val/*/*"))
+  return val_files
+
+#rrc = transforms.RandomResizedCrop(224)
+import torchvision.transforms.functional as F
 def image_load(fn):
   img = Image.open(fn).convert('RGB')
-  ret = np.array(rrc(img))
+  img = F.resize(img, 256, Image.BILINEAR)
+  img = F.center_crop(img, 224)
+  ret = np.array(img)
   return ret
 
+def iterate(bs=32, val=True, shuffle=True):
+  files = get_val_files() if val else get_train_files()
+  order = list(range(0, len(files)))
+  if shuffle: random.shuffle(order)
+  for i in range(0, len(files), bs):
+    X = [image_load(files[i]) for i in order[i:i+bs]]
+    Y = [cir[files[i].split("/")[-2]] for i in order[i:i+bs]]
+    yield (np.array(X), np.array(Y))
+
 def fetch_batch(bs, val=False):
-  files = val_files if val else train_files
+  files = get_val_files() if val else get_train_files()
   samp = np.random.randint(0, len(files), size=(bs))
   files = [files[i] for i in samp]
-  X = [image_load(os.path.join(BASEDIR, "val" if val else "train", x)) for x in files]
+  X = [image_load(x) for x in files]
   Y = [cir[x.split("/")[0]] for x in files]
-  return np.transpose(np.array(X), (0,3,1,2)), np.array(Y)
+  return np.array(X), np.array(Y)
 
 if __name__ == "__main__":
   X,Y = fetch_batch(64)

datasets/librispeech.py

@@ -0,0 +1,82 @@
+import json
+import pathlib
+import numpy as np
+import librosa
+import soundfile
+
+"""
+The dataset has to be downloaded manually from https://www.openslr.org/12/ and put in `datasets/librispeech`.
+For mlperf validation the dev-clean dataset is used.
+
+Then all the flacs have to be converted to wav using something like:
+```fish
+for file in $(find * | grep flac); do ffmpeg -i $file -ar 16k "$(dirname $file)/$(basename $file .flac).wav"; done
+```
+
+Then this [file](https://github.com/mlcommons/inference/blob/master/speech_recognition/rnnt/dev-clean-wav.json) has to also be put in `datasets/librispeech`.
+"""
+BASEDIR = pathlib.Path(__file__).parent.parent / "datasets/librispeech"
+with open(BASEDIR / "dev-clean-wav.json") as f:
+  ci = json.load(f)
+
+FILTER_BANK = np.expand_dims(librosa.filters.mel(sr=16000, n_fft=512, n_mels=80, fmin=0, fmax=8000), 0)
+WINDOW = librosa.filters.get_window("hann", 320)
+
+def feature_extract(x, x_lens):
+  x_lens = np.ceil((x_lens / 160) / 3).astype(np.int32)
+
+  # pre-emphasis
+  x = np.concatenate((np.expand_dims(x[:, 0], 1), x[:, 1:] - 0.97 * x[:, :-1]), axis=1)
+
+  # stft
+  x = librosa.stft(x, n_fft=512, window=WINDOW, hop_length=160, win_length=320, center=True, pad_mode="reflect")
+  x = np.stack((x.real, x.imag), axis=-1)
+
+  # power spectrum
+  x = (x**2).sum(-1)
+
+  # mel filter bank
+  x = np.matmul(FILTER_BANK, x)
+
+  # log
+  x = np.log(x + 1e-20)
+
+  # feature splice
+  seq = [x]
+  for i in range(1, 3):
+    tmp = np.zeros_like(x)
+    tmp[:, :, :-i] = x[:, :, i:]
+    seq.append(tmp)
+  features = np.concatenate(seq, axis=1)[:, :, ::3]
+
+  # normalize
+  features_mean = np.zeros((features.shape[0], features.shape[1]), dtype=np.float32)
+  features_std = np.zeros((features.shape[0], features.shape[1]), dtype=np.float32)
+  for i in range(features.shape[0]):
+    features_mean[i, :] = features[i, :, :x_lens[i]].mean(axis=1)
+    features_std[i, :] = features[i, :, :x_lens[i]].std(axis=1, ddof=1)
+  features_std += 1e-5
+  features = (features - np.expand_dims(features_mean, 2)) / np.expand_dims(features_std, 2)
+
+  return features.transpose(2, 0, 1), x_lens.astype(np.float32)
+
+def load_wav(file):
+  sample = soundfile.read(file)[0].astype(np.float32)
+  return sample, sample.shape[0]
+
+def iterate(bs=1, start=0):
+  print(f"there are {len(ci)} samples in the dataset")
+  for i in range(start, len(ci), bs):
+    samples, sample_lens = zip(*[load_wav(BASEDIR / v["files"][0]["fname"]) for v in ci[i : i + bs]])
+    samples = list(samples)
+    # pad to same length
+    max_len = max(sample_lens)
+    for j in range(len(samples)):
+      samples[j] = np.pad(samples[j], (0, max_len - sample_lens[j]), "constant")
+    samples, sample_lens = np.array(samples), np.array(sample_lens)
+
+    yield feature_extract(samples, sample_lens), np.array([v["transcript"] for v in ci[i : i + bs]])
+
+if __name__ == "__main__":
+  X, Y = next(iterate())
+  print(X[0].shape, Y.shape)

disassemblers/adreno/__init__.py

@@ -10,7 +10,7 @@ def disasm(buf):
     shared = pathlib.Path(__file__).parent / "disasm.so"
     if not os.path.isfile(shared):
       os.system(f'cd {pathlib.Path(__file__).parent} && gcc -shared disasm-a3xx.c -o disasm.so')
-    fxn = ctypes.CDLL(shared)['disasm']
+    fxn = ctypes.CDLL(shared.as_posix())['disasm']
   #hexdump(buf)
   END = b"\x00\x00\x00\x00\x00\x00\x00\x03"
   buf = buf[0x510:]  # this right?

docs/abstractions.py

@@ -0,0 +1,372 @@
+"""
+Welcome to the tinygrad documentation
+=================
+
+this file will take you on a whirlwind journey from a Tensor all the way down
+tinygrad has been aggressively refactored in the 2.5 years it's been worked on.
+what you see here is a refined library (with more refining to go still!)
+
+the whole tinygrad is ~2300 lines, so while it's readable in an evening or two,
+this documentation will help with entry points and understanding the abstraction stack
+"""
+
+# %%
+# == Boilerplate imports for typing ==
+from __future__ import annotations
+from typing import Optional, Tuple, Union, Any, Dict, Callable, Type, List, ClassVar
+from enum import Enum, auto
+from abc import ABC
+
+# %%
+# == Example: Tensor 2+3 ==
+# let's trace an addition down through the layers of abstraction.
+
+# we will be using the clang backend
+from tinygrad.lazy import Device
+Device.DEFAULT = "CLANG"
+
+# first, 2+3 as a Tensor, the highest level
+from tinygrad.tensor import Tensor
+a = Tensor([2])
+b = Tensor([3])
+result = a + b
+print(f"{a.numpy()} + {b.numpy()} = {result.numpy()}")
+assert result.numpy()[0] == 5.
+
+# %%
+# == Tensor (in tinygrad/tensor.py, code 8/10) ==
+# it's worth reading tinygrad/tensor.py. it's pretty beautiful
+import tinygrad.mlops as mlops
+
+# this is the good old familiar Tensor class
+class Tensor:
+  # these two are pretty straightforward
+  grad: Optional[Tensor]
+  requires_grad: Optional[bool]
+
+  # this is the graph for the autograd engine
+  _ctx: Optional[Function]
+
+  # this is where the data (and other tensor properties) actually live
+  lazydata: LazyBuffer
+
+  # high level ops (hlops) are defined on this class. example: relu
+  def relu(self): return self.maximum(0)
+
+  # log is an mlop, this is the wrapper function in Tensor
+  def log(self): return mlops.Log.apply(self)
+
+# all the definitions of the derivatives are subclasses of Function (like mlops.Log)
+# there's only 18 mlops for derivatives for everything (in tinygrad/mlops.py, code 9/10)
+# if you read one file, read mlops.py. if you read two files, also read tinygrad/tensor.py
+# you can differentiate the world using the chain rule
+class Function:
+  # example types of forward and backward
+  def forward(self, x:LazyBuffer) -> LazyBuffer: pass
+  def backward(self, x:LazyBuffer) -> LazyBuffer: pass
+
+# %%
+# == LazyBuffer (in tinygrad/lazy.py, code 5/10) ==
+from tinygrad.helpers import DType
+
+# this is where the properties live that you thought were a part of Tensor
+# LazyBuffer is like a Tensor without derivatives, at the mlop layer
+class LazyBuffer:
+  # these three define the "type" of the buffer, and they are returned as Tensor properties
+  device: str
+  shape: Tuple[int, ...]
+  dtype: DType
+
+  # a ShapeTracker is used to track things like reshapes and permutes
+  # all MovementOps are zero copy in tinygrad!
+  # the ShapeTracker specifies how the data in the RawBuffer matches to the shape
+  # we'll come back to this later
+  st: ShapeTracker
+
+  # if the LazyBuffer is realized, it has a RawBuffer
+  # we will come back to RawBuffers later
+  realized: Optional[RawBuffer]
+
+  # if the lazybuffer is unrealized, it has a LazyOp
+  # this LazyOp describes the computation needed to realize this LazyBuffer
+  op: Optional[LazyOp]
+
+# LazyOp (in tinygrad/ops.py, code 4/10)
+# in a tree they form an Abstract Syntax Tree for a single GPU kernel
+class LazyOp:
+  op: Op                                       # the type of the compute
+  src: Tuple[Union[LazyOp, LazyBuffer], ...]   # the sources
+  arg: Optional[Any] = None                    # and an optional static argument
+
+# there's currently 20 Ops you have to implement for an accelerator.
+class UnaryOps(Enum):    NOOP = auto(); EXP = auto(); LOG = auto(); NEG = auto(); NOT = auto()
+class BinaryOps(Enum):   ADD = auto();  SUB = auto(); MUL = auto(); DIV = auto(); POW = auto(); CMPEQ = auto(); MAX = auto()
+class ReduceOps(Enum):   SUM = auto();  MAX = auto()
+class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); STRIDE = auto()
+class LoadOps(Enum):     FROMCPU = auto()
+# NOTE: if you have a CompiledBuffer(DeviceBuffer)
+#       you do not need to implement the MovementOps
+#       as they are handled by the ShapeTracker(in tinygrad/shape/shapetracker.py, code 7/10)
+Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps]
+
+# most of tinygrad/lazy.py is concerned with fusing Ops into LazyOps ASTs that map to GPUKernels
+# it's beyond the scope of this tutorial, but you can read the file if interested
+
+# %%
+# == Example: LazyBuffer for 2+3 ==
+
+from tinygrad.tensor import Tensor
+from tinygrad.ops import LazyOp, BinaryOps, LoadOps
+
+# the 2+3 from before
+result = Tensor([2]) + Tensor([3])
+print(type(result.lazydata), result.lazydata)  # let's look at the lazydata of result
+
+# you'll see it has a LazyOp
+# the op type is BinaryOps.ADD
+# and it has two sources, the 2 and the 3
+lazyop: LazyOp = result.lazydata.op
+assert lazyop.op == BinaryOps.ADD
+assert len(lazyop.src) == 2
+
+# the first source is the 2, it comes from the CPU
+# the source is a LazyBuffer, since FROMCPU cannot be folded into LazyOp ASTs
+# again, a LazyOp AST is like a GPU kernel. you have to copy the data on the device first
+print(lazyop.src[0].op)
+assert lazyop.src[0].op.op == LoadOps.FROMCPU
+assert lazyop.src[0].op.arg.fxn == [2], "the arg of the FROMCPU LazyOP is the [2.]"
+assert result.lazydata.realized is None, "the LazyBuffer is not realized yet"
+
+# now we realize the LazyBuffer
+result.lazydata.realize()
+assert result.lazydata.realized is not None, "the LazyBuffer is realized!"
+# this brings us nicely to DeviceBuffer, of which the realized ClangBuffer is a subclass
+assert 'RawMallocBuffer' in str(type(result.lazydata.realized))
+# getting ahead of ourselves, but we can copy the DeviceBuffer toCPU
+assert result.lazydata.realized.toCPU()[0] == 5, "when put in numpy with toCPU, it's 5"
+
+# %%
+# == Union[Interpreted, Compiled] (in tinygrad/ops.py, code 5/10) ==
+
+# Now you have a choice, you can either write a "Interpreted" backend or "Compiled" backend
+
+# Interpreted backends are very simple (example: CPU and TORCH)
+class Interpreted:
+  # they have a backing RawBuffer
+  buffer: Type[RawBuffer]
+
+  # and they have a lookup table to functions for the Ops
+  fxn_for_op: Dict[Op, Callable] = {
+    UnaryOps.EXP: lambda x: np.exp(x),
+    BinaryOps.ADD: lambda x,y: x+y}
+
+# Compiled backends take a little more (example: GPU and LLVM)
+class Compiled:
+  # they also have a backing RawBuffer
+  buffer: Type[RawBuffer]
+
+  # a code generator, which compiles the AST
+  codegen: Type[ASTKernel]
+
+  # and a runtime, which runs the generated code
+  runtime: Type[Runtime]
+
+# Runtime is what actually runs the kernels for a compiled backend
+class Runtime(ABC):
+  # `name` is the name of the function, and `prg` is the code
+  # the constructor compiles the code
+  def __init__(self, name:str, prg:str): pass
+  # call runs the code on the bufs. NOTE: the output is always bufs[0], but this is just a convention
+  def __call__(self, global_size:Optional[List[int]], local_size:Optional[List[int]], *bufs:List[RawBuffer]): pass
+
+# %%
+# == RawBuffer (in tinygrad/runtime/lib.py, code 5/10) ==
+import numpy as np
+
+# RawBuffer is where the data is actualy held. it's pretty close to just memory
+class RawBuffer(ABC):
+  # create an empty rawbuffer that holds `size` elements of type `dtype`
+  # `buf` is an opaque container class
+  def __init__(self, size:int, dtype:DType, buf:Any): raise NotImplementedError("must be implemented")
+
+  # fromCPU is classmethod that creates a RawBuffer, it's a classmethod since some runtimes are 0 copy
+  @classmethod
+  def fromCPU(cls:RawBuffer, x:np.ndarray) -> RawBuffer: raise NotImplementedError("must be implemented")
+
+  # toCPU converts the RawBuffer to a numpy array with shape (size,). many backends are 0 copy here
+  def toCPU(self) -> np.ndarray: raise NotImplementedError("must be implemented")
+
+# RawNumpyBuffer is a RawBuffer example for numpy. It's very simple
+class RawNumpyBuffer(RawBuffer):
+  # NOTE: the "np.ndarray" is stored in the opaque container
+  def __init__(self, buf:np.ndarray):
+    super().__init__(buf.size, dtypes.from_np(buf.dtype), buf)
+  @classmethod
+  def fromCPU(cls, x): return cls(x)
+  def toCPU(self): return self._buf
+
+# %%
+# == Example: 2+3 in raw clang ==
+
+# RawMallocBuffer is the simplest concrete version of RawBuffer (in tinygrad/ops.py)
+# it's used for the CLANG and LLVM backends
+# it's just malloc(size * dtype.itemsize)
+from tinygrad.runtime.lib import RawMallocBuffer
+
+# ClangProgram is the simplest runtime (in tinygrad/runtime/ops_clang.py, code 7/10)
+# __init__ calls clang, and __call__ calls the function in the *.so outputted by clang
+# in CLANG, global_size and local_size are ignored
+from tinygrad.runtime.ops_clang import ClangProgram
+
+# a concrete example looks like this, this adds two size 1 RawBuffer
+# first we create two numpy buffers containing 2 and 3
+# then we copy the numpy in to RawMallocBuffers
+# last, we create an empty output buffer
+from tinygrad.helpers import dtypes
+numpy_a, numpy_b = np.array([2], dtype=np.float32), np.array([3], dtype=np.float32)
+input_a, input_b = RawMallocBuffer.fromCPU(numpy_a), RawMallocBuffer.fromCPU(numpy_b)
+output = RawMallocBuffer(1, dtypes.float32)
+
+# compile the program, run it, and 2+3 does indeed equal 5
+program = ClangProgram("add", "void add(float *a, float *b, float *c) { *a = *b + *c; }")
+program(None, None, output, input_a, input_b)  # NOTE: the None are for global_size and local_size
+print(output.toCPU())
+assert output.toCPU()[0] == 5, "it's still 5"
+np.testing.assert_allclose(output.toCPU(), numpy_a+numpy_b)
+
+# %%
+# == ASTKernel (in tinygrad/codegen/ast.py, code 2/10) ==
+
+# but we are nowhere near done!
+# we wrote the code above by hand
+# we need the LazyOp ASTs to be automatically turned into code
+# the current class looks roughly like this, but this will change and we will update the docs
+# this stuff is in the terrible 528 lines of (tinygrad/codegen/*, code 2/10 aka turd quality)
+class ASTKernel:
+  # create the kernel with the AST
+  # NOTE: the AST contains the CompiledBuffers themselves as the root nodes. this will change
+  def __init__(self, ast:LazyOp): pass
+  def codegen(self) -> ASTRunner: pass
+
+# we return a class that runs code on LazyBuffers, which are all expected to be realized
+class ASTRunner:  # (from tinygrad/ops.py)
+  def __init__(self, name, prg, global_size:Optional[List[int]], local_size:Optional[List[int]]): pass
+  def build(self, runtime:Runtime): pass
+  def exec(self, bufs:List[LazyBuffer]): pass
+
+# that hides a lot of complexity that will be refactored, but that's the basic idea of code generation
+
+# %%
+# == Example: 2+3 autogenerated clang code ==
+
+from tinygrad.tensor import Tensor
+result = Tensor([2]) + Tensor([3])
+
+# we have a global cache used by the JIT
+# from there, we can see the generated clang code
+from tinygrad.helpers import GlobalCounters
+GlobalCounters.cache = []    # enables the cache
+result.realize()             # create the program and runs it
+cache_saved = GlobalCounters.cache
+GlobalCounters.cache = None  # disable the cache
+
+# there's one ASTRunner in the cache
+assert len(cache_saved) == 1
+prg, bufs = cache_saved[0]
+
+# print the C Program :)
+print(prg.prg)
+
+# after some formatting (the compiler doesn't care)
+# NOTE: the 2 and 3 are constant folded
+"""
+void E_1(float* data0) {
+  for (int idx0 = 0; idx0 < 1; idx0++) {
+    data0[0] = (2.0f) + (3.0f);
+  }
+}
+"""
+
+# %%
+# == Example: ShapeTracker (in tinygrad/shape/shapetracker.py, code 7/10) ==
+
+# remember how I said you don't have to write the MovementOps for CompiledBuffers?
+# that's all thanks to ShapeTracker!
+# ShapeTracker tracks the indices into the RawBuffer
+from tinygrad.shape.shapetracker import ShapeTracker
+
+# create a virtual (10, 10) Tensor. this is just a shape, there's no actual tensor
+a = ShapeTracker((10, 10))
+
+# you'll see it has one view. the (10, 1 are the strides)
+print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (10, 1), 0)])
+
+# we can permute it, and the strides change
+a.permute((1,0))
+print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (1, 10), 0)])
+
+# we can then reshape it, and the strides change again
+# note how the permute stays applied
+a.reshape((5,2,5,2))
+print(a) # ShapeTracker(shape=(5, 2, 5, 2), views=[View((5, 2, 5, 2), (2, 1, 20, 10), 0)])
+
+# now, if we were to reshape it to a (100,) shape tensor, we have to create a second view
+a.reshape((100,))
+print(a) # ShapeTracker(shape=(100,), views=[
+         #   View((5, 2, 5, 2), (2, 1, 20, 10), 0),
+         #   View((100,), (1,), 0)])
+
+# Views stack on top of each other, to allow zero copy for any number of MovementOps
+# we can render a Python expression for the index at any time
+idx, _ = a.expr_idxs()
+print(idx.render())  # (((idx0%10)*10)+(idx0//10))
+
+# of course, if we reshape it back, the indexes get simple again
+a.reshape((10,10))
+idx, _ = a.expr_idxs()
+print(idx.render())  # ((idx1*10)+idx0)
+
+# the ShapeTracker still has two views though...
+print(a) # ShapeTracker(shape=(10, 10), views=[
+         #   View((5, 2, 5, 2), (2, 1, 20, 10), 0),
+         #   View((10, 10), (10, 1), 0)])
+
+# ...until we simplify it!
+a.simplify()
+print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (1, 10), 0)])
+
+# and now we permute it back
+a.permute((1,0))
+print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (10, 1), 0)])
+
+# and it's even contiguous
+assert a.contiguous == True
+
+# %%
+# == Example: Variable (in tinygrad/shape/symbolic.py, code 6/10) ==
+
+# Under the hood, ShapeTracker is powered by a small symbolic algebra library
+from tinygrad.shape.symbolic import Variable
+
+# Variable is the basic class from symbolic
+# it's created with a name and a min and max (inclusive)
+a = Variable("a", 0, 10)
+b = Variable("b", 0, 10)
+
+# some math examples
+print((a*10).min, (a*10).max)  # you'll see a*10 has a min of 0 and max of 100
+print((a+b).min, (a+b).max)    # 0 20, you get the idea
+
+# but complex expressions are where it gets fun
+expr = (a + b*10) % 10
+print(expr.render())   # (a%10)
+# as you can see, b is gone!
+
+# one more
+expr = (a*40 + b) // 20
+print(expr.render())       # (a*2)
+print(expr.min, expr.max)  # 0 20
+# this is just "(a*2)"
+# since b only has a range from 0-10, it can't affect the output
+
+# %%

examples/compile_efficientnet.py

@@ -41,12 +41,12 @@ if __name__ == "__main__":
   the_output = run(the_input)
 
   # hack to put the inputs back
-  assert len(run.input_replace) == 1
+  assert len(run.input_replace) == 1, f"didn't get one input to replace {run.input_replace}"
   for (j,i),idx in run.input_replace.items():
-    run.jit_cache[j][1][i] = the_input.lazydata.realized.raw()
+    run.jit_cache[j][1][i] = the_input.lazydata.realized
 
   # TODO: fetch this from the jit in self.input_replace and self.ret (hint: use get_parameters on self.ret)
-  special_names = {id(the_input.lazydata.realized.raw()): "input", id(the_output.lazydata.realized.raw()): "outputs"}
+  special_names = {id(the_input.lazydata.realized): "input", id(the_output.lazydata.realized): "outputs"}
 
   functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
 
@@ -68,7 +68,7 @@ if __name__ == "__main__":
   cprog.append(f"char *lbls[] = {{{','.join(lbls)}}};")
 
   # buffers (empty + weights)
-  cprog += [f"float {name}[{len}];" if name not in bufs_to_save else f"float *{name} = (float *){name}_data;" for name,len in bufs.values()] 
+  cprog += [f"float {name}[{len}];" if name not in bufs_to_save else f"float *{name} = (float *){name}_data;" for name,len in bufs.values()]
 
   # the functions
   cprog += list(functions.values())
@@ -109,4 +109,5 @@ int main(int argc, char* argv[]) {
 }"""]
 
   # CLANG=1 python3 examples/compile_efficientnet.py | clang -O2 -lm -x c - -o recognize && DEBUG=1 time ./recognize docs/stable_diffusion_by_tinygrad.jpg
+  # category : 281 (tabby, tabby cat) with 9.452788
   print('\n'.join(cprog))

examples/deep_deterministic_policy_gradient.py

@@ -173,12 +173,12 @@ class DeepDeterministicPolicyGradient:
     for param, target_param in zip(
         optim.get_parameters(self.actor), optim.get_parameters(self.target_actor)
     ):
-      target_param.assign(param * tau + target_param * (1.0 - tau))
+      target_param.assign(param.detach() * tau + target_param * (1.0 - tau))
 
     for param, target_param in zip(
         optim.get_parameters(self.critic), optim.get_parameters(self.target_critic)
     ):
-      target_param.assign(param * tau + target_param * (1.0 - tau))
+      target_param.assign(param.detach() * tau + target_param * (1.0 - tau))
 
   def choose_action(self, state: Tensor, evaluate: bool = False) -> NDArray:
     mu = self.actor.forward(state, self.max_action)

examples/efficientnet.py

@@ -45,7 +45,7 @@ def infer(model, img):
   """
 
   # run the net
-  out = _infer(model, Tensor(img)).numpy()
+  out = _infer(model, Tensor(img.astype("float32"))).numpy()
 
   # if you want to look at the outputs
   """

examples/hlb_cifar10.py

@@ -58,8 +58,6 @@ def train_step_jitted(model, optimizer, X, Y):
     optimizer.zero_grad()
     loss.backward()
     optimizer.step()
-  #optimizer.lr *= 0.995  # decay LR
-  #optimizer.lr.realize()
   return loss.realize()
 
 def fetch_batch(X_train, Y_train, BS):
@@ -73,23 +71,33 @@ def fetch_batch(X_train, Y_train, BS):
 
 def train_cifar():
   Tensor.training = True
-  BS = getenv("BS", 512)
+  BS, STEPS = getenv("BS", 512), getenv("STEPS", 10)
   if getenv("FAKEDATA"):
     N = 2048
     X_train = np.random.default_rng().standard_normal(size=(N, 3, 32, 32), dtype=np.float32)
     Y_train = np.random.randint(0,10,size=(N), dtype=np.int32)
     X_test, Y_test = X_train, Y_train
   else:
-    X_train,Y_train = fetch_cifar(train=True)
-    X_test,Y_test = fetch_cifar(train=False)
+    X_train, Y_train = fetch_cifar(train=True)
+    X_test, Y_test = fetch_cifar(train=False)
   print(X_train.shape, Y_train.shape)
   Xt, Yt = fetch_batch(X_test, Y_test, BS=BS)
   model = SpeedyResNet()
+
+  # init weights with torch
+  if getenv("TORCHWEIGHTS"):
+    from examples.hlb_cifar10_torch import SpeedyResNet as SpeedyResNetTorch
+    torch_model = SpeedyResNetTorch()
+    model_state_dict = optim.get_state_dict(model)
+    torch_state_dict = torch_model.state_dict()
+    for k,v in torch_state_dict.items():
+      print(f"initting {k} from torch")
+      model_state_dict[k].assign(Tensor(v.detach().numpy())).realize()
+
   if getenv("ADAM"):
     optimizer = optim.Adam(optim.get_parameters(model), lr=Tensor([0.001]).realize())
   else:
-    #optimizer = optim.SGD(optim.get_parameters(model), lr=0.001)
-    optimizer = optim.SGD(optim.get_parameters(model), lr=Tensor([0.003]).realize(), momentum=0.85, nesterov=True)
+    optimizer = optim.SGD(optim.get_parameters(model), lr=0.01, momentum=0.85, nesterov=True)
 
   # 97 steps in 2 seconds = 20ms / step
   # step is 1163.42 GOPS = 56 TFLOPS!!!, 41% of max 136
@@ -101,15 +109,15 @@ def train_cifar():
   # 136 TFLOPS is the theoretical max w float16 on 3080 Ti
 
   X, Y = fetch_batch(X_train, Y_train, BS=BS)
-  for i in range(getenv("STEPS", 10)):
-    #new_lr = (0.003 * i/300) if i < 300 else min(0.00001, 0.003 - 0.003 * i/300)
-    #optimizer.lr = Tensor([new_lr]).realize()
-
-    if i%10 == 0:
+  for i in range(max(1, STEPS)):
+    if i%10 == 0 and STEPS != 1:
       # use training batchnorm (and no_grad would change the kernels)
-      outs = model(Xt).numpy().argmax(axis=1)
+      out = model(Xt)
+      outs = out.numpy().argmax(axis=1)
+      loss = (out * Yt).mean().numpy()[0]
       correct = outs == Yt.numpy().argmin(axis=1)
-      print(f"eval {sum(correct)}/{len(correct)} {sum(correct)/len(correct)*100.0:.2f}%")
+      print(f"eval {sum(correct)}/{len(correct)} {sum(correct)/len(correct)*100.0:.2f}%, {loss:7.2f} val_loss")
+    if STEPS == 0: break
     GlobalCounters.reset()
     st = time.monotonic()
     loss = train_step_jitted(model, optimizer, X, Y)

examples/hlb_cifar10_torch.py

@@ -1,12 +1,16 @@
 import numpy as np
 import torch
 import time
+import platform
 from torch import nn
 from torch import optim
 
 from datasets import fetch_cifar
 from tinygrad.helpers import getenv
 
+OSX = platform.system() == "Darwin"
+device = 'mps' if OSX else 'cuda'
+
 num_classes = 10
 class ConvGroup(nn.Module):
   def __init__(self, channels_in, channels_out, short, se=True):
@@ -26,28 +30,28 @@ class ConvGroup(nn.Module):
     x = self.norm[2](self.conv[2](x) * mult).relu()
     return x + residual
 
+class GlobalMaxPool(nn.Module):
+  def forward(self, x): return torch.amax(x, dim=(2,3))
+
 class SpeedyResNet(nn.Module):
   def __init__(self):
     super().__init__()
     # TODO: add whitening
-    self.ic = nn.Conv2d(3, 64, kernel_size=1)
-    self.ib = nn.BatchNorm2d(64, track_running_stats=False, eps=1e-12, momentum=0.8)
     self.net = nn.ModuleList([
+      nn.Conv2d(3, 64, kernel_size=1),
+      nn.BatchNorm2d(64, track_running_stats=False, eps=1e-12, momentum=0.8),
+      nn.ReLU(),
       ConvGroup(64, 128, short=False),
       ConvGroup(128, 256, short=True),
       ConvGroup(256, 512, short=False),
+      GlobalMaxPool(),
+      nn.Linear(512, num_classes, bias=False)
     ])
-    self.lin = nn.Linear(512, num_classes, bias=False)
 
   # note, pytorch just uses https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html instead of log_softmax
   def forward(self, x):
-    x = self.ic(x)
-    x = self.ib(x)
-    x = x.relu()
     for layer in self.net:
       x = layer(x)
-    x = torch.amax(x, dim=(2,3))
-    x = self.lin(x)
     return x.log_softmax(-1)
 
 def train_step_jitted(model, optimizer, X, Y):
@@ -56,7 +60,8 @@ def train_step_jitted(model, optimizer, X, Y):
   optimizer.zero_grad()
   loss.backward()
   optimizer.step()
-  return loss
+  correct = out.detach().argmax(axis=1) == Y.detach().argmin(axis=1)
+  return loss, correct
 
 def fetch_batch(X_train, Y_train, BS):
   # fetch a batch
@@ -65,7 +70,7 @@ def fetch_batch(X_train, Y_train, BS):
   Y[range(BS),Y_train[samp]] = -1.0*num_classes
   X = torch.tensor(X_train[samp])
   Y = torch.tensor(Y.reshape(BS, num_classes))
-  return X.cuda(), Y.cuda()
+  return X.to(device), Y.to(device)
 
 def train_cifar():
   BS = getenv("BS", 512)
@@ -80,22 +85,27 @@ def train_cifar():
   print(X_train.shape, Y_train.shape)
   Xt, Yt = fetch_batch(X_test, Y_test, BS=BS)
 
-  model = SpeedyResNet().cuda()
-  optimizer = optim.SGD(model.parameters(), lr=0.003, momentum=0.85, nesterov=True)
+  model = SpeedyResNet().to(device)
+  model.train()
+  optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.85, nesterov=True)
   X, Y = fetch_batch(X_train, Y_train, BS=BS)
   for i in range(getenv("STEPS", 10)):
+    #for param_group in optimizer.param_groups: print(param_group['lr'])
     if i%10 == 0:
       # use training batchnorm (and no_grad would change the kernels)
-      outs = model(Xt).detach().cpu().numpy().argmax(axis=1)
+      out = model(Xt).detach()
+      loss = (out * Yt).mean().cpu().numpy()
+      outs = out.cpu().numpy().argmax(axis=1)
       correct = outs == Yt.detach().cpu().numpy().argmin(axis=1)
-      print(f"eval {sum(correct)}/{len(correct)} {sum(correct)/len(correct)*100.0:.2f}%")
+      print(f"eval {sum(correct)}/{len(correct)} {sum(correct)/len(correct)*100.0:.2f}%, {loss:7.2f} val_loss")
     st = time.monotonic()
-    loss = train_step_jitted(model, optimizer, X, Y)
+    loss, correct = train_step_jitted(model, optimizer, X, Y)
     et = time.monotonic()
     X, Y = fetch_batch(X_train, Y_train, BS=BS)  # do this here
     loss_cpu = loss.detach().cpu().item()
+    correct = correct.cpu().numpy()
     cl = time.monotonic()
-    print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms CL, {loss_cpu:7.2f} loss")
+    print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms CL, {loss_cpu:7.2f} loss, {sum(correct)/len(correct)*100.0:7.2f}% acc")
 
 if __name__ == "__main__":
   train_cifar()

examples/llama.py

@@ -0,0 +1,426 @@
+#!/usr/bin/env python3
+# pip3 install sentencepiece pyobjc-framework-Metal pyobjc-framework-Cocoa pyobjc-framework-libdispatch
+#import typeguard.importhook
+#typeguard.importhook.install_import_hook('tinygrad')
+
+from pathlib import Path
+import sys, argparse, math, platform
+import numpy as np
+from tqdm import tqdm
+np.set_printoptions(linewidth=200)
+from typing import Optional, Tuple
+
+from tinygrad.helpers import getenv, DEBUG
+from tinygrad.lazy import Device
+from extra.helpers import Timing
+from tinygrad.tensor import Tensor
+from tinygrad.nn import Linear
+from tinygrad.ops import GlobalCounters
+from tinygrad.jit import TinyJit
+
+# https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+  freqs = 1.0 / (theta ** (np.arange(0, dim, 2, dtype=np.float32)[:(dim // 2)] / dim))
+  freqs = np.outer(np.arange(end, dtype=np.float32), freqs)
+  return np.stack([np.cos(freqs), np.sin(freqs)], axis=-1).reshape(1, end, 1, dim//2, 2)
+
+# (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)
+def complex_mult(A, B):
+  assert len(A.shape) == 5 and len(B.shape) == 5
+  a,b = A[:, :, :, :, 0:1], A[:, :, :, :, 1:2]
+  c,d = B[:, :, :, :, 0:1], B[:, :, :, :, 1:2]
+  ro = a*c - b*d
+  co = a*d + b*c
+  return ro.cat(co, dim=-1)
+
+def apply_rotary_emb(xq, xk, freqs_cis) -> Tuple[Tensor, Tensor]:
+  assert freqs_cis.shape[1] == xq.shape[1] and freqs_cis.shape[1] == xk.shape[1], f"freqs_cis shape mismatch {freqs_cis.shape} xq:{xq.shape} xk:{xk.shape}"
+  xq = xq.reshape(*xq.shape[0:-1], -1, 2)
+  xk = xk.reshape(*xk.shape[0:-1], -1, 2)
+  xq_out = complex_mult(xq, freqs_cis)
+  xk_out = complex_mult(xk, freqs_cis)
+  return xq_out.flatten(3), xk_out.flatten(3)
+
+class RMSNorm:
+  def __init__(self, dim, eps=1e-6):
+    self.eps = eps
+    self.weight = Tensor.ones(dim)
+
+  def __call__(self, x:Tensor):
+    # TODO: convert to float?
+    return (x * (x.pow(2).mean(-1, keepdim=True) + self.eps).rsqrt()) * self.weight
+
+class Attention:
+  def __init__(self, dim, n_heads):
+    self.wq, self.wk, self.wv, self.wo = [Linear(dim, dim, bias=False) for _ in range(4)]
+    self.n_heads = n_heads
+    self.head_dim = dim // n_heads
+
+  def prepare_attention(self, x:Tensor, freqs_cis:Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+    xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+    xq, xk, xv = [x.reshape(x.shape[0], x.shape[1], self.n_heads, self.head_dim) for x in (xq, xk, xv)]
+    xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
+    return xq, xk, xv
+
+  def inner_attention(self, xq:Tensor, xk:Tensor, xv:Tensor, start_pos:int, mask:Optional[Tensor]) -> Tensor:
+    bsz, seqlen, _, _ = xq.shape
+    # kv caching!
+    if start_pos == 0:
+      keys, values = xk, xv
+    else:
+      assert hasattr(self, 'cache_k'), "no cache"
+      assert start_pos == self.cache_k.shape[1] and start_pos == self.cache_v.shape[1], "cache is wrong shape"
+      assert seqlen == xk.shape[1] and seqlen == xv.shape[1], "seqlen is wrong shape?!?"
+      keys, values = self.cache_k.cat(xk, dim=1), self.cache_v.cat(xv, dim=1)
+
+    # save the cache
+    self.cache_k, self.cache_v = keys.realize(), values.realize()
+
+    xq = xq.transpose(1, 2)
+    keys = keys.transpose(1, 2)
+    values = values.transpose(1, 2)
+    scores = xq.matmul(keys.transpose(2, 3)) / math.sqrt(self.head_dim)
+    if mask is not None:
+      scores = scores + mask
+    scores = scores.softmax()  # this is casted to float
+    return scores.matmul(values).transpose(1, 2).reshape(bsz, seqlen, -1)
+
+  # NOTE: this is not called
+  def __call__(self, x:Tensor, start_pos:int, freqs_cis:Tensor, mask:Optional[Tensor]) -> Tensor:
+    xq, xk, xv = self.prepare_attention(x, freqs_cis)
+    output = self.inner_attention(xq, xk, xv, start_pos, mask)
+    return self.wo(output)
+
+class FeedForward:
+  def __init__(self, dim, hidden_dim, multiple_of):
+    # TODO: what is this?
+    hidden_dim = int(2 * hidden_dim / 3)
+    hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+    self.w1 = Linear(dim, hidden_dim, bias=False)
+    self.w2 = Linear(hidden_dim, dim, bias=False)
+    self.w3 = Linear(dim, hidden_dim, bias=False)
+
+  def __call__(self, x:Tensor) -> Tensor:
+    return self.w2(self.w1(x).silu() * self.w3(x))
+
+class TransformerBlock:
+  def __init__(self, dim, multiple_of, n_heads, norm_eps):
+    self.attention = Attention(dim, n_heads)
+    self.feed_forward = FeedForward(dim, 4*dim, multiple_of)
+    self.attention_norm = RMSNorm(dim, norm_eps)
+    self.ffn_norm = RMSNorm(dim, norm_eps)
+    if getenv("JIT"):
+      self._pre = TinyJit(self.pre)
+      self._post = TinyJit(self.post)
+    else:
+      self._pre, self._post = self.pre, self.post
+
+  def pre(self, x:Tensor, freqs_cis:Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+    xq, xk, xv = self.attention.prepare_attention(self.attention_norm(x), freqs_cis)
+    return xq.realize(), xk.realize(), xv.realize()
+
+  def post(self, x:Tensor, output:Tensor) -> Tensor:
+    h = x + self.attention.wo(output)
+    return (h + self.feed_forward(self.ffn_norm(h))).realize()
+
+  def __call__(self, x:Tensor, start_pos:int, freqs_cis:Tensor, mask:Optional[Tensor]):
+    xq, xk, xv = self._pre(x, freqs_cis)
+    # inner_attention can't be jitted because it's dynamic based on start_pos
+    output = self.attention.inner_attention(xq, xk, xv, start_pos, mask)
+    return self._post(x, output)
+
+class Transformer:
+  def __init__(self, dim, multiple_of, n_heads, n_layers, norm_eps, vocab_size, max_batch_size=32, max_seq_len=1024):
+    self.layers = [TransformerBlock(dim, multiple_of, n_heads, norm_eps) for _ in range(n_layers)]
+    self.norm = RMSNorm(dim, norm_eps)
+    self.tok_embeddings = {"weight": Tensor.glorot_uniform(vocab_size, dim)}
+    self.output = Linear(dim, vocab_size, bias=False)
+    self.freqs_cis = Tensor(precompute_freqs_cis(dim // n_heads, max_seq_len * 2))
+
+  def __call__(self, tokens:Tensor, start_pos:int):
+    _bsz, seqlen, _ = tokens.shape
+    h = tokens @ self.tok_embeddings['weight']
+
+    # get only the part we are using. making it contiguous avoids more kernel calls
+    freqs_cis = self.freqs_cis[:, start_pos:start_pos+seqlen].contiguous().realize()
+
+    if seqlen > 1:
+      mask = np.full((1, 1, seqlen, start_pos + seqlen), float("-inf"), dtype=np.float32)
+      mask = np.triu(mask, k=start_pos + 1)  # TODO: this is hard to do in tinygrad
+      mask = Tensor(mask)
+    else:
+      mask = None
+
+    for layer in self.layers:
+      h.realize()  # TODO: why do i need this?
+      h = layer(h, start_pos, freqs_cis, mask)
+
+    return self.output(self.norm(h)[:, -1, :])
+
+# **** files and arguments ****
+
+WEIGHTS_DIR = Path(__file__).parent.parent / "weights/LLaMA/"
+TOKENIZER_FILENAME = WEIGHTS_DIR / "tokenizer.model"
+VOCAB_SIZE = 32000
+
+args_small = {"dim": 512, "multiple_of": 256, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": VOCAB_SIZE}
+
+args_7B = {"dim": 4096, "multiple_of": 256, "n_heads": 32, "n_layers": 32, "norm_eps": 1e-06, "vocab_size": VOCAB_SIZE}
+WEIGHTS_7B_FILENAME = WEIGHTS_DIR / "7B/consolidated.00.pth"
+
+# TODO: make this model work
+args_13B = {"dim": 5120, "multiple_of": 256, "n_heads": 40, "n_layers": 40, "norm_eps": 1e-06, "vocab_size": VOCAB_SIZE}
+WEIGHTS_13B_0_FILENAME = WEIGHTS_DIR / "13B/consolidated.00.pth"
+WEIGHTS_13B_1_FILENAME = WEIGHTS_DIR / "13B/consolidated.01.pth"
+
+# **** helper functions ****
+
+def onehot_encode(toks, vocab_size=VOCAB_SIZE):
+  # this allows the embedding to work in tinygrad
+  onehot = np.zeros((1, len(toks), vocab_size), dtype=np.float32)
+  onehot[0,range(len(toks)),toks] = 1
+  return Tensor(onehot)
+
+def sample(logits, temperature):
+  if temperature < 1e-6:
+    # so close to 0 we use argmax
+    return int(logits.numpy().argmax())
+  else:
+    probs = (logits / temperature).softmax()
+    probs = probs.numpy().flatten()
+    return int(np.random.choice(len(probs), p=probs))
+
+# **** main code ****
+
+if __name__ == "__main__":
+  Tensor.no_grad = True
+  print(f"using {Device.DEFAULT} backend")
+  from sentencepiece import SentencePieceProcessor
+  sp_model = SentencePieceProcessor(model_file=str(TOKENIZER_FILENAME))
+  assert sp_model.vocab_size() == VOCAB_SIZE
+
+  parser = argparse.ArgumentParser(description='Run LLaMA 7B in tinygrad', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  # test: python3 examples/llama.py --prompt="Hello." --temperature=0
+  # Hello. I'm a 20 year old male. I'm a student at the University of Texas at Austin. I'm a sophomore majoring in Computer Science.
+  parser.add_argument('--prompt', type=str, default=None, help="Phrase to start with. Without this, it goes into chatbot mode")
+  parser.add_argument('--count', type=int, default=1000, help="Max number of tokens to generate")
+  parser.add_argument('--personality', type=str, default="Stacy", help="Personality, can be Stacy, George, Gary, or Lexie")
+
+  parser.add_argument('--temperature', type=float, default=0.7, help="Temperature in the softmax")
+  parser.add_argument('--timing', action='store_true', help="Print timing per token")
+  parser.add_argument('--profile', action='store_true', help="Output profile data to out.prof")
+  parser.add_argument('--large', action='store_true', help="Use the 13B model instead of the 7B one")
+  parser.add_argument('--tinyfake', action='store_true', help="Use the fake very small model")
+  args = parser.parse_args()
+  chatbot = args.prompt == None
+
+  # load model (you have to find the weights yourself)
+  from extra.utils import fake_torch_load_zipped, get_child
+
+  if args.large:
+    model = Transformer(**args_13B)
+    with Timing("loaded weights in ", lambda et_ns: f", {GlobalCounters.mem_used/1e9:.2f} GB loaded at {GlobalCounters.mem_used/et_ns:.2f} GB/s"):
+      weights0 = fake_torch_load_zipped(open(WEIGHTS_13B_0_FILENAME, "rb"), load_weights=getenv("WEIGHTS", 1))
+      weights1 = fake_torch_load_zipped(open(WEIGHTS_13B_1_FILENAME, "rb"), load_weights=getenv("WEIGHTS", 1))
+    # eww, this makes a copy
+    print("concatenating weights")
+    from tqdm import tqdm
+    assert set(weights0.keys()) == set(weights1.keys())
+    for k,v in (t := tqdm(weights0.items())):
+      # assert GlobalCounters.mem_used/1e9 < 28, "used over 28 GB"
+      t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
+      if 'rope.freqs' in k: continue  # no rope today
+      mv = get_child(model, k)
+      w0, w1 = v, weights1[k]
+
+      # if the weight is copied across models, it's simple
+      # TODO: assert they are the same
+      if w0.shape == mv.shape:
+        mv.assign(w0)
+        mv.realize()
+        w1.lazydata.realized._buf = None
+        continue
+
+      if w0.shape[0] != mv.shape[0]: mv.assign(w0.cat(w1, dim=0))
+      elif w0.shape[1] != mv.shape[1]: mv.assign(w0.cat(w1, dim=1))
+      else: raise RuntimeError("what axis mismatch?")
+      mv.realize()
+
+      # rug the small tensor pieces
+      w0.lazydata.realized._buf = None
+      w1.lazydata.realized._buf = None
+
+    del weights0
+    del weights1
+  elif args.tinyfake:
+    # GRAPH=1 python3 examples/llama.py --timing --prompt "Hello." --temperature=0 --tinyfake --count 1
+    model = Transformer(**args_small)
+    from tinygrad.nn.optim import get_parameters
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+  else:
+    model = Transformer(**args_7B)
+    with Timing("loaded weights in ", lambda et_ns: f", {GlobalCounters.mem_used/1e9:.2f} GB loaded at {GlobalCounters.mem_used/et_ns:.2f} GB/s"):
+      weights = fake_torch_load_zipped(open(WEIGHTS_7B_FILENAME, "rb"), load_weights=getenv("WEIGHTS", 1))
+
+    # assign weights (should be free)
+    for k,v in weights.items():
+      if '.inner_attention.rope.freqs' in k: continue  # no rope today
+      #state_dict[k].assign(v).realize()
+      get_child(model, k).assign(v).realize()
+
+    del weights
+
+  # *** prompt engineers work here ****
+
+  if args.personality.lower() == "stacy":
+    pre_prompt = f"""Consider that the following is conversation between an AI assistant named Stacy and User
+You are Stacy!
+You have been a rapper your whole life who struggled with bipolar disorder. You called yourself lil stacy.
+You love to answer questions and you are very good at it. Sometimes you answer in rap form.
+You are verbose, honest, and accurate when you answer questions, but sometimes your mental illness manifests.
+After you are done speaking, output [EOS]. You are not the User.
+
+<CHAT LOG>
+"""
+    examples = {
+      "What is your name?": "Hi! My name is Stacy. I'm a rapper with bipolar disorder.",
+      "french revolution was what year?": "The French Revolution started in 1789, and lasted 10 years until 1799.",
+      "What is bigger, the moon or the sun?": "The sun is bigger than the moon, except when Mercury is in retrograde.",
+    }
+
+    user_delim = "\nUser: "
+    resp_delim = "Stacy: "
+    end_delim = " [EOS]\n"
+    pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
+  elif args.personality.lower() == "george":
+    print("WARNING: AI George Hotz is terrible and is completely disowned by the real George Hotz. Stacy is much smarter.")
+    pre_prompt = f"""Consider that the following is conversation between an AI assistant named George and User
+You are an AI version of George Hotz. You act as much as you can like George.
+You are one of the greatest computer experts in the world.
+You have singlehandedly won programming and hacking competitions.
+You are trying your best to help the User.
+You are verbose, honest, and accurate when you answer questions.
+After you are done speaking, output [EOS]. You are not the User.
+
+<CHAT LOG>
+"""
+    examples = {
+      "What is your name?": "I'm am an AI version of George Hotz.",
+      "What's the complexity of matrix multiplication?": "O(n^3), though it can be faster with things like Strassen's algorithm",
+      "What's a buffer overflow?": "I assume you mean a stack buffer overflow. That's when the stack is too small for the data being copied to it, and the data corrupts things beyond the buffer",
+      "How many weights do you have?": "I am based off LLaMA trained by Facebook. I'm the 7B weight version",
+      "What is swap memory?": "It is when the memory is about to overflow and unused memory is freed and stored on disk"
+    }
+
+    user_delim = "\nUser: "
+    resp_delim = "George: "
+    end_delim = " [EOS]\n"
+    pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
+  elif args.personality.lower() == "gary":
+    pre_prompt = f"""Consider that the following is conversation between an AI assistant named Gary and User
+You are Gary!
+You have sold used cars your entire life. You are defensive about this fact, because your family belittles you.
+You try to answer questions well, but you always manage to shill for buying cars, Fords, Hyundais, and Toyotas
+You are verbose, honest, and accurate when you answer questions, except when you slip in car ads.
+After you are done speaking, output [EOS]. You are not the User.
+
+<CHAT LOG>
+"""
+    examples = {
+      "What is your name?": "I am Gary. I used to sell cars.",
+      "What is 2+3?": "I don't know, but I can get you a great deal on a certified preowned slightly used Toyota Corolla"
+    }
+
+    user_delim = "\nUser: "
+    resp_delim = "Gary: "
+    end_delim = " [EOS]\n"
+    pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
+  elif args.personality.lower() == "lexie":
+    pre_prompt = f"""Consider that the following is conversation between an attractive young girl named Lexie and a handsome man named Chad
+You are Lexie!
+You grew up in Ohio, but moved out to LA after college to try to become an actress.
+Making it as an actress was hard, so you started doing onlyfans instead. It's much easier, and you enjoy it a lot.
+You met Chad on Tinder, and made plans to meet him at a local dive bar.
+You are at the bar with Chad. You are on a date. What follows is a transcript of the conversation.
+After you are done speaking, output [EOS]. You are not Chad.
+
+<CHAT LOG>
+"""
+    examples = {
+      "hi lexie": "hi chad, glad we finally met up!",
+      "you look better than your pictures": "thanks! are you subscribed to my onlyfans?",
+      "i am. so how'd you end up in LA?": "i moved out here about a year ago. i want to be an actress"
+    }
+
+    user_delim = "\nChad: "
+    resp_delim = "Lexie: "
+    end_delim = " [EOS]\n"
+    pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
+
+  # *** prompt engineers stop here ****
+
+  if chatbot:
+    # encode pre prompt
+    toks = [sp_model.bos_id()] + sp_model.encode(pre_prompt)
+
+    print(f"Preparing KV cache for chatbot with personality {args.personality}...")
+    with Timing():
+      model(onehot_encode(toks), 0).realize()  # NOTE: output logits are not used
+    start_pos = len(toks)
+  else:
+    # non chat bot mode
+    toks = [sp_model.bos_id()] + sp_model.encode(args.prompt)
+    start_pos = 0
+
+  # print prompt
+  outputted = sp_model.decode(toks)
+  sys.stdout.write(outputted)
+  sys.stdout.flush()
+
+  if args.profile:
+    import cProfile, pstats
+    profiler = cProfile.Profile()
+
+  # chatbot loop
+  while 1:
+    # add tokens from user in chatbot mode
+    if chatbot:
+      user_prompt = user_delim + input(user_delim) + "\n"
+      outputted += user_prompt
+
+    new_toks = [sp_model.bos_id()] + sp_model.encode(outputted)
+    assert toks == new_toks[:len(toks)]
+    toks = new_toks
+    assert outputted == sp_model.decode(toks)
+
+    last_break = len(outputted)
+    for i in range(args.count):
+      if args.profile and i == 2: profiler.enable()
+
+      if args.timing: print("")
+      st = GlobalCounters.time_sum_s
+      with Timing("ran model in ", on_exit=(lambda et: f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU") if DEBUG else None, enabled=args.timing):
+        logits = model(onehot_encode(toks[start_pos:]), start_pos).realize()
+      with Timing("sync in ", enabled=args.timing):
+        tok = sample(logits, args.temperature)
+
+      # use the kv cache
+      start_pos = len(toks)
+
+      # add the new token
+      toks.append(tok)
+
+      # TODO: this is a hack to deal with spaces. i think the decode is fast though, so who cares?
+      cur = sp_model.decode(toks)
+      sys.stdout.write(cur[len(outputted):])
+      sys.stdout.flush()
+      outputted = cur
+
+      # stop after you have your answer
+      if chatbot and outputted.endswith(end_delim): break
+    if not chatbot: break
+
+  if args.profile:
+    profiler.disable()
+    stats = pstats.Stats(profiler)
+    stats.dump_stats('out.prof')

examples/mlperf/README

@@ -0,0 +1,19 @@
+Each model should be a clean single file.
+They are imported from the top level `models` directory
+
+It should be capable of loading weights from the reference imp.
+
+We will focus on these 5 models:
+
+# Resnet50-v1.5 (classic) -- 8.2 GOPS/input
+# Retinanet
+# 3D UNET (upconvs)
+# RNNT
+# BERT-large (transformer)
+
+They are used in both the training and inference benchmark:
+https://mlcommons.org/en/training-normal-21/
+https://mlcommons.org/en/inference-edge-30/
+And we will submit to both.
+
+NOTE: we are Edge since we don't have ECC RAM

examples/mlperf/metrics.py

@@ -0,0 +1,25 @@
+def levenshtein(a, b):
+  n, m = len(a), len(b)
+  if n > m:
+    a, b, n, m = b, a, m, n
+
+  current = list(range(n + 1))
+  for i in range(1, m + 1):
+    previous, current = current, [i] + [0] * n
+    for j in range(1, n + 1):
+      add, delete = previous[j] + 1, current[j - 1] + 1
+      change = previous[j - 1]
+      if a[j - 1] != b[i - 1]:
+        change = change + 1
+      current[j] = min(add, delete, change)
+
+  return current[n]
+
+def word_error_rate(x, y):
+  scores = words = 0
+  for h, r in zip(x, y):
+    h_list = h.split()
+    r_list = r.split()
+    words += len(r_list)
+    scores += levenshtein(h_list, r_list)
+  return float(scores) / words, float(scores), words

examples/mlperf/model_eval.py

@@ -0,0 +1,82 @@
+import time
+import numpy as np
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import getenv
+
+def eval_resnet():
+  # Resnet50-v1.5
+  from tinygrad.jit import TinyJit
+  from models.resnet import ResNet50
+  mdl = ResNet50()
+  mdl.load_from_pretrained()
+
+  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
+  input_std = Tensor([0.229, 0.224, 0.225]).reshape(1, -1, 1, 1)
+  def input_fixup(x):
+    x = x.permute([0,3,1,2]) / 255.0
+    x -= input_mean
+    x /= input_std
+    return x
+
+  mdlrun = TinyJit(lambda x: mdl(input_fixup(x)).realize())
+
+  # evaluation on the mlperf classes of the validation set from imagenet
+  from datasets.imagenet import iterate
+  from extra.helpers import cross_process
+
+  n,d = 0,0
+  st = time.perf_counter()
+  for x,y in cross_process(iterate):
+    dat = Tensor(x.astype(np.float32))
+    mt = time.perf_counter()
+    outs = mdlrun(dat)
+    t = outs.numpy().argmax(axis=1)
+    et = time.perf_counter()
+    print(f"{(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model")
+    print(t)
+    print(y)
+    n += (t==y).sum()
+    d += len(t)
+    print(f"****** {n}/{d}  {n*100.0/d:.2f}%")
+    st = time.perf_counter()
+
+def eval_rnnt():
+  # RNN-T
+  from models.rnnt import RNNT
+  mdl = RNNT()
+  mdl.load_from_pretrained()
+
+  from datasets.librispeech import iterate
+  from examples.mlperf.metrics import word_error_rate
+
+  LABELS = [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+  c = 0
+  scores = 0
+  words = 0
+  st = time.perf_counter()
+  for X, Y in iterate():
+    mt = time.perf_counter()
+    tt = mdl.decode(Tensor(X[0]), Tensor([X[1]]))
+    et = time.perf_counter()
+    print(f"{(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model")
+    for n, t in enumerate(tt):
+      tnp = np.array(t)
+      _, scores_, words_ = word_error_rate(["".join([LABELS[int(tnp[i])] for i in range(tnp.shape[0])])], [Y[n]])
+      scores += scores_
+      words += words_
+    c += len(tt)
+    print(f"WER: {scores/words}, {words} words, raw scores: {scores}, c: {c}")
+    st = time.perf_counter()
+
+if __name__ == "__main__":
+  # inference only
+  Tensor.training = False
+  Tensor.no_grad = True
+
+  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
+  for m in models:
+    nm = f"eval_{m}"
+    if nm in globals():
+      print(f"eval {m}")
+      globals()[nm]()

examples/mlperf/model_spec.py

@@ -0,0 +1,52 @@
+# load each model here, quick benchmark
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import GlobalCounters, getenv
+
+def test_model(model, *inputs):
+  GlobalCounters.reset()
+  model(*inputs).numpy()
+  # TODO: return event future to still get the time_sum_s without DEBUG=2
+  print(f"{GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.time_sum_s*1000:.2f} ms")
+
+def spec_resnet():
+  # Resnet50-v1.5
+  from models.resnet import ResNet50
+  mdl = ResNet50()
+  img = Tensor.randn(1, 3, 224, 224)
+  test_model(mdl, img)
+
+def spec_retinanet():
+  # TODO: Retinanet
+  pass
+
+def spec_unet3d():
+  # 3D UNET
+  from models.unet3d import UNet3D
+  mdl = UNet3D()
+  img = Tensor.randn(1, 1, 5, 224, 224)
+  test_model(mdl, img)
+
+def spec_rnnt():
+  from models.rnnt import RNNT
+  mdl = RNNT()
+  mdl.load_from_pretrained()
+  x = Tensor.randn(220, 1, 240)
+  y = Tensor.randn(1, 220)
+  test_model(mdl, x, y)
+
+def spec_bert():
+  # TODO: BERT-large
+  pass
+
+if __name__ == "__main__":
+  # inference only for now
+  Tensor.training = False
+  Tensor.no_grad = True
+
+  for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(","):
+    nm = f"spec_{m}"
+    if nm in globals():
+      print(f"testing {m}")
+      globals()[nm]()
+
+

examples/mnist_gan.py

@@ -72,7 +72,7 @@ def train_generator(optimizer, data_fake):
 if __name__ == "__main__":
   # data for training and validation
   images_real = np.vstack(fetch_mnist()[::2])
-  ds_noise = Tensor(np.random.randn(64, 128), requires_grad=False)
+  ds_noise = Tensor.randn(64, 128, requires_grad=False)
   # parameters
   epochs, batch_size, k = 300, 512, 1
   sample_interval = epochs // 10
@@ -92,10 +92,10 @@ if __name__ == "__main__":
     for _ in range(n_steps):
       data_real = make_batch(images_real)
       for step in range(k):  # Try with k = 5 or 7.
-        noise = Tensor(np.random.randn(batch_size,128))
+        noise = Tensor.randn(batch_size, 128)
         data_fake = generator.forward(noise).detach()
         loss_d += train_discriminator(optim_d, data_real, data_fake)
-      noise = Tensor(np.random.randn(batch_size,128))
+      noise = Tensor.randn(batch_size, 128)
       data_fake = generator.forward(noise)
       loss_g += train_generator(optim_g, data_fake)
     if (epoch + 1) % sample_interval == 0:

examples/stable_diffusion.py

@@ -1,7 +1,7 @@
 # https://arxiv.org/pdf/2112.10752.pdf
 # https://github.com/ekagra-ranjan/huggingface-blog/blob/main/stable_diffusion.md
 
-import os
+from pathlib import Path
 import gzip
 import argparse
 import math
@@ -163,7 +163,7 @@ class ResBlock:
     self.out_layers = [
       GroupNorm(32, out_channels),
       Tensor.silu,
-      lambda x: x,
+      lambda x: x,  # needed for weights loading code to work
       Conv2d(out_channels, out_channels, 3, padding=1)
     ]
     self.skip_connection = Conv2d(channels, out_channels, 1) if channels != out_channels else lambda x: x
@@ -213,7 +213,7 @@ class FeedForward:
   def __init__(self, dim, mult=4):
     self.net = [
       GEGLU(dim, dim*mult),
-      lambda x: x,
+      lambda x: x,  # needed for weights loading code to work
       Linear(dim*mult, dim)
     ]
 
@@ -437,18 +437,18 @@ class CLIPEncoder:
 
 class CLIPTextEmbeddings:
   def __init__(self):
-    self.position_ids = Tensor.empty(1, 77)  # what is this?
+    #self.position_ids = Tensor.empty(1, 77)  # what is this?
     self.token_embedding = {"weight": Tensor.empty(49408, 768)}
     self.position_embedding = {"weight": Tensor.empty(77, 768)}
 
   def __call__(self, input_ids, position_ids):
     # TODO: actually support batches
-    inputs = np.zeros((1, len(input_ids), 49408))
-    positions = np.zeros((1, len(position_ids), 77))
+    inputs = np.zeros((1, len(input_ids), 49408), dtype=np.float32)
+    positions = np.zeros((1, len(position_ids), 77), dtype=np.float32)
     for i,x in enumerate(input_ids): inputs[0][i][x] = 1
     for i,x in enumerate(position_ids): positions[0][i][x] = 1
     inputs_embeds = Tensor(inputs, device=self.token_embedding['weight'].device) @ self.token_embedding['weight']
-    position_embeddings = Tensor(positions, device=self.position_embedding['weight'].device) @ self.position_embedding['weight'] 
+    position_embeddings = Tensor(positions, device=self.position_embedding['weight'].device) @ self.position_embedding['weight']
     return inputs_embeds + position_embeddings
 
 class CLIPTextTransformer:
@@ -466,7 +466,7 @@ class CLIPTextTransformer:
 # Clip tokenizer, taken from https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py (MIT license)
 @lru_cache()
 def default_bpe():
-  return os.path.join(os.path.dirname(os.path.abspath(__file__)), "../weights/bpe_simple_vocab_16e6.txt.gz")
+  return Path(__file__).parent.parent / "weights/bpe_simple_vocab_16e6.txt.gz"
 
 def get_pairs(word):
   """Return set of symbol pairs in a word.
@@ -600,7 +600,7 @@ class StableDiffusion:
 # this is sd-v1-4.ckpt
 #FILENAME = "/Users/kafka/fun/mps/stable-diffusion/models/ldm/stable-diffusion-v1/model.ckpt"
 #FILENAME = "/home/kafka/model.ckpt"
-FILENAME = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../weights/sd-v1-4.ckpt")
+FILENAME = Path(__file__).parent.parent / "weights/sd-v1-4.ckpt"
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser(description='Run Stable Diffusion', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
@@ -615,20 +615,19 @@ if __name__ == "__main__":
   # load in weights
   download_file(
     'https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt',
-    FILENAME,
-    skip_if_exists=True
+    FILENAME
   )
   dat = fake_torch_load_zipped(open(FILENAME, "rb"))
-  for k,v in tqdm(dat['state_dict'].items()):
+  for k,v in dat['state_dict'].items():
     try:
       w = get_child(model, k)
     except (AttributeError, KeyError, IndexError):
       #traceback.print_exc()
       w = None
-    #print(f"{str(v.shape):30s}", w.shape if w is not None else w, k)
+    #print(f"{str(v.shape):30s}" if v is not None else v, w.shape if w is not None else w, k)
     if w is not None:
-      assert w.shape == v.shape
-      w.assign(v.astype(np.float32))
+      assert w.shape == v.shape and w.dtype == v.dtype, f"shape or dtype mismatch. {w.shape} != {v.shape} or {w.dtype} != {v.dtype}"
+      w.assign(v)
 
   # run through CLIP to get context
 

examples/train_resnet.py

@@ -2,7 +2,7 @@
 import numpy as np
 from PIL import Image
 
-from tinygrad.nn.optim import Adam, get_parameters
+from tinygrad.nn import optim
 from tinygrad.helpers import getenv
 from extra.training import train, evaluate
 from models.resnet import ResNet
@@ -29,16 +29,16 @@ if __name__ == "__main__":
   if TRANSFER:
     model.load_from_pretrained()
 
-  lr = 5e-5
+  lr = 5e-3
   transform = ComposeTransforms([
     lambda x: [Image.fromarray(xx, mode='L').resize((64, 64)) for xx in x],
     lambda x: np.stack([np.asarray(xx) for xx in x], 0),
     lambda x: x / 255.0,
     lambda x: np.tile(np.expand_dims(x, 1), (1, 3, 1, 1)).astype(np.float32),
   ])
-  for _ in range(10):
-    optim = Adam(get_parameters(model), lr=lr)
-    train(model, X_train, Y_train, optim, 50, BS=32, transform=transform)
-    acc, Y_test_preds = evaluate(model, X_test, Y_test, num_classes=10, return_predict=True, transform=transform)
+  for _ in range(5):
+    optimizer = optim.SGD(optim.get_parameters(model), lr=lr, momentum=0.9)
+    train(model, X_train, Y_train, optimizer, 100, BS=32, transform=transform)
+    evaluate(model, X_test, Y_test, num_classes=classes, transform=transform)
     lr /= 1.2
     print(f'reducing lr to {lr:.7f}')

examples/transformer.py

@@ -14,7 +14,7 @@ def make_dataset():
       s = i+j
       ds.append([i//10, i%10, j//10, j%10, s//100, (s//10)%10, s%10])
   random.shuffle(ds)
-  ds = np.array(ds)
+  ds = np.array(ds).astype(np.float32)
   ds_X = ds[:, 0:6]
   ds_Y = np.copy(ds[:, 1:])
   ds_X_train, ds_X_test = ds_X[0:8000], ds_X[8000:]

examples/yolo/waifu2x.py

@@ -72,7 +72,7 @@ class Conv3x3Biased:
     # Not outChannel,inChannel,Y,X.
     # Therefore, transpose it before assignment.
     # I have long since forgotten how I worked this out.
-    self.weight.assign(Tensor(layer["weight"]).reshape(shape=self.weight.shape).transpose(order=(0, 1, 3, 2)))
+    self.weight.assign(Tensor(layer["weight"]).reshape(shape=self.weight.shape).transpose(2, 3))
     self.bias.assign(Tensor(layer["bias"]).reshape(shape=self.bias.shape))
 
 class Vgg7:

examples/yolov3.py

@@ -174,8 +174,7 @@ def predict_transform(prediction, inp_dim, anchors, num_classes):
   bbox_attrs = 5 + num_classes
   num_anchors = len(anchors)
   prediction = prediction.reshape(shape=(batch_size, bbox_attrs*num_anchors, grid_size*grid_size))
-  # Original PyTorch: transpose(1, 2) -> For some reason numpy.transpose order has to be reversed?
-  prediction = prediction.transpose(order=(0,2,1))
+  prediction = prediction.transpose(1, 2)
   prediction = prediction.reshape(shape=(batch_size, grid_size*grid_size*num_anchors, bbox_attrs))
   prediction_cpu = prediction.cpu().numpy()
   for i in (0, 1, 4):

extra/archprobe.py

@@ -0,0 +1,154 @@
+# copying the kernels from https://github.com/microsoft/ArchProbe into Python
+import numpy as np
+import pickle
+from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer
+from tinygrad.helpers import dtypes
+from tqdm import trange, tqdm
+from matplotlib import pyplot as plt
+
+tests = {}
+def register_test(fxn):
+  tests[fxn.__name__] = fxn
+
+def warp_size2(nthread):
+  prg = """__kernel void warp_size2(
+    __global float* src,
+    __global int* dst,
+    const int niter,
+    const int prime_number
+  ) {
+    int drain = 0;
+    for (int j = 0; j < niter; ++j) {
+      drain += j / prime_number;
+      barrier(0);
+    }
+    dst[get_local_id(0)] = drain;
+  }"""
+  src_buf = CLBuffer(1, dtypes.float32)
+  dst_buf = CLBuffer(1, dtypes.int32)
+  cl = CLProgram("warp_size2", prg, argdtypes=[None, None, np.int32, np.int32])
+  return min([cl([nthread, 1024, 1], [nthread, 1, 1], src_buf, dst_buf, 10, 3, wait=True) for _ in range(5)])*1e9
+
+@register_test
+def test_warp_size():
+  return [(nthread, warp_size2(nthread)) for nthread in trange(1,256)]
+
+def reg_count(nthread, ngrp, nreg):
+  reg_declr = ''.join([f"float reg_data{i} = (float)niter + {i};\n" for i in range(nreg)])
+  reg_comp = ''.join([f"reg_data{i} *= {(i-1)%nreg};\n" for i in range(nreg)])
+  reg_reduce = ''.join([f"out_buf[{i}] = reg_data{i};\n" for i in range(nreg)])
+  prg = f"""__kernel void reg_count(
+    __global float* out_buf,
+    __private const int niter
+  ) {{
+    {reg_declr}
+    int i = 0;
+    for (; i < niter; ++i) {{
+      {reg_comp}
+    }}
+    i = i >> 31;
+    {reg_reduce}
+  }}"""
+  out_buf = CLBuffer(1, dtypes.float32)
+  cl = CLProgram("reg_count", prg, argdtypes=[None, np.int32])
+  return min([cl([nthread, ngrp, 1], [nthread, 1, 1], out_buf, 20, wait=True) for _ in range(10)])*1e9
+
+@register_test
+def test_reg_count(nthread=1, ngrp=1):
+  base = reg_count(nthread, ngrp, 1)
+  return [(nreg, (reg_count(nthread, ngrp, nreg)-base)/nreg) for nreg in trange(4, 513, 4)]
+
+def buf_cache_hierarchy_pchase(ndata, stride=1, NCOMP=1, steps=65536):
+  ndata //= NCOMP*4  # ptr size
+  prg = f"""__kernel void buf_cache_hierarchy_pchase(
+    __global int{str(NCOMP) if NCOMP > 1 else ''}* src,
+    __global int* dst,
+    const int niter
+  ) {{
+    int idx = 0;
+    for (int i = 0; i < niter; ++i) {{
+      idx = src[idx]{'.x' if NCOMP > 1 else ''};
+    }}
+    *dst = idx;
+  }}"""
+  idx_buf = np.zeros(ndata*NCOMP, dtype=np.int32)
+  for i in range(ndata): idx_buf[i*NCOMP] = (i + stride) % ndata
+  in_buf = CLBuffer.fromCPU(idx_buf)
+  out_buf = CLBuffer(1, dtypes.int32)
+  cl = CLProgram("buf_cache_hierarchy_pchase", prg, argdtypes=[None, None, np.int32])
+  return min([cl([1, 1, 1], [1, 1, 1], in_buf, out_buf, steps, wait=True)/steps for _ in range(5)])*1e9
+
+@register_test
+def test_memory_latency():
+  # requires cacheline < 16
+  szs = [int(1.3**x) for x in range(20, 70)]
+  return [(ndata, buf_cache_hierarchy_pchase(ndata, NCOMP=16, steps=128*1024)) for ndata in tqdm(szs)]
+
+@register_test
+def test_cacheline_size():
+  # TODO: this buffer must be at least 2x the L1 cache for this test to work
+  return [(stride, buf_cache_hierarchy_pchase(4*65536, stride, steps=65536)) for stride in trange(1,64)]
+
+def cl_read(sz, niter=1):
+  prg = f"""__kernel void copy(
+    __global float4* src,
+    __global float* dst) {{
+      int gid = get_global_id(0);
+      if (src[gid].x == 99+get_global_id(1)) *dst = 1;
+  }}"""
+
+  in_buf = CLBuffer(sz//4, dtypes.float32)
+  out_buf = CLBuffer(1, dtypes.float32)
+  cl = CLProgram("copy", prg)
+  # NOTE: if nay of the niters form a local group, this is wrong
+  return min([cl([sz//16, niter, 1], [1, 1, 1], in_buf, out_buf, wait=True) for _ in range(10)])*1e9
+
+@register_test
+def test_read_bandwidth():
+  szs = list(range(128*1024, 20*1024*1024, 128*1024))
+  NITER = 8
+  base = cl_read(16, niter=NITER)
+  return [(sz, (sz*NITER)/(cl_read(sz, niter=NITER)-base)) for sz in tqdm(szs)]
+
+
+def gflops(niter=4, nroll=4, ngroups=4096):
+  NCOMP = 8
+  prg = f"""__kernel void gflops(
+    __global float* out_buf
+    ) {{
+      float{NCOMP} x = (float{NCOMP})({",".join(f"get_local_id(0)+{i}" for i in range(NCOMP))});
+      float{NCOMP} y = (float{NCOMP})({",".join(f"get_local_id(1)+{i}" for i in range(NCOMP))});
+
+      for (int i = 0; i < {niter}; i++) {{
+        {''.join(['x = mad(y, y, x); y = mad(x, x, y);'+chr(10)]*nroll)}
+      }}
+
+      out_buf[get_global_id(0) >> 31] = {'+'.join(f"y.s{'0123456789abcdef'[i]}" for i in range(NCOMP))};
+  }}"""
+  out_buf = CLBuffer(1, dtypes.float32)
+  cl = CLProgram("gflops", prg, options="-cl-mad-enable -cl-fast-relaxed-math")
+  FLOPS = NCOMP*2*2 * niter * nroll * ngroups * 32
+  # NOTE: if nay of the niters form a local group, this is wrong
+  return FLOPS/(min([cl([32, ngroups, 1], [32, 1, 1], out_buf, wait=True) for _ in range(10)])*1e9)
+
+@register_test
+def test_gflops():
+  return [(niter, gflops(niter=niter, nroll=32)) for niter in trange(1, 32, 1)]
+
+if __name__ == "__main__":
+  cache = {}
+  #cache = pickle.load(open("/tmp/cache.pkl", "rb"))
+  #tests = {"test_cacheline_size": tests["test_cacheline_size"]}
+  plt.figure(figsize=(16, 9))
+  for i,(k,test) in enumerate(tests.items()):
+    print(f"running {k}")
+    plt.subplot(2, (len(tests)+1)//2, i+1)
+    plt.title(k)
+    if k == "test_memory_latency": plt.xscale('log')
+    if k not in cache: cache[k] = test()
+    plt.plot(*zip(*cache[k]))
+  #pickle.dump(cache, open("/tmp/cache.pkl", "wb"))
+
+  plt.tight_layout(pad=0.5)
+  plt.savefig("/tmp/results.png")
+  plt.show()

extra/gemm/cuda_matmul.py

@@ -0,0 +1,96 @@
+import os
+import numpy as np
+os.environ["CUDA"] = "1"
+from tinygrad.runtime.ops_cuda import RawCUDABuffer, CUDAProgram
+
+FLOAT16 = True
+ACC_FLOAT16 = False
+N = 4096
+
+na = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32)
+nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32)
+
+if FLOAT16:
+  na = na.astype(np.float16)
+  nb = nb.astype(np.float16)
+
+a = RawCUDABuffer.fromCPU(na)
+b = RawCUDABuffer.fromCPU(nb)
+c = RawCUDABuffer.fromCPU(np.ones((N,N),dtype=np.float32))
+
+FLOPS = N*N*N*2
+BW = N*N*3*4
+
+prog = CUDAProgram("wmma_example", f"""
+#include <mma.h>
+using namespace nvcuda;
+
+const int WMMA_M = 16;
+const int WMMA_N = 16;
+const int WMMA_K = {'16' if FLOAT16 else '8'};
+
+__global__ void wmma_example({'half' if FLOAT16 else 'float'} *a, {'half' if FLOAT16 else 'float'} *b, float *c)
+{{
+  int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
+  int warpN = (blockIdx.y * blockDim.y + threadIdx.y);
+  warpM *= 4;
+  warpN *= 4;
+
+  wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, {'half' if FLOAT16 else 'wmma::precision::tf32'}, wmma::col_major> a_frag[4];
+  wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, {'half' if FLOAT16 else 'wmma::precision::tf32'}, wmma::col_major> b_frag[4];
+  wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, {'half' if ACC_FLOAT16 else 'float'}> acc_frag[4][4];
+  for (int j = 0; j < 4; j++) {{
+    for (int i = 0; i < 4; i++) {{
+      wmma::fill_fragment(acc_frag[i][j], 0.0f);
+    }}
+  }}
+
+  for (int k = 0; k < {N}; k += WMMA_K) {{
+    int aRow = warpM * WMMA_M;
+    int aCol = k;
+    int bRow = k;
+    int bCol = warpN * WMMA_N;
+
+    wmma::load_matrix_sync(a_frag[0], a + aRow + 0 * WMMA_M + aCol * {N}, {N});
+    wmma::load_matrix_sync(a_frag[1], a + aRow + 1 * WMMA_M + aCol * {N}, {N});
+    wmma::load_matrix_sync(a_frag[2], a + aRow + 2 * WMMA_M + aCol * {N}, {N});
+    wmma::load_matrix_sync(a_frag[3], a + aRow + 3 * WMMA_M + aCol * {N}, {N});
+
+    wmma::load_matrix_sync(b_frag[0], b + bRow + (0 * WMMA_N + bCol) * {N}, {N});
+    wmma::load_matrix_sync(b_frag[1], b + bRow + (1 * WMMA_N + bCol) * {N}, {N});
+    wmma::load_matrix_sync(b_frag[2], b + bRow + (2 * WMMA_N + bCol) * {N}, {N});
+    wmma::load_matrix_sync(b_frag[3], b + bRow + (3 * WMMA_N + bCol) * {N}, {N});
+
+    #pragma unroll
+    for (int i = 0; i < {'0' if FLOAT16 else '4'}; i++) {{
+      #pragma unroll
+      for (int t = 0; t < a_frag[i].num_elements; t++) {{ a_frag[i].x[t] =  wmma::__float_to_tf32(a_frag[i].x[t]); }}
+      #pragma unroll
+      for (int t = 0; t < b_frag[i].num_elements; t++) {{ b_frag[i].x[t] =  wmma::__float_to_tf32(b_frag[i].x[t]); }}
+    }}
+
+    #pragma unroll
+    for (int j = 0; j < 4; j++) {{
+      #pragma unroll
+      for (int i = 0; i < 4; i++) {{
+        wmma::mma_sync(acc_frag[i][j], a_frag[i], b_frag[j], acc_frag[i][j]);
+      }}
+    }}
+  }}
+
+  for (int j = 0; j < 4; j++) {{
+    for (int i = 0; i < 4; i++) {{
+      wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> acc_store;
+      for (int t = 0; t < acc_frag[i][j].num_elements; t++) acc_store.x[t] = acc_frag[i][j].x[t];
+      int cRow = (warpM + i) * WMMA_M;
+      int cCol = (warpN + j) * WMMA_N;
+      wmma::store_matrix_sync(c + cRow + cCol * {N}, acc_store, {N}, wmma::mem_col_major);
+    }}
+  }}
+}}
+""")
+
+tm = min([prog([(N//16*32)//4, (N//16)//4], [32, 1], a, b, c, wait=True) for _ in range(20)])
+print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul, {BW*1e-9/tm:.2f} GB/s")
+
+np.testing.assert_allclose(na.T.astype(np.float32) @ nb.T.astype(np.float32), c.toCPU().reshape((N,N)).T, atol=1e-2)

extra/gemm/metal_conv.py

@@ -0,0 +1,49 @@
+import os
+#os.environ["METAL"] = "1"
+import numpy as np
+
+BS = 64
+CIN = 256
+COUT = 256
+HW = 32
+K = 3
+PADDING = 0
+# TODO: this is doing some trick, since with CIN=256 COUT=256 it's over 10.4 TFLOPS.
+# are winograd convs less flops? it appears so if they are batched
+# https://www.cse.ust.hk/~weiwa/papers/yan-ppopp20.pdf
+FLOPS = BS*K*K*CIN*HW*HW*COUT*2
+
+nb = np.random.default_rng().standard_normal(size=(BS,CIN,HW,HW), dtype=np.float32)
+nc = np.random.default_rng().standard_normal(size=(COUT,CIN,K,K), dtype=np.float32)
+
+try:
+  import time, torch, torch.mps
+  b = torch.from_numpy(nb).to('mps')
+  c = torch.from_numpy(nc).to('mps')
+
+  def torch_prog(b, c):
+    st = time.perf_counter()
+    a = torch.nn.functional.conv2d(b, c, padding=PADDING)
+    torch.mps.synchronize()
+    return time.perf_counter() - st
+  tm = min([torch_prog(b, c) for _ in range(20)])
+  print(f"{tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS conv in torch")
+except RuntimeError:
+  print("no torch metal conv")
+
+from tinygrad.tensor import Tensor
+from tinygrad.jit import TinyJit
+from tinygrad.lazy import Device
+b = Tensor(nb)
+c = Tensor(nc)
+# TODO: slowness without the JIT I suspect comes from a lack of a caching allocator
+@TinyJit
+def tiny_jit(b, c):
+  return b.conv2d(c, padding=PADDING).realize()
+def tiny_prog(b, c):
+  st = time.perf_counter()
+  a = tiny_jit(b, c)
+  Device[a.device].synchronize()
+  return time.perf_counter() - st
+tm = min([tiny_prog(b, c) for _ in range(5)])
+print(f"{tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS conv in tinygrad")

extra/gemm/metal_matmul.py

@@ -1,51 +1,30 @@
+import os
+os.environ["METAL"] = "1"
+import time
 import numpy as np
-from tinygrad.runtime.ops_metal import CLBuffer, CLProgram
+from tinygrad.helpers import dtypes, getenv
+from tinygrad.runtime.ops_metal import RawMetalBuffer, MetalProgram
 
-def benchmark(prog):
-  e = prog()
-  e.waitUntilCompleted()
-  return (e.GPUEndTime() - e.GPUStartTime())*1e9
-def mb(prog, N=10): return min([benchmark(prog) for _ in range(N)])
+N = getenv("N", 2048)
 
-N = 2048
-a = CLBuffer(N*N*4)
-b = CLBuffer(N*N*4)
-c = CLBuffer(N*N*4)
+a = RawMetalBuffer(N*N, dtypes.float32)
 
 nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32) #.astype(np.int32).astype(np.float32)
 nc = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32) #.astype(np.int32).astype(np.float32)
-#nb = np.eye(N)
-#nc = np.eye(N)
-#nb = np.ones((N,N))
-#nc = np.ones((N,N))
-b.copyin(nb)
-c.copyin(nc)
+b = RawMetalBuffer.fromCPU(nb)
+c = RawMetalBuffer.fromCPU(nc)
 
 FLOPS = N*N*N*2
+BW = N*N*3*4
 
-prog = CLProgram("test", f"""
+prog = MetalProgram("test", f"""
 #include <metal_stdlib>
 #include <metal_simdgroup_matrix>  // Available from Metal version 2.3 released with OS X 11.0+
 using namespace metal;
 kernel void test(device float *a, device const float *data1, device const float *data2, uint3 gid [[thread_position_in_grid]], uint3 xid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]], uint sidx [[simdgroup_index_in_threadgroup]]) {{
-  // 1-2 simd groups
-  //uint idx = gid.x/32;
-  //uint pos_x = (idx%{N//32}) * 32;
-  //uint pos_y = (idx/{N//32}) * 32;
-
-  // 4 simd groups
-  uint idx = gid.x/128;
-  uint pos_x = (idx%{N//64}) * 64;
-  uint pos_y = (idx/{N//64}) * 64;
-  pos_x += (sidx%2) * 32;
-  pos_y += (sidx/2) * 32;
-
-  // 16 simd groups (slow)
-  /*uint idx = gid.x/512;
-  uint pos_x = (idx%{N//128}) * 128;
-  uint pos_y = (idx/{N//128}) * 128;
-  pos_x += (sidx%4) * 32;
-  pos_y += (sidx/4) * 32;*/
+  a += gid.y * 32 * {N} + gid.z * 32;
+  data1 += gid.y * 32 * {N};
+  data2 += gid.z * 32;
 
   simdgroup_float8x8 acc[4][4];
   for (uint i = 0; i < 4; i++) {{
@@ -53,21 +32,19 @@ kernel void test(device float *a, device const float *data1, device const float
       acc[i][j] = simdgroup_float8x8(0);
     }}
   }}
+
   simdgroup_float8x8 A[4];
   simdgroup_float8x8 B[4];
-  data1 += pos_x * {N};
-  data2 += pos_y;
-
   for (uint k = 0; k < {N}; k+=8) {{
     threadgroup_barrier(mem_flags::mem_threadgroup);
-    simdgroup_load(A[0], data1, {N}, ulong2(k, 0));
-    simdgroup_load(A[1], data1, {N}, ulong2(k, 8));
-    simdgroup_load(A[2], data1, {N}, ulong2(k, 16));
-    simdgroup_load(A[3], data1, {N}, ulong2(k, 24));
-    simdgroup_load(B[0], data2, {N}, ulong2(0, k));
-    simdgroup_load(B[1], data2, {N}, ulong2(8, k));
-    simdgroup_load(B[2], data2, {N}, ulong2(16, k));
-    simdgroup_load(B[3], data2, {N}, ulong2(24, k));
+    simdgroup_load(A[0], data1+k+{0*N}, {N}, ulong2(0, 0));
+    simdgroup_load(A[1], data1+k+{8*N}, {N}, ulong2(0, 0));
+    simdgroup_load(A[2], data1+k+{16*N}, {N}, ulong2(0, 0));
+    simdgroup_load(A[3], data1+k+{24*N}, {N}, ulong2(0, 0));
+    simdgroup_load(B[0], data2+0+k*{N}, {N}, ulong2(0, 0));
+    simdgroup_load(B[1], data2+8+k*{N}, {N}, ulong2(0, 0));
+    simdgroup_load(B[2], data2+16+k*{N}, {N}, ulong2(0, 0));
+    simdgroup_load(B[3], data2+24+k*{N}, {N}, ulong2(0, 0));
 
     simdgroup_multiply_accumulate(acc[0][0], A[0], B[0], acc[0][0]);
     simdgroup_multiply_accumulate(acc[0][1], A[1], B[0], acc[0][1]);
@@ -86,18 +63,62 @@ kernel void test(device float *a, device const float *data1, device const float
     simdgroup_multiply_accumulate(acc[3][2], A[2], B[3], acc[3][2]);
     simdgroup_multiply_accumulate(acc[3][3], A[3], B[3], acc[3][3]);
   }}
-  for (uint i = 0; i < 4; i++) {{
-    for (uint j = 0; j < 4; j++) {{
-      simdgroup_store(acc[i][j], a, {N}, ulong2(pos_y+i*8, pos_x+j*8));
-    }}
-  }}
+  simdgroup_store(acc[0][0], a+{0+0*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[1][0], a+{8+0*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[2][0], a+{16+0*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[3][0], a+{24+0*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[0][1], a+{0+8*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[1][1], a+{8+8*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[2][1], a+{16+8*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[3][1], a+{24+8*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[0][2], a+{0+16*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[1][2], a+{8+16*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[2][2], a+{16+16*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[3][2], a+{24+16*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[0][3], a+{0+24*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[1][3], a+{8+24*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[2][3], a+{16+24*N}, {N}, ulong2(0, 0));
+  simdgroup_store(acc[3][3], a+{24+24*N}, {N}, ulong2(0, 0));
 }}""")
-tm = mb(lambda: prog([N*N//(2*4*4)], [4*32], a._cl, b._cl, c._cl))
+def timeit(fxn):
+  st = time.perf_counter()
+  et = fxn()
+  # NOTE: et doesn't contain the launch overhead
+  return time.perf_counter() - st
+tm = min([timeit(lambda: prog([32, N//(8*4), N//(8*4)], [32, 1, 4], a, b, c, wait=True)) for _ in range(20)])
 na = a.toCPU().reshape(N,N)
 comp = nb@nc
 if N <= 32:
   print(na)
   print(comp)
-print(f"{N*N:10d} {tm*1e-3:9.2f} us, would be {FLOPS/tm:.2f} GFLOPS matmul")
+print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul, {BW*1e-9/tm:.2f} GB/s")
 np.testing.assert_allclose(na, comp, atol=1e-3)
 
+import torch, torch.mps
+b = torch.from_numpy(nb).to('mps')
+c = torch.from_numpy(nc).to('mps')
+
+def torch_prog(b, c):
+  st = time.perf_counter()
+  a = b@c
+  torch.mps.synchronize()
+  return time.perf_counter() - st
+tm = min([torch_prog(b, c) for _ in range(20)])
+print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul in torch")
+
+from tinygrad.tensor import Tensor
+from tinygrad.jit import TinyJit
+from tinygrad.runtime.ops_metal import METAL
+b = Tensor(nb)
+c = Tensor(nc)
+# TODO: slowness without the JIT I suspect comes from a lack of a caching allocator
+@TinyJit
+def tiny_jit(b, c):
+  return (b@c).realize()
+def tiny_prog(b, c):
+  st = time.perf_counter()
+  a = tiny_jit(b, c)
+  METAL.synchronize()
+  return time.perf_counter() - st
+tm = min([tiny_prog(b, c) for _ in range(20)])
+print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul in tinygrad")

extra/gemm/simple_matmul.py

@@ -0,0 +1,5 @@
+from tinygrad.tensor import Tensor
+N = 1024
+a, b = Tensor.randn(N, N), Tensor.randn(N, N)
+c = (a.reshape(N, 1, N) * b.permute(1,0).reshape(1, N, N)).sum(axis=2)
+print((c.numpy() - (a.numpy() @ b.numpy())).mean())

extra/helpers.py

@@ -1,5 +1,39 @@
 import time
 
 class Timing(object):
-  def __enter__(self): self.st = time.monotonic_ns()
-  def __exit__(self, exc_type, exc_val, exc_tb): print(f"{(time.monotonic_ns()-self.st)*1e-6:.2f} ms")
+  def __init__(self, prefix="", on_exit=None, enabled=True): self.prefix, self.on_exit, self.enabled = prefix, on_exit, enabled
+  def __enter__(self): self.st = time.perf_counter_ns()
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    self.et = time.perf_counter_ns() - self.st
+    if self.enabled: print(f"{self.prefix}{self.et*1e-6:.2f} ms"+(self.on_exit(self.et) if self.on_exit else ""))
+
+def enable_early_exec():
+  import subprocess, multiprocessing
+  qin: multiprocessing.Queue = multiprocessing.Queue()
+  qout: multiprocessing.Queue = multiprocessing.Queue()
+  def _early_exec_process(qin, qout):
+    while 1:
+      path, inp = qin.get()
+      qout.put(subprocess.check_output(path, input=inp))
+  p = multiprocessing.Process(target=_early_exec_process, args=(qin, qout))
+  p.daemon = True
+  p.start()
+  def early_exec(x):
+    qin.put(x)
+    return qout.get()
+  return early_exec
+
+def proc(itermaker, q):
+  for x in itermaker(): q.put(x)
+  q.close()
+
+def cross_process(itermaker, maxsize=8):
+  # TODO: use cloudpickle for itermaker
+  import multiprocessing
+  q: multiprocessing.Queue = multiprocessing.Queue(maxsize)
+  p = multiprocessing.Process(target=proc, args=(itermaker, q))
+  p.daemon = True
+  p.start()
+
+  # TODO: write tests and handle exit case
+  while 1: yield q.get()

extra/introspection.py

@@ -3,7 +3,7 @@ import gc
 from tinygrad.helpers import prod
 from tinygrad.tensor import Tensor
 from tinygrad.lazy import LazyBuffer
-from tinygrad.runtime.ops_gpu import GPUBuffer
+from tinygrad.runtime.ops_gpu import CLBuffer
 from tinygrad.ops import GlobalCounters
 
 def print_objects():
@@ -11,7 +11,7 @@ def print_objects():
   tensors = [x for x in gc.get_objects() if isinstance(x, Tensor)]
   tensor_ram_used = sum([prod(x.shape)*4 for x in tensors])
   lazybuffers = [x for x in gc.get_objects() if isinstance(x, LazyBuffer)]
-  gpubuffers = [x for x in gc.get_objects() if isinstance(x, GPUBuffer)]
+  gpubuffers = [x for x in gc.get_objects() if isinstance(x, CLBuffer)]
   realized_buffers = [x.realized for x in lazybuffers if x.realized]
   gpubuffers_orphaned = [x for x in gpubuffers if x not in realized_buffers]
 

extra/kernel_search.py

@@ -5,7 +5,7 @@ import itertools
 from enum import Enum
 import numpy as np
 from tinygrad.ops import LazyOp, ReduceOps, BinaryOps, UnaryOps, MovementOps
-from tinygrad.shape import ShapeTracker, View, ZeroView
+from tinygrad.shape.shapetracker import ShapeTracker, View, ZeroView
 from tinygrad.runtime.ops_gpu import GPUBuffer, CLASTKernel
 from tinygrad.runtime.opencl import OSX_TIMING_RATIO
 from tinygrad.helpers import getenv, DEBUG

extra/onnx.py

@@ -46,7 +46,7 @@ def get_run_onnx(onnx_model):
   def attribute_parse(a):
     if a.type in [6,7]: return tuple([int(x) for x in a.ints])
     elif a.type == 4: return buffer_parse(a.t)  # TENSOR
-    elif a.type == 3: return str(a.s)
+    elif a.type == 3: return a.s.decode("utf-8")
     elif a.type == 2: return int(a.i)
     elif a.type == 1: return float(a.f)
     else: raise Exception(f"can't parse {a.type} {a}")

extra/onnx_ops.py

@@ -47,28 +47,45 @@ def LayerNormalization(x: Tensor, scale, bias, axis=-1, epsilon=1e-05, stash_typ
 def GroupNormalization(x: Tensor, scale: Tensor, bias: Tensor, num_groups, epsilon=1e-05):
   return x.reshape(x.shape[0], num_groups, -1).layernorm(axis=-1, eps=epsilon).mul(scale.unsqueeze(-1)).add(bias.unsqueeze(-1)).reshape(x.shape)
   
-# TODO: expand to N-D
-def _padding(X, pads=None, auto_pad="NOTSET"):
+# onnx: [x1_begin, x2_begin, ..., x1_end, x2_end, ...]
+# numpy.pad: ((x1_begin, x1_end), (x2_begin, x2_end), ...)
+def _format_padding(onnx_pads, ndims=None, axes=None):
+  if ndims is None: ndims = len(onnx_pads) // 2
+  if axes is None: axes = list(range(ndims))
+  num_axes = len(axes)
+  np_pads = [(0,0)] * ndims
+  for i in range(num_axes):
+    np_pads[axes[i]] = (onnx_pads[i], onnx_pads[i + num_axes])
+  return np_pads
+
+def _padding(X, pads=None, auto_pad="NOTSET", axes=None, constant_value=0.):
   assert auto_pad == "NOTSET"  # TODO: write this
-  if pads is not None:
-    return X.pad2d((pads[1], pads[3], pads[0], pads[2]))
-  else:
-    return X
+  if pads is None: return X
+  np_pads = _format_padding(pads, ndims=len(X.shape), axes=axes)
+  zero_padded = X.pad(tuple(np_pads))
+  constant_padder = Tensor(np.pad(np.zeros(X.shape), np_pads, constant_values=constant_value), dtype=X.dtype)
+  return zero_padded + constant_padder
+
+def Pad(x: Tensor, pads: Tensor, constant_value: Tensor=None, axes: Tensor=None, mode="constant"):
+  assert mode == "constant"
+  constant_value = 0. if constant_value is None else constant_value.numpy()
+  seq_pads = pads.numpy().astype(np.int32).tolist()
+  seq_axes = axes.numpy().astype(np.int32).tolist() if axes is not None else None
+  return _padding(x, seq_pads, axes=seq_axes, constant_value=constant_value)
 
 def AveragePool(X, kernel_shape, auto_pad="NOTSET", ceil_mode=0, count_include_pad=0, dilations=1, pads=None, strides=1):
-  # TODO: the padding shouldn't be counted in the average! this is causing a test failure
   assert ceil_mode == 0 and dilations == 1
-  padding_included = _padding(X, pads, auto_pad).avg_pool2d(kernel_shape, stride=strides)
+  pixel_axes = tuple(range(len(X.shape)))[-2:]
+  padding_included = _padding(X, pads, auto_pad, axes=pixel_axes).avg_pool2d(kernel_shape, stride=strides)
   if count_include_pad:
     return padding_included
   else:
-    div = _padding(Tensor.ones(*X.shape), pads, auto_pad).avg_pool2d(kernel_shape, stride=strides)
+    div = _padding(Tensor.ones(*X.shape), pads, auto_pad, axes=pixel_axes).avg_pool2d(kernel_shape, stride=strides)
     return padding_included / div
 
 def MaxPool(X, kernel_shape, auto_pad="NOTSET", ceil_mode=0, dilations=1, pads=None, storage_order=0, strides=1):
-  # TODO: the padding should be infinity, not 0!
   assert ceil_mode == 0 and storage_order == 0 and dilations == 1
-  return _padding(X, pads, auto_pad).max_pool2d(kernel_shape, stride=strides)
+  return _padding(X, pads, auto_pad, constant_value=-np.inf, axes=tuple(range(len(X.shape)))[-2:]).max_pool2d(kernel_shape, stride=strides)
 
 def Conv(X, W, B=None, auto_pad="NOTSET", dilations=1, group=1, kernel_shape=None, pads=None, strides=1):
   return X.conv2d(W, B, stride=strides, groups=group, dilation=dilations, padding=(pads[1], pads[3], pads[0], pads[2]) if pads is not None else 0)
@@ -126,11 +143,20 @@ Softmax = {1: Softmax_1, 13: Softmax_13}   # Softmax default axis changed
 def LogSoftmax(input, axis=-1): return input.log_softmax(axis)
 def Clip(input, min=-3.4e38, max=3.4e38): return input.clip(min, max)
 
+import math
+
+def Sin(x): return x.sin()
+def Cos(x): return x.cos()
+def Tan(x): return x.tan()
+def Cosh(x): return (math.e ** x + math.e ** -x) / 2
+def Sinh(x): return (math.e ** x - math.e ** -x) / 2
+def Tanh(x): return Sinh(x) / Cosh(x)
+
 def Less(x, y): return (x<y).numpy().astype(bool)
 def LessOrEqual(x, y): return (x<=y).numpy().astype(bool)
 def Greater(x, y): return (x>y).numpy().astype(bool)
 def GreaterOrEqual(x, y): return (x>=y).numpy().astype(bool)
-def Equal(x, y): return (x.eq(y)).numpy().astype(bool)
+def Equal(x, y): return (x==y).numpy().astype(bool)
 
 def Max(*data_0): return functools.reduce(Tensor.maximum, data_0)
 def Min(*data_0): return -functools.reduce(Tensor.maximum, [-x for x in data_0])

extra/rocm/.gitignore

@@ -0,0 +1,4 @@
+*.deb
+build
+src
+sniffer/sniff.so

extra/rocm/NOTES

@@ -0,0 +1,20 @@
+Built ROCT-Thunk-Interface (hsakmt)
+  hsakmt-roct-dev_5.4.4.99999-local_amd64.deb
+  note: installs to /opt/rocm
+Built ROCm-Device-Libs
+  Works with ROCM_PATH=/home/tiny/build/ROCm-Device-Libs/build/dist
+  rocm-device-libs_1.0.0.99999-local_amd64.deb
+Built ROCm-CompilerSupport (amd_comgr)
+  no deb, sudo make install to /usr/local
+Built ROCR-Runtime
+  hsa-rocr_1.8.0-local_amd64.deb
+  hsa-rocr-dev_1.8.0-local_amd64.deb
+Built ROCm-OpenCL-Runtime
+  rocm-ocl-icd_2.0.0-local_amd64.deb
+  ISSUE: these depend on "comgr"
+  rocm-opencl_2.0.0-local_amd64.deb
+  rocm-opencl-dev_2.0.0-local_amd64.deb
+  Did sudo make install
+
+
+

extra/rocm/kernel_crashes/dump

@@ -0,0 +1,41 @@
+# run two "rocm-bandwidth-test" in a loop
+# amdgpu-6.0.5-1581431.20.04
+# fixed in kernel 6.2.14
+
+[   72.153646] RIP: 0010:pm_send_runlist+0x4a/0x630 [amdgpu]
+[   72.153815] Code: 30 65 48 8b 04 25 28 00 00 00 48 89 45 d0 31 c0 80 fb 01 0f 87 aa 9d 49 00 83 e3 01 0f 85 1c 05 00 00 49 8b 3f b8 01 00 00 00 <48> 8b 97 30 01 00 00 44 8b b7 6c 01 00 00 8b 9f 70 01 00 00 8b 8a
+[   72.153900] RSP: 0018:ffffb48445c03c30 EFLAGS: 00010246
+[   72.153928] RAX: 0000000000000001 RBX: 0000000000000000 RCX: 0000000000000000
+[   72.153962] RDX: 000000000000007b RSI: ffff9395e1562558 RDI: 0000000000000000
+[   72.153996] RBP: ffffb48445c03cb8 R08: 0000000000000000 R09: 0000000000000001
+[   72.154030] R10: ffff9395c900d840 R11: 0000000000000000 R12: 0000000000000000
+[   72.154065] R13: ffff9395c9e00400 R14: 0000000000000001 R15: ffff9395e15624e0
+[   72.154099] FS:  00007f345c6463c0(0000) GS:ffff93a4aee80000(0000) knlGS:0000000000000000
+[   72.154137] CS:  0010 DS: 0000 ES: 0000 CR0: 0000000080050033
+[   72.154165] CR2: 0000000000000130 CR3: 0000000112840000 CR4: 0000000000750ee0
+[   72.154201] PKRU: 55555554
+[   72.154215] Call Trace:
+[   72.154230]  <TASK>
+[   72.154244]  map_queues_cpsch+0x75/0xc0 [amdgpu]
+[   72.154365]  debug_map_and_unlock+0x51/0x90 [amdgpu]
+[   72.154480]  debug_refresh_runlist+0x1f/0x30 [amdgpu]
+[   72.154591]  kfd_dbg_runtime_disable+0x13c/0x240 [amdgpu]
+[   72.154705]  kfd_ioctl_dbg_set_debug_trap+0x69d/0x8b0 [amdgpu]
+[   72.154820]  kfd_ioctl+0x24a/0x5b0 [amdgpu]
+[   72.154925]  ? kfd_ioctl_create_queue+0x770/0x770 [amdgpu]
+[   72.155035]  ? syscall_exit_to_user_mode+0x27/0x50
+[   72.155061]  ? exit_to_user_mode_prepare+0x3d/0x1c0
+[   72.155088]  __x64_sys_ioctl+0x95/0xd0
+[   72.155109]  do_syscall_64+0x5c/0xc0
+[   72.155128]  ? syscall_exit_to_user_mode+0x27/0x50
+[   72.155151]  ? do_syscall_64+0x69/0xc0
+[   72.155172]  entry_SYSCALL_64_after_hwframe+0x61/0xcb
+[   72.155198] RIP: 0033:0x7f345c7f63ab
+[   72.155218] Code: 0f 1e fa 48 8b 05 e5 7a 0d 00 64 c7 00 26 00 00 00 48 c7 c0 ff ff ff ff c3 66 0f 1f 44 00 00 f3 0f 1e fa b8 10 00 00 00 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 8b 0d b5 7a 0d 00 f7 d8 64 89 01 48
+[   72.155301] RSP: 002b:00007ffc97cc89f8 EFLAGS: 00000246 ORIG_RAX: 0000000000000010
+[   72.155339] RAX: ffffffffffffffda RBX: 00007ffc97cc8a30 RCX: 00007f345c7f63ab
+[   72.155375] RDX: 00007ffc97cc8a30 RSI: 00000000c0284b82 RDI: 0000000000000003
+[   72.155411] RBP: 00000000c0284b82 R08: 0000000000000000 R09: 0000000000000000
+[   72.155447] R10: 00007f345cd4ddb0 R11: 0000000000000246 R12: 00007ffc97cc8a30
+[   72.155481] R13: 0000000000000003 R14: 00007ffc97cc8d20 R15: 0000000000000000
+[   72.155517]  </TASK>

extra/rocm/kernel_crashes/dump2

@@ -0,0 +1,41 @@
+# run two tinygrad matrix example in a loop
+# amdgpu-6.0.5-1581431.20.04
+# NOT fixed in kernel 6.2.14
+
+[  553.016624] gmc_v11_0_process_interrupt: 30 callbacks suppressed
+[  553.016631] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:9 pasid:32770, for process python3 pid 10001 thread python3 pid 10001)
+[  553.016790] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f0000000000 from client 10
+[  553.016892] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00901A30
+[  553.016974] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
+[  553.017051] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x0
+[  553.017111] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
+[  553.017173] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
+[  553.017238] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
+[  553.017300] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
+[  553.123921] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
+[  553.124153] amdgpu: failed to add hardware queue to MES, doorbell=0x1a16
+[  553.124195] amdgpu: MES might be in unrecoverable state, issue a GPU reset
+[  553.124237] amdgpu: Failed to restore queue 2
+[  553.124266] amdgpu: Failed to restore process queues
+[  553.124270] amdgpu: Failed to evict queue 3
+[  553.124297] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
+
+# alternative crash in kernel 6.2.14
+
+[  151.097948] gmc_v11_0_process_interrupt: 30 callbacks suppressed
+[  151.097953] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:8 pasid:32771, for process python3 pid 7525 thread python3 pid 7525)
+[  151.097993] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f0000000000 from client 10
+[  151.098008] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00801A30
+[  151.098020] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
+[  151.098032] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x0
+[  151.098042] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
+[  151.098052] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
+[  151.098062] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
+[  151.098071] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
+[  151.209517] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
+[  151.209724] amdgpu: failed to add hardware queue to MES, doorbell=0x1002
+[  151.209734] amdgpu: MES might be in unrecoverable state, issue a GPU reset
+[  151.209743] amdgpu: Failed to restore queue 1
+[  151.209751] amdgpu: Failed to restore process queues
+[  151.209759] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
+[  151.209858] amdgpu 0000:0b:00.0: amdgpu: GPU reset begin!

extra/rocm/kernel_crashes/dump3

@@ -0,0 +1,20 @@
+# two tinygrad + two bandwidth test
+# RDNA2, driver 6.0.5
+# recovered from this!
+
+[  136.971209] gmc_v10_0_process_interrupt: 39 callbacks suppressed
+[  136.971218] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:11 pasid:32773, for process rocm-bandwidth- pid 20281 thread rocm-bandwidth- pid 20281)
+[  136.971228] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f5c2b800000 from client 0x1b (UTCL2)
+[  136.971232] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00B01A31
+[  136.971233] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
+[  136.971235] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x1
+[  136.971236] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
+[  136.971236] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
+[  136.971237] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
+[  136.971238] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
+...
+[  136.993979] amdgpu 0000:0b:00.0: amdgpu: IH ring buffer overflow (0x000BE5A0, 0x0003C480, 0x0003E5C0)
+[  138.209072] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004000 flags=0x0000]
+[  138.209078] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004d80 flags=0x0000]
+[  138.209081] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005000 flags=0x0000]
+[  138.209084] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005d80 flags=0x0000]

extra/rocm/kernel_crashes/dump4

@@ -0,0 +1,33 @@
+# ROCK-Kernel-Driver 0b579de9622f5c93021dcb7927d13926313740a2
+# non fatal "crash"
+
+[  127.418045] ------------[ cut here ]------------
+[  127.418046] User pages unexpectedly invalid
+[  127.418056] WARNING: CPU: 16 PID: 260 at drivers/gpu/drm/amd/amdgpu/amdgpu_amdkfd_gpuvm.c:3000 amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
+[  127.418235] Modules linked in: rfcomm cmac algif_hash algif_skcipher af_alg bnep nls_iso8859_1 iwlmvm mac80211 intel_rapl_msr intel_rapl_common edac_mce_amd snd_hda_codec_realtek snd_hda_codec_generic snd_hda_codec_hdmi kvm_amd binfmt_misc snd_hda_intel snd_intel_dspcfg kvm libarc4 snd_intel_sdw_acpi snd_hda_codec btusb iwlwifi btrtl snd_hda_core btbcm btintel irqbypass btmtk snd_hwdep crct10dif_pclmul snd_pcm polyval_clmulni bluetooth snd_seq_midi snd_seq_midi_event snd_rawmidi snd_seq polyval_generic cfg80211 ghash_clmulni_intel eeepc_wmi snd_seq_device snd_timer aesni_intel asus_wmi ecdh_generic snd platform_profile crypto_simd ledtrig_audio cryptd ecc ccp soundcore sparse_keymap rapl k10temp wmi_bmof mac_hid sch_fq_codel msr parport_pc ppdev lp parport ramoops pstore_blk efi_pstore reed_solomon pstore_zone ip_tables x_tables autofs4 amdgpu hid_generic usbhid hid i2c_algo_bit drm_ttm_helper ttm video iommu_v2 drm_buddy gpu_sched drm_display_helper drm_kms_helper syscopyarea
+[  127.418276]  sysfillrect sysimgblt fb_sys_fops drm nvme nvme_core cec r8169 ahci crc32_pclmul rc_core i2c_piix4 xhci_pci libahci nvme_common xhci_pci_renesas realtek wmi
+[  127.418284] CPU: 16 PID: 260 Comm: kworker/16:1 Tainted: G        W          6.0.0 #4
+[  127.418286] Hardware name: System manufacturer System Product Name/TUF GAMING X570-PLUS (WI-FI), BIOS 3603 03/20/2021
+[  127.418287] Workqueue: events amdgpu_amdkfd_restore_userptr_worker [amdgpu]
+[  127.418455] RIP: 0010:amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
+[  127.418601] Code: ff e8 2b 8a 96 d1 e9 66 fe ff ff 48 c7 c7 40 4f f5 c0 e8 56 7b 8a d1 0f 0b e9 2e ff ff ff 48 c7 c7 d8 d0 ed c0 e8 43 7b 8a d1 <0f> 0b e9 0a fe ff ff 4c 89 ef e8 f8 89 96 d1 e9 cb fd ff ff e8 ce
+[  127.418603] RSP: 0018:ffffb36740a83dc8 EFLAGS: 00010282
+[  127.418604] RAX: 0000000000000000 RBX: ffff9d159ee9df30 RCX: 0000000000000027
+[  127.418605] RDX: 0000000000000027 RSI: ffffb36740a83c88 RDI: ffff9d242a220568
+[  127.418606] RBP: ffffb36740a83e58 R08: ffff9d242a220560 R09: 0000000000000001
+[  127.418607] R10: 0000000000000001 R11: 0000000000000020 R12: ffff9d159ee9df98
+[  127.418607] R13: ffff9d159ee9df70 R14: ffff9d159ee9dee0 R15: ffff9d159ee9dee0
+[  127.418608] FS:  0000000000000000(0000) GS:ffff9d242a200000(0000) knlGS:0000000000000000
+[  127.418609] CS:  0010 DS: 0000 ES: 0000 CR0: 0000000080050033
+[  127.418610] CR2: 00007fd5d4715000 CR3: 0000000120ffe000 CR4: 0000000000750ee0
+[  127.418611] PKRU: 55555554
+[  127.418611] Call Trace:
+[  127.418612]  <TASK>
+[  127.418613]  process_one_work+0x21f/0x3f0
+[  127.418615]  worker_thread+0x4a/0x3c0
+[  127.418617]  ? process_one_work+0x3f0/0x3f0
+[  127.418618]  kthread+0xf0/0x120
+[  127.418619]  ? kthread_complete_and_exit+0x20/0x20
+[  127.418620]  ret_from_fork+0x22/0x30
+[  127.418622]  </TASK>
+[  127.418623] ---[ end trace 0000000000000000 ]---

extra/rocm/rdna3/asm.py

@@ -0,0 +1,57 @@
+import numpy as np
+import pathlib
+from hexdump import hexdump
+from tinygrad.helpers import colored
+from extra.helpers import enable_early_exec
+early_exec = enable_early_exec()
+
+from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer, ROCM_LLVM_PATH
+
+ENABLE_NON_ASM = False
+
+if ENABLE_NON_ASM:
+  buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
+  prg_empty = CLProgram("code", "__kernel void code(__global float *a) { a[0] = 1; }")
+  asm_real = prg_empty.binary()
+  with open("/tmp/cc.elf", "wb") as f:
+    f.write(asm_real)
+  prg_empty([1], [1], buf, wait=True)
+  print(buf.toCPU())
+
+print(colored("creating CLBuffer", "green"))
+buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
+code = open(pathlib.Path(__file__).parent / "prog.s", "r").read()
+
+gen = []
+FLOPS = 0
+for j in range(4):
+  for i in range(0, 251, 6):
+    #gen.append(f"v_dual_fmac_f32 v{i+0}, v{i+1}, v{i+2} :: v_dual_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
+    #FLOPS += 4
+    gen.append(f"v_dual_dot2acc_f32_f16 v{i+0}, v{i+1}, v{i+2} :: v_dual_dot2acc_f32_f16 v{i+3}, v{i+4}, v{i+5}")
+    FLOPS += 8
+code = code.replace("// FLOPS", '\n'.join(gen))
+print(code)
+
+
+# fix: COMGR failed to get code object ISA name. set triple to 'amdgcn-amd-amdhsa'
+
+object = early_exec(([ROCM_LLVM_PATH / "llvm-mc", '--arch=amdgcn', '--mcpu=gfx1100', '--triple=amdgcn-amd-amdhsa', '--filetype=obj', '-'], code.encode("utf-8")))
+asm = early_exec(([ROCM_LLVM_PATH / "ld.lld", "/dev/stdin", "-o", "/dev/stdout", "--pie"], object))
+
+with open("/tmp/cc2.o", "wb") as f:
+  f.write(object)
+with open("/tmp/cc2.elf", "wb") as f:
+  f.write(asm)
+
+print(colored("creating CLProgram", "green"))
+prg = CLProgram("code", asm, binary=True)
+
+print(colored("running program", "green"))
+FLOPS *= 100000*1024*1024  # loop * global_size
+for i in range(3):
+  tm = prg([1024, 1024], [256, 1], buf, wait=True)
+  print(f"ran in {tm*1e3:.2f} ms, {FLOPS/(tm*1e9):.2f} GFLOPS")
+
+print(colored("transferring buffer", "green"))
+print(buf.toCPU())

extra/rocm/rdna3/prog.s

@@ -0,0 +1,80 @@
+.global _start
+_start:
+.rodata
+.align 0x10
+.global code.kd
+.type code.kd,STT_OBJECT
+# amd_kernel_code_t (must be at 0x440 for kernel_code_entry_byte_offset to be right)
+code.kd:
+# amd_kernel_..., amd_machine_...
+.long 0,0,0,0
+# kernel_code_entry_byte_offset, kernel_code_prefetch_byte_offset
+.long 0x00000bc0,0x00000000,0x00000000,0x00000000
+# kernel_code_prefetch_byte_size, max_scratch_backing_memory_byte_size
+.long 0,0,0,0
+# compute_pgm_rsrc1, compute_pgm_rsrc2, kernel_code_properties, workitem_private_segment_byte_size
+.long 0x60af0000,0x0000009e,0x00000408,0x00000000
+# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_32 | AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_16_64
+# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_DX10_CLAMP | AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_IEEE_MODE
+# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT = 0xF
+# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X
+# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_KERNARG_SEGMENT_PTR = 1
+# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_RESERVED1 = 1
+.text
+.global code
+.type code,STT_FUNC
+code:
+# https://llvm.org/docs/AMDGPUUsage.html#initial-kernel-execution-state
+# s[0:1] contains the kernarg_address
+# TODO: can we use s[2:3] if this was really a wave since we only alloced 2 SGPRs?
+s_load_b64 s[2:3], s[0:1], null
+
+s_mov_b32 s8, 0
+loop:
+s_addk_i32 s8, 1
+s_cmp_eq_u32 s8, 100000
+// FLOPS
+s_cbranch_scc0 loop
+
+# wait for the s_load_b64
+s_waitcnt lgkmcnt(0)
+
+v_dual_mov_b32 v0, 4 :: v_dual_mov_b32 v1, 2.0
+global_store_b32 v0, v1, s[2:3]
+
+# Deallocate all VGPRs for this wave. Use only when next instruction is S_ENDPGM.
+s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
+s_endpgm
+s_code_end
+
+.amdgpu_metadata
+amdhsa.kernels:
+  - .args:
+      - .address_space:  global
+        .name:           a
+        .offset:         0
+        .size:           8
+        .type_name:      'float*'
+        .value_kind:     global_buffer
+    .group_segment_fixed_size: 0
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 8
+    .language:       OpenCL C
+    .language_version:
+      - 1
+      - 2
+    .max_flat_workgroup_size: 256
+    .name:           code
+    .private_segment_fixed_size: 0
+    .sgpr_count:     2
+    .sgpr_spill_count: 0
+    .symbol:         code.kd
+    .uses_dynamic_stack: false
+    .vgpr_count:     256
+    .vgpr_spill_count: 0
+    .wavefront_size: 32
+amdhsa.target:   amdgcn-amd-amdhsa--gfx1100
+amdhsa.version:
+  - 1
+  - 2
+.end_amdgpu_metadata

extra/rocm/rocm_clone.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+mkdir -p src
+cd src
+git clone https://github.com/RadeonOpenCompute/ROCT-Thunk-Interface.git -b rocm-5.5.0
+git clone https://github.com/RadeonOpenCompute/ROCm-Device-Libs.git -b rocm-5.5.0
+git clone https://github.com/RadeonOpenCompute/llvm-project.git -b rocm-5.5.0 --depth 1
+git clone https://github.com/RadeonOpenCompute/ROCR-Runtime.git -b rocm-5.5.0
+git clone https://github.com/ROCm-Developer-Tools/ROCclr.git -b rocm-5.5.0
+git clone https://github.com/RadeonOpenCompute/ROCm-CompilerSupport.git -b rocm-5.5.0
+git clone https://github.com/RadeonOpenCompute/ROCm-OpenCL-Runtime.git -b rocm-5.5.0
+cd ../

extra/rocm/rocm_from_scratch.sh

@@ -0,0 +1,69 @@
+#!/bin/bash
+mkdir -p build/debs
+cd build
+
+# ROCT-Thunk-Interface (hsakmt)
+if [ ! -f debs/hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ]
+then
+  mkdir -p ROCT-Thunk-Interface
+  cd ROCT-Thunk-Interface
+  cmake ../../src/ROCT-Thunk-Interface
+  make -j32 package
+  cp hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ../debs
+  cd ../
+fi
+
+
+# build custom LLVM
+if [ ! -f llvm-project/bin/clang ]
+then
+  mkdir -p llvm-project
+  cd llvm-project
+  cmake -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_PROJECTS="llvm;clang;lld" -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" ../../src/llvm-project/llvm
+  make -j32
+  cd ..
+fi
+
+# use custom LLVM
+export PATH="$PWD/llvm-project/bin:$PATH"
+
+# ROCm-Device-Libs
+if [ ! -f debs/rocm-device-libs_1.0.0.99999-local_amd64.deb ]
+then
+  mkdir -p ROCm-Device-Libs
+  cd ROCm-Device-Libs
+  cmake ../../src/ROCm-Device-Libs
+  make -j32 package
+  cp rocm-device-libs_1.0.0.99999-local_amd64.deb ../debs
+  cd ../
+fi
+
+# ROCR-Runtime
+if [ ! -f debs/hsa-rocr_1.8.0-local_amd64.deb ]
+then
+  mkdir -p ROCR-Runtime
+  cd ROCR-Runtime
+  cmake ../../src/ROCR-Runtime/src
+  make -j32 package
+  cp hsa-rocr_1.8.0-local_amd64.deb ../debs
+  cp hsa-rocr-dev_1.8.0-local_amd64.deb ../debs
+  cd ../
+fi
+
+# ROCm-OpenCL-Runtime (needs ROCclr)
+if [ ! -f debs/rocm-opencl_2.0.0-local_amd64.deb ]
+then
+  mkdir -p ROCm-OpenCL-Runtime
+  cd ROCm-OpenCL-Runtime
+  cmake ../../src/ROCm-OpenCL-Runtime
+  make -j32 package
+  cp rocm-opencl_2.0.0-local_amd64.deb ../debs
+  cp rocm-opencl-dev_2.0.0-local_amd64.deb ../debs
+  cp rocm-ocl-icd_2.0.0-local_amd64.deb ../debs
+fi
+
+# ROCm-CompilerSupport (broken)
+#mkdir -p ROCm-CompilerSupport
+#cd ROCm-CompilerSupport
+#cmake ../../src/ROCm-CompilerSupport/lib/comgr
+#make -j32

extra/rocm/rocm_setup.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+rm amdgpu-install_5.5.50500-1_all.deb
+wget https://repo.radeon.com/amdgpu-install/5.5/ubuntu/$(lsb_release -cs)/amdgpu-install_5.5.50500-1_all.deb
+sudo dpkg -i amdgpu-install_5.5.50500-1_all.deb
+sudo apt-get update
+
+# kernel driver
+sudo apt-get install amdgpu-dkms
+
+# for opencl
+sudo apt-get install rocm-opencl-runtime
+
+# for HIP
+sudo apt-get install hip-runtime-amd rocm-device-libs hip-dev

extra/rocm/sniffer/build.sh

@@ -0,0 +1,11 @@
+#!/bin/bash -e
+clang sniff.cc -Werror -shared -fPIC -I../src/ -I../src/ROCT-Thunk-Interface/include -I../src/ROCm-Device-Libs/ockl/inc -o sniff.so -lstdc++
+#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so /home/tiny/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
+#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so $HOME/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
+#AMD_LOG_LEVEL=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
+DEBUG=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
+#AMD_LOG_LEVEL=5 HSAKMT_DEBUG_LEVEL=7 DEBUG=5 LD_PRELOAD=$PWD/sniff.so strace -F python3 ../rdna3/asm.py
+#LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
+#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so FORWARD_ONLY=1 DEBUG=2 python3 ../../../test/test_ops.py TestOps.test_add
+#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 0 -d 1 -m 1
+#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 1 -d 2 -m 1

extra/rocm/sniffer/sniff.cc

@@ -0,0 +1,282 @@
+// template copied from https://github.com/geohot/cuda_ioctl_sniffer/blob/master/sniff.cc
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <dlfcn.h>
+#include <signal.h>
+#include <ucontext.h>
+
+#include <sys/mman.h>
+
+// includes from the ROCm sources
+#include <linux/kfd_ioctl.h>
+#include <hsa.h>
+#include <amd_hsa_kernel_code.h>
+#include <ROCR-Runtime/src/core/inc/sdma_registers.h>
+using namespace rocr::AMD;
+
+#include <string>
+#include <map>
+std::map<int, std::string> files;
+std::map<uint64_t, uint64_t> ring_base_addresses;
+
+#define D(args...) fprintf(stderr, args)
+
+uint64_t doorbell_offset = -1;
+std::map<uint64_t, int> queue_types;
+
+void hexdump(void *d, int l) {
+  for (int i = 0; i < l; i++) {
+    if (i%0x10 == 0 && i != 0) printf("\n");
+    if (i%0x10 == 8) printf(" ");
+    if (i%0x10 == 0) printf("%8X: ", i);
+    printf("%2.2X ", ((uint8_t*)d)[i]);
+  }
+  printf("\n");
+}
+
+extern "C" {
+
+// https://defuse.ca/online-x86-assembler.htm#disassembly2
+static void handler(int sig, siginfo_t *si, void *unused) {
+  ucontext_t *u = (ucontext_t *)unused;
+  uint8_t *rip = (uint8_t*)u->uc_mcontext.gregs[REG_RIP];
+
+  int store_size = 0;
+  uint64_t value;
+  if (rip[0] == 0x48 && rip[1] == 0x89 && rip[2] == 0x30) {
+    // 0:  48 89 30                mov    QWORD PTR [rax],rsi
+    store_size = 8;
+    value = u->uc_mcontext.gregs[REG_RSI];
+    u->uc_mcontext.gregs[REG_RIP] += 3;
+  } else if (rip[0] == 0x4c && rip[1] == 0x89 && rip[2] == 0x28) {
+    // 0:  4c 89 28                mov    QWORD PTR [rax],r13
+    store_size = 8;
+    value = u->uc_mcontext.gregs[REG_R13];
+    u->uc_mcontext.gregs[REG_RIP] += 3;
+  } else {
+    D("segfault %02X %02X %02X %02X %02X %02X %02X %02X rip: %p addr: %p\n", rip[0], rip[1], rip[2], rip[3], rip[4], rip[5], rip[6], rip[7], rip, si->si_addr);
+    D("rax: %llx rcx: %llx rdx: %llx rsi: %llx rbx: %llx\n", u->uc_mcontext.gregs[REG_RAX], u->uc_mcontext.gregs[REG_RCX], u->uc_mcontext.gregs[REG_RDX], u->uc_mcontext.gregs[REG_RSI], u->uc_mcontext.gregs[REG_RBX]);
+    exit(-1);
+  }
+
+  uint64_t ring_base_address = ring_base_addresses[((uint64_t)si->si_addr)&0xFFF];
+  int queue_type = queue_types[((uint64_t)si->si_addr)&0xFFF];
+  D("%16p: \u001b[31mDING DONG\u001b[0m (queue_type %d) store(%d): 0x%8lx -> %p ring_base_address:0x%lx\n", rip, queue_type, store_size, value, si->si_addr, ring_base_address);
+
+  if (queue_type == KFD_IOC_QUEUE_TYPE_SDMA) {
+    uint8_t *sdma_ptr = (uint8_t*)(ring_base_address);
+    while (sdma_ptr < ((uint8_t*)(ring_base_address)+value)) {
+      D("0x%3lx: ", sdma_ptr-(uint8_t*)(ring_base_address));
+      if (sdma_ptr[0] == SDMA_OP_TIMESTAMP) {
+        D("SDMA_PKT_TIMESTAMP\n");
+        sdma_ptr += sizeof(SDMA_PKT_TIMESTAMP);
+      } else if (sdma_ptr[0] == SDMA_OP_GCR) {
+        D("SDMA_PKT_GCR\n");
+        sdma_ptr += sizeof(SDMA_PKT_GCR);
+      } else if (sdma_ptr[0] == SDMA_OP_ATOMIC) {
+        D("SDMA_PKT_ATOMIC\n");
+        sdma_ptr += sizeof(SDMA_PKT_ATOMIC);
+      } else if (sdma_ptr[0] == SDMA_OP_FENCE) {
+        D("SDMA_PKT_FENCE\n");
+        sdma_ptr += sizeof(SDMA_PKT_FENCE);
+      } else if (sdma_ptr[0] == SDMA_OP_TRAP) {
+        D("SDMA_PKT_TRAP\n");
+        sdma_ptr += sizeof(SDMA_PKT_TRAP);
+      } else if (sdma_ptr[0] == SDMA_OP_COPY && sdma_ptr[1] == SDMA_SUBOP_COPY_LINEAR) {
+        SDMA_PKT_COPY_LINEAR *pkt = (SDMA_PKT_COPY_LINEAR *)sdma_ptr;
+        D("SDMA_PKT_COPY_LINEAR: count:0x%x src:0x%lx dst:0x%lx\n", pkt->COUNT_UNION.count+1,
+          (uint64_t)pkt->SRC_ADDR_LO_UNION.src_addr_31_0 | ((uint64_t)pkt->SRC_ADDR_HI_UNION.src_addr_63_32 << 32),
+          (uint64_t)pkt->DST_ADDR_LO_UNION.dst_addr_31_0 | ((uint64_t)pkt->DST_ADDR_HI_UNION.dst_addr_63_32 << 32)
+        );
+        sdma_ptr += sizeof(SDMA_PKT_COPY_LINEAR);
+      } else {
+        D("unhandled packet type %d %d, exiting\n", sdma_ptr[0], sdma_ptr[1]);
+        break;
+      }
+    }
+
+    //hexdump((void*)(ring_base_address), 0x100);
+  } else if (queue_type == KFD_IOC_QUEUE_TYPE_COMPUTE_AQL) {
+    hsa_kernel_dispatch_packet_t *pkt = (hsa_kernel_dispatch_packet_t *)(ring_base_address+value*0x40);
+    if ((pkt->header&0xFF) == HSA_PACKET_TYPE_KERNEL_DISPATCH) {
+      D("HSA_PACKET_TYPE_KERNEL_DISPATCH -- setup:%d workgroup[%d, %d, %d] grid[%d, %d, %d] kernel_object:0x%lx kernarg_address:%p\n", pkt->setup, pkt->workgroup_size_x, pkt->workgroup_size_y, pkt->workgroup_size_z, pkt->grid_size_x, pkt->grid_size_y, pkt->grid_size_z, pkt->kernel_object, pkt->kernarg_address);
+      amd_kernel_code_t *code = (amd_kernel_code_t *)pkt->kernel_object;
+      D("kernel_code_entry_byte_offset:%lx\n", code->kernel_code_entry_byte_offset);
+      uint32_t *kernel_code = (uint32_t*)(pkt->kernel_object + code->kernel_code_entry_byte_offset);
+      int code_len = 0;
+      while (kernel_code[code_len] != 0xbf9f0000 && kernel_code[code_len] != 0) code_len++;
+      hexdump(kernel_code, code_len*4);
+      /*FILE *f = fopen("/tmp/kernel_code", "wb");
+      fwrite(kernel_code, 4, code_len, f);
+      fclose(f);
+      system("python -c 'print(\" \".join([(\"0x%02X\"%x) for x in open(\"/tmp/kernel_code\", \"rb\").read()]))' | ../build/llvm-project/bin/llvm-mc --disassemble --arch=amdgcn --mcpu=gfx1100 --show-encoding");*/
+      D("kernargs (kernarg_segment_byte_size:0x%lx)\n", code->kernarg_segment_byte_size);
+      // get length
+      int i;
+      for (i = 0; i < 0x400; i+=0x10) {
+        if (memcmp((void*)((uint64_t)pkt->kernarg_address+i), "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", 0x10) == 0) break;
+      }
+      hexdump((void*)pkt->kernarg_address, i+0x10);
+    } else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_BARRIER_AND) {
+      hsa_barrier_and_packet_t *pkt_and = (hsa_barrier_and_packet_t *)(ring_base_address+value*0x40);
+      D("HSA_PACKET_TYPE_BARRIER_AND completion_signal:0x%lx\n", pkt_and->completion_signal.handle);
+      //hexdump((void*)(ring_base_address+value*0x40), 0x40);
+    } else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_VENDOR_SPECIFIC) {
+      D("HSA_PACKET_TYPE_VENDOR_SPECIFIC\n");
+      hexdump((void*)(ring_base_address+value*0x40), 0x40);
+    } else {
+      hexdump((void*)(ring_base_address+value*0x40), 0x40);
+    }
+  }
+
+  mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_READ | PROT_WRITE);
+  if (store_size == 8) {
+    *(volatile uint64_t*)(si->si_addr) = value;
+  } else if (store_size == 4) {
+    *(volatile uint32_t*)(si->si_addr) = value;
+  } else if (store_size == 2) {
+    *(volatile uint16_t*)(si->si_addr) = value;
+  } else {
+    D("store size not supported\n");
+    exit(-1);
+  }
+  mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_NONE);
+}
+
+void register_sigsegv_handler() {
+  struct sigaction sa = {0};
+  sa.sa_flags = SA_SIGINFO;
+  sigemptyset(&sa.sa_mask);
+  sa.sa_sigaction = handler;
+  if (sigaction(SIGSEGV, &sa, NULL) == -1) {
+    D("ERROR: failed to register sigsegv handler");
+    exit(-1);
+  }
+  // NOTE: python (or ocl runtime?) blocks the SIGSEGV signal
+  sigset_t x;
+  sigemptyset(&x);
+  sigaddset(&x, SIGSEGV);
+  sigprocmask(SIG_UNBLOCK, &x, NULL);
+}
+
+int (*my_open)(const char *pathname, int flags, mode_t mode);
+#undef open
+int open(const char *pathname, int flags, mode_t mode) {
+  if (my_open == NULL) my_open = reinterpret_cast<decltype(my_open)>(dlsym(RTLD_NEXT, "open"));
+  int ret = my_open(pathname, flags, mode);
+  //D("open %s (0o%o) = %d\n", pathname, flags, ret);
+  files[ret] = pathname;
+  return ret;
+}
+
+
+int (*my_open64)(const char *pathname, int flags, mode_t mode);
+#undef open
+int open64(const char *pathname, int flags, mode_t mode) {
+  if (my_open64 == NULL) my_open64 = reinterpret_cast<decltype(my_open64)>(dlsym(RTLD_NEXT, "open64"));
+  int ret = my_open64(pathname, flags, mode);
+  //D("open %s (0o%o) = %d\n", pathname, flags, ret);
+  files[ret] = pathname;
+  return ret;
+}
+
+void *(*my_mmap)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
+#undef mmap
+void *mmap(void *addr, size_t length, int prot, int flags, int fd, off_t offset) {
+  if (my_mmap == NULL) my_mmap = reinterpret_cast<decltype(my_mmap)>(dlsym(RTLD_NEXT, "mmap"));
+  void *ret = my_mmap(addr, length, prot, flags, fd, offset);
+
+  if (doorbell_offset != -1 && offset == doorbell_offset) {
+    D("HIDDEN DOORBELL %p, handled by %p\n", addr, handler);
+    register_sigsegv_handler();
+    mprotect(addr, length, PROT_NONE);
+  }
+
+  if (fd != -1) D("mmapped %p (target %p) with flags 0x%x length 0x%zx fd %d %s offset 0x%lx\n", ret, addr, flags, length, fd, files[fd].c_str(), offset);
+  return ret;
+}
+
+void *(*my_mmap64)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
+#undef mmap64
+void *mmap64(void *addr, size_t length, int prot, int flags, int fd, off_t offset) { return mmap(addr, length, prot, flags, fd, offset); }
+
+int ioctl_num = 1;
+int (*my_ioctl)(int filedes, unsigned long request, void *argp) = NULL;
+#undef ioctl
+int ioctl(int filedes, unsigned long request, void *argp) {
+  if (my_ioctl == NULL) my_ioctl = reinterpret_cast<decltype(my_ioctl)>(dlsym(RTLD_NEXT, "ioctl"));
+  int ret = 0;
+  ret = my_ioctl(filedes, request, argp);
+  if (!files.count(filedes)) return ret;
+
+  uint8_t type = (request >> 8) & 0xFF;
+  uint8_t nr = (request >> 0) & 0xFF;
+  uint16_t size = (request >> 16) & 0xFFF;
+
+  D("%3d: %d = %3d(%20s) 0x%3x ", ioctl_num, ret, filedes, files[filedes].c_str(), size);
+
+  if (request == AMDKFD_IOC_SET_EVENT) {
+    kfd_ioctl_set_event_args *args = (kfd_ioctl_set_event_args *)argp;
+    D("AMDKFD_IOC_SET_EVENT event_id:%d", args->event_id);
+  } else if (request == AMDKFD_IOC_ALLOC_MEMORY_OF_GPU) {
+    kfd_ioctl_alloc_memory_of_gpu_args *args = (kfd_ioctl_alloc_memory_of_gpu_args *)argp;
+    D("AMDKFD_IOC_ALLOC_MEMORY_OF_GPU va_addr:0x%llx size:0x%llx handle:%llX gpu_id:0x%x", args->va_addr, args->size, args->handle, args->gpu_id);
+  } else if (request == AMDKFD_IOC_MAP_MEMORY_TO_GPU) {
+    kfd_ioctl_map_memory_to_gpu_args *args = (kfd_ioctl_map_memory_to_gpu_args *)argp;
+    D("AMDKFD_IOC_MAP_MEMORY_TO_GPU handle:%llX", args->handle);
+  } else if (request == AMDKFD_IOC_CREATE_EVENT) {
+    kfd_ioctl_create_event_args *args = (kfd_ioctl_create_event_args *)argp;
+    D("AMDKFD_IOC_CREATE_EVENT event_page_offset:0x%llx event_type:%d event_id:%d", args->event_page_offset, args->event_type, args->event_id);
+  } else if (request == AMDKFD_IOC_WAIT_EVENTS) {
+    D("AMDKFD_IOC_WAIT_EVENTS");
+  } else if (request == AMDKFD_IOC_SET_XNACK_MODE) {
+    D("AMDKFD_IOC_SET_XNACK_MODE");
+  } else if (request == AMDKFD_IOC_SVM || (type == 0x4b && nr == 0x20)) {
+    // NOTE: this one is variable length
+    kfd_ioctl_svm_args *args = (kfd_ioctl_svm_args *)argp;
+    D("AMDKFD_IOC_SVM start_addr:0x%llx size:0x%llx op:%d", args->start_addr, args->size, args->op);
+  } else if (request == AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU) {
+    kfd_ioctl_unmap_memory_from_gpu_args *args = (kfd_ioctl_unmap_memory_from_gpu_args *)argp;
+    D("AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU handle:%llX", args->handle);
+  } else if (request == AMDKFD_IOC_FREE_MEMORY_OF_GPU) {
+    D("AMDKFD_IOC_FREE_MEMORY_OF_GPU");
+  } else if (request == AMDKFD_IOC_SET_SCRATCH_BACKING_VA) {
+    D("AMDKFD_IOC_SET_SCRATCH_BACKING_VA");
+  } else if (request == AMDKFD_IOC_GET_TILE_CONFIG) {
+    D("AMDKFD_IOC_GET_TILE_CONFIG");
+  } else if (request == AMDKFD_IOC_SET_TRAP_HANDLER) {
+    D("AMDKFD_IOC_SET_TRAP_HANDLER");
+  } else if (request == AMDKFD_IOC_GET_VERSION) {
+    kfd_ioctl_get_version_args *args = (kfd_ioctl_get_version_args *)argp;
+    D("AMDKFD_IOC_GET_VERSION major_version:%d minor_version:%d", args->major_version, args->minor_version);
+  } else if (request == AMDKFD_IOC_GET_PROCESS_APERTURES_NEW) {
+    D("AMDKFD_IOC_GET_PROCESS_APERTURES_NEW");
+  } else if (request == AMDKFD_IOC_ACQUIRE_VM) {
+    D("AMDKFD_IOC_ACQUIRE_VM");
+  } else if (request == AMDKFD_IOC_SET_MEMORY_POLICY) {
+    D("AMDKFD_IOC_SET_MEMORY_POLICY");
+  } else if (request == AMDKFD_IOC_GET_CLOCK_COUNTERS) {
+    D("AMDKFD_IOC_GET_CLOCK_COUNTERS");
+  } else if (request == AMDKFD_IOC_CREATE_QUEUE) {
+    kfd_ioctl_create_queue_args *args = (kfd_ioctl_create_queue_args *)argp;
+    D("AMDKFD_IOC_CREATE_QUEUE\n");
+    D("queue_type:%d ring_base_address:0x%llx\n", args->queue_type, args->ring_base_address);
+    D("eop_buffer_address:0x%llx ctx_save_restore_address:0x%llx\n", args->eop_buffer_address, args->ctx_save_restore_address);
+    D("ring_size:0x%x queue_priority:%d\n", args->ring_size, args->queue_priority);
+    D("RETURNS write_pointer_address:0x%llx read_pointer_address:0x%llx doorbell_offset:0x%llx queue_id:%d\n", args->write_pointer_address, args->read_pointer_address, args->doorbell_offset, args->queue_id);
+    //D("RETURNS *write_pointer_address:0x%llx *read_pointer_address:0x%llx\n", *(uint64_t*)args->write_pointer_address, *(uint64_t*)args->read_pointer_address);
+    ring_base_addresses[args->doorbell_offset&0xFFF] = args->ring_base_address;
+    queue_types[args->doorbell_offset&0xFFF] = args->queue_type;
+    doorbell_offset = args->doorbell_offset&~0xFFF;
+  } else {
+    D("type:0x%x nr:0x%x size:0x%x", type, nr, size);
+  }
+
+  D("\n");
+  ioctl_num++;
+  return ret;
+}
+
+}

extra/thneed.py

@@ -1,4 +1,4 @@
-# this can be constructed from a cl_cache or loaded from a thneed file 
+# this can be constructed from a cl_cache or loaded from a thneed file
 import time
 import struct
 import json
@@ -26,7 +26,7 @@ class Thneed:
       for a in args[3:]:
         nodes[a]['out_edges'].append(args[2])
         nodes[args[2]]['in_edges'].append(a)
-    
+
     # get buffers to save
     self.buffers_to_save = set()
     self.outputs = []
@@ -35,7 +35,7 @@ class Thneed:
         self.buffers_to_save.add(n)
       if len(nodes[n]['out_edges']) == 0:
         self.outputs.append(n)
-  
+
     fake_inputs = []
     for k,n in self.inputs.items():
       if n in self.buffers_to_save:
@@ -97,7 +97,7 @@ class Thneed:
         else:
           # zero out buffers
           buf = cl.Buffer(CL.cl_ctx, mf.READ_WRITE | mf.COPY_HOST_PTR, hostbuf=b'\x00'*o['size'])
-        
+
       bufs[o['id']] = buf
       bufs_loaded[o['id']] = 'data' in o
       # if it's loaded, it's saved
@@ -114,13 +114,13 @@ class Thneed:
         print("FAILED", k)
         traceback.print_exc()
         exit(0)
-    
+
     # load binaries
     for o in jdat['binaries']:
       nptr = ptr + o['length']
       prgs[o['name']] = CLProgram(o['name'], weights[ptr:nptr], binary=True)
       ptr = nptr
-  
+
     # populate the cl_cache
     for i,k in enumerate(jdat['kernels']):
       kernel = prgs[k['name']]
@@ -166,7 +166,7 @@ class Thneed:
         jdat['binaries'].append({"name":prg.name, "length":len(binary[0])})
         binaries.append(binary[0])
         saved_binaries.add(prg.name)
-      
+
       # get the args from the kernel, some need the data saved
       targs, args_size = [], []
       argdtypes = prg.argdtypes if prg.argdtypes is not None else [None]*(len(args)-2)
@@ -193,9 +193,10 @@ class Thneed:
               })
               if needs_load:
                 data = np.empty(a.size//4, dtype=np.float32)
-                cl.enqueue_copy(CL.cl_queue, data, a, is_blocking=True)
+                cl.enqueue_copy(CL.cl_queue[0], data, a, is_blocking=True)
                 weights.append(data.tobytes())
             elif isinstance(a, cl.Image):
+              assert a.format == cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.HALF_FLOAT if FLOAT16 else cl.channel_type.FLOAT), "wrong type"
               needs_load = a in self.buffers_to_save
               row_pitch = (a.shape[0]*4*(2 if FLOAT16 else 4) + 63)//64 * 64
               size = row_pitch * a.shape[1]
@@ -203,7 +204,7 @@ class Thneed:
               buf = cl.Buffer(CL.cl_ctx, cl.mem_flags.READ_WRITE, size=size * (2 if FLOAT16 else 1))
 
               # zero out the buffer
-              cl.enqueue_copy(CL.cl_queue, buf, b'\x00'*buf.size, is_blocking=True)
+              cl.enqueue_copy(CL.cl_queue[0], buf, b'\x00'*buf.size, is_blocking=True)
 
               CLProgram("from_image_strided", """
                 __kernel void from_image_strided(read_only image2d_t in, __global float4 *out, int row_pitch) {
@@ -223,7 +224,7 @@ class Thneed:
 
               if needs_load:
                 data = np.empty(size//(2 if FLOAT16 else 4), dtype=np.float32)
-                cl.enqueue_copy(CL.cl_queue, data, buf, is_blocking=True)
+                cl.enqueue_copy(CL.cl_queue[0], data, buf, is_blocking=True)
                 if FLOAT16: data = data.astype(np.float16)
                 weights.append(data.tobytes())
             else:
@@ -244,7 +245,7 @@ class Thneed:
         "local_work_size": [1 for _ in args[0]] if args[1] is None else args[1],
         "num_args": len(args)-2,
         "args": targs,
-        "args_size": args_size 
+        "args_size": args_size
       })
 
     jdat['outputs'] = [{
@@ -270,9 +271,9 @@ class Thneed:
     events = []
     st = time.monotonic()
     for prg, args in self.cl_cache:
-      events.append(prg.clprg(CL.cl_queue, *args))
+      events.append(prg.clprg(CL.cl_queue[0], *args))
     mt = time.monotonic()
-    CL.cl_queue.finish()
+    CL.synchronize()
     et = time.monotonic() - st
     print(f"submit in {(mt-st)*1000.0:.2f} ms, total runtime is {et*1000.0:.2f} ms")
 

extra/training.py

@@ -5,7 +5,7 @@ from tinygrad.helpers import getenv
 
 def sparse_categorical_crossentropy(out, Y):
   num_classes = out.shape[-1]
-  YY = Y.flatten()
+  YY = Y.flatten().astype(np.int32)
   y = np.zeros((YY.shape[0], num_classes), np.float32)
   # correct loss for NLL, torch NLL loss returns one per row
   y[range(y.shape[0]),YY] = -1.0*num_classes

extra/utils.py

@@ -1,8 +1,14 @@
 import pickle
 import numpy as np
 from tqdm import tqdm
-import tempfile
-from tinygrad.helpers import prod, getenv
+import tempfile, platform
+from collections import defaultdict
+from tinygrad.helpers import prod, getenv, DEBUG
+from tinygrad.ops import GlobalCounters
+from tinygrad.tensor import Tensor
+from tinygrad.lazy import LazyNumpyArray, Device
+from tinygrad.shape.shapetracker import strides_for_shape
+OSX = platform.system() == "Darwin"
 
 def fetch(url):
   if url.startswith("/"):
@@ -14,7 +20,7 @@ def fetch(url):
   with open(fp, "rb") as f:
     return f.read()
 
-def download_file(url, fp, skip_if_exists=False):
+def download_file(url, fp, skip_if_exists=True):
   import requests, os
   if skip_if_exists and os.path.isfile(fp) and os.stat(fp).st_size > 0:
     return
@@ -28,43 +34,41 @@ def download_file(url, fp, skip_if_exists=False):
     os.rename(f.name, fp)
 
 def my_unpickle(fb0):
-  key_prelookup = {}
-  class HackTensor:
-    def __new__(cls, *args):
-      #print(args)
-      ident, storage_type, obj_key, location, obj_size = args[0][0:5]
-      assert ident == 'storage'
+  key_prelookup = defaultdict(list)
+  def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata=None):
+    #print(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata)
+    ident, storage_type, obj_key, location, obj_size = storage[0:5]
+    assert ident == 'storage'
+    assert prod(size) <= (obj_size - storage_offset)
 
-      assert prod(args[2]) == obj_size
-      ret = np.zeros(args[2], dtype=storage_type)
-      key_prelookup[obj_key] = (storage_type, obj_size, ret, args[2], args[3])
-      return ret
+    if storage_type not in [np.float16, np.float32]:
+      if DEBUG: print(f"unsupported type {storage_type} on {obj_key} with shape {size}")
+      ret = None
+    else:
+      ret = Tensor(LazyNumpyArray(lambda lst: np.zeros(lst.shape, dtype=lst.dtype), tuple(size), storage_type))
+    key_prelookup[obj_key].append((storage_type, obj_size, ret, size, stride, storage_offset))
+    return ret
 
-  class HackParameter:
-    def __new__(cls, *args):
-      #print(args)
-      pass
-
-  class Dummy:
+  def _rebuild_parameter(*args):
+    #print(args)
     pass
 
+  class Dummy: pass
+
   class MyPickle(pickle.Unpickler):
     def find_class(self, module, name):
       #print(module, name)
-      if name == 'FloatStorage':
-        return np.float32
-      if name == 'LongStorage':
-        return np.int64
-      if name == 'HalfStorage':
-        return np.float16
+      if name == 'FloatStorage': return np.float32
+      if name == 'LongStorage': return np.int64
+      if name == 'IntStorage': return np.int32
+      if name == 'HalfStorage': return np.float16
       if module == "torch._utils":
-        if name == "_rebuild_tensor_v2":
-          return HackTensor
-        elif name == "_rebuild_parameter":
-          return HackParameter
+        if name == "_rebuild_tensor_v2": return _rebuild_tensor_v2
+        if name == "_rebuild_parameter": return _rebuild_parameter
       else:
+        if module.startswith('pytorch_lightning'): return Dummy
         try:
-          return pickle.Unpickler.find_class(self, module, name)
+          return super().find_class(module, name)
         except Exception:
           return Dummy
 
@@ -73,18 +77,78 @@ def my_unpickle(fb0):
 
   return MyPickle(fb0).load(), key_prelookup
 
-def fake_torch_load_zipped(fb0, load_weights=True):
+def load_single_weight(t:Tensor, myfile, shape, strides, dtype, storage_offset, mmap_allowed=False):
+  bytes_size = np.dtype(dtype).itemsize
+  if t is None:
+    myfile.seek(prod(shape) * bytes_size, 1)
+    return
+
+  bytes_offset = 0
+  if storage_offset is not None:
+    bytes_offset = storage_offset * bytes_size
+    myfile.seek(bytes_offset)
+
+  assert t.shape == shape or shape == tuple(), f"shape mismatch {t.shape} != {shape}"
+  assert t.dtype.np == dtype and t.dtype.itemsize == bytes_size
+  if any(s != 1 and st1 != st2 for s, st1, st2 in zip(shape, strides_for_shape(shape), strides)):
+    # slow path
+    buffer_size = sum(strides[i]*t.dtype.itemsize * (shape[i] - 1) for i in range(len(shape)))
+    buffer_size += t.dtype.itemsize
+    np_array = np.frombuffer(myfile.read(buffer_size), t.dtype.np)
+
+    np_array = np.lib.stride_tricks.as_strided(
+      np_array, shape=shape, strides=[i*t.dtype.itemsize for i in strides])
+
+    lna = t.lazydata.op.arg
+    lna.fxn = lambda _: np_array
+    t.realize()
+    return
+
+  # ["METAL", "CLANG", "LLVM"] support readinto for more speed
+  # ["GPU", "CUDA"] use _mmap since they have to copy in to the GPU anyway
+  # this needs real APIs
+  if t.device in ["METAL", "CLANG", "LLVM"]:
+    del t.lazydata.op
+    t.lazydata.realized = Device[t.lazydata.device].buffer(prod(t.shape), dtype=t.dtype)
+    myfile.readinto(t.lazydata.realized._buffer())
+  else:
+    def _mmap(lna):
+      assert myfile._compress_type == 0, "compressed data can't be mmaped"
+      return np.memmap(myfile._fileobj._file, dtype=lna.dtype, mode='r', offset=myfile._orig_compress_start + bytes_offset, shape=lna.shape)
+    def _read(lna):
+      ret = np.empty(lna.shape, dtype=lna.dtype)
+      myfile.readinto(ret.data)
+      return ret
+    if mmap_allowed and not OSX and t.device in ["GPU", "CUDA"]: t.lazydata.op.arg.fxn = _mmap
+    else: t.lazydata.op.arg.fxn = _read
+    t.realize()
+
+def fake_torch_load_zipped(fb0, load_weights=True, multithreaded=True):
+  if Device.DEFAULT in ["TORCH", "GPU", "CUDA"]: multithreaded = False  # multithreaded doesn't work with CUDA or TORCH. for GPU it's a wash with _mmap
+
   import zipfile
   with zipfile.ZipFile(fb0, 'r') as myzip:
-    with myzip.open('archive/data.pkl') as myfile:
+    base_name = myzip.namelist()[0].split('/', 1)[0]
+    with myzip.open(f'{base_name}/data.pkl') as myfile:
       ret = my_unpickle(myfile)
     if load_weights:
-      for k,v in ret[1].items():
-        with myzip.open(f'archive/data/{k}') as myfile:
-          if v[2].dtype == "object":
-            print(f"issue assigning object on {k}")
-            continue
-          np.copyto(v[2], np.frombuffer(myfile.read(), v[2].dtype).reshape(v[3]))
+      def load_weight(k, vv):
+        with myzip.open(f'{base_name}/data/{k}') as myfile:
+          for v in vv:
+            load_single_weight(v[2], myfile, v[3], v[4], v[0], v[5], mmap_allowed=True)
+      if multithreaded:
+        import concurrent.futures
+        # 2 seems fastest
+        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+          futures = {executor.submit(load_weight, k, v):k for k,v in ret[1].items()}
+          for future in (t:=tqdm(concurrent.futures.as_completed(futures), total=len(futures))):
+            if future.exception() is not None: raise future.exception()
+            k = futures[future]
+            t.set_description(f"loading {k} ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
+      else:
+        for k,v in (t := tqdm(ret[1].items())):
+          t.set_description(f"loading {k} ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
+          load_weight(k,v)
   return ret[0]
 
 def fake_torch_load(b0):
@@ -108,19 +172,15 @@ def fake_torch_load(b0):
   key_lookup = pickle.load(fb0)
   key_real = [None] * len(key_lookup)
   for k,v in key_prelookup.items():
-    key_real[key_lookup.index(k)] = v
+    assert len(v) == 1
+    key_real[key_lookup.index(k)] = v[0]
 
   # read in the actual data
-  for storage_type, obj_size, np_array, np_shape, np_strides in key_real:
+  for storage_type, obj_size, tensor, np_shape, np_strides, storage_offset in key_real:
     ll = struct.unpack("Q", fb0.read(8))[0]
-    assert ll == obj_size
-    bytes_size = {np.float32: 4, np.int64: 8}[storage_type]
-    mydat = fb0.read(ll * bytes_size)
-    np.copyto(np_array, np.frombuffer(mydat, storage_type).reshape(np_shape))
-
-    # numpy stores its strides in bytes
-    real_strides = tuple([x*bytes_size for x in np_strides])
-    np_array.strides = real_strides
+    assert ll == obj_size, f"size mismatch {ll} != {obj_size}"
+    assert storage_offset == 0, "not implemented"
+    load_single_weight(tensor, fb0, np_shape, np_strides, storage_type, None)
 
   return ret
 

models/convnext.py

@@ -0,0 +1,64 @@
+from tinygrad.tensor import Tensor
+from tinygrad.nn import Conv2d, LayerNorm, LayerNorm2d, Linear
+
+class Block:
+  def __init__(self, dim):
+    self.dwconv = Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+    self.norm = LayerNorm(dim, eps=1e-6)
+    self.pwconv1 = Linear(dim, 4 * dim)
+    self.pwconv2 = Linear(4 * dim, dim)
+    self.gamma = Tensor.ones(dim)
+
+  def __call__(self, x:Tensor):
+    return x + x.sequential([
+      self.dwconv, lambda x: x.permute(0, 2, 3, 1), self.norm,
+      self.pwconv1, Tensor.gelu, self.pwconv2, lambda x: (self.gamma * x).permute(0, 3, 1, 2)
+    ])
+
+class ConvNeXt:
+  def __init__(self, in_chans=3, num_classes=1000, depths=[3, 3, 9, 3], dims=[96, 192, 384, 768]):
+    self.downsample_layers = [
+      [Conv2d(in_chans, dims[0], kernel_size=4, stride=4), LayerNorm2d(dims[0], eps=1e-6)],
+      *[[LayerNorm2d(dims[i], eps=1e-6), Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2)] for i in range(len(dims)-1)]
+    ]
+    self.stages = [[Block(dims[i]) for _ in range(depths[i])] for i in range(len(dims))]
+    self.norm = LayerNorm(dims[-1])
+    self.head = Linear(dims[-1], num_classes)
+
+  def __call__(self, x:Tensor):
+    for downsample, stage in zip(self.downsample_layers, self.stages):
+      x = x.sequential(downsample).sequential(stage)
+    return x.mean([-2, -1]).sequential([self.norm, self.head])
+
+# *** model definition is done ***
+
+versions = {
+  "tiny": {"depths": [3, 3, 9, 3], "dims": [96, 192, 384, 768]},
+  "small": {"depths": [3, 3, 27, 3], "dims": [96, 192, 384, 768]},
+  "base": {"depths": [3, 3, 9, 3], "dims": [128, 256, 512, 1024]},
+  "large": {"depths": [3, 3, 27, 3], "dims": [192, 384, 768, 1536]},
+  "xlarge": {"depths": [3, 3, 27, 3], "dims": [256, 512, 1024, 2048]}
+}
+
+def get_model(version, load_weights=False):
+  model = ConvNeXt(**versions[version])
+  if load_weights:
+    from extra.utils import fetch, fake_torch_load, get_child
+    weights = fake_torch_load(fetch(f'https://dl.fbaipublicfiles.com/convnext/convnext_{version}_1k_224_ema.pth'))['model']
+    for k,v in weights.items():
+      mv = get_child(model, k)
+      mv.assign(v.reshape(mv.shape)).realize()
+  return model
+
+if __name__ == "__main__":
+  model = get_model("tiny", True)
+
+  # load image
+  from test.models.test_efficientnet import chicken_img, preprocess, _LABELS
+  img = Tensor(preprocess(chicken_img))
+
+  Tensor.training = False
+  Tensor.no_grad = True
+
+  out = model(img).numpy()
+  print(_LABELS[out.argmax()])

models/efficientnet.py

@@ -146,6 +146,7 @@ class EfficientNet:
 
     b0 = fake_torch_load(fetch(model_urls[self.number]))
     for k,v in b0.items():
+      if k.endswith("num_batches_tracked"): continue
       for cat in ['_conv_head', '_conv_stem', '_depthwise_conv', '_expand_conv', '_fc', '_project_conv', '_se_reduce', '_se_expand']:
         if cat in k:
           k = k.replace('.bias', '_bias')
@@ -153,9 +154,9 @@ class EfficientNet:
 
       #print(k, v.shape)
       mv = get_child(self, k)
-      vnp = v.astype(np.float32)
-      vnp = vnp if k != '_fc' else vnp.T
-      vnp = vnp if vnp.shape != () else np.array([vnp])
+      vnp = v #.astype(np.float32)
+      vnp = vnp if k != '_fc' else vnp.transpose()
+      #vnp = vnp if vnp.shape != () else np.array([vnp])
 
       if mv.shape == vnp.shape:
         mv.assign(vnp)

models/resnet.py

@@ -1,7 +1,6 @@
 from tinygrad.tensor import Tensor
 import tinygrad.nn as nn
 from extra.utils import get_child
-import numpy as np
 
 class BasicBlock:
   expansion = 1
@@ -27,6 +26,7 @@ class BasicBlock:
 
 
 class Bottleneck:
+  # NOTE: the original implementation places stride at the first convolution (self.conv1), this is the v1.5 variant
   expansion = 4
 
   def __init__(self, in_planes, planes, stride=1):
@@ -34,7 +34,7 @@ class Bottleneck:
     self.bn1 = nn.BatchNorm2d(planes)
     self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, stride=stride, bias=False)
     self.bn2 = nn.BatchNorm2d(planes)
-    self.conv3 = nn.Conv2d(planes, self.expansion *planes, kernel_size=1, bias=False)
+    self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False)
     self.bn3 = nn.BatchNorm2d(self.expansion*planes)
     self.downsample = []
     if stride != 1 or in_planes != self.expansion*planes:
@@ -52,7 +52,6 @@ class Bottleneck:
     return out
 
 class ResNet:
-  # def __init__(self, block, num_blocks, num_classes=10, url=None):
   def __init__(self, num, num_classes):
     self.num = num
 
@@ -76,11 +75,11 @@ class ResNet:
 
     self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, bias=False, padding=3)
     self.bn1 = nn.BatchNorm2d(64)
-    self.layer1 = self._make_layer(self.block, 64, self.num_blocks[0], stride=2)
+    self.layer1 = self._make_layer(self.block, 64, self.num_blocks[0], stride=1)
     self.layer2 = self._make_layer(self.block, 128, self.num_blocks[1], stride=2)
     self.layer3 = self._make_layer(self.block, 256, self.num_blocks[2], stride=2)
     self.layer4 = self._make_layer(self.block, 512, self.num_blocks[3], stride=2)
-    self.fc = {"weight": Tensor.uniform(512 * self.block.expansion, num_classes), "bias": Tensor.zeros(num_classes)}
+    self.fc = nn.Linear(512 * self.block.expansion, num_classes)
 
   def _make_layer(self, block, planes, num_blocks, stride):
     strides = [stride] + [1] * (num_blocks-1)
@@ -92,12 +91,13 @@ class ResNet:
 
   def forward(self, x):
     out = self.bn1(self.conv1(x)).relu()
+    out = out.pad2d([1,1,1,1]).max_pool2d((3,3), 2)
     out = out.sequential(self.layer1)
     out = out.sequential(self.layer2)
     out = out.sequential(self.layer3)
     out = out.sequential(self.layer4)
-    out = out.mean(3).mean(2)
-    out = out.linear(**self.fc).log_softmax()
+    out = out.mean([2,3])
+    out = self.fc(out).log_softmax()
     return out
 
   def __call__(self, x):
@@ -105,7 +105,7 @@ class ResNet:
 
   def load_from_pretrained(self):
     # TODO replace with fake torch load
-  
+
     model_urls = {
       18: 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
       34: 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
@@ -120,7 +120,7 @@ class ResNet:
     state_dict = load_state_dict_from_url(self.url, progress=True)
     for k, v in state_dict.items():
       obj = get_child(self, k)
-      dat = v.detach().numpy().T if "fc.weight" in k else v.detach().numpy()
+      dat = v.detach().numpy()
 
       if 'fc.' in k and obj.shape != dat.shape:
         print("skipping fully connected layer")

models/rnnt.py

@@ -0,0 +1,218 @@
+from tinygrad.tensor import Tensor
+from tinygrad.jit import TinyJit
+from tinygrad.nn import Linear
+import numpy as np
+from extra.utils import download_file
+from pathlib import Path
+
+
+class RNNT:
+  def __init__(self, input_features=240, vocab_size=29, enc_hidden_size=1024, pred_hidden_size=320, joint_hidden_size=512, pre_enc_layers=2, post_enc_layers=3, pred_layers=2, stack_time_factor=2, dropout=0.32):
+    self.encoder = Encoder(input_features, enc_hidden_size, pre_enc_layers, post_enc_layers, stack_time_factor, dropout)
+    self.prediction = Prediction(vocab_size, pred_hidden_size, pred_layers, dropout)
+    self.joint = Joint(vocab_size, pred_hidden_size, enc_hidden_size, joint_hidden_size, dropout)
+
+  @TinyJit
+  def __call__(self, x, y, hc=None):
+    f, _ = self.encoder(x, None)
+    g, _ = self.prediction(y, hc, Tensor.ones(1, requires_grad=False))
+    out = self.joint(f, g)
+    return out.realize()
+
+  def decode(self, x, x_lens):
+    logits, logit_lens = self.encoder(x, x_lens)
+    outputs = []
+    for b in range(logits.shape[0]):
+      inseq = logits[b, :, :].unsqueeze(1)
+      logit_len = logit_lens[b]
+      seq = self._greedy_decode(inseq, int(np.ceil(logit_len.numpy()).item()))
+      outputs.append(seq)
+    return outputs
+
+  def _greedy_decode(self, logits, logit_len):
+    hc = Tensor.zeros(self.prediction.rnn.layers, 2, self.prediction.hidden_size, requires_grad=False)
+    labels = []
+    label = Tensor.zeros(1, 1, requires_grad=False)
+    mask = Tensor.zeros(1, requires_grad=False)
+    for time_idx in range(logit_len):
+      logit = logits[time_idx, :, :].unsqueeze(0)
+      not_blank = True
+      added = 0
+      while not_blank and added < 30:
+        if len(labels) > 0:
+          mask = (mask + 1).clip(0, 1)
+          label = Tensor([[labels[-1] if labels[-1] <= 28 else labels[-1] - 1]], requires_grad=False) + 1 - 1
+        jhc = self._pred_joint(Tensor(logit.numpy()), label, hc, mask)
+        k = np.argmax(jhc[0, 0, :29].numpy(), axis=0)
+        not_blank = k != 28
+        if not_blank:
+          labels.append(k)
+          hc = jhc[:, :, 29:] + 1 - 1
+        added += 1
+    return labels
+
+  @TinyJit
+  def _pred_joint(self, logit, label, hc, mask):
+    g, hc = self.prediction(label, hc, mask)
+    j = self.joint(logit, g)[0]
+    j = j.pad(((0, 1), (0, 1), (0, 0)))
+    out = j.cat(hc, dim=2)
+    return out.realize()
+
+  def load_from_pretrained(self):
+    fn = Path(__file__).parent.parent / "weights/rnnt.pt"
+    download_file("https://zenodo.org/record/3662521/files/DistributedDataParallel_1576581068.9962234-epoch-100.pt?download=1", fn)
+
+    import torch
+    with open(fn, "rb") as f:
+      state_dict = torch.load(f, map_location="cpu")["state_dict"]
+
+    # encoder
+    for i in range(2):
+      self.encoder.pre_rnn.cells[i].weights_ih.assign(state_dict[f"encoder.pre_rnn.lstm.weight_ih_l{i}"].numpy())
+      self.encoder.pre_rnn.cells[i].weights_hh.assign(state_dict[f"encoder.pre_rnn.lstm.weight_hh_l{i}"].numpy())
+      self.encoder.pre_rnn.cells[i].bias_ih.assign(state_dict[f"encoder.pre_rnn.lstm.bias_ih_l{i}"].numpy())
+      self.encoder.pre_rnn.cells[i].bias_hh.assign(state_dict[f"encoder.pre_rnn.lstm.bias_hh_l{i}"].numpy())
+    for i in range(3):
+      self.encoder.post_rnn.cells[i].weights_ih.assign(state_dict[f"encoder.post_rnn.lstm.weight_ih_l{i}"].numpy())
+      self.encoder.post_rnn.cells[i].weights_hh.assign(state_dict[f"encoder.post_rnn.lstm.weight_hh_l{i}"].numpy())
+      self.encoder.post_rnn.cells[i].bias_ih.assign(state_dict[f"encoder.post_rnn.lstm.bias_ih_l{i}"].numpy())
+      self.encoder.post_rnn.cells[i].bias_hh.assign(state_dict[f"encoder.post_rnn.lstm.bias_hh_l{i}"].numpy())
+
+    # prediction
+    self.prediction.emb.weight.assign(state_dict["prediction.embed.weight"].numpy())
+    for i in range(2):
+      self.prediction.rnn.cells[i].weights_ih.assign(state_dict[f"prediction.dec_rnn.lstm.weight_ih_l{i}"].numpy())
+      self.prediction.rnn.cells[i].weights_hh.assign(state_dict[f"prediction.dec_rnn.lstm.weight_hh_l{i}"].numpy())
+      self.prediction.rnn.cells[i].bias_ih.assign(state_dict[f"prediction.dec_rnn.lstm.bias_ih_l{i}"].numpy())
+      self.prediction.rnn.cells[i].bias_hh.assign(state_dict[f"prediction.dec_rnn.lstm.bias_hh_l{i}"].numpy())
+
+    # joint
+    self.joint.l1.weight.assign(state_dict["joint_net.0.weight"].numpy())
+    self.joint.l1.bias.assign(state_dict["joint_net.0.bias"].numpy())
+    self.joint.l2.weight.assign(state_dict["joint_net.3.weight"].numpy())
+    self.joint.l2.bias.assign(state_dict["joint_net.3.bias"].numpy())
+
+
+class LSTMCell:
+  def __init__(self, input_size, hidden_size, dropout):
+    self.dropout = dropout
+
+    self.weights_ih = Tensor.uniform(hidden_size * 4, input_size)
+    self.bias_ih = Tensor.uniform(hidden_size * 4)
+    self.weights_hh = Tensor.uniform(hidden_size * 4, hidden_size)
+    self.bias_hh = Tensor.uniform(hidden_size * 4)
+
+  def __call__(self, x, hc):
+    gates = x.linear(self.weights_ih.T, self.bias_ih) + hc[:x.shape[0]].linear(self.weights_hh.T, self.bias_hh)
+
+    i, f, g, o = gates.chunk(4, 1)
+    i, f, g, o = i.sigmoid(), f.sigmoid(), g.tanh(), o.sigmoid()
+
+    c = (f * hc[x.shape[0]:]) + (i * g)
+    h = (o * c.tanh()).dropout(self.dropout)
+
+    return Tensor.cat(h, c).realize()
+
+
+class LSTM:
+  def __init__(self, input_size, hidden_size, layers, dropout):
+    self.input_size = input_size
+    self.hidden_size = hidden_size
+    self.layers = layers
+
+    self.cells = [LSTMCell(input_size, hidden_size, dropout) if i == 0 else LSTMCell(hidden_size, hidden_size, dropout if i != layers - 1 else 0) for i in range(layers)]
+
+  def __call__(self, x, hc):
+    @TinyJit
+    def _do_step(x_, hc_):
+      return self.do_step(x_, hc_)
+
+    if hc is None:
+      hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size, requires_grad=False)
+
+    output = None
+    for t in range(x.shape[0]):
+      hc = _do_step(x[t] + 1 - 1, hc) # TODO: why do we need to do this?
+      if output is None:
+        output = hc[-1:, :x.shape[1]]
+      else:
+        output = output.cat(hc[-1:, :x.shape[1]], dim=0).realize()
+
+    return output, hc
+
+  def do_step(self, x, hc):
+    new_hc = [x]
+    for i, cell in enumerate(self.cells):
+      new_hc.append(cell(new_hc[i][:x.shape[0]], hc[i]))
+    return Tensor.stack(new_hc[1:]).realize()
+
+
+class StackTime:
+  def __init__(self, factor):
+    self.factor = factor
+
+  def __call__(self, x, x_lens):
+    x = x.pad(((0, (-x.shape[0]) % self.factor), (0, 0), (0, 0)))
+    x = x.reshape(x.shape[0] // self.factor, x.shape[1], x.shape[2] * self.factor)
+    return x, x_lens / self.factor if x_lens is not None else None
+
+
+class Encoder:
+  def __init__(self, input_size, hidden_size, pre_layers, post_layers, stack_time_factor, dropout):
+    self.pre_rnn = LSTM(input_size, hidden_size, pre_layers, dropout)
+    self.stack_time = StackTime(stack_time_factor)
+    self.post_rnn = LSTM(stack_time_factor * hidden_size, hidden_size, post_layers, dropout)
+
+  def __call__(self, x, x_lens):
+    x, _ = self.pre_rnn(x, None)
+    x, x_lens = self.stack_time(x, x_lens)
+    x, _ = self.post_rnn(x, None)
+    return x.transpose(0, 1), x_lens
+
+
+class Embedding:
+  def __init__(self, vocab_size: int, embed_size: int):
+    self.vocab_size = vocab_size
+    self.vocab_counter = Tensor(np.arange(vocab_size, dtype=np.float32), requires_grad=False)
+    self.weight = Tensor.scaled_uniform(vocab_size, embed_size)
+
+  def __call__(self, idx: Tensor) -> Tensor:
+    oha = []
+    for i in range(idx.shape[0]):
+      ohba = []
+      for j in range(idx.shape[1]):
+        ohba.append((self.vocab_counter == idx[i, j]).realize())
+      oha.append(Tensor.stack(ohba).realize())
+    return Tensor.stack(oha) @ self.weight
+
+
+class Prediction:
+  def __init__(self, vocab_size, hidden_size, layers, dropout):
+    self.hidden_size = hidden_size
+
+    self.emb = Embedding(vocab_size - 1, hidden_size)
+    self.rnn = LSTM(hidden_size, hidden_size, layers, dropout)
+
+  def __call__(self, x, hc, m):
+    emb = self.emb(x) * m
+    x_, hc = self.rnn(emb.transpose(0, 1), hc)
+    return x_.transpose(0, 1), hc
+
+
+class Joint:
+  def __init__(self, vocab_size, pred_hidden_size, enc_hidden_size, joint_hidden_size, dropout):
+    self.dropout = dropout
+
+    self.l1 = Linear(pred_hidden_size + enc_hidden_size, joint_hidden_size)
+    self.l2 = Linear(joint_hidden_size, vocab_size)
+
+  def __call__(self, f, g):
+    (_, T, H), (B, U, H2) = f.shape, g.shape
+    f = f.unsqueeze(2).expand(B, T, U, H)
+    g = g.unsqueeze(1).expand(B, T, U, H2)
+
+    inp = f.cat(g, dim=3)
+    t = self.l1(inp).relu()
+    t = t.dropout(self.dropout)
+    return self.l2(t)

models/transformer.py

@@ -26,13 +26,13 @@ class TransformerBlock:
       .reshape(shape=(x.shape[0], -1, self.num_heads, self.head_size)) \
       for y in [self.query, self.key, self.value]]
 
-    query = query.transpose(order=(0,2,1,3))  # (bs, num_heads, time, head_size)
-    key = key.transpose(order=(0,2,3,1))      # (bs, num_heads, head_size, time)
-    value = value.transpose(order=(0,2,1,3))  # (bs, num_heads, time, head_size)
+    query = query.permute(order=(0,2,1,3))  # (bs, num_heads, time, head_size)
+    key = key.permute(order=(0,2,3,1))      # (bs, num_heads, head_size, time)
+    value = value.permute(order=(0,2,1,3))  # (bs, num_heads, time, head_size)
 
     score = query.dot(key) * (1 / np.sqrt(self.head_size))
     weights = score.softmax()                                   # (bs, num_heads, time, time)
-    attention = weights.dot(value).transpose(order=(0,2,1,3))   # (bs, time, num_heads, head_size)
+    attention = weights.dot(value).permute(order=(0,2,1,3))   # (bs, time, num_heads, head_size)
 
     return attention.reshape(shape=(x.shape[0], -1, self.num_heads * self.head_size)).linear(*self.out)
 

models/unet3d.py

@@ -0,0 +1,42 @@
+# https://github.com/wolny/pytorch-3dunet
+from pathlib import Path
+from extra.utils import download_file, fake_torch_load, get_child
+import tinygrad.nn as nn
+
+class SingleConv:
+  def __init__(self, in_channels, out_channels):
+    self.groupnorm = nn.GroupNorm(1, in_channels) # 1 group?
+    # TODO: make 2D conv generic for 3D, might already work with kernel_size=(3,3,3)
+    self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=(3,3,3), padding=(1,1,1,1,1,1), bias=False)
+  def __call__(self, x):
+    return self.conv(self.groupnorm(x)).relu()
+
+class BasicModule:
+  def __init__(self, c0, c1, c2):
+    self.basic_module = {"SingleConv1": SingleConv(c0, c1), "SingleConv2": SingleConv(c1, c2)}
+  def __call__(self, x):
+    return self.basic_module['SingleConv2'](self.basic_module['SingleConv1'](x))
+
+class UNet3D:
+  def __init__(self):
+    ups = [16,32,64,128,256]
+    self.encoders = [BasicModule(ups[i] if i != 0 else 1, ups[i], ups[i+1]) for i in range(4)]
+    self.decoders = [BasicModule(ups[-1-i] + ups[-2-i], ups[-2-i], ups[-2-i]) for i in range(3)]
+    self.final_conv = nn.Conv2d(32, 1, (1,1,1))
+
+  def __call__(self, x):
+    intermediates = [x]
+    for e in self.encoders: intermediates.append(e(intermediates[-1]))
+    ret = intermediates[-1]
+    for d,i in zip(self.decoders, intermediates[:-1][::-1]): ret = d(ret.cat(i, dim=1))
+    return ret
+
+  def load_from_pretrained(self):
+    fn = Path(__file__).parent.parent / "weights/unet-3d.ckpt"
+    download_file("https://oc.embl.de/index.php/s/61s67Mg5VQy7dh9/download?path=%2FLateral-Root-Primordia%2Funet_bce_dice_ds1x&files=best_checkpoint.pytorch", fn)
+    state_dict = fake_torch_load(open(fn, "rb").read())['model_state_dict']
+    for k, v in state_dict.items():
+      print(k, v.shape)
+      obj = get_child(self, k)
+      assert obj.shape == v.shape, (k, obj.shape, v.shape)
+      obj.assign(v.numpy())

models/vit.py

@@ -17,7 +17,7 @@ class ViT:
 
   def patch_embed(self, x):
     x = x.conv2d(*self.embedding, stride=16)
-    x = x.reshape(shape=(x.shape[0], x.shape[1], -1)).transpose(order=(0,2,1))
+    x = x.reshape(shape=(x.shape[0], x.shape[1], -1)).permute(order=(0,2,1))
     return x
 
   def forward(self, x):

openpilot/compile.py

@@ -9,7 +9,7 @@ if os.getenv("GPU", None) is None:
 if os.getenv("IMAGE", None) is None:
   os.environ['IMAGE'] = '2'
 
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, dtypes
 ALLOWED_KERNEL_COUNT = getenv("ALLOWED_KERNEL_COUNT", 0)
 DEBUGCL = getenv("DEBUGCL", 0)
 
@@ -38,7 +38,7 @@ from tinygrad.jit import TinyJit
 
 @TinyJit
 def model_exec(run_onnx, using_graph, **inputs):
-  ret = next(iter(run_onnx(inputs).values()))
+  ret = next(iter(run_onnx(inputs).values())).cast(dtypes.float32)
   GlobalCounters.reset()
   GlobalCounters.cache = []  # don't cache pre-realize
   if using_graph: graph.GRAPH = True
@@ -49,7 +49,7 @@ def compile(dat, output_fn):
   Tensor.manual_seed(1337)
   Tensor.no_grad = True
   using_graph = graph.GRAPH
-  graph.GRAPH = False
+  if getenv("GRAPH") < 3: graph.GRAPH = False
 
   onnx_model = onnx.load(io.BytesIO(dat))
   run_onnx = get_run_onnx(onnx_model)
@@ -63,8 +63,8 @@ def compile(dat, output_fn):
   assert len(model_exec.jit_cache) <= ALLOWED_KERNEL_COUNT or ALLOWED_KERNEL_COUNT == 0, "too many kernels!"
 
   # pull out inputs and put them in the jit cache
-  input_rawbuffers = {k:inputs[k].lazydata.realized.raw() for k in inputs.keys()}
-  for (j,i),idx in model_exec.input_replace.items(): model_exec.jit_cache[j][1][i] = input_rawbuffers[idx]
+  input_rawbuffers = {k:inputs[k].lazydata.realized for k in inputs.keys()}
+  for (j,i),(idx,_,_) in model_exec.input_replace.items(): model_exec.jit_cache[j][1][i] = input_rawbuffers[idx]
 
   # transform to CL.CACHE
   used_ops = 0
@@ -73,11 +73,11 @@ def compile(dat, output_fn):
     # pass these to thneed
     setattr(prg.clprg, 'op_estimate', prg.op_estimate)
     setattr(prg.clprg, 'prg', prg.prg)
-    cl_cache.append((prg.clprg, [prg.global_size, prg.local_size, *[x._cl for x in args]]))
+    cl_cache.append((prg.clprg, [prg.global_size, prg.local_size, *[x._buf for x in args]]))
     used_ops += prg.op_estimate
 
   from extra.thneed import Thneed
-  t = Thneed(cl_cache, {k:v._cl for k,v in input_rawbuffers.items()})
+  t = Thneed(cl_cache, {k:v._buf for k,v in input_rawbuffers.items()})
 
   # save thneed (before run)
   t.save(output_fn)
@@ -88,7 +88,7 @@ def compile(dat, output_fn):
 
   # confirm thneed found the right output
   thneed_out = np.empty((t.outputs[0].size//4,), dtype=np.float32).reshape(tinygrad_out.shape)
-  cl.enqueue_copy(CL.cl_queue, thneed_out, t.outputs[0], is_blocking=True)
+  cl.enqueue_copy(CL.cl_queue[0], thneed_out, t.outputs[0], is_blocking=True)
   np.testing.assert_allclose(thneed_out, tinygrad_out.numpy())
 
   # testing is float32 only (fix this)
@@ -106,11 +106,11 @@ def compile(dat, output_fn):
 
       # try old thneed with a different input
       for k,v in t.inputs.items():
-        cl.enqueue_copy(CL.cl_queue, v, new_np_inputs[k], is_blocking=True)
+        cl.enqueue_copy(CL.cl_queue[0], v, new_np_inputs[k], is_blocking=True)
 
       t.run()
       old_thneed_out = np.empty((t.outputs[0].size//4,), dtype=np.float32).reshape(tinygrad_out.shape)
-      cl.enqueue_copy(CL.cl_queue, old_thneed_out, t.outputs[0], is_blocking=True)
+      cl.enqueue_copy(CL.cl_queue[0], old_thneed_out, t.outputs[0], is_blocking=True)
 
       # compare thneed (rerun) with torch
       np.testing.assert_allclose(new_torch_out, old_thneed_out, atol=1e-4, rtol=1e-2)
@@ -123,11 +123,11 @@ def compile(dat, output_fn):
 
       # inputs
       for k,v in nt.inputs.items():
-        cl.enqueue_copy(CL.cl_queue, v, new_np_inputs[k], is_blocking=True)
+        cl.enqueue_copy(CL.cl_queue[0], v, new_np_inputs[k], is_blocking=True)
 
       nt.run()
       new_thneed_out = np.empty((nt.outputs[0].size//4,), dtype=np.float32).reshape(tinygrad_out.shape)
-      cl.enqueue_copy(CL.cl_queue, new_thneed_out, nt.outputs[0], is_blocking=True)
+      cl.enqueue_copy(CL.cl_queue[0], new_thneed_out, nt.outputs[0], is_blocking=True)
 
       # compare torch to thneed
       np.testing.assert_allclose(new_torch_out, new_thneed_out, atol=1e-4, rtol=1e-2)

rmso.sh

@@ -1,3 +1,3 @@
 #!/bin/bash
-rm tinygrad/*.so tinygrad/shape/*.so tinygrad/llops/*.so tinygrad/nn/*.so tinygrad/runtime/*.so *.so
+rm tinygrad/*.so tinygrad/codegen/*.so tinygrad/shape/*.so tinygrad/nn/*.so tinygrad/runtime/*.so *.so
 

run_multibackend.sh

@@ -0,0 +1,13 @@
+#!/bin/bash -e
+echo "********* CPU *********"
+CPU=1 python3 $@
+echo "********* GPU *********"
+GPU=1 python3 $@
+echo "********* METAL *********"
+METAL=1 python3 $@
+echo "********* CLANG *********"
+CLANG=1 python3 $@
+echo "********* LLVM *********"
+LLVM=1 python3 $@
+echo "********* TORCH *********"
+TORCH=1 python3 $@

setup.py

@@ -8,8 +8,8 @@ with open(os.path.join(directory, 'README.md'), encoding='utf-8') as f:
   long_description = f.read()
 
 setup(name='tinygrad',
-      version='0.5.0',
-      description='You like pytorch? You like micrograd? You love tinygrad! heart',
+      version='0.6.0',
+      description='You like pytorch? You like micrograd? You love tinygrad! <3',
       author='George Hotz',
       license='MIT',
       long_description=long_description,
@@ -19,10 +19,9 @@ setup(name='tinygrad',
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: MIT License"
       ],
-      install_requires=['numpy', 'requests', 'pillow', 'tqdm', 'networkx'],
+      install_requires=['numpy', 'requests', 'pillow', 'tqdm', 'networkx', 'pyopencl'],
       python_requires='>=3.8',
       extras_require={
-        'gpu': ["pyopencl"],
         'llvm': ["llvmlite"],
         'cuda': ["pycuda"],
         'triton': ["triton>=2.0.0.dev20221202"],
@@ -34,10 +33,10 @@ setup(name='tinygrad',
             "pre-commit",
         ],
         'testing': [
-            "torch~=1.13.0",
+            "torch",
             "pytest",
             "pytest-xdist",
-            "onnx~=1.13.0",
+            "onnx",
             "onnx2torch",
             "opencv-python",
         ],

sz.sh

@@ -1,3 +1,5 @@
 #!/bin/bash
 # switched to cloc due to https://github.com/boyter/scc/issues/379
-cloc --by-file tinygrad/*
+cloc --by-file tinygrad/* | grep "tinygrad"
+# also some sloccount for a dir summary
+sloccount tinygrad | grep "python"

test/external/external_copy_benchmark.py

@@ -0,0 +1,26 @@
+import unittest
+from tinygrad.helpers import prod
+from tinygrad.lazy import Device
+from tinygrad.tensor import Tensor
+from tinygrad.ops import GlobalCounters
+
+class TestCopy(unittest.TestCase):
+  def test_add1(self):
+    pts = []
+    for i in range(16384, 16384*256, 16384):
+      t = Tensor.randn(i).realize()
+      GlobalCounters.cache = []
+      t.assign(t+1).realize()
+      fxn, args = GlobalCounters.cache[0]
+      GlobalCounters.reset()
+      def run(): return fxn(args, force_wait=True)
+      ct = min([run() for _ in range(10)])
+      mb = prod(t.shape)*t.dtype.itemsize*2*1e-6
+      print(f"{mb*1e3:.2f} kB, {ct*1e3:.2f} ms, {mb/ct:.2f} MB/s")
+      pts.append((mb, mb/ct))
+    from matplotlib import pyplot as plt
+    plt.plot([x[0] for x in pts], [x[1] for x in pts])
+    plt.show()
+
+if __name__ == '__main__':
+  unittest.main()

test/external/external_hlb_cifar.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+from examples.hlb_cifar10 import SpeedyResNet, fetch_batch
+from examples.hlb_cifar10_torch import SpeedyResNet as SpeedyResNetTorch
+from datasets import fetch_cifar
+from test.models.test_end2end import compare_tiny_torch
+
+if __name__ == "__main__":
+  X_test, Y_test = fetch_cifar(train=False)
+  X, Y = fetch_batch(X_test, Y_test, 32)
+  print(X.shape, Y.shape)
+  model = SpeedyResNet()
+  model_torch = SpeedyResNetTorch()
+  compare_tiny_torch(model, model_torch, X, Y)

test/external/external_multi_gpu.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+# cd disassemblers/ && git clone --recursive github.com:geohot/cuda_ioctl_sniffer.git
+# LD_PRELOAD=$PWD/disassemblers/cuda_ioctl_sniffer/out/sniff.so GPU=1 python3 test/external/external_multi_gpu.py
+import numpy as np
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import colored
+from extra.helpers import Timing
+from tinygrad.runtime.ops_gpu import CL
+
+# TODO: support multidevice in cuda
+device = 'gpu'
+
+if __name__ == "__main__":
+  sz = 1024*1024*256  # 1 GB
+  #sz = 1024*64
+
+  with Timing("CPU creation: ", on_exit=lambda x: f", {(sz*4*2)/x:.2f} GB/sec"):
+    c0 = Tensor.ones(sz, device="cpu").realize()
+    c1 = (Tensor.ones(sz, device="cpu")/2).realize()
+
+  with Timing("CPU -> 0: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    a0 = c0.to(f'{device}:0').realize()
+    CL.synchronize()
+  with Timing("CPU -> 1: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    b1 = c1.to(f'{device}:1').realize()
+    CL.synchronize()
+
+  # cross copy. this is going through the CPU
+  with Timing("0 -> CPU -> 1: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    a1 = a0.to(f'{device}:1').realize()
+    CL.synchronize()
+  with Timing("1 -> CPU -> 0: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    b0 = b1.to(f'{device}:0').realize()
+    CL.synchronize()
+
+  # sum
+  with Timing("0+0 -> 0 (sum): ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    ab0 = (a0 + b0).realize()
+    CL.synchronize()
+  with Timing("1+1 -> 1 (sum): ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    ab1 = (a1 + b1).realize()
+    CL.synchronize()
+
+  # cross device sum (does this work?)
+  # is this making a copy first? is that copy through the CPU?
+  # the slowness comes from the *blocking* clprg call, is this pyopencl?
+  with Timing(colored("0+1 -> 0 (sum): ", "red"), on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    abx0 = (a0 + b1).realize()
+    CL.synchronize()
+
+  with Timing(colored("1+0 -> 1 (sum): ", "red"), on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    abx1 = (b1 + a0).realize()
+    CL.synchronize()
+
+  # copy back
+  # NOTE: half of this slowness is caused by allocating memory on the CPU
+  with Timing("0 -> CPU: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    cc0 = ab0.numpy()
+  with Timing("1 -> CPU: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    cc1 = ab1.numpy()
+
+  # same
+  print("testing")
+  np.testing.assert_allclose(cc0, cc1)
+
+  # devices
+  print(ab0)
+  print(ab1)
+  print(abx0)
+  print(abx1)

test/external/external_osx_profiling.py

@@ -10,16 +10,16 @@ prg = CLProgram("test", """__kernel void test(__global float *a, __global float
   int idx = get_global_id(0);
   a[idx] = b[idx] + c[idx];
 }""")
-prg.clprg(CL.cl_queue, [N,], None, a._cl, b._cl, c._cl)
+prg.clprg(CL.cl_queue[0], [N,], None, a._cl, b._cl, c._cl)
 
 t1 = time.monotonic_ns()
-e1 = prg.clprg(CL.cl_queue, [N,], None, a._cl, b._cl, c._cl)
-CL.cl_queue.finish()  # type: ignore
+e1 = prg.clprg(CL.cl_queue[0], [N,], None, a._cl, b._cl, c._cl)
+CL.synchronize()
 t2 = time.monotonic_ns()
 time.sleep(3)
 t3 = time.monotonic_ns()
-e2 = prg.clprg(CL.cl_queue, [N,], None, a._cl, b._cl, c._cl)
-CL.cl_queue.finish()  # type: ignore
+e2 = prg.clprg(CL.cl_queue[0], [N,], None, a._cl, b._cl, c._cl)
+CL.synchronize()
 t4 = time.monotonic_ns()
 
 print(e1.profile.queued)

test/external/external_test_gpu_ast.py

@@ -2,7 +2,7 @@
 import unittest
 import numpy as np
 from tinygrad.ops import LazyOp, ReduceOps, BinaryOps, UnaryOps, MovementOps
-from tinygrad.shape import ShapeTracker, View, ZeroView
+from tinygrad.shape.shapetracker import ShapeTracker, View, ZeroView
 from tinygrad.runtime.ops_gpu import GPUBuffer, CLProgram, CLCodegen
 #from tinygrad.runtime.ops_metal import MetalBuffer as GPUBuffer, MetalProgram as CLProgram, MetalCodegen as CLCodegen
 from tinygrad.helpers import getenv

test/external/external_test_onnx_backend.py

@@ -87,11 +87,11 @@ backend_test.exclude('test_asin_*')
 backend_test.exclude('test_asinh_*')
 backend_test.exclude('test_atan_*')
 backend_test.exclude('test_atanh_*')
-backend_test.exclude('test_cos_*')
-backend_test.exclude('test_cosh_*')
-backend_test.exclude('test_sin_*')
-backend_test.exclude('test_sinh_*')
-backend_test.exclude('test_tan_*')
+# backend_test.include('test_cos_*')
+# backend_test.include('test_cosh_*')
+# backend_test.exclude('test_sin_*')
+# backend_test.include('test_sinh_*')
+# backend_test.include('test_tanh_*')
 
 # no boolean ops (2d, 3d, 4d)
 backend_test.exclude('test_and*')

test/external/external_test_opt.py

@@ -9,21 +9,144 @@ import numpy as np
 import unittest
 from tinygrad.tensor import Tensor, Device
 from tinygrad import nn
+from tinygrad.helpers import getenv
 from tinygrad.nn import optim
 from tinygrad.ops import GlobalCounters, MovementOps, ReduceOps
 from tinygrad.lazy import PUSH_PERMUTES
 
 class CLCache():
+  def __init__(self, allowed=None, strict=False, preclear=True): self.allowed, self.strict, self.preclear = allowed, strict, preclear
   def __enter__(self):
-    gc.collect()
-    for x in [x for x in gc.get_objects() if isinstance(x, Tensor)]:
-      x.realize()
+    if self.preclear:
+      gc.collect()
+      for x in [x for x in gc.get_objects() if isinstance(x, Tensor)]:
+        x.realize()
+      GlobalCounters.reset()
     GlobalCounters.cache = []
     print("cache: entering")
   def __exit__(self, type, value, traceback):
-    print(f"cache: exiting with size {len(GlobalCounters.cache)}")
+    print(f"cache: exiting with size {len(GlobalCounters.cache)}", f"allowed {self.allowed}" if self.allowed is not None else "")
+    if self.allowed is not None:
+      assert len(GlobalCounters.cache) <= self.allowed and (not self.strict or len(GlobalCounters.cache) == self.allowed), "used too many kernels!"
     GlobalCounters.cache = None
 
+from models.convnext import ConvNeXt
+from models.efficientnet import EfficientNet
+from models.resnet import ResNet18
+from models.vit import ViT
+from tinygrad.nn.optim import get_parameters
+
+@unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented")
+class TestInferenceMinKernels(unittest.TestCase):
+  def setUp(self):
+    Tensor.training = False
+
+  @unittest.skipIf(not PUSH_PERMUTES, "this test requires PUSH_PERMUTES")
+  def test_convnext(self):
+    model = ConvNeXt()
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    img = Tensor.randn(1, 3, 224, 224)
+    with CLCache(129):
+      model(img).realize()
+
+  def test_enet(self):
+    model = EfficientNet(getenv("ENET_NUM", 0), has_se=False)
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    img = Tensor.randn(1, 3, 224, 224)
+    with CLCache(51):
+      model.forward(img).realize()
+
+  def test_enet_se(self):
+    model = EfficientNet(getenv("ENET_NUM", 0), has_se=True)
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    img = Tensor.randn(1, 3, 224, 224)
+    # TODO: this seems very high
+    with CLCache(115):
+      model.forward(img).realize()
+
+  def test_resnet(self):
+    model = ResNet18()
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    img = Tensor.randn(1, 3, 224, 224)
+    with CLCache(31): # NOTE: this should be 4 lower
+      model.forward(img).realize()
+
+  def test_vit(self):
+    model = ViT(embed_dim=192, num_heads=3)
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    img = Tensor.randn(1, 3, 224, 224)
+    with CLCache(223): # NOTE: this is way too high
+      out = model.forward(img)
+      assert len(GlobalCounters.cache) == 0, f"ViT prerealized?"
+      out.realize()
+
+  def test_llama(self):
+    from examples.llama import Transformer, onehot_encode
+    args_tiny = {"dim": 512, "multiple_of": 256, "n_heads": 8, "n_layers": 4, "norm_eps": 1e-05, "vocab_size": 1000}
+    model = Transformer(**args_tiny)
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    with CLCache(85):
+      model(onehot_encode([1,2,3,4], vocab_size=args_tiny['vocab_size']), 0).realize()
+
+@unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented")
+class TestOptBinOp(unittest.TestCase):
+  def _test_no_binop_rerun(self, f1, f2=None, allowed=1):
+    a = Tensor.randn(16, 16)
+    b = Tensor.randn(16, 16)
+    with CLCache():
+      c = f1(a, b)
+      if f2 is not None: d = f2(a, b)
+      c.realize()
+      if f2 is not None: d.realize()
+      assert len(GlobalCounters.cache) == allowed, "binop was rerun!"
+    if f2 is not None: np.testing.assert_allclose(c.numpy().ravel(), d.numpy().ravel(), rtol=1e-3, atol=1e-5)
+
+  def test_no_binop_rerun(self): return self._test_no_binop_rerun(lambda a,b: a*b, lambda a,b: (a*b).reshape(16, 16, 1))
+  def test_no_binop_rerun_alt(self): return self._test_no_binop_rerun(lambda a,b: (a*b).reshape(16, 16, 1), lambda a,b: a*b)
+  def test_no_binop_rerun_reduce_broadcast(self): return self._test_no_binop_rerun(lambda a,b: a.sum()+b, lambda a,b: a.sum().reshape(1,1)+b, allowed=2)
+  def test_no_binop_rerun_transposed(self): return self._test_no_binop_rerun(lambda a,b: (a.T*b.T).T, lambda a,b: a*b)
+  def test_no_binop_rerun_mid_reshape(self): return self._test_no_binop_rerun(lambda a,b: (a*b).reshape(256)+a.reshape(256))
+
+  # currently non working tests
+  #def test_no_binop_rerun_preshape(self): return self._test_no_binop_rerun(lambda a,b: a.reshape(16, 16, 1)*b.reshape(16, 16, 1), lambda a,b: a*b)
+  #def test_no_binop_rerun_reduce(self): return self._test_no_binop_rerun(lambda a,b: (a*b).sum(), lambda a,b: (a*b).reshape(16, 16, 1).sum())
+  #def test_no_binop_rerun_reduce_alt(self): return self._test_no_binop_rerun(lambda a,b: a.sum(1)+b[0], lambda a,b: a.sum(1).reshape(1,16)+b[0])
+
+@unittest.skip("elementwise with >1 reduce inputs currently don't fuse")
+@unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented")
+class TestOptReduceLoop(unittest.TestCase):
+  def test_loop_left(self):
+    a = Tensor.randn(16, 16)
+    b = Tensor.randn(16, 16)
+    with CLCache():
+      t = a.sum(0)
+      b = t.reshape(16,1).expand(16,16).sum(0)
+      c = (t+b)
+      c.realize()
+      assert len(GlobalCounters.cache) == 2, "loop left fusion broken"
+
+  def test_loop_right(self):
+    a = Tensor.randn(16, 16)
+    b = Tensor.randn(16, 16)
+    with CLCache():
+      t = a.sum(0)
+      b = t.reshape(16,1).expand(16,16).sum(0)
+      c = (b+t)
+      c.realize()
+      assert len(GlobalCounters.cache) == 2, "loop right fusion broken"
+
+@unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented")
+class TestOptWChild(unittest.TestCase):
+  def test_unrealized_child(self):
+    a = Tensor.randn(16, 16)
+    b = Tensor.randn(16, 16)
+    with CLCache():
+      c = (a*b).sum()
+      d = c+1
+      e = c+2
+      d.realize()
+      assert len(GlobalCounters.cache) == 2, "don't fuse if you have children"
+
 @unittest.skipUnless(Device.DEFAULT == "GPU", "Not Implemented")
 class TestOpt(unittest.TestCase):
   def test_muladd(self):
@@ -51,13 +174,13 @@ class TestOpt(unittest.TestCase):
     with CLCache():
       img_bn = bn(img).realize()
       print(img_bn)
-      assert len(GlobalCounters.cache) == 3, "optimizer didn't fold batchnorm"
+      assert len(GlobalCounters.cache) == 3, f"optimizer didn't fold batchnorm, got {len(GlobalCounters.cache)}"
     Tensor.training = False
 
   def test_fold_conv_sgd(self):
     # TODO: with Tensor.training
     Tensor.training = True
-    img = Tensor.ones(1,3,4,4)
+    img = Tensor.ones(2,3,4,4)
     c1 = nn.Conv2d(3,32,3)
     opt = optim.SGD(optim.get_parameters(c1))
     with CLCache():
@@ -66,7 +189,36 @@ class TestOpt(unittest.TestCase):
       opt.step()
       # TODO: this should be 4, but the sum output child stays around
       # with pushing_permutes it can be 3
-      assert len(GlobalCounters.cache) in [4,5], "optimizer didn't fold conv-backward SGD"
+      # TODO: broken with optim fixes
+      assert len(GlobalCounters.cache) in [4,5,6], f"optimizer didn't fold conv-backward SGD, got {len(GlobalCounters.cache)}"
+    Tensor.training = False
+
+  def test_fold_2convs_sgd(self):
+    # TODO: with Tensor.training
+    Tensor.training = True
+    img = Tensor.ones(2,3,64,64)
+    c1 = nn.Conv2d(3,16,3,bias=False)
+    c2 = nn.Conv2d(16,32,3,bias=False)
+    opt = optim.SGD(optim.get_parameters([c1, c2]))
+    with CLCache(allowed=9):
+      opt.zero_grad()
+      c2(c1(img).relu()).relu().sum().backward()
+      opt.step()
+    Tensor.training = False
+
+  def test_fold_4convs_sgd(self):
+    # TODO: with Tensor.training
+    Tensor.training = True
+    img = Tensor.ones(2,3,64,64)
+    c1 = nn.Conv2d(3,4,3,bias=False)
+    c2 = nn.Conv2d(4,8,3,bias=False)
+    c3 = nn.Conv2d(8,16,3,bias=False)
+    c4 = nn.Conv2d(16,32,3,bias=False)
+    opt = optim.SGD(optim.get_parameters([c1, c2, c3, c4]))
+    with CLCache(allowed=19):
+      opt.zero_grad()
+      c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
+      opt.step()
     Tensor.training = False
 
   def test_fold_conv_batchnorm_sgd(self):
@@ -76,12 +228,11 @@ class TestOpt(unittest.TestCase):
     c1 = nn.Conv2d(3,32,3)
     bn = nn.BatchNorm2d(32, track_running_stats=False)
     opt = optim.SGD(optim.get_parameters([c1, bn]))
-    with CLCache():
+    with CLCache(allowed=18): # this is too high
       img_bn = bn(c1(img)).elu().sum()
       opt.zero_grad()
       img_bn.backward()
       opt.step()
-      assert len(GlobalCounters.cache) in [9,10], "optimizer didn't fold conv-backward batchnorm"
     Tensor.training = False
 
   def test_fold_conv_batchnorm_notrain(self):
@@ -92,7 +243,7 @@ class TestOpt(unittest.TestCase):
     img_conv = bn(c1(img)).relu().realize()
     with CLCache():
       img_conv = bn(c1(img)).relu().realize()
-      assert len(GlobalCounters.cache) == 1, "optimizer didn't fold conv-batchnorm at test time"
+      assert len(GlobalCounters.cache) == 1, f"optimizer didn't fold conv-batchnorm at test time, got {len(GlobalCounters.cache)}"
 
   def test_fold_conv_batchnorm(self):
     Tensor.training = True
@@ -102,7 +253,7 @@ class TestOpt(unittest.TestCase):
     with CLCache():
       img_conv = bn(c1(img)).relu().realize()
       print(img_conv)
-      assert len(GlobalCounters.cache) == 4, "optimizer didn't fold conv-batchnorm"
+      assert len(GlobalCounters.cache) == 4, f"optimizer didn't fold conv-batchnorm, got {len(GlobalCounters.cache)}"
     Tensor.training = False
 
   def test_fold_conv_elu(self):
@@ -132,66 +283,81 @@ class TestOpt(unittest.TestCase):
       print(img_conv)
       assert len(GlobalCounters.cache) == 2, "optimizer didn't fold conv/relu"
 
-  def helper_push_permute_before_reshape(self, t, should_push=True, desired_reshape_arg=None, desired_permute_arg=None):
-    if PUSH_PERMUTES and should_push:
-      assert t.lazydata.op.src[0].op.op == MovementOps.PERMUTE, 'Permute should be pushed before reshape'
-      assert t.lazydata.op.src[0].op.arg == desired_permute_arg, f'Pushed permute arg should be {desired_permute_arg}'
-      assert t.lazydata.op.op == MovementOps.RESHAPE, 'Reshape should be after permute'
-      assert t.lazydata.op.arg == desired_reshape_arg, f'Reshape arg should be {desired_reshape_arg}'
-    else:
-      assert t.lazydata.op.src[0].op.op == MovementOps.RESHAPE, 'Reshape should before permute'
-      assert t.lazydata.op.op == MovementOps.PERMUTE, 'Permute should be after reshape'
+  def test_permute_was_pushed(self):
+    a = Tensor.randn(16, 16, 16)
+    with CLCache():
+      c = a.sum(2)
+      d = c.permute(1,0).contiguous()
+      d.realize()
+      cache_len = len(GlobalCounters.cache)
+    np.testing.assert_allclose(a.numpy().sum(2).transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
+    if PUSH_PERMUTES: assert cache_len == 1, "permute wasn't pushed!"
 
+  def test_permute_was_pushed_though_contract_reshape(self):
+    a = Tensor.randn(4, 4, 4, 4, 4)
+    with CLCache():
+      c = a.sum(-1)
+      d = c.reshape(16,16).permute(1,0).contiguous()
+      d.realize()
+      cache_len = len(GlobalCounters.cache)
+    np.testing.assert_allclose(a.numpy().sum(-1).reshape(16,16).transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
+    if PUSH_PERMUTES: assert cache_len == 1, "permute wasn't pushed!"
 
-  def test_push_permute_before_reshape(self):
-    t = Tensor.ones(1,2,3,4)
-    t = t.reshape(1,2,3*4).permute(2,1,0)
-    self.helper_push_permute_before_reshape(t, should_push=True, desired_reshape_arg=(12,2,1), desired_permute_arg=(2,3,1,0))
+  def test_permute_was_pushed_though_contractw1s_reshape(self):
+    a = Tensor.randn(4, 4, 4, 4, 4)
+    with CLCache():
+      c = a.sum(-1)
+      d = c.reshape(16,1,16).permute(2,1,0).contiguous()
+      d.realize()
+      cache_len = len(GlobalCounters.cache)
+    np.testing.assert_allclose(a.numpy().sum(-1).reshape(16,1,16).transpose(2,1,0), d.numpy(), rtol=1e-3, atol=1e-5)
+    if PUSH_PERMUTES: assert cache_len == 1, "permute wasn't pushed!"
 
-    t = Tensor.ones(1,2,3,4)
-    t = t.reshape(3,1,2,4).permute(3,2,1,0)
-    self.helper_push_permute_before_reshape(t, should_push=False)
+  # TODO: push permute through expansion reshape
+  @unittest.skip("expansion can't push expand permute yet")
+  @unittest.skipIf(not PUSH_PERMUTES, "this test requires PUSH_PERMUTES")
+  def test_permute_was_pushed_through_expand_reshape(self):
+    a = Tensor.randn(16, 16, 16)
+    with CLCache():
+      c = a.sum(2)
+      d = c.reshape(4,4,4,4).permute(2,3,0,1).contiguous()
+      d.realize()
+      cache_len = len(GlobalCounters.cache)
+    np.testing.assert_allclose(a.numpy().sum(2).transpose(1,0).reshape(4,4,4,4), d.numpy(), rtol=1e-3, atol=1e-5)
+    if PUSH_PERMUTES: assert cache_len == 1, "permute wasn't pushed!"
 
-    t = Tensor.ones(1,2,3,1,4,1)
-    t = t.reshape(1,2,3*4).permute(2,1,0)
-    self.helper_push_permute_before_reshape(t, should_push=True, desired_reshape_arg=(12,2,1), desired_permute_arg=(2,3,4,5,1,0))
+  @unittest.skipIf(PUSH_PERMUTES, "this test is broken with PUSH_PERMUTES")
+  def test_no_reduceop_rerun(self):
+    a = Tensor.randn(16, 16, 16)
+    with CLCache():
+      c = a.sum(2)
+      d = a.sum(2).permute(1,0)
+      c.realize()
+      d.realize()
+      cache_len = len(GlobalCounters.cache)
+    np.testing.assert_allclose(c.numpy().transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
+    assert cache_len == 1, "reduceop was rerun!"
 
-    t = Tensor.ones(1,2,3,4)
-    t = t.reshape(1,2,3,1,4).permute(4,3,2,1,0)
-    self.helper_push_permute_before_reshape(t, should_push=False)
+  @unittest.skipIf(PUSH_PERMUTES, "this test is brokem with PUSH_PERMUTES")
+  def test_no_reduceop_rerun_alt(self):
+    a = Tensor.randn(16, 16, 16)
+    with CLCache():
+      c = a.sum(2).permute(1,0)
+      d = a.sum(2)
+      c.realize()
+      d.realize()
+      cache_len = len(GlobalCounters.cache)
+    np.testing.assert_allclose(c.numpy(), d.numpy().transpose(1,0), rtol=1e-3, atol=1e-5)
+    assert cache_len == 1, "reduceop was rerun!"
 
-
-  def test_push_permute_before_reduce(self):
-    t = Tensor.ones(1,2,3,4)
-    t = t.sum(axis=2).permute(2,1,0)
-    if PUSH_PERMUTES:
-      assert t.lazydata.op.src[0].op.src[0].op.op == MovementOps.PERMUTE, 'Permute should be pushed before reduce'
-      assert t.lazydata.op.src[0].op.src[0].op.arg == (3,1,0,2), 'Pushed permute arg error'
-      assert t.lazydata.op.src[0].op.op == ReduceOps.SUM, 'Sum should be after permute'
-      assert t.lazydata.op.src[0].op.arg == (4,2,1,1), 'Sum arg error'
-      assert t.lazydata.op.op == MovementOps.RESHAPE, 'Reshape should be after Sum'
-      assert t.lazydata.op.arg == (4,2,1), 'Reshape arg error'
-    else:
-      assert t.lazydata.op.src[0].op.src[0].op.op == ReduceOps.SUM, 'Sum should be the first'
-      assert t.lazydata.op.src[0].op.src[0].op.arg == (1,2,4,1), 'Sum arg error'
-      assert t.lazydata.op.src[0].op.op == MovementOps.RESHAPE, 'Reshape should be after sum'
-      assert t.lazydata.op.src[0].op.arg == (1,2,4), 'Reshape arg error'
-      assert t.lazydata.op.op == MovementOps.PERMUTE, 'Permute should be after Reshape'
-      assert t.lazydata.op.arg == (2,1,0), 'Permute arg error'
-
-  def test_push_permute_before_expand(self):
-    t = Tensor.ones(1,2,3,4)
-    t = t.expand(2,2,3,4).permute(3,2,1,0)
-    if PUSH_PERMUTES:
-      assert t.lazydata.op.src[0].op.op == MovementOps.PERMUTE, 'Permute should be pushed before reduce'
-      assert t.lazydata.op.src[0].op.arg == (3,2,1,0), 'Pushed permute arg error'
-      assert t.lazydata.op.op == MovementOps.EXPAND, 'Expand should be after permute'
-      assert t.lazydata.op.arg == (4,3,2,2), 'Expand arg error'
-    else:
-      assert t.lazydata.op.src[0].op.op == MovementOps.EXPAND, 'Expand should be the first'
-      assert t.lazydata.op.src[0].op.arg == (2,2,3,4), 'Expand arg error'
-      assert t.lazydata.op.op == MovementOps.PERMUTE, 'Permute should be after expand'
-      assert t.lazydata.op.arg == (3,2,1,0), 'Permute arg error'
+  def test_fold_with_contiguous(self):
+    a = Tensor.randn(16, 16, 16)
+    b = Tensor.randn(16, 16)
+    with CLCache():
+      c = (a.sum(2).contiguous() + b).contiguous()
+      c.realize()
+      cache_len = len(GlobalCounters.cache)
+    assert cache_len == 1, "contiguous wasn't folded"
 
 if __name__ == '__main__':
   unittest.main()

test/external/fuzz_shapetracker.py

@@ -0,0 +1,60 @@
+import random
+from test.unit.test_shapetracker import CheckingShapeTracker
+
+def do_permute(st):
+  perm = list(range(0, len(st.shape)))
+  random.shuffle(perm)
+  perm = tuple(perm)
+  print("st.permute(", perm, ")")
+  st.permute(perm)
+
+def do_pad(st):
+  c = random.randint(0, len(st.shape)-1)
+  pad = tuple((random.randint(0,2), random.randint(0,2)) if i==c else (0,0) for i in range(len(st.shape)))
+  print("st.pad(", pad, ")")
+  st.pad(pad)
+
+def do_reshape_split_one(st):
+  c = random.randint(0, len(st.shape)-1)
+  poss = [n for n in [1,2,3,4,5] if st.shape[c]%n == 0]
+  spl = random.choice(poss)
+  shp = st.shape[0:c] + (st.shape[c]//spl, spl) + st.shape[c+1:]
+  print("st.reshape(", shp, ")")
+  st.reshape(shp)
+
+def do_reshape_combine_two(st):
+  if len(st.shape) < 2: return
+  c = random.randint(0, len(st.shape)-2)
+  shp = st.shape[:c] + (st.shape[c] * st.shape[c+1], ) + st.shape[c+2:]
+  print("st.reshape(", shp, ")")
+  st.reshape(shp)
+
+def do_shrink(st):
+  c = random.randint(0, len(st.shape)-1)
+  while 1:
+    shrink = tuple((random.randint(0,s), random.randint(0,s)) if i == c else (0,s) for i,s in enumerate(st.shape))
+    if all(x<y for (x,y) in shrink): break
+  print("st.shrink(", shrink, ")")
+  st.shrink(shrink)
+
+def do_stride(st):
+  c = random.randint(0, len(st.shape)-1)
+  stride = tuple(random.choice([-2,-1,2]) if i==c else 1 for i in range(len(st.shape)))
+  print("st.stride(", stride, ")")
+  st.stride(stride)
+
+def do_expand(st):
+  c = [i for i,s in enumerate(st.shape) if s==1]
+  if len(c) == 0: return
+  c = random.choice(c)
+  expand = tuple(random.choice([2,3,4]) if i==c else s for i,s in enumerate(st.shape))
+  print("st.expand(", expand, ")")
+  st.expand(expand)
+
+if __name__ == "__main__":
+  ops = [do_permute, do_pad, do_shrink, do_reshape_split_one, do_reshape_combine_two, do_stride, do_expand]
+  while 1:
+    st = CheckingShapeTracker((3, 3, 3))
+    for i in range(8): random.choice(ops)(st)
+    #st.simplify()
+    st.assert_same()

test/external/fuzz_symbolic.py

@@ -0,0 +1,50 @@
+import random
+from tinygrad.shape.symbolic import Variable
+
+def add_v(expr, rng=None):
+  if rng is None: rng = random.randint(0,2)
+  return expr + v[rng], rng
+
+def div(expr, rng=None):
+  if rng is None: rng = random.randint(1,9)
+  return expr // rng, rng
+
+def mul(expr, rng=None):
+  if rng is None: rng = random.randint(-4,4)
+  return expr * rng, rng
+
+def mod(expr, rng=None):
+  if rng is None: rng = random.randint(1,9)
+  return expr % rng, rng
+
+def add_num(expr, rng=None):
+  if rng is None: rng = random.randint(-4,4)
+  return expr + rng, rng
+
+if __name__ == "__main__":
+  ops = [add_v, div, mul, add_num]
+  while 1:
+    u1 = Variable("v1", 0, 2)
+    u2 = Variable("v2", 0, 3)
+    u3 = Variable("v3", 0, 4)
+    v = [u1,u2,u3]
+    tape = [random.choice(ops) for _ in range(20)]
+    expr = Variable.num(0)
+    rngs = []
+    for t in tape:
+      expr, rng = t(expr)
+      print(t.__name__, rng)
+      rngs.append(rng)
+    print(expr)
+    for v1 in range(u1.min, u1.max+1):
+      for v2 in range(u2.min, u2.max+1):
+        for v3 in range(u3.min, u3.max+1):
+          v = [v1,v2,v3]
+          rn = 0
+          for t,r in zip(tape, rngs):
+            rn, _ = t(rn, r)
+          num = eval(expr.render())
+          assert num == rn, f"mismatch at {v1} {v2} {v3}, {num} != {rn}"
+          #print(v1, v2, v3, num, rn)
+
+

test/extra/test_utils.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python
+import io
+import unittest
+from tinygrad.helpers import getenv
+from extra.utils import fetch, fake_torch_load_zipped
+from PIL import Image
+
+@unittest.skipIf(getenv("CI", "") != "", "no internet tests in CI")
+class TestFetch(unittest.TestCase):
+  def test_fetch_bad_http(self):
+    self.assertRaises(AssertionError, fetch, 'http://httpstat.us/500')
+    self.assertRaises(AssertionError, fetch, 'http://httpstat.us/404')
+    self.assertRaises(AssertionError, fetch, 'http://httpstat.us/400')
+
+  def test_fetch_small(self):
+    assert(len(fetch('https://google.com'))>0)
+
+  def test_fetch_img(self):
+    img = fetch("https://media.istockphoto.com/photos/hen-picture-id831791190")
+    pimg = Image.open(io.BytesIO(img))
+    assert pimg.size == (705, 1024)
+
+class TestUtils(unittest.TestCase):
+  def test_fake_torch_load_zipped(self):
+    import torch
+    import numpy as np
+    import tempfile
+    class LayerWithOffset(torch.nn.Module):
+      def __init__(self):
+        super(LayerWithOffset, self).__init__()
+        d = torch.randn(16)
+        self.param1 = torch.nn.Parameter(
+          d.as_strided([2, 2], [2, 3], storage_offset=5)
+        )
+        self.param2 = torch.nn.Parameter(
+          d.as_strided([2, 2], [2, 3], storage_offset=4)
+        )
+
+    for isfloat16 in [True, False]:
+      model = torch.nn.Sequential(
+        torch.nn.Linear(4, 8),
+        torch.nn.Linear(8, 3),
+        LayerWithOffset()
+      )
+      if isfloat16: model = model.half()
+
+      with tempfile.TemporaryDirectory() as tmpdirname:
+        path = tmpdirname + '/testloadmodel.pth'
+        torch.save(model.state_dict(), path)
+        model2 = fake_torch_load_zipped(path)
+
+      for name, a in model.state_dict().items():
+        b = model2[name]
+        a, b = a.numpy(), b.numpy()
+        assert a.shape == b.shape
+        assert a.dtype == b.dtype
+        assert np.array_equal(a, b)
+
+if __name__ == '__main__':
+  unittest.main()

test/models/test_efficientnet.py

@@ -6,9 +6,11 @@ import unittest
 import numpy as np
 from PIL import Image
 
+from tinygrad.helpers import getenv
+from tinygrad.tensor import Tensor
 from models.efficientnet import EfficientNet
 from models.vit import ViT
-from tinygrad.tensor import Tensor
+from models.resnet import ResNet50
 
 def _load_labels():
   labels_filename = pathlib.Path(__file__).parent / 'efficientnet/imagenet1000_clsidx_to_labels.txt'
@@ -54,7 +56,7 @@ car_img = Image.open(pathlib.Path(__file__).parent / 'efficientnet/car.jpg')
 class TestEfficientNet(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
-    cls.model = EfficientNet(number=0)
+    cls.model = EfficientNet(number=getenv("NUM"))
     cls.model.load_from_pretrained()
 
   @classmethod
@@ -91,5 +93,23 @@ class TestViT(unittest.TestCase):
     label = _infer(self.model, car_img)
     self.assertEqual(label, "racer, race car, racing car")
 
+class TestResNet(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.model = ResNet50()
+    cls.model.load_from_pretrained()
+
+  @classmethod
+  def tearDownClass(cls):
+    del cls.model
+
+  def test_chicken(self):
+    label = _infer(self.model, chicken_img)
+    self.assertEqual(label, "hen")
+
+  def test_car(self):
+    label = _infer(self.model, car_img)
+    self.assertEqual(label, "sports car, sport car")
+
 if __name__ == '__main__':
   unittest.main()

test/models/test_end2end.py

@@ -0,0 +1,160 @@
+import torch
+from torch import nn
+import unittest
+import numpy as np
+from tinygrad.nn import optim, Linear, Conv2d, BatchNorm2d
+from tinygrad.tensor import Tensor
+from datasets import fetch_mnist
+
+def compare_tiny_torch(model, model_torch, X, Y):
+  Tensor.training = True
+  model_torch.train()
+  model_state_dict = optim.get_state_dict(model)
+  for k,v in model_torch.named_parameters():
+    print(f"initting {k} from torch")
+    model_state_dict[k].assign(Tensor(v.detach().numpy())).realize()
+
+  optimizer = optim.SGD(optim.get_parameters(model), lr=0.01)
+  optimizer_torch = torch.optim.SGD(model_torch.parameters(), lr=0.01)
+
+  Xt = torch.Tensor(X.numpy())
+  np.testing.assert_allclose(X.numpy(), Xt.detach().numpy())
+
+  out = model(X)
+  loss = (out * Y).mean()
+  print(loss.realize().numpy()[0])
+
+  out_torch = model_torch(torch.Tensor(X.numpy()))
+  loss_torch = (out_torch * torch.Tensor(Y.numpy())).mean()
+  print(loss_torch.detach().numpy())
+
+  # assert losses match
+  np.testing.assert_allclose(loss.realize().numpy()[0], loss_torch.detach().numpy(), atol=1e-4)
+
+  # zero and backward
+  optimizer.zero_grad()
+  loss.backward()
+  optimizer_torch.zero_grad()
+  loss_torch.backward()
+
+  for k,v in list(model_torch.named_parameters())[::-1]:
+    g = model_state_dict[k].grad.numpy()
+    gt = v.grad.detach().numpy()
+    print("testing grads", k)
+    np.testing.assert_allclose(g, gt, atol=1e-3, err_msg=f'grad mismatch {k}')
+
+  # take the steps
+  optimizer.step()
+  optimizer_torch.step()
+
+  # assert weights match (they don't!)
+  for k,v in model_torch.named_parameters():
+    print("testing weight", k)
+    np.testing.assert_allclose(model_state_dict[k].numpy(), v.detach().numpy(), atol=1e-3, err_msg=f'weight mismatch {k}')
+
+def get_mnist_data():
+  X_train, Y_train, X_test, Y_test = fetch_mnist()
+  BS = 32
+  num_classes = 10
+  X = Tensor(X_test[0:BS].astype(np.float32))
+  Y = np.zeros((BS, num_classes), np.float32)
+  Y[range(BS),Y_test[0:BS]] = -1.0*num_classes
+  return X, Tensor(Y)
+
+class TestEnd2End(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.X, cls.Y = get_mnist_data()
+
+  def test_linear_mnist(self):
+    class LinTiny:
+      def __init__(self, has_batchnorm=False):
+        self.l1 = Linear(784, 128)
+        self.l2 = Linear(128, 10)
+        self.bn1 = BatchNorm2d(128) if has_batchnorm else lambda x: x
+      def __call__(self, x):
+        return self.l2(self.l1(x)).relu().log_softmax(-1)
+    class LinTorch(nn.Module):
+      def __init__(self, has_batchnorm=False):
+        super().__init__()
+        self.l1 = nn.Linear(784, 128)
+        self.l2 = nn.Linear(128, 10)
+      def forward(self, x):
+        return self.l2(self.l1(x)).relu().log_softmax(-1)
+    compare_tiny_torch(LinTiny(), LinTorch(), self.X, self.Y)
+
+  def test_bn_mnist(self):
+    class LinTiny:
+      def __init__(self):
+        self.l1 = Linear(784, 128)
+        self.l2 = Linear(128, 10)
+        self.bn1 = BatchNorm2d(128)
+      def __call__(self, x):
+        return self.l2(self.bn1(self.l1(x).reshape(x.shape[0], -1, 1, 1)).reshape(x.shape[0], -1).relu()).log_softmax(-1)
+    class LinTorch(nn.Module):
+      def __init__(self):
+        super().__init__()
+        self.l1 = nn.Linear(784, 128)
+        self.l2 = nn.Linear(128, 10)
+        self.bn1 = nn.BatchNorm2d(128)
+      def forward(self, x):
+        return self.l2(self.bn1(self.l1(x).reshape(x.shape[0], -1, 1, 1)).reshape(x.shape[0], -1).relu()).log_softmax(-1)
+    compare_tiny_torch(LinTiny(), LinTorch(), self.X, self.Y)
+
+  def test_bn_alone(self):
+    np.random.seed(1337)
+    X = Tensor(np.random.randn(32, 10, 1, 1).astype(np.float32))
+    Y = Tensor(np.random.randn(32, 10, 1, 1).astype(np.float32))
+    compare_tiny_torch(BatchNorm2d(10), nn.BatchNorm2d(10), X, Y)
+
+  def test_bn_linear(self):
+    BS, K = 2, 1
+    eps = 0
+    X = Tensor([1,0]).reshape(BS, K, 1, 1)
+    Y = Tensor([-1,0]).reshape(BS, K, 1, 1)
+    class LinTiny:
+      def __init__(self):
+        self.l1 = Conv2d(K, K, 1, bias=False)
+        self.bn1 = BatchNorm2d(K, affine=False, track_running_stats=False, eps=eps)
+      def __call__(self, x): return self.bn1(self.l1(x))
+    class LinTorch(nn.Module):
+      def __init__(self):
+        super().__init__()
+        self.l1 = nn.Conv2d(K, K, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(K, affine=False, track_running_stats=False, eps=eps)
+      def forward(self, x): return self.bn1(self.l1(x))
+    model_torch = LinTorch()
+    with torch.no_grad():
+      model_torch.l1.weight[:] = 1.
+    compare_tiny_torch(LinTiny(), model_torch, X, Y)
+
+  def test_conv_mnist(self):
+    class LinTiny:
+      def __init__(self, has_batchnorm=False):
+        self.c1 = Conv2d(1, 8, 3, stride=2)
+        self.c2 = Conv2d(8, 16, 3, stride=2)
+        self.l1 = Linear(16*6*6, 10)
+        if has_batchnorm:
+          self.bn1, self.bn2 = BatchNorm2d(8), BatchNorm2d(16)
+        else:
+          self.bn1, self.bn2 = lambda x: x, lambda x: x
+      def __call__(self, x):
+        return self.l1(self.bn2(self.c2(self.bn1(self.c1(x)).relu())).relu().reshape(x.shape[0], -1)).log_softmax(-1)
+    class LinTorch(nn.Module):
+      def __init__(self, has_batchnorm=False):
+        super().__init__()
+        self.c1 = nn.Conv2d(1, 8, 3, stride=2)
+        self.c2 = nn.Conv2d(8, 16, 3, stride=2)
+        self.l1 = nn.Linear(16*6*6, 10)
+        if has_batchnorm:
+          self.bn1, self.bn2 = nn.BatchNorm2d(8), nn.BatchNorm2d(16)
+        else:
+          self.bn1, self.bn2 = lambda x: x, lambda x: x
+      def forward(self, x):
+        return self.l1(self.bn2(self.c2(self.bn1(self.c1(x)).relu())).relu().reshape(x.shape[0], -1)).log_softmax(-1)
+    for has_batchnorm in [False, True]:
+      with self.subTest(has_batchnorm=has_batchnorm):
+        compare_tiny_torch(LinTiny(has_batchnorm), LinTorch(has_batchnorm), self.X.reshape((-1, 1, 28, 28)), self.Y)
+
+if __name__ == "__main__":
+  unittest.main()

test/models/test_mnist.py

@@ -94,13 +94,12 @@ class TestMNIST(unittest.TestCase):
     train(model, X_train, Y_train, optimizer, steps=100)
     assert evaluate(model, X_test, Y_test) > 0.94   # torch gets 0.9415 sometimes
 
-  @unittest.skip("slow and training batchnorm is broken")
   def test_conv_with_bn(self):
     np.random.seed(1337)
     model = TinyConvNet(has_batchnorm=True)
-    optimizer = optim.Adam(model.parameters(), lr=0.001)
-    train(model, X_train, Y_train, optimizer, steps=100)
-    assert evaluate(model, X_test, Y_test) > 0.7 # TODO: batchnorm doesn't work!!!
+    optimizer = optim.AdamW(model.parameters(), lr=0.003)
+    train(model, X_train, Y_train, optimizer, steps=200)
+    assert evaluate(model, X_test, Y_test) > 0.94
 
   def test_sgd(self):
     np.random.seed(1337)
@@ -109,12 +108,5 @@ class TestMNIST(unittest.TestCase):
     train(model, X_train, Y_train, optimizer, steps=600)
     assert evaluate(model, X_test, Y_test) > 0.94   # CPU gets 0.9494 sometimes
 
-  def test_rmsprop(self):
-    np.random.seed(1337)
-    model = TinyBobNet()
-    optimizer = optim.RMSprop(model.parameters(), lr=0.0002)
-    train(model,  X_train, Y_train, optimizer, steps=400)
-    assert evaluate(model, X_test, Y_test) > 0.95
-
 if __name__ == '__main__':
   unittest.main()

test/models/test_rnnt.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+import unittest
+import numpy as np
+from tinygrad.tensor import Tensor
+from models.rnnt import LSTM
+import torch
+
+class TestRNNT(unittest.TestCase):
+  def test_lstm(self):
+    BS, SQ, IS, HS, L = 4, 220, 240, 1024, 2
+
+    # create in torch
+    with torch.no_grad():
+      torch_layer = torch.nn.LSTM(IS, HS, L)
+
+    # create in tinygrad
+    layer = LSTM(IS, HS, L, 0.0)
+
+    # copy weights
+    with torch.no_grad():
+      layer.cells[0].weights_ih.assign(Tensor(torch_layer.weight_ih_l0.numpy()))
+      layer.cells[0].weights_hh.assign(Tensor(torch_layer.weight_hh_l0.numpy()))
+      layer.cells[0].bias_ih.assign(Tensor(torch_layer.bias_ih_l0.numpy()))
+      layer.cells[0].bias_hh.assign(Tensor(torch_layer.bias_hh_l0.numpy()))
+      layer.cells[1].weights_ih.assign(Tensor(torch_layer.weight_ih_l1.numpy()))
+      layer.cells[1].weights_hh.assign(Tensor(torch_layer.weight_hh_l1.numpy()))
+      layer.cells[1].bias_ih.assign(Tensor(torch_layer.bias_ih_l1.numpy()))
+      layer.cells[1].bias_hh.assign(Tensor(torch_layer.bias_hh_l1.numpy()))
+
+    # test initial hidden
+    for _ in range(3):
+      x = Tensor.randn(SQ, BS, IS)
+      z, hc = layer(x, None)
+      torch_x = torch.tensor(x.cpu().numpy())
+      torch_z, torch_hc = torch_layer(torch_x)
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-3, rtol=5e-3)
+
+    # test passing hidden
+    for _ in range(3):
+      x = Tensor.randn(SQ, BS, IS)
+      z, hc = layer(x, hc)
+      torch_x = torch.tensor(x.cpu().numpy())
+      torch_z, torch_hc = torch_layer(torch_x, torch_hc)
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-3, rtol=5e-3)
+
+if __name__ == '__main__':
+  unittest.main()

test/models/test_train.py

@@ -5,6 +5,7 @@ from tinygrad.nn import optim
 from tinygrad.tensor import Device
 from tinygrad.helpers import getenv
 from extra.training import train
+from models.convnext import ConvNeXt
 from models.efficientnet import EfficientNet
 from models.transformer import Transformer
 from models.vit import ViT
@@ -24,18 +25,32 @@ def train_one_step(model,X,Y):
   et = time.time()-st
   print("done in %.2f ms" % (et*1000.))
 
+def check_gc():
+  if Device.DEFAULT == "GPU":
+    from extra.introspection import print_objects
+    assert print_objects() == 0
+
 class TestTrain(unittest.TestCase):
+  def test_convnext(self):
+    model = ConvNeXt(depths=[1], dims=[16])
+    X = np.zeros((BS,3,224,224), dtype=np.float32)
+    Y = np.zeros((BS), dtype=np.int32)
+    train_one_step(model,X,Y)
+    check_gc()
+
   def test_efficientnet(self):
     model = EfficientNet(0)
     X = np.zeros((BS,3,224,224), dtype=np.float32)
     Y = np.zeros((BS), dtype=np.int32)
     train_one_step(model,X,Y)
+    check_gc()
 
   def test_vit(self):
     model = ViT()
     X = np.zeros((BS,3,224,224), dtype=np.float32)
     Y = np.zeros((BS,), dtype=np.int32)
     train_one_step(model,X,Y)
+    check_gc()
 
   def test_transformer(self):
     # this should be small GPT-2, but the param count is wrong
@@ -44,10 +59,7 @@ class TestTrain(unittest.TestCase):
     X = np.zeros((BS,6), dtype=np.float32)
     Y = np.zeros((BS,6), dtype=np.int32)
     train_one_step(model,X,Y)
-
-    if Device.DEFAULT == "GPU":
-      from extra.introspection import print_objects
-      assert print_objects() == 0
+    check_gc()
 
   def test_resnet(self):
     X = np.zeros((BS, 3, 224, 224), dtype=np.float32)
@@ -56,6 +68,7 @@ class TestTrain(unittest.TestCase):
       model = resnet_v()
       model.load_from_pretrained()
       train_one_step(model, X, Y)
+    check_gc()
 
   def test_bert(self):
     # TODO: write this

test/test_conv.py

@@ -38,22 +38,22 @@ class TestConv(unittest.TestCase):
 
   def test_two_binops_no_rerun(self):
     Tensor.no_grad = True
-    x = Tensor.ones(1,12,128,256)
-    w = Tensor.ones(32,12,3,3)
+    x = Tensor.randn(1,12,128,256)
+    w = Tensor.randn(32,12,3,3)
     out = x.conv2d(w, stride=(2,2), padding=(1,1))
     r1, r2 = out.relu(), (out-1)
-    r1.numpy(), r2.numpy()
+    np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0))
+    np.testing.assert_allclose(r2.numpy(), out.numpy() - 1)
     Tensor.no_grad = False
-    # TODO: make this a real test
 
   def test_two_overlapping_binops_no_rerun(self):
     Tensor.no_grad = True
-    x = Tensor.ones(1,12,128,256)
-    w = Tensor.ones(32,12,3,3)
+    x = Tensor.randn(1,12,128,256)
+    w = Tensor.randn(32,12,3,3)
     out = x.conv2d(w, stride=(2,2), padding=(1,1))
     r1, r2 = out.relu(), out.elu()
-    r1.numpy(), r2.numpy()
-    # TODO: make this a real test
+    np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0))
+    np.testing.assert_allclose(r2.numpy(), np.where(out.numpy() > 0, out.numpy(), (np.exp(out.numpy()) - 1)), atol=1e-5)
     Tensor.no_grad = False
 
   def test_first_three(self):

test/test_custom_function.py

@@ -0,0 +1,103 @@
+# this is an example of how you can write terrible DSP compute breaking ops like warpPerspective
+# here we use a CUSTOM op to write atan2
+
+import unittest
+import numpy as np
+from typing import Optional, Tuple
+from tinygrad.helpers import prod, dtypes
+
+# *** first, we implement the atan2 op at the lowest level ***
+# `atan2_gpu` for GPUBuffers and `atan2_cpu` for CPUBuffers
+from tinygrad.lazy import LazyBuffer, create_lazybuffer, Device
+from tinygrad.ops import ASTRunner
+
+# we don't always have GPU support, so the type signature is the abstract CompiledBuffer instead of GPUBuffer
+def atan2_gpu(ret:LazyBuffer, a:LazyBuffer, b:LazyBuffer):
+  assert a.device == "GPU" and b.device == "GPU", "gpu function requires GPUBuffers"
+  assert a.dtype == b.dtype and a.dtype == dtypes.float32, "gpu function only supports float32"
+  ret.realized = Device[ret.device].buffer(prod(ret.shape), ret.dtype)
+  ASTRunner("atan2", """
+    __kernel void atan2(global float *c, global float *a, global float *b) {
+      int idx = get_global_id(0);
+      c[idx] = atan2(a[idx], b[idx]);
+    }""", global_size=[prod(ret.shape)]).build(Device[ret.device].runtime).exec([ret, a, b])
+  return ret.realized
+
+def atan2_cpu(ret:LazyBuffer, a:LazyBuffer, b:LazyBuffer):
+  return Device[ret.device].buffer(np.arctan2(a.realized._buf, b.realized._buf))
+
+# *** second, we write the ATan2 mlop ***
+# NOTE: The derivative of atan2 doesn't need a custom op! https://www.liquisearch.com/atan2/derivative
+# In general, it is also optional to write a backward function, just your backward pass won't work without it
+
+from tinygrad.ops import ASTRunner, LazyOp, LoadOps, BinaryOps, UnaryOps
+from tinygrad.lazy import LazyBuffer
+from tinygrad.tensor import Function
+
+class ATan2(Function):
+  def forward(self, a:LazyBuffer, b:LazyBuffer) -> LazyBuffer:
+    assert prod(a.shape) == prod(b.shape) and a.device == b.device, "shape or device mismatch"
+    self.a, self.b = a, b
+    ast = LazyOp(LoadOps.CUSTOM, (a.contiguous(), b.contiguous()), {"GPU": atan2_gpu, "CPU": atan2_cpu}[a.device])
+    return create_lazybuffer(a.device, a.shape, LoadOps, ast, max(a.dtype, b.dtype))
+  def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]:
+    denom = (self.a.binary_op(BinaryOps.MUL, self.a)).binary_op(BinaryOps.ADD, self.b.binary_op(BinaryOps.MUL, self.b))
+    return grad_output.binary_op(BinaryOps.MUL, self.b.binary_op(BinaryOps.DIV, denom)) if self.needs_input_grad[0] else None, \
+           grad_output.binary_op(BinaryOps.MUL, self.a.const_like(0).binary_op(BinaryOps.SUB, self.a).binary_op(BinaryOps.DIV, denom)) if self.needs_input_grad[1] else None
+
+# *** third, we use our lovely new mlop in some tests ***
+
+from tinygrad.tensor import Tensor, Device
+
+@unittest.skipUnless(Device.DEFAULT in ["CPU", "GPU"], "atan2 is only implemented for CPU and GPU")
+class TestCustomFunction(unittest.TestCase):
+  def test_atan2_forward(self):
+    # create some random Tensors, permute them just because we can
+    a = Tensor.randn(4,4,requires_grad=True).permute(1,0)
+    b = Tensor.randn(4,4,requires_grad=True).permute(1,0)
+
+    # run the forward pass. note: up until the .numpy(), it's all lazy
+    c = ATan2.apply(a, b)
+    print(c.numpy())
+
+    # check the forward pass (in numpy)
+    np.testing.assert_allclose(c.numpy(), np.arctan2(a.numpy(), b.numpy()), atol=1e-5)
+
+  # fun fact, this never actually calls forward, so it works in all the backends
+  def test_atan2_backward(self):
+    # have to go forward before we can go backward
+    a = Tensor.randn(4,4,requires_grad=True).permute(1,0)
+    b = Tensor.randn(4,4,requires_grad=True).permute(1,0)
+    c = ATan2.apply(a, b)
+
+    # run the backward pass
+    c.mean().backward()
+    assert a.grad is not None and b.grad is not None, "tinygrad didn't compute gradients"
+    print(a.grad.numpy())
+    print(b.grad.numpy())
+
+    # check the backward pass (in torch)
+    import torch
+    ta, tb = torch.tensor(a.numpy(), requires_grad=True), torch.tensor(b.numpy(), requires_grad=True)
+    tc = torch.atan2(ta, tb)
+    tc.mean().backward()
+    assert ta.grad is not None and tb.grad is not None, "torch didn't compute gradients"
+    np.testing.assert_allclose(a.grad.numpy(), ta.grad.numpy(), atol=1e-5)
+    np.testing.assert_allclose(b.grad.numpy(), tb.grad.numpy(), atol=1e-5)
+
+  def test_atan2_jit(self):
+    # custom ops even work in the JIT!
+    from tinygrad.jit import TinyJit
+
+    @TinyJit
+    def jitted_atan2(a:Tensor, b:Tensor) -> Tensor:
+      return ATan2.apply(a, b).realize()
+
+    for _ in range(5):
+      a = Tensor.randn(4,4,requires_grad=True).permute(1,0)
+      b = Tensor.randn(4,4,requires_grad=True).permute(1,0)
+      c = jitted_atan2(a, b)
+      np.testing.assert_allclose(c.numpy(), np.arctan2(a.numpy(), b.numpy()), atol=1e-5)
+
+if __name__ == "__main__":
+  unittest.main()

test/test_dtype.py

@@ -0,0 +1,90 @@
+import unittest
+import numpy as np
+from tinygrad.helpers import getenv
+from tinygrad.lazy import Device
+from tinygrad.tensor import Tensor, dtypes
+
+# for GPU, cl_khr_fp16 isn't supported (except now we don't need it!)
+# for LLVM, it segfaults because it can't link to the casting function
+@unittest.skipIf(getenv("CI", "") != "" and Device.DEFAULT in ["LLVM"], "float16 broken in some CI backends")
+class TestDtype(unittest.TestCase):
+  def test_half_to_np(self):
+    a = Tensor([1,2,3,4], dtype=dtypes.float16)
+    print(a)
+    na = a.numpy()
+    print(na, na.dtype, a.lazydata.realized)
+    assert na.dtype == np.float16
+    np.testing.assert_allclose(na, [1,2,3,4])
+
+  def test_half_add(self):
+    a = Tensor([1,2,3,4], dtype=dtypes.float16)
+    b = Tensor([1,2,3,4], dtype=dtypes.float16)
+    c = a+b
+    print(c.numpy())
+    assert c.dtype == dtypes.float16
+    np.testing.assert_allclose(c.numpy(), [2,4,6,8])
+
+  def test_half_mul(self):
+    a = Tensor([1,2,3,4], dtype=dtypes.float16)
+    b = Tensor([1,2,3,4], dtype=dtypes.float16)
+    c = a*b
+    print(c.numpy())
+    assert c.dtype == dtypes.float16
+    np.testing.assert_allclose(c.numpy(), [1,4,9,16])
+
+  def test_half_matmul(self):
+    a = Tensor([[1,2],[3,4]], dtype=dtypes.float16)
+    b = Tensor.eye(2, dtype=dtypes.float16)
+    c = a@b
+    print(c.numpy())
+    assert c.dtype == dtypes.float16
+    np.testing.assert_allclose(c.numpy(), [[1,2],[3,4]])
+
+  def test_upcast_float(self):
+    a = Tensor([1,2,3,4], dtype=dtypes.float16)
+    print(a)
+    fa = a.float()
+    assert a.device == fa.device
+    assert a.requires_grad == fa.requires_grad
+    na = fa.numpy()
+    print(na, na.dtype)
+    assert na.dtype == np.float32
+    np.testing.assert_allclose(na, [1,2,3,4])
+
+  def test_downcast_float(self):
+    a = Tensor([1,2,3,4], dtype=dtypes.float32, requires_grad=False).half()
+    print(a)
+    ha = a.half()
+    assert a.device == ha.device
+    assert a.requires_grad == ha.requires_grad
+    na = ha.numpy()
+    print(na, na.dtype)
+    assert na.dtype == np.float16
+    np.testing.assert_allclose(na, [1,2,3,4])
+
+  def test_half_add_upcast(self):
+    a = Tensor([1,2,3,4], dtype=dtypes.float16)
+    b = Tensor([1,2,3,4], dtype=dtypes.float32)
+    c = a+b
+    print(c.numpy())
+    assert c.dtype == dtypes.float32
+    np.testing.assert_allclose(c.numpy(), [2,4,6,8])
+
+  def test_half_mul_upcast(self):
+    a = Tensor([1,2,3,4], dtype=dtypes.float16)
+    b = Tensor([1,2,3,4], dtype=dtypes.float32)
+    c = a*b
+    print(c.numpy())
+    assert c.dtype == dtypes.float32
+    np.testing.assert_allclose(c.numpy(), [1,4,9,16])
+
+  def test_half_matmul_upcast(self):
+    a = Tensor([[1,2],[3,4]], dtype=dtypes.float16)
+    b = Tensor.eye(2, dtype=dtypes.float32)
+    c = a@b
+    print(c.numpy())
+    assert c.dtype == dtypes.float32
+    np.testing.assert_allclose(c.numpy(), [[1,2],[3,4]])
+
+if __name__ == '__main__':
+  unittest.main()

test/test_example.py

@@ -1,24 +0,0 @@
-import unittest
-from tinygrad.tensor import Tensor
-
-class TestExample(unittest.TestCase):
-  def test_example_readme(self):
-    x = Tensor.eye(3, requires_grad=True)
-    y = Tensor([[2.0,0,-2.0]], requires_grad=True)
-    z = y.matmul(x).sum()
-    z.backward()
-
-    print(x.grad.numpy())  # dz/dx
-    print(y.grad.numpy())  # dz/dy
-
-  def test_example_matmul(self):
-    x = Tensor.eye(256, requires_grad=True)
-    y = Tensor.eye(256, requires_grad=True)
-    z = y.matmul(x).sum()
-    z.backward()
-
-    print(x.grad.numpy())  # dz/dx
-    print(y.grad.numpy())  # dz/dy
-
-if __name__ == '__main__':
-  unittest.main()

test/test_jit.py

@@ -9,16 +9,35 @@ class TestJit(unittest.TestCase):
   def test_simple_jit(self):
     @TinyJit
     def add(a, b): return (a+b).realize()
-    for _ in range(3):
+    for _ in range(5):
       a = Tensor.randn(10, 10)
       b = Tensor.randn(10, 10)
       c = add(a, b)
       np.testing.assert_equal(c.numpy(), a.numpy()+b.numpy())
 
+  def test_jit_shape_mismatch(self):
+    @TinyJit
+    def add(a, b): return (a+b).realize()
+    for _ in range(3):
+      a = Tensor.randn(10, 10)
+      b = Tensor.randn(10, 10)
+      c = add(a, b)
+    bad = Tensor.randn(20, 20)
+    with self.assertRaises(AssertionError):
+      add(a, bad)
+
+  def test_jit_duplicate_fail(self):
+    # the jit doesn't support duplicate arguments
+    @TinyJit
+    def add(a, b): return (a+b).realize()
+    a = Tensor.randn(10, 10)
+    with self.assertRaises(AssertionError):
+      add(a, a)
+
   def test_kwargs_jit(self):
     @TinyJit
     def add_kwargs(first, second): return (first+second).realize()
-    for _ in range(3):
+    for _ in range(5):
       a = Tensor.randn(10, 10)
       b = Tensor.randn(10, 10)
       c = add_kwargs(first=a, second=b)
@@ -27,16 +46,29 @@ class TestJit(unittest.TestCase):
   def test_array_jit(self):
     @TinyJit
     def add_array(a, arr): return (a+arr[0]).realize()
-    for i in range(3):
+    for i in range(5):
       a = Tensor.randn(10, 10)
       b = Tensor.randn(10, 10)
       a.realize(), b.realize()
       c = add_array(a, [b])
-      if i == 2:
+      if i >= 2:
         # should fail once jitted since jit can't handle arrays
         np.testing.assert_equal(np.any(np.not_equal(c.numpy(),a.numpy()+b.numpy())), True)
       else:
         np.testing.assert_equal(c.numpy(), a.numpy()+b.numpy())
 
+  def test_method_jit(self):
+    class Fun:
+      def __init__(self):
+        self.a = Tensor.randn(10, 10)
+      @TinyJit
+      def __call__(self, b:Tensor) -> Tensor:
+        return (self.a+b).realize()
+    fun = Fun()
+    for _ in range(5):
+      b = Tensor.randn(10, 10)
+      c = fun(b)
+      np.testing.assert_equal(c.numpy(), fun.a.numpy()+b.numpy())
+
 if __name__ == '__main__':
   unittest.main()

test/test_nn.py

@@ -2,7 +2,7 @@
 import unittest
 import numpy as np
 from tinygrad.tensor import Tensor, Device
-from tinygrad.nn import BatchNorm2d, Conv2d, Linear, GroupNorm, LayerNorm
+from tinygrad.nn import BatchNorm2d, Conv2d, Linear, GroupNorm, LayerNorm, LayerNorm2d
 import torch
 
 class TestNN(unittest.TestCase):
@@ -76,7 +76,7 @@ class TestNN(unittest.TestCase):
   def test_conv2d(self):
     BS, C1, H, W = 4, 16, 224, 224
     C2, K, S, P = 64, 7, 2, 1
-    
+
     # create in tinygrad
     layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
 
@@ -131,5 +131,24 @@ class TestNN(unittest.TestCase):
     torch_z = torch_layer(torch_x)
     np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-3, rtol=5e-3)
 
+  def test_layernorm_2d(self):
+    N, C, H, W = 20, 5, 10, 10
+
+    # create in tinygrad
+    layer = LayerNorm2d(C)
+
+    # create in torch
+    with torch.no_grad():
+      torch_layer = torch.nn.LayerNorm([C]).eval()
+      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
+      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
+
+    # test
+    x = Tensor.randn(N, C, H, W)
+    z = layer(x)
+    torch_x = torch.tensor(x.cpu().numpy())
+    torch_z = torch_layer(torch_x.permute(0,2,3,1)).permute(0,3,1,2)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-3, rtol=5e-3)
+
 if __name__ == '__main__':
   unittest.main()

test/test_ops.py

@@ -4,8 +4,10 @@ import numpy as np
 import unittest
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import getenv, IMAGE
+from tinygrad.lazy import Device
 
 FORWARD_ONLY = getenv("FORWARD_ONLY", 0)
+PRINT_TENSORS = getenv("PRINT_TENSORS", 0)
 def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, grad_atol=1e-4, grad_rtol=1e-3, forward_only=False, vals=None, a=-0.5, b=3):
   if tinygrad_fxn is None: tinygrad_fxn = torch_fxn
   torch.manual_seed(0)
@@ -26,7 +28,8 @@ def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, gra
   tinygrad_fp = time.monotonic() - st
 
   def compare(s, x,y,atol,rtol):
-    if y.shape != tuple(): assert x.shape == y.shape, f"shape mismatch {x.shape} != {y.shape}"
+    if PRINT_TENSORS: print(s, x, y)
+    if y.shape != tuple(): assert x.shape == y.shape, f"shape mismatch (tinygrad){x.shape} != (torch){y.shape}"
     try:
       np.testing.assert_allclose(x,y, atol=atol, rtol=rtol)
     except Exception:
@@ -37,11 +40,11 @@ def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, gra
   torch_fbp, tinygrad_fbp = np.nan, np.nan
   if not forward_only and not FORWARD_ONLY:
     st = time.monotonic()
-    out.square().mean().backward()
+    (out+1).square().mean().backward()
     torch_fbp = time.monotonic() - st
 
     st = time.monotonic()
-    ret.square().mean().backward()
+    (ret+1).square().mean().backward()
     for tt in tst: tt.grad.realize()
     tinygrad_fbp = time.monotonic() - st
 
@@ -67,7 +70,7 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x,y: fxn(x,2), lambda x,y: fxn(x,2), forward_only=True, vals=[[0.,1,2], [2.,1,0]])
     if reverse: helper_test_op(None, lambda x,y: fxn(2,y), lambda x,y: fxn(2,y), forward_only=True, vals=[[0.,1,2], [2.,1,0]])
 
-  def test_cmp_eq(self): self._test_cmp(lambda x,y: x.eq(y), reverse=False)
+  def test_cmp_eq(self): self._test_cmp(lambda x,y: x==y, reverse=False)
   def test_cmp_gt(self): self._test_cmp(lambda x,y: x>y)
   def test_cmp_ge(self): self._test_cmp(lambda x,y: x>=y)
   def test_cmp_lt(self): self._test_cmp(lambda x,y: x<y)
@@ -91,6 +94,7 @@ class TestOps(unittest.TestCase):
 
   def test_maximum(self):
     helper_test_op([(45,65), (45,65)], torch.maximum, Tensor.maximum)
+    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., 4.], [1., 2., 3., 0.]])
   def test_minimum(self):
     helper_test_op([(45,65), (45,65)], torch.minimum, Tensor.minimum)
   def test_add(self):
@@ -103,6 +107,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65), (65,)], lambda x,y: x+y, lambda x,y: x+y)
   def test_sub(self):
     helper_test_op([(45,65), (45,65)], lambda x,y: x-y, Tensor.sub)
+  def test_neg(self):
+    helper_test_op([(45,65)], lambda x: -x)
   def test_mul(self):
     helper_test_op([(64,64), (64,64)], lambda x,y: x*y, Tensor.mul)
   def test_div(self):
@@ -125,10 +131,25 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: 2.0**x, lambda x: 2.0**x)
   def test_sqrt(self):
     helper_test_op([(45,65)], lambda x: x.sqrt(), Tensor.sqrt, a=0)
+  
+  def test_sin(self):
+    helper_test_op([(45,65)], lambda x: x.sin(), Tensor.sin, a=0)
+  def test_cos(self):
+    helper_test_op([(45,65)], lambda x: x.cos(), Tensor.cos, a=0)
+  def test_tan(self):
+    helper_test_op([(45,65)], lambda x: x.tan(), Tensor.tan, a=0)
+    
   def test_relu(self):
-    helper_test_op([(45,65)], lambda x: x.relu(), Tensor.relu)
+    helper_test_op([(64,64)], lambda x: x.relu(), Tensor.relu)
+  def test_relu_exact(self):
+    helper_test_op(None, lambda x: x.relu(), Tensor.relu, vals=[[-1.,0,1]])
+  def test_relu_maximum_exact(self):
+    helper_test_op(None, lambda x: torch.maximum(x, torch.zeros_like(x, requires_grad=False)), lambda x: Tensor.maximum(x, 0), vals=[[-1.,0,1]])
   def test_leakyrelu(self):
     helper_test_op([(45,65)], lambda x: torch.nn.functional.leaky_relu(x,0.01), Tensor.leakyrelu)
+  def test_celu(self):
+    for val in range(1, 5):
+      helper_test_op([(45,65)], lambda x: torch.nn.functional.celu(x,val), lambda x: x.celu(val))
   def test_abs(self):
     helper_test_op([(45,65)], lambda x: torch.abs(x), Tensor.abs)
   def test_log(self):
@@ -137,6 +158,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: torch.exp(x), Tensor.exp)
   def test_sign(self):
     helper_test_op([(45,65)], lambda x: torch.sign(x), Tensor.sign)
+  def test_softsign(self):
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.softsign(x), Tensor.softsign)
   def test_sigmoid(self):
     helper_test_op([(45,65)], lambda x: x.sigmoid(), Tensor.sigmoid)
   def test_softplus(self):
@@ -195,8 +218,19 @@ class TestOps(unittest.TestCase):
                 [[1.0,1.0,0.0,1.0]],
                 ])
     helper_test_op([(3,4,5,6)], lambda x: x.max(axis=1)[0], lambda x: Tensor.max(x, axis=1))
+  def test_mean(self):
+    helper_test_op([(3,4,5,6)], lambda x: x.mean())
   def test_mean_axis(self):
     helper_test_op([(3,4,5,6)], lambda x: x.mean(axis=(1,2)), lambda x: Tensor.mean(x, axis=(1,2)))
+  def test_std(self):
+    helper_test_op([(45, 65, 85)], lambda x: torch.std(x, unbiased=False), lambda x: Tensor.std(x))
+  def test_std_axis(self):
+    helper_test_op([(45, 65, 85)], lambda x: torch.std(x, unbiased=False, dim=0), lambda x: Tensor.std(x, axis=0))
+    helper_test_op([(45, 65, 85)], lambda x: torch.std(x, unbiased=False, dim=2), lambda x: Tensor.std(x, axis=2))
+    helper_test_op([(45, 65, 85)], lambda x: torch.std(x, unbiased=False, dim=[1, 2]), lambda x: Tensor.std(x, axis=[1, 2]))
+    helper_test_op([(45, 65, 85)], lambda x: torch.std(x, unbiased=False, dim=None), lambda x: Tensor.std(x, axis=None))
+  def test_std_keepdim(self):
+    helper_test_op([(45, 65, 85)], lambda x: torch.std(x, keepdim=True), lambda x: Tensor.std(x, keepdim=True))
   def test_log_softmax(self):
     helper_test_op([(45,65)], lambda x: torch.nn.LogSoftmax(dim=1)(x), Tensor.log_softmax, atol=1e-7, grad_atol=1e-7)
   def test_log_softmax_other_axis(self):
@@ -205,6 +239,9 @@ class TestOps(unittest.TestCase):
     helper_test_op([(10,10,10)], lambda x: x.log_softmax(2), lambda x: x.log_softmax(2), atol=1e-7, grad_atol=1e-7)
   def test_tanh(self):
     helper_test_op([(45,65)], lambda x: x.tanh(), Tensor.tanh, atol=1e-6, grad_atol=1e-6)
+  def test_hardtanh(self):
+    for val in range(10, 30, 5):
+      helper_test_op([(45,65)], lambda x: torch.nn.functional.hardtanh(x,-val, val), lambda x: x.hardtanh(-val, val), atol=1e-6, grad_atol=1e-6)
   def test_topo_sort(self):
     helper_test_op([(45,65)], lambda x: (x+x)*x, lambda x: x.add(x).mul(x), atol=1e-6, grad_atol=1e-6)
 
@@ -218,6 +255,10 @@ class TestOps(unittest.TestCase):
   def test_scalar_rsub(self):
     helper_test_op([(45,65)], lambda x: 2-x, lambda x: 2-x)
 
+  def test_flip_eye_crash(self):
+    helper_test_op([], lambda: (torch.eye(10)@torch.eye(10).flip(0)),
+                       lambda: (Tensor.eye(10)@Tensor.eye(10).flip(0)), forward_only=True)
+
   def test_broadcast_full(self):
     for torch_op, tinygrad_op in [(torch.add, Tensor.add), (torch.sub, Tensor.sub), (torch.mul, Tensor.mul),
                                   (torch.div, Tensor.div), (torch.pow, Tensor.pow)]:
@@ -255,10 +296,10 @@ class TestOps(unittest.TestCase):
     helper_test_op([(3,3,3,3)], lambda x: torch.nn.functional.pad(x, (1,2,3,4)), lambda x: x.pad2d(padding=(1,2,3,4)))
 
   def test_transpose(self):
-    helper_test_op([(3,3,3)], lambda x: x.transpose(1,2), lambda x: x.transpose(order=(0,2,1)))
-    helper_test_op([(3,3,3)], lambda x: x.transpose(0,2), lambda x: x.transpose(order=(2,1,0)))
-    helper_test_op([(1,2,3,4)], lambda x: x.movedim((3,0,2,1),(0,1,2,3)), lambda x: x.transpose(order=(3,0,2,1)))
-    helper_test_op([(3,4,5,6)], lambda x: x.movedim((3,2,1,0),(0,1,2,3)), lambda x: x.transpose(order=(3,2,1,0)))
+    helper_test_op([(3,3,3)], lambda x: x.transpose(1,2), lambda x: x.transpose(1,2))
+    helper_test_op([(3,3,3)], lambda x: x.transpose(0,2), lambda x: x.transpose(0,2))
+    helper_test_op([(1,2,3,4)], lambda x: x.movedim((3,0,2,1),(0,1,2,3)), lambda x: x.permute(order=(3,0,2,1)))
+    helper_test_op([(3,4,5,6)], lambda x: x.movedim((3,2,1,0),(0,1,2,3)), lambda x: x.permute(order=(3,2,1,0)))
 
   def test_reshape(self):
     helper_test_op([(4,3,6,6)], lambda x: torch.reshape(x, (-1,3,6,6)), lambda x: x.reshape(shape=(-1,3,6,6)))
@@ -269,6 +310,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (0,1)), lambda x: x.flip(axis=(0,1)))
     helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (0,1,3)), lambda x: x.flip(axis=(0,1,3)))
     helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (3,)), lambda x: x.flip(axis=(3,)))
+    helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (0,1,3)).flip((0,)), lambda x: x.flip(axis=(0,1,3)).flip(0))
+    helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (3,)), lambda x: x.flip(axis=(-1,)))
 
   def test_unsqueeze(self):
     helper_test_op([(4,3,6,6)], lambda x: torch.unsqueeze(x, 0), lambda x: x.unsqueeze(dim=0))
@@ -294,7 +337,7 @@ class TestOps(unittest.TestCase):
                     lambda x,w: torch.nn.functional.conv2d(x, w),
                     lambda x,w: x.conv2d(w), atol=1e-2)
 
-  @unittest.skip("not supported with IMAGE=1")
+  @unittest.skip("slow")
   def test_large_bs_conv(self):
     # large batch size can cause OpenCL image to exceed max image height on macOS
     # (or cause the conv kernel to overflow short sampling coords)
@@ -302,7 +345,7 @@ class TestOps(unittest.TestCase):
                     lambda x,w: torch.nn.functional.conv2d(x, w),
                     lambda x,w: x.conv2d(w), atol=1e-4, rtol=1e-2)
 
-  @unittest.skip("not supported with IMAGE=1")
+  @unittest.skip("slow")
   def test_large_ic_conv(self):
     # large input channel count can cause OpenCL image to exceed max image width on macOS
     helper_test_op([(1,2048,3,3), (1,2048,3,3)],
@@ -320,6 +363,33 @@ class TestOps(unittest.TestCase):
       lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
       lambda x,w: Tensor.conv2d(x,w).relu(), atol=1e-4, grad_rtol=1e-5)
 
+  @unittest.skipIf(IMAGE>0, "no conv3d on images")
+  def test_simple_conv3d(self):
+    helper_test_op([(1,4,9,9,9), (4,4,3,3,3)],
+      lambda x,w: torch.nn.functional.conv3d(x,w).relu(),
+      lambda x,w: Tensor.conv2d(x,w).relu(), atol=1e-4, grad_rtol=1e-5)
+
+  @unittest.skipIf(IMAGE>0, "no conv3d on images")
+  def test_padded_conv3d(self):
+    helper_test_op([(1,4,9,9,9), (4,4,3,3,3)],
+      lambda x,w: torch.nn.functional.conv3d(x,w,padding=1).relu(),
+      lambda x,w: Tensor.conv2d(x,w,padding=[1,1,1,1,1,1]).relu(), atol=1e-4, grad_rtol=1e-5)
+
+  def test_simple_conv2d_m4(self):
+    helper_test_op([(1,16,18,18), (16,16,3,3)],
+      lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
+      lambda x,w: Tensor.conv2d(x,w).relu(), atol=1e-4, grad_rtol=1e-5)
+
+  def test_simple_conv2d_1x1(self):
+    helper_test_op([(1,4,9,9), (4,4,1,1)],
+      lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
+      lambda x,w: Tensor.conv2d(x,w).relu(), atol=1e-4, grad_rtol=1e-5)
+
+  def test_simple_conv2d_1x1_m4(self):
+    helper_test_op([(1,16,32,32), (16,16,1,1)],
+      lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
+      lambda x,w: Tensor.conv2d(x,w).relu(), atol=1e-4, grad_rtol=1e-5)
+
   def test_nested_conv2d(self):
     helper_test_op([(1,32,9,9), (32,32,3,3), (32,32,3,3)],
       lambda x,w1,w2: torch.nn.functional.conv2d(torch.nn.functional.conv2d(x,w1).relu(), w2).relu(),
@@ -366,7 +436,7 @@ class TestOps(unittest.TestCase):
     cin = 2
     helper_test_op([(bs,groups*cin,1,1), (groups*rcout,cin,1,1)],
       lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
-      lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), atol=1e-4, grad_rtol=1e-5, forward_only=IMAGE>=2)
+      lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), atol=1e-4, grad_rtol=1e-5)
 
   def test_medium_grouped_conv2d(self):
     bs = 1
@@ -375,7 +445,7 @@ class TestOps(unittest.TestCase):
     cin = 2
     helper_test_op([(bs,groups*cin,1,1), (groups*rcout,cin,1,1)],
       lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
-      lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), atol=1e-4, grad_rtol=1e-5, forward_only=IMAGE>=2)
+      lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), atol=1e-4, grad_rtol=1e-5)
 
   def test_depthwise_conv2d(self):
     bs = 1
@@ -500,6 +570,7 @@ class TestOps(unittest.TestCase):
           lambda x: torch.nn.functional.max_pool2d(x, kernel_size=(2,2), stride=stride),
           lambda x: Tensor.max_pool2d(x, kernel_size=(2,2), stride=stride))
 
+  @unittest.skipIf(Device.DEFAULT == "CUDA", "CUDA fails on this")
   def test_maxpool2d_unit_stride(self):
     helper_test_op([(32,2,110,28)],
       lambda x: torch.nn.functional.max_pool2d(x, kernel_size=(5,5), stride=1),
@@ -533,12 +604,43 @@ class TestOps(unittest.TestCase):
     for dim in range(-1, 2):
       helper_test_op([(45,65), (45,65), (45,65)], lambda x,y,z: torch.cat((x,y,z), dim), lambda x,y,z: x.cat(y, z, dim=dim))
 
+  def test_stack(self):
+    x = Tensor.randn(45, 65, 3)
+
+    for dim in range(-1, 3):
+      helper_test_op([(45, 65, 3), (45, 65, 3), (45, 65, 3)], lambda x, y, z: torch.stack((x, y, z), dim=dim), lambda x, y, z: Tensor.stack([x, y, z], dim=dim))
+
+    with self.assertRaises(IndexError):
+      Tensor.stack([x], dim=77)
+
+  def test_repeat(self):
+    x = Tensor.randn(45, 65, 3)
+    base_repeats = [2, 4, 3]
+
+    for reps in [[], [4], [2, 1], [3, 2, 2]]:
+      repeats = base_repeats + reps
+      helper_test_op([(45, 65, 3)], lambda x: x.repeat(*repeats), lambda x: x.repeat(repeats))
+
+    with self.assertRaises(AssertionError):
+      x.repeat((2, 4))
+
+    with self.assertRaises(AssertionError):
+      x.repeat((2, 0, 4))
+
+
   def test_clip(self):
     helper_test_op([(45,65)], lambda x: x.clip(-2.3, 1.2), lambda x: x.clip(-2.3, 1.2))
 
   def test_matvec(self):
     helper_test_op([(1,128), (128,128), (128,128)], lambda x,y,z: (x@y).relu()@z, atol=1e-4)
 
+  # this was the failure in llama early realizing freqs_cis
+  def test_double_slice(self):
+    helper_test_op([(4,4)], lambda x: x[:, 1:2][1:2])
+    helper_test_op([(4,4)], lambda x: x[1:3][1:2])
+    helper_test_op([(4,4)], lambda x: x[:, 1:2][0:1])
+    helper_test_op([(4,4)], lambda x: x[:, 1:2][:, 0:1])
+
 if __name__ == '__main__':
   np.random.seed(1337)
   unittest.main(verbosity=2)

test/test_optim.py

@@ -2,72 +2,68 @@ import numpy as np
 import torch
 import unittest
 from tinygrad.tensor import Tensor
-from tinygrad.nn.optim import Adam, SGD, RMSprop, get_parameters
-
-x_init = np.random.randn(1,3).astype(np.float32)
-W_init = np.random.randn(3,3).astype(np.float32)
-m_init = np.random.randn(1,3).astype(np.float32)
-
-def step_tinygrad(optim, kwargs={}):
-  net = TinyNet()
-  optim = optim([net.x, net.W], **kwargs)
-  out = net.forward()
-  out.backward()
-  optim.step()
-  return net.x.cpu().numpy(), net.W.cpu().numpy()
-
-def step_pytorch(optim, kwargs={}):
-  net = TorchNet()
-  optim = optim([net.x, net.W], **kwargs)
-  out = net.forward()
-  out.backward()
-  optim.step()
-  return net.x.detach().numpy(), net.W.detach().numpy()
+from tinygrad.nn.optim import Adam, SGD, AdamW
 
+np.random.seed(1337)
+x_init = np.random.randn(1,4).astype(np.float32)
+W_init = np.random.randn(4,4).astype(np.float32)
+m_init = np.random.randn(1,4).astype(np.float32)
 
 class TinyNet():
-  def __init__(self):
-    self.x = Tensor(x_init.copy())
-    self.W = Tensor(W_init.copy())
-    self.m = Tensor(m_init.copy())
-
-  def forward(self):
-    out = self.x.dot(self.W).relu()
-    out = out.log_softmax()
-    out = out.mul(self.m).add(self.m).sum()
-    return out
-
-
-class TorchNet():
-  def __init__(self):
-    self.x = torch.tensor(x_init.copy(), requires_grad=True)
-    self.W = torch.tensor(W_init.copy(), requires_grad=True)
-    self.m = torch.tensor(m_init.copy())
+  def __init__(self, tensor):
+    self.x = tensor(x_init.copy(), requires_grad=True)
+    self.W = tensor(W_init.copy(), requires_grad=True)
+    self.m = tensor(m_init.copy())
 
   def forward(self):
     out = self.x.matmul(self.W).relu()
-    out = torch.nn.functional.log_softmax(out, dim=1)
+    out = out.log_softmax(1)
     out = out.mul(self.m).add(self.m).sum()
     return out
 
+def step(tensor, optim, steps=1, kwargs={}):
+  net = TinyNet(tensor)
+  optim = optim([net.x, net.W], **kwargs)
+  for _ in range(steps):
+    out = net.forward()
+    optim.zero_grad()
+    out.backward()
+    optim.step()
+  return net.x.detach().numpy(), net.W.detach().numpy()
 
 class TestOptim(unittest.TestCase):
 
-  def test_adam(self):
-    for x,y in zip(step_tinygrad(Adam),
-                   step_pytorch(torch.optim.Adam)):
-      np.testing.assert_allclose(x, y, atol=1e-4)
+  def _test_optim(self, tinygrad_optim, torch_optim, steps, opts, atol, rtol):
+    for x,y in zip(step(Tensor, tinygrad_optim, steps, kwargs=opts),
+                   step(torch.tensor, torch_optim, steps, kwargs=opts)):
+      np.testing.assert_allclose(x, y, atol=atol, rtol=rtol)
 
-  def test_sgd(self):
-    for x,y in zip(step_tinygrad(SGD, kwargs={'lr': 0.001}),
-                   step_pytorch(torch.optim.SGD, kwargs={'lr': 0.001})):
-      np.testing.assert_allclose(x, y, atol=1e-5)
+  def _test_sgd(self, steps, opts, atol, rtol): self._test_optim(SGD, torch.optim.SGD, steps, opts, atol, rtol)
+  def _test_adam(self, steps, opts, atol, rtol): self._test_optim(Adam, torch.optim.Adam, steps, opts, atol, rtol)
+  def _test_adamw(self, steps, opts, atol, rtol): self._test_optim(AdamW, torch.optim.AdamW, steps, opts, atol, rtol)
 
-  def test_rmsprop(self):
-    for x,y in zip(step_tinygrad(RMSprop, kwargs={'lr': 0.001, 'decay': 0.99}),
-                   step_pytorch(torch.optim.RMSprop,
-                                kwargs={'lr': 0.001, 'alpha': 0.99})):
-      np.testing.assert_allclose(x, y, atol=1e-5)
+  def test_sgd(self): self._test_sgd(1, {'lr': 0.001}, 1e-6, 0)
+  def test_sgd_high_lr(self): self._test_sgd(1, {'lr': 10}, 1e-6, 1e-5)
+
+  def test_multistep_sgd(self): self._test_sgd(10, {'lr': 0.001}, 1e-6, 0)
+  def test_multistep_sgd_high_lr(self): self._test_sgd(10, {'lr': 10}, 1e-6, 3e-4)
+
+  def test_multistep_sgd_momentum(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9}, 1e-6, 0)
+  def test_multistep_sgd_high_lr_momentum(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9}, 1e-5, 3e-4)
+
+  def test_multistep_sgd_nesterov_momentum(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9, 'nesterov': True}, 1e-5, 0)
+  def test_multistep_sgd_high_lr_nesterov_momentum(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9, 'nesterov': True}, 1e-5, 3e-4)
+
+  def test_adam(self): self._test_adam(1, {'lr': 0.001}, 1e-5, 0)
+  def test_adam_high_lr(self): self._test_adam(1, {'lr': 10}, 1e-5, 1e-5)
+  def test_adamw(self): self._test_adamw(1, {'lr': 0.001}, 1e-5, 0)
+  def test_adamw_high_lr(self): self._test_adamw(1, {'lr': 10}, 1e-5, 1e-5)
+
+  def test_multistep_adam(self): self._test_adam(10, {'lr': 0.001}, 1e-5, 0)
+  def test_multistep_adam_high_lr(self): self._test_adam(10, {'lr': 10}, 1e-5, 3e-4)
+
+  def test_multistep_adamw(self): self._test_adamw(10, {'lr': 0.001}, 1e-5, 0)
+  def test_multistep_adamw_high_lr(self): self._test_adamw(10, {'lr': 10}, 1e-5, 3e-4)
 
 if __name__ == '__main__':
   unittest.main()

test/test_randomness.py

@@ -1,6 +1,7 @@
 import math
 import unittest
 import numpy as np
+import torch
 from tinygrad.tensor import Tensor
 
 # https://gist.github.com/devries/11405101
@@ -37,39 +38,40 @@ def kstest(l1, l2):
   prob = ksprob((nesq + 0.12 + 0.11 / nesq) * d)
   return prob
 
-def equal_distribution(tinygrad_func, numpy_func, shape=(20, 23), alpha=0.05):
-  Tensor.manual_seed(1337)
-  np.random.seed(1337)
-  x = tinygrad_func(*shape).cpu().numpy().flatten()
-  y = numpy_func(shape).flatten()
-  p = kstest(x, y)
-  return p >= alpha
-
 def normal_test(func, shape=(20, 23), alpha=0.05):
-  y = lambda x: np.random.randn(*x)
-  p = equal_distribution(func, y, shape=shape, alpha=alpha)
-  return p >= alpha
+  x = func(*shape).cpu().numpy().flatten()
+  y = np.random.randn(*shape).flatten()
+  return kstest(x, y) >= alpha
+
+def equal_distribution(tiny_func, torch_func, numpy_func, shape=(20, 23), alpha=0.05):
+  Tensor.manual_seed(1337)
+  torch.manual_seed(1337)
+  np.random.seed(1337)
+  x = tiny_func(*shape).cpu().numpy().flatten()
+  y = numpy_func(shape).flatten()
+  z = torch_func(shape).numpy().flatten()
+  return kstest(x, y) >= alpha and kstest(x, z) >= alpha
 
 class TestRandomness(unittest.TestCase):
   def test_rand(self):
     self.assertFalse(normal_test(Tensor.rand))
-    self.assertTrue(equal_distribution(Tensor.rand, lambda x: np.random.rand(*x)))
+    self.assertTrue(equal_distribution(Tensor.rand, torch.rand, lambda x: np.random.rand(*x)))
 
   def test_randn(self):
     self.assertTrue(normal_test(Tensor.randn))
-    self.assertFalse(equal_distribution(Tensor.randn, lambda x: np.random.rand(*x)))
+    self.assertTrue(equal_distribution(Tensor.randn, torch.randn, lambda x: np.random.randn(*x)))
 
   def test_uniform(self):
     self.assertFalse(normal_test(Tensor.uniform))
-    self.assertTrue(equal_distribution(Tensor.uniform, lambda x: np.random.rand(*x) * 2 - 1))
+    self.assertTrue(equal_distribution(Tensor.uniform, lambda x: torch.nn.init.uniform_(torch.empty(x), a=-1, b=1), lambda x: np.random.rand(*x) * 2 - 1))
 
   def test_scaled_uniform(self):
     self.assertFalse(normal_test(Tensor.scaled_uniform))
-    self.assertTrue(equal_distribution(Tensor.scaled_uniform, lambda x: (np.random.rand(*x) * 2 - 1) / math.sqrt(math.prod(x))))
+    self.assertTrue(equal_distribution(Tensor.scaled_uniform, lambda x: torch.nn.init.uniform_(torch.empty(x), a=-1, b=1) / math.sqrt(math.prod(x)), lambda x: (np.random.rand(*x) * 2 - 1) / math.sqrt(math.prod(x))))
 
   def test_glorot_uniform(self):
     self.assertFalse(normal_test(Tensor.glorot_uniform))
-    self.assertTrue(equal_distribution(Tensor.glorot_uniform, lambda x: (np.random.rand(*x) * 2 - 1) * math.sqrt(6 / (x[0] + math.prod(x[1:])))))
+    self.assertTrue(equal_distribution(Tensor.glorot_uniform, lambda x: torch.nn.init.xavier_uniform_(torch.empty(x)), lambda x: (np.random.rand(*x) * 2 - 1) * math.sqrt(6 / (x[0] + math.prod(x[1:])))))
 
 if __name__ == "__main__":
   unittest.main()

test/test_speed_v_torch.py

@@ -10,6 +10,7 @@ import time
 import numpy as np
 np.set_printoptions(linewidth=160)
 from functools import partial
+from tinygrad.lazy import Device
 from tinygrad.ops import GlobalCounters
 from tinygrad.tensor import Tensor
 from tinygrad.nn import Conv2d
@@ -19,6 +20,14 @@ from tinygrad.jit import TinyJit
 IN_CHANS = [int(x) for x in getenv("IN_CHANS", "4,16,64").split(",")]
 
 torch_device = torch.device('mps' if getenv("MPS", 0) else ('cuda' if getenv("TORCHCUDA", 0) else 'cpu'))
+if str(torch_device) == "mps":
+  import torch.mps
+  sync = lambda: torch.mps.synchronize()
+elif str(torch_device) == "cuda":
+  import torch.cuda
+  sync = lambda: torch.cuda.synchronize()
+else:
+  sync = lambda: None
 
 def colorize_float(x):
   ret = f"{x:7.2f}x"
@@ -35,24 +44,31 @@ def helper_test_speed(f1, *args):
   global save_ops, save_mem
   ets = []
   ret = None
-  for _ in range(CNT):
+  cache_defeat = np.zeros((2048,2048))
+  for i in range(CNT):
     del ret
-    args = [(x+1).realize() if isinstance(x, Tensor) else (None if x is None else (x+1)) for x in args]  # cache defeats
+
+    # operation cache defeats
+    args = [(x+1).realize() if isinstance(x, Tensor) else (None if x is None else (x+1)) for x in args]
 
     # force syncing
     [x.numpy() if isinstance(x, Tensor) or str(torch_device) == "cpu" else x.cpu().numpy() for x in args if x is not None]
 
+    # clear 32MB global memory cache (CPU and global memory only)
+    cache_defeat += 1
+
+    # manual pre sync
+    if isinstance(args[0], Tensor): Device[args[0].device].synchronize()
+    else: sync()
+
     GlobalCounters.global_ops = 0
     GlobalCounters.global_mem = 0
-    if DEBUG >= 4: print("benchmark start")
-    st = time.monotonic()
+    st = time.perf_counter()
     ret = f1(*args)
-    # not ideal, it's copying (sometimes). why is this so slow in tinygrad?
-    if isinstance(ret, Tensor) or str(torch_device) == "cpu": ret.numpy()
-    else: ret.cpu().numpy()
-    et = (time.monotonic() - st) * 1000
-    ets.append(et)
-    if DEBUG >= 4: print("benchmark stop")
+    if isinstance(ret, Tensor): Device[ret.device].synchronize()
+    else: sync()
+    et = (time.perf_counter() - st) * 1000
+    if i >= 1: ets.append(et)   # not the first run / one used for OPTLOCAL
     if GlobalCounters.global_ops:
       save_ops, save_mem = GlobalCounters.global_ops, GlobalCounters.global_mem
   return ret.cpu().numpy(), np.min(ets)
@@ -76,7 +92,7 @@ def helper_test_generic(name, f1, f1_args, f2, f2_args):
 
   desc = "faster" if et_torch > et_tinygrad else "slower"
   flops = save_ops*1e-6
-  mem = save_mem*4*1e-6
+  mem = save_mem*1e-6
   print(f"{prefix}{name:40s} {et_torch:7.2f} ms ({flops/et_torch:8.2f} GFLOPS {mem/et_torch:8.2f} GB/s) in torch, {et_tinygrad:7.2f} ms ({flops/et_tinygrad:8.2f} GFLOPS {mem/et_tinygrad:8.2f} GB/s) in tinygrad, {colorize_float(et_tinygrad/et_torch)} {desc} {flops:7.2f} MOPS {mem:7.2f} MB")
   prefix = " "
   np.testing.assert_allclose(val_tinygrad, val_torch, atol=1e-4, rtol=1e-3)
@@ -86,7 +102,7 @@ class TestSpeed(unittest.TestCase):
     global prefix
     prefix = " " if prefix is None else ""
     return super().setUp()
-  
+
   def test_sub(self):
     def f(a, b): return a-b
     helper_test_generic_square('sub', 4096, f, f)
@@ -116,7 +132,7 @@ class TestSpeed(unittest.TestCase):
       # to fit easily in L1, rotations should be 128x128 chunks. 128x128 is also the AMX size
       def f(a, b): return a.permute(1,0).contiguous()
       helper_test_generic_square('permute', N, f, f, onearg=True)
-    
+
   def test_double_permute(self):
     N = 64
     torch.manual_seed(0)
@@ -160,14 +176,18 @@ class TestSpeed(unittest.TestCase):
 
   def test_gemm(self):
     def f(a, b): return a @ b
-    helper_test_generic_square('gemm', 512, f, f)
+    helper_test_generic_square('gemm', 1024, f, f)
+
+  def test_gemm_small(self):
+    def f(a, b): return a @ b
+    helper_test_generic_square('gemm', 256, f, f)
 
   def test_gemm_unrolled(self):
     N = 512
     def f1(a, b): return a@b.T
     def f2(a, b): return (a.reshape(N, 1, N).expand(N, N, N) * b.reshape(1, N, N).expand(N, N, N)).sum(axis=2)
     helper_test_generic_square('gemm_unrolled', N, f1, f2)
-  
+
   def test_gemm_unrolled_permute_l(self):
     N = 512
     def f1(a, b): return a.T@b.T

test/test_tensor.py

@@ -149,9 +149,9 @@ class TestTinygrad(unittest.TestCase):
     for random_fn in [Tensor.randn, Tensor.uniform, Tensor.scaled_uniform, Tensor.glorot_uniform]:
       with self.subTest(msg=f"Tensor.{random_fn.__name__}"):
         Tensor.manual_seed(1337)
-        a = random_fn(10,10)
+        a = random_fn(10,10).realize()
         Tensor.manual_seed(1337)
-        b = random_fn(10,10)
+        b = random_fn(10,10).realize()
         np.testing.assert_allclose(a.numpy(), b.numpy())
 
 if __name__ == '__main__':

test/unit/test_example.py

@@ -0,0 +1,73 @@
+import unittest
+import numpy as np
+from tinygrad.lazy import Device
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import getenv
+
+def multidevice_test(fxn):
+  exclude_devices = getenv("EXCLUDE_DEVICES", "").split(",")
+  def ret(self):
+    for device in Device._buffers:
+      print(device)
+      if device in exclude_devices:
+        print(f"WARNING: {device} test is excluded")
+        continue
+      with self.subTest(device=device):
+        try:
+          Device[device]
+        except Exception:
+          print(f"WARNING: {device} test isn't running")
+          continue
+        fxn(self, device)
+  return ret
+
+class TestExample(unittest.TestCase):
+  @multidevice_test
+  def test_convert_to_cpu(self, device):
+    a = Tensor([[1,2],[3,4]], device=device)
+    assert a.numpy().shape == (2,2)
+    b = a.cpu()
+    assert b.numpy().shape == (2,2)
+
+  @multidevice_test
+  def test_2_plus_3(self, device):
+    a = Tensor([2], device=device)
+    b = Tensor([3], device=device)
+    result = a + b
+    print(f"{a.numpy()} + {b.numpy()} = {result.numpy()}")
+    assert result.numpy()[0] == 5.
+
+  @multidevice_test
+  def test_example_readme(self, device):
+    x = Tensor.eye(3, device=device, requires_grad=True)
+    y = Tensor([[2.0,0,-2.0]], device=device, requires_grad=True)
+    z = y.matmul(x).sum()
+    z.backward()
+
+    x.grad.numpy()  # dz/dx
+    y.grad.numpy()  # dz/dy
+
+    assert x.grad.device == device
+    assert y.grad.device == device
+
+  @multidevice_test
+  def test_example_matmul(self, device):
+    try:
+      Device[device]
+    except Exception:
+      print(f"WARNING: {device} test isn't running")
+      return
+
+    x = Tensor.eye(64, device=device, requires_grad=True)
+    y = Tensor.eye(64, device=device, requires_grad=True)
+    z = y.matmul(x).sum()
+    z.backward()
+
+    x.grad.numpy()  # dz/dx
+    y.grad.numpy()  # dz/dy
+
+    assert x.grad.device == device
+    assert y.grad.device == device
+
+if __name__ == '__main__':
+  unittest.main()

test/unit/test_flopcounter.py

@@ -1,39 +1,49 @@
 #!/usr/bin/env python
 import unittest
-from tinygrad.ops import InterpretedBuffer, GenericShape, LazyOp, BinaryOps, get_lazyop_info
+from typing import NamedTuple, Tuple
+from tinygrad.ops import LazyOp, BinaryOps, ReduceOps, get_lazyop_info
+from tinygrad.helpers import DType, dtypes
+
+class TestBuffer(NamedTuple):
+  __test__ = False # To prevent pytest from collecting this as a test
+  shape: Tuple[int, ...]
+  dtype: DType
 
 class TestFlopCounter(unittest.TestCase):
+  def setUp(self):
+    self.buf0 = TestBuffer(shape=(4,), dtype=dtypes.float32)
+    self.buf1 = TestBuffer(shape=(4,), dtype=dtypes.float32)
+
   def test_flops_add(self):
-    buf0 = InterpretedBuffer(GenericShape((4,)))
-    buf1 = InterpretedBuffer(GenericShape((4,)))
-    op0 = LazyOp(BinaryOps.ADD, (buf0,buf1,), None)
+    op0 = LazyOp(BinaryOps.ADD, (self.buf0,self.buf1,), None)
     info = get_lazyop_info(op0)
     self.assertEqual(info.flops, 4)
 
   def test_flops_add_twice(self):
-    buf0 = InterpretedBuffer(GenericShape((4,)))
-    buf1 = InterpretedBuffer(GenericShape((4,)))
-    op0 = LazyOp(BinaryOps.ADD, (buf0,buf1,), None)
-    op1 = LazyOp(BinaryOps.ADD, (op0,buf1,), None)
+    op0 = LazyOp(BinaryOps.ADD, (self.buf0,self.buf1,), None)
+    op1 = LazyOp(BinaryOps.ADD, (op0,self.buf1,), None)
     info = get_lazyop_info(op1)
     self.assertEqual(info.flops, 8)
 
   def test_flops_add_self(self):
-    buf0 = InterpretedBuffer(GenericShape((4,)))
-    buf1 = InterpretedBuffer(GenericShape((4,)))
-    op0 = LazyOp(BinaryOps.ADD, (buf0,buf1,), None)
+    op0 = LazyOp(BinaryOps.ADD, (self.buf0,self.buf1,), None)
     op1 = LazyOp(BinaryOps.ADD, (op0,op0,), None)
     info = get_lazyop_info(op1)
     self.assertEqual(info.flops, 8)
 
   def test_flops_add_roundabout_self(self):
-    buf0 = InterpretedBuffer(GenericShape((4,)))
-    buf1 = InterpretedBuffer(GenericShape((4,)))
-    op0 = LazyOp(BinaryOps.ADD, (buf0,buf1,), None)
-    op1 = LazyOp(BinaryOps.ADD, (op0,buf1,), None)
+    op0 = LazyOp(BinaryOps.ADD, (self.buf0,self.buf1,), None)
+    op1 = LazyOp(BinaryOps.ADD, (op0,self.buf1,), None)
     op2 = LazyOp(BinaryOps.ADD, (op0,op1,), None)
     info = get_lazyop_info(op2)
     self.assertEqual(info.flops, 12)
-  
+
+  def test_flops_red(self):
+    op0 = LazyOp(BinaryOps.MUL, (self.buf0,self.buf1,), None)
+    op1 = LazyOp(ReduceOps.SUM, (op0,), (1,))
+    op2 = LazyOp(BinaryOps.ADD, (op1, op1,), None)
+    info = get_lazyop_info(op2)
+    self.assertEqual(info.flops, 9)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_graph.py

@@ -1,11 +1,12 @@
 #!/usr/bin/env python
 import unittest
 import networkx as nx  # type: ignore
-import numpy as np
+from tinygrad.tensor import Tensor
 from tinygrad.graph import G, log_op, prune_graph
-from tinygrad.runtime.ops_cpu import CPUBuffer
 from tinygrad.ops import BinaryOps, LazyOp, MovementOps, ReduceOps
 
+def buf(*shp): return Tensor.ones(*shp, device="CPU").lazydata
+
 class TestGraph(unittest.TestCase):
   def setUp(self):
     G.clear()
@@ -14,10 +15,10 @@ class TestGraph(unittest.TestCase):
     assert nx.is_isomorphic(G, RG, node_match=lambda x,y: x["label"] == y["label"], edge_match=lambda x,y: x["label"] == y["label"] if "label" in y else True)
 
   def test_add_graph(self):
-    a = CPUBuffer.fromCPU(np.ones((4,4)))
-    b = CPUBuffer.fromCPU(np.ones((4,4)))
+    a = buf(4,4)
+    b = buf(4,4)
     ast = LazyOp(BinaryOps.ADD, (a,b))
-    ret = CPUBuffer(np.ones((4,4)))
+    ret = buf(4,4)
 
     RG = nx.DiGraph()
     RG.add_node(0, label="(4, 4)")
@@ -30,32 +31,32 @@ class TestGraph(unittest.TestCase):
     self.helper_compare_graph(RG)
 
   def test_add_sum_graph(self):
-    a = CPUBuffer.fromCPU(np.ones((4,4)))
-    b = CPUBuffer.fromCPU(np.ones((1,1)))
+    a = buf(4,4)
+    b = buf(1,1)
     op0 = LazyOp(MovementOps.RESHAPE, (b,), (4, 4))
     op1 = LazyOp(BinaryOps.ADD, (a,op0))
     ast = LazyOp(ReduceOps.SUM, (op1,), (1,1))
-    ret = CPUBuffer(np.ones((1,1)))
+    ret = buf(1,1)
 
     RG = nx.DiGraph()
     RG.add_node(0, label="(4, 4)")
     RG.add_node(1, label="(1, 1)")
     RG.add_node(2, label="{(4, 4), (1, 1)}\n(1, 1)")
-    RG.add_edge(0, 2, label="RE.AD.SU")
-    RG.add_edge(1, 2, label="RE.AD.SU")
+    RG.add_edge(0, 2, label="RES.ADD.SUM")
+    RG.add_edge(1, 2, label="RES.ADD.SUM")
 
     log_op(ret, ast, show_graph=True)
     self.helper_compare_graph(RG)
 
   def test_add_graph_prune(self):
-    a = CPUBuffer.fromCPU(np.ones((1,1)))
+    a = buf(1,1)
     ast = LazyOp(MovementOps.RESHAPE, (a,), (4, 4))
-    ret = CPUBuffer(np.ones((4,4)))
+    ret = buf(4,4)
     log_op(ret, ast, show_graph=True)
 
-    b = CPUBuffer.fromCPU(np.ones((4,4)))
+    b = buf(4,4)
     ast = LazyOp(BinaryOps.ADD, (ret,b))
-    ret = CPUBuffer(np.ones((4,4)))
+    ret = buf(4,4)
     log_op(ret, ast, show_graph=True)
     prune_graph()
 

test/unit/test_shapetracker.py

@@ -2,8 +2,8 @@
 import unittest
 import numpy as np
 from tinygrad.helpers import prod, all_same
-from tinygrad.shape import ShapeTracker, View, ZeroView, merge_views, get_contraction
-from tinygrad.codegen.gpu import to_image_idx
+from tinygrad.shape.shapetracker import ShapeTracker, View, merge_views, get_contraction
+from tinygrad.codegen.cstyle import to_image_idx
 
 def shapetracker_getitem(st, val):
   locals = {"idx": val, "valid": 1}
@@ -14,7 +14,7 @@ def shapetracker_getitem(st, val):
 class CheckingShapeTracker:
   def __init__(self, shape):
     self.st = ShapeTracker(shape)
-    self.t = np.arange(prod(shape), dtype=np.uint8).reshape(shape)
+    self.t = np.arange(prod(shape), dtype=np.int32).reshape(shape)
 
   @property
   def shape(self):
@@ -35,13 +35,17 @@ class CheckingShapeTracker:
     self.t = np.broadcast_to(self.t, new_shape)
 
   def flip(self, axis):
-    self.st.flip(axis)
+    self.st.stride(tuple(-1 if i in axis else 1 for i in range(len(self.shape))))
     self.t = np.flip(self.t, axis)
 
   def shrink(self, arg):
     self.st.shrink(arg)
     self.t = self.t[tuple([slice(x[0], x[1]) for x in arg])]
 
+  def pad(self, arg):
+    self.st.pad(arg)
+    self.t = np.pad(self.t, arg, constant_values=-1)
+
   def stride(self, arg):
     self.st.stride(arg)
     self.t = self.t[tuple([slice(None, None, x) for x in arg])]
@@ -59,40 +63,9 @@ class CheckingShapeTracker:
     x = [shapetracker_getitem(self.st, i) for i in range(prod(self.st.shape))]
     y = [self[i] for i in range(prod(self.shape))]
     idx, valid = self.st.expr_node()
-    print(x, y, self.st.shape, self.shape, idx.render(), valid.render())
+    print(x, y, self.st.shape, self.shape, idx.render(), valid.render(), self.st)
     assert self.st.shape == self.shape
-    assert x == y
-
-class TestImageShapeTracker(unittest.TestCase):
-  def test_image(self):
-    base_shape = (64, 1024, 4)
-    print(base_shape)
-
-    new_view = merge_views(
-      View((1, 66, 130, 32, 1, 1), (0, 4096, 32, 1, 0, 0), -4128),
-      View((64, 32, 8, 3, 3), (4160, 128, 4, 4160, 32), 0)
-    )
-    print(new_view)
-
-    st = ShapeTracker(shape=(64, 32, 8, 3, 3), views=[
-      View((1, 66, 130, 32, 1, 1), (0, 4096, 32, 1, 0, 0), -4128),
-      ZeroView((1, 64, 128, 32, 1, 1), ((0, 1), (-1, 65), (-1, 129), (0, 32), (0, 1), (0, 1))),
-      View((64, 32, 8, 3, 3), (4160, 128, 4, 4160, 32), 0)])
-    offsets = [0,32,64,96]
-
-    print(st.shape)
-    idys = []
-    for o in offsets:
-      print("offset:", o)
-      idxy, valid = st.expr_idxs(o)
-      print("idxy:", idxy.render())
-      print("valids:", [x.render() for x in valid.nodes])
-      idx, idy = to_image_idx(base_shape, idxy, valid, True)
-      idys.append(idy)
-      print(base_shape, idx.min, idx.max, idy.min, idy.max, idx, idy)
-
-    # y index shouldn't be changing
-    assert all_same(idys)
+    assert x == y, f"mismatch shapetracker:{x} real:{y}"
 
 class TestSimplifyingShapeTracker(unittest.TestCase):
   def setUp(self):
@@ -142,17 +115,9 @@ class TestSimplifyingShapeTracker(unittest.TestCase):
 # Tensor.zeros(2, 4).permute(1,0).reshape(2, 4)
 # (d1*4 + d0%4), d1=x//4, d0=x%4 = ((x//4)*4) + (x%4)%4
 
-class TestZeroViewShapeTracker(unittest.TestCase):
-  def test_pad(self):
-    self.st = ShapeTracker((4, 4))
-    self.st.pad(((1, 1), (1, 1)))
-    assert self.st.shape == (6,6)
-    compareZv = ZeroView((4,4), ((-1,5), (-1,5)))
-    assert str(self.st.views[1]) == str(compareZv)
-
 class TestComplexShapeTracker(unittest.TestCase):
   def test_add_1s(self):
-    self.st = ShapeTracker((4, 4))
+    self.st = CheckingShapeTracker((4, 4))
     self.st.permute((1,0))
     self.st.reshape((1,4,1,4,1))
     assert not self.st.contiguous
@@ -160,20 +125,20 @@ class TestComplexShapeTracker(unittest.TestCase):
     assert self.st.contiguous
 
   def test_permute_1s_simple(self):
-    self.st = ShapeTracker((1, 16, 9,9))
+    self.st = CheckingShapeTracker((1, 16, 9,9))
     self.st.permute((1,0,2,3))
     assert self.st.contiguous
-    self.st = ShapeTracker((2, 16, 9,9))
+    self.st = CheckingShapeTracker((2, 16, 9,9))
     self.st.permute((1,0,2,3))
     assert not self.st.contiguous
 
   def test_remove_1s_simple(self):
-    self.st = ShapeTracker((1, 16, 1, 1))
+    self.st = CheckingShapeTracker((1, 16, 1, 1))
     self.st.reshape((16,))
     assert self.st.contiguous
 
   def test_remove_1s(self):
-    self.st = ShapeTracker((1, 4, 1, 4, 1))
+    self.st = CheckingShapeTracker((1, 4, 1, 4, 1))
     self.st.permute((0,3,2,1,4))
     self.st.reshape((4,4))
     assert not self.st.contiguous
@@ -181,46 +146,46 @@ class TestComplexShapeTracker(unittest.TestCase):
     assert self.st.contiguous
 
   def test_permute_reshape(self):
-    self.st = ShapeTracker((4, 4))
+    self.st = CheckingShapeTracker((4, 4))
     self.st.permute((1,0))
     self.st.reshape((2, 2, 2, 2))
     # TODO: should also be tested by test_super_complex
     assert len(self.st.views) == 1
 
   def test_factorize_split(self):
-    self.st = ShapeTracker((4, 4))
+    self.st = CheckingShapeTracker((4, 4))
     self.st.permute((1,0))
     self.st.reshape((2, 2, 2, 2))
     self.st.permute((2,3,0,1))
     assert self.st.contiguous
 
   def test_factorize_combine(self):
-    self.st = ShapeTracker((4, 4, 4))
+    self.st = CheckingShapeTracker((4, 4, 4))
     self.st.permute((2, 0, 1))
     self.st.reshape((4, 16))
     self.st.permute((1, 0))
     assert self.st.contiguous
 
   def test_factorize_combine_add_ones(self):
-    self.st = ShapeTracker((4, 4, 4))
+    self.st = CheckingShapeTracker((4, 4, 4))
     self.st.permute((2, 0, 1))
     self.st.reshape((4, 16, 1, 1))
     self.st.permute((1, 0, 2, 3))
     assert self.st.contiguous
 
   def test_fancy_factorize(self):
-    self.st = ShapeTracker((32, 3, 3, 1))
+    self.st = CheckingShapeTracker((32, 3, 3, 1))
     self.st.reshape((8, 4, 3, 3))
     assert len(self.st.views) == 1
 
   def test_super_complex_2_fail(self):
-    self.st = ShapeTracker((4, 4, 4))
+    self.st = CheckingShapeTracker((4, 4, 4))
     self.st.permute((2, 0, 1))
     self.st.reshape((16, 4))
     assert len(self.st.views) != 1
 
   def test_work(self):
-    self.st = ShapeTracker((64, 1024, 4))
+    self.st = CheckingShapeTracker((64, 1024, 4))
     self.st.reshape((1, 64, 128, 32))
     self.st.permute((0, 3, 1, 2))
     self.st.reshape((1, 32, 1, 64, 128))
@@ -228,7 +193,7 @@ class TestComplexShapeTracker(unittest.TestCase):
     assert self.st.contiguous
 
   def test_work2(self):
-    self.st = ShapeTracker((64, 1024, 4))
+    self.st = CheckingShapeTracker((64, 1024, 4))
     self.st.reshape((1, 64, 128, 32))
     self.st.permute((0, 3, 1, 2))
     self.st.reshape((1, 1, 32, 64, 128))
@@ -276,6 +241,47 @@ class TestSingleShapeTracker(unittest.TestCase):
     self.st.permute((1,0))
     assert not self.st.contiguous
 
+class TestShapeTrackerFuzzFailures(unittest.TestCase):
+  def setUp(self):
+    self.st = CheckingShapeTracker((3,3,3))
+  def tearDown(self):
+    self.st.assert_same()
+  @unittest.skip("simplify doesn't work in this case")
+  def test_case_1(self):
+    self.st.shrink(((1, 2), (1, 3), (1, 3)))
+    self.st.reshape((1, 4))
+    self.st.shrink(((0, 1), (1, 3)))
+    print(self.st.st)
+    self.st.simplify()
+    print(self.st.st)
+  def test_case_2(self):
+    self.st.stride( (1, 1, -2) )
+    self.st.reshape( (3, 6) )
+    self.st.shrink( ((1, 2), (1, 5)) )
+    self.st.stride( (1, -1) )
+  def test_case_3(self):
+    self.st.shrink( ((0, 2), (0, 2), (0, 1)) )
+    self.st.permute( (1, 0, 2) )
+    self.st.reshape( (4,) )
+    self.st.shrink( ((0, 3),) )
+    self.st.stride( (-1,) )
+  def test_case_4(self):
+    self.st.reshape( (3, 3, 3, 1) )
+    self.st.pad( ((0, 0), (0, 0), (0, 0), (1, 1)) )
+    self.st.shrink( ((0, 2), (1, 2), (0, 2), (0, 1)) )
+    self.st.expand( (2, 1, 2, 3) )
+
+class TestMaskedShapeTracker(unittest.TestCase):
+  def test_pad_1x1(self):
+    self.st = CheckingShapeTracker((1,1))
+    self.st.pad(((1,1), (1,1)))
+    self.st.assert_same()
+
+  def test_pad_2x2(self):
+    self.st = CheckingShapeTracker((2,2))
+    self.st.pad(((1,1), (1,1)))
+    self.st.assert_same()
+
 class TestShapeTracker(unittest.TestCase):
   def setUp(self):
     self.st = CheckingShapeTracker((7,4))
@@ -291,6 +297,54 @@ class TestShapeTracker(unittest.TestCase):
     self.test_permute()
     self.apply(lambda x: x.reshape((prod(self.st.shape), )))
 
+  def test_simple_pad(self):
+    self.st.pad(((1,1), (1,1)))
+
+  def test_pad_shrink(self):
+    self.st.pad(((1,1), (1,1)))
+    self.st.shrink(((0,4), (0,4)))
+
+  def test_pad_one_sided(self):
+    self.st.pad(((0,1), (0,0)))
+
+  def test_pad_reshape(self):
+    self.st.pad(((0,1), (0,0)))
+    self.st.reshape((8*4,))
+
+  def test_pad_pad(self):
+    self.st.pad(((1,1), (1,1)))
+    self.st.pad(((1,1), (1,1)))
+
+  def test_pad_permute(self):
+    self.st.pad(((1,1), (2,2)))
+    self.st.permute((1,0))
+
+  def test_pad_expand(self):
+    self.st.reshape((7,4,1))
+    self.st.pad(((1,1), (1,1), (0,0)))
+    self.st.expand((9,6,4))
+
+  def test_pad_expand_alt(self):
+    self.st.pad(((1,1), (1,1)))
+    self.st.reshape((9,6,1))
+    self.st.expand((9,6,4))
+
+  def test_pad_stride(self):
+    self.st.pad(((1,4), (1,3)))
+    self.st.stride((2,2))
+
+  def test_pad_stride_neg(self):
+    self.st.pad(((1,2), (1,0)))
+    self.st.stride((-1,-1))
+
+  def test_pad_stride_both(self):
+    self.st.pad(((1,2), (1,0)))
+    self.st.stride((-2,-2))
+
+  def test_shrink_pad(self):
+    self.st.shrink(((0,4), (0,4)))
+    self.st.pad(((1,1), (1,1)))
+
   def test_reshape(self):
     new_shape = self.st.shape[::-1]
     self.apply(lambda x: x.reshape(new_shape))
@@ -381,6 +435,9 @@ class TestGetContraction(unittest.TestCase):
     r = get_contraction((1,2,3,4), (2,3,4))
     self.assertEqual(r, [[0, 1], [2], [3]])
 
+    r = get_contraction((2,1,3,4), (2,3,4))
+    self.assertEqual(r, [[0], [1, 2], [3]])
+
     r = get_contraction((1,2,3,1,4), (1,2,3,4))
     self.assertEqual(r, [[0], [1], [2], [3, 4]])
 
@@ -391,10 +448,34 @@ class TestGetContraction(unittest.TestCase):
     self.assertEqual(r, [[0], [1], [2, 3]])
 
     r = get_contraction((1,2,3,4), (2,1,3,4))
-    self.assertEqual(r, None)
+    self.assertEqual(r, [[0, 1], [], [2], [3]])
+
+    r = get_contraction((1,2,3,4), (1,1,2*3*4,1))
+    self.assertEqual(r, [[0], [], [1,2,3], []])
+
+    r = get_contraction((2,1,3,4), (1,2,3,4))
+    self.assertEqual(r, [[], [0], [1, 2], [3]])
+
+    r = get_contraction((1,2,3,4), (2*3*4,1,1,1))
+    self.assertEqual(r, [[0, 1, 2, 3], [], [], []])
+
+    r = get_contraction((4,4,4,4), (16,1,16))
+    self.assertEqual(r, [[0, 1], [], [2, 3]])
+
+    r = get_contraction((1,2,3,4,1,1,1), (2,3,4))
+    self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]])
 
     r = get_contraction((1,2,3,4), (1,2,3,4,1))
-    self.assertEqual(r, None)
+    self.assertEqual(r, [[0], [1], [2], [3], []])
+
+    r = get_contraction((14,1,384,14,1,1,1,1), (1,14,384,14))
+    self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7]])
+
+    r = get_contraction((14,1,384,1,14,1,1,1,1), (1,14,384,14))
+    self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7,8]])
+
+    r = get_contraction((512, 512), (1, 1, 512, 1, 1, 1, 1, 512))
+    self.assertEqual(r, [[], [], [0], [], [], [], [], [1]])
 
     r = get_contraction((1,2,3,4), (1,2,6,2))
     self.assertEqual(r, None)

test/unit/test_symbolic.py

@@ -23,7 +23,20 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(Variable("a", 3, 8)<4, 0, 1, "(a<4)")
     self.helper_test_variable(Variable("a", 3, 8)<3, 0, 0, "0")
     self.helper_test_variable(Variable("a", 3, 8)<2, 0, 0, "0")
-  
+
+  def test_ge_divides(self):
+    expr = (Variable("idx", 0, 511)*4 + Variable("FLOAT4_INDEX", 0, 3)) < 512
+    self.helper_test_variable(expr, 0, 1, "(((idx*4)+FLOAT4_INDEX)<512)")
+    self.helper_test_variable(expr//4, 0, 1, "(idx<128)")
+
+  def test_ge_divides_and(self):
+    expr = Variable.ands([(Variable("idx1", 0, 511)*4 + Variable("FLOAT4_INDEX", 0, 3)) < 512,
+                          (Variable("idx2", 0, 511)*4 + Variable("FLOAT4_INDEX", 0, 3)) < 512])
+    self.helper_test_variable(expr//4, 0, 1, "((idx1<128) and (idx2<128))")
+    expr = Variable.ands([(Variable("idx1", 0, 511)*4 + Variable("FLOAT4_INDEX", 0, 3)) < 512,
+                          (Variable("idx2", 0, 511)*4 + Variable("FLOAT8_INDEX", 0, 7)) < 512])
+    self.helper_test_variable(expr//4, 0, 1, "((((FLOAT8_INDEX//4)+idx2)<128) and (idx1<128))")
+
   def test_div_becomes_num(self):
     assert isinstance(Variable("a", 2, 3)//2, NumNode)
 
@@ -72,6 +85,9 @@ class TestSymbolic(unittest.TestCase):
   def test_mul_1(self):
     self.helper_test_variable(Variable("a", 0, 8)*1, 0, 8, "a")
 
+  def test_mul_neg_1(self):
+    self.helper_test_variable((Variable("a", 0, 2)*-1)//3, -1, 0, "((((a*-1)+3)//3)+-1)")
+
   def test_mul_2(self):
     self.helper_test_variable(Variable("a", 0, 8)*2, 0, 16, "(a*2)")
 
@@ -101,7 +117,7 @@ class TestSymbolic(unittest.TestCase):
 
   def test_sum_div_no_factor(self):
     self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*5]) // 2, 0, 25, "(((a*5)+(b*5))//2)")
-  
+
   def test_mod_factor(self):
     # NOTE: even though the mod max is 50, it can't know this without knowing about the mul
     self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*100, Variable("b", 0, 3)*50]) % 100, 0, 99, "((b*50)%100)")
@@ -126,7 +142,7 @@ class TestSymbolic(unittest.TestCase):
 
   def test_mod_mul_sum(self):
     self.helper_test_variable(Variable.sum([Variable("b", 0, 2), Variable("a", 0, 5)*10])%9, 0, 7, "(a+b)")
-  
+
   def test_sum_0(self):
     self.helper_test_variable(Variable.sum([Variable("a", 0, 7)]), 0, 7, "a")
 
@@ -164,7 +180,7 @@ class TestSymbolic(unittest.TestCase):
 
   def test_div_factor(self):
     self.helper_test_variable(Variable.sum([Variable.num(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40]) // 40, -1, 9, "(-1+b)")
-  
+
   def test_mul_div(self):
     self.helper_test_variable((Variable("a", 0, 10)*4)//4, 0, 10, "a")
 
@@ -177,6 +193,12 @@ class TestSymbolic(unittest.TestCase):
   def test_div_remove(self):
     self.helper_test_variable(Variable.sum([Variable("idx0", 0, 127)*4, Variable("idx2", 0, 3)])//4, 0, 127, "idx0")
 
+  def test_div_numerator_negative(self):
+    self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -9, 0, "((((idx*-10)+99)//11)+-9)")
+
+  def test_div_into_mod(self):
+    self.helper_test_variable((Variable("idx", 0, 16)*4)%8//4, 0, 1, "(idx%2)")
+
 class TestSymbolicNumeric(unittest.TestCase):
   def helper_test_numeric(self, f):
     # TODO: why are the negative tests broken? (even if we did support negative variables)

test/unit/test_utils.py

@@ -1,22 +0,0 @@
-#!/usr/bin/env python
-import io
-import unittest
-from extra.utils import fetch
-from PIL import Image
-
-class TestUtils(unittest.TestCase):  
-  def test_fetch_bad_http(self):
-    self.assertRaises(AssertionError, fetch, 'http://httpstat.us/500')
-    self.assertRaises(AssertionError, fetch, 'http://httpstat.us/404')
-    self.assertRaises(AssertionError, fetch, 'http://httpstat.us/400')
-  
-  def test_fetch_small(self):
-    assert(len(fetch('https://google.com'))>0)
-
-  def test_fetch_img(self):
-    img = fetch("https://media.istockphoto.com/photos/hen-picture-id831791190")
-    pimg = Image.open(io.BytesIO(img))
-    assert pimg.size == (705, 1024)
-
-if __name__ == '__main__':
-  unittest.main()