lazy opencl works
MODIFY
accel/lazy/ops_lazy.py
accel/lazy/ops_lazy.py
from __future__ import annotations
+import os
from typing import Union, NamedTuple, List, Any, Tuple, Dict
from tinygrad.shapetracker import ShapeTracker
import functools, operator
@@ -150,8 +151,10 @@ def elementwise_op(op, srcs:Tuple[LazyBuffer]) -> LazyBuffer:
return LazyBuffer(out_shape, BinaryOps, LazyOp(op, srcs))
-
# these functions determines the backing buffer
+if int(os.getenv("LAZY_OPENCL", 0)) == 1:
+ import tinygrad.llops.ops_opencl as gops
+else:
import tinygrad.llops.ops_gpu as gops
def _realize_binary_op(self:LazyBuffer) -> Tuple[gops.GPUBuffer, List[gops.GPUBuffer]]:

MODIFY
accel/opencl/ops_opencl.py
accel/opencl/ops_opencl.py
@@ -52,9 +52,10 @@ class OpenCLBuffer(GPUBuffer):
@property
def cl(self):
if self._buf is None:
+ if self.st.contiguous:
self._buf = cl.Buffer(get_cl_ctx(), cl.mem_flags.READ_WRITE, 4*roundup(prod(self.shape)))
if self._image is not None:
- assert prod(self.shape) == prod(self._image.shape)*4
+ self._buf = cl.Buffer(get_cl_ctx(), cl.mem_flags.READ_WRITE, 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(
@@ -77,8 +78,9 @@ class OpenCLBuffer(GPUBuffer):
def image(self):
if self._image is None:
assert self.shape[2] == 4 and len(self.shape) == 3
- #fmt = cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.FLOAT)
- fmt = cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.HALF_FLOAT)
+ fmt = cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.FLOAT)
+ # HALF_FLOAT breaks tests
+ #fmt = cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.HALF_FLOAT)
self._image = cl.Image(get_cl_ctx(), cl.mem_flags.READ_WRITE, fmt, shape=(self.shape[1], self.shape[0]))
if self._buf is not None:
assert prod(self.shape) == prod(self._image.shape)*4
@@ -97,18 +99,11 @@ class OpenCLBuffer(GPUBuffer):
self._buf = None
return self._image
- def movement_op(x, op, arg):
- x.cl
- # TODO: call super after x.cl
- ret = type(x)(x.st, x)
- ret.shape = ret.st.movement_op(op, arg).shape
- return ret
-
- def processing_op(x, op, w, C:ConvArgs):
- assert op == ProcessingOps.CONV, f"{op} isn't supported"
- return type(x)(C.out_shape)._processing_op_cl([("input", x), ("weight", w)], "acc", C)
+ def _processing_op(ret, bufs: List[Tuple[str, OpenCLBuffer]]=[], code:str="acc", C=None):
+ if C is None:
+ # TODO: handle an opencl conv without the conv part
+ return super()._processing_op(bufs, code, C)
- def _processing_op_cl(ret, bufs: List[Tuple[str, OpenCLBuffer]]=[], code:str="acc", C=None):
assert bufs[0][0] == "input" and bufs[1][0] == "weight"
x,w = bufs[0][1], bufs[1][1]
ewbufs = bufs[2:]
@@ -140,29 +135,14 @@ class OpenCLBuffer(GPUBuffer):
conv_src = CONV_SRC
for k,v in replacements.items():
conv_src = conv_src.replace(k, v)
- #if any(ewimages):
#print(conv_src)
- # exit(0)
conv_prg = CLProgram("image_conv", conv_src,
options=tuple(options),
argdtypes=tuple([None, None, None] + [np.int16]*15 + [None]*len(ewbufs))
)
conv_args = [max(1, C.cin//4), C.groups*C.cin//4, max(1, C.rcout//4), C.cout//4, C.ox, C.oy, C.iy, C.W, C.H, C.px, C.py, C.xs, C.ys, C.dx, C.dy]
global_work_size = [C.cout//4, (C.ox+3)//4, C.bs*C.oy]
- local_work_size = [1, (C.ox+3)//4, 1]
-
- MAX_LOCAL_WORK_SIZE = 256
-
- n = 1
- while global_work_size[0] % n == 0 and n*local_work_size[1]*local_work_size[2] <= MAX_LOCAL_WORK_SIZE:
- local_work_size[0] = n
- n *= 2
- n = 1
- while global_work_size[2] % n == 0 and n*local_work_size[0]*local_work_size[1] <= MAX_LOCAL_WORK_SIZE:
- local_work_size[2] = n
- n *= 2
- #print(global_work_size, local_work_size)
- conv_prg(global_work_size, local_work_size, x.image, w.image, ret.image, *conv_args, *[buf.image if is_img else buf.cl for is_img, (_, buf) in zip(ewimages, ewbufs)])
+ conv_prg(global_work_size, None, x.image, w.image, ret.image, *conv_args, *[buf.image if is_img else buf.cl for is_img, (_, buf) in zip(ewimages, ewbufs)])
return ret
GPUBuffer = OpenCLBuffer

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -70,7 +70,7 @@ class GPUBuffer:
def contiguous_op(x): return x if x.st.contiguous else x.unary_op(UnaryOps.NOOP)
def movement_op(x, op:MovementOps, arg) -> GPUBuffer:
- ret = GPUBuffer(x.st, x)
+ ret = type(x)(x.st, x)
ret.shape = ret.st.movement_op(op, arg).shape
return ret
@@ -79,7 +79,7 @@ class GPUBuffer:
return type(x)(C.out_shape)._processing_op([("input", x.contiguous_op()), ("weight", w.contiguous_op())], "acc", C)
def reduce_op(x, op:ReduceOps, new_shape:Tuple[int]):
- ret = GPUBuffer(new_shape)
+ ret = type(x)(new_shape)
if op == ReduceOps.SUM: code, start = "out += a", "0.0"
elif op == ReduceOps.MAX: code, start = "out = max(a,out)", "-INFINITY"
else: raise Exception(f"{op} isn't supported")

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

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
+import os
import numpy as np # TODO: remove this, it's used for np.prod and np.argsort
from tinygrad.helpers import prod, reduce_shape, get_conv_args
from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps
@@ -167,10 +168,20 @@ class Slice(Function):
# ************* processing ops *************
class Conv2D(Function):
+ # TODO: this does NOT belong here
+ def _conv(ctx, x, w, C):
+ if "OPENCL" in ctx.device or int(os.getenv("LAZY_OPENCL", 0)):
+ from accel.opencl.preprocessing import preprocessing_op, postprocessing_op
+ x,w,Cmod = preprocessing_op(ctx, x, w, C)
+ ret = ctx.processing_op(ProcessingOps.CONV, x, w, Cmod)
+ return postprocessing_op(ctx, ret, Cmod, C)
+ else:
+ return ctx.processing_op(ProcessingOps.CONV, x, w, C)
+
def forward(ctx, x, w, stride=1, groups=1, dilation=1, padding=0):
C = get_conv_args(x.shape, w.shape, stride, groups, dilation=dilation, padding=padding)
ctx.save_for_backward(x,w,C)
- return ctx.processing_op(ProcessingOps.CONV, x, w, C)
+ return ctx._conv(x, w, C)
def backward(ctx, grad_output):
x, w, C = ctx.saved_tensors
@@ -188,7 +199,7 @@ class Conv2D(Function):
wt = ctx.movement_op(MovementOps.RESHAPE, wt, (C.groups*C.cin, C.rcout, C.H, C.W))
Cdx = get_conv_args(xt.shape, wt.shape, dilation=(C.dy, C.dx), padding=((C.H-1)*C.dy-C.py,(C.W-1)*C.dx-C.px), groups=C.groups)
# TODO: this shape can be wrong. support asymmetric padding to remove the slice
- dx = ctx.processing_op(ProcessingOps.CONV, xt, wt, Cdx)
+ dx = ctx._conv(xt, wt, Cdx)
dx = ctx.movement_op(MovementOps.SLICE, dx, [(0,s) for s in x.shape])
if ctx.needs_input_grad[1]:
@@ -199,7 +210,7 @@ class Conv2D(Function):
grad_output_dw = ctx.movement_op(MovementOps.PERMUTE, grad_output, (1,0,2,3))
grad_output_dw = ctx.movement_op(MovementOps.RESHAPE, grad_output_dw, (C.cout, C.bs, C.oy, C.ox))
Cdw = get_conv_args(xdw.shape, grad_output_dw.shape, padding=(C.py, C.px), stride=(C.dy, C.dx), dilation=(C.ys, C.xs), groups=C.groups)
- grad_weight = ctx.processing_op(ProcessingOps.CONV, xdw, grad_output_dw, Cdw)
+ grad_weight = ctx._conv(xdw, grad_output_dw, Cdw)
grad_weight = ctx.movement_op(MovementOps.PERMUTE, grad_weight, (1,0,2,3))
# TODO: remove this slice using asymmetric padding
dw = ctx.movement_op(MovementOps.SLICE, grad_weight, ((0, grad_weight.shape[0]), (0, grad_weight.shape[1]), (0, w.shape[2]), (0, w.shape[3])))

lazy opencl works