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: static 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: static 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: PETSC_EXTERN PetscErrorCode MatCreate_MPIAIJViennaCL(Mat A)
     47: {
     48:   PetscFunctionBegin;
     49:   PetscCall(MatCreate_MPIAIJ(A));
     50:   A->boundtocpu = PETSC_FALSE;
     51:   PetscCall(PetscFree(A->defaultvectype));
     52:   PetscCall(PetscStrallocpy(VECVIENNACL, &A->defaultvectype));
     53:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatMPIAIJSetPreallocation_C", MatMPIAIJSetPreallocation_MPIAIJViennaCL));
     54:   A->ops->assemblyend = MatAssemblyEnd_MPIAIJViennaCL;
     55:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATMPIAIJVIENNACL));
     56:   PetscFunctionReturn(PETSC_SUCCESS);
     57: }

     59: /*@C
     60:   MatCreateAIJViennaCL - Creates a sparse matrix in `MATAIJ` (compressed row) format
     61:   (the default parallel PETSc format).  This matrix will ultimately be pushed down
     62:   to GPUs and use the ViennaCL library for calculations.

     64:   Collective

     66:   Input Parameters:
     67: + comm  - MPI communicator, set to `PETSC_COMM_SELF`
     68: . m     - number of rows, or `PETSC_DECIDE` if `M` is provided
     69: . n     - number of columns, or `PETSC_DECIDE` if `N` is provided
     70: . M     - number of global rows in the matrix, or `PETSC_DETERMINE` if `m` is provided
     71: . N     - number of global columns in the matrix, or `PETSC_DETERMINE` if `n` is provided
     72: . d_nz  - number of nonzeros per row in DIAGONAL portion of local submatrix
     73:            (same value is used for all local rows)
     74: . d_nnz - array containing the number of nonzeros in the various rows of the
     75:            DIAGONAL portion of the local submatrix (possibly different for each row)
     76:            or `NULL`, if `d_nz` is used to specify the nonzero structure.
     77:            The size of this array is equal to the number of local rows, i.e `m`.
     78:            For matrices you plan to factor you must leave room for the diagonal entry and
     79:            put in the entry even if it is zero.
     80: . o_nz  - number of nonzeros per row in the OFF-DIAGONAL portion of local
     81:            submatrix (same value is used for all local rows).
     82: - o_nnz - array containing the number of nonzeros in the various rows of the
     83:            OFF-DIAGONAL portion of the local submatrix (possibly different for
     84:            each row) or `NULL`, if `o_nz` is used to specify the nonzero
     85:            structure. The size of this array is equal to the number
     86:            of local rows, i.e `m`.

     88:   Output Parameter:
     89: . A - the matrix

     91:   Level: intermediate

     93:   Notes:
     94:   It is recommended that one use the `MatCreate()`, `MatSetType()` and/or `MatSetFromOptions()`,
     95:   MatXXXXSetPreallocation() paradigm instead of this routine directly.
     96:   [MatXXXXSetPreallocation() is, for example, `MatSeqAIJSetPreallocation()`]

     98:   The AIJ format, also called
     99:   compressed row storage), is fully compatible with standard Fortran
    100:   storage.  That is, the stored row and column indices can begin at
    101:   either one (as in Fortran) or zero.

    103: .seealso: `Mat`, `MatCreate()`, `MatCreateAIJ()`, `MatCreateAIJCUSPARSE()`, `MatSetValues()`, `MatSeqAIJSetColumnIndices()`, `MatCreateSeqAIJWithArrays()`,
    104:           `MATMPIAIJVIENNACL`, `MATAIJVIENNACL`
    105: @*/
    106: 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)
    107: {
    108:   PetscMPIInt size;

    110:   PetscFunctionBegin;
    111:   PetscCall(MatCreate(comm, A));
    112:   PetscCall(MatSetSizes(*A, m, n, M, N));
    113:   PetscCallMPI(MPI_Comm_size(comm, &size));
    114:   if (size > 1) {
    115:     PetscCall(MatSetType(*A, MATMPIAIJVIENNACL));
    116:     PetscCall(MatMPIAIJSetPreallocation(*A, d_nz, d_nnz, o_nz, o_nnz));
    117:   } else {
    118:     PetscCall(MatSetType(*A, MATSEQAIJVIENNACL));
    119:     PetscCall(MatSeqAIJSetPreallocation(*A, d_nz, d_nnz));
    120:   }
    121:   PetscFunctionReturn(PETSC_SUCCESS);
    122: }

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

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

    130:    This matrix type is identical to `MATSEQAIJVIENNACL` when constructed with a single process communicator,
    131:    and `MATMPIAIJVIENNACL` otherwise.  As a result, for single process communicators,
    132:    `MatSeqAIJSetPreallocation()` is supported, and similarly `MatMPIAIJSetPreallocation()` is supported
    133:    for communicators controlling multiple processes.  It is recommended that you call both of
    134:    the above preallocation routines for simplicity.

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

    139:   Level: beginner

    141: .seealso: `Mat`, `MatType`, `MatCreateAIJViennaCL()`, `MATSEQAIJVIENNACL`, `MatCreateSeqAIJVIENNACL()`
    142: M*/