diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index a8a1dcc0a..2b1a685a1 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -48,7 +48,7 @@ come from - **No silent CPU fallback** for folder kernels: on the CuPy backend a folder kernel without CUDA version raises. The other pyccel kernels (B-splines, field evaluation, DOF kernels, used at setup) are wrapped in `cunumpy.kernels.PyccelKernel` and copy their arrays to the host and back; see [Open questions](#open-questions). -- **Same cunumpy as struphy.** `cunumpy >= 0.5.0, < 0.6`; names are imported from the submodules +- **Same cunumpy as struphy.** `cunumpy >= 0.6.1` (the 3D stencil kernels need its 6D array views); names are imported from the submodules (`cunumpy.kernels`, `cunumpy.arguments`, `cunumpy.cuda`, `cunumpy.mpi`, `cunumpy.kernel_testing`); cunumpy 0.6 removed the old top-level names, and `CudaKernel`/`CudaKernelVariants` are in `cunumpy.kernels` since 0.6. - **Small steps.** Every PR keeps the NumPy path working and tested. @@ -117,16 +117,11 @@ CuPy. There are no backend branches and no host staging at these call sites any of generated source is gone. - **Signature changes.** The inner kernels add their result to an argument `res` (`res[0] += ...`, the caller zeroes it; `inner` passes the MPI send buffer) instead of returning it, because a CUDA kernel cannot return a - value; on the GPU each block reduces in shared memory and adds its sum with one `atomicAdd`. The transpose - kernels take `e_in` (end of the row range of `mat`) after `s_in`; pyccel does not need it, the 3D CUDA kernel - needs it for the shape of `mat`. -- **6D matrix data.** cunumpy's array views stop at `Array4D`, so the 3D kernels take the six-axis matrix data as - raw pointers (C-contiguity checked by `CudaKernel`) and derive their shape from the other arguments: rows from - `out` (dot) or from `s/e/p` (transpose), and `2 * p + 1` diagonals. This is exactly what the precompiled pyccel - kernels read, and they are wrong for other data too, so `StencilMatrix.set_backend` records whether the matrix - has no shifts and `2 * p + 1` diagonals (`_kernel_shapes_ok`), and `dot`, `vdot` and `transpose` raise - `NotImplementedError` otherwise, on both backends (no test or call site in the test suite builds such a matrix - for these kernels). Array views up to 6D in cunumpy would remove this restriction (see [Open questions](#open-questions)). + value; on the GPU each block reduces in shared memory and adds its sum with one `atomicAdd`. (The transpose + kernels had an argument `e_in` for the shape of the 3D matrix data; it is gone, see the 6D views below.) +- **6D matrix data:** see [6D matrix views](#6d-matrix-views-stencil-6d-views). (In #88 the 3D kernels took the + six-axis data as raw pointers and assumed `2 * p + 1` diagonals; `dot`, `vdot` and `transpose` raised for other + matrices.) - **Launch sizes** are declared in each folder's `__init__.py` (`n_threads_from`): one thread per entry of `out`, `matT`, `v1` or `x`, the last axis varying fastest; threads outside the owned rows or diagonals return without writing, as the pyccel loops do. The inner kernels use blocks of 256 threads. @@ -168,6 +163,37 @@ every CG iteration preconditioned by struphy's `MassMatrixPreconditioner` (strup device against host solves under `assert_no_transfers`, serial and with MPI. Before this change a 3D solve on the fake CuPy made 3 `to_host` copies (one per direction) and 3 uncounted copies back. +## 6D matrix views (`stencil-6d-views`) + +With cunumpy 0.6.1 (`Array5D`/`Array6D`, `CArray5D`/`CArray6D` in `cunumpy/array_view.cuh`) the 3D kernels take +the matrix data as views, like the 1D/2D kernels (`Array2D`/`Array4D`): + +- `stencil_dot_3d(Array6D mat, Array3D x, Array3D out, s_in, p_in, add, s_out, e_out, + p_out)` and `stencil_transpose_3d(Array6D mat, Array6D matT, s_in, p_in, add, s_out, e_out, + p_out)`; the pyccel versions take `float[:, :, :, :, :, :]` with the same arguments in the same order. Strided + views (`Array6D`, not `CArray6D`), as in 1D/2D, so a non-contiguous matrix is not refused. +- **No `e_in`.** The transpose kernels (all dimensions, the argument list is the same for all) lost `e_in`: it was + only needed for the row extents of the raw pointer in 3D. +- **The number of diagonals comes from the data.** All six kernels (1D, 2D, 3D, pyccel and CUDA) read the number + of diagonals `n_k` of each direction from the shape of the matrix data instead of assuming `2 * p_in + 1`: the + pads of the matrix are `q_k = (n_k - 1) // 2` (at most the pads of the spaces, `StencilMatrix(V, W, pads=...)`), + diagonal `d` of row `i` is the column `i - q + d` (in `x` at `i - q + d - s_in + p_in`), and the last owned row + uses `n_k - 1 + add[k]` diagonals. With `q = p` this is the old loop, in the same order and with bitwise the same + results (checked against the kernels of #88 run as Python); the pyccel kernels are now one loop nest with the + number of diagonals picked per row instead of the eight (3D) spelled-out combinations. +- **`StencilMatrix`.** `dot`, `vdot` and `transpose` no longer raise for matrices with fewer diagonals (blocks + between spaces of different degree, derivative-type stencils); `transpose(out=...)` asserts that `out` has the + pads of the matrix. The one restriction left is spaces with **shifts > 1**, which still raise + `NotImplementedError` (`_check_kernel_shifts`): the stencil kernels have never handled them, and there is no + reference to match, since the general psydac kernels (`matvec_d`, `transpose_d`) disagree with `toarray()` + and with each other (the transpose is not the adjoint of the product) for shifts > 1, and psydac never tested + them ("TODO: verify for s>1"). +- **Tests.** `MATRIX_CASES` (parity and CPU emulation) has matrices with fewer diagonals and non-periodic + rectangular blocks between spaces of different size per direction, including pads 0 (one diagonal); + `test_device_matvec.py` compares `dot`, `vdot` and `transpose` of such matrices with dense `toarray()` + references, and `test_mpi_device.py` checks a matrix with fewer diagonals and the adjoint identity of its + transpose against the global field on any number of ranks. + ## Kernel folders ``` @@ -214,8 +240,9 @@ Conventions, as in struphy: `PyccelKernel` (host copies). They run at setup, not in the time loop; they move into kernel folders with CUDA versions when a profile shows they matter. - **GPU CI.** No GPU runner yet; GPU tests are run by hand. -- **6D array views in cunumpy.** With `Array5D`/`Array6D` (also needed by struphy's matrix accumulations), the 3D - kernels could take the matrix data as views, drop the shape assumptions and the `e_in` argument. +- **Shifts > 1.** `StencilMatrix.dot`/`transpose` raise for spaces with shifts > 1 (see + [6D matrix views](#6d-matrix-views-stencil-6d-views)); supporting them needs a verified definition of the data + layout first. - **Complex data on the device.** Not needed by struphy so far; would need a second CUDA kernel per folder (or dtype dispatch in `cunumpy.kernels.Kernel`). - **`inner` reduction.** One `atomicAdd` per block of 256 threads, then a copy of the 8-byte result to the host in diff --git a/feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_cuda.cu b/feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_cuda.cu index f3960ba05..95d4650bc 100644 --- a/feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_cuda.cu +++ b/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 mat, Array1D 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; } diff --git a/feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_kernels.py b/feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_kernels.py index 885b7b41b..4be219aad 100644 --- a/feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_kernels.py +++ b/feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_kernels.py @@ -4,6 +4,10 @@ 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``. """ @@ -17,19 +21,18 @@ def stencil_dot_1d(mat: 'float[:, :]', e_out: int, p_out: int): - for i1 in range(s_out, e_out): # global row index + 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(2*p_in + 1): - val += mat[p_out + i1_loc, d1] * x[i1 + d1 - s_in] + for d1 in range(nd1): + val += mat[p_out + i1_loc, d1] * x[i1 + d1 + off1] out[p_out + i1_loc] = val - - # last row treated separately - i1 = e_out - i1_loc = i1 - s_out # local row index - val = 0. - for d1 in range(2*p_in + add): - val += mat[p_out + i1_loc, d1] * x[i1 + d1 - s_in] - - out[p_out + i1_loc] = val diff --git a/feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_cuda.cu b/feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_cuda.cu index 8ba154b49..1327dfd78 100644 --- a/feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_cuda.cu +++ b/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 mat, Array2D 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; } diff --git a/feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_kernels.py b/feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_kernels.py index 4b86a3e9f..e2e378c0a 100644 --- a/feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_kernels.py +++ b/feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_kernels.py @@ -4,6 +4,10 @@ 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]``. """ @@ -17,83 +21,27 @@ def stencil_dot_2d(mat: 'float[:, :, :, :]', e_out: 'int[:]', p_out: 'int[:]'): - ##################################### - ##################################### - # without last row in 1st direction # - ##################################### - ##################################### - for i1 in range(s_out[0], e_out[0]): - i1_loc = i1 - s_out[0] + 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] - ##################################### - # without last row in 2nd direction # - ##################################### - for i2 in range(s_out[1], e_out[1]): + 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(2 * p_in[0] + 1): - for d2 in range(2 * p_in[1] + 1): - 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]] - out[p_out[0] + i1_loc, - p_out[1] + i2_loc] = val - - ############################################## - # treat last row in 2nd direction separately # - ############################################## - i2 = e_out[1] - i2_loc = i2 - s_out[1] - - val = 0. - for d1 in range(2 * p_in[0] + 1): - for d2 in range(2 * p_in[1] + add[1]): - 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]] - out[p_out[0] + i1_loc, - p_out[1] + i2_loc] = val - - ############################################## - ############################################## - # treat last row in 1st direction separately # - ############################################## - ############################################## - i1 = e_out[0] - i1_loc = i1 - s_out[0] - - ##################################### - # without last row in 2nd direction # - ##################################### - for i2 in range(s_out[1], e_out[1]): - i2_loc = i2 - s_out[1] - - val = 0. - for d1 in range(2 * p_in[0] + add[0]): - for d2 in range(2 * p_in[1] + 1): - 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]] - out[p_out[0] + i1_loc, - p_out[1] + i2_loc] = val - - ############################################## - # treat last row in 2nd direction separately # - ############################################## - i2 = e_out[1] - i2_loc = i2 - s_out[1] - - val = 0. - for d1 in range(2 * p_in[0] + add[0]): - for d2 in range(2 * p_in[1] + add[1]): - 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 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 + out[p_out[0] + i1_loc, p_out[1] + i2_loc] = val diff --git a/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_cuda.cu b/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_cuda.cu index 09e8412f6..066fa86d3 100644 --- a/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_cuda.cu +++ b/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_cuda.cu @@ -7,22 +7,16 @@ * * 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 eight combinations, this - * kernel picks its own per thread. The diagonals are summed in the same order (d1, d2, d3 innermost). + * 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). * - * `mat` has six axes and cunumpy's array views stop at four, so it is a raw pointer to C-contiguous data - * (checked by cunumpy at the call) with the shape (out.shape[0], out.shape[1], out.shape[2], - * 2 * p_in[0] + 1, 2 * p_in[1] + 1, 2 * p_in[2] + 1): the rows of the matrix are those of `out`, and it - * has exactly the diagonals the pyccel kernel reads. StencilMatrix checks the diagonals before calling - * the precompiled kernels (see StencilMatrix.set_backend). - * - * @param mat matrix data, C-contiguous, shape as above + * @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(const double* mat, Array3D x, Array3D out, +extern "C" __global__ void stencil_dot_3d(Array6D mat, Array3D x, Array3D 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) @@ -39,26 +33,20 @@ extern "C" __global__ void stencil_dot_3d(const double* mat, Array3D x, const long long i2 = s_out[1] + i2_loc; const long long i3 = s_out[2] + i3_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 n_diags3 = (i3 == e_out[2]) ? 2 * p_in[2] + add[2] : 2 * p_in[2] + 1; - - // element strides of the C-contiguous matrix data - const long long stride_d3 = 1; - const long long stride_d2 = (2 * p_in[2] + 1) * stride_d3; - const long long stride_d1 = (2 * p_in[1] + 1) * stride_d2; - const long long stride_i3 = (2 * p_in[0] + 1) * stride_d1; - const long long stride_i2 = out.shape[2] * stride_i3; - const long long stride_i1 = out.shape[1] * stride_i2; - const double* row = mat + (p_out[0] + i1_loc) * stride_i1 + (p_out[1] + i2_loc) * stride_i2 + - (p_out[2] + i3_loc) * stride_i3; + 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 < n_diags1; ++d1) - for (long long d2 = 0; d2 < n_diags2; ++d2) - for (long long d3 = 0; d3 < n_diags3; ++d3) - val += row[d1 * stride_d1 + d2 * stride_d2 + d3 * stride_d3] * - x(i1 + d1 - s_in[0], i2 + d2 - s_in[1], i3 + d3 - s_in[2]); + 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; } diff --git a/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_kernels.py b/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_kernels.py index c3ddc63fc..0b7f7896a 100644 --- a/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_kernels.py +++ b/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_kernels.py @@ -4,6 +4,10 @@ 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]``. """ @@ -17,190 +21,43 @@ def stencil_dot_3d(mat: 'float[:, :, :, :, :, :]', e_out: 'int[:]', p_out: 'int[:]'): - ##################################### - ##################################### - # without last row in 1st direction # - ##################################### - ##################################### - for i1 in range(s_out[0], e_out[0]): - i1_loc = i1 - s_out[0] - - ##################################### - # without last row in 2nd direction # - ##################################### - for i2 in range(s_out[1], e_out[1]): + 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] - # without last row in 3rd direction - for i3 in range(s_out[2], e_out[2]): + 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(2*p_in[0] + 1): - for d2 in range(2*p_in[1] + 1): - for d3 in range(2*p_in[2] + 1): - + 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 - s_in[0], - i2 + d2 - s_in[1], - i3 + d3 - s_in[2]] + 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 - - # treat last row in 3rd direction separately - i3 = e_out[2] - i3_loc = i3 - s_out[2] - val = 0. - for d1 in range(2*p_in[0] + 1): - for d2 in range(2*p_in[1] + 1): - for d3 in range(2*p_in[2] + add[2]): - - val += mat[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] * x[i1 + d1 - s_in[0], - i2 + d2 - s_in[1], - i3 + d3 - s_in[2]] - out[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc] = val - - ############################################## - # treat last row in 2nd direction separately # - ############################################## - i2 = e_out[1] - i2_loc = i2 - s_out[1] - - # without last row in 3rd direction - for i3 in range(s_out[2], e_out[2]): - i3_loc = i3 - s_out[2] - - val = 0. - for d1 in range(2*p_in[0] + 1): - for d2 in range(2*p_in[1] + add[1]): - for d3 in range(2*p_in[2] + 1): - - val += mat[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] * x[i1 + d1 - s_in[0], - i2 + d2 - s_in[1], - i3 + d3 - s_in[2]] - out[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc] = val - - # treat last row in 3rd direction separately - i3 = e_out[2] - i3_loc = i3 - s_out[2] - val = 0. - for d1 in range(2*p_in[0] + 1): - for d2 in range(2*p_in[1] + add[1]): - for d3 in range(2*p_in[2] + add[2]): - - val += mat[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] * x[i1 + d1 - s_in[0], - i2 + d2 - s_in[1], - i3 + d3 - s_in[2]] - out[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc] = val - - ############################################## - ############################################## - # treat last row in 1st direction separately # - ############################################## - ############################################## - i1 = e_out[0] - i1_loc = i1 - s_out[0] - - ##################################### - # without last row in 2nd direction # - ##################################### - for i2 in range(s_out[1], e_out[1]): - i2_loc = i2 - s_out[1] - - # without last row in 3rd direction - for i3 in range(s_out[2], e_out[2]): - i3_loc = i3 - s_out[2] - - val = 0. - for d1 in range(2*p_in[0] + add[0]): - for d2 in range(2*p_in[1] + 1): - for d3 in range(2*p_in[2] + 1): - - val += mat[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] * x[i1 + d1 - s_in[0], - i2 + d2 - s_in[1], - i3 + d3 - s_in[2]] - out[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc] = val - - # treat last row in 3rd direction separately - i3 = e_out[2] - i3_loc = i3 - s_out[2] - val = 0. - for d1 in range(2*p_in[0] + add[0]): - for d2 in range(2*p_in[1] + 1): - for d3 in range(2*p_in[2] + add[2]): - - val += mat[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] * x[i1 + d1 - s_in[0], - i2 + d2 - s_in[1], - i3 + d3 - s_in[2]] - out[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc] = val - - ############################################## - # treat last row in 2nd direction separately # - ############################################## - i2 = e_out[1] - i2_loc = i2 - s_out[1] - - # without last row in 3rd direction - for i3 in range(s_out[2], e_out[2]): - i3_loc = i3 - s_out[2] - - val = 0. - for d1 in range(2*p_in[0] + add[0]): - for d2 in range(2*p_in[1] + add[1]): - for d3 in range(2*p_in[2] + 1): - - val += mat[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] * x[i1 + d1 - s_in[0], - i2 + d2 - s_in[1], - i3 + d3 - s_in[2]] - out[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc] = val - - # treat last row in 3rd direction separately - i3 = e_out[2] - i3_loc = i3 - s_out[2] - val = 0. - for d1 in range(2*p_in[0] + add[0]): - for d2 in range(2*p_in[1] + add[1]): - for d3 in range(2*p_in[2] + add[2]): - - val += mat[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] * x[i1 + d1 - s_in[0], - i2 + d2 - s_in[1], - i3 + d3 - s_in[2]] - out[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc] = val diff --git a/feectools/linalg/kernels/stencil_transpose_1d/stencil_transpose_1d_cuda.cu b/feectools/linalg/kernels/stencil_transpose_1d/stencil_transpose_1d_cuda.cu index 01242ad31..32ec75f58 100644 --- a/feectools/linalg/kernels/stencil_transpose_1d/stencil_transpose_1d_cuda.cu +++ b/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 mat, Array2D 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 mat, Array2D 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); } diff --git a/feectools/linalg/kernels/stencil_transpose_1d/stencil_transpose_1d_kernels.py b/feectools/linalg/kernels/stencil_transpose_1d/stencil_transpose_1d_kernels.py index 9009a0d21..3353b08a5 100644 --- a/feectools/linalg/kernels/stencil_transpose_1d/stencil_transpose_1d_kernels.py +++ b/feectools/linalg/kernels/stencil_transpose_1d/stencil_transpose_1d_kernels.py @@ -3,35 +3,36 @@ 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. ``e_in`` (the end of the row range of ``mat``) is not -needed here; the CUDA version needs it for the row extents of ``mat``, which it reads through a raw -pointer in 3D. +``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 - e_in: int, # refers to matT; used by the CUDA version only p_in: int, add: int, s_out: int, e_out: int, p_out: int): - for i1 in range(s_out, e_out): # global row index of matT = global column index of mat - i1_loc = i1 - s_out # local row index of matT - for d1 in range(2*p_in + 1): - j1 = i1 - p_in + d1 # global column index of matT - j1_loc = j1 - s_in # local column index of matT = local row index of mat + n_diags1 = matT.shape[1] + qT1 = (n_diags1 - 1) // 2 # pads of matT + q1 = (mat.shape[1] - 1) // 2 # pads of mat - matT[p_out + i1_loc, d1] = mat[p_in + j1_loc, p_out + i1 - j1] + 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 - # last row treated separately - i1 = e_out - i1_loc = i1 - s_out # local row index of matT - for d1 in range(2*p_in + add): - j1 = i1 - p_in + d1 # global column index of matT - j1_loc = j1 - s_in # local column index of matT = local row index of mat + 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, p_out + i1 - j1] + matT[p_out + i1_loc, d1] = mat[p_in + j1_loc, q1 + i1 - j1] diff --git a/feectools/linalg/kernels/stencil_transpose_2d/stencil_transpose_2d_cuda.cu b/feectools/linalg/kernels/stencil_transpose_2d/stencil_transpose_2d_cuda.cu index 5ffefcc04..782dbd69b 100644 --- a/feectools/linalg/kernels/stencil_transpose_2d/stencil_transpose_2d_cuda.cu +++ b/feectools/linalg/kernels/stencil_transpose_2d/stencil_transpose_2d_cuda.cu @@ -8,18 +8,17 @@ * * One thread per entry (i1, i2, d1, d2) of `matT` (n_threads = matT.size, see __init__.py), the last axis * varying fastest. A thread outside the owned rows of matT, or at a diagonal the pyccel kernel does not - * write (per direction k: d_k >= 2 * p_in[k] + 1, or >= 2 * p_in[k] + add[k] in the last owned row), - * returns without writing; every other thread copies one entry, as in the 1D kernel. + * write (per direction k: d_k >= matT.shape[2 + k] - 1 + add[k] in the last owned row), returns without + * writing; every other thread copies one entry, as in the 1D kernel. * * @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, e_in, p_in, add, s_out, e_out, p_out per direction (length 2), as in the 1D kernel - * (e_in is not needed in 2D, `mat` is a view) + * @param s_in, p_in, add, s_out, e_out, p_out per direction (length 2), as in the 1D kernel */ extern "C" __global__ void stencil_transpose_2d(Array4D mat, Array4D matT, const long long* s_in, - const long long* e_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) + 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, matT.size()); @@ -34,14 +33,14 @@ extern "C" __global__ void stencil_transpose_2d(Array4D mat, Array4D e_out[1] - s_out[1]) return; const long long i1 = s_out[0] + i1_loc; // global row indices of matT = global column indices of mat const long long i2 = s_out[1] + i2_loc; - if (d1 >= ((i1 == e_out[0]) ? 2 * p_in[0] + add[0] : 2 * p_in[0] + 1)) return; - if (d2 >= ((i2 == e_out[1]) ? 2 * p_in[1] + add[1] : 2 * p_in[1] + 1)) return; + if (i1 == e_out[0] && d1 >= matT.shape[2] - 1 + add[0]) return; + if (i2 == e_out[1] && d2 >= matT.shape[3] - 1 + add[1]) return; - const long long j1 = i1 - p_in[0] + d1; // global column indices of matT - const long long j2 = i2 - p_in[1] + d2; + const long long j1 = i1 - (matT.shape[2] - 1) / 2 + d1; // global column indices of matT + const long long j2 = i2 - (matT.shape[3] - 1) / 2 + d2; const long long j1_loc = j1 - s_in[0]; // local column indices of matT = local row indices of mat const long long j2_loc = j2 - s_in[1]; matT(p_out[0] + i1_loc, p_out[1] + i2_loc, d1, d2) = - mat(p_in[0] + j1_loc, p_in[1] + j2_loc, p_out[0] + i1 - j1, p_out[1] + i2 - j2); + mat(p_in[0] + j1_loc, p_in[1] + j2_loc, (mat.shape[2] - 1) / 2 + i1 - j1, (mat.shape[3] - 1) / 2 + i2 - j2); } diff --git a/feectools/linalg/kernels/stencil_transpose_2d/stencil_transpose_2d_kernels.py b/feectools/linalg/kernels/stencil_transpose_2d/stencil_transpose_2d_kernels.py index 961eab6b5..fea54c2a8 100644 --- a/feectools/linalg/kernels/stencil_transpose_2d/stencil_transpose_2d_kernels.py +++ b/feectools/linalg/kernels/stencil_transpose_2d/stencil_transpose_2d_kernels.py @@ -3,114 +3,50 @@ The host version of ``stencil_transpose_2d``. Moved from ``feectools.linalg.stencil_transpose_kernels.transpose_2d_kernel``. The CUDA version in -``stencil_transpose_2d_cuda.cu`` takes the same arguments in the same order. ``e_in`` (the end of the row range of ``mat``) is not -needed here; the CUDA version needs it for the row extents of ``mat``, which it reads through a raw -pointer in 3D. +``stencil_transpose_2d_cuda.cu`` takes the same arguments in the same order. + +The numbers of diagonals are read from the matrices, per direction as in ``stencil_transpose_1d``: diagonal ``d_k`` +of row ``i_k`` of ``matT`` is the column ``j_k = i_k - qT_k + d_k``, which is row ``j_k`` of ``mat`` at the diagonal +``q_k + i_k - j_k`` (``qT_k``, ``q_k``: the pads of ``matT`` and ``mat``, from their shapes). """ def stencil_transpose_2d(mat: 'float[:, :, :, :]', matT: 'float[:, :, :, :]', s_in: 'int[:]', # refers to matT - e_in: 'int[:]', # refers to matT; used by the CUDA version only p_in: 'int[:]', add: 'int[:]', s_out: 'int[:]', e_out: 'int[:]', p_out: 'int[:]'): - ##################################### - ##################################### - # without last row in 1st direction # - ##################################### - ##################################### - for i1 in range(s_out[0], e_out[0]): # global row index of matT = global column index of mat - i1_loc = i1 - s_out[0] # local row index of matT - - ##################################### - # without last row in 2nd direction # - ##################################### - for i2 in range(s_out[1], e_out[1]): # global row index of matT = global column index of mat - i2_loc = i2 - s_out[1] # local row index of matT - - for d1 in range(2 * p_in[0] + 1): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2 * p_in[1] + 1): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - d1, d2] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2] - - ############################################## - # treat last row in 2nd direction separately # - ############################################## - i2 = e_out[1] - i2_loc = i2 - s_out[1] # local row index of matT - - for d1 in range(2 * p_in[0] + 1): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2 * p_in[1] + add[1]): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - d1, d2] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2] - - ############################################## - ############################################## - # treat last row in 1st direction separately # - ############################################## - ############################################## - i1 = e_out[0] - i1_loc = i1 - s_out[0] # local row index of matT - - ##################################### - # without last row in 2nd direction # - ##################################### - for i2 in range(s_out[1], e_out[1]): - i2_loc = i2 - s_out[1] # local row index of matT - - for d1 in range(2 * p_in[0] + add[0]): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2 * p_in[1] + 1): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - d1, d2] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2] - - ############################################## - # treat last row in 2nd direction separately # - ############################################## - i2 = e_out[1] - i2_loc = i2 - s_out[1] # local row index of matT - - for d1 in range(2 * p_in[0] + add[0]): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2 * p_in[1] + add[1]): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - d1, d2] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2] + n_diags1 = matT.shape[2] + n_diags2 = matT.shape[3] + qT1 = (n_diags1 - 1) // 2 # pads of matT + qT2 = (n_diags2 - 1) // 2 + q1 = (mat.shape[2] - 1) // 2 # pads of mat + q2 = (mat.shape[3] - 1) // 2 + + for i1 in range(s_out[0], e_out[0] + 1): # global row indices of matT = global column indices of mat + i1_loc = i1 - s_out[0] # local row indices of matT + 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 d1 in range(nd1): + j1 = i1 - qT1 + d1 # global column indices of matT + j1_loc = j1 - s_in[0] # local column indices of matT = local row indices of mat + for d2 in range(nd2): + j2 = i2 - qT2 + d2 + j2_loc = j2 - s_in[1] + + matT[p_out[0] + i1_loc, p_out[1] + i2_loc, d1, d2] = mat[p_in[0] + j1_loc, + p_in[1] + j2_loc, + q1 + i1 - j1, + q2 + i2 - j2] diff --git a/feectools/linalg/kernels/stencil_transpose_3d/stencil_transpose_3d_cuda.cu b/feectools/linalg/kernels/stencil_transpose_3d/stencil_transpose_3d_cuda.cu index 9c121ea4c..fb60b6168 100644 --- a/feectools/linalg/kernels/stencil_transpose_3d/stencil_transpose_3d_cuda.cu +++ b/feectools/linalg/kernels/stencil_transpose_3d/stencil_transpose_3d_cuda.cu @@ -1,4 +1,5 @@ // CUDA version of stencil_transpose_3d (stencil_transpose_3d_kernels.py), with the same arguments in the same order. +#include "cunumpy/array_view.cuh" #include "cunumpy/index.cuh" /** @@ -6,71 +7,51 @@ * stencil_transpose_3d. * * One thread per entry (i1, i2, i3, d1, d2, d3) of `matT` (n_threads = matT.size, see __init__.py), the last - * axis varying fastest. A thread outside the owned rows of matT, or at a diagonal the pyccel kernel does - * not write (per direction k: d_k >= 2 * p_in[k] + 1, or >= 2 * p_in[k] + add[k] in the last owned row), - * returns without writing; every other thread copies one entry, as in the 1D kernel. + * axis varying fastest. A thread outside the owned rows of matT, or at a diagonal the pyccel kernel does not + * write (per direction k: d_k >= matT.shape[3 + k] - 1 + add[k] in the last owned row), returns without + * writing; every other thread copies one entry, as in the 1D kernel. Both matrices are 6D views, so their + * shapes (rows and diagonals) come with them. * - * Both matrices have six axes and cunumpy's array views stop at four, so they are raw pointers to - * C-contiguous data (checked by cunumpy at the call), with the shapes of stencil matrices without shifts: - * matT has the rows e_out - s_out + 1 + 2 * p_out and the diagonals 2 * p_in + 1, mat the rows - * e_in - s_in + 1 + 2 * p_in and the diagonals 2 * p_out + 1 (per direction). StencilMatrix checks this - * before calling the precompiled kernels (see StencilMatrix.set_backend). - * - * @param mat matrix data, C-contiguous, shape as above - * @param matT data of the transposed matrix, C-contiguous, shape as above; the owned rows are written - * @param s_in, e_in, p_in, add, s_out, e_out, p_out per direction (length 3), as in the 1D kernel; e_in is - * the global end (inclusive) of the rows of mat, used for the shape of mat + * @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, p_in, add, s_out, e_out, p_out per direction (length 3), as in the 1D kernel */ -extern "C" __global__ void stencil_transpose_3d(const double* mat, double* matT, const long long* s_in, - const long long* e_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) +extern "C" __global__ void stencil_transpose_3d(Array6D mat, Array6D matT, 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) { - // shapes of matT (rows, diagonals) and mat (rows, diagonals), per direction - long long rows_T[3], diags_T[3], rows[3], diags[3]; - for (int k = 0; k < 3; ++k) { - rows_T[k] = e_out[k] - s_out[k] + 1 + 2 * p_out[k]; - diags_T[k] = 2 * p_in[k] + 1; - rows[k] = e_in[k] - s_in[k] + 1 + 2 * p_in[k]; - diags[k] = 2 * p_out[k] + 1; - } - const long long size_T = rows_T[0] * rows_T[1] * rows_T[2] * diags_T[0] * diags_T[1] * diags_T[2]; - CUNUMPY_THREAD_1D(thread, size_T); + CUNUMPY_THREAD_1D(thread, matT.size()); long long rest = thread; - const long long d3 = rest % diags_T[2]; - rest /= diags_T[2]; - const long long d2 = rest % diags_T[1]; - rest /= diags_T[1]; - const long long d1 = rest % diags_T[0]; - rest /= diags_T[0]; - const long long i3_loc = rest % rows_T[2] - p_out[2]; // local row indices of matT - rest /= rows_T[2]; - const long long i2_loc = rest % rows_T[1] - p_out[1]; - const long long i1_loc = rest / rows_T[1] - p_out[0]; + const long long d3 = rest % matT.shape[5]; + rest /= matT.shape[5]; + const long long d2 = rest % matT.shape[4]; + rest /= matT.shape[4]; + const long long d1 = rest % matT.shape[3]; + rest /= matT.shape[3]; + const long long i3_loc = rest % matT.shape[2] - p_out[2]; // local row indices of matT + rest /= matT.shape[2]; + const long long i2_loc = rest % matT.shape[1] - p_out[1]; + const long long i1_loc = rest / matT.shape[1] - p_out[0]; 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 of matT = global column indices of mat const long long i2 = s_out[1] + i2_loc; const long long i3 = s_out[2] + i3_loc; - if (d1 >= ((i1 == e_out[0]) ? 2 * p_in[0] + add[0] : 2 * p_in[0] + 1)) return; - if (d2 >= ((i2 == e_out[1]) ? 2 * p_in[1] + add[1] : 2 * p_in[1] + 1)) return; - if (d3 >= ((i3 == e_out[2]) ? 2 * p_in[2] + add[2] : 2 * p_in[2] + 1)) return; + if (i1 == e_out[0] && d1 >= matT.shape[3] - 1 + add[0]) return; + if (i2 == e_out[1] && d2 >= matT.shape[4] - 1 + add[1]) return; + if (i3 == e_out[2] && d3 >= matT.shape[5] - 1 + add[2]) return; - const long long j1 = i1 - p_in[0] + d1; // global column indices of matT - const long long j2 = i2 - p_in[1] + d2; - const long long j3 = i3 - p_in[2] + d3; + const long long j1 = i1 - (matT.shape[3] - 1) / 2 + d1; // global column indices of matT + const long long j2 = i2 - (matT.shape[4] - 1) / 2 + d2; + const long long j3 = i3 - (matT.shape[5] - 1) / 2 + d3; const long long j1_loc = j1 - s_in[0]; // local column indices of matT = local row indices of mat const long long j2_loc = j2 - s_in[1]; const long long j3_loc = j3 - s_in[2]; - const long long index_T = - ((((((p_out[0] + i1_loc) * rows_T[1] + p_out[1] + i2_loc) * rows_T[2] + p_out[2] + i3_loc) * diags_T[0] + - d1) * diags_T[1] + d2) * diags_T[2] + d3); - const long long index = - ((((((p_in[0] + j1_loc) * rows[1] + p_in[1] + j2_loc) * rows[2] + p_in[2] + j3_loc) * diags[0] + - p_out[0] + i1 - j1) * diags[1] + p_out[1] + i2 - j2) * diags[2] + p_out[2] + i3 - j3); - - matT[index_T] = mat[index]; + matT(p_out[0] + i1_loc, p_out[1] + i2_loc, p_out[2] + i3_loc, d1, d2, d3) = + mat(p_in[0] + j1_loc, p_in[1] + j2_loc, p_in[2] + j3_loc, + (mat.shape[3] - 1) / 2 + i1 - j1, (mat.shape[4] - 1) / 2 + i2 - j2, (mat.shape[5] - 1) / 2 + i3 - j3); } diff --git a/feectools/linalg/kernels/stencil_transpose_3d/stencil_transpose_3d_kernels.py b/feectools/linalg/kernels/stencil_transpose_3d/stencil_transpose_3d_kernels.py index 08c670a5d..864426b63 100644 --- a/feectools/linalg/kernels/stencil_transpose_3d/stencil_transpose_3d_kernels.py +++ b/feectools/linalg/kernels/stencil_transpose_3d/stencil_transpose_3d_kernels.py @@ -3,250 +3,67 @@ The host version of ``stencil_transpose_3d``. Moved from ``feectools.linalg.stencil_transpose_kernels.transpose_3d_kernel``. The CUDA version in -``stencil_transpose_3d_cuda.cu`` takes the same arguments in the same order. ``e_in`` (the end of the row range of ``mat``) is not -needed here; the CUDA version needs it for the row extents of ``mat``, which it reads through a raw -pointer in 3D. +``stencil_transpose_3d_cuda.cu`` takes the same arguments in the same order. + +The numbers of diagonals are read from the matrices, per direction as in ``stencil_transpose_1d``: diagonal ``d_k`` +of row ``i_k`` of ``matT`` is the column ``j_k = i_k - qT_k + d_k``, which is row ``j_k`` of ``mat`` at the diagonal +``q_k + i_k - j_k`` (``qT_k``, ``q_k``: the pads of ``matT`` and ``mat``, from their shapes). """ def stencil_transpose_3d(mat: 'float[:, :, :, :, :, :]', matT: 'float[:, :, :, :, :, :]', s_in: 'int[:]', # refers to matT - e_in: 'int[:]', # refers to matT; used by the CUDA version only p_in: 'int[:]', add: 'int[:]', s_out: 'int[:]', e_out: 'int[:]', p_out: 'int[:]'): - ##################################### - ##################################### - # without last row in 1st direction # - ##################################### - ##################################### - for i1 in range(s_out[0], e_out[0]): # global row index of matT = global column index of mat - i1_loc = i1 - s_out[0] # local row index of matT - - ##################################### - # without last row in 2nd direction # - ##################################### - for i2 in range(s_out[1], e_out[1]): # global row index of matT = global column index of mat - i2_loc = i2 - s_out[1] # local row index of matT - - # without last row in 3rd direction - for i3 in range(s_out[2], e_out[2]): # global row index of matT = global column index of mat - i3_loc = i3 - s_out[2] # local row index of matT - - for d1 in range(2*p_in[0] + 1): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2*p_in[1] + 1): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - for d3 in range(2*p_in[2] + 1): - j3 = i3 - p_in[2] + d3 # global column index of matT - j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat - + n_diags1 = matT.shape[3] + n_diags2 = matT.shape[4] + n_diags3 = matT.shape[5] + qT1 = (n_diags1 - 1) // 2 # pads of matT + qT2 = (n_diags2 - 1) // 2 + qT3 = (n_diags3 - 1) // 2 + q1 = (mat.shape[3] - 1) // 2 # pads of mat + q2 = (mat.shape[4] - 1) // 2 + q3 = (mat.shape[5] - 1) // 2 + + for i1 in range(s_out[0], e_out[0] + 1): # global row indices of matT = global column indices of mat + i1_loc = i1 - s_out[0] # local row indices of matT + 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] + + for d1 in range(nd1): + j1 = i1 - qT1 + d1 # global column indices of matT + j1_loc = j1 - s_in[0] # local column indices of matT = local row indices of mat + for d2 in range(nd2): + j2 = i2 - qT2 + d2 + j2_loc = j2 - s_in[1] + for d3 in range(nd3): + j3 = i3 - qT3 + d3 + j3_loc = j3 - s_in[2] + matT[p_out[0] + i1_loc, p_out[1] + i2_loc, p_out[2] + i3_loc, d1, d2, d3] = mat[p_in[0] + j1_loc, p_in[1] + j2_loc, p_in[2] + j3_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2, - p_out[2] + i3 - j3] - - # treat last row in 3rd direction separately - i3 = e_out[2] - i3_loc = i3 - s_out[2] # local row index of matT - - for d1 in range(2*p_in[0] + 1): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2*p_in[1] + 1): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - for d3 in range(2*p_in[2] + add[2]): - j3 = i3 - p_in[2] + d3 # global column index of matT - j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_in[2] + j3_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2, - p_out[2] + i3 - j3] - - ############################################## - # treat last row in 2nd direction separately # - ############################################## - i2 = e_out[1] - i2_loc = i2 - s_out[1] # local row index of matT - - # without last row in 3rd direction - for i3 in range(s_out[2], e_out[2]): # global row index of matT = global column index of mat - i3_loc = i3 - s_out[2] # local row index of matT - - for d1 in range(2*p_in[0] + 1): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2*p_in[1] + add[1]): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - for d3 in range(2*p_in[2] + 1): - j3 = i3 - p_in[2] + d3 # global column index of matT - j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_in[2] + j3_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2, - p_out[2] + i3 - j3] - - # treat last row in 3rd direction separately - i3 = e_out[2] - i3_loc = i3 - s_out[2] # local row index of matT - - for d1 in range(2*p_in[0] + 1): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2*p_in[1] + add[1]): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - for d3 in range(2*p_in[2] + add[2]): - j3 = i3 - p_in[2] + d3 # global column index of matT - j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_in[2] + j3_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2, - p_out[2] + i3 - j3] - - ############################################## - ############################################## - # treat last row in 1st direction separately # - ############################################## - ############################################## - i1 = e_out[0] - i1_loc = i1 - s_out[0] # local row index of matT - - ##################################### - # without last row in 2nd direction # - ##################################### - for i2 in range(s_out[1], e_out[1]): # global row index of matT = global column index of mat - i2_loc = i2 - s_out[1] # local row index of matT - - # without last row in 3rd direction - for i3 in range(s_out[2], e_out[2]): # global row index of matT = global column index of mat - i3_loc = i3 - s_out[2] # local row index of matT - - for d1 in range(2*p_in[0] + add[0]): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2*p_in[1] + 1): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - for d3 in range(2*p_in[2] + 1): - j3 = i3 - p_in[2] + d3 # global column index of matT - j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_in[2] + j3_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2, - p_out[2] + i3 - j3] - - # treat last row in 3rd direction separately - i3 = e_out[2] - i3_loc = i3 - s_out[2] # local row index of matT - - for d1 in range(2*p_in[0] + add[0]): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2*p_in[1] + 1): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - for d3 in range(2*p_in[2] + add[2]): - j3 = i3 - p_in[2] + d3 # global column index of matT - j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_in[2] + j3_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2, - p_out[2] + i3 - j3] - - ############################################## - # treat last row in 2nd direction separately # - ############################################## - i2 = e_out[1] - i2_loc = i2 - s_out[1] # local row index of matT - - # without last row in 3rd direction - for i3 in range(s_out[2], e_out[2]): # global row index of matT = global column index of mat - i3_loc = i3 - s_out[2] # local row index of matT - - for d1 in range(2*p_in[0] + add[0]): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2*p_in[1] + add[1]): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - for d3 in range(2*p_in[2] + 1): - j3 = i3 - p_in[2] + d3 # global column index of matT - j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_in[2] + j3_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2, - p_out[2] + i3 - j3] - - # treat last row in 3rd direction separately - i3 = e_out[2] - i3_loc = i3 - s_out[2] # local row index of matT - - for d1 in range(2*p_in[0] + add[0]): - j1 = i1 - p_in[0] + d1 # global column index of matT - j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat - for d2 in range(2*p_in[1] + add[1]): - j2 = i2 - p_in[1] + d2 # global column index of matT - j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat - for d3 in range(2*p_in[2] + add[2]): - j3 = i3 - p_in[2] + d3 # global column index of matT - j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat - - matT[p_out[0] + i1_loc, - p_out[1] + i2_loc, - p_out[2] + i3_loc, - d1, d2, d3] = mat[p_in[0] + j1_loc, - p_in[1] + j2_loc, - p_in[2] + j3_loc, - p_out[0] + i1 - j1, - p_out[1] + i2 - j2, - p_out[2] + i3 - j3] + q1 + i1 - j1, + q2 + i2 - j2, + q3 + i3 - j3] diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index 64ca4c095..ee7092fc9 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -1131,7 +1131,7 @@ def dot(self, v, out=None): if not v.ghost_regions_in_sync: v.update_ghost_regions() - self._check_kernel_shapes() + self._check_kernel_shifts() # zeros, not empty: the kernel only writes the interior (non-padding) # region, so the padding must be initialized to avoid leaking stale @@ -1145,18 +1145,18 @@ def dot(self, v, out=None): return out # ... - def _check_kernel_shapes(self, *others): + def _check_kernel_shifts(self, *others): """ - Raise if the precompiled dot/transpose kernels cannot be used with the - data of this matrix (or of `others`), see `set_backend`. + Raise if the stencil kernels cannot be used with the data of this + matrix (or of `others`): they handle any number of diagonals (pads of + the matrix up to those of the spaces) and rectangular matrices, but + only spaces without shifts (shifts == 1), see `set_backend`. """ for M in (self, *others): - if not M._kernel_shapes_ok: + if not M._kernel_shifts_ok: raise NotImplementedError( - 'The precompiled stencil kernels need a matrix without shifts and with ' - '2 * p + 1 diagonals per direction (p: the pads of the domain), ' - f'got data of shape {M._data_shape} for the pads {M.domain.pads}, ' - f'domain shifts {M.domain.shifts} and codomain shifts {M.codomain.shifts}.') + 'The stencil kernels need spaces without shifts (shifts == 1), ' + f'got domain shifts {M.domain.shifts} and codomain shifts {M.codomain.shifts}.') # ... def vdot( self, v, out=None): @@ -1191,7 +1191,7 @@ def vdot( self, v, out=None): if not v.ghost_regions_in_sync: v.update_ghost_regions() - self._check_kernel_shapes() + self._check_kernel_shifts() # Instead of computing A_*x, this function computes (A*x_)_ # zeros, not empty: see comment in dot() above. @@ -1228,11 +1228,13 @@ def transpose(self, conjugate=False, out=None): assert isinstance(out, StencilMatrix) assert out.codomain == M.domain assert out.domain == M.codomain - + # the kernels map diagonal d of `out` to diagonal 2 q - d of M (q: the pads of the matrices) + assert tuple(out._pads) == tuple(M._pads) + else : out = StencilMatrix(M.codomain, M.domain, pads=self._pads, backend=self._backend, precompiled=self._precompiled) - M._check_kernel_shapes(out) + M._check_kernel_shifts(out) # Call low-level '_transpose' function: the kernel of the active backend # (a cunumpy Kernel, see set_backend) @@ -1993,7 +1995,7 @@ def set_backend(self, backend, precompiled): self._args = self._dotargs_null.copy() # The kernels are called with the values of self._args (and of # self._transpose_args) as positional arguments, in this order. - self._kernel_shapes_ok = True + self._kernel_shifts_ok = True if self._backend is None: for key, arg in self._args.items(): @@ -2005,13 +2007,13 @@ def set_backend(self, backend, precompiled): ('starts', 'nrows', 'nrows_extra', 'dm', 'cm', 'pad_imp', 'ndiags', 'gpads')} elif precompiled: - # The precompiled kernels read the diagonals 0 <= d < 2 * p + 1 of the - # matrix data and its rows at p + i_loc (no shifts); the CUDA versions - # take the shape of 3D matrix data (6 axes, a raw pointer) from that. + # The precompiled kernels take the number of diagonals of each + # direction from the shape of the matrix data (pads of the matrix up + # to those of the spaces), but assume spaces without shifts: their + # rows start at index p of the data and diagonal d of row i is the + # column i - q + d. Shifts > 1 raise in dot, vdot and transpose. shifts = (*self.domain.shifts, *self.codomain.shifts) - diags = tuple(2 * int(p) + 1 for p in self.domain.pads) - self._kernel_shapes_ok = (all(int(m) == 1 for m in shifts) - and tuple(self._data_shape[self._ndim:]) == diags) + self._kernel_shifts_ok = all(int(m) == 1 for m in shifts) # matvec kernel: stencil_dot_d, pyccel or CUDA by backend self._func = stencil_kernels['dot'][self._ndim] @@ -2044,7 +2046,6 @@ def set_backend(self, backend, precompiled): self._transpose_args = {} if self._ndim == 1: self._transpose_args['s_in'] = int(self.codomain.starts[0]) - self._transpose_args['e_in'] = int(self.codomain.ends[0]) self._transpose_args['p_in'] = int(self.codomain.pads[0]) self._transpose_args['add'] = int(add[0]) self._transpose_args['s_out'] = int(self.domain.starts[0]) @@ -2052,7 +2053,6 @@ def set_backend(self, backend, precompiled): self._transpose_args['p_out'] = int(self.domain.pads[0]) else: self._transpose_args['s_in'] = xp.array(self.codomain.starts) - self._transpose_args['e_in'] = xp.array(self.codomain.ends) self._transpose_args['p_in'] = xp.array(self.codomain.pads) self._transpose_args['add'] = xp.array(add) self._transpose_args['s_out'] = xp.array(self.domain.starts) diff --git a/feectools/linalg/tests/kernel_test_args.py b/feectools/linalg/tests/kernel_test_args.py index a6dd1d218..3fa3d6f8a 100644 --- a/feectools/linalg/tests/kernel_test_args.py +++ b/feectools/linalg/tests/kernel_test_args.py @@ -10,12 +10,28 @@ from feectools.ddm.cart import DomainDecomposition, CartDecomposition from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix -# (npts of the domain, npts of the codomain, pads): square matrices, and rectangular ones whose spaces -# differ by one point in a direction, which makes `add` zero (or one) there, as for derivative operators. +# (npts of the domain, npts of the codomain, pads of the spaces, pads of the matrix, periods): square matrices, and +# rectangular ones whose spaces differ by one point in a direction, which makes `add` zero (or one) there, as for +# derivative operators. A matrix whose pads are None has those of the spaces (2 * p + 1 diagonals), periods None +# means periodic in every direction. The cases after the first three have fewer diagonals than 2 * p + 1 in some +# direction (pads of the matrix smaller than those of the spaces) and/or non-periodic directions: blocks between +# spaces of different degree per direction, and derivative-type (one-sided, offset) stencils. MATRIX_CASES = { - 1: [((24,), (24,), (2,)), ((23,), (24,), (2,)), ((24,), (23,), (3,))], - 2: [((10, 12), (10, 12), (2, 3)), ((11, 10), (12, 10), (2, 2)), ((12, 9), (11, 10), (1, 2))], - 3: [((7, 8, 9), (7, 8, 9), (1, 2, 3)), ((8, 8, 9), (9, 8, 10), (1, 2, 2)), ((6, 5, 7), (5, 6, 7), (2, 1, 1))], + 1: [((24,), (24,), (2,), None, None), ((23,), (24,), (2,), None, None), ((24,), (23,), (3,), None, None), + ((24,), (24,), (3,), (1,), (False,)), ((13,), (12,), (2,), (1,), (False,)), ((12,), (13,), (3,), (0,), (False,))], + 2: [((10, 12), (10, 12), (2, 3), None, None), ((11, 10), (12, 10), (2, 2), None, None), + ((12, 9), (11, 10), (1, 2), None, None), + ((10, 9), (10, 9), (3, 2), (1, 2), (False, True)), ((9, 8), (8, 8), (2, 3), (2, 1), (False, False))], + 3: [((7, 8, 9), (7, 8, 9), (1, 2, 3), None, None), ((8, 8, 9), (9, 8, 10), (1, 2, 2), None, None), + ((6, 5, 7), (5, 6, 7), (2, 1, 1), None, None), + # square, fewer diagonals in two directions (e.g. a mass matrix of a space of lower degree) + ((6, 7, 8), (6, 7, 8), (2, 3, 2), (1, 1, 2), (True, False, True)), + # rectangular between non-periodic spaces of different size per direction, all diagonals + ((8, 7, 9), (7, 7, 8), (2, 2, 3), None, (False, True, False)), + # derivative-type: rectangular, non-periodic, fewer diagonals in every direction + ((7, 8, 6), (8, 7, 6), (3, 2, 2), (1, 1, 1), (False, False, True)), + # one diagonal in two directions (pads 0), the codomain larger in both + ((7, 6, 5), (8, 6, 6), (1, 2, 1), (0, 2, 0), (False, True, False))], } # (npts, pads) of the vector spaces @@ -26,10 +42,10 @@ } -def make_space(npts, pads, dtype=float): - """A serial StencilVectorSpace with periodic directions.""" +def make_space(npts, pads, dtype=float, periods=None): + """A serial StencilVectorSpace, periodic in every direction unless `periods` says otherwise.""" ndim = len(npts) - D = DomainDecomposition(list(npts), periods=[True] * ndim) + D = DomainDecomposition(list(npts), periods=list(periods) if periods is not None else [True] * ndim) global_starts, global_ends = [], [] for axis in range(ndim): ee = D.global_element_ends[axis].copy() @@ -46,11 +62,11 @@ def random_like(array, rng): return xp.asarray(rng.random(shape).astype(array.dtype)) -def stencil_matrix(npts_domain, npts_codomain, pads, rng): +def stencil_matrix(npts_domain, npts_codomain, pads, matrix_pads, periods, rng): """A StencilMatrix with random entries (spurious entries removed), its domain and its codomain.""" - V = make_space(npts_domain, pads) - W = V if npts_domain == npts_codomain else make_space(npts_codomain, pads) - A = StencilMatrix(V, W) + V = make_space(npts_domain, pads, periods=periods) + W = V if npts_domain == npts_codomain else make_space(npts_codomain, pads, periods=periods) + A = StencilMatrix(V, W, pads=matrix_pads) A._data[...] = random_like(A._data, rng) A.remove_spurious_entries() return A, V, W @@ -64,20 +80,21 @@ def stencil_vector(V, rng): return v -def dot_arguments(npts_domain, npts_codomain, pads, seed): +def dot_arguments(npts_domain, npts_codomain, pads, matrix_pads, periods, seed): """The arguments of ``stencil_dot_d`` as ``StencilMatrix.dot`` passes them.""" rng = np.random.default_rng(seed) - A, V, W = stencil_matrix(npts_domain, npts_codomain, pads, rng) + A, V, W = stencil_matrix(npts_domain, npts_codomain, pads, matrix_pads, periods, rng) v = stencil_vector(V, rng) out = StencilVector(W) return (A._data, v._data, out._data, *A._args.values()) -def transpose_arguments(npts_domain, npts_codomain, pads, seed): +def transpose_arguments(npts_domain, npts_codomain, pads, matrix_pads, periods, seed): """The arguments of ``stencil_transpose_d`` as ``StencilMatrix.transpose`` passes them.""" rng = np.random.default_rng(seed) - A, V, W = stencil_matrix(npts_domain, npts_codomain, pads, rng) - out = StencilMatrix(W, V) + A, V, W = stencil_matrix(npts_domain, npts_codomain, pads, matrix_pads, periods, rng) + A.update_ghost_regions() + out = StencilMatrix(W, V, pads=matrix_pads) return (A._data, out._data, *A._transpose_args.values()) diff --git a/feectools/linalg/tests/test_cuda_parity.py b/feectools/linalg/tests/test_cuda_parity.py index 7f49fdf6f..6032e92d4 100644 --- a/feectools/linalg/tests/test_cuda_parity.py +++ b/feectools/linalg/tests/test_cuda_parity.py @@ -89,7 +89,7 @@ def test_solver_loop_has_no_host_transfers(): with xp.use_backend("cupy"): rng = np.random.default_rng(0) - A, V, W = stencil_matrix((12, 10, 8), (12, 10, 8), (2, 2, 3), rng) + A, V, W = stencil_matrix((12, 10, 8), (12, 10, 8), (2, 2, 3), None, None, rng) x = stencil_vector(V, rng) y = A.dot(x) # warm up: compiles the CUDA kernels V.axpy(0.5, y, x) diff --git a/feectools/linalg/tests/test_device_matvec.py b/feectools/linalg/tests/test_device_matvec.py index 79771dd09..a216f3cd6 100644 --- a/feectools/linalg/tests/test_device_matvec.py +++ b/feectools/linalg/tests/test_device_matvec.py @@ -23,9 +23,9 @@ from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix, stencil_kernels # =============================================================================== -def make_space(npts, pads, dtype): +def make_space(npts, pads, dtype, periods=None, shifts=None): ndim = len(npts) - D = DomainDecomposition(list(npts), periods=[True] * ndim) + D = DomainDecomposition(list(npts), periods=list(periods) if periods else [True] * ndim) global_starts, global_ends = [], [] for axis in range(ndim): ee = D.global_element_ends[axis].copy() @@ -33,7 +33,7 @@ def make_space(npts, pads, dtype): global_ends.append(ee) global_starts.append(xp.array([0] + (ee[:-1] + 1).tolist())) C = CartDecomposition(D, list(npts), global_starts, global_ends, - pads=list(pads), shifts=[1] * ndim) + pads=list(pads), shifts=list(shifts) if shifts else [1] * ndim) return StencilVectorSpace(C, dtype=dtype) @@ -55,15 +55,18 @@ def reference_matvec(A, v, out): """ out = A @ v, as a sum of shifted elementwise products. - Interior rows use 2 * p_in + 1 diagonals along each direction, the last row - along a direction uses 2 * p_in + add, so every combination of - (interior, last) over the directions is accumulated separately. + The matrix has n_k diagonals along direction k (from its data, pads + q_k = (n_k - 1) // 2 up to the pads p_in of the domain); interior rows use + all of them, the last row along a direction uses n_k - 1 + add, so every + combination of (interior, last) over the directions is accumulated + separately. Diagonal d of row i is the column i - q + d. """ ndim = v.space.ndim a = dot_args(A, ndim) n = [e - s + 1 for s, e in zip(a['s_out'], a['e_out'])] - off = [so - si for so, si in zip(a['s_out'], a['s_in'])] p_out, p_in, add = a['p_out'], a['p_in'], a['add'] + n_diags = [int(n) for n in A._data.shape[ndim:]] + off = [so - si + pi - (nd - 1) // 2 for so, si, pi, nd in zip(a['s_out'], a['s_in'], p_in, n_diags)] out._data[...] = 0 @@ -79,7 +82,7 @@ def reference_matvec(A, v, out): nrow = [r.stop - r.start for r in rows] base = [r.start - p_out[k] for k, r in enumerate(rows)] - bounds = [2 * p_in[k] + (add[k] if last[k] else 1) for k in range(ndim)] + bounds = [n_diags[k] - 1 + add[k] if last[k] else n_diags[k] for k in range(ndim)] for d in itertools.product(*[range(b) for b in bounds]): src = tuple(slice(off[k] + base[k] + d[k], @@ -101,11 +104,11 @@ def fill_like(arr, seed): # =============================================================================== -def build(npts_domain, npts_codomain, pads, dtype): - V = make_space(npts_domain, pads, dtype) - W = V if npts_domain == npts_codomain else make_space(npts_codomain, pads, dtype) +def build(npts_domain, npts_codomain, pads, dtype, matrix_pads=None, periods=None): + V = make_space(npts_domain, pads, dtype, periods) + W = V if npts_domain == npts_codomain else make_space(npts_codomain, pads, dtype, periods) - A = StencilMatrix(V, W) + A = StencilMatrix(V, W, pads=matrix_pads) A._data[...] = fill_like(A._data, 1) A.remove_spurious_entries() @@ -119,6 +122,8 @@ def build(npts_domain, npts_codomain, pads, dtype): # =============================================================================== # Square matrices, and rectangular ones whose spaces differ by one point in a # direction -- the case that makes `add` zero there, as for derivative operators. +# The cases with matrix pads have fewer diagonals than 2 * p + 1 (blocks between +# spaces of different degree, derivative-type stencils), on non-periodic spaces. CASES = [ ('1d-square', (24,), (24,), (2,)), ('1d-rect', (23,), (24,), (2,)), @@ -129,6 +134,18 @@ def build(npts_domain, npts_codomain, pads, dtype): ('3d-rect', (8, 8, 9), (9, 8, 10), (1, 2, 2)), ] +# (name, npts of the domain, npts of the codomain, pads of the spaces, pads of the matrix, periods) +CASES_FEWER_DIAGONALS = [ + ('1d-fewer', (24,), (24,), (3,), (1,), (False,)), + ('1d-rect-fewer', (13,), (12,), (2,), (1,), (False,)), + ('2d-fewer', (10, 9), (10, 9), (3, 2), (1, 2), (False, True)), + ('2d-rect-fewer', (9, 8), (8, 8), (2, 3), (2, 1), (False, False)), + ('3d-fewer', (6, 7, 8), (6, 7, 8), (2, 3, 2), (1, 1, 2), (True, False, True)), + ('3d-rect-nonperiodic', (8, 7, 9), (7, 7, 8), (2, 2, 3), None, (False, True, False)), + ('3d-rect-fewer', (7, 8, 6), (8, 7, 6), (3, 2, 2), (1, 1, 1), (False, False, True)), + ('3d-rect-one-diagonal', (7, 6, 5), (8, 6, 6), (1, 2, 1), (0, 2, 0), (False, True, False)), +] + @pytest.mark.parametrize('name, npts_d, npts_c, pads', CASES, ids=[c[0] for c in CASES]) @@ -188,15 +205,41 @@ def test_dot_calls_the_folder_kernel(ndim): # =============================================================================== -def test_kernels_reject_matrices_with_other_diagonals(): - """A matrix with fewer diagonals than 2 * p + 1 (pads smaller than those of - the space) cannot use the precompiled kernels: they would read past the - diagonal axis, so dot and transpose raise instead.""" - V = make_space((8, 9), (2, 2), float) - A = StencilMatrix(V, V, pads=(1, 1)) - with pytest.raises(NotImplementedError, match='2 \\* p \\+ 1 diagonals'): +@pytest.mark.parametrize('name, npts_d, npts_c, pads, matrix_pads, periods', CASES_FEWER_DIAGONALS, + ids=[c[0] for c in CASES_FEWER_DIAGONALS]) +def test_fewer_diagonals_match_dense_reference(name, npts_d, npts_c, pads, matrix_pads, periods): + """Matrices the kernels used to reject (pads of the matrix smaller than + those of the spaces, i.e. fewer than 2 * p + 1 diagonals): dot, vdot and + transpose agree with the dense matrix, and dot with the shifted-view + reference.""" + V, W, A, v = build(npts_d, npts_c, pads, float, matrix_pads, periods) + dense = xp.to_numpy(A.toarray()) + x = xp.to_numpy(v.toarray()) + + got = A.dot(v) + assert np.allclose(xp.to_numpy(got.toarray()), dense @ x, rtol=0.0, atol=1e-12) + assert xp.allclose(got._data, reference_matvec(A, v, StencilVector(W))._data, rtol=0.0, atol=1e-12) + assert np.allclose(xp.to_numpy(A.vdot(v).toarray()), dense @ x, rtol=0.0, atol=1e-12) + + AT = A.transpose() + assert AT.pads == A.pads + assert np.array_equal(xp.to_numpy(AT.toarray()), dense.T) + # transpose into an existing matrix gives the same + out = StencilMatrix(W, V, pads=matrix_pads) + A.transpose(out=out) + assert np.array_equal(xp.to_numpy(out.toarray()), dense.T) + + +# =============================================================================== +def test_kernels_reject_spaces_with_shifts(): + """The stencil kernels assume spaces without shifts (shifts == 1); a matrix + between spaces with shifts raises in dot and transpose instead of giving + wrong results.""" + V = make_space((8, 9), (2, 2), float, shifts=(2, 1)) + A = StencilMatrix(V, V) + with pytest.raises(NotImplementedError, match='without shifts'): A.dot(StencilVector(V)) - with pytest.raises(NotImplementedError, match='2 \\* p \\+ 1 diagonals'): + with pytest.raises(NotImplementedError, match='without shifts'): A.transpose() diff --git a/feectools/linalg/tests/test_mpi_device.py b/feectools/linalg/tests/test_mpi_device.py index be243c4f1..70768c213 100644 --- a/feectools/linalg/tests/test_mpi_device.py +++ b/feectools/linalg/tests/test_mpi_device.py @@ -245,6 +245,43 @@ def test_inner_matches_global_reference(): assert abs(float(x.inner(x)) - float((glob * glob).sum())) <= 1e-9 +# =============================================================================== +def test_fewer_diagonals_dot_and_transpose_match_global_reference(): + """A matrix with fewer diagonals than its spaces' pads allow (pads (1, 1) + on a space with pads (1, 2)), which the stencil kernels used to reject: + `dot` equals the stencil applied to the global field, and the transpose of + a non-symmetric stencil is its adjoint, (A^T u, v) = (u, A v), checked + through global inner products so the check is decomposition-free.""" + comm = MPI.COMM_WORLD + V = make_space(NPTS, PADS, comm=comm) + glob = global_field(NPTS) + v = scatter(V, glob) + + A = StencilMatrix(V, V, pads=(1, 1)) + A[:, :, 0, 0] = 4.5 + for axis in range(2): + for shift in (-1, 1): + key = [slice(None)] * 2 + [0, 0] + key[2 + axis] = shift + A[tuple(key)] = -1.0 + A.remove_spurious_entries() + assert A._data.shape[2:] == (3, 3) + + w = A.dot(v) + ref = reference_apply(glob) + expected = float((ref * glob).sum()) + assert abs(float(w.inner(v)) - expected) <= 1e-9 * abs(expected) + + # non-symmetric: one-sided (derivative-type) entries + A[:, :, 1, 0] = -2.0 + A[:, :, 0, -1] = 0.5 + A[:, :, 1, 1] = 0.25 + u = scatter(V, np.flipud(glob).copy()) + lhs = float(A.transpose().dot(u).inner(v)) + rhs = float(u.inner(A.dot(v))) + assert abs(lhs - rhs) <= 1e-12 * abs(rhs) + + # =============================================================================== if __name__ == "__main__": import sys