opencl tests passed
MODIFY
accel/lazy/ops_lazy.py
accel/lazy/ops_lazy.py
@@ -191,7 +191,7 @@ def _realize(self:LazyBuffer) -> Tuple[gops.GPUBuffer, List[gops.GPUBuffer]]:
return gops.GPUBuffer.fromCPU(self.op.arg), []
elif self.optype == LoadOps and self.op.op == LoadOps.CONTIGUOUS:
real_src = self.op.src[0].realize()
- return real_src.contiguous(), [real_src]
+ return real_src.contiguous_op(), [real_src]
elif self.optype == ReduceOps:
real_src = self.op.src[0].realize()
return real_src.reduce_op(self.op.op, self.op.arg), [real_src]

MODIFY
accel/opencl/conv.cl
accel/opencl/conv.cl
#define NUM_OUTPUTS 4
- __kernel void conv(
+//PREFIX
+
+__kernel void image_conv(
read_only image2d_t input,
read_only image2d_t weights,
write_only image2d_t output,
short filterSizeX, short filterSizeY,
short paddingX, short paddingY,
short strideX, short strideY,
- short dilationX, short dilationY) {
+ short dilationX, short dilationY
+ //ARGS
+ ) {
const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;
outputLocation.y = outputRow;
for (short i = 0; i < NUM_OUTPUTS; ++i) {
outputLocation.x = mad24(outputColumn, totalNumPackedOutputChannels, packedOutputChannel);
+ //BINOP
if (outputColumn < numOutputColumns) {
write_imagef(output, outputLocation, outputValues[i]);
}

MODIFY
accel/opencl/ops_opencl.py
accel/opencl/ops_opencl.py
-# this is focused on speed
-# it may not run everything
+from __future__ import annotations
+import pyopencl as cl
+from tinygrad.llops.ops_gpu import GPUBuffer, get_cl_ctx, get_cl_queue, CLProgram, code_for_op
+from tinygrad.ops import ProcessingOps
+from tinygrad.helpers import prod, ConvArgs
+from typing import List, Tuple, Optional, Dict
+import numpy as np
import pathlib
-import numpy as np
-from tinygrad.ops import MovementOps, ProcessingOps
-from tinygrad.llops.ops_gpu import require_init_gpu, clbuild, get_cl_queue, get_cl_ctx
-from tinygrad.llops.ops_gpu import contiguous
-from tinygrad.llops.ops_gpu import unary_op as unary_op_gpu, binary_op as binary_op_gpu, reduce_op as reduce_op_gpu
-from tinygrad.helpers import prod
-from tinygrad.shapetracker import ShapeTracker
-import pyopencl as cl
-from copy import deepcopy
+def load(x):
+ with open(x) as f:
+ ret = f.read()
+ return ret
+CONV_SRC = load(pathlib.Path(__file__).parent.parent.parent / 'accel/opencl/conv.cl')
+
+def get_replacements(prg_src:str, opencl_type:List[str]) -> Dict[str, str]:
+ middle_code = []
+
+ """
+ vv = "xyzw"
+ for i in range(4):
+ acc = f"outputValues[i].{vv[i%4]}"
+ args = [x.split(" ")[-1].replace("*", "") for x in opencl_type]
+ args = [f"(outputRow * get_image_width(output) + outputLocation.x)*4+{i}", acc]+args
+ middle_code.append(f"{acc} = _ewop("+', '.join(args)+");\n")
+ """
+ acc = f"outputValues[i]"
+ args = [x.split(" ")[-1].replace("*", "") for x in opencl_type]
+ args = ["smp", "outputLocation", "(outputLocation.y * get_image_width(output) + outputLocation.x)*4", acc]+args
+ middle_code.append(f"{acc} = _ewop("+', '.join(args)+");\n")
+
+ replacements = {}
+ if len(opencl_type) != 0:
+ replacements["//PREFIX"] = prg_src
+ replacements["//ARGS"] = ","+','.join(opencl_type)
+ replacements["//BINOP"] = ''.join(middle_code)
+ return replacements
def roundup(x, n=4): return (x+(n-1))//n * n
-def flip(x): return (x[1], x[0])
-class OpenCLBuffer:
- def __init__(self, shape, hostbuf=None, _buf=None, _image=None):
- require_init_gpu()
- self.shapetracker = deepcopy(shape) if isinstance(shape, ShapeTracker) else ShapeTracker(*shape)
- self._buf = _buf
- self._image = _image
- self.dtype = np.float32
- if hostbuf is not None:
- # TODO: lazy?
- self._buf = cl.Buffer(get_cl_ctx(), cl.mem_flags.READ_WRITE, 4*roundup(prod(shape)))
- cl.enqueue_copy(get_cl_queue(), self._buf, hostbuf.astype(np.float32).ravel())
-
- def clone(self):
- return OpenCLBuffer(self.shapetracker, _buf=self._buf, _image=self._image)
-
- @property
- def shape(self): return self.shapetracker.shape
+class OpenCLBuffer(GPUBuffer):
+ def __init__(self, shape, hostbuf:Optional[OpenCLBuffer]=None):
+ super().__init__(shape, hostbuf)
+ self._image = hostbuf._image if hostbuf is not None else None
@staticmethod
def fromCPU(x):
- return OpenCLBuffer(x.shape, x)
-
- def toCPU(self):
- data = np.empty(self.shape, dtype=np.float32)
- if self.shapetracker.contiguous == False:
- tmp = OpenCLBuffer(self.shapetracker.shape)
- contiguous(None, self, self.shapetracker, tmp)
- else:
- tmp = self
- cl.enqueue_copy(get_cl_queue(), data, tmp.cl, is_blocking=True)
- return data
+ ret = OpenCLBuffer(x.shape)
+ # TODO: this is blocking even though we told it not to
+ cl.enqueue_copy(get_cl_queue(), ret.cl, x.view(np.ndarray).astype(np.float32).ravel(), is_blocking=False)
+ return ret
@property
def cl(self):
@@ -52,8 +55,8 @@ class OpenCLBuffer:
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
- print(f"converting {self.shape} back to buffer, image shape is {self._image.shape}")
- clbuild("from_image", """
+ #print(f"converting {self.shape} back to buffer, image shape is {self._image.shape}")
+ CLProgram("from_image", """
__kernel void from_image(
read_only image2d_t in,
__global float4 *out) {
@@ -68,16 +71,19 @@ class OpenCLBuffer:
self._image = None
return self._buf
+ def is_image(self): return self._image is not None
+
@property
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)
- self._image = cl.Image(get_cl_ctx(), cl.mem_flags.READ_WRITE, fmt, shape=flip(self.shape))
+ #fmt = cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.FLOAT)
+ 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
- print(f"converting {self.shape} to image with shape {self._image.shape}")
- clbuild("to_image", """
+ #print(f"converting {self.shape} to image with shape {self._image.shape}")
+ CLProgram("to_image", """
__kernel void to_image(
__global const float4 *in,
write_only image2d_t out) {
@@ -91,67 +97,72 @@ class OpenCLBuffer:
self._buf = None
return self._image
-def unary_op(ctx, op, x):
- # TODO: this doesn't actually have to be contiguous
- x = contiguous(ctx, x, x.shapetracker) if not x.shapetracker.contiguous else x
- return unary_op_gpu(ctx, op, x)
-
-def binary_op(ctx, op, x, y):
- x = contiguous(ctx, x, x.shapetracker) if not x.shapetracker.contiguous else x
- y = contiguous(ctx, y, y.shapetracker) if not y.shapetracker.contiguous else y
- return binary_op_gpu(ctx, op, x, y)
+ 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 reduce_op(ctx, op, x, new_shape):
- x = contiguous(ctx, x, x.shapetracker) if not x.shapetracker.contiguous else x
- return reduce_op_gpu(ctx, op, x, new_shape)
+ 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_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:]
+ ewimages = [buf.is_image() and buf.shape == ret.shape and buf.st.contiguous for _,buf in ewbufs]
+ ewtypes = [f"read_only image2d_t {name}_g" if is_img else f"__global const float *{name}_g" for is_img, (name, buf) in zip(ewimages, ewbufs)]
+
+ getters = []
+ for is_img, (name, buf) in zip(ewimages, ewbufs):
+ if is_img:
+ getters.append(f"inline float4 get4_{name}(read_only image2d_t x, const sampler_t smp, int2 loc, int gid) {{ return read_imagef(x, smp, loc); }}")
+ else:
+ getters.append(buf.contiguous_view(name))
+ getters.append(f"inline float4 get4_{name}(__global const float *x, const sampler_t smp, int2 loc, int gid) {{ return (float4)(get_{name}(x,gid+0), get_{name}(x,gid+1), get_{name}(x,gid+2), get_{name}(x,gid+3)); }}")
-def movement_op(ctx, op, x, arg=None):
- xc = x.clone()
- # convert from image if the buffer can change shape
- if op in [MovementOps.EXPAND, MovementOps.SLICE]: xc.cl
- xc.shapetracker.movement_op(op, arg)
- if not xc.shapetracker.contiguous: return contiguous(ctx, xc, xc.shapetracker)
- else: return xc
+ elementwise_prefix = '\n'.join(getters)+ \
+ "\n\ninline float4 _ewop("+','.join(["const sampler_t smp", "int2 loc", "int gid", "float4 acc"]+ewtypes)+") {\n"+ \
+ ''.join([f"float4 {name} = get4_{name}({name}_g, smp, loc, gid);\n" for name, _ in ewbufs])+ \
+ f"return {code}; }}"
-def load(x):
- with open(x) as f:
- ret = f.read()
- return ret
+ replacements = get_replacements(elementwise_prefix, ewtypes)
-def conv(x,w,ret,C):
- print(x.shapetracker.expr(), w.shapetracker.expr())
- print(x.shape, w.shape, ret.shape)
+ x, w = x.contiguous_op(), w.contiguous_op()
options = []
if C.cin == 1: options.append("-DDEPTHWISE")
if C.bs > 1:
options.append("-DBATCH")
assert C.py == 0, "batched conv doesn't work with y-padding"
- conv_prg = clbuild("conv", load(pathlib.Path(__file__).parent.parent.parent / 'accel/opencl/conv.cl'), tuple(options))
assert C.cout%4 == 0
- kernel_args = [C.cout//4, (C.ox+3)//4, C.bs*C.oy]
+ 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]
- print(conv_args, kernel_args)
- conv_prg(kernel_args, None, x.image, w.image, ret.image, *[np.int16(x) for x in conv_args])
-
-
-def processing_op(ctx,op,x,w,out_shape,C):
- assert op == ProcessingOps.CONV, f"{op} isn't supported"
- ret = ctx.buffer((C.bs*C.oy, C.ox*C.cout//4, 4))
- conv(x, w, ret, C)
+ 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)])
return ret
-
-def test_image():
- hostbuf = np.random.randn(5,8,4).astype(np.float32)
- x = OpenCLBuffer((5,8,4), hostbuf)
- assert np.allclose(x.toCPU(), hostbuf)
- print(x.image)
- assert np.allclose(x.toCPU(), hostbuf)
-
-if __name__ == "__main__":
- test_image()
-
-
-
-
-
+GPUBuffer = OpenCLBuffer

opencl tests passed (#347)