xref: /petsc/src/mat/impls/localref/mlocalref.c (revision 6e191704724e79c1d402f837edaaafddcafef25c)
1 
2 #include <petsc-private/matimpl.h>          /*I "petscmat.h" I*/
3 
4 typedef struct {
5   Mat Top;
6   PetscErrorCode (*SetValues)(Mat,PetscInt,const PetscInt[],PetscInt,const PetscInt[],const PetscScalar[],InsertMode);
7   PetscErrorCode (*SetValuesBlocked)(Mat,PetscInt,const PetscInt[],PetscInt,const PetscInt[],const PetscScalar[],InsertMode);
8 } Mat_LocalRef;
9 
10 /* These need to be macros because they use sizeof */
11 #define IndexSpaceGet(buf,nrow,ncol,irowm,icolm) do {                   \
12     if (nrow + ncol > (PetscInt)(sizeof(buf)/sizeof(buf[0]))) {         \
13       ierr = PetscMalloc2(nrow,&irowm,ncol,&icolm);CHKERRQ(ierr); \
14     } else {                                                            \
15       irowm = &buf[0];                                                  \
16       icolm = &buf[nrow];                                               \
17     }                                                                   \
18 } while (0)
19 
20 #define IndexSpaceRestore(buf,nrow,ncol,irowm,icolm) do {       \
21     if (nrow + ncol > (PetscInt)(sizeof(buf)/sizeof(buf[0]))) { \
22       ierr = PetscFree2(irowm,icolm);CHKERRQ(ierr);             \
23     }                                                           \
24 } while (0)
25 
26 static void BlockIndicesExpand(PetscInt n,const PetscInt idx[],PetscInt bs,PetscInt idxm[])
27 {
28   PetscInt i,j;
29   for (i=0; i<n; i++) {
30     for (j=0; j<bs; j++) {
31       idxm[i*bs+j] = idx[i]*bs + j;
32     }
33   }
34 }
35 
36 #undef __FUNCT__
37 #define __FUNCT__ "MatSetValuesBlockedLocal_LocalRef_Block"
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,*icolm;
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 #undef __FUNCT__
55 #define __FUNCT__ "MatSetValuesBlockedLocal_LocalRef_Scalar"
56 static PetscErrorCode MatSetValuesBlockedLocal_LocalRef_Scalar(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv)
57 {
58   Mat_LocalRef   *lr = (Mat_LocalRef*)A->data;
59   PetscErrorCode ierr;
60   PetscInt       rbs,cbs,buf[4096],*irowm,*icolm;
61 
62   PetscFunctionBegin;
63   ierr = MatGetBlockSizes(A,&rbs,&cbs);CHKERRQ(ierr);
64   IndexSpaceGet(buf,nrow*rbs,ncol*cbs,irowm,icolm);
65   BlockIndicesExpand(nrow,irow,rbs,irowm);
66   BlockIndicesExpand(ncol,icol,cbs,icolm);
67   ierr = ISLocalToGlobalMappingApplyBlock(A->rmap->mapping,nrow*rbs,irowm,irowm);CHKERRQ(ierr);
68   ierr = ISLocalToGlobalMappingApplyBlock(A->cmap->mapping,ncol*cbs,icolm,icolm);CHKERRQ(ierr);
69   ierr = (*lr->SetValues)(lr->Top,nrow*rbs,irowm,ncol*cbs,icolm,y,addv);CHKERRQ(ierr);
70   IndexSpaceRestore(buf,nrow*rbs,ncol*cbs,irowm,icolm);
71   PetscFunctionReturn(0);
72 }
73 
74 #undef __FUNCT__
75 #define __FUNCT__ "MatSetValuesLocal_LocalRef_Scalar"
76 static PetscErrorCode MatSetValuesLocal_LocalRef_Scalar(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv)
77 {
78   Mat_LocalRef   *lr = (Mat_LocalRef*)A->data;
79   PetscErrorCode ierr;
80   PetscInt       buf[4096],*irowm,*icolm;
81 
82   PetscFunctionBegin;
83   IndexSpaceGet(buf,nrow,ncol,irowm,icolm);
84   ierr = ISLocalToGlobalMappingApplyBlock(A->rmap->mapping,nrow,irow,irowm);CHKERRQ(ierr);
85   ierr = ISLocalToGlobalMappingApplyBlock(A->cmap->mapping,ncol,icol,icolm);CHKERRQ(ierr);
86   ierr = (*lr->SetValues)(lr->Top,nrow,irowm,ncol,icolm,y,addv);CHKERRQ(ierr);
87   IndexSpaceRestore(buf,nrow,ncol,irowm,icolm);
88   PetscFunctionReturn(0);
89 }
90 
91 #undef __FUNCT__
92 #define __FUNCT__ "ISL2GCompose"
93 /* Compose an IS with an ISLocalToGlobalMapping to map from IS source indices to global indices */
94 static PetscErrorCode ISL2GCompose(IS is,ISLocalToGlobalMapping ltog,ISLocalToGlobalMapping *cltog)
95 {
96   PetscErrorCode ierr;
97   const PetscInt *idx;
98   PetscInt       m,*idxm;
99   PetscInt       bs;
100   PetscBool      isblock;
101 
102   PetscFunctionBegin;
103   PetscValidHeaderSpecific(is,IS_CLASSID,1);
104   PetscValidHeaderSpecific(ltog,IS_LTOGM_CLASSID,2);
105   PetscValidPointer(cltog,3);
106   ierr = PetscObjectTypeCompare((PetscObject)is,ISBLOCK,&isblock);CHKERRQ(ierr);
107   if (isblock) {
108     PetscInt lbs;
109 
110     ierr = ISGetBlockSize(is,&bs);CHKERRQ(ierr);
111     ierr = ISLocalToGlobalMappingGetBlockSize(ltog,&lbs);CHKERRQ(ierr);
112     if (bs == lbs) {
113       ierr = ISGetLocalSize(is,&m);CHKERRQ(ierr);
114       m    = m/bs;
115       ierr = ISBlockGetIndices(is,&idx);CHKERRQ(ierr);
116       ierr = PetscMalloc1(m,&idxm);CHKERRQ(ierr);
117       ierr = ISLocalToGlobalMappingApplyBlock(ltog,m,idx,idxm);CHKERRQ(ierr);
118       ierr = ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is),bs,m,idxm,PETSC_OWN_POINTER,cltog);CHKERRQ(ierr);
119       ierr = ISBlockRestoreIndices(is,&idx);CHKERRQ(ierr);
120       PetscFunctionReturn(0);
121     }
122   }
123   ierr = ISGetLocalSize(is,&m);CHKERRQ(ierr);
124   ierr = ISGetIndices(is,&idx);CHKERRQ(ierr);
125   ierr = ISGetBlockSize(is,&bs);CHKERRQ(ierr);
126   ierr = PetscMalloc1(m,&idxm);CHKERRQ(ierr);
127   if (ltog) {
128     ierr = ISLocalToGlobalMappingApply(ltog,m,idx,idxm);CHKERRQ(ierr);
129   } else {
130     ierr = PetscMemcpy(idxm,idx,m*sizeof(PetscInt));CHKERRQ(ierr);
131   }
132   ierr = ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is),bs,m,idxm,PETSC_OWN_POINTER,cltog);CHKERRQ(ierr);
133   ierr = ISRestoreIndices(is,&idx);CHKERRQ(ierr);
134   PetscFunctionReturn(0);
135 }
136 
137 #undef __FUNCT__
138 #define __FUNCT__ "ISL2GComposeBlock"
139 static PetscErrorCode ISL2GComposeBlock(IS is,ISLocalToGlobalMapping ltog,ISLocalToGlobalMapping *cltog)
140 {
141   PetscErrorCode ierr;
142   const PetscInt *idx;
143   PetscInt       m,*idxm,bs;
144 
145   PetscFunctionBegin;
146   PetscValidHeaderSpecific(is,IS_CLASSID,1);
147   PetscValidHeaderSpecific(ltog,IS_LTOGM_CLASSID,2);
148   PetscValidPointer(cltog,3);
149   ierr = ISBlockGetLocalSize(is,&m);CHKERRQ(ierr);
150   ierr = ISBlockGetIndices(is,&idx);CHKERRQ(ierr);
151   ierr = ISLocalToGlobalMappingGetBlockSize(ltog,&bs);CHKERRQ(ierr);
152   ierr = PetscMalloc1(m,&idxm);CHKERRQ(ierr);
153   if (ltog) {
154     ierr = ISLocalToGlobalMappingApplyBlock(ltog,m,idx,idxm);CHKERRQ(ierr);
155   } else {
156     ierr = PetscMemcpy(idxm,idx,m*sizeof(PetscInt));CHKERRQ(ierr);
157   }
158   ierr = ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is),bs,m,idxm,PETSC_OWN_POINTER,cltog);CHKERRQ(ierr);
159   ierr = ISBlockRestoreIndices(is,&idx);CHKERRQ(ierr);
160   PetscFunctionReturn(0);
161 }
162 
163 #undef __FUNCT__
164 #define __FUNCT__ "MatDestroy_LocalRef"
165 static PetscErrorCode MatDestroy_LocalRef(Mat B)
166 {
167   PetscErrorCode ierr;
168 
169   PetscFunctionBegin;
170   ierr = PetscFree(B->data);CHKERRQ(ierr);
171   PetscFunctionReturn(0);
172 }
173 
174 
175 #undef __FUNCT__
176 #define __FUNCT__ "MatCreateLocalRef"
177 /*@
178    MatCreateLocalRef - Gets a logical reference to a local submatrix, for use in assembly
179 
180    Not Collective
181 
182    Input Arguments:
183 + A - Full matrix, generally parallel
184 . isrow - Local index set for the rows
185 - iscol - Local index set for the columns
186 
187    Output Arguments:
188 . newmat - New serial Mat
189 
190    Level: developer
191 
192    Notes:
193    Most will use MatGetLocalSubMatrix() which returns a real matrix corresponding to the local
194    block if it available, such as with matrix formats that store these blocks separately.
195 
196    The new matrix forwards MatSetValuesLocal() and MatSetValuesBlockedLocal() to the global system.
197    In general, it does not define MatMult() or any other functions.  Local submatrices can be nested.
198 
199 .seealso: MatSetValuesLocal(), MatSetValuesBlockedLocal(), MatGetLocalSubMatrix(), MatCreateSubMatrix()
200 @*/
201 PetscErrorCode  MatCreateLocalRef(Mat A,IS isrow,IS iscol,Mat *newmat)
202 {
203   PetscErrorCode ierr;
204   Mat_LocalRef   *lr;
205   Mat            B;
206   PetscInt       m,n;
207   PetscBool      islr;
208 
209   PetscFunctionBegin;
210   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
211   PetscValidHeaderSpecific(isrow,IS_CLASSID,2);
212   PetscValidHeaderSpecific(iscol,IS_CLASSID,3);
213   PetscValidPointer(newmat,4);
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     ierr = MatSetLocalToGlobalMapping(B,rltog,cltog);CHKERRQ(ierr);
254     ierr = ISLocalToGlobalMappingDestroy(&rltog);CHKERRQ(ierr);
255     ierr = ISLocalToGlobalMappingDestroy(&cltog);CHKERRQ(ierr);
256 
257     ierr = MatGetBlockSizes(A,&arbs,&acbs);CHKERRQ(ierr);
258     ierr = ISGetBlockSize(isrow,&rbs);CHKERRQ(ierr);
259     ierr = ISGetBlockSize(iscol,&cbs);CHKERRQ(ierr);
260     /* Always support block interface insertion on submatrix */
261     ierr = PetscLayoutSetBlockSize(B->rmap,rbs);CHKERRQ(ierr);
262     ierr = PetscLayoutSetBlockSize(B->cmap,cbs);CHKERRQ(ierr);
263     if (arbs != rbs || acbs != cbs || (arbs == 1 && acbs == 1)) {
264       /* Top-level matrix has different block size, so we have to call its scalar insertion interface */
265       B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Scalar;
266     } else {
267       /* Block sizes match so we can forward values to the top level using the block interface */
268       B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Block;
269 
270       ierr = ISL2GComposeBlock(isrow,A->rmap->mapping,&rltog);CHKERRQ(ierr);
271       if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) {
272         ierr  =  PetscObjectReference((PetscObject)rltog);CHKERRQ(ierr);
273         cltog = rltog;
274       } else {
275         ierr = ISL2GComposeBlock(iscol,A->cmap->mapping,&cltog);CHKERRQ(ierr);
276       }
277       ierr = MatSetLocalToGlobalMapping(B,rltog,cltog);CHKERRQ(ierr);
278       ierr = ISLocalToGlobalMappingDestroy(&rltog);CHKERRQ(ierr);
279       ierr = ISLocalToGlobalMappingDestroy(&cltog);CHKERRQ(ierr);
280     }
281   }
282   *newmat = B;
283   PetscFunctionReturn(0);
284 }
285