WIP: Merge branch 'master' into gpu

Author: kaushikcfd

.github/workflows/ci.yml

@@ -15,40 +15,82 @@ jobs:
     runs-on: ubuntu-latest
     env:
       CC: mpicc
-      PETSC_CONFIGURE_OPTIONS: --with-debugging=1
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v2
+      PETSC_DIR: ${{ github.workspace }}/petsc
+      PETSC_ARCH: default
+      PETSC_CONFIGURE_OPTIONS: --with-debugging=1 --with-shared-libraries=1 --with-c2html=0 --with-fortran-bindings=0
 
+    steps:
       - name: Install system dependencies
         shell: bash
         run: |
-          sudo apt install  \
-            build-essential python3-dev git python3-pip libopenmpi-dev  \
-            openmpi-bin libblas-dev liblapack-dev gfortran
+          sudo apt update
+          sudo apt install build-essential mpich libmpich-dev \
+              libblas-dev liblapack-dev gfortran
 
       - name: Set correct Python version
         uses: actions/setup-python@v2
         with:
           python-version: '3.6'
 
-      # xargs is used to force installation of requirements in the order we specified.
-      # pip 20.2 needed for loopy install to work.
-      - name: Install Python dependencies
+      - name: Clone PETSc
+        uses: actions/checkout@v2
+        with:
+          repository: firedrakeproject/petsc
+          path: ${{ env.PETSC_DIR }}
+
+      - name: Build and install PETSc
         shell: bash
+        working-directory: ${{ env.PETSC_DIR }}
         run: |
-          pip install pip==20.2
-          xargs -l1 pip install < requirements-ext.txt
-          xargs -l1 pip install < requirements-git.txt
-          pip install pulp
-          pip install -U flake8
+          ./configure ${PETSC_CONFIGURE_OPTIONS}
+          make
+
+      - name: Build and install petsc4py
+        shell: bash
+        working-directory: ${{ env.PETSC_DIR }}/src/binding/petsc4py
+        run: |
+          python -m pip install --upgrade cython numpy
+          python -m pip install --no-deps .
+
+      - name: Checkout PyOP2
+        uses: actions/checkout@v2
+        with:
+          path: PyOP2
 
       - name: Install PyOP2
         shell: bash
-        run: pip install -e .
+        working-directory: PyOP2
+        run: |
+          python -m pip install pip==20.2  # pip 20.2 needed for loopy install to work.
 
-      - name: Do tests
+          # xargs is used to force installation of requirements in the order we specified.
+          xargs -l1 python -m pip install < requirements-ext.txt
+          xargs -l1 python -m pip install < requirements-git.txt
+          python -m pip install pulp
+          python -m pip install -U flake8
+          python -m pip install .
+
+      - name: Run linting
         shell: bash
+        working-directory: PyOP2
+        run: make lint
+
+      - name: Run tests
+        shell: bash
+        working-directory: PyOP2
+        run: pytest test -v --tb=native
+
+      - name: Build documentation
+        shell: bash
+        working-directory: PyOP2
         run: |
-          make lint
-          py.test test -v --tb=native
+          python -m pip install sphinx
+          make -C doc/sphinx html
+      - name: Upload to github pages
+        if: ${{ github.ref == 'refs/heads/master' && github.event_name == 'push' }}
+        uses: crazy-max/[email protected]
+        with:
+          build_dir: PyOP2/doc/sphinx/build/html
+          jekyll: false
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/gh-pages.yml

@@ -6,11 +6,9 @@ PyOP2.egg-info
 *.py[cdo]
 
 # Extension modules
-computeind.c
-computeind.so
-sparsity.cpp
 sparsity.so
-
+sparsity.c
+sparsity.cpython*.so
 # Docs
 pyop2.coffee.rst
 pyop2.rst

.gitignore

