Vectorisation sprint

pyop2/codegen/c/inverse.c

@@ -6,8 +6,10 @@
 #define BUF_SIZE 30
 static PetscBLASInt ipiv_buffer[BUF_SIZE];
 static PetscScalar work_buffer[BUF_SIZE*BUF_SIZE];
+static PetscScalar Aout_proxy_buffer[BUF_SIZE*BUF_SIZE];
 #endif
 
+
 #ifndef PYOP2_INV_LOG_EVENTS
 #define PYOP2_INV_LOG_EVENTS
 PetscLogEvent ID_inv_memcpy = -1;
@@ -16,32 +18,58 @@ PetscLogEvent ID_inv_getri = -1;
 static PetscBool log_active_inv = 0;
 #endif
 
-void inverse(PetscScalar* __restrict__ Aout, const PetscScalar* __restrict__ A, PetscBLASInt N)
+static void inverse(PetscScalar* __restrict__ Aout, const PetscScalar* __restrict__ A, PetscBLASInt N,
+                    PetscBLASInt incA, PetscBLASInt incAout)
 {
     PetscLogIsActive(&log_active_inv);
-    if (log_active_inv){PetscLogEventBegin(ID_inv_memcpy,0,0,0,0);}
     PetscBLASInt info;
     PetscBLASInt *ipiv = N <= BUF_SIZE ? ipiv_buffer : malloc(N*sizeof(*ipiv));
     PetscScalar *Awork = N <= BUF_SIZE ? work_buffer : malloc(N*N*sizeof(*Awork));
-    memcpy(Aout, A, N*N*sizeof(PetscScalar));
+
+    PetscInt N_sq = N * N;
+    PetscInt one = 1;
+
+    // Aout_proxy: 'Aout', but stored contiguously
+    PetscScalar *Aout_proxy;
+    if (incAout == 1)
+      Aout_proxy = Aout;
+    else
+    {
+      // TODO: Must see if allocating has a significant performance impact
+      Aout_proxy = N_sq <= BUF_SIZE ? Aout_proxy_buffer : malloc(N*N*sizeof(*Aout));
+    }
+
+    if (log_active_inv){PetscLogEventBegin(ID_inv_memcpy,0,0,0,0);}
+    BLAScopy_(&N_sq, A, &incA, Aout_proxy, &one);
     if (log_active_inv){PetscLogEventEnd(ID_inv_memcpy,0,0,0,0);}
 
     if (log_active_inv){PetscLogEventBegin(ID_inv_getrf,0,0,0,0);}
-    LAPACKgetrf_(&N, &N, Aout, &N, ipiv, &info);
+    LAPACKgetrf_(&N, &N, Aout_proxy, &N, ipiv, &info);
     if (log_active_inv){PetscLogEventEnd(ID_inv_getrf,0,0,0,0);}
 
     if(info == 0){
         if (log_active_inv){PetscLogEventBegin(ID_inv_getri,0,0,0,0);}
-        LAPACKgetri_(&N, Aout, &N, ipiv, Awork, &N, &info);
+        LAPACKgetri_(&N, Aout_proxy, &N, ipiv, Awork, &N, &info);
         if (log_active_inv){PetscLogEventEnd(ID_inv_getri,0,0,0,0);}
+
+        // Copy Aout_proxy back to Aout
+        if (Aout != Aout_proxy)
+        {
+          if (log_active_inv){PetscLogEventBegin(ID_inv_memcpy,0,0,0,0);}
+          BLAScopy_(&N_sq, Aout_proxy, &one, Aout, &incAout);
+          if (log_active_inv){PetscLogEventEnd(ID_inv_memcpy,0,0,0,0);}
+        }
     }
 
     if(info != 0){
         fprintf(stderr, "Getri throws nonzero info.");
         abort();
     }
-    if ( N > BUF_SIZE ) {
+
+    if (Awork != work_buffer)
         free(Awork);
+    if (ipiv != ipiv_buffer)
         free(ipiv);
-    }
+    if ((Aout_proxy != Aout) && (Aout_proxy != Aout_proxy_buffer))
+        free(Aout_proxy);
 }

pyop2/codegen/c/solve.c

@@ -8,6 +8,8 @@ static PetscBLASInt ipiv_buffer[BUF_SIZE];
 static PetscScalar work_buffer[BUF_SIZE*BUF_SIZE];
 #endif
 
