Actual source code: mpiaijviennacl.cxx

  1: #define PETSC_SKIP_IMMINTRIN_H_CUDAWORKAROUND 1

  3: #include <petscconf.h>
  4: #include <../src/mat/impls/aij/mpi/mpiaij.h>
  5: #include <../src/mat/impls/aij/seq/seqviennacl/viennaclmatimpl.h>

  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;

 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: }

 29: PetscErrorCode MatAssemblyEnd_MPIAIJViennaCL(Mat A, MatAssemblyType mode)
 30: {
 31:   Mat_MPIAIJ *b = (Mat_MPIAIJ *)A->data;
 32:   PetscBool   v;

 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: }

 46: PetscErrorCode MatDestroy_MPIAIJViennaCL(Mat A)
 47: {
 48:   PetscFunctionBegin;
 49:   PetscCall(MatDestroy_MPIAIJ(A));
 50:   PetscFunctionReturn(PETSC_SUCCESS);
 51: }

 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: }

 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.

 71:    Collective

 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`.

 95:    Output Parameter:
 96: .  A - the matrix

 98:    Level: intermediate

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()`]

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.

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;

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: }

131: /*MC
132:    MATAIJVIENNACL - MATMPIAIJVIENNACL= "aijviennacl" = "mpiaijviennacl" - A matrix type to be used for sparse matrices.

134:    A matrix type (CSR format) whose data resides on GPUs.
135:    All matrix calculations are performed using the ViennaCL library.

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.

143:    Options Database Keys:
144: .  -mat_type mpiaijviennacl - sets the matrix type to `MATAIJVIENNACL` during a call to `MatSetFromOptions()`

146:   Level: beginner

148: .seealso: `Mat`, `MatType`, `MatCreateAIJViennaCL()`, `MATSEQAIJVIENNACL`, `MatCreateSeqAIJVIENNACL()`
149: M*/