xref: /petsc/src/mat/impls/aij/mpi/mpiviennacl/mpiaijviennacl.cxx (revision 72a35d6d63ea98b60d8ef4c5aaefda5ff933c0fe)
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. For good matrix
70    assembly performance the user should preallocate the matrix storage by setting
71    the parameter nz (or the array nnz).  By setting these parameters accurately,
72    performance during matrix assembly can be increased substantially.
73 
74    Collective
75 
76    Input Parameters:
77 +  comm - MPI communicator, set to `PETSC_COMM_SELF`
78 .  m - number of rows
79 .  n - number of columns
80 .  nz - number of nonzeros per row (same for all rows)
81 -  nnz - array containing the number of nonzeros in the various rows
82          (possibly different for each row) or NULL
83 
84    Output Parameter:
85 .  A - the matrix
86 
87    It is recommended that one use the `MatCreate()`, `MatSetType()` and/or `MatSetFromOptions()`,
88    MatXXXXSetPreallocation() paradigm instead of this routine directly.
89    [MatXXXXSetPreallocation() is, for example, `MatSeqAIJSetPreallocation()`]
90 
91    Notes:
92    If nnz is given then nz is ignored
93 
94    The AIJ format, also called
95    compressed row storage), is fully compatible with standard Fortran
96    storage.  That is, the stored row and column indices can begin at
97    either one (as in Fortran) or zero.  See the users' manual for details.
98 
99    Specify the preallocated storage with either nz or nnz (not both).
100    Set nz = `PETSC_DEFAULT` and nnz = NULL for PETSc to control dynamic memory
101    allocation.  For large problems you MUST preallocate memory or you
102    will get TERRIBLE performance, see the users' manual chapter on matrices.
103 
104    Level: intermediate
105 
106 .seealso: `MatCreate()`, `MatCreateAIJ()`, `MatCreateAIJCUSPARSE()`, `MatSetValues()`, `MatSeqAIJSetColumnIndices()`, `MatCreateSeqAIJWithArrays()`, `MatCreateAIJ()`, `MATMPIAIJVIENNACL`, `MATAIJVIENNACL`
107 @*/
108 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)
109 {
110   PetscMPIInt size;
111 
112   PetscFunctionBegin;
113   PetscCall(MatCreate(comm, A));
114   PetscCall(MatSetSizes(*A, m, n, M, N));
115   PetscCallMPI(MPI_Comm_size(comm, &size));
116   if (size > 1) {
117     PetscCall(MatSetType(*A, MATMPIAIJVIENNACL));
118     PetscCall(MatMPIAIJSetPreallocation(*A, d_nz, d_nnz, o_nz, o_nnz));
119   } else {
120     PetscCall(MatSetType(*A, MATSEQAIJVIENNACL));
121     PetscCall(MatSeqAIJSetPreallocation(*A, d_nz, d_nnz));
122   }
123   PetscFunctionReturn(PETSC_SUCCESS);
124 }
125 
126 /*MC
127    MATAIJVIENNACL - MATMPIAIJVIENNACL= "aijviennacl" = "mpiaijviennacl" - A matrix type to be used for sparse matrices.
128 
129    A matrix type (CSR format) whose data resides on GPUs.
130    All matrix calculations are performed using the ViennaCL library.
131 
132    This matrix type is identical to `MATSEQAIJVIENNACL` when constructed with a single process communicator,
133    and `MATMPIAIJVIENNACL` otherwise.  As a result, for single process communicators,
134    `MatSeqAIJSetPreallocation()` is supported, and similarly `MatMPIAIJSetPreallocation()` is supported
135    for communicators controlling multiple processes.  It is recommended that you call both of
136    the above preallocation routines for simplicity.
137 
138    Options Database Keys:
139 .  -mat_type mpiaijviennacl - sets the matrix type to `MATAIJVIENNACL` during a call to `MatSetFromOptions()`
140 
141   Level: beginner
142 
143  .seealso: `MatCreateAIJViennaCL()`, `MATSEQAIJVIENNACL`, `MatCreateSeqAIJVIENNACL()`
144 M*/
145