factor out KOPT code
MODIFY
datasets/__init__.py
datasets/__init__.py
@@ -14,9 +14,13 @@ def fetch_mnist():
Y_test = parse(os.path.dirname(__file__)+"/mnist/t10k-labels-idx1-ubyte.gz")[8:]
return X_train, Y_train, X_test, Y_test
-def fetch_cifar():
+def fetch_cifar(train=True):
tt = tarfile.open(fileobj=io.BytesIO(fetch('https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz')), mode='r:gz')
+ if train:
+ # TODO: data_batch 2-5
db = pickle.load(tt.extractfile('cifar-10-batches-py/data_batch_1'), encoding="bytes")
+ else:
+ db = pickle.load(tt.extractfile('cifar-10-batches-py/test_batch'), encoding="bytes")
X = db[b'data'].reshape((-1, 3, 32, 32))
Y = np.array(db[b'labels'])
return X, Y

MODIFY
examples/hlb_cifar10.py
examples/hlb_cifar10.py
#!/usr/bin/env python3
# tinygrad implementation of https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
# https://myrtle.ai/learn/how-to-train-your-resnet-8-bag-of-tricks/
+# https://siboehm.com/articles/22/CUDA-MMM
+# TODO: gelu is causing nans!
+import numpy as np
+import time
from datasets import fetch_cifar
from tinygrad import nn
from tinygrad.nn import optim
-from extra.training import train
+from tinygrad.tensor import Tensor
+from extra.training import train, evaluate
from extra.utils import get_parameters
+from tinygrad.ops import GlobalCounters
+
+num_classes = 10
class ConvGroup:
- def __init__(self, channels_in, channels_out, short):
- self.short = short
+ def __init__(self, channels_in, channels_out, short, se=False):
+ self.short, self.se = short, se and not short
self.conv = [nn.Conv2d(channels_in if i == 0 else channels_out, channels_out, kernel_size=3, padding=1, bias=False) for i in range(1 if short else 3)]
self.norm = [nn.BatchNorm2D(channels_out) for _ in range(1 if short else 3)]
-
- if not self.short:
- self.se1 = nn.Linear(channels_out, channels_out//16)
- self.se2 = nn.Linear(channels_out//16, channels_out)
+ if self.se: self.se1, self.se2 = nn.Linear(channels_out, channels_out//16), nn.Linear(channels_out//16, channels_out)
def __call__(self, x):
x = self.conv[0](x).max_pool2d(2)
- x = self.norm[0](x).gelu()
+ x = self.norm[0](x).relu()
if self.short: return x
residual = x
- mult = self.se2(self.se1(residual.mean((2,3))).gelu()).sigmoid().reshape(x.shape[0], x.shape[1], 1, 1)
- x = self.norm[1](self.conv[1](x)).gelu()
- x = self.norm[2](self.conv[2](x) * mult).gelu()
+ mult = self.se2(self.se1(residual.mean((2,3))).relu()).sigmoid().reshape(x.shape[0], x.shape[1], 1, 1) if self.se else 1.0
+ x = self.norm[1](self.conv[1](x)).relu()
+ x = self.norm[2](self.conv[2](x) * mult).relu()
return x + residual
class SpeedyResNet:
@@ -33,21 +38,81 @@ class SpeedyResNet:
self.net = [
nn.Conv2d(3, 64, kernel_size=1),
nn.BatchNorm2D(64),
- lambda x: x.gelu(),
+ lambda x: x.relu(),
ConvGroup(64, 128, short=False),
ConvGroup(128, 256, short=True),
ConvGroup(256, 512, short=False),
lambda x: x.max((2,3)),
- nn.Linear(512, 1000, bias=False)
+ nn.Linear(512, num_classes, bias=False)
]
- def __call__(self, x): return x.sequential(self.net)
+ # note, pytorch just uses https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html instead of logsoftmax
+ def __call__(self, x): return x.sequential(self.net).logsoftmax()
+
+# TODO: this will become @tinygrad.jit
+first, cl_cache, loss = True, None, None
+from tinygrad.llops.ops_gpu import CL
+def train_step_jitted(model, optimizer, X, Y, enable_jit=False):
+ global cl_cache, first, loss
+
+ if not cl_cache:
+ if not first:
+ CL.CACHE = []
+ if enable_jit: first = False
+ out = model(X)
+ loss = out.mul(Y).mean()
+ optimizer.zero_grad()
+ loss.backward()
+ optimizer.step()
+ if not first:
+ cl_cache = CL.CACHE
+ CL.CACHE = None
+
+ if cl_cache:
+ for prg, args in cl_cache:
+ prg.clprg(CL().cl_queue, *args)
+ GlobalCounters.global_ops += prg.op_estimate
+ return loss
+
+def fetch_batch(X_train, Y_train, BS):
+ # fetch a batch
+ samp = np.random.randint(0, X_train.shape[0], size=(BS))
+ X = Tensor(X_train[samp])
+ Y = np.zeros((BS, num_classes), np.float32)
+ Y[range(BS),Y_train[samp]] = -1.0*num_classes
+ Y = Tensor(Y.reshape(BS, num_classes))
+ return X.realize(), Y.realize()
def train_cifar():
- X,Y = fetch_cifar()
+ Tensor.training = True
+ X_train,Y_train = fetch_cifar(train=True)
+ #X_test,Y_test = fetch_cifar(train=False)
model = SpeedyResNet()
- optimizer = optim.SGD(get_parameters(model))
- train(model, X, Y, optimizer, steps=X.shape[0]//512, BS=512)
+ optimizer = optim.SGD(get_parameters(model), lr=0.001)
+
+ # 97 steps in 2 seconds = 20ms / step
+ # step is 1163.42 GOPS = 56 TFLOPS!!!, 41% of max 136
+ # 4 seconds for tfloat32 ~ 28 TFLOPS, 41% of max 68
+ # 6.4 seconds for float32 ~ 17 TFLOPS, 50% of max 34.1
+ # 4.7 seconds for float32 w/o channels last. 24 TFLOPS. we get 50ms then i'll be happy. only 64x off
+
+ # https://www.anandtech.com/show/16727/nvidia-announces-geforce-rtx-3080-ti-3070-ti-upgraded-cards-coming-in-june
+ # 136 TFLOPS is the theoretical max w float16 on 3080TI
+
+ for i in range(10):
+ X, Y = fetch_batch(X_train, Y_train, BS=512)
+ CL.time_sum, CL.kernel_count = 0, -1
+ CL.ops_sum = 0 # TODO: this should be GlobalCounters.global_ops
+ GlobalCounters.global_ops = 0
+ st = time.monotonic()
+ loss = train_step_jitted(model, optimizer, X, Y)
+ et = time.monotonic()
+ loss_cpu = loss.detach().cpu().data[0]
+ cl = time.monotonic()
+ print(f"{(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, {CL.mem_used/1e9:.2f} GB used, {GlobalCounters.global_ops*1e-9/(cl-st):9.2f} GFLOPS")
+
+ #train(model, X, Y, optimizer, steps=X.shape[0]//BS, BS=BS)
+ #evaluate(model, X_test, Y_test)
if __name__ == "__main__":
train_cifar()

MODIFY
extra/kernel_search.py
extra/kernel_search.py
@@ -8,8 +8,12 @@ from tinygrad.ops import LazyOp, ReduceOps, BinaryOps, UnaryOps, MovementOps
from tinygrad.shape import ShapeTracker, View, ZeroView
from tinygrad.llops.ops_gpu import GPUBuffer, CLASTKernel, CL
from tinygrad.runtime.opencl import OSX_TIMING_RATIO
+from tinygrad.ops import DEBUG
from extra.lib_test_ast import test_ast
+import pickle, dbm
+intervention_cache = dbm.open('/tmp/kopt.db', 'c')
+
Interventions = Enum("Interventions", ["SWAP", "UPCAST", "SHIFT", "REDUCE"])
def get_interventions(k):
p1, p2, p3, p4, p5 = [], [], [], [], []
@@ -96,9 +100,22 @@ def search_one(ast, winning_interventions=[]):
baseline = options[0]
options = sorted(options, key=lambda x: x[0]*x[2])
best = options[0]
- print(f"{name:30s} {baseline[0]/1e3:8.2f} ms -> {best[0]/1e3:8.2f} ms {baseline[0]/best[0]:7.2f}x *with* {winning_interventions} + {best[1]}")
+ print(f"{name:30s} {baseline[0]/1e3:8.2f} us -> {best[0]/1e3:8.2f} us {baseline[0]/best[0]:7.2f}x *with* {winning_interventions} + {best[1]}")
return best
+def apply_optimization(k, ast, max_interventions=1, cache=True):
+ from extra.kernel_search import search_one, apply_intervention
+ if k.key not in intervention_cache or cache == False:
+ winning_interventions = []
+ for i in range(max_interventions): # NOTE: multiple interventions is breaking the ASTs
+ oo = search_one(ast, winning_interventions)
+ if oo[1] is None: break
+ winning_interventions.append(oo[1])
+ intervention_cache[k.key] = pickle.dumps(winning_interventions)
+ ic = pickle.loads(intervention_cache[k.key])
+ if DEBUG >= 3: print("intervention", ic)
+ for w in ic: apply_intervention(k, *w)
+
def search(ast):
k = CLASTKernel(ast)
best_time = baseline = run_and_time(k)

MODIFY
extra/training.py
extra/training.py
@@ -50,7 +50,8 @@ def evaluate(model, X_test, Y_test, num_classes=None, BS=128, return_predict=Fal
Y_test_preds_out = np.zeros(list(Y_test.shape)+[num_classes])
for i in trange((len(Y_test)-1)//BS+1, disable=os.getenv('CI') is not None):
x = Tensor(transform(X_test[i*BS:(i+1)*BS]))
- Y_test_preds_out[i*BS:(i+1)*BS] = model.forward(x).cpu().data
+ out = model.forward(x) if hasattr(model, 'forward') else model(x)
+ Y_test_preds_out[i*BS:(i+1)*BS] = out.cpu().data
Y_test_preds = np.argmax(Y_test_preds_out, axis=-1)
Y_test = target_transform(Y_test)
return (Y_test == Y_test_preds).mean(), Y_test_preds

MODIFY
test/external_test_gpu_ast.py
test/external_test_gpu_ast.py
@@ -9,6 +9,9 @@ from extra.lib_test_ast import test_ast
def compile_and_test_ast(ast):
k = CLASTKernel(ast)
+ if int(os.getenv("KOPT", "0")):
+ from extra.kernel_search import apply_optimization
+ apply_optimization(k, ast, 10, int(os.getenv("KCACHE", "0")))
k.codegen()(*k.bufs)
if not int(os.getenv("NOTEST", "0")): test_ast(k)
@@ -20,6 +23,13 @@ class TestAST(unittest.TestCase):
ast = LazyOp(UnaryOps.RELU, (op1,), None)
compile_and_test_ast(ast)
+ def test_cifar_conv(self):
+ buf0 = GPUBuffer(shape=ShapeTracker(shape=(512, 1, 128, 32, 32, 64, 3, 3), views=[View((512, 64, 34, 34), (65536, 1024, 32, 1), -33), ZeroView((512, 64, 32, 32), ((0, 512), (0, 64), (-1, 33), (-1, 33))), View((512, 1, 128, 32, 32, 64, 3, 3), (73984, 73984, 0, 34, 1, 1156, 34, 1), 0)]), hostbuf=GPUBuffer(shape=(512, 64, 32, 32), force_create=True))
+ buf1 = GPUBuffer(shape=ShapeTracker(shape=(512, 1, 128, 32, 32, 64, 3, 3), views=[View((512, 1, 128, 32, 32, 64, 3, 3), (0, 0, 576, 0, 0, 9, 3, 1), 0)]), hostbuf=GPUBuffer(shape=(128, 64, 3, 3), force_create=True))
+ op0 = LazyOp(BinaryOps.MUL, (buf0,buf1,), None)
+ ast = LazyOp(ReduceOps.SUM, (op0,), (512, 1, 128, 32, 32, 1, 1, 1))
+ compile_and_test_ast(ast)
+
def test_first_op_conv(self):
buf0 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 8, 4, 3, 3, 3, 4), views=[View((1, 130, 258, 1, 12), (393216, 3072, 12, 12, 1), -3084), ZeroView((1, 128, 256, 1, 12), ((0, 1), (-1, 129), (-1, 257), (0, 1), (0, 12))), View((1, 64, 128, 8, 4, 3, 3, 3, 4), (0, 6192, 24, 0, 0, 3096, 12, 4, 1), 0)]), hostbuf=GPUBuffer(shape=(128, 768, 4), force_create=True))
buf1 = GPUBuffer(shape=ShapeTracker(shape=(1, 64, 128, 8, 4, 3, 3, 3, 4), views=[View((1, 64, 128, 8, 4, 3, 3, 3, 4), (0, 0, 0, 432, 4, 144, 16, 48, 1), 0)]), hostbuf=GPUBuffer(shape=(8, 108, 4), force_create=True))

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
from __future__ import annotations
import os
import numpy as np
-from typing import List, Tuple, Optional, Dict, Union, Set, Any
+from typing import List, Tuple, Optional, Dict, Union, Set
from tinygrad.helpers import prod
from tinygrad.ops import DEBUG, UnaryOps, BinaryOps, ReduceOps, MovementOps, LazyOp, Op, ExplicitExecAST, GlobalCounters
from tinygrad.ast import ASTKernel, Token, Types
@@ -20,10 +20,6 @@ KOPT = int(os.getenv("KOPT", "0"))
PRINT_AST = os.getenv("PRINT_AST", "0")
TEST_AST = int(os.getenv("TEST_AST", "0"))
-if KOPT:
- import pickle, dbm
- intervention_cache = dbm.open('/tmp/kopt.db', 'c')
-
def group_float4(x):
assert all(y.typ == Types.FLOAT for y in x) and len(x)%4 == 0
return [Token(f"(float4)({','.join([x[i+j].tok for j in range(4)])})", Types.FLOAT4) for i in range(0, len(x), 4)]
@@ -361,17 +357,8 @@ class GPUBuffer(ExplicitExecAST):
def exec_ast(cls, ast:LazyOp):
k = CLASTKernel(ast)
if KOPT:
- from extra.kernel_search import search_one, apply_intervention
- if k.key not in intervention_cache:
- winning_interventions : List[Any] = []
- for i in range(1): # NOTE: multiple interventions is breaking the ASTs
- oo = search_one(ast, winning_interventions)
- if oo[1] is None: break
- winning_interventions.append(oo[1])
- intervention_cache[k.key] = pickle.dumps(winning_interventions)
- ic = pickle.loads(intervention_cache[k.key])
- if DEBUG >= 3: print(ic)
- for w in ic: apply_intervention(k, *w)
+ from extra.kernel_search import apply_optimization
+ apply_optimization(k, ast)
k.codegen()(*k.bufs)
if PRINT_AST == "1" or (hasattr(k, "fxn") and PRINT_AST == k.fxn.name):
print(k.fxn.name)

MODIFY
tinygrad/runtime/opencl.py
tinygrad/runtime/opencl.py
@@ -75,6 +75,8 @@ class CLProgram:
def __call__(self, *args):
CL.kernel_count += 1
if DEBUG >= 4: print(args[0], args[1], self.prg)
+ # print the PTX for NVIDIA. TODO: probably broken for everything else
+ if DEBUG >= 5: print(self.clprogram.get_info(cl.program_info.BINARIES)[0].decode('utf-8'))
if CL.CACHE is not None: CL.CACHE.append((self, args))
else: e = self.clprg(CL().cl_queue, *args)
if DEBUG >= 2 and CL.CACHE is None:

factor out KOPT code