+static PetscScalar out_proxy_buffer[BUF_SIZE];
+
 #ifndef PYOP2_SOLVE_LOG_EVENTS
 #define PYOP2_SOLVE_LOG_EVENTS
 PetscLogEvent ID_solve_memcpy = -1;
@@ -16,15 +18,32 @@ PetscLogEvent ID_solve_getrs = -1;
 static PetscBool log_active_solve = 0;
 #endif
 
-void solve(PetscScalar* __restrict__ out, const PetscScalar* __restrict__ A, const PetscScalar* __restrict__ B, PetscBLASInt N)
+
+/*
+ * @param[incA]: Stride value while accessing elements of 'A'.
+ * @param[incB]: Stride value while accessing elements of 'B'.
+ * @param[incOut]: Stride value while accessing elements of 'out'.
+ */
+void solve(PetscScalar* __restrict__ out, const PetscScalar* __restrict__ A, const PetscScalar* __restrict__ B, PetscBLASInt N,
+           PetscBLASInt incA, PetscBLASInt incB, PetscBLASInt incOut)
 {
+    PetscScalar* out_proxy;  /// output laid-out with unit stride, expected by LAPACK
+    PetscInt N_sq = N*N;
+    PetscInt one = 1;
     PetscLogIsActive(&log_active_solve);
     if (log_active_solve){PetscLogEventBegin(ID_solve_memcpy,0,0,0,0);}
     PetscBLASInt info;
     PetscBLASInt *ipiv = N <= BUF_SIZE ? ipiv_buffer : malloc(N*sizeof(*ipiv));
-    memcpy(out,B,N*sizeof(PetscScalar));
-    PetscScalar *Awork = N <= BUF_SIZE ? work_buffer : malloc(N*N*sizeof(*Awork));
-    memcpy(Awork,A,N*N*sizeof(PetscScalar));
+
+    if (incOut == 1)
+      out_proxy = out;
+    else
+      out_proxy = (N <= BUF_SIZE) ? out_proxy_buffer : malloc(N*sizeof(*out));
+
+    BLAScopy_(&N, B, &incB, out_proxy, &one);
+
+    PetscScalar *Awork = N <= BUF_SIZE ? work_buffer : malloc(N_sq*sizeof(*Awork));
+    BLAScopy_(&N_sq, A, &incA, Awork, &one);
     if (log_active_solve){PetscLogEventEnd(ID_solve_memcpy,0,0,0,0);}
 
     PetscBLASInt NRHS = 1;
@@ -35,7 +54,11 @@ void solve(PetscScalar* __restrict__ out, const PetscScalar* __restrict__ A, con
 
     if(info == 0){
         if (log_active_solve){PetscLogEventBegin(ID_solve_getrs,0,0,0,0);}
-        LAPACKgetrs_(&T, &N, &NRHS, Awork, &N, ipiv, out, &N, &info);
+        LAPACKgetrs_(&T, &N, &NRHS, Awork, &N, ipiv, out_proxy, &N, &info);
+
+        if (out != out_proxy)
+            BLAScopy_(&N, out_proxy, &one, out, &incOut);
+
         if (log_active_solve){PetscLogEventEnd(ID_solve_getrs,0,0,0,0);}
     }
 
@@ -44,8 +67,12 @@ void solve(PetscScalar* __restrict__ out, const PetscScalar* __restrict__ A, con
         abort();
     }
 
-    if ( N > BUF_SIZE ) {
-        free(ipiv);
-        free(Awork);
-    }
+    if (ipiv != ipiv_buffer)
+      free(ipiv);
+
+    if (Awork != work_buffer)
+      free(Awork);
+
+    if ((out_proxy != out) && (out_proxy != out_proxy_buffer))
+      free(out_proxy);
 }

pyop2/codegen/rep2loopy.py

@@ -143,6 +143,7 @@ def with_types(self, arg_id_to_dtype, callables_table):
                 callables_table)
 
     def emit_call_insn(self, insn, target, expression_to_code_mapper):
+        from loopy.codegen import UnvectorizableError
         assert self.is_ready_for_codegen()
         assert isinstance(insn, loopy.CallInstruction)
 
