Improve compatibility with pyopencl

[haesleinhuepf]

At the moment, when using pyopencl and clesperanto together, pyclesperanto breaks pyopencl in some contexts, e.g.

• working with complex numbers (see also #115)
• adding images with > 3 dimensions

Related: https://forum.image.sc/t/migrating-from-clij-to-clesperanto/54985/50

Example code adding two 5D-images:

import numpy as np
import pyopencl as cl
import pyopencl.array as cla

# initialize GPU driver and create a queue for running operations.
context = cl.create_some_context()
queue = cl.CommandQueue(context)

# copy data to GPU memory
cl_a = cla.to_device(queue, np.asarray([[[[1, 2, 3]]]]))
cl_b = cla.to_device(queue, np.asarray([[[[4, 4, 5]]]]))  

# execute operations on the data
cl_c = cl_a + cl_b

# show result
cl_c

This works. After I import pyclesperanto_prototype, the + operator gets [overwritten]() and repeating the same code fails:

import pyclesperanto_prototype as cle
cl_c = cl_a + cl_b

Error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_5296/3440715320.py in <module>
      1 import pyclesperanto_prototype as cle
----> 2 cl_c = cl_a + cl_b

c:\structure\code\pyclesperanto_prototype\pyclesperanto_prototype\_tier0\_pycl.py in add(x1, x2)
    182         else:
    183             from .._tier1 import add_images_weighted
--> 184             return add_images_weighted(x1, x2)
    185     cls.__add__ = add
    186 

c:\structure\code\pyclesperanto_prototype\pyclesperanto_prototype\_tier0\_plugin_function.py in worker_function(*args, **kwargs)
     69 
     70         # call the decorated function
---> 71         return function(*bound.args, **bound.kwargs)
     72 
     73     return worker_function

c:\structure\code\pyclesperanto_prototype\pyclesperanto_prototype\_tier1\_add_images_weighted.py in add_images_weighted(summand1, summand2, destination, factor1, factor2)
     50     }
     51 
---> 52     execute(__file__, '../clij-opencl-kernels/kernels/add_images_weighted_' + str(len(destination.shape)) + 'd_x.cl', 'add_images_weighted_' + str(len(destination.shape)) + 'd', destination.shape, parameters)
     53 
     54     return destination

c:\structure\code\pyclesperanto_prototype\pyclesperanto_prototype\_tier0\_execute.py in execute(anchor, opencl_kernel_filename, kernel_name, global_size, parameters, prog, constants, image_size_independent_kernel_compilation, device)
    155             depth_height_width = [1, 1, 1]
    156             depth_height_width[-len(value.shape) :] = value.shape
--> 157             depth, height, width = depth_height_width
    158             ndim = value.ndim
    159 

ValueError: too many values to unpack (expected 3)

[tlambert03]

Oh no! This feels very high priority. definitely can't directly monkeypatch the pyopencl Array class. You can use a subclass instead of wrapOclArray if you want to add those magic methods for use in this lib.

Curious though, why didn't the builtin __add__ and stuff work?

[haesleinhuepf]

You can use a subclass instead of wrapOclArray if you want to add those magic methods for use in this lib.

I'm on it. I think I will struggle with constructing a cl-array. But let's see. I'll send it as PR to you for review asap.

Curious though, why didn't the builtin __add__ and stuff work?

pyopencl doesn't support cl_array + np_array, cl_array + cupy_array and cl_array + dask_array. pyclesperanto does.

[tlambert03]

pyopencl doesn't support

I think the best way to fix that might be to add some special numpy dunder methods to your subclass (not necessarily just to call pyclesperanto methods) see here: https://numpy.org/doc/stable/user/basics.dispatch.html

will write a short example

[tlambert03]

edited

from numpy.lib.arraysetops import isin
import pyopencl as cl
import numpy as np
import pyopencl.array as cla
import numpy.lib.mixins

context = cl.create_some_context(interactive=False)
queue = cl.CommandQueue(context)

class OcA(cla.Array, numpy.lib.mixins.NDArrayOperatorsMixin):
    def __array__(self, dtype=None):
        return self.get().astype(dtype)

    @classmethod
    def to_device(cls, ary, queue):
        if isinstance(ary, OcA):
            return ary
        cl_a = OcA(queue, ary.shape, ary.dtype, strides=ary.strides)
        cl_a.set(ary, queue=queue)
        return cl_a

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        if method == "__call__":
            func = getattr(OcA, f'__{ufunc.__name__}__', None)
            if func is not None:
                return func(*[OcA.to_device(i, self.queue) for i in inputs], **kwargs)
        return NotImplemented

np_a = np.random.rand(4, 4)
cl_a = OcA.to_device(np_a, queue=queue)

# now this will work
cl_b = cl_a + np_a

assert isinstance(cl_a, OcA)
assert isinstance(np_a, np.ndarray)
assert isinstance(cl_b, OcA)

more details at the link above

[tlambert03]

i updated the code above to give an example of how a subclass of cla.Array can overload the + operator to take np.arrays as input while keeping the data on the GPU