begin fixing up opencl
MODIFY
accel/opencl/ops_opencl.py
accel/opencl/ops_opencl.py
from __future__ import annotations
import os
-from tinygrad.llops.ops_gpu import GPUBuffer, CL, CLProgram, code_for_op
+from tinygrad.llops.ops_gpu import GPUBuffer, CL, CLProgram, CLBuffer
from tinygrad.ops import ProcessingOps
from tinygrad.helpers import prod, ConvArgs
from typing import List, Tuple, Optional, Dict
@@ -60,13 +60,15 @@ class OpenCLBuffer(GPUBuffer):
CL.enqueue_copy(ret.cl, x.view(np.ndarray).astype(np.float32).ravel(), is_blocking=False)
return ret
+ def __repr__(self): return f"<OpenCLBuffer with shape {self.shape!r}>"
+
@property
def cl(self):
if self._buf is None:
if self.st.contiguous:
- self._buf = CL.malloc(4*roundup(prod(self.shape)))
+ self._buf = CLBuffer(4*roundup(prod(self.shape)))
if self._image is not None:
- self._buf = CL.malloc(4*roundup(prod(self._image.shape)*4))
+ self._buf = CLBuffer(4*roundup(prod(self._image.shape)*4))
#print(f"converting {self.shape} back to buffer, image shape is {self._image.shape}")
CLProgram("from_image", """
__kernel void from_image(
@@ -79,9 +81,9 @@ class OpenCLBuffer(GPUBuffer):
int W = get_image_width(in);
out[l.y*W + l.x] = read_imagef(in, smp, l);
}
- """)(self._image.shape, None, self._image, self._buf)
+ """)(self._image.shape, None, self._image, self._buf.cl)
self._image = None
- return self._buf
+ return self._buf.cl
def is_image(self): return self._image is not None
@@ -113,9 +115,9 @@ class OpenCLBuffer(GPUBuffer):
# TODO: handle an opencl conv without the conv part
return super()._processing_op(bufs, code, C)
- assert bufs[0][0] == "input" and bufs[1][0] == "weight"
- x,w = bufs[0][1], bufs[1][1]
- ewbufs = bufs[2:]
+ x = [x for x in bufs if x[0] == "input"][0][1]
+ w = [x for x in bufs if x[0] == "weight"][0][1]
+ ewbufs = [x for x in bufs if x[0] not in ["input", "weight"]]
if tuple(bufs[0:2]) in OpenCLBuffer.seen:
print("WARNING: recomputing CONV with", bufs[0], bufs[1])

MODIFY
accel/opencl/preprocessing.py
accel/opencl/preprocessing.py
-from tinygrad.ops import MovementOps
+from tinygrad.ops import MovementOps, ProcessingOps
# input format is N, H x W, C//4 x 4
# dweight format is oc//4 x ch, cw x 4(oc)
# weight format is oc//4 x ch, ic//4, cw, 4(oc) x 4(ic)
-def preprocessing_op(ctx,x,w,C):
- w = ctx.movement_op(MovementOps.RESHAPE, w, (C.groups, C.rcout, C.cin, C.H, C.W))
+def preprocessing_op(x,w,C):
+ w = w.movement_op(MovementOps.RESHAPE, (C.groups, C.rcout, C.cin, C.H, C.W))
#print(x.shape, w.shape)
if C.bs > 1 and C.py > 0:
@@ -12,7 +12,7 @@ def preprocessing_op(ctx,x,w,C):
# N C H W
xs = [(0, s) for s in x.shape]
xs[2] = (-C.py, x.shape[2]+C.py)
- x = ctx.movement_op(MovementOps.SLICE, x, xs)
+ x = x.movement_op(MovementOps.SLICE, xs)
C = C._replace(iy=C.iy + C.py*2, py=0)
# hack for non multiples of 4 on C.cin
@@ -20,55 +20,60 @@ def preprocessing_op(ctx,x,w,C):
to_add = 4 - (C.cin % 4)
ws = [(0, s) for s in w.shape]
ws[2] = (0, w.shape[2]+to_add)
- w = ctx.movement_op(MovementOps.SLICE, w, ws)
+ w = w.movement_op(MovementOps.SLICE, ws)
- x = ctx.movement_op(MovementOps.RESHAPE, x, (C.bs, C.groups, C.cin, C.iy, C.ix))
+ x = x.movement_op(MovementOps.RESHAPE, (C.bs, C.groups, C.cin, C.iy, C.ix))
xs = [(0, s) for s in x.shape]
xs[2] = (0, x.shape[2]+to_add)
- x = ctx.movement_op(MovementOps.SLICE, x, xs)
+ x = x.movement_op(MovementOps.SLICE, xs)
C = C._replace(cin = C.cin + to_add)
- x = ctx.movement_op(MovementOps.RESHAPE, x, (C.bs, C.groups*C.cin, C.iy, C.ix))
+ x = x.movement_op(MovementOps.RESHAPE, (C.bs, C.groups*C.cin, C.iy, C.ix))
# hack for non multiples of 4 on C.rcout
if C.rcout % 4 != 0 and not (C.rcout == 1 and C.groups%4 == 0):
added_output_channels = 4 - (C.rcout % 4)
ws = [(0, s) for s in w.shape]
ws[1] = (0, w.shape[1]+added_output_channels)
- w = ctx.movement_op(MovementOps.SLICE, w, ws)
+ w = w.movement_op(MovementOps.SLICE, ws)
C = C._replace(rcout = C.rcout + added_output_channels, cout = C.groups * (C.rcout + added_output_channels))
# packed
assert (C.groups*C.cin) % 4 == 0
#print(x.shape)
- x = ctx.movement_op(MovementOps.PERMUTE, x, (0,2,3,1))
- x = ctx.movement_op(MovementOps.RESHAPE, x, (C.bs*C.iy, C.ix*C.groups*C.cin//4, 4))
+ x = x.movement_op(MovementOps.PERMUTE, (0,2,3,1))
+ x = x.movement_op(MovementOps.RESHAPE, (C.bs*C.iy, C.ix*C.groups*C.cin//4, 4))
assert C.cout % 4 == 0
if C.cin == 1:
# depthwise
- w = ctx.movement_op(MovementOps.RESHAPE, w, (C.cout//4,4,C.H*C.W))
- w = ctx.movement_op(MovementOps.PERMUTE, w, (0,2,1))
+ w = w.movement_op(MovementOps.RESHAPE, (C.cout//4,4,C.H*C.W))
+ w = w.movement_op(MovementOps.PERMUTE, (0,2,1))
else:
- w = ctx.movement_op(MovementOps.RESHAPE, w, (C.cout//4,4,C.cin//4,4,C.H,C.W))
- w = ctx.movement_op(MovementOps.PERMUTE, w, (0,4,2,5,1,3))
- w = ctx.movement_op(MovementOps.RESHAPE, w, (C.cout//4, C.H * C.cin//4 * C.W * 4, 4))
+ w = w.movement_op(MovementOps.RESHAPE, (C.cout//4,4,C.cin//4,4,C.H,C.W))
+ w = w.movement_op(MovementOps.PERMUTE, (0,4,2,5,1,3))
+ w = w.movement_op(MovementOps.RESHAPE, (C.cout//4, C.H * C.cin//4 * C.W * 4, 4))
C = C._replace(out_shape = (C.bs*C.oy, C.ox*C.cout//4, 4))
#x = contiguous(ctx, x, x.shapetracker) if not x.shapetracker.contiguous else x
#w = contiguous(ctx, w, w.shapetracker) if not w.shapetracker.contiguous else w
return x,w,C
-def postprocessing_op(ctx, ret, C, C_initial):
+def postprocessing_op(ret, C, C_initial):
added_output_channels = C.rcout - C_initial.rcout
# undo hack for non multiples of 4 on C.rcout
if added_output_channels != 0:
- ret = ctx.movement_op(MovementOps.RESHAPE, ret, (C.bs, C.oy, C.ox, C.groups, C.rcout))
+ ret = ret.movement_op(MovementOps.RESHAPE, (C.bs, C.oy, C.ox, C.groups, C.rcout))
xs = [(0, s) for s in ret.shape]
xs[4] = (0, ret.shape[4]-added_output_channels)
- ret = ctx.movement_op(MovementOps.SLICE, ret, xs)
+ ret = ret.movement_op(MovementOps.SLICE, xs)
C = C._replace(rcout = C.rcout - added_output_channels, cout = C.groups * (C.rcout - added_output_channels))
- ret = ctx.movement_op(MovementOps.RESHAPE, ret, (C.bs, C.oy, C.ox, C.cout))
- ret = ctx.movement_op(MovementOps.PERMUTE, ret, (0,3,1,2))
+ ret = ret.movement_op(MovementOps.RESHAPE, (C.bs, C.oy, C.ox, C.cout))
+ ret = ret.movement_op(MovementOps.PERMUTE, (0,3,1,2))
return ret
+
+def processed_conv(x, w, C):
+ x,w,Cn = preprocessing_op(x,w,C)
+ ret = x.processing_op(ProcessingOps.CONV, w, Cn)
+ return postprocessing_op(ret, Cn, C)

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -86,7 +86,7 @@ class GPUBuffer:
return self._buf.cl
def __repr__(self): return f"<GPUBuffer with shape {self.shape!r}>"
- def shapeTrackerView(self, st:ShapeTracker): return GPUBuffer(ShapeTracker(st), hostbuf=self)
+ def shapeTrackerView(x, st:ShapeTracker): return type(x)(ShapeTracker(st), hostbuf=x)
@staticmethod
def fromCPU(x): return GPUBuffer(x.shape, backing=x.view(np.ndarray).astype(np.float32).ravel())

ADD
tinygrad/llops/ops_opencl.py
+../../accel/opencl/ops_opencl.py
\ No newline at end of file

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -159,6 +159,10 @@ class Conv2D(Function):
def _conv(ctx, x, w, C):
# TODO: this does NOT belong here
# was pre/post processing for opencl
+ if x.device == "OPENCL":
+ from accel.opencl.preprocessing import processed_conv
+ return processed_conv(x, w, C)
+ else:
return x.processing_op(ProcessingOps.CONV, w, C)
def forward(ctx, x, w, stride=1, groups=1, dilation=1, padding=0):

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -231,7 +231,7 @@ class LazyBuffer:
def movement_op(x:LazyBuffer, op:MovementOps, arg) -> LazyBuffer:
# TODO: look into why that copy is needed
- arg = copy(arg)
+ arg = tuple(copy(arg))
# TODO: SHUFFLE_SLICE_OPS is okay if it's a shrink
if SHUFFLE_MOVEMENT_OPS and x.optype == BinaryOps and x.realized is None and (SHUFFLE_SLICE_OPS or op != MovementOps.SLICE):

begin fixing up opencl