@@ -151,6 +152,9 @@ def emit_call_insn(self, insn, target, expression_to_code_mapper):
         parameters = list(parameters)
         par_dtypes = [self.arg_id_to_dtype[i] for i, _ in enumerate(parameters)]
 
+        if expression_to_code_mapper.codegen_state.vectorization_info:
+            raise UnvectorizableError("LACallable: cannot take in vector arrays")
+
         parameters.append(insn.assignees[-1])
         par_dtypes.append(self.arg_id_to_dtype[0])
 
@@ -177,6 +181,46 @@ class INVCallable(LACallable):
     """
     name = "inverse"
 
+    def with_descrs(self, arg_id_to_descr, callables_table):
+        a_descr = arg_id_to_descr.get(0)
+        a_inv_descr = arg_id_to_descr.get(-1)
+
+        if a_descr is None or a_inv_descr is None:
+            # shapes aren't specialized enough to be resolved
+            return self, callables_table
+
+        assert len(a_descr.shape) == 2
+        assert a_descr.shape == a_inv_descr.shape
+        assert a_descr.shape[1] == a_descr.shape[0]
+
+        return self.copy(arg_id_to_descr=arg_id_to_descr), callables_table
+
+    def emit_call_insn(self, insn, target, expression_to_code_mapper):
+        from loopy.codegen import UnvectorizableError
+
+        # Override codegen to emit stride info. to the blas calls.
+        in_descr = self.arg_id_to_descr[0]
+        out_descr = self.arg_id_to_descr[-1]
+        ecm = expression_to_code_mapper
+
+        # see pyop2/codegen/c/inverse.c for the func. signature
+        inc_a = in_descr.dim_tags[1].stride
+        inc_a_out = out_descr.dim_tags[1].stride
+        n = in_descr.shape[0]
+
+        a, = insn.expression.parameters
+        a_out, = insn.assignees
+
+        if ecm.codegen_state.vectorization_info is not None:
+            raise UnvectorizableError("cannot vectorize 'inverse'.")
+
+        c_parameters = [ecm(a_out).expr,
+                        ecm(a).expr,
+                        n,
+                        inc_a,
+                        inc_a_out]
+        return var(self.name_in_target)(*c_parameters), False
+
     def generate_preambles(self, target):
         assert isinstance(target, type(target))
         yield ("inverse", inverse_preamble)
@@ -189,19 +233,65 @@ class SolveCallable(LACallable):
     """
     name = "solve"
 
+    def with_descrs(self, arg_id_to_descr, callables_table):
+        a_descr = arg_id_to_descr.get(0)
+        b_descr = arg_id_to_descr.get(1)
+        x_descr = arg_id_to_descr.get(-1)
+
+        if a_descr is None or b_descr is None:
+            # shapes aren't specialized enough to be resolved
+            return self, callables_table
+
+        assert len(a_descr.shape) == 2
+        assert len(x_descr.shape) == 1
+        assert b_descr.shape == x_descr.shape
+
+        return self.copy(arg_id_to_descr=arg_id_to_descr), callables_table
+
+    def emit_call_insn(self, insn, target, expression_to_code_mapper):
+        from loopy.codegen import UnvectorizableError
+
+        # Override codegen to emit stride info. to the blas calls.
+        a_descr = self.arg_id_to_descr[0]
+        b_descr = self.arg_id_to_descr[1]
+        out_descr = self.arg_id_to_descr[-1]
+        ecm = expression_to_code_mapper
+
+        # see pyop2/codegen/c/solve.c for the func. signature
+        inc_a = a_descr.dim_tags[1].stride
+        inc_b = b_descr.dim_tags[0].stride
+        inc_out = out_descr.dim_tags[0].stride
+        n = a_descr.shape[0]
+
+        a, b = insn.expression.parameters
+        out, = insn.assignees
+
+        if ecm.codegen_state.vectorization_info is not None:
+            raise UnvectorizableError("cannot vectorize 'inverse'.")
+
+        c_parameters = [ecm(out).expr,
+                        ecm(a).expr,
+                        ecm(b).expr,
+                        n,
+                        inc_a,
+                        inc_b,
+                        inc_out]
+        return var(self.name_in_target)(*c_parameters), False
+
     def generate_preambles(self, target):
         assert isinstance(target, type(target))
         yield ("solve", solve_preamble)
 
 
 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):
@@ -537,7 +627,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

