Actual source code: mhypre_kernels.hpp
1: #pragma once
3: #include <../src/mat/impls/hypre/mhypre.h>
5: // Zero the specified n rows in rows[] of the hypre CSRMatrix (i, j, a) and replace the diagonal entry with diag
6: __global__ static void ZeroRows(PetscInt n, const PetscInt rows[], const HYPRE_Int i[], const HYPRE_Int j[], HYPRE_Complex a[], HYPRE_Complex diag)
7: {
8: PetscInt k = blockDim.x * blockIdx.x + threadIdx.x; // k-th entry in rows[]
9: PetscInt c = blockDim.y * blockIdx.y + threadIdx.y; // c-th nonzero in row rows[k]
10: PetscInt gridx = gridDim.x * blockDim.x;
11: PetscInt gridy = gridDim.y * blockDim.y;
12: for (; k < n; k += gridx) {
13: PetscInt r = rows[k]; // r-th row of the matrix
14: PetscInt nz = i[r + 1] - i[r];
15: for (; c < nz; c += gridy) {
16: if (r == j[i[r] + c]) a[i[r] + c] = diag;
17: else a[i[r] + c] = 0.0;
18: }
19: }
20: }