[WARNING]: Should be avlbl anywhere but local branches; explicit transfers work

Author: kaushikcfd

pyop2/configuration.py

@@ -71,6 +71,10 @@ class Configuration(dict):
         cdim > 1 be built as block sparsities, or dof sparsities.  The
         former saves memory but changes which preconditioners are
         available for the resulting matrices.  (Default yes)
+    :param only_explicit_host_device_data_transfers: Flag to set host<->device
+        transfers mode. If set *True*, the user has to invoke all the
+        host<->device transfers. If set *False* (default), Firedrake automatically
+        figures out the data transfers, however this might lead to sub-optimality.
     """
     # name, env variable, type, default, write once
     DEFAULTS = {
@@ -112,6 +116,7 @@ class Configuration(dict):
         "print_summary": ("PYOP2_PRINT_SUMMARY", bool, False),
         "matnest": ("PYOP2_MATNEST", bool, True),
         "block_sparsity": ("PYOP2_BLOCK_SPARSITY", bool, True),
+        "only_explicit_host_device_data_transfers": ("EXPLICIT_TRNSFRS", bool, False)
     }
     """Default values for PyOP2 configuration parameters"""
 

pyop2/gpu/cuda.py

@@ -119,19 +119,6 @@ class Dat(petsc_Dat):
     """
     Dat for GPU.
     """
-    @validate_type(('dataset', (base.DataCarrier, DataSet, Set), DataSetTypeError),
-                   ('name', str, NameTypeError))
-    @validate_dtype(('dtype', None, DataTypeError))
-    def __init__(self, dataset, data=None, dtype=None, name=None, uid=None):
-
-        if isinstance(dataset, petsc_Dat) and not isinstance(dataset, Dat):
-            self.__init__(dataset.dataset, None, dtype=dataset.dtype,
-                          name="copy_of_%s" % dataset.name)
-            self._data[...] = dataset.data
-            return
-
-        super(Dat, self).__init__(dataset, data, dtype, name, uid)
-
     @cached_property
     def _vec(self):
         assert self.dtype == PETSc.ScalarType, \
@@ -149,47 +136,6 @@ def _vec(self):
 
         return cuda_vec
 
-    @cached_property
-    def device_handle(self):
-        if self.dtype == PETSc.ScalarType:
-            with self.vec as v:
-                return v.getCUDAHandle()
-        elif self.dtype == PETSc.IntType:
-            m_gpu = cuda.mem_alloc(int(self._data.nbytes))
-            cuda.memcpy_htod(m_gpu, self._data)
-            return m_gpu
-        else:
-            raise NotImplementedError("Unknown type: %s." % self.dtype)
-
-    @cached_property
-    def _kernel_args_(self):
-        return (self.device_handle, )
-
-    @collective
-    @property
-    def data(self):
-
-        with self.vec as v:
-            v.restoreCUDAHandle(self.device_handle)
-            return v.array
-
-    ## TODO: fail when trying to acess elems from data_ro
-    @collective
-    @property
-    def data_ro(self):
-        with self.vec_ro as v:
-            v.restoreCUDAHandle(self.device_handle)
-            return v.array
-
-
-    def move_to_host(self):
-        with self.vec_ro as v:
-            v.restoreCUDAHandle(self.device_handle)
-            self._data = v.array
-
-        return petsc_Dat(self._dataset, self._data, self.dtype, 'copy_of_%s' %
-                self.name)
-
 
 class Global(petsc_Global):
 
@@ -226,6 +172,7 @@ def __init__(self, kernel, iterset, *args, **kwargs):
            otherwise they (and the :class:`~.Dat`\s, :class:`~.Map`\s
            and :class:`~.Mat`\s they reference) will never be collected.
         """
+
         # Return early if we were in the cache.
         if self._initialized:
             return
@@ -326,6 +273,12 @@ def ith_added_global_arg_i(self, i):
 
     @collective
     def __call__(self, *args):
+        #FIXME: Should prolly get rid of this once the implementation is
+        # finalized.
+        from pyop2.op2 import device
+        import pyop2.gpu.cuda
+        assert device.target == pyop2.gpu.cuda
+
         if self._initialized:
             grid, block = self.grid_size(args[0], args[1])
             extra_global_args = self.get_args_marked_for_globals
@@ -417,13 +370,11 @@ def argtypes(self):
         argtypes = (index_type, index_type)
         argtypes += self._iterset._argtypes_
         for arg in self._args:
-            assert isinstance(arg.data, Dat)
             argtypes += arg._argtypes_
         seen = set()
         for arg in self._args:
             maps = arg.map_tuple
             for map_ in maps:
-                assert isinstance(map_, Map)
                 for k, t in zip(map_._kernel_args_, map_._argtypes_):
                     if k in seen:
                         continue

