xref: /petsc/src/mat/impls/localref/mlocalref.c (revision d244ea2e0b0868d12f8d635edc743b2190e533e6) ! 
1 
2 #include <petsc/private/matimpl.h>          /*I "petscmat.h" I*/
3 
4 typedef struct {
5   Mat Top;
6   PetscBool rowisblock;
7   PetscBool colisblock;
8   PetscErrorCode (*SetValues)(Mat,PetscInt,const PetscInt[],PetscInt,const PetscInt[],const PetscScalar[],InsertMode);
9   PetscErrorCode (*SetValuesBlocked)(Mat,PetscInt,const PetscInt[],PetscInt,const PetscInt[],const PetscScalar[],InsertMode);
10 } Mat_LocalRef;
11 
12 /* These need to be macros because they use sizeof */
13 #define IndexSpaceGet(buf,nrow,ncol,irowm,icolm) do {                   \
14     if (nrow + ncol > (PetscInt)(sizeof(buf)/sizeof(buf[0]))) {         \
15       ierr = PetscMalloc2(nrow,&irowm,ncol,&icolm);CHKERRQ(ierr); \
16     } else {                                                            \
17       irowm = &buf[0];                                                  \
18       icolm = &buf[nrow];                                               \
19     }                                                                   \
20 } while (0)
21 
22 #define IndexSpaceRestore(buf,nrow,ncol,irowm,icolm) do {       \
23     if (nrow + ncol > (PetscInt)(sizeof(buf)/sizeof(buf[0]))) { \
24       ierr = PetscFree2(irowm,icolm);CHKERRQ(ierr);             \
25     }                                                           \
26 } while (0)
27 
28 static void BlockIndicesExpand(PetscInt n,const PetscInt idx[],PetscInt bs,PetscInt idxm[])
29 {
30   PetscInt i,j;
31   for (i=0; i<n; i++) {
32     for (j=0; j<bs; j++) {
33       idxm[i*bs+j] = idx[i]*bs + j;
34     }
35   }
36 }
37 
38 #undef __FUNCT__
39 #define __FUNCT__ "MatSetValuesBlockedLocal_LocalRef_Block"
40 static PetscErrorCode MatSetValuesBlockedLocal_LocalRef_Block(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv)
41 {
42   Mat_LocalRef   *lr = (Mat_LocalRef*)A->data;
43   PetscErrorCode ierr;
44   PetscInt       buf[4096],*irowm=NULL,*icolm; /* suppress maybe-uninitialized warning */
45 
46   PetscFunctionBegin;
47   if (!nrow || !ncol) PetscFunctionReturn(0);
48   IndexSpaceGet(buf,nrow,ncol,irowm,icolm);
49   ierr = ISLocalToGlobalMappingApplyBlock(A->rmap->mapping,nrow,irow,irowm);CHKERRQ(ierr);
50   ierr = ISLocalToGlobalMappingApplyBlock(A->cmap->mapping,ncol,icol,icolm);CHKERRQ(ierr);
51   ierr = (*lr->SetValuesBlocked)(lr->Top,nrow,irowm,ncol,icolm,y,addv);CHKERRQ(ierr);
52   IndexSpaceRestore(buf,nrow,ncol,irowm,icolm);
53   PetscFunctionReturn(0);
54 }
55 
56 #undef __FUNCT__
57 #define __FUNCT__ "MatSetValuesBlockedLocal_LocalRef_Scalar"
58 static PetscErrorCode MatSetValuesBlockedLocal_LocalRef_Scalar(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv)
59 {
60   Mat_LocalRef   *lr = (Mat_LocalRef*)A->data;
61   PetscErrorCode ierr;
62   PetscInt       rbs,cbs,buf[4096],*irowm,*icolm;
63 
64   PetscFunctionBegin;
65   ierr = MatGetBlockSizes(A,&rbs,&cbs);CHKERRQ(ierr);
66   IndexSpaceGet(buf,nrow*rbs,ncol*cbs,irowm,icolm);
67   BlockIndicesExpand(nrow,irow,rbs,irowm);
68   BlockIndicesExpand(ncol,icol,cbs,icolm);
69   ierr = ISLocalToGlobalMappingApplyBlock(A->rmap->mapping,nrow*rbs,irowm,irowm);CHKERRQ(ierr);
70   ierr = ISLocalToGlobalMappingApplyBlock(A->cmap->mapping,ncol*cbs,icolm,icolm);CHKERRQ(ierr);
71   ierr = (*lr->SetValues)(lr->Top,nrow*rbs,irowm,ncol*cbs,icolm,y,addv);CHKERRQ(ierr);
72   IndexSpaceRestore(buf,nrow*rbs,ncol*cbs,irowm,icolm);
73   PetscFunctionReturn(0);
74 }
75 
76 #undef __FUNCT__
77 #define __FUNCT__ "MatSetValuesLocal_LocalRef_Scalar"
78 static PetscErrorCode MatSetValuesLocal_LocalRef_Scalar(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv)
79 {
80   Mat_LocalRef   *lr = (Mat_LocalRef*)A->data;
81   PetscErrorCode ierr;
82   PetscInt       buf[4096],*irowm,*icolm;
83 
84   PetscFunctionBegin;
85   IndexSpaceGet(buf,nrow,ncol,irowm,icolm);
86   /* If the row IS defining this submatrix was an ISBLOCK, then the unblocked LGMapApply is the right one to use.  If
87    * instead it was (say) an ISSTRIDE with a block size > 1, then we need to use LGMapApplyBlock */
88   if (lr->rowisblock) {
89     ierr = ISLocalToGlobalMappingApply(A->rmap->mapping,nrow,irow,irowm);CHKERRQ(ierr);
90   } else {
91     ierr = ISLocalToGlobalMappingApplyBlock(A->rmap->mapping,nrow,irow,irowm);CHKERRQ(ierr);
92   }
93   /* As above, but for the column IS. */
94   if (lr->colisblock) {
95     ierr = ISLocalToGlobalMappingApply(A->cmap->mapping,ncol,icol,icolm);CHKERRQ(ierr);
96   } else {
97     ierr = ISLocalToGlobalMappingApplyBlock(A->cmap->mapping,ncol,icol,icolm);CHKERRQ(ierr);
98   }
99   ierr = (*lr->SetValues)(lr->Top,nrow,irowm,ncol,icolm,y,addv);CHKERRQ(ierr);
100   IndexSpaceRestore(buf,nrow,ncol,irowm,icolm);
101   PetscFunctionReturn(0);
102 }
103 
104 #undef __FUNCT__
105 #define __FUNCT__ "ISL2GCompose"
106 /* Compose an IS with an ISLocalToGlobalMapping to map from IS source indices to global indices */
107 static PetscErrorCode ISL2GCompose(IS is,ISLocalToGlobalMapping ltog,ISLocalToGlobalMapping *cltog)
108 {
109   PetscErrorCode ierr;
110   const PetscInt *idx;
111   PetscInt       m,*idxm;
112   PetscInt       bs;
113   PetscBool      isblock;
114 
115   PetscFunctionBegin;
116   PetscValidHeaderSpecific(is,IS_CLASSID,1);
117   PetscValidHeaderSpecific(ltog,IS_LTOGM_CLASSID,2);
118   PetscValidPointer(cltog,3);
119   ierr = PetscObjectTypeCompare((PetscObject)is,ISBLOCK,&isblock);CHKERRQ(ierr);
120   if (isblock) {
121     PetscInt lbs;
122 
123     ierr = ISGetBlockSize(is,&bs);CHKERRQ(ierr);
124     ierr = ISLocalToGlobalMappingGetBlockSize(ltog,&lbs);CHKERRQ(ierr);
125     if (bs == lbs) {
126       ierr = ISGetLocalSize(is,&m);CHKERRQ(ierr);
127       m    = m/bs;
128       ierr = ISBlockGetIndices(is,&idx);CHKERRQ(ierr);
129       ierr = PetscMalloc1(m,&idxm);CHKERRQ(ierr);
130       ierr = ISLocalToGlobalMappingApplyBlock(ltog,m,idx,idxm);CHKERRQ(ierr);
131       ierr = ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is),bs,m,idxm,PETSC_OWN_POINTER,cltog);CHKERRQ(ierr);
132       ierr = ISBlockRestoreIndices(is,&idx);CHKERRQ(ierr);
133       PetscFunctionReturn(0);
134     }
135   }
136   ierr = ISGetLocalSize(is,&m);CHKERRQ(ierr);
137   ierr = ISGetIndices(is,&idx);CHKERRQ(ierr);
138   ierr = ISGetBlockSize(is,&bs);CHKERRQ(ierr);
139   ierr = PetscMalloc1(m,&idxm);CHKERRQ(ierr);
140   if (ltog) {
141     ierr = ISLocalToGlobalMappingApply(ltog,m,idx,idxm);CHKERRQ(ierr);
142   } else {
143     ierr = PetscMemcpy(idxm,idx,m*sizeof(PetscInt));CHKERRQ(ierr);
144   }
145   ierr = ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is),bs,m,idxm,PETSC_OWN_POINTER,cltog);CHKERRQ(ierr);
146   ierr = ISRestoreIndices(is,&idx);CHKERRQ(ierr);
147   PetscFunctionReturn(0);
148 }
149 
150 #undef __FUNCT__
151 #define __FUNCT__ "ISL2GComposeBlock"
152 static PetscErrorCode ISL2GComposeBlock(IS is,ISLocalToGlobalMapping ltog,ISLocalToGlobalMapping *cltog)
153 {
154   PetscErrorCode ierr;
155   const PetscInt *idx;
156   PetscInt       m,*idxm,bs;
157 
158   PetscFunctionBegin;
159   PetscValidHeaderSpecific(is,IS_CLASSID,1);
160   PetscValidHeaderSpecific(ltog,IS_LTOGM_CLASSID,2);
161   PetscValidPointer(cltog,3);
162   ierr = ISBlockGetLocalSize(is,&m);CHKERRQ(ierr);
163   ierr = ISBlockGetIndices(is,&idx);CHKERRQ(ierr);
164   ierr = ISLocalToGlobalMappingGetBlockSize(ltog,&bs);CHKERRQ(ierr);
165   ierr = PetscMalloc1(m,&idxm);CHKERRQ(ierr);
166   if (ltog) {
167     ierr = ISLocalToGlobalMappingApplyBlock(ltog,m,idx,idxm);CHKERRQ(ierr);
168   } else {
169     ierr = PetscMemcpy(idxm,idx,m*sizeof(PetscInt));CHKERRQ(ierr);
170   }
171   ierr = ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is),bs,m,idxm,PETSC_OWN_POINTER,cltog);CHKERRQ(ierr);
172   ierr = ISBlockRestoreIndices(is,&idx);CHKERRQ(ierr);
173   PetscFunctionReturn(0);
174 }
175 
176 #undef __FUNCT__
177 #define __FUNCT__ "MatDestroy_LocalRef"
178 static PetscErrorCode MatDestroy_LocalRef(Mat B)
179 {
180   PetscErrorCode ierr;
181 
182   PetscFunctionBegin;
183   ierr = PetscFree(B->data);CHKERRQ(ierr);
184   PetscFunctionReturn(0);
185 }
186 
187 
188 #undef __FUNCT__
189 #define __FUNCT__ "MatCreateLocalRef"
190 /*@
191    MatCreateLocalRef - Gets a logical reference to a local submatrix, for use in assembly
192 
193    Not Collective
194 
195    Input Arguments:
196 + A - Full matrix, generally parallel
197 . isrow - Local index set for the rows
198 - iscol - Local index set for the columns
199 
200    Output Arguments:
201 . newmat - New serial Mat
202 
203    Level: developer
204 
205    Notes:
206    Most will use MatGetLocalSubMatrix() which returns a real matrix corresponding to the local
207    block if it available, such as with matrix formats that store these blocks separately.
208 
209    The new matrix forwards MatSetValuesLocal() and MatSetValuesBlockedLocal() to the global system.
210    In general, it does not define MatMult() or any other functions.  Local submatrices can be nested.
211 
212 .seealso: MatSetValuesLocal(), MatSetValuesBlockedLocal(), MatGetLocalSubMatrix(), MatCreateSubMatrix()
213 @*/
214 PetscErrorCode  MatCreateLocalRef(Mat A,IS isrow,IS iscol,Mat *newmat)
215 {
216   PetscErrorCode ierr;
217   Mat_LocalRef   *lr;
218   Mat            B;
219   PetscInt       m,n;
220   PetscBool      islr;
221 
222   PetscFunctionBegin;
223   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
224   PetscValidHeaderSpecific(isrow,IS_CLASSID,2);
225   PetscValidHeaderSpecific(iscol,IS_CLASSID,3);
226   PetscValidPointer(newmat,4);
227   if (!A->rmap->mapping) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONGSTATE,"Matrix must have local to global mapping provided before this call");
228   *newmat = 0;
229 
230   ierr = MatCreate(PETSC_COMM_SELF,&B);CHKERRQ(ierr);
231   ierr = ISGetLocalSize(isrow,&m);CHKERRQ(ierr);
232   ierr = ISGetLocalSize(iscol,&n);CHKERRQ(ierr);
233   ierr = MatSetSizes(B,m,n,m,n);CHKERRQ(ierr);
234   ierr = PetscObjectChangeTypeName((PetscObject)B,MATLOCALREF);CHKERRQ(ierr);
235   ierr = MatSetUp(B);CHKERRQ(ierr);
236 
237   B->ops->destroy = MatDestroy_LocalRef;
238 
239   ierr    = PetscNewLog(B,&lr);CHKERRQ(ierr);
240   B->data = (void*)lr;
241 
242   ierr = PetscObjectTypeCompare((PetscObject)A,MATLOCALREF,&islr);CHKERRQ(ierr);
243   if (islr) {
244     Mat_LocalRef *alr = (Mat_LocalRef*)A->data;
245     lr->Top = alr->Top;
246   } else {
247     /* This does not increase the reference count because MatLocalRef is not allowed to live longer than its parent */
248     lr->Top = A;
249   }
250   {
251     ISLocalToGlobalMapping rltog,cltog;
252     PetscInt               arbs,acbs,rbs,cbs;
253 
254     /* We will translate directly to global indices for the top level */
255     lr->SetValues        = MatSetValues;
256     lr->SetValuesBlocked = MatSetValuesBlocked;
257 
258     B->ops->setvalueslocal = MatSetValuesLocal_LocalRef_Scalar;
259 
260     ierr = ISL2GCompose(isrow,A->rmap->mapping,&rltog);CHKERRQ(ierr);
261     if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) {
262       ierr  = PetscObjectReference((PetscObject)rltog);CHKERRQ(ierr);
263       cltog = rltog;
264     } else {
265       ierr = ISL2GCompose(iscol,A->cmap->mapping,&cltog);CHKERRQ(ierr);
266     }
267     /* Remember if the ISes we used to pull out the submatrix are of type ISBLOCK.  Will be used later in
268      * MatSetValuesLocal_LocalRef_Scalar. */
269     ierr = PetscObjectTypeCompare((PetscObject)isrow,ISBLOCK,&lr->rowisblock);CHKERRQ(ierr);
270     ierr = PetscObjectTypeCompare((PetscObject)iscol,ISBLOCK,&lr->colisblock);CHKERRQ(ierr);
271     ierr = MatSetLocalToGlobalMapping(B,rltog,cltog);CHKERRQ(ierr);
272     ierr = ISLocalToGlobalMappingDestroy(&rltog);CHKERRQ(ierr);
273     ierr = ISLocalToGlobalMappingDestroy(&cltog);CHKERRQ(ierr);
274 
275     ierr = MatGetBlockSizes(A,&arbs,&acbs);CHKERRQ(ierr);
276     ierr = ISGetBlockSize(isrow,&rbs);CHKERRQ(ierr);
277     ierr = ISGetBlockSize(iscol,&cbs);CHKERRQ(ierr);
278     /* Always support block interface insertion on submatrix */
279     ierr = PetscLayoutSetBlockSize(B->rmap,rbs);CHKERRQ(ierr);
280     ierr = PetscLayoutSetBlockSize(B->cmap,cbs);CHKERRQ(ierr);
281     if (arbs != rbs || acbs != cbs || (arbs == 1 && acbs == 1)) {
282       /* Top-level matrix has different block size, so we have to call its scalar insertion interface */
283       B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Scalar;
284     } else {
285       /* Block sizes match so we can forward values to the top level using the block interface */
286       B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Block;
287 
288       ierr = ISL2GComposeBlock(isrow,A->rmap->mapping,&rltog);CHKERRQ(ierr);
289       if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) {
290         ierr  =  PetscObjectReference((PetscObject)rltog);CHKERRQ(ierr);
291         cltog = rltog;
292       } else {
293         ierr = ISL2GComposeBlock(iscol,A->cmap->mapping,&cltog);CHKERRQ(ierr);
294       }
295       ierr = MatSetLocalToGlobalMapping(B,rltog,cltog);CHKERRQ(ierr);
296       ierr = ISLocalToGlobalMappingDestroy(&rltog);CHKERRQ(ierr);
297       ierr = ISLocalToGlobalMappingDestroy(&cltog);CHKERRQ(ierr);
298     }
299   }
300   *newmat = B;
301   PetscFunctionReturn(0);
302 }
303