proper llops
RENAME
tinygrad/llops/cpu.py
tinygrad/llops/ops_cpu.py
RENAME
tinygrad/llops/opencl.py
tinygrad/llops/ops_gpu.py
-# llops don't know about derivatives or memory allocation
import functools
import numpy as np
import pyopencl as cl

RENAME
tinygrad/llops/torch.py
tinygrad/llops/ops_torch.py
-import os
import torch
import numpy as np
-from ..tensor import Function
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class Buffer(torch.Tensor):
@@ -21,7 +19,7 @@ class Buffer(torch.Tensor):
# ************* unary+binary+reduce+movement ops *************
-from tinygrad.llops.cpu import unary_op, binary_op, reduce_op, reshape, perm_axis, inner_slice, matmul
+from tinygrad.llops.ops_cpu import unary_op, binary_op, reduce_op, reshape, perm_axis, inner_slice, matmul
# ************* processing ops *************

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -11,7 +11,7 @@ def select_llops(ops):
# ************* unary ops *************
-class UnaryOp(Function):
+class _UnaryOp(Function):
def forward(ctx, input):
ctx.save_for_backward(input)
return ctx.op.unary_op(ctx.fop, input, ctx.op.Buffer(input.shape))
@@ -20,7 +20,7 @@ class UnaryOp(Function):
input, = ctx.saved_tensors
return ctx.op.binary_op(ctx.bop, input, grad_output, ctx.op.Buffer(input.shape))
-class ReLU(UnaryOp):
+class ReLU(_UnaryOp):
fop = UnaryOps.RELU
def backward(ctx, grad_output):
@@ -30,11 +30,11 @@ class ReLU(UnaryOp):
ctx.op.unary_op(UnaryOps.RELU, ret, ret)
return ctx.op.binary_op(BinaryOps.MUL, ret, grad_output, ret)
-class Log(UnaryOp):
+class Log(_UnaryOp):
fop = UnaryOps.LOG
bop = BinaryOps.DIV
-class Exp(UnaryOp):
+class Exp(_UnaryOp):
def forward(ctx, input):
ret = ctx.op.unary_op(UnaryOps.EXP, input, ctx.op.Buffer(input.shape))
ctx.save_for_backward(ret) # we save the output here, not the input

DELETE
tinygrad/ops/ops_cpu.py
-from ..mlops import *
-Buffer = select_llops("cpu")
\ No newline at end of file

DELETE
tinygrad/ops/ops_gpu.py
-from ..mlops import *
-Buffer = select_llops("opencl")

DELETE
tinygrad/ops/ops_torch.py
-from ..mlops import *
-Buffer = select_llops("torch")
\ No newline at end of file

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -65,14 +65,19 @@ class ProfileOp:
# **** enumerate supported devices ****
class Device:
- _ops = sorted(os.listdir(os.path.join(os.path.dirname(os.path.realpath(__file__)), "ops")))
+ _ops = sorted(os.listdir(os.path.join(os.path.dirname(os.path.realpath(__file__)), "llops")))
imports = dict(enumerate([os.path.splitext(x)[0] for x in _ops if x.startswith("ops_")]))
DEFAULT = None
- buffers = {}
+ buffers, llops = {}, {}
for i,op in imports.items():
name = op[len("ops_"):].upper()
vars()[name] = i
DEFAULT = i if os.environ.get(name, 0) == "1" else DEFAULT
+ try:
+ llops[i] = importlib.import_module('tinygrad.llops.'+op)
+ buffers[i] = llops[i].Buffer
+ except ImportError as e:
+ print(op, "not available", e)
DEFAULT = CPU if DEFAULT is None else DEFAULT
# **** start with two base classes, Tensor and Function ****
@@ -80,7 +85,7 @@ class Device:
class Tensor:
did_float_warning = False
training = False
- ops = defaultdict(dict)
+ ops = {}
def __init__(self, data, device=Device.DEFAULT, requires_grad=True):
self.device, self.data = device, self._move_data(data, device)
@@ -370,15 +375,8 @@ class Function:
self.requires_grad = any(t.requires_grad for t in tensors)
self.saved_tensors = []
- @property
- def op(f):
- # uhhh, obviously stupid
- if f.device == 0:
- return importlib.import_module(f".cpu", f"tinygrad.llops")
- elif f.device == 1:
- return importlib.import_module(f".opencl", f"tinygrad.llops")
- elif f.device == 2:
- return importlib.import_module(f".torch", f"tinygrad.llops")
+ buffer = property(lambda self: Device.buffers[self.device])
+ op = property(lambda self: Device.llops[self.device])
def save_for_backward(self, *x):
if self.requires_grad:
@@ -400,15 +398,17 @@ class Function:
device=ctx.device, requires_grad=any(ctx.needs_input_grad))
po.output = [ret]
if ret.requires_grad:
- ret._ctx = ctx
+ ret._ctx = ctx # used by autograd engine
return ret
-def register(name, fxn, device=Device.CPU):
- Tensor.ops[device][name] = fxn
+def register(name, fxn):
+ Tensor.ops[name] = fxn
def dispatch(*x, **kwargs):
+ # get first tensor in args to determine device
tt = [arg for arg in x if isinstance(arg, Tensor)][0]
+ # create tensors from number arguments
x = [Tensor(np.array([arg], dtype=tt.dtype), device=tt.device, requires_grad=False) if not isinstance(arg, Tensor) else arg for arg in x]
- f = Tensor.ops[tt.device][name]
+ f = Tensor.ops[name] # get the function by device and name
f.device = tt.device
return f.apply(f, *x, **kwargs)
if getattr(Tensor, name, None) is not None:
@@ -420,18 +420,11 @@ def register(name, fxn, device=Device.CPU):
setattr(Tensor, f"__i{name}__", lambda self,x: self.assign(dispatch(self,x)))
setattr(Tensor, f"__r{name}__", lambda self,x: dispatch(x,self))
+# register functions to move between devices
for device in [device for device in Device.__dict__.keys() if device[0] != "_"]:
setattr(Tensor, f"{device.lower()}", functools.partialmethod(Tensor.to, Device.__dict__[device]))
setattr(Tensor, f"{device.lower()}_", functools.partialmethod(Tensor.to_, Device.__dict__[device]))
-# this registers all the operations
-def _register_ops(namespace, device=Device.CPU):
- for name, cls in inspect.getmembers(namespace, inspect.isclass):
- if name.endswith("Buffer"): Device.buffers[device] = cls
- elif name[0] != "_": register(name.lower(), cls, device=device)
-
-for d,ops in Device.imports.items():
- try:
- _register_ops(importlib.import_module('tinygrad.ops.'+ops), d)
- except ImportError:
- pass
+# this registers all the mlops "math" operations
+for name, cls in inspect.getmembers(importlib.import_module('tinygrad.mlops'), inspect.isclass):
+ if name[0] != "_" and name != "Function": register(name.lower(), cls)

proper llops