reduce line count by simplifying DeviceBuffer
MODIFY
tinygrad/device.py
tinygrad/device.py
-import os
-
def get_available_llops():
- import importlib, inspect
+ import os, importlib, inspect
_buffers, DEFAULT = {}, "CPU"
for op in [os.path.splitext(x)[0] for x in sorted(os.listdir(os.path.join(os.path.dirname(os.path.realpath(__file__)), "llops"))) if x.startswith("ops_")]:
name = op[len("ops_"):].upper()

MODIFY
tinygrad/lazy.py
tinygrad/lazy.py
@@ -132,7 +132,7 @@ class LazyBuffer:
else:
# movement ops aren't an AST, just run them
real_src = src.realize(self.device)
- self.realized = real_src.movement_op(self.op.op, self.op.arg)
+ self.realized = real_src.movement_op(self.op.op, self.op.arg) # movement_op stays
ast = LazyOp(self.op.op, (real_src, ))
elif self.optype == ProcessingOps: ast = self.op # no ast modifications for ProcessingOps
elif self.optype == ReduceOps: ast = _ast_reduceops(self)

MODIFY
tinygrad/llops/ops_cpu.py
tinygrad/llops/ops_cpu.py
@@ -12,11 +12,21 @@ base_fxn_for_op : Dict[Op, Callable] = {
MovementOps.SHRINK: lambda x, arg: x[tuple(slice(p[0], p[1], None) for p in arg)],
}
+def numpy_strided(x,arg): return np.lib.stride_tricks.as_strided(x.ravel().reshape(x.shape), shape=[y[0] for y in arg], strides=[y[1]*x.dtype.itemsize for y in arg])
+def numpy_conv(x,w,C):
+ assert C.px == 0 and C.px_ == 0 and C.py == 0 and C.py_ == 0, "padding in conv is not supported"
+ tx = numpy_strided(x, (
+ (C.bs, C.groups*C.cin*x.shape[2]*x.shape[3]), (C.groups, C.cin*x.shape[2]*x.shape[3]),
+ (C.oy, C.sy*x.shape[3]), (C.ox, C.sx), (C.cin, x.shape[2]*x.shape[3]), (C.H, C.dy*x.shape[3]), (C.W, C.dx)))
+ tw = w.reshape(C.groups, C.rcout, C.cin, C.H, C.W)
+ out = np.einsum("nGhwCHW, GkCHW -> nGkhw", tx.ravel().reshape(tx.shape), tw.ravel().reshape(tw.shape))
+ return out.reshape(C.bs, C.groups*C.rcout, C.oy, C.ox)
+
numpy_fxn_for_op : Dict[Op, Callable] = {**base_fxn_for_op, **{
UnaryOps.RELU: lambda x: np.maximum(x, 0), UnaryOps.EXP: lambda x: np.exp(x), UnaryOps.LOG: lambda x: np.log(x), BinaryOps.CMPEQ: lambda x,y: (x==y).astype(np.float32),
MovementOps.FLIP: lambda x, axis: np.flip(x, axis), MovementOps.PERMUTE: lambda x, order: x.transpose(order),
MovementOps.PAD: lambda x, padding: np.pad(x, padding), MovementOps.EXPAND: lambda x, new_shape: np.broadcast_to(x, new_shape),
- MovementOps.STRIDED: lambda x, arg: np.lib.stride_tricks.as_strided(x.ravel().reshape(x.shape), shape=[y[0] for y in arg], strides=[y[1]*x.dtype.itemsize for y in arg])
+ MovementOps.STRIDED: numpy_strided, ProcessingOps.CONV: numpy_conv
}}
class CPUBuffer(GenericExecAST):
@@ -25,13 +35,3 @@ class CPUBuffer(GenericExecAST):
@staticmethod
def fromCPU(x): return CPUBuffer(x)
def toCPU(x): return x.buf
-
- def processing_op(x,op,w,C):
- assert op == ProcessingOps.CONV, f"{op} isn't supported"
- assert C.px == 0 and C.px_ == 0 and C.py == 0 and C.py_ == 0, "padding in conv is not supported"
- tx = x.movement_op(MovementOps.STRIDED, (
- (C.bs, C.groups*C.cin*x.shape[2]*x.shape[3]), (C.groups, C.cin*x.shape[2]*x.shape[3]),
- (C.oy, C.sy*x.shape[3]), (C.ox, C.sx), (C.cin, x.shape[2]*x.shape[3]), (C.H, C.dy*x.shape[3]), (C.W, C.dx)))
- tw = w.movement_op(MovementOps.RESHAPE, (C.groups, C.rcout, C.cin, C.H, C.W))
- out = np.einsum("nGhwCHW, GkCHW -> nGkhw", tx.buf.ravel().reshape(tx.shape), tw.buf.ravel().reshape(tw.shape))
- return CPUBuffer(out.reshape(C.bs, C.groups*C.rcout, C.oy, C.ox))
\ No newline at end of file

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -328,7 +328,6 @@ class GPUBuffer(ExplicitExecAST):
def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], hostbuf:Optional[GPUBuffer]=None, backing:Optional[np.ndarray]=None, force_create=False):
super().__init__(shape, hostbuf)
self._buf : Optional[Union[CLImage, CLBuffer]] = hostbuf._buf if hostbuf is not None else None
- self._base_shape : Tuple[int, ...] = hostbuf._base_shape if hostbuf is not None else self.shape
self._backing : Optional[np.ndarray] = hostbuf._backing if hostbuf is not None else backing
# early copy in for large buffers
if (self._backing is not None and self._backing.shape != (1,)) or force_create:
@@ -352,8 +351,8 @@ class GPUBuffer(ExplicitExecAST):
def fromCPU(x): return GPUBuffer(x.shape, backing=x.view(np.ndarray).astype(np.float32).ravel())
def toCPU(self) -> np.ndarray:
- cl_buf = self.unary_op(UnaryOps.NOOP) if not self.st.contiguous or prod(self._base_shape) != prod(self.shape) else self
- cl_buf = cl_buf if isinstance(cl_buf._buf, CLBuffer) else self.movement_op(MovementOps.RESHAPE, tuple(list(self.shape)+[1])).unary_op(UnaryOps.NOOP)
+ cl_buf = self.contiguous()
+ cl_buf = cl_buf if isinstance(cl_buf._buf, CLBuffer) else self.movement_op(MovementOps.RESHAPE, tuple(list(self.shape)+[1])).contiguous()
assert prod(cl_buf._base_shape) == prod(self.shape), f"shape product mismatch {cl_buf._base_shape} vs {self.shape}"
data = np.empty(self.shape, dtype=np.float32)
cl_buf._buf.copyout(data)

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -61,27 +61,23 @@ class GenericExecAST(DeviceBuffer): # pylint: disable=abstract-method
fxn_for_op : ClassVar = shape_fxn_for_op
# TODO: use generic types here to remove __init__ in specialized classes
def __init__(self, lbuf:Any): self.buf, self.shape = lbuf, tuple(lbuf.shape)
- def contiguous(self): return self.unary_op(UnaryOps.NOOP)
- def unary_op(self, op): return type(self)(self.fxn_for_op[op](self.buf))
- def binary_op(self, op, y): return type(self)(self.fxn_for_op[op](self.buf, y.buf))
- def reduce_op(self, op, new_shape): return type(self)(self.fxn_for_op[op](self.buf, new_shape))
- def movement_op(self, op, arg=None): return type(self)(self.fxn_for_op[op](self.buf, arg)) if op in self.fxn_for_op else type(self)(getattr(self.buf, op.name.lower())(arg))
- def processing_op(self, op, w, C): return type(self)(self.fxn_for_op[op](self.buf, w.buf, C))
+ def contiguous(self): return type(self).exec_ast(LazyOp(op=UnaryOps.NOOP, src=(self,)))
+ def movement_op(self, op:MovementOps, arg=None): return type(self)(self.fxn_for_op[op](self.buf, arg)) if op in self.fxn_for_op else type(self)(getattr(self.buf, op.name.lower())(arg))
@classmethod
def exec_ast(cls, ast:LazyOp, output_buffer:Optional[GenericExecAST]=None, preprocess=lambda x: x):
srcs = [cls.exec_ast(x, preprocess=preprocess) if isinstance(x, LazyOp) else preprocess(x) for x in ast.src]
if ast.op in UnaryOps:
- ret = srcs[0].unary_op(ast.op)
+ ret = type(srcs[0])(srcs[0].fxn_for_op[ast.op](srcs[0].buf))
elif ast.op in BinaryOps:
assert srcs[0].shape == srcs[1].shape, f"BinaryOps shape mismatch {srcs[0].shape} != {srcs[1].shape}"
- ret = srcs[0].binary_op(ast.op, srcs[1])
+ ret = type(srcs[0])(srcs[0].fxn_for_op[ast.op](srcs[0].buf, srcs[1].buf))
elif ast.op in ReduceOps:
assert all(r == n or n == 1 for r,n in zip(srcs[0].shape, ast.arg)), f"ReduceOps can't reduce {srcs[0].shape} -> {ast.arg}"
- ret = srcs[0].reduce_op(ast.op, ast.arg)
+ ret = type(srcs[0])(srcs[0].fxn_for_op[ast.op](srcs[0].buf, ast.arg))
elif ast.op in MovementOps:
ret = srcs[0].movement_op(ast.op, ast.arg)
elif ast.op in ProcessingOps:
- ret = srcs[0].processing_op(ast.op, srcs[1], ast.arg)
+ ret = type(srcs[0])(srcs[0].fxn_for_op[ast.op](srcs[0].buf, srcs[1].buf, ast.arg))
else:
raise TypeError("unknown op")
if output_buffer is not None:
@@ -107,15 +103,8 @@ class ExplicitExecAST(DeviceBuffer): # pylint: disable=abstract-method
def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], hostbuf=None):
self.st = shape if isinstance(shape, ShapeTracker) else ShapeTracker(tuple(shape))
self.shape = self.st.shape
-
- # universal
- def unary_op(self, op:UnaryOps): return type(self)(self.shape).exec_ast(LazyOp(op=op, src=(self,)))
- def binary_op(self, op:BinaryOps, y): return type(self)(self.shape).exec_ast(LazyOp(op=op, src=(self, y)))
- def reduce_op(self, op:ReduceOps, new_shape:Tuple[int, ...]): return type(self)(new_shape).exec_ast(LazyOp(op=op, src=(self,), arg=new_shape))
+ self._base_shape : Tuple[int, ...] = hostbuf._base_shape if hostbuf is not None else self.shape
# universal for shape tracked
+ def contiguous(self): return self if self.st.contiguous and prod(self._base_shape) == prod(self.shape) else type(self).exec_ast(LazyOp(op=UnaryOps.NOOP, src=(self,)))
def movement_op(self, op:MovementOps, arg): return type(self)(ShapeTracker(self.st).movement_op(op, arg), self)
-
- # TODO: creating a new object is making a copy, breaking the thneed compiler
- def contiguous(self): return self if self.st.contiguous else self.unary_op(UnaryOps.NOOP)
- #def contiguous(self): return type(self)(self.shape, hostbuf=self) if self.st.contiguous else self.unary_op(UnaryOps.NOOP)

reduce line count by simplifying DeviceBuffer