Multi gpu support

pyclesperanto_prototype/_tier0/__init__.py

@@ -35,7 +35,7 @@
 from ._plugin_function import plugin_function
 from ._types import Image
 from ._cl_info import cl_info
-from ._device import get_device, select_device, set_device_scoring_key
+from ._device import get_device, select_device, set_device_scoring_key, new_device, Device
 from ._cl_image import create_image, empty_image_like, empty_image
 from ._available_device_names import available_device_names
 from ._set_wait_for_kernel_finish import set_wait_for_kernel_finish

pyclesperanto_prototype/_tier0/_create.py

@@ -1,8 +1,8 @@
 from ._pycl import OCLArray
 import numpy as np
+from ._device import Device, get_device
 
-
-def create(dimensions, dtype=np.float32):
+def create(dimensions, dtype=np.float32, device:Device = None):
 
     """
     Convenience method for creating images on the GPU. This method basically does the same as in CLIJ:
@@ -18,129 +18,131 @@ def create(dimensions, dtype=np.float32):
         if isinstance(dimensions, OCLArray)
         else tuple(dimensions)  # reverses a list/tuple
     )
-    return OCLArray.empty(dimensions, dtype)
+    if device is None:
+        device = get_device()
+    return device.empty(dimensions, dtype)
 
-def create_zyx(dimensions):
-    return create(dimensions[::-1])
+def create_zyx(dimensions, device:Device = None):
+    return create(dimensions[::-1], device=device)
 
-def create_like(*args):
+def create_like(*args, device:Device = None):
     dimensions = args[0]
     if isinstance(dimensions, OCLArray):
         dimensions = dimensions.shape
     elif isinstance(dimensions, np.ndarray):
         dimensions = dimensions.shape[::-1]
-    return create(dimensions)
+    return create(dimensions, device=device)
 
-def create_binary_like(*args):
+def create_binary_like(*args, device:Device = None):
     dimensions = args[0]
     if isinstance(dimensions, OCLArray):
         dimensions = dimensions.shape
     elif isinstance(dimensions, np.ndarray):
         dimensions = dimensions.shape[::-1]
-    return create(dimensions, np.uint8)
+    return create(dimensions, np.uint8, device=device)
 
-def create_labels_like(*args):
+def create_labels_like(*args, device:Device = None):
     dimensions = args[0]
     if isinstance(dimensions, OCLArray):
         dimensions = dimensions.shape
     elif isinstance(dimensions, np.ndarray):
         dimensions = dimensions.shape[::-1]
-    return create(dimensions, np.uint32)
+    return create(dimensions, np.uint32, device=device)
 
-def create_pointlist_from_labelmap(input:OCLArray, *args):
+def create_pointlist_from_labelmap(input:OCLArray, *args, device:Device = None):
     from .._tier2 import maximum_of_all_pixels
     number_of_labels = int(maximum_of_all_pixels(input))
     number_of_dimensions = len(input.shape)
 
-    return create([number_of_dimensions, number_of_labels])
+    return create([number_of_dimensions, number_of_labels], device=device)
 
-def create_vector_from_labelmap(input: OCLArray, *args):
+def create_vector_from_labelmap(input: OCLArray, *args, device:Device = None):
     from .._tier2 import maximum_of_all_pixels
     number_of_labels = int(maximum_of_all_pixels(input)) + 1
 
-    return create([1, number_of_labels])
+    return create([1, number_of_labels], device=device)
 
-def create_matrix_from_pointlists(pointlist1:OCLArray, pointlist2:OCLArray):
+def create_matrix_from_pointlists(pointlist1:OCLArray, pointlist2:OCLArray, device:Device = None):
     width = pointlist1.shape[1] + 1
     height = pointlist2.shape[1] + 1
 
-    return create([width, height])
+    return create([width, height], device=device)
 
-def create_from_pointlist(pointlist: OCLArray, *args):
+def create_from_pointlist(pointlist: OCLArray, *args, device:Device = None):
     from .._tier1 import maximum_x_projection
     from .._tier0 import pull
 
     max_pos = pull(maximum_x_projection(pointlist)).T.astype(int)
     max_pos = max_pos[0]
 
     if len(max_pos) == 3:  # 3D image requested
-        destination = create([max_pos[2] + 1, max_pos[1] + 1, max_pos[0] + 1])
+        destination = create([max_pos[2] + 1, max_pos[1] + 1, max_pos[0] + 1], device=device)
     elif len(max_pos) == 2:  # 2D image requested
-        destination = create([max_pos[1] + 1, max_pos[0] + 1])
+        destination = create([max_pos[1] + 1, max_pos[0] + 1], device=device)
     else:
         raise Exception("Size not supported: " + str(max_pos))
     return destination
 
-def create_square_matrix_from_pointlist(pointlist1:OCLArray):
+def create_square_matrix_from_pointlist(pointlist1:OCLArray, device:Device = None):
     width = pointlist1.shape[1] + 1
 
-    return create([width, width])
+    return create([width, width], device=device)
 
 
-def create_square_matrix_from_labelmap(labelmap: OCLArray):
+def create_square_matrix_from_labelmap(labelmap: OCLArray, device:Device = None):
     from .._tier2 import maximum_of_all_pixels
     width = int(maximum_of_all_pixels(labelmap) + 1)
 
-    return create([width, width])
+    return create([width, width], device=device)
 
 
-def create_square_matrix_from_two_labelmaps(labelmap1: OCLArray, labelmap2: OCLArray):
+def create_square_matrix_from_two_labelmaps(labelmap1: OCLArray, labelmap2: OCLArray, device:Device = None):
     from .._tier2 import maximum_of_all_pixels
     width = int(maximum_of_all_pixels(labelmap1) + 1)
     height = int(maximum_of_all_pixels(labelmap2) + 1)
 
-    return create([height, width])
+    return create([height, width],device=device)
 
 
-def create_vector_from_square_matrix(square_matrix : OCLArray, *args):
-    return create([1, square_matrix.shape[0]])
+def create_vector_from_square_matrix(square_matrix : OCLArray, *args, device:Device = None):
+    return create([1, square_matrix.shape[0]], device=device)
 
 
-def create_2d_xy(input):
+def create_2d_xy(input, device:Device = None):
     if len(input.shape) == 3:
-        return create([input.shape[2], input.shape[1]])
+        return create([input.shape[2], input.shape[1]], device=device)
     else:
-        return create([input.shape[1], input.shape[0]])
+        return create([input.shape[1], input.shape[0]], device=device)
 
-def create_2d_yx(input):
+def create_2d_yx(input, device:Device = None):
     if len(input.shape) == 3:
-        return create([input.shape[1], input.shape[2]])
+        return create([input.shape[1], input.shape[2]], device=device)
     else:
-        return create([input.shape[0], 1])
+        return create([input.shape[0], 1], device=device)
 
-def create_2d_zy(input):
+def create_2d_zy(input, device:Device = None):
     if len(input.shape) == 3:
-        return create([input.shape[0], input.shape[1]])
+        return create([input.shape[0], input.shape[1]], device=device)
     else:
-        return create([1, input.shape[0]])
+        return create([1, input.shape[0]], device=device)
 
-def create_2d_yz(input):
+def create_2d_yz(input, device:Device = None):
     if len(input.shape) == 3:
-        return create([input.shape[1], input.shape[0]])
+        return create([input.shape[1], input.shape[0]], device=device)
     else:
-        return create([input.shape[0], 1])
+        return create([input.shape[0], 1], device=device)
 
-def create_2d_zx(input):
+def create_2d_zx(input, device:Device = None):
     if len(input.shape) == 3:
-        return create([input.shape[0], input.shape[2]])
+        return create([input.shape[0], input.shape[2]], device=device)
     else:
-        return create([1, input.shape[1]])
+        return create([1, input.shape[1]], device=device)
 
-def create_2d_xz(input):
+def create_2d_xz(input, device:Device = None):
     if len(input.shape) == 3:
-        return create([input.shape[2], input.shape[0]])
+        return create([input.shape[2], input.shape[0]], device=device)
     else:
-        return create([input.shape[1], 1])
+        return create([input.shape[1], 1], device=device)
 
-def create_none(*args):
+def create_none(*args, device:Device = None):
     return None

pyclesperanto_prototype/_tier0/_device.py

@@ -1,6 +1,9 @@
 import pyopencl as cl
+from pyopencl import array
 from typing import Callable, List, Optional
 from functools import lru_cache
+from ._utils import prepare
+import numpy as np
 
 # TODO: we should discuss whether this collection is actually the best thing to pass
 # around. might be better to work lower level with contexts...
@@ -25,6 +28,11 @@ def program_from_source(self, source):
         from ._program import OCLProgram
         return OCLProgram(src_str=source, dev=self)
 
+    def from_array(self, arr, *args, **kwargs):
+        return array.to_device(self.queue, prepare(arr), *args, **kwargs)
+
+    def empty(self, shape, dtype=np.float32):
+        return array.empty(self.queue, shape, dtype)
 
 def score_device(dev: cl.Device) -> float:
     score = 4e12 if dev.type == cl.device_type.GPU else 2e12
@@ -110,4 +118,3 @@ def set_device_scoring_key(func: Callable[[cl.Device], int]) -> None:
     except Exception as e:
         raise ValueError(f"Scoring algorithm invalid: {e}")
     _current_device.score_key = func
-

pyclesperanto_prototype/_tier0/_execute.py

@@ -98,7 +98,6 @@ def execute(anchor, opencl_kernel_filename, kernel_name, global_size, parameters
     # time_stamp = time.time()
     defines = ["#define MAX_ARRAY_SIZE 1000"]
 
-
     if image_size_independent_kernel_compilation:
         defines.extend([
             "#define GET_IMAGE_WIDTH(image_key) image_size_ ## image_key ## _width",

pyclesperanto_prototype/_tier0/_plugin_function.py

@@ -7,6 +7,7 @@
 from ._create import create_like
 from ._types import Image, is_image
 from ._push import push
+from ._device import get_device
 
 
 @curry
@@ -48,6 +49,8 @@ def plugin_function(
     function.categories = categories
     function.priority = priority
 
+
+
     @wraps(function)
     def worker_function(*args, **kwargs):
         sig = inspect.signature(function)
@@ -59,13 +62,20 @@ def worker_function(*args, **kwargs):
         # https://docs.python.org/3/library/inspect.html#inspect.BoundArguments.apply_defaults
         bound.apply_defaults()
 
+        # determine on which GPU the operation should be executed and
+        # potentially, output images should be created on
+        if 'device' in kwargs.keys():
+            device = kwargs['device']
+        else:
+            device = get_device()
+
         # copy images to GPU, and create output array if necessary
         for key, value in bound.arguments.items():
             if is_image(value):
-                bound.arguments[key] = push(value)
+                bound.arguments[key] = push(value, device=device)
             if key in sig.parameters and sig.parameters[key].annotation is Image and value is None:
-                sig = inspect.signature(output_creator)
-                bound.arguments[key] = output_creator(*bound.args[:len(sig.parameters)])
+                sig = inspect.signature(output_creator)            # -1 because we add device by hand
+                bound.arguments[key] = output_creator(*bound.args[:len(sig.parameters) - 1], device=device)
 
         # call the decorated function
         return function(*bound.args, **bound.kwargs)

pyclesperanto_prototype/_tier0/_push.py

@@ -1,8 +1,8 @@
 import numpy as np
 from ._pycl import OCLArray
+from ._device import Device, get_device
 
-
-def push(any_array):
+def push(any_array, device : Device = None):
     """Copies an image to GPU memory and returns its handle
 
     .. deprecated:: 0.6.0