pyop2/op2.py

@@ -135,6 +135,9 @@ def __getattr__(self, attr):
     def set(self, target):
         self.target = target
 
+    def copy(self):
+        return Target(self.target)
+
     def __repr__(self):
         return "Target(%r)" % self.target
 

pyop2/petsc_base.py

@@ -45,6 +45,7 @@
 from pyop2.base import _make_object, Subset
 from pyop2.mpi import collective
 from pyop2.profiling import timed_region
+from pyop2.configuration import configuration
 
 
 class DataSet(base.DataSet):
@@ -333,6 +334,18 @@ def vec_ro(self):
 
 
 class Dat(base.Dat, VecAccessMixin):
+
+    def __init__(self, *args, **kwargs):
+        super(Dat, self).__init__(*args, **kwargs)
+
+        import pyop2.sequential
+        import pyop2.gpu.cuda
+
+        self._PETSC_VEC_TYPE_TO_TARGET_ = {
+                'seq': pyop2.sequential,
+                'seqcuda': pyop2.gpu.cuda,
+                }
+
     @utils.cached_property
     def _vec(self):
         print(75*'=')
@@ -363,17 +376,65 @@ def vec_context(self, access):
             self.halo_valid = False
 
     def move_to_device(self, device):
-        """
-        Returns an instance of :class:`pyop2.base.Dat`, which is located on
-        device.
-        """
-        from pyop2.op2 import Target
+        from pyop2 import gpu
+        from pyop2.op2 import Target, host
         if isinstance(device, Target):
             device = device.target
-        return device.Dat(self)
+        with self.vec as petsc_vec:
+
+            if self._PETSC_VEC_TYPE_TO_TARGET_[petsc_vec.type] == device:
+                return
+
+            print("We want {} to {}.".format(petsc_vec.type, device))
+
+            if device == host.target:
+                self.move_to_host()
+            elif device == gpu.cuda:
+                host_data = petsc_vec.array.copy()
+                petsc_vec.setType('seqcuda')
+                petsc_vec.setArray(host_data)
+            elif device == gpu.opencl:
+                raise NotImplementedError("OpenCL target not yet implemented")
+            else:
+                raise NotImplementedError("Unknown target '%s'." % device)
 
     def move_to_host(self):
-        raise ValueError("Cannot transfer Dat to host which is already on host.")
+        with self.vec as petsc_vec:
+            size = self.dataset.layout_vec.getSizes()
+            #FIXME: This is probably more involved
+            self._data = self.data.copy()
+            petsc_vec.setType('seq')
+            petsc_vec.setArray(self._data[:size[0]])
+
+    @property
+    def _kernel_args_(self):
+        #FIXME: This should be cached_property, but because of the in-place
+        # updates this is no longer possible; improve the caching mechanism.
+        from pyop2.op2 import device
+        with self.vec as petsc_vec:
+            if self._PETSC_VEC_TYPE_TO_TARGET_[petsc_vec.type] != device.target:
+                if configuration['only_explicit_host_device_data_transfers']:
+                    raise RuntimeError("Memory location mismatch.")
+                else:
+                    self.move_to_device(device)
+            if petsc_vec.type == 'seq':
+                return (self._data.ctypes.data, )
+            elif petsc_vec.type == 'seqcuda':
+                return (petsc_vec.getCUDAHandle(), )
+            else:
+                raise NotImplementedError()
+
+    @collective
+    @property
+    def data(self):
+        with self.vec as v:
+            if v.type == 'seq':
+                return v.array
+            elif v.type == 'seqcuda':
+                v.restoreCUDAHandle(v.getCUDAHandle())
+                return v.array
+            else:
+                raise NotImplementedError("Unknown vec type %s." % v.type)
 
 
 class MixedDat(base.MixedDat, VecAccessMixin):

pyop2/sequential.py

@@ -84,7 +84,6 @@ def __init__(self, kernel, iterset, *args, **kwargs):
            and :class:`~.Mat`\s they reference) will never be collected.
         """
         # Return early if we were in the cache.
-
         if self._initialized:
             return
         self.comm = iterset.comm
@@ -104,6 +103,10 @@ def __init__(self, kernel, iterset, *args, **kwargs):
 
     @collective
     def __call__(self, *args):
+        #FIXME: Should prolly get rid of this once the implementation is
+        # finalized.
+        from pyop2.op2 import device, host
+        assert host == device
         return self._fun(*args)
 
     @cached_property
@@ -112,6 +115,8 @@ def _wrapper_name(self):
 
     @cached_property
     def code_to_compile(self):
+        if self._wrapper_name == 'wrap_copy':
+            import pudb; pu.db
         print(75*"=")
         print("pyop2.sequential:115: Compiling %s on Host." % self._wrapper_name)
         print(75*"=")