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