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