feectools 0.5.0__tar.gz → 0.6.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.5.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>

{feectools-0.5.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.
 
         """
 

{feectools-0.5.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.

{feectools-0.5.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.5.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.5.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.5.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

feectools-0.6.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;
+}

feectools-0.6.0/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_kernels.py

@@ -0,0 +1,63 @@
+"""Stencil matrix-vector product ``out = mat @ x`` on the owned rows (3D).
+
+The host version of ``stencil_dot_3d``.
+
+Moved from ``feectools.linalg.stencil_dot_kernels.matvec_3d_kernel``. The CUDA version in ``stencil_dot_3d_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[3 + 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_3d(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[3]
+    n_diags2 = mat.shape[4]
+    n_diags3 = mat.shape[5]
+    # 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]
+    off3 = p_in[2] - (n_diags3 - 1) // 2 - s_in[2]
+
+    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]
+
+            for i3 in range(s_out[2], e_out[2] + 1):
+                i3_loc = i3 - s_out[2]
+                nd3 = n_diags3
+                if i3 == e_out[2]:
+                    nd3 = n_diags3 - 1 + add[2]
+
+                val = 0.
+                for d1 in range(nd1):
+                    for d2 in range(nd2):
+                        for d3 in range(nd3):
+                            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

{feectools-0.5.0 → feectools-0.6.0}/feectools/linalg/kernels/stencil_transpose_1d/stencil_transpose_1d_cuda.cu

@@ -6,24 +6,24 @@
  * Transpose of a stencil matrix, matT = mat.T on the owned rows of matT, as the pyccel kernel
  * stencil_transpose_1d.
  *
- * One thread per entry (row, diagonal) of `matT` (n_threads = matT.size, see __init__.py). A thread outside
- * the owned rows of matT, or at a diagonal d1 the pyccel kernel does not write (d1 >= 2 * p_in + 1, or
- * >= 2 * p_in + add in the last owned row), returns without writing. Every other thread copies one entry:
- * row i1 of matT, column j1 = i1 - p_in + d1, is row j1 of mat at diagonal p_out + i1 - j1.
+ * One thread per entry (row, diagonal) of `matT` (n_threads = matT.size, see __init__.py). matT has
+ * n = matT.shape[1] diagonals (pads qT = (n - 1) / 2), mat has the pads q = (mat.shape[1] - 1) / 2. A thread
+ * outside the owned rows of matT, or at a diagonal d1 the pyccel kernel does not write (d1 >= n - 1 + add in
+ * the last owned row), returns without writing. Every other thread copies one entry: row i1 of matT, column
+ * j1 = i1 - qT + d1, is row j1 of mat at diagonal q + i1 - j1.
  *
  * @param mat   matrix data, shape (rows of the codomain of mat, diagonals)
  * @param matT  data of the transposed matrix; the owned rows are written
  * @param s_in  global start of the rows of mat (= columns of matT) of this process
- * @param e_in  global end of the rows of mat; not needed in 1D (`mat` is a view), kept for the signature
  * @param p_in  padding of the rows of mat
- * @param add   1 if the last row of matT uses all 2 * p_in + 1 diagonals, else 0
+ * @param add   1 if the last row of matT uses all diagonals, else 0
  * @param s_out global start of the rows of matT of this process
  * @param e_out global end (inclusive) of the rows of matT of this process
  * @param p_out padding of the rows of matT
  */
 extern "C" __global__ void stencil_transpose_1d(Array2D<double> mat, Array2D<double> matT, long long s_in,
-                                                long long e_in, long long p_in, long long add,
-                                                long long s_out, long long e_out, long long p_out)
+                                                long long p_in, long long add, long long s_out,
+                                                long long e_out, long long p_out)
 {
     CUNUMPY_THREAD_1D(thread, matT.size());
 
@@ -31,10 +31,10 @@ extern "C" __global__ void stencil_transpose_1d(Array2D<double> mat, Array2D<dou
     const long long i1_loc = thread / matT.shape[1] - p_out;  // local row index of matT
     if (i1_loc < 0 || i1_loc > e_out - s_out) return;
     const long long i1 = s_out + i1_loc;  // global row index of matT = global column index of mat
-    if (d1 >= ((i1 == e_out) ? 2 * p_in + add : 2 * p_in + 1)) return;
+    if (i1 == e_out && d1 >= matT.shape[1] - 1 + add) return;
 
-    const long long j1 = i1 - p_in + d1;  // global column index of matT
-    const long long j1_loc = j1 - s_in;   // local column index of matT = local row index of mat
+    const long long j1 = i1 - (matT.shape[1] - 1) / 2 + d1;  // global column index of matT
+    const long long j1_loc = j1 - s_in;  // local column index of matT = local row index of mat
 
-    matT(p_out + i1_loc, d1) = mat(p_in + j1_loc, p_out + i1 - j1);
+    matT(p_out + i1_loc, d1) = mat(p_in + j1_loc, (mat.shape[1] - 1) / 2 + i1 - j1);
 }

feectools-0.6.0/feectools/linalg/kernels/stencil_transpose_1d/stencil_transpose_1d_kernels.py

@@ -0,0 +1,38 @@
+"""Transpose of a stencil matrix, ``matT = mat.T`` on the owned rows of ``matT`` (1D).
+
+The host version of ``stencil_transpose_1d``.
+
+Moved from ``feectools.linalg.stencil_transpose_kernels.transpose_1d_kernel``. The CUDA version in
+``stencil_transpose_1d_cuda.cu`` takes the same arguments in the same order.
+
+The numbers of diagonals are read from the matrices: ``matT`` has ``n = matT.shape[1]`` diagonals, so its pads are
+``qT = (n - 1) // 2`` and diagonal ``d`` of row ``i`` is the column ``j = i - qT + d``; that entry is row ``j`` of
+``mat`` at the diagonal ``q + i - j``, with ``q = (mat.shape[1] - 1) // 2``. Interior rows of ``matT`` get all ``n``
+diagonals, the last owned row ``n - 1 + add``.
+"""
+
+
+def stencil_transpose_1d(mat: 'float[:, :]',
+                         matT: 'float[:, :]',
+                         s_in: int,  # refers to matT
+                         p_in: int,
+                         add: int,
+                         s_out: int,
+                         e_out: int,
+                         p_out: int):
+
+    n_diags1 = matT.shape[1]
+    qT1 = (n_diags1 - 1) // 2  # pads of matT
+    q1 = (mat.shape[1] - 1) // 2  # pads of mat
+
+    for i1 in range(s_out, e_out + 1):  # global row index of matT = global column index of mat
+        i1_loc = i1 - s_out  # local row index of matT
+        nd1 = n_diags1
+        if i1 == e_out:
+            nd1 = n_diags1 - 1 + add
+
+        for d1 in range(nd1):
+            j1 = i1 - qT1 + d1  # global column index of matT
+            j1_loc = j1 - s_in  # local column index of matT = local row index of mat
+
+            matT[p_out + i1_loc, d1] = mat[p_in + j1_loc, q1 + i1 - j1]