xref: /petsc/src/mat/impls/aij/mpi/mpiviennacl/mpiaijviennacl.cxx (revision bcee047adeeb73090d7e36cc71e39fc287cdbb97)
1 #define PETSC_SKIP_IMMINTRIN_H_CUDAWORKAROUND 1
2 
3 #include <petscconf.h>
4 #include <../src/mat/impls/aij/mpi/mpiaij.h> /*I "petscmat.h" I*/
5 #include <../src/mat/impls/aij/seq/seqviennacl/viennaclmatimpl.h>
6 
7 PetscErrorCode MatMPIAIJSetPreallocation_MPIAIJViennaCL(Mat B, PetscInt d_nz, const PetscInt d_nnz[], PetscInt o_nz, const PetscInt o_nnz[])
8 {
9   Mat_MPIAIJ *b = (Mat_MPIAIJ *)B->data;
10 
11   PetscFunctionBegin;
12   PetscCall(PetscLayoutSetUp(B->rmap));
13   PetscCall(PetscLayoutSetUp(B->cmap));
14   if (!B->preallocated) {
15     /* Explicitly create the two MATSEQAIJVIENNACL matrices. */
16     PetscCall(MatCreate(PETSC_COMM_SELF, &b->A));
17     PetscCall(MatSetSizes(b->A, B->rmap->n, B->cmap->n, B->rmap->n, B->cmap->n));
18     PetscCall(MatSetType(b->A, MATSEQAIJVIENNACL));
19     PetscCall(MatCreate(PETSC_COMM_SELF, &b->B));
20     PetscCall(MatSetSizes(b->B, B->rmap->n, B->cmap->N, B->rmap->n, B->cmap->N));
21     PetscCall(MatSetType(b->B, MATSEQAIJVIENNACL));
22   }
23   PetscCall(MatSeqAIJSetPreallocation(b->A, d_nz, d_nnz));
24   PetscCall(MatSeqAIJSetPreallocation(b->B, o_nz, o_nnz));
25   B->preallocated = PETSC_TRUE;
26   PetscFunctionReturn(PETSC_SUCCESS);
27 }
28 
29 PetscErrorCode MatAssemblyEnd_MPIAIJViennaCL(Mat A, MatAssemblyType mode)
30 {
31   Mat_MPIAIJ *b = (Mat_MPIAIJ *)A->data;
32   PetscBool   v;
33 
34   PetscFunctionBegin;
35   PetscCall(MatAssemblyEnd_MPIAIJ(A, mode));
36   PetscCall(PetscObjectTypeCompare((PetscObject)b->lvec, VECSEQVIENNACL, &v));
37   if (!v) {
38     PetscInt m;
39     PetscCall(VecGetSize(b->lvec, &m));
40     PetscCall(VecDestroy(&b->lvec));
41     PetscCall(VecCreateSeqViennaCL(PETSC_COMM_SELF, m, &b->lvec));
42   }
43   PetscFunctionReturn(PETSC_SUCCESS);
44 }
45 
46 PetscErrorCode MatDestroy_MPIAIJViennaCL(Mat A)
47 {
48   PetscFunctionBegin;
49   PetscCall(MatDestroy_MPIAIJ(A));
50   PetscFunctionReturn(PETSC_SUCCESS);
51 }
52 
53 PETSC_EXTERN PetscErrorCode MatCreate_MPIAIJViennaCL(Mat A)
54 {
55   PetscFunctionBegin;
56   PetscCall(MatCreate_MPIAIJ(A));
57   A->boundtocpu = PETSC_FALSE;
58   PetscCall(PetscFree(A->defaultvectype));
59   PetscCall(PetscStrallocpy(VECVIENNACL, &A->defaultvectype));
60   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatMPIAIJSetPreallocation_C", MatMPIAIJSetPreallocation_MPIAIJViennaCL));
61   A->ops->assemblyend = MatAssemblyEnd_MPIAIJViennaCL;
62   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATMPIAIJVIENNACL));
63   PetscFunctionReturn(PETSC_SUCCESS);
64 }
65 
66 /*@C
67    MatCreateAIJViennaCL - Creates a sparse matrix in `MATAIJ` (compressed row) format
68    (the default parallel PETSc format).  This matrix will ultimately be pushed down
69    to GPUs and use the ViennaCL library for calculations.
70 
71    Collective
72 
73    Input Parameters:
74 +  comm - MPI communicator, set to `PETSC_COMM_SELF`
75 .  m - number of rows, or `PETSC_DECIDE` if `M` is provided
76 .  n - number of columns, or `PETSC_DECIDE` if `N` is provided
77 .  M - number of global rows in the matrix, or `PETSC_DETERMINE` if `m` is provided
78 .  N - number of global columns in the matrix, or `PETSC_DETERMINE` if `n` is provided
79 .  d_nz  - number of nonzeros per row in DIAGONAL portion of local submatrix
80            (same value is used for all local rows)
81 .  d_nnz - array containing the number of nonzeros in the various rows of the
82            DIAGONAL portion of the local submatrix (possibly different for each row)
83            or `NULL`, if `d_nz` is used to specify the nonzero structure.
84            The size of this array is equal to the number of local rows, i.e `m`.
85            For matrices you plan to factor you must leave room for the diagonal entry and
86            put in the entry even if it is zero.
87 .  o_nz  - number of nonzeros per row in the OFF-DIAGONAL portion of local
88            submatrix (same value is used for all local rows).
89 -  o_nnz - array containing the number of nonzeros in the various rows of the
90            OFF-DIAGONAL portion of the local submatrix (possibly different for
91            each row) or `NULL`, if `o_nz` is used to specify the nonzero
92            structure. The size of this array is equal to the number
93            of local rows, i.e `m`.
94 
95    Output Parameter:
96 .  A - the matrix
97 
98    Level: intermediate
99 
100    Notes:
101    It is recommended that one use the `MatCreate()`, `MatSetType()` and/or `MatSetFromOptions()`,
102    MatXXXXSetPreallocation() paradigm instead of this routine directly.
103    [MatXXXXSetPreallocation() is, for example, `MatSeqAIJSetPreallocation()`]
104 
105    The AIJ format, also called
106    compressed row storage), is fully compatible with standard Fortran
107    storage.  That is, the stored row and column indices can begin at
108    either one (as in Fortran) or zero.
109 
110 .seealso: `Mat`, `MatCreate()`, `MatCreateAIJ()`, `MatCreateAIJCUSPARSE()`, `MatSetValues()`, `MatSeqAIJSetColumnIndices()`, `MatCreateSeqAIJWithArrays()`,
111           `MatCreateAIJ()`, `MATMPIAIJVIENNACL`, `MATAIJVIENNACL`
112 @*/
113 PetscErrorCode MatCreateAIJViennaCL(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscInt d_nz, const PetscInt d_nnz[], PetscInt o_nz, const PetscInt o_nnz[], Mat *A)
114 {
115   PetscMPIInt size;
116 
117   PetscFunctionBegin;
118   PetscCall(MatCreate(comm, A));
119   PetscCall(MatSetSizes(*A, m, n, M, N));
120   PetscCallMPI(MPI_Comm_size(comm, &size));
121   if (size > 1) {
122     PetscCall(MatSetType(*A, MATMPIAIJVIENNACL));
123     PetscCall(MatMPIAIJSetPreallocation(*A, d_nz, d_nnz, o_nz, o_nnz));
124   } else {
125     PetscCall(MatSetType(*A, MATSEQAIJVIENNACL));
126     PetscCall(MatSeqAIJSetPreallocation(*A, d_nz, d_nnz));
127   }
128   PetscFunctionReturn(PETSC_SUCCESS);
129 }
130 
131 /*MC
132    MATAIJVIENNACL - MATMPIAIJVIENNACL= "aijviennacl" = "mpiaijviennacl" - A matrix type to be used for sparse matrices.
133 
134    A matrix type (CSR format) whose data resides on GPUs.
135    All matrix calculations are performed using the ViennaCL library.
136 
137    This matrix type is identical to `MATSEQAIJVIENNACL` when constructed with a single process communicator,
138    and `MATMPIAIJVIENNACL` otherwise.  As a result, for single process communicators,
139    `MatSeqAIJSetPreallocation()` is supported, and similarly `MatMPIAIJSetPreallocation()` is supported
140    for communicators controlling multiple processes.  It is recommended that you call both of
141    the above preallocation routines for simplicity.
142 
143    Options Database Keys:
144 .  -mat_type mpiaijviennacl - sets the matrix type to `MATAIJVIENNACL` during a call to `MatSetFromOptions()`
145 
146   Level: beginner
147 
148 .seealso: `Mat`, `MatType`, `MatCreateAIJViennaCL()`, `MATSEQAIJVIENNACL`, `MatCreateSeqAIJVIENNACL()`
149 M*/
150