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 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; 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 static 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 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 } 58 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. 63 64 Collective 65 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`. 87 88 Output Parameter: 89 . A - the matrix 90 91 Level: intermediate 92 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()`] 97 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. 102 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; 109 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 } 123 124 /*MC 125 MATAIJVIENNACL - MATMPIAIJVIENNACL= "aijviennacl" = "mpiaijviennacl" - A matrix type to be used for sparse matrices. 126 127 A matrix type (CSR format) whose data resides on GPUs. 128 All matrix calculations are performed using the ViennaCL library. 129 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. 135 136 Options Database Keys: 137 . -mat_type mpiaijviennacl - sets the matrix type to `MATAIJVIENNACL` during a call to `MatSetFromOptions()` 138 139 Level: beginner 140 141 .seealso: `Mat`, `MatType`, `MatCreateAIJViennaCL()`, `MATSEQAIJVIENNACL`, `MatCreateSeqAIJVIENNACL()` 142 M*/ 143