Actual source code: isblock.c

      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: }