1 | /*!\file: MatPartition.cpp
|
---|
2 | * \brief partition matrix according to node sets
|
---|
3 | */
|
---|
4 |
|
---|
5 | #ifdef HAVE_CONFIG_H
|
---|
6 | #include <config.h>
|
---|
7 | #else
|
---|
8 | #error "Cannot compile with HAVE_CONFIG_H symbol! run configure first!"
|
---|
9 | #endif
|
---|
10 |
|
---|
11 | /*Petsc includes: */
|
---|
12 | #include "petscmat.h"
|
---|
13 | #include "petscvec.h"
|
---|
14 | #include "petscksp.h"
|
---|
15 |
|
---|
16 | #include "../../../shared/shared.h"
|
---|
17 | #include "../../mpi/patches/mpipatches.h"
|
---|
18 |
|
---|
19 | int MatPartition(Mat* poutmatrix,Mat matrixA,double* row_partition_vector,int row_partition_vector_size ,
|
---|
20 | double* col_partition_vector,int col_partition_vector_size){
|
---|
21 |
|
---|
22 | int i;
|
---|
23 |
|
---|
24 | /*Petsc matrix*/
|
---|
25 | int d_nz;
|
---|
26 | int o_nz;
|
---|
27 | int* node_rows=NULL;
|
---|
28 | int* node_cols=NULL;
|
---|
29 | int count;
|
---|
30 | IS col_index=NULL;
|
---|
31 | IS row_index=NULL;
|
---|
32 | int lower_row,upper_row,range;
|
---|
33 | int MA,NA; //matrixA dimensions
|
---|
34 | const char* type=NULL;
|
---|
35 | int csize;
|
---|
36 |
|
---|
37 | /*output*/
|
---|
38 | Mat outmatrix=NULL;
|
---|
39 |
|
---|
40 | /*get input matrix size: */
|
---|
41 | MatGetSize(matrixA,&MA,&NA);
|
---|
42 |
|
---|
43 | /*If one of the partitioning row or col vectors has a 0 dimension, we return a NILL matrix. Return it with the same type as matrixA.*/
|
---|
44 |
|
---|
45 | if ((row_partition_vector_size==0) || (col_partition_vector==0)){
|
---|
46 | MatGetType(matrixA,&type);
|
---|
47 | if (strcmp(type,"mpiaij")==0){
|
---|
48 | d_nz=0;
|
---|
49 | o_nz=0;
|
---|
50 | #if _PETSC_MAJOR_ == 3 && _PETSC_MINOR_ > 2
|
---|
51 | MatCreateAIJ(MPI_COMM_WORLD,PETSC_DETERMINE,PETSC_DETERMINE, 0,0,d_nz,PETSC_NULL,o_nz,PETSC_NULL,&outmatrix);
|
---|
52 | #else
|
---|
53 | MatCreateMPIAIJ(MPI_COMM_WORLD,PETSC_DETERMINE,PETSC_DETERMINE, 0,0,d_nz,PETSC_NULL,o_nz,PETSC_NULL,&outmatrix);
|
---|
54 | #endif
|
---|
55 | }
|
---|
56 | else if (strcmp(type,"mpidense")==0){
|
---|
57 | #if _PETSC_MAJOR_ == 3 && _PETSC_MINOR_ > 2
|
---|
58 | MatCreateDense(MPI_COMM_WORLD,PETSC_DETERMINE,PETSC_DETERMINE, 0,0,PETSC_NULL,&outmatrix);
|
---|
59 | #else
|
---|
60 | MatCreateMPIDense(MPI_COMM_WORLD,PETSC_DETERMINE,PETSC_DETERMINE, 0,0,PETSC_NULL,&outmatrix);
|
---|
61 | #endif
|
---|
62 | }
|
---|
63 | else{
|
---|
64 | _error2_("MatType " << type << " not supported yet");
|
---|
65 | }
|
---|
66 | /*Assemble*/
|
---|
67 | MatAssemblyBegin(outmatrix,MAT_FINAL_ASSEMBLY);
|
---|
68 | MatAssemblyEnd(outmatrix,MAT_FINAL_ASSEMBLY);
|
---|
69 | }
|
---|
70 | else{
|
---|
71 | /*Both vectors are non nill, use MatGetSubMatrix to condense out*/
|
---|
72 | /*Figure out which rows each node is going to get from matrix A.*/
|
---|
73 | MatGetOwnershipRange(matrixA,&lower_row,&upper_row);
|
---|
74 | upper_row--;
|
---|
75 | range=upper_row-lower_row+1;
|
---|
76 |
|
---|
77 | count=0;
|
---|
78 | if (range){
|
---|
79 | node_rows=xNew<int>(range); //this is the maximum number of rows one node can extract.
|
---|
80 |
|
---|
81 | for (i=0;i<row_partition_vector_size;i++){
|
---|
82 | if ( ((int)(*(row_partition_vector+i))>=(lower_row+1)) && ((int)(*(row_partition_vector+i))<=(upper_row+1)) ){
|
---|
83 | *(node_rows+count)=(int)*(row_partition_vector+i)-1;
|
---|
84 | count++;
|
---|
85 | }
|
---|
86 | }
|
---|
87 | }
|
---|
88 | else{
|
---|
89 | count=0;
|
---|
90 | }
|
---|
91 |
|
---|
92 | /*Now each node has a node_rows vectors holding which rows they should extract from matrixA. Create an Index Set from node_rows.*/
|
---|
93 | #if _PETSC_MAJOR_ < 3 || (_PETSC_MAJOR_ == 3 && _PETSC_MINOR_ < 2)
|
---|
94 | ISCreateGeneral(MPI_COMM_WORLD,count,node_rows,&row_index);
|
---|
95 | #else
|
---|
96 | ISCreateGeneral(MPI_COMM_WORLD,count,node_rows,PETSC_COPY_VALUES,&row_index);
|
---|
97 | #endif
|
---|
98 |
|
---|
99 | /*Same deal for columns*/
|
---|
100 | node_cols=xNew<int>(col_partition_vector_size);
|
---|
101 | for (i=0;i<col_partition_vector_size;i++){
|
---|
102 | *(node_cols+i)=(int)*(col_partition_vector+i)-1;
|
---|
103 | }
|
---|
104 | #if _PETSC_MAJOR_ < 3 || (_PETSC_MAJOR_ == 3 && _PETSC_MINOR_ < 2)
|
---|
105 | ISCreateGeneral(MPI_COMM_WORLD,col_partition_vector_size,node_cols,&col_index);
|
---|
106 | #else
|
---|
107 | ISCreateGeneral(MPI_COMM_WORLD,col_partition_vector_size,node_cols,PETSC_COPY_VALUES,&col_index);
|
---|
108 | #endif
|
---|
109 |
|
---|
110 | /*Call MatGetSubMatrix*/
|
---|
111 | csize=DetermineLocalSize(col_partition_vector_size);
|
---|
112 | if(col_partition_vector_size==row_partition_vector_size){
|
---|
113 | #if _PETSC_MAJOR_ >= 3
|
---|
114 | MatGetSubMatrix(matrixA,row_index,col_index,MAT_INITIAL_MATRIX,&outmatrix);
|
---|
115 | #else
|
---|
116 | MatGetSubMatrix(matrixA,row_index,col_index,count,MAT_INITIAL_MATRIX,&outmatrix);
|
---|
117 | #endif
|
---|
118 | }
|
---|
119 | else{
|
---|
120 | #if _PETSC_MAJOR_ >= 3
|
---|
121 | MatGetSubMatrix(matrixA,row_index,col_index,MAT_INITIAL_MATRIX,&outmatrix);
|
---|
122 | #else
|
---|
123 | MatGetSubMatrix(matrixA,row_index,col_index,csize,MAT_INITIAL_MATRIX,&outmatrix);
|
---|
124 | #endif
|
---|
125 | }
|
---|
126 |
|
---|
127 | }
|
---|
128 |
|
---|
129 | /*Free ressources:*/
|
---|
130 | xDelete<int>(node_rows);
|
---|
131 | xDelete<int>(node_cols);
|
---|
132 | ISFree(&col_index);
|
---|
133 | ISFree(&row_index);
|
---|
134 |
|
---|
135 | /*Assign output pointers:*/
|
---|
136 | *poutmatrix=outmatrix;
|
---|
137 | return 1;
|
---|
138 | }
|
---|