Vectorisation sprint

pyop2/codegen/rep2loopy.py

@@ -191,13 +191,14 @@ def generate_preambles(self, target):
 
 
 class _PreambleGen(ImmutableRecord):
-    fields = set(("preamble", ))
+    fields = {"preamble", "idx"}
 
-    def __init__(self, preamble):
+    def __init__(self, preamble, idx="0"):
         self.preamble = preamble
+        self.idx = idx
 
     def __call__(self, preamble_info):
-        yield ("0", self.preamble)
+        yield (self.idx, self.preamble)
 
 
 class PyOP2KernelCallable(loopy.ScalarCallable):
@@ -533,7 +534,9 @@ def renamer(expr):
                                 options=options,
                                 assumptions=assumptions,
                                 lang_version=(2018, 2),
-                                name=wrapper_name)
+                                name=wrapper_name,
+                                # TODO, should these really be silenced?
+                                silenced_warnings=["write_race*", "data_dep*"])
 
     # prioritize loops
     for indices in context.index_ordering:

pyop2/compilation.py

@@ -219,7 +219,6 @@ def workaround_cflags(self):
             if version.StrictVersion("7.3") <= ver <= version.StrictVersion("7.5"):
                 # GCC bug https://gcc.gnu.org/bugzilla/show_bug.cgi?id=90055
                 # See also https://github.com/firedrakeproject/firedrake/issues/1442
-                # And https://github.com/firedrakeproject/firedrake/issues/1717
                 # Bug also on skylake with the vectoriser in this
                 # combination (disappears without
                 # -fno-tree-loop-vectorize!)
@@ -370,7 +369,7 @@ class MacCompiler(Compiler):
 
     def __init__(self, cppargs=[], ldargs=[], cpp=False, comm=None):
         machine = platform.uname().machine
-        opt_flags = ["-O3", "-ffast-math"]
+        opt_flags = ["-O3", "-ffast-math", "-fopenmp-simd"]
         if machine == "arm64":
             # See https://stackoverflow.com/q/65966969
             opt_flags.append("-mcpu=apple-a14")
@@ -405,7 +404,7 @@ class LinuxCompiler(Compiler):
     :kwarg comm: Optional communicator to compile the code on (only
     rank 0 compiles code) (defaults to COMM_WORLD)."""
     def __init__(self, cppargs=[], ldargs=[], cpp=False, comm=None):
-        opt_flags = ['-march=native', '-O3', '-ffast-math']
+        opt_flags = ['-march=native', '-O3', '-ffast-math', '-fopenmp-simd']
         if configuration['debug']:
             opt_flags = ['-O0', '-g']
         cc = "mpicc"
@@ -431,7 +430,7 @@ class LinuxIntelCompiler(Compiler):
         rank 0 compiles code) (defaults to COMM_WORLD).
     """
     def __init__(self, cppargs=[], ldargs=[], cpp=False, comm=None):
-        opt_flags = ['-Ofast', '-xHost']
+        opt_flags = ['-march=native', '-Ofast', '-xHost', '-qopenmp-simd']
         if configuration['debug']:
             opt_flags = ['-O0', '-g']
         cc = "mpicc"

pyop2/configuration.py

@@ -39,6 +39,22 @@
 from pyop2.exceptions import ConfigurationError
 
 
