feectools 0.5.0__tar.gz → 0.7.0__tar.gz

{feectools-0.5.0/feectools.egg-info → feectools-0.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.5.0
+Version: 0.7.0
 Summary: Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies.
 Author-email: Psydac development team <psydac@googlegroups.com>
 Maintainer-email: Stefan Possanner <stefan.possanner@ipp.mpg.de>, Max Lindqvist <max.lindqvist@ipp.mpg.de>, Yaman Güçlü <yaman.guclu@gmail.com>, Martin Campos Pinto <martin.campos-pinto@ipp.mpg.de>, Ahmed Ratnani <ratnaniahmed@gmail.com>

{feectools-0.5.0 → feectools-0.7.0}/feectools/accelerate/compile_psydac.mk

@@ -37,7 +37,7 @@ all: $(OUTPUTS)
 	@for dep in $^ ; do \
 		echo $$dep ; \
     done
-	pyccel compile -v $(FLAGS)$(FLAGS_openmp) $<
+	pyccel compile $(FLAGS)$(FLAGS_openmp) $<
 	@echo ""
 
 #--------------------------------------

{feectools-0.5.0 → feectools-0.7.0}/feectools/api/fem_bilinear_form.py

@@ -594,9 +594,9 @@ class DiscreteBilinearForm:
                     if is_conformal:
                         matrix[k1, k2] = global_mats[k1, k2]
                     elif use_restriction:
-                        matrix.multiplicants[-1][k1, k2] = global_mats[k1, k2]
+                        matrix.multiplicands[-1][k1, k2] = global_mats[k1, k2]
                     elif use_prolongation:
-                        matrix.multiplicants[0][k1, k2] = global_mats[k1, k2]
+                        matrix.multiplicands[0][k1, k2] = global_mats[k1, k2]
 
         else: # case of scalar equation
             if is_broken: # multi-patch

{feectools-0.5.0 → feectools-0.7.0}/feectools/api/fem_common.py

@@ -277,7 +277,7 @@ def extract_stencil_mats(mats):
         if isinstance(M, (StencilInterfaceMatrix, StencilMatrix)):
             new_mats.append(M)
         elif isinstance(M, ComposedLinearOperator):
-            new_mats += [i for i in M.multiplicants if isinstance(i, (StencilInterfaceMatrix, StencilMatrix))]
+            new_mats += [i for i in M.multiplicands if isinstance(i, (StencilInterfaceMatrix, StencilMatrix))]
     return new_mats
 
 #==============================================================================

{feectools-0.5.0 → feectools-0.7.0}/feectools/feec/derivatives.py

@@ -136,6 +136,21 @@ class DirectionalDerivativeOperator(LinearOperator):
     def dtype( self ):
         return self.domain.dtype
 
+    @property
+    def diffdir(self) -> int:
+        """Direction (axis) of the derivative."""
+        return self._diffdir
+
+    @property
+    def negative(self) -> bool:
+        """Whether the operator is the negative derivative."""
+        return self._negative
+
+    @property
+    def transposed(self) -> bool:
+        """Whether the operator is the transposed derivative."""
+        return self._transposed
+
     def __truediv__(self, a):
         """ Divide by scalar. """
         return self * (1.0 / a)

{feectools-0.5.0 → feectools-0.7.0}/feectools/linalg/basic.py

@@ -8,6 +8,7 @@ provides the fundamental classes for linear algebra operations.
 """
 
 import itertools
+import warnings
 from abc import ABC, abstractmethod
 from types import LambdaType 
 from inspect import signature
@@ -93,9 +94,11 @@ class VectorSpace(ABC):
 
         Returns
         -------
-        float | complex
+        float | complex | cupy.ndarray
             The scalar product of the two vectors. Note that inner(x, x) is
             a non-negative real number which is zero if and only if x = 0.
+            For vectors with device (CuPy) data, a 0-d device array: the
+            result stays on the device.
 
         """
 
@@ -1055,27 +1058,27 @@ class ComposedLinearOperator(LinearOperator):
         for i in range(len(args)-1):
             assert args[i].domain == args[i+1].codomain
 
-        multiplicants = ()
+        multiplicands = ()
         tmp_vectors = []
         for a in args[:-1]:
             if isinstance(a, ComposedLinearOperator):
-                multiplicants = (*multiplicants, *a.multiplicants)
+                multiplicands = (*multiplicands, *a.multiplicands)
                 tmp_vectors.extend(a.tmp_vectors)
                 tmp_vectors.append(a.domain.zeros())
             else:
-                multiplicants = (*multiplicants, a)
+                multiplicands = (*multiplicands, a)
                 tmp_vectors.append(a.domain.zeros())
 
         last = args[-1]
         if isinstance(last, ComposedLinearOperator):
-            multiplicants = (*multiplicants, *last.multiplicants)
+            multiplicands = (*multiplicands, *last.multiplicands)
             tmp_vectors.extend(last.tmp_vectors)
         else:
-            multiplicants = (*multiplicants, last)
+            multiplicands = (*multiplicands, last)
 
         self._domain = domain
         self._codomain = codomain
-        self._multiplicants = multiplicants
+        self._multiplicands = multiplicands
         self._tmp_vectors = tuple(tmp_vectors)
 
     @property
@@ -1096,34 +1099,41 @@ class ComposedLinearOperator(LinearOperator):
         return self._codomain
 
     @property
-    def multiplicants(self):
+    def multiplicands(self):
         r"""
-        A tuple $(A_1,\dots,A_n)$ containing the multiplicants of the linear operator 
+        A tuple $(A_1,\dots,A_n)$ containing the multiplicands of the linear operator 
         $self = A_n\circ\dots\circ A_1$.
         
         """
-        return self._multiplicants
+        return self._multiplicands
+
+    @property
+    def multiplicants(self):
+        """Deprecated alias of ``multiplicands``."""
+        warnings.warn("ComposedLinearOperator.multiplicants is deprecated, use multiplicands instead.",
+                      DeprecationWarning, stacklevel=2)
+        return self._multiplicands
 
     @property
     def dtype(self):
         return None
 
     def tosparse(self):
-        mats = [M.tosparse() for M in self._multiplicants]
+        mats = [M.tosparse() for M in self._multiplicands]
         M = mats[0]
         for Mi in mats[1:]:
             M = M @ Mi
         return coo_matrix(M)
 
     def transpose(self, conjugate=False):
-        t_multiplicants = ()
-        for a in self._multiplicants:
-            t_multiplicants = (a.transpose(conjugate=conjugate), *t_multiplicants)
+        t_multiplicands = ()
+        for a in self._multiplicands:
+            t_multiplicands = (a.transpose(conjugate=conjugate), *t_multiplicands)
         new_dom = self.codomain
         new_cod = self.domain
         assert isinstance(new_dom, VectorSpace)
         assert isinstance(new_cod, VectorSpace)
-        return ComposedLinearOperator(self.codomain, self.domain, *t_multiplicants)
+        return ComposedLinearOperator(self.codomain, self.domain, *t_multiplicands)
 
     def dot(self, v, out=None):
         assert isinstance(v, Vector)
@@ -1135,11 +1145,11 @@ class ComposedLinearOperator(LinearOperator):
         x = v
         for i in range(len(self._tmp_vectors)):
             y = self._tmp_vectors[-1-i]
-            A = self._multiplicants[-1-i]
+            A = self._multiplicands[-1-i]
             A.dot(x, out=y)
             x = y
 
-        A = self._multiplicants[0]
+        A = self._multiplicands[0]
         if out is not None:
 
             A.dot(x, out=out)
@@ -1148,11 +1158,11 @@ class ComposedLinearOperator(LinearOperator):
         return out
 
     def exchange_assembly_data(self):
-        for op in self._multiplicants:
+        for op in self._multiplicands:
             op.exchange_assembly_data()
 
     def set_backend(self, backend, precompiled=False):
-        for op in self._multiplicants:
+        for op in self._multiplicands:
             op.set_backend(backend)
 
 #===============================================================================

{feectools-0.5.0 → feectools-0.7.0}/feectools/linalg/block.py

@@ -112,9 +112,11 @@ class BlockVectorSpace(VectorSpace):
 
         Returns
         -------
-        float | complex
+        float | complex | cupy.ndarray
             The scalar product of the two vectors. Note that inner(x, x) is
             a non-negative real number which is zero if and only if x = 0.
+            For vectors with device (CuPy) data, a 0-d device array: the
+            result stays on the device.
 
         """
 
@@ -134,7 +136,8 @@ class BlockVectorSpace(VectorSpace):
         Parameters
         ----------
         a : scalar
-            The scaling coefficient needed for the operation.
+            The scaling coefficient needed for the operation (a 0-d device
+            array is accepted, see `StencilVectorSpace.axpy`).
 
         x : BlockVector
             The vector which is not modified by this function.

{feectools-0.5.0 → feectools-0.7.0}/feectools/linalg/direct_solvers.py

@@ -4,13 +4,14 @@
 # for full license details.                                                 #
 #---------------------------------------------------------------------------#
 from abc                 import abstractmethod
+import numpy as np
 import cunumpy as xp
 from cunumpy.xp import array_backend
 from scipy.sparse        import spmatrix, dia_matrix
 
 from feectools.linalg.basic    import LinearSolver
 
-__all__ = ('to_bnd', 'BandedSolver', 'SparseSolver')
+__all__ = ('to_bnd', 'DenseInverse', 'BandedSolver', 'SparseSolver')
 
 #===============================================================================
 def to_bnd(A):
@@ -28,6 +29,89 @@ def to_bnd(A):
 
     return A_bnd, la, ua
 
+#===============================================================================
+class DenseInverse:
+    """
+    The inverse of a 1D solver's matrix as a dense matrix, for solves on the device.
+
+    LAPACK and SuperLU run on the host only. For device data the matrix is inverted
+    once on the host, by the solver itself (it solves for the identity), and the inverse
+    is applied with one matrix product per solve (a cuBLAS GEMM on CuPy). A device solve
+    then makes no host/device copy. The 1D matrices of Kronecker solvers are small
+    (n ~ tens to hundreds), so the n x n inverse is cheap to store, and one GEMM over all
+    right-hand sides is faster on a GPU than banded triangular solves.
+
+    Right-hand sides are rows, as in `BandedSolver.solve`: for ``rhs`` of shape
+    ``(m, n)`` the solution is ``rhs @ M`` with ``M = solver.solve(I)``, whose row k is
+    the solution for the k-th unit vector, i.e. ``M = op(A)^{-T}``. This holds for any
+    linear 1D solver with the row convention, transposed or not.
+
+    Parameters
+    ----------
+    solver : LinearSolver
+        A 1D solver whose ``solve`` accepts host (NumPy) arrays of shape ``(m, n)``.
+
+    n : int
+        The size of the 1D matrix.
+
+    dtype : dtype
+        The dtype of the right-hand sides.
+    """
+    def __init__(self, solver, n, dtype):
+        eye = np.eye(int(n), dtype=dtype)
+        # host bookkeeping at setup: the host solver applied to the identity
+        self._host = np.ascontiguousarray(xp.to_numpy(solver.solve(eye)))
+        self._device = None
+
+    @property
+    def shape(self):
+        return self._host.shape
+
+    @property
+    def host_matrix(self):
+        """``M`` on the host."""
+        return self._host
+
+    def device_matrix(self):
+        """``M`` on the device, copied there once (on the first call)."""
+        if self._device is None:
+            self._device = xp.to_cupy(self._host)
+        return self._device
+
+    def matrix_for(self, array):
+        """``M`` on the device for a device array, on the host otherwise."""
+        return self.device_matrix() if xp.is_gpu(array) else self._host
+
+    def solve(self, rhs, out=None):
+        """
+        Solves for the right-hand sides ``rhs`` (rows) where they live, with one matrix product.
+
+        The right-hand sides are the rows of ``rhs``, as in `BandedSolver.solve`, so the
+        solutions are the rows of ``rhs @ op(A)^{-T}``; ``M`` stores ``op(A)^{-T}`` (no
+        symmetry is assumed). Multiplying from the right keeps the row-contiguous work
+        arrays of `KroneckerLinearSolver` as they are: one GEMM, no transpose copy.
+
+        ``out`` may be ``rhs`` (in-place solve).
+        """
+        assert rhs.shape[-1] == self._host.shape[0]
+        result = rhs @ self.matrix_for(rhs)
+        if out is None:
+            return result
+        assert out.shape == rhs.shape
+        out[...] = result
+        return out
+
+#===============================================================================
+def _device_inverse(solver, n, dtype):
+    """The `DenseInverse` of a 1D solver for device solves, built on first use and cached per dtype."""
+    cache = getattr(solver, '_dense_inverses', None)
+    if cache is None:
+        cache = solver._dense_inverses = {}
+    key = np.dtype(dtype)
+    if key not in cache:
+        cache[key] = DenseInverse(solver, n, key)
+    return cache[key]
+
 #===============================================================================
 class BandedSolver(LinearSolver):
     """
@@ -117,6 +201,7 @@ class BandedSolver(LinearSolver):
         obj._space = self._space
         obj._dtype = self._dtype
         obj._transposed = not self._transposed
+        obj._dense_inverses = {}
 
         return obj
 
@@ -140,11 +225,19 @@ class BandedSolver(LinearSolver):
 
         transposed = self._transposed
 
+        # LAPACK is host-only: device data is solved with the dense inverse on the
+        # device (see DenseInverse). Decided by the array itself, not the global
+        # backend, since host arrays may be passed on the CuPy backend too.
+        if xp.is_gpu(rhs):
+            if out is not None:
+                assert out.shape == rhs.shape
+                assert out.dtype == rhs.dtype
+            return _device_inverse(self, self._bmat.shape[1], rhs.dtype).solve(rhs, out=out)
+
         if out is None:
-            # LAPACK is host-only: solve on the host, return on the caller's backend.
-            preout, self._sinfo = self._solver_function(self._bmat, self._l, self._u, xp.to_numpy(rhs).T,
+            preout, self._sinfo = self._solver_function(self._bmat, self._l, self._u, rhs.T,
                                                         self._ipiv, trans=transposed)
-            out = xp.asarray(preout.T) if xp.is_gpu(rhs) else preout.T
+            out = preout.T
 
         else:
             assert out.shape == rhs.shape
@@ -157,16 +250,7 @@ class BandedSolver(LinearSolver):
             # TODO: handle non-contiguous views?
 
             # we want FORTRAN-contiguous data (default is assumed to be C contiguous).
-            # LAPACK is host-only: a device array is solved in a host copy. Decided by
-            # the array itself, not the global backend, since host arrays may be passed
-            # on the CuPy backend too.
-            if xp.is_gpu(out):
-                out_cpu = xp.to_numpy(out)
-                _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out_cpu.T, self._ipiv, overwrite_b=True,
-                                                   trans=transposed)
-                out[...] = xp.asarray(out_cpu)
-            else:
-                _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out.T, self._ipiv, overwrite_b=True,
+            _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out.T, self._ipiv, overwrite_b=True,
                                                    trans=transposed)
 
         return out
