bugfix: switching GPUs

closes #110

Author: haesleinhuepf

pyclesperanto_prototype/_tier0/_device.py

@@ -1,5 +1,6 @@
 import pyopencl as cl
 from typing import Callable, List, Optional
+from functools import lru_cache
 
 # TODO: we should discuss whether this collection is actually the best thing to pass
 # around. might be better to work lower level with contexts...
@@ -19,6 +20,11 @@ def __repr__(self) -> str:
     def name(self) -> str:
         return self.device.name
 
+    @lru_cache(maxsize=128)
+    def program_from_source(self, source):
+        from ._program import OCLProgram
+        return OCLProgram(src_str=source, dev=self)
+
 
 def score_device(dev: cl.Device) -> float:
     score = 4e12 if dev.type == cl.device_type.GPU else 2e12
@@ -57,6 +63,11 @@ def select_device(name: str = None, dev_type: str = None, score_key=None) -> Dev
     _current_device._instance = Device(device, context, queue)
     return _current_device._instance
 
+def new_device(name: str = None, dev_type: str = None, score_key=None) -> Device:
+    device = filter_devices(name, dev_type, score_key)[-1]
+    context = cl.Context(devices=[device])
+    queue = cl.CommandQueue(context)
+    return Device(device, context, queue)
 
 def filter_devices(
     name: str = None, dev_type: str = None, score_key=None
@@ -99,3 +110,4 @@ 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

@@ -5,7 +5,9 @@
 import numpy as np
 
 import pyopencl as cl
-from ._pycl import OCLProgram, _OCLImage
+from ._pycl import _OCLImage
+from ._device import Device
+from ._program import OCLProgram
 
 if not os.getenv("PYOPENCL_COMPILER_OUTPUT"):
     import warnings
@@ -68,7 +70,7 @@ def get_ocl_source(anchor, opencl_kernel_filename):
 #define WRITE_{key}_IMAGE(a,b,c) write_image{typeId}(a,b,c)
 """
 
-def execute(anchor, opencl_kernel_filename, kernel_name, global_size, parameters, prog : OCLProgram = None, constants = None, image_size_independent_kernel_compilation : bool = True):
+def execute(anchor, opencl_kernel_filename, kernel_name, global_size, parameters, prog : OCLProgram = None, constants = None, image_size_independent_kernel_compilation : bool = True, device: Device = None):
     """
     Convenience method for calling opencl kernel files
 
@@ -274,8 +276,12 @@ def execute(anchor, opencl_kernel_filename, kernel_name, global_size, parameters
         # time_stamp = time.time()
 
         defines.append(get_ocl_source(anchor, opencl_kernel_filename))
-        
-        prog = OCLProgram.from_source("\n".join(defines))
+
+        if device is None:
+            from ._device import get_device
+            device = get_device()
+        prog = device.program_from_source("\n".join(defines))
+        #prog = OCLProgram.from_source("\n".join(defines))
 
         # Todo: the order of the arguments matters; fix that
         # print("Compilation " + opencl_kernel_filename + " took " + str((time.time() - time_stamp) * 1000) + " ms")

pyclesperanto_prototype/_tier0/_program.py

@@ -0,0 +1,44 @@
+
+import pyopencl as cl
+import warnings
+from functools import lru_cache
+
+class OCLProgram(cl.Program):
+    """ a wrapper class representing a CPU/GPU Program
+    example:
+         prog = OCLProgram("mykernels.cl",build_options=["-D FLAG"])
+    """
+    _wait_for_kernel_finish = None
+
+    def __init__(self, file_name=None, src_str=None, build_options=[], dev=None):
+        if file_name is not None:
+            with open(file_name, "r") as f:
+                src_str = f.read()
+
+        if src_str is None:
+            raise ValueError("empty src_str! ")
+
+        if dev is None:
+            from ._device import get_device
+            dev = get_device()
+
+        self._dev = dev
+        self._kernel_dict = {}
+        super().__init__(self._dev.context, src_str)
+        self.build(options=build_options)
+
+    def run_kernel(self, name, global_size, local_size, *args, **kwargs):
+        if name not in self._kernel_dict:
+            self._kernel_dict[name] = getattr(self, name)
+
+        self._kernel_dict[name](
+            self._dev.queue, global_size, local_size, *args, **kwargs
+        )
+        if OCLProgram._wait_for_kernel_finish:
+            self._dev.queue.finish()
+
+    @classmethod
+    @lru_cache(maxsize=128)
+    def from_source(cls, source):
+        warnings.warn("OCLProgram.from_source is deprecated. Use Device.program_from_source instead.")
+        return cls(src_str=source)

pyclesperanto_prototype/_tier0/_pycl.py

@@ -1,10 +1,10 @@
 import os
 import sys
+
 import numpy as np
 import pyopencl as cl
 from pyopencl import characterize
 from pyopencl import array
-from functools import lru_cache
 from ._device import get_device
 
 """ Below here, vendored from GPUtools
@@ -38,43 +38,7 @@
 """
 
 
-class OCLProgram(cl.Program):
-    """ a wrapper class representing a CPU/GPU Program
-    example:
-         prog = OCLProgram("mykernels.cl",build_options=["-D FLAG"])
-    """
-    _wait_for_kernel_finish = None
-
-    def __init__(self, file_name=None, src_str=None, build_options=[], dev=None):
-        if file_name is not None:
-            with open(file_name, "r") as f:
-                src_str = f.read()
-
-        if src_str is None:
-            raise ValueError("empty src_str! ")
-
-        if dev is None:
-            dev = get_device()
 
-        self._dev = dev
-        self._kernel_dict = {}
-        super().__init__(self._dev.context, src_str)
-        self.build(options=build_options)
-
-    def run_kernel(self, name, global_size, local_size, *args, **kwargs):
-        if name not in self._kernel_dict:
-            self._kernel_dict[name] = getattr(self, name)
-
-        self._kernel_dict[name](
-            self._dev.queue, global_size, local_size, *args, **kwargs
-        )
-        if OCLProgram._wait_for_kernel_finish:
-            self._dev.queue.finish()
-
-    @classmethod
-    @lru_cache(maxsize=128)
-    def from_source(cls, source):
-        return cls(src_str=source)
 
 
 cl_image_datatype_dict = {

tests/test_gpu_switch.py

@@ -0,0 +1,19 @@
+import numpy as np
+import pyclesperanto_prototype as cle
+
+# this test may only work on a machine with multiple GPUs
+def test_gpu_switch():
+    image = np.random.random((100,100))
+
+    # select NVidia
+    print(cle.select_device("RTX"))
+    cle.gaussian_blur(image, sigma_x=1, sigma_y=1)
+
+    # select AMD
+    print(cle.select_device("gfx"))
+    cle.gaussian_blur(image, sigma_x=1, sigma_y=1)
+
+    # select Intel
+    print(cle.select_device("Intel"))
+    cle.gaussian_blur(image, sigma_x=1, sigma_y=1)
+