Cache cl.Context on OpenCL

Author: szymonlopaciuk

xobjects/context.py

@@ -668,14 +668,15 @@ def get_context_from_string(ctxstr):
 
     if ctxstr is None:
         return xo.ContextCpu()
+
+    ll = ctxstr.split(":")
+    if len(ll) <= 1:
+        ctxtype = ll[0]
+        option = []
     else:
-        ll = ctxstr.split(":")
-        if len(ll) <= 1:
-            ctxtype = ll[0]
-            option = []
-        else:
-            ctxtype, options = ctxstr.split(":")
-            option = options.split(",")
+        ctxtype, options = ctxstr.split(":")
+        option = options.split(",")
+
     if ctxtype == "ContextCpu":
         if len(option) == 0:
             return xo.ContextCpu()

xobjects/context_pyopencl.py

@@ -81,6 +81,8 @@ def _build_view(cls, a):
 
 
 class ContextPyopencl(XContext):
+    context_cache = {}
+
     @property
     def nplike_array_type(self):
         return cla.Array
@@ -128,20 +130,29 @@ def __init__(
         super().__init__()
 
         # TODO assume one device only
-        if device is None:
+        if device in self.context_cache:
+            self.platform, self.device, self.context = self.context_cache[
+                device
+            ]
+        elif device is None:
             self.context = cl.create_some_context(interactive=False)
             self.device = self.context.devices[0]
             self.platform = self.device.platform
         else:
             if isinstance(device, str):
-                platform, device = map(int, device.split("."))
+                platform, device_ = map(int, device.split("."))
                 self.platform = cl.get_platforms()[platform]
-                self.device = self.platform.get_devices()[device]
+                self.device = self.platform.get_devices()[device_]
             else:
                 self.device = device
                 self.platform = device.platform
 
             self.context = cl.Context([self.device])
+            self.context_cache[device] = (
+                self.platform,
+                self.device,
+                self.context,
+            )
 
         self.queue = cl.CommandQueue(self.context)
 

xobjects/general.py

@@ -18,9 +18,8 @@ def __call__(self, *args, **kwargs):
 
 
 def assert_allclose(a, b, rtol=0, atol=0, max_outliers=0):
-
     try:
-        if a == b: # if this passes we return without raising
+        if a == b:  # if this passes we return without raising
             return
     except:
         pass