@@ -36,7 +36,9 @@ def push(any_array):
         any_array = np.asarray(any_array.get())
 
     float_arr = any_array.astype(np.float32)
-    return OCLArray.from_array(float_arr)
+    if device is None:
+        device = get_device()
+    return device.from_array(float_arr)
 
 
 def push_zyx(any_array):

pyclesperanto_prototype/_tier0/_pycl.py

@@ -6,6 +6,7 @@
 from pyopencl import characterize
 from pyopencl import array
 from ._device import get_device
+from ._utils import prepare
 
 """ Below here, vendored from GPUtools
 Copyright (c) 2016, Martin Weigert
@@ -105,8 +106,6 @@ def _wrap_OCLArray(cls):
     WRAPPER
     """
 
-    def prepare(arr):
-        return np.require(arr, None, "C")
 
     @classmethod
     def from_array(cls, arr, *args, **kwargs):

pyclesperanto_prototype/_tier0/_utils.py

@@ -0,0 +1,3 @@
+import numpy as np
+def prepare(arr):
+    return np.require(arr, None, "C")

pyclesperanto_prototype/_tier1/_copy.py

@@ -1,9 +1,9 @@
 from .._tier0 import execute
 from .._tier0 import plugin_function
-from .._tier0 import Image
+from .._tier0 import Image, Device
 
 @plugin_function
-def copy(source : Image, destination : Image = None):
+def copy(source : Image, destination : Image = None, device: Device = None):
     """Copies an image.
 
     <pre>f(x) = x</pre> 
@@ -12,6 +12,8 @@ def copy(source : Image, destination : Image = None):
     ----------
     source : Image
     destination : Image
+    device : Device, optional
+        OpenCL-device to operate on
 
     Returns
     -------
@@ -33,5 +35,5 @@ def copy(source : Image, destination : Image = None):
         "src":source
     }
 
-    execute(__file__, '../clij-opencl-kernels/kernels/copy_' + str(len(destination.shape)) + 'd_x.cl', 'copy_' + str(len(destination.shape)) + 'd', destination.shape, parameters)
+    execute(__file__, '../clij-opencl-kernels/kernels/copy_' + str(len(destination.shape)) + 'd_x.cl', 'copy_' + str(len(destination.shape)) + 'd', destination.shape, parameters, device=device)
     return destination