1: /* Routines to be used by MatIncreaseOverlap() for BAIJ and SBAIJ matrices */
  2: #include <petscis.h>
  3: #include <petscbt.h>
  4: #include <petsc/private/hashmapi.h>

  6: /*@
  7:   ISCompressIndicesGeneral - convert the indices of an array of `IS` into an array of `ISGENERAL` of block indices

  9:   Input Parameters:
 10: + n     - maximum possible length of the index set
 11: . nkeys - expected number of keys when using `PETSC_USE_CTABLE`
 12: . bs    - the size of block
 13: . imax  - the number of index sets
 14: - is_in - the non-blocked array of index sets

 16:   Output Parameter:
 17: . is_out - the blocked new index set, as `ISGENERAL`, not as `ISBLOCK`

 19:   Level: intermediate

 21: .seealso: [](sec_scatter), `IS`, `ISGENERAL`, `ISExpandIndicesGeneral()`
 22: @*/
 23: PetscErrorCode ISCompressIndicesGeneral(PetscInt n, PetscInt nkeys, PetscInt bs, PetscInt imax, const IS is_in[], IS is_out[])
 24: {
 25:   PetscInt        isz, len, i, j, ival, bbs;
 26:   const PetscInt *idx;
 27:   PetscBool       isblock;
 28: #if defined(PETSC_USE_CTABLE)
 29:   PetscHMapI    gid1_lid1 = NULL;
 30:   PetscInt      tt, gid1, *nidx;
 31:   PetscHashIter tpos;
 32: #else
 33:   PetscInt *nidx;
 34:   PetscInt  Nbs;
 35:   PetscBT   table;
 36: #endif

 38:   PetscFunctionBegin;
 39: #if defined(PETSC_USE_CTABLE)
 40:   PetscCall(PetscHMapICreateWithSize(nkeys / bs, &gid1_lid1));
 41: #else
 42:   Nbs = n / bs;
 43:   PetscCall(PetscMalloc1(Nbs, &nidx));
 44:   PetscCall(PetscBTCreate(Nbs, &table));
 45: #endif
 46:   for (i = 0; i < imax; i++) {
 47:     PetscCall(ISGetLocalSize(is_in[i], &len));
 48:     /* special case where IS is already block IS of the correct size */
 49:     PetscCall(PetscObjectTypeCompare((PetscObject)is_in[i], ISBLOCK, &isblock));
 50:     if (isblock) {
 51:       PetscCall(ISGetBlockSize(is_in[i], &bbs));
 52:       if (bs == bbs) {
 53:         len = len / bs;
 54:         PetscCall(ISBlockGetIndices(is_in[i], &idx));
 55:         PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)is_in[i]), len, idx, PETSC_COPY_VALUES, is_out + i));
 56:         PetscCall(ISBlockRestoreIndices(is_in[i], &idx));
 57:         continue;
 58:       }
 59:     }
 60:     isz = 0;
 61: #if defined(PETSC_USE_CTABLE)
 62:     PetscCall(PetscHMapIClear(gid1_lid1));
 63: #else
 64:     PetscCall(PetscBTMemzero(Nbs, table));
 65: #endif
 66:     PetscCall(ISGetIndices(is_in[i], &idx));
 67:     for (j = 0; j < len; j++) {
 68:       ival = idx[j] / bs; /* convert the indices into block indices */
 69: #if defined(PETSC_USE_CTABLE)
 70:       PetscCall(PetscHMapIGetWithDefault(gid1_lid1, ival + 1, 0, &tt));
 71:       if (!tt) {
 72:         PetscCall(PetscHMapISet(gid1_lid1, ival + 1, isz + 1));
 73:         isz++;
 74:       }
 75: #else
 76:       PetscCheck(ival <= Nbs, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "index greater than mat-dim");
 77:       if (!PetscBTLookupSet(table, ival)) nidx[isz++] = ival;
 78: #endif
 79:     }
 80:     PetscCall(ISRestoreIndices(is_in[i], &idx));

 82: #if defined(PETSC_USE_CTABLE)
 83:     PetscCall(PetscMalloc1(isz, &nidx));
 84:     PetscHashIterBegin(gid1_lid1, tpos);
 85:     j = 0;
 86:     while (!PetscHashIterAtEnd(gid1_lid1, tpos)) {
 87:       PetscHashIterGetKey(gid1_lid1, tpos, gid1);
 88:       PetscHashIterGetVal(gid1_lid1, tpos, tt);
 89:       PetscCheck(tt-- <= isz, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "index greater than array-dim");
 90:       nidx[tt] = gid1 - 1;
 91:       j++;
 92:       PetscHashIterNext(gid1_lid1, tpos);
 93:     }
 94:     PetscCheck(j == isz, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "table error: jj != isz");
 95:     PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)is_in[i]), isz, nidx, PETSC_OWN_POINTER, is_out + i));
 96: #else
 97:     PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)is_in[i]), isz, nidx, PETSC_COPY_VALUES, is_out + i));
 98: #endif
 99:   }
100: #if defined(PETSC_USE_CTABLE)
101:   PetscCall(PetscHMapIDestroy(&gid1_lid1));
102: #else
103:   PetscCall(PetscBTDestroy(&table));
104:   PetscCall(PetscFree(nidx));
105: #endif
106:   PetscFunctionReturn(PETSC_SUCCESS);
107: }

109: /*@
110:   ISExpandIndicesGeneral - convert the indices of an array `IS` into non-block indices in an array of `ISGENERAL`

112:   Input Parameters:
113: + n     - the length of the index set (not being used)
114: . nkeys - expected number of keys when `PETSC_USE_CTABLE` is used
115: . bs    - the size of block
116: . imax  - the number of index sets
117: - is_in - the blocked array of index sets, must be as large as `imax`

119:   Output Parameter:
120: . is_out - the non-blocked new index set, as `ISGENERAL`, must be as large as `imax`

122:   Level: intermediate

124: .seealso: [](sec_scatter), `IS`, `ISGENERAL`, `ISCompressIndicesGeneral()`
125: @*/
126: PetscErrorCode ISExpandIndicesGeneral(PetscInt n, PetscInt nkeys, PetscInt bs, PetscInt imax, const IS is_in[], IS is_out[])
127: {
128:   PetscInt        len, i, j, k, *nidx;
129:   const PetscInt *idx;
130:   PetscInt        maxsz;

132:   PetscFunctionBegin;
133:   /* Check max size of is_in[] */
134:   maxsz = 0;
135:   for (i = 0; i < imax; i++) {
136:     PetscCall(ISGetLocalSize(is_in[i], &len));
137:     if (len > maxsz) maxsz = len;
138:   }
139:   PetscCall(PetscMalloc1(maxsz * bs, &nidx));

141:   for (i = 0; i < imax; i++) {
142:     PetscCall(ISGetLocalSize(is_in[i], &len));
143:     PetscCall(ISGetIndices(is_in[i], &idx));
144:     for (j = 0; j < len; ++j) {
145:       for (k = 0; k < bs; k++) nidx[j * bs + k] = idx[j] * bs + k;
146:     }
147:     PetscCall(ISRestoreIndices(is_in[i], &idx));
148:     PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)is_in[i]), len * bs, nidx, PETSC_COPY_VALUES, is_out + i));
149:   }
150:   PetscCall(PetscFree(nidx));
151:   PetscFunctionReturn(PETSC_SUCCESS);
152: }