@@ -206,6 +290,7 @@ class SparseSolver (LinearSolver):
         obj._space = self._space
         obj._splu = self._splu
         obj._transposed = not self._transposed
+        obj._dense_inverses = {}
 
         return obj
 
@@ -229,19 +314,22 @@ class SparseSolver (LinearSolver):
         assert rhs.T.shape[0] == self._splu.shape[1]
         transposed = self._transposed
 
+        # SuperLU is host-only: device data is solved with the dense inverse on the
+        # device (see DenseInverse); decided by the array, not the global backend.
+        if xp.is_gpu(rhs):
+            if out is not None:
+                assert out.shape == rhs.shape
+                assert out.dtype == rhs.dtype
+            return _device_inverse(self, self._splu.shape[1], rhs.dtype).solve(rhs, out=out)
+
         if out is None:
-            # SuperLU is host-only: solve on the host, return on the caller's backend.
-            out = self._splu.solve(xp.to_numpy(rhs).T, trans='T' if transposed else 'N').T
-            if xp.is_gpu(rhs):
-                out = xp.asarray(out)
+            out = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T
 
         else:
             assert out.shape == rhs.shape
             assert out.dtype == rhs.dtype
 
-            # currently no in-place solve exposed. SuperLU is host-only; decided by
-            # the arrays themselves, not the global backend.
-            result = self._splu.solve(xp.to_numpy(rhs).T, trans='T' if transposed else 'N').T
-            out[:] = xp.asarray(result) if xp.is_gpu(out) else result
+            # currently no in-place solve exposed.
+            out[:] = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T
 
         return out

