feectools 0.4.0__tar.gz → 0.6.0__tar.gz

{feectools-0.4.0/feectools.egg-info → feectools-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.4.0
+Version: 0.6.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>
@@ -42,6 +42,7 @@ Requires-Dist: matplotlib
 Requires-Dist: pyyaml>=5.1
 Requires-Dist: packaging
 Requires-Dist: pyevtk
+Requires-Dist: maybempi>=0.1.2
 Requires-Dist: cunumpy>=0.6.1
 Requires-Dist: pyccel>=2.1.0
 Requires-Dist: h5py

feectools-0.6.0/feectools/__init__.py

@@ -0,0 +1,7 @@
+# Bind this rank to its own GPU (by its rank within the node) and create the CUDA context before MPI starts, as
+# CUDA-aware MPI requires. Every feectools module imports this package first, and MPI starts only when a module
+# imports `from maybempi import MPI` (maybempi uses mpi4py under an MPI launcher and a serial stand-in otherwise;
+# MAYBEMPI=0/1 overrides the detection). A no-op on the NumPy backend.
+from cunumpy.cuda import bind_local_device
+
+bind_local_device()

{feectools-0.4.0 → feectools-0.6.0}/feectools/ddm/blocking_data_exchanger.py

@@ -3,8 +3,7 @@
 import cunumpy as xp
 from cunumpy.mpi import synchronize_for_mpi
 import numpy as np
-from feectools.ddm.mpi import mpi as MPI
-
+from maybempi import MPI
 from .cart import CartDecomposition, find_mpi_type
 from .basic import CartDataExchanger
 

{feectools-0.4.0 → feectools-0.6.0}/feectools/ddm/cart.py

@@ -10,16 +10,8 @@ from itertools import product
 
 from cunumpy.xp import array_backend, to_numpy
 
-# Bind this rank to its own GPU (by its rank within the node) and create the
-# CUDA context before MPI is initialized, as CUDA-aware MPI requires. Must stay
-# above the feectools.ddm.mpi import, which initializes MPI as a side effect.
-# A no-op on the NumPy backend.
-from cunumpy.cuda import bind_local_device
-
-bind_local_device()
-
-from feectools.ddm.mpi import mpi as MPI
-from feectools.ddm.mpi import MockMPI
+# The GPU of this rank is bound in feectools/__init__.py, before MPI starts.
+from maybempi import MPI, is_serial
 from feectools.ddm.partition import compute_dims, partition_procs_per_patch
 
 
@@ -46,7 +38,7 @@ def find_mpi_type( dtype ):
         MPI datatype to be used for communication.
 
     """
-    if not isinstance(MPI, MockMPI):
+    if not is_serial(MPI):
         if isinstance( dtype, MPI.Datatype ):
             mpi_type = dtype
         else:
@@ -92,7 +84,7 @@ class MultiPatchDomainDecomposition:
     def __init__(self, ncells, periods, comm=None, num_threads=None):
 
         assert len( ncells ) == len( periods )
-        if not isinstance(MPI, MockMPI) and comm is not None:
+        if not is_serial(MPI) and comm is not None:
             assert isinstance( comm, MPI.Comm )
         num_threads = num_threads if num_threads else int(os.environ.get('OMP_NUM_THREADS', 1))
 
@@ -236,7 +228,7 @@ class DomainDecomposition:
         assert len( ncells ) == len( periods )
         assert all( n >=1 for n in ncells )
         assert all( isinstance( period, bool ) for period in periods )
-        if isinstance(MPI, MockMPI): 
+        if is_serial(MPI):
             comm = None
         else:
             if comm is not None:

{feectools-0.4.0 → feectools-0.6.0}/feectools/ddm/interface_data_exchanger.py

@@ -1,8 +1,7 @@
 # coding: utf-8
 
 from cunumpy.mpi import synchronize_for_mpi
-from feectools.ddm.mpi import mpi as MPI
-
+from maybempi import MPI
 from .cart import InterfaceCartDecomposition, find_mpi_type
 
 __all__ = ('InterfaceCartDataExchanger',)

{feectools-0.4.0 → feectools-0.6.0}/feectools/ddm/nonblocking_data_exchanger.py

@@ -5,7 +5,7 @@ from cunumpy.mpi import synchronize_for_mpi
 import numpy as np
 from itertools import product
 
-from feectools.ddm.mpi import mpi as MPI
+from maybempi import MPI
 from .cart import CartDecomposition, find_mpi_type
 from .basic import CartDataExchanger
 

{feectools-0.4.0 → feectools-0.6.0}/feectools/ddm/tests/test_cart_1d.py

@@ -14,7 +14,7 @@ from feectools.ddm.nonblocking_data_exchanger import NonBlockingCartDataExchange
 def run_cart_1d( data_exchanger_type, verbose=False ):
 
     import cunumpy as xp
-    from feectools.ddm.mpi import mpi as MPI
+    from maybempi import MPI
     from feectools.ddm.cart import DomainDecomposition, CartDecomposition
 
     #---------------------------------------------------------------------------

{feectools-0.4.0 → feectools-0.6.0}/feectools/ddm/tests/test_cart_2d.py

@@ -9,7 +9,7 @@ from feectools.ddm.nonblocking_data_exchanger import NonBlockingCartDataExchange
 def run_cart_2d( data_exchanger_type, verbose=False , nprocs=None, reverse_axis=None):
 
     import cunumpy as xp
-    from feectools.ddm.mpi import mpi as MPI  
+    from maybempi import MPI
     from feectools.ddm.cart import DomainDecomposition, CartDecomposition
 
     #---------------------------------------------------------------------------

{feectools-0.4.0 → feectools-0.6.0}/feectools/ddm/tests/test_cart_3d.py

@@ -9,7 +9,7 @@ from feectools.ddm.nonblocking_data_exchanger import NonBlockingCartDataExchange
 def run_cart_3d( data_exchanger_type, verbose=False ):
 
     import cunumpy as xp
-    from feectools.ddm.mpi import mpi as MPI 
+    from maybempi import MPI
     from feectools.ddm.cart import DomainDecomposition, CartDecomposition
 
     #---------------------------------------------------------------------------

{feectools-0.4.0 → feectools-0.6.0}/feectools/ddm/tests/test_coarsen.py

@@ -2,9 +2,7 @@ import numpy as np
 import pytest
 
 from feectools.ddm.cart import DomainDecomposition
-from feectools.ddm.mpi import mpi as MPI
-
-
+from maybempi import MPI
 def _comm(parallel):
     return MPI.COMM_WORLD if parallel else None
 

{feectools-0.4.0 → feectools-0.6.0}/feectools/ddm/tests/test_multicart_2d.py

@@ -41,7 +41,7 @@ def get_plus_starts_ends(minus_starts, minus_ends, minus_npts, plus_npts, minus_
 def run_carts_2d():
     import cunumpy as xp
     
-    from feectools.ddm.mpi import mpi as MPI
+    from maybempi import MPI
     from feectools.ddm.cart import MultiPatchDomainDecomposition, CartDecomposition, create_interfaces_cart
     from feectools.ddm.blocking_data_exchanger  import BlockingCartDataExchanger
     from feectools.ddm.interface_data_exchanger import InterfaceCartDataExchanger

{feectools-0.4.0 → feectools-0.6.0}/feectools/fem/tensor.py

@@ -5,8 +5,7 @@ We assume here that a tensor space is the product of fem spaces whom basis are
 of compact support
 
 """
-from feectools.ddm.mpi import mpi as MPI
-    
+from maybempi import MPI
 import cunumpy as xp
 import numpy as np
 import itertools

{feectools-0.4.0 → feectools-0.6.0}/feectools/fem/tests/test_spline_interpolation.py

@@ -5,8 +5,7 @@
 #---------------------------------------------------------------------------#
 import time
 
-from feectools.ddm.mpi import mpi as MPI
-    
+from maybempi import MPI
 import cunumpy as xp
 import pytest
 

{feectools-0.4.0 → feectools-0.6.0}/feectools/fem/tests/test_splines_par.py

@@ -7,8 +7,7 @@ from feectools.fem.splines import SplineSpace
 from feectools.fem.tensor  import TensorFemSpace
 from feectools.fem.vector  import VectorFemSpace
 from feectools.ddm.cart    import DomainDecomposition
-from feectools.ddm.mpi import mpi as MPI
-
+from maybempi import MPI
 import cunumpy as xp
 
 def test_2d_1():

{feectools-0.4.0 → feectools-0.6.0}/feectools/fem/tests/test_tensor.py

@@ -19,7 +19,7 @@ import numpy as np
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 from PIL import Image
-from mpi4py import MPI
+from maybempi import MPI
 
 from sympde.topology.domain import Square
 from sympde.topology.space import ScalarFunctionSpace

{feectools-0.4.0 → feectools-0.6.0}/feectools/linalg/basic.py

@@ -93,9 +93,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.
 
         """
 
@@ -385,7 +387,7 @@ class LinearOperator(ABC):
 
         if not is_sparse:
             if comm is not None:
-                from feectools.ddm.mpi import mpi as MPI
+                from maybempi import MPI
                 comm.Allreduce(MPI.IN_PLACE, out, op=MPI.SUM)
             return out
 

{feectools-0.4.0 → feectools-0.6.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.
@@ -1076,7 +1079,7 @@ class BlockLinearOperator(LinearOperator):
         if not self.codomain.parallel:
             return blocks, blocks_T
 
-        from feectools.ddm.mpi import mpi as MPI
+        from maybempi import MPI
         from feectools.linalg.stencil import StencilInterfaceMatrix
 
         if not isinstance(self.codomain, BlockVectorSpace):

{feectools-0.4.0 → feectools-0.6.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.4.0 → feectools-0.6.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.4.0 → feectools-0.6.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.6.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.4.0 → feectools-0.6.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.6.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