@@ -139,6 +139,8 @@ def sniff_compiler(exe):
                 compiler = MacClangARMCompiler
             elif machine == "x86_64":
                 compiler = MacClangCompiler
+        elif name == "GNU":
+            compiler = MacGNUCompiler
         else:
             compiler = AnonymousCompiler
     else:
@@ -456,21 +458,26 @@ class MacClangCompiler(Compiler):
     _cxxflags = ("-fPIC", "-Wall", "-framework", "Accelerate")
     _ldflags = ("-dynamiclib",)
 
-    _optflags = ("-O3", "-ffast-math", "-march=native")
+    _optflags = ("-O3", "-ffast-math", "-march=native", "-fopenmp-simd")
     _debugflags = ("-O0", "-g")
 
 
 class MacClangARMCompiler(MacClangCompiler):
     """A compiler for building a shared library on ARM based Mac systems."""
     # See https://stackoverflow.com/q/65966969
-    _optflags = ("-O3", "-ffast-math", "-mcpu=apple-a14")
+    _optflags = ("-O3", "-ffast-math", "-mcpu=apple-a14", "-fopenmp-simd")
     # Need to pass -L/opt/homebrew/opt/gcc/lib/gcc/11 to prevent linker error:
     # ld: file not found: @rpath/libgcc_s.1.1.dylib for architecture arm64 This
     # seems to be a homebrew configuration issue somewhere. Hopefully this
     # requirement will go away at some point.
     _ldflags = ("-dynamiclib", "-L/opt/homebrew/opt/gcc/lib/gcc/11")
 
 
+class MacGNUCompiler(MacClangCompiler):
+    """A compiler for building a shared library on Mac systems with a GNU compiler."""
+    _name = "Mac GNU"
+
+
 class LinuxGnuCompiler(Compiler):
     """The GNU compiler for building a shared library on Linux systems."""
     _name = "GNU"
@@ -551,7 +558,7 @@ class LinuxIntelCompiler(Compiler):
     _cxxflags = ("-fPIC", "-no-multibyte-chars")
     _ldflags = ("-shared",)
 
-    _optflags = ("-Ofast", "-xHost")
+    _optflags = ("-Ofast", "-xHost", "-qopenmp-simd")
     _debugflags = ("-O0", "-g")
 
 

pyop2/configuration.py

@@ -40,6 +40,24 @@
 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_simd_extension = [t for t in get_cpu_info()["flags"] if t.startswith('avx')][-1]
+    if longest_simd_extension in avx_to_width.keys():
+        return avx_to_width[longest_simd_extension]
+    elif longest_simd_extension[:6] in avx_to_width.keys():
+        return avx_to_width[longest_simd_extension[:6]]
+    elif longest_simd_extension[:4] in avx_to_width.keys():
+        return avx_to_width[longest_simd_extension[:4]]
+    else:
+        raise ConfigurationError(f"The vector extension of your architecture is unknown.\
+                                   Must be one of {str(avx_to_width.keys())}.\
+                                   We advise to disable vectorisation."
+                                 )
+
+
 class Configuration(dict):
     r"""PyOP2 configuration parameters
 
@@ -74,6 +92,14 @@ 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 vectorization_strategy: A :class:`str` describing the
+        vectorization strategy that must to be applied to the kernels. Can
+        be one of the following --
+    :param alignment: A :class:`int` which specifies a size to which all temporaries
+        are aligned in memory.
+
+        - ``cross-element``: Cross-element vectorization strategy of
+          `<https://doi.org/10.1177/1094342020945005>`__.
     """
     # name, env variable, type, default, write once
     cache_dir = os.path.join(gettempdir(), "pyop2-cache-uid%s" % os.getuid())
@@ -91,7 +117,11 @@ class Configuration(dict):
         "ldflags":
             ("PYOP2_LDFLAGS", str, ""),
         "simd_width":
-            ("PYOP2_SIMD_WIDTH", int, 4),
+            ("PYOP2_SIMD_WIDTH", int, default_simd_width()),
+        "vectorization_strategy":
+            ("PYOP2_VECT_STRATEGY", str, "cross-element"),
+        "alignment":
+            ("PYOP2_ALIGNMENT", int, 64),
         "debug":
             ("PYOP2_DEBUG", bool, False),
         "compute_kernel_flops":