GPU

pyop3/codegen/ir.py

@@ -22,6 +22,7 @@
 from pyop3 import utils
 from pyop3.axes import Axis, AxisComponent, AxisTree, AxisVariable
 from pyop3.axes.tree import ContextSensitiveAxisTree
+from pyop3.device import CPUDevice, CUDADevice, OpenCLDevice, offloading_device
 from pyop3.distarray import DistributedArray, MultiArray
 from pyop3.distarray.multiarray import ContextSensitiveMultiArray
 from pyop3.distarray.petsc import IndexedPetscMat, PetscMat, PetscObject
@@ -66,10 +67,20 @@
     strictly_all,
 )
 
-# FIXME this needs to be synchronised with TSFC, tricky
-# shared base package? or both set by Firedrake - better solution
-LOOPY_TARGET = lp.CWithGNULibcTarget()
-LOOPY_LANG_VERSION = (2018, 2)
+
+def loopy_target():
+    if isinstance(offloading_device, CPUDevice):
+        return lp.CWithGNULibcTarget()
+    elif isinstance(offloading_device, CUDADevice):
+        return lp.CudaTarget()
+    elif isinstance(offloading_device, OpenCLDevice):
+        return lp.PyOpenCLTarget()
+    else:
+        raise AssertionError
+
+
+def loopy_lang_version():
+    return (2018, 2)
 
 
 class OpaqueType(lp.types.OpaqueType):
@@ -320,8 +331,8 @@ def compile(expr: LoopExpr, name="mykernel"):
         ctx.instructions,
         ctx.arguments,
         name=name,
-        target=LOOPY_TARGET,
-        lang_version=LOOPY_LANG_VERSION,
+        target=loopy_target(),
+        lang_version=loopy_lang_version(),
         preambles=preambles,
         # options=lp.Options(check_dep_resolution=False),
     )