+def default_simd_width():
+    from cpuinfo import get_cpu_info
+    avx_to_width = {'avx': 2, 'avx1': 2, 'avx128': 2, 'avx2': 4,
+                    'avx256': 4, 'avx3': 8, 'avx512': 8}
+    longest_ext = [t for t in get_cpu_info()["flags"] if t.startswith('avx')][-1]
+    if longest_ext not in avx_to_width.keys():
+        if longest_ext[:6] not in avx_to_width.keys():
+            assert longest_ext[:4] in avx_to_width.keys(), \
+                "The vector extension of your architecture is unknown. Disable vectorisation!"
+            return avx_to_width[longest_ext[:4]]
+        else:
+            return avx_to_width[longest_ext[:6]]
+    else:
+        return avx_to_width[longest_ext]
+
+
 class Configuration(dict):
     r"""PyOP2 configuration parameters
 
@@ -78,7 +94,10 @@ class Configuration(dict):
     # name, env variable, type, default, write once
     DEFAULTS = {
         "compiler": ("PYOP2_BACKEND_COMPILER", str, "gcc"),
-        "simd_width": ("PYOP2_SIMD_WIDTH", int, 4),
+        "simd_width": ("PYOP2_SIMD_WIDTH", int, default_simd_width()),
+        "vectorization_strategy": ("PYOP2_VECT_STRATEGY", str, "ve"),
+        "alignment": ("PYOP2_ALIGNMENT", int, 64),
+        "time": ("PYOP2_TIME", bool, False),
         "debug": ("PYOP2_DEBUG", bool, False),
         "cflags": ("PYOP2_CFLAGS", str, ""),
         "ldflags": ("PYOP2_LDFLAGS", str, ""),

pyop2/global_kernel.py

@@ -15,6 +15,8 @@
 from pyop2.datatypes import IntType, as_ctypes
 from pyop2.types import IterationRegion
 from pyop2.utils import cached_property, get_petsc_dir
+from pyop2 import configuration
+from pyop2 import op2
 
 
 # We set eq=False to force identity-based hashing. This is required for when
@@ -331,14 +333,68 @@ def code_to_compile(self):
         from pyop2.codegen.rep2loopy import generate
 
         wrapper = generate(self.builder)
+        if self._iterset._extruded:
+            iname = "layer"
+        else:
+            iname = "n"
+
+        has_matrix = any(arg._is_mat for arg in self._args)
+        has_rw = any(arg.access == op2.RW for arg in self._args)
+        is_cplx = any(arg.dtype.name == 'complex128' for arg in self._args)
+        vectorisable = not (has_matrix or has_rw) and (configuration["vectorization_strategy"])
+
+        if (isinstance(self._kernel.code, lp.LoopKernel) and vectorisable):
+            wrapper = lp.inline_callable_kernel(wrapper, self._kernel.name)
+            if not is_cplx:
+                wrapper = self.vectorise(wrapper, iname, configuration["simd_width"])
         code = lp.generate_code_v2(wrapper)
 
         if self.local_kernel.cpp:
-            from loopy.codegen.result import process_preambles
+            from lp.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()
+
+    def vectorise(wrapper, iname, batch_size):
+        """Return a vectorised version of wrapper, vectorising over iname.
+
+        :arg wrapper: A loopy kernel to vectorise.
+        :arg iname: The iteration index to vectorise over.
+        :arg batch_size: The vector width."""
+        if batch_size == 1:
+            return wrapper
+
+        # create constant zero vectors
+        wrapper = wrapper.copy(target=lp.CVecTarget(batch_size))
+        kernel = wrapper.root_kernel
+
+        # split iname and vectorize the inner loop
+        slabs = (1, 1)
+        inner_iname = iname + "_batch"
+
+        if configuration["vectorization_strategy"] == "ve":
+            kernel = lp.split_iname(kernel, iname, batch_size, slabs=slabs, inner_tag="vec", inner_iname=inner_iname)
+
+        alignment = configuration["alignment"]
+        tmps = dict((name, tv.copy(alignment=alignment)) for name, tv in kernel.temporary_variables.items())
+        kernel = kernel.copy(temporary_variables=tmps)
+
+        from lp.preprocess import check_cvec_vectorizability, cvec_retag_and_privatize, realize_ilp
+        from lp.kernel.data import OpenMPSIMDTag, VectorizeTag
+        from lp.transform.iname import tag_inames
+
+        kernel = realize_ilp(kernel)  # FIXME: do we also need to realize the reductions first?
+
+        # try to vectorise with vector extensionn
+        vector_inst, pragma_inst_to_tag, unr_inst_to_tag = check_cvec_vectorizability(kernel)
+
+        # if not possible fall back to OpenMP SIMD pragmas or unrolling by retagging, then privatize
+        kernel = cvec_retag_and_privatize(kernel, vector_inst, pragma_inst_to_tag, unr_inst_to_tag)
+
+        wrapper = wrapper.with_root_kernel(kernel)
+
+        return wrapper
 
     @PETSc.Log.EventDecorator()
     @mpi.collective

requirements-ext.txt

@@ -4,6 +4,7 @@ pytest>=2.3
 flake8>=2.1.0
 pycparser>=2.10
 mpi4py>=1.3.1
+py-cpuinfo
 decorator<=4.4.2
 dataclasses
 cachetools

requirements-git.txt

@@ -1,2 +1,2 @@
 git+https://github.com/coneoproject/COFFEE.git#egg=coffee
-git+https://github.com/firedrakeproject/loopy.git@main#egg=loopy
+git+https://github.com/firedrakeproject/loopy.git@vectorisation-sprint#egg=loopy