{feectools-0.5.0 → feectools-0.7.0}/feectools/linalg/fft.py

@@ -36,6 +36,9 @@ class DistributedFFTBase(LinearOperator):
         function : Callable
             The given function.
         """
+        # KroneckerLinearSolver: keep applying the function (not its dense matrix) to device data
+        dense_on_device = False
+
         def __init__(self, function):
             self._function = function
 

{feectools-0.5.0 → feectools-0.7.0}/feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_cuda.cu

@@ -7,15 +7,16 @@
  *
  * One thread per entry of `out` (n_threads = out.size, see __init__.py). A thread outside the owned rows
  * (local row index i1_loc outside [0, e_out - s_out]) returns without writing, so the padding of `out` is
- * left as it is, as in pyccel. Interior rows use 2 * p_in + 1 diagonals, the last owned row (i1 == e_out)
- * uses 2 * p_in + add, which is how a rectangular matrix is handled.
+ * left as it is, as in pyccel. The matrix has n = mat.shape[1] diagonals (its pads are q = (n - 1) / 2) and
+ * diagonal d1 of row i1 is the column i1 - q + d1. Interior rows use all n diagonals, the last owned row
+ * (i1 == e_out) uses n - 1 + add, which is how a rectangular matrix is handled.
  *
- * @param mat   matrix data, shape (rows of `out`, 2 * p_in + 1)
+ * @param mat   matrix data, shape (rows of `out`, diagonals)
  * @param x     data of the domain vector, ghost regions included
  * @param out   data of the codomain vector; the owned rows are written
  * @param s_in  global start of the domain of this process
  * @param p_in  padding of the domain
- * @param add   1 if the last row uses all 2 * p_in + 1 diagonals, else 0
+ * @param add   1 if the last row uses all diagonals, else 0
  * @param s_out global start of the codomain of this process
  * @param e_out global end (inclusive) of the codomain of this process
  * @param p_out padding of the codomain: the owned rows start at index p_out of `mat` and `out`
@@ -30,11 +31,13 @@ extern "C" __global__ void stencil_dot_1d(Array2D<double> mat, Array1D<double> x
     if (i1_loc < 0 || i1_loc > e_out - s_out) return;
     const long long i1 = s_out + i1_loc;  // global row index
 
-    const long long n_diags1 = (i1 == e_out) ? 2 * p_in + add : 2 * p_in + 1;
+    const long long n_diags1 = mat.shape[1];
+    const long long nd1 = (i1 == e_out) ? n_diags1 - 1 + add : n_diags1;
+    const long long off1 = p_in - (n_diags1 - 1) / 2 - s_in;  // x index of diagonal 0 minus i1
 
     double val = 0.;
-    for (long long d1 = 0; d1 < n_diags1; ++d1)
-        val += mat(p_out + i1_loc, d1) * x(i1 + d1 - s_in);
+    for (long long d1 = 0; d1 < nd1; ++d1)
+        val += mat(p_out + i1_loc, d1) * x(i1 + d1 + off1);
 
     out(p_out + i1_loc) = val;
 }

feectools-0.7.0/feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_kernels.py

@@ -0,0 +1,38 @@
+"""Stencil matrix-vector product ``out = mat @ x`` on the owned rows (1D).
+
+The host version of ``stencil_dot_1d``.
+
+Moved from ``feectools.linalg.stencil_dot_kernels.matvec_1d_kernel``. The CUDA version in ``stencil_dot_1d_cuda.cu``
+takes the same arguments in the same order.
+
+The number of diagonals is read from ``mat``: ``n = mat.shape[1]`` diagonals, the pads of the matrix are
+``q = (n - 1) // 2`` (at most the pads ``p_in`` of the domain), and diagonal ``d`` of row ``i`` is the column
+``i - q + d``. Interior rows use all ``n`` diagonals, the last owned row ``n - 1 + add``.
+"""
+
+
+def stencil_dot_1d(mat: 'float[:, :]',
+                   x: 'float[:]',
+                   out: 'float[:]',
+                   s_in: int,
+                   p_in: int,
+                   add: int,
+                   s_out: int,
+                   e_out: int,
+                   p_out: int):
+
+    n_diags1 = mat.shape[1]
+    # x index of diagonal 0 minus the global row index: column i1 - q + d1 is at x[i1 - q + d1 - s_in + p_in]
+    off1 = p_in - (n_diags1 - 1) // 2 - s_in
+
+    for i1 in range(s_out, e_out + 1):  # global row index
+        i1_loc = i1 - s_out  # local row index
+        nd1 = n_diags1
+        if i1 == e_out:
+            nd1 = n_diags1 - 1 + add
+
+        val = 0.
+        for d1 in range(nd1):
+            val += mat[p_out + i1_loc, d1] * x[i1 + d1 + off1]
+
+        out[p_out + i1_loc] = val

{feectools-0.5.0 → feectools-0.7.0}/feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_cuda.cu

@@ -7,11 +7,11 @@
  *
  * One thread per entry of `out` (n_threads = out.size, see __init__.py), the last axis varying fastest. A
  * thread outside the owned rows returns without writing, so the padding of `out` is left as it is, as in
- * pyccel. Along each direction k, interior rows use 2 * p_in[k] + 1 diagonals and the last owned row
- * (i_k == e_out[k]) uses 2 * p_in[k] + add[k]; the pyccel kernel spells out the four combinations, this
- * kernel picks its own per thread. The diagonals are summed in the same order (d1 outer, d2 inner).
+ * pyccel. Along each direction k the matrix has n_k = mat.shape[2 + k] diagonals; interior rows use all of
+ * them, the last owned row (i_k == e_out[k]) uses n_k - 1 + add[k], as in the 1D kernel. The diagonals are
+ * summed in the same order as in pyccel (d1 outer, d2 inner).
  *
- * @param mat   matrix data, shape (rows of `out`..., 2 * p_in + 1...)
+ * @param mat   matrix data, shape (rows of `out`..., diagonals...)
  * @param x     data of the domain vector, ghost regions included
  * @param out   data of the codomain vector; the owned rows are written
  * @param s_in, p_in, add, s_out, e_out, p_out  per direction (length 2), as in the 1D kernel
@@ -30,13 +30,16 @@ extern "C" __global__ void stencil_dot_2d(Array4D<double> mat, Array2D<double> x
     const long long i1 = s_out[0] + i1_loc;  // global row indices
     const long long i2 = s_out[1] + i2_loc;
 
-    const long long n_diags1 = (i1 == e_out[0]) ? 2 * p_in[0] + add[0] : 2 * p_in[0] + 1;
-    const long long n_diags2 = (i2 == e_out[1]) ? 2 * p_in[1] + add[1] : 2 * p_in[1] + 1;
+    const long long nd1 = (i1 == e_out[0]) ? mat.shape[2] - 1 + add[0] : mat.shape[2];
+    const long long nd2 = (i2 == e_out[1]) ? mat.shape[3] - 1 + add[1] : mat.shape[3];
+    // x index of diagonal 0 minus the global row index, per direction
+    const long long off1 = p_in[0] - (mat.shape[2] - 1) / 2 - s_in[0];
+    const long long off2 = p_in[1] - (mat.shape[3] - 1) / 2 - s_in[1];
 
     double val = 0.;
-    for (long long d1 = 0; d1 < n_diags1; ++d1)
-        for (long long d2 = 0; d2 < n_diags2; ++d2)
-            val += mat(p_out[0] + i1_loc, p_out[1] + i2_loc, d1, d2) * x(i1 + d1 - s_in[0], i2 + d2 - s_in[1]);
+    for (long long d1 = 0; d1 < nd1; ++d1)
+        for (long long d2 = 0; d2 < nd2; ++d2)
+            val += mat(p_out[0] + i1_loc, p_out[1] + i2_loc, d1, d2) * x(i1 + d1 + off1, i2 + d2 + off2);
 
     out(p_out[0] + i1_loc, p_out[1] + i2_loc) = val;
 }

feectools-0.7.0/feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_kernels.py

@@ -0,0 +1,47 @@
+"""Stencil matrix-vector product ``out = mat @ x`` on the owned rows (2D).
+
+The host version of ``stencil_dot_2d``.
+
+Moved from ``feectools.linalg.stencil_dot_kernels.matvec_2d_kernel``. The CUDA version in ``stencil_dot_2d_cuda.cu``
+takes the same arguments in the same order.
+
+The number of diagonals is read from ``mat``, per direction as in ``stencil_dot_1d``: ``n_k = mat.shape[2 + k]``,
+the pads of the matrix are ``q_k = (n_k - 1) // 2`` and diagonal ``d_k`` of row ``i_k`` is the column
+``i_k - q_k + d_k``. Interior rows use all ``n_k`` diagonals, the last owned row ``n_k - 1 + add[k]``.
+"""
+
+
+def stencil_dot_2d(mat: 'float[:, :, :, :]',
+                   x: 'float[:, :]',
+                   out: 'float[:, :]',
+                   s_in: 'int[:]',
+                   p_in: 'int[:]',
+                   add: 'int[:]',
+                   s_out: 'int[:]',
+                   e_out: 'int[:]',
+                   p_out: 'int[:]'):
+
+    n_diags1 = mat.shape[2]
+    n_diags2 = mat.shape[3]
+    # x index of diagonal 0 minus the global row index, per direction
+    off1 = p_in[0] - (n_diags1 - 1) // 2 - s_in[0]
+    off2 = p_in[1] - (n_diags2 - 1) // 2 - s_in[1]
+
+    for i1 in range(s_out[0], e_out[0] + 1):  # global row indices
+        i1_loc = i1 - s_out[0]  # local row indices
+        nd1 = n_diags1
+        if i1 == e_out[0]:
+            nd1 = n_diags1 - 1 + add[0]
+
+        for i2 in range(s_out[1], e_out[1] + 1):
+            i2_loc = i2 - s_out[1]
+            nd2 = n_diags2
+            if i2 == e_out[1]:
+                nd2 = n_diags2 - 1 + add[1]
+
+            val = 0.
+            for d1 in range(nd1):
+                for d2 in range(nd2):
+                    val += mat[p_out[0] + i1_loc, p_out[1] + i2_loc, d1, d2] * x[i1 + d1 + off1, i2 + d2 + off2]
+
+            out[p_out[0] + i1_loc, p_out[1] + i2_loc] = val

feectools-0.7.0/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_cuda.cu

@@ -0,0 +1,52 @@
+// CUDA version of stencil_dot_3d (stencil_dot_3d_kernels.py), with the same arguments in the same order.
+#include "cunumpy/array_view.cuh"
+#include "cunumpy/index.cuh"
+
+/**
+ * Stencil matrix-vector product out = mat @ x on the owned rows, as the pyccel kernel stencil_dot_3d.
+ *
+ * One thread per entry of `out` (n_threads = out.size, see __init__.py), the last axis varying fastest. A
+ * thread outside the owned rows returns without writing, so the padding of `out` is left as it is, as in
+ * pyccel. Along each direction k the matrix has n_k = mat.shape[3 + k] diagonals; interior rows use all of
+ * them, the last owned row (i_k == e_out[k]) uses n_k - 1 + add[k], as in the 1D kernel. The diagonals are
+ * summed in the same order as in pyccel (d1, d2, d3 innermost).
+ *
+ * @param mat   matrix data (a 6D view), shape (rows of `out`..., diagonals...)
+ * @param x     data of the domain vector, ghost regions included
+ * @param out   data of the codomain vector; the owned rows are written
+ * @param s_in, p_in, add, s_out, e_out, p_out  per direction (length 3), as in the 1D kernel
+ */
+extern "C" __global__ void stencil_dot_3d(Array6D<double> mat, Array3D<double> x, Array3D<double> out,
+                                          const long long* s_in, const long long* p_in, const long long* add,
+                                          const long long* s_out, const long long* e_out,
+                                          const long long* p_out)
+{
+    CUNUMPY_THREAD_1D(thread, out.size());
+
+    const long long i1_loc = thread / (out.shape[1] * out.shape[2]) - p_out[0];  // local row indices
+    const long long i2_loc = (thread / out.shape[2]) % out.shape[1] - p_out[1];
+    const long long i3_loc = thread % out.shape[2] - p_out[2];
+    if (i1_loc < 0 || i1_loc > e_out[0] - s_out[0]) return;
+    if (i2_loc < 0 || i2_loc > e_out[1] - s_out[1]) return;
+    if (i3_loc < 0 || i3_loc > e_out[2] - s_out[2]) return;
+    const long long i1 = s_out[0] + i1_loc;  // global row indices
+    const long long i2 = s_out[1] + i2_loc;
+    const long long i3 = s_out[2] + i3_loc;
+
+    const long long nd1 = (i1 == e_out[0]) ? mat.shape[3] - 1 + add[0] : mat.shape[3];
+    const long long nd2 = (i2 == e_out[1]) ? mat.shape[4] - 1 + add[1] : mat.shape[4];
+    const long long nd3 = (i3 == e_out[2]) ? mat.shape[5] - 1 + add[2] : mat.shape[5];
+    // x index of diagonal 0 minus the global row index, per direction
+    const long long off1 = p_in[0] - (mat.shape[3] - 1) / 2 - s_in[0];
+    const long long off2 = p_in[1] - (mat.shape[4] - 1) / 2 - s_in[1];
+    const long long off3 = p_in[2] - (mat.shape[5] - 1) / 2 - s_in[2];
+
+    double val = 0.;
+    for (long long d1 = 0; d1 < nd1; ++d1)
+        for (long long d2 = 0; d2 < nd2; ++d2)
+            for (long long d3 = 0; d3 < nd3; ++d3)
+                val += mat(p_out[0] + i1_loc, p_out[1] + i2_loc, p_out[2] + i3_loc, d1, d2, d3) *
+                       x(i1 + d1 + off1, i2 + d2 + off2, i3 + d3 + off3);
+
+    out(p_out[0] + i1_loc, p_out[1] + i2_loc, p_out[2] + i3_loc) = val;
+}