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. 70 71 Collective 72 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`. 94 95 Output Parameter: 96 . A - the matrix 97 98 Level: intermediate 99 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()`] 104 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. 109 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; 116 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 } 130 131 /*MC 132 MATAIJVIENNACL - MATMPIAIJVIENNACL= "aijviennacl" = "mpiaijviennacl" - A matrix type to be used for sparse matrices. 133 134 A matrix type (CSR format) whose data resides on GPUs. 135 All matrix calculations are performed using the ViennaCL library. 136 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. 142 143 Options Database Keys: 144 . -mat_type mpiaijviennacl - sets the matrix type to `MATAIJVIENNACL` during a call to `MatSetFromOptions()` 145 146 Level: beginner 147 148 .seealso: `Mat`, `MatType`, `MatCreateAIJViennaCL()`, `MATSEQAIJVIENNACL`, `MatCreateSeqAIJVIENNACL()` 149 M*/ 150