@@ -7,25 +7,26 @@ class AbstractComputeBackend:
     Abstract class to record all the backend specific implementation of
     :mod:`pyop2`'s data structures.
     """
-    ParLoop = _not_implemented()
-    Set = _not_implemented()
-    ExtrudedSet = _not_implemented()
-    MixedSet = _not_implemented()
-    Subset = _not_implemented()
-    DataSet = _not_implemented()
-    MixedDataSet = _not_implemented()
-    Map = _not_implemented()
-    MixedMap = _not_implemented()
-    Dat = _not_implemented()
-    MixedDat = _not_implemented()
-    DatView = _not_implemented()
-    Mat = _not_implemented()
-    Global = _not_implemented()
-    GlobalDataSet = _not_implemented()
+    ParLoop = _not_implemented
+    Set = _not_implemented
+    ExtrudedSet = _not_implemented
+    MixedSet = _not_implemented
+    Subset = _not_implemented
+    DataSet = _not_implemented
+    MixedDataSet = _not_implemented
+    Map = _not_implemented
+    MixedMap = _not_implemented
+    Dat = _not_implemented
+    MixedDat = _not_implemented
+    DatView = _not_implemented
+    Mat = _not_implemented
+    Global = _not_implemented
+    GlobalDataSet = _not_implemented
+    PETScVecType = _not_implemented
 
     def _getattr_(self, key):
         val = super(AbstractComputeBackend, self)._getattr_(key)
-        if isinstance(val, _not_implemented):
+        if val is _not_implemented:
             raise NotImplementedError("'{}' is not implemented for backend"
                                       " '{}'.".format(val, self.__name__))
         return val

pyop2/backend.py renamed to pyop2/backends/__init__.py

@@ -0,0 +1,169 @@
+from pyop2.types.dat import Dat as BaseDat, MixedDat, DatView
+from pyop2.types.set import Set, ExtrudedSet, Subset, MixedSet
+from pyop2.types.dataset import DataSet, GlobalDataSet, MixedDataSet
+from pyop2.types.map import Map, MixedMap
+from pyop2.parloop import AbstractParLoop, AbstractJITModule
+from pyop2.types.mat import Mat
+from pyop2.glob import Global
+from pyop2.backends import AbstractComputeBackend
+from petsc4py import PETSc
+from . import (
+    compilation,
+    configuration as conf,
+    datatypes as dtypes,
+    mpi,
+    utils
+)
+
+import ctypes
+import os
+import loopy as lp
+
+
+class Dat(BaseDat):
+    @utils.cached_property
+    def _vec(self):
+        assert self.dtype == PETSc.ScalarType, \
+            "Can't create Vec with type %s, must be %s" % (self.dtype, PETSc.ScalarType)
+        # Can't duplicate layout_vec of dataset, because we then
+        # carry around extra unnecessary data.
+        # But use getSizes to save an Allreduce in computing the
+        # global size.
+        size = self.dataset.layout_vec.getSizes()
+        data = self._data[:size[0]]
+        vec = PETSc.Vec().createWithArray(data, size=size, bsize=self.cdim, comm=self.comm)
+        return vec
+
+
+class JITModule(AbstractJITModule):
+    @utils.cached_property
+    def code_to_compile(self):
+        from pyop2.codegen.builder import WrapperBuilder
+        from pyop2.codegen.rep2loopy import generate
+
+        builder = WrapperBuilder(kernel=self._kernel,
+                                 iterset=self._iterset,
+                                 iteration_region=self._iteration_region,
+                                 pass_layer_to_kernel=self._pass_layer_arg)
+        for arg in self._args:
+            builder.add_argument(arg)
+
+        wrapper = generate(builder)
+        code = lp.generate_code_v2(wrapper)
+
+        if self._kernel._cpp:
+            from loopy.codegen.result import process_preambles
+            preamble = "".join(process_preambles(getattr(code, "device_preambles", [])))
+            device_code = "\n\n".join(str(dp.ast) for dp in code.device_programs)
+            return preamble + "\nextern \"C\" {\n" + device_code + "\n}\n"
+        return code.device_code()
+
+    @PETSc.Log.EventDecorator()
+    @mpi.collective
+    def compile(self):
+        # If we weren't in the cache we /must/ have arguments
+        if not hasattr(self, '_args'):
+            raise RuntimeError("JITModule has no args associated with it, should never happen")
+
+        compiler = conf.configuration["compiler"]
+        extension = "cpp" if self._kernel._cpp else "c"
+        cppargs = self._cppargs
+        cppargs += ["-I%s/include" % d for d in utils.get_petsc_dir()] + \
+                   ["-I%s" % d for d in self._kernel._include_dirs] + \
+                   ["-I%s" % os.path.abspath(os.path.dirname(__file__))]
+        ldargs = ["-L%s/lib" % d for d in utils.get_petsc_dir()] + \
+                 ["-Wl,-rpath,%s/lib" % d for d in utils.get_petsc_dir()] + \
+                 ["-lpetsc", "-lm"] + self._libraries
+        ldargs += self._kernel._ldargs
+
+        self._fun = compilation.load(self,
+                                     extension,
+                                     self._wrapper_name,
+                                     cppargs=cppargs,
+                                     ldargs=ldargs,
+                                     restype=ctypes.c_int,
+                                     compiler=compiler,
+                                     comm=self.comm)
+        # Blow away everything we don't need any more
+        del self._args
+        del self._kernel
+        del self._iterset
+
+    @utils.cached_property
+    def argtypes(self):
+        index_type = dtypes.as_ctypes(dtypes.IntType)
+        argtypes = (index_type, index_type)
+        argtypes += self._iterset._argtypes_
+        for arg in self._args:
+            argtypes += arg._argtypes_
+        seen = set()
+        for arg in self._args:
+            maps = arg.map_tuple
+            for map_ in maps:
+                for k, t in zip(map_._kernel_args_, map_._argtypes_):
+                    if k in seen:
+                        continue
+                    argtypes += (t,)
+                    seen.add(k)
+        return argtypes
+    ...
+
+
+class ParLoop(AbstractParLoop):
+
+    def prepare_arglist(self, iterset, *args):
+        arglist = iterset._kernel_args_
+        for arg in args:
+            arglist += arg._kernel_args_
+        seen = set()
+        for arg in args:
+            maps = arg.map_tuple
+            for map_ in maps:
+                if map_ is None:
+                    continue
+                for k in map_._kernel_args_:
+                    if k in seen:
+                        continue
+                    arglist += (k,)
+                    seen.add(k)
+        return arglist
+
+    @utils.cached_property
+    def _jitmodule(self):
+        return JITModule(self.kernel, self.iterset, *self.args,
+                         iterate=self.iteration_region,
+                         pass_layer_arg=self._pass_layer_arg)
+
+    @mpi.collective
+    def _compute(self, part, fun, *arglist):
+        with self._compute_event:
+            self.log_flops(part.size * self.num_flops)
+            fun(part.offset, part.offset + part.size, *arglist)
+
+
+class CPUBackend(AbstractComputeBackend):
+    ParLoop = ParLoop
+    Set = Set
+    ExtrudedSet = ExtrudedSet
+    MixedSet = MixedSet
+    Subset = Subset
+    DataSet = DataSet
+    MixedDataSet = MixedDataSet
+    Map = Map
+    MixedMap = MixedMap
+    Dat = Dat
+    MixedDat = MixedDat
+    DatView = DatView
+    Mat = Mat
+    Global = Global
+    GlobalDataSet = GlobalDataSet
+    PETScVecType = 'standard'
+
+    def turn_on_offloading(self):
+        pass
+
+    def turn_off_offloading(self):
+        pass
+
+
+cpu_backend = CPUBackend()