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:   PetscLayoutSetUp(B->rmap);
 12:   PetscLayoutSetUp(B->cmap);
 13:   if (!B->preallocated) {
 14:     /* Explicitly create the two MATSEQAIJVIENNACL matrices. */
 15:     MatCreate(PETSC_COMM_SELF, &b->A);
 16:     MatSetSizes(b->A, B->rmap->n, B->cmap->n, B->rmap->n, B->cmap->n);
 17:     MatSetType(b->A, MATSEQAIJVIENNACL);
 18:     MatCreate(PETSC_COMM_SELF, &b->B);
 19:     MatSetSizes(b->B, B->rmap->n, B->cmap->N, B->rmap->n, B->cmap->N);
 20:     MatSetType(b->B, MATSEQAIJVIENNACL);
 21:   }
 22:   MatSeqAIJSetPreallocation(b->A, d_nz, d_nnz);
 23:   MatSeqAIJSetPreallocation(b->B, o_nz, o_nnz);
 24:   B->preallocated = PETSC_TRUE;
 25:   return 0;
 26: }

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

 33:   MatAssemblyEnd_MPIAIJ(A, mode);
 34:   PetscObjectTypeCompare((PetscObject)b->lvec, VECSEQVIENNACL, &v);
 35:   if (!v) {
 36:     PetscInt m;
 37:     VecGetSize(b->lvec, &m);
 38:     VecDestroy(&b->lvec);
 39:     VecCreateSeqViennaCL(PETSC_COMM_SELF, m, &b->lvec);
 40:   }
 41:   return 0;
 42: }

 44: PetscErrorCode MatDestroy_MPIAIJViennaCL(Mat A)
 45: {
 46:   MatDestroy_MPIAIJ(A);
 47:   return 0;
 48: }

 50: PETSC_EXTERN PetscErrorCode MatCreate_MPIAIJViennaCL(Mat A)
 51: {
 52:   MatCreate_MPIAIJ(A);
 53:   A->boundtocpu = PETSC_FALSE;
 54:   PetscFree(A->defaultvectype);
 55:   PetscStrallocpy(VECVIENNACL, &A->defaultvectype);
 56:   PetscObjectComposeFunction((PetscObject)A, "MatMPIAIJSetPreallocation_C", MatMPIAIJSetPreallocation_MPIAIJViennaCL);
 57:   A->ops->assemblyend = MatAssemblyEnd_MPIAIJViennaCL;
 58:   PetscObjectChangeTypeName((PetscObject)A, MATMPIAIJVIENNACL);
 59:   return 0;
 60: }

 62: /*@C
 63:    MatCreateAIJViennaCL - Creates a sparse matrix in `MATAIJ` (compressed row) format
 64:    (the default parallel PETSc format).  This matrix will ultimately be pushed down
 65:    to GPUs and use the ViennaCL library for calculations. For good matrix
 66:    assembly performance the user should preallocate the matrix storage by setting
 67:    the parameter nz (or the array nnz).  By setting these parameters accurately,
 68:    performance during matrix assembly can be increased substantially.

 70:    Collective

 72:    Input Parameters:
 73: +  comm - MPI communicator, set to `PETSC_COMM_SELF`
 74: .  m - number of rows
 75: .  n - number of columns
 76: .  nz - number of nonzeros per row (same for all rows)
 77: -  nnz - array containing the number of nonzeros in the various rows
 78:          (possibly different for each row) or NULL

 80:    Output Parameter:
 81: .  A - the matrix

 83:    It is recommended that one use the `MatCreate()`, `MatSetType()` and/or `MatSetFromOptions()`,
 84:    MatXXXXSetPreallocation() paradigm instead of this routine directly.
 85:    [MatXXXXSetPreallocation() is, for example, `MatSeqAIJSetPreallocation()`]

 87:    Notes:
 88:    If nnz is given then nz is ignored

 90:    The AIJ format, also called
 91:    compressed row storage), is fully compatible with standard Fortran 77
 92:    storage.  That is, the stored row and column indices can begin at
 93:    either one (as in Fortran) or zero.  See the users' manual for details.

 95:    Specify the preallocated storage with either nz or nnz (not both).
 96:    Set nz = `PETSC_DEFAULT` and nnz = NULL for PETSc to control dynamic memory
 97:    allocation.  For large problems you MUST preallocate memory or you
 98:    will get TERRIBLE performance, see the users' manual chapter on matrices.

100:    Level: intermediate

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

108:   MatCreate(comm, A);
109:   MatSetSizes(*A, m, n, M, N);
110:   MPI_Comm_size(comm, &size);
111:   if (size > 1) {
112:     MatSetType(*A, MATMPIAIJVIENNACL);
113:     MatMPIAIJSetPreallocation(*A, d_nz, d_nnz, o_nz, o_nnz);
114:   } else {
115:     MatSetType(*A, MATSEQAIJVIENNACL);
116:     MatSeqAIJSetPreallocation(*A, d_nz, d_nnz);
117:   }
118:   return 0;
119: }

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

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

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

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

136:   Level: beginner

138:  .seealso: `MatCreateAIJViennaCL()`, `MATSEQAIJVIENNACL`, `MatCreateSeqAIJVIENNACL()`
139: M*/