impls/sliced/sliced.c

#include <petscdmsliced.h> /*I      "petscdmsliced.h" I*/
#include <petscmat.h>
#include <petsc/private/dmimpl.h>

/* CSR storage of the nonzero structure of a bs*bs matrix */
typedef struct {
  PetscInt bs, nz, *i, *j;
} DMSlicedBlockFills;

typedef struct {
  PetscInt            bs, n, N, Nghosts, *ghosts;
  PetscInt            d_nz, o_nz, *d_nnz, *o_nnz;
  DMSlicedBlockFills *dfill, *ofill;
} DM_Sliced;

PetscErrorCode DMCreateMatrix_Sliced(DM dm, Mat *J)
{
  PetscInt              *globals, *sd_nnz, *so_nnz, rstart, bs, i;
  ISLocalToGlobalMapping lmap;
  void (*aij)(void) = NULL;
  DM_Sliced *slice  = (DM_Sliced *)dm->data;

  PetscFunctionBegin;
  bs = slice->bs;
  PetscCall(MatCreate(PetscObjectComm((PetscObject)dm), J));
  PetscCall(MatSetSizes(*J, slice->n * bs, slice->n * bs, PETSC_DETERMINE, PETSC_DETERMINE));
  PetscCall(MatSetBlockSize(*J, bs));
  PetscCall(MatSetType(*J, dm->mattype));
  PetscCall(MatSeqBAIJSetPreallocation(*J, bs, slice->d_nz, slice->d_nnz));
  PetscCall(MatMPIBAIJSetPreallocation(*J, bs, slice->d_nz, slice->d_nnz, slice->o_nz, slice->o_nnz));
  /* In general, we have to do extra work to preallocate for scalar (AIJ) matrices so we check whether it will do any
  * good before going on with it. */
  PetscCall(PetscObjectQueryFunction((PetscObject)*J, "MatMPIAIJSetPreallocation_C", &aij));
  if (!aij) PetscCall(PetscObjectQueryFunction((PetscObject)*J, "MatSeqAIJSetPreallocation_C", &aij));
  if (aij) {
    if (bs == 1) {
      PetscCall(MatSeqAIJSetPreallocation(*J, slice->d_nz, slice->d_nnz));
      PetscCall(MatMPIAIJSetPreallocation(*J, slice->d_nz, slice->d_nnz, slice->o_nz, slice->o_nnz));
    } else if (!slice->d_nnz) {
      PetscCall(MatSeqAIJSetPreallocation(*J, slice->d_nz * bs, NULL));
      PetscCall(MatMPIAIJSetPreallocation(*J, slice->d_nz * bs, NULL, slice->o_nz * bs, NULL));
    } else {
      /* The user has provided preallocation per block-row, convert it to per scalar-row respecting DMSlicedSetBlockFills() if applicable */
      PetscCall(PetscMalloc2(slice->n * bs, &sd_nnz, (!!slice->o_nnz) * slice->n * bs, &so_nnz));
      for (i = 0; i < slice->n * bs; i++) {
        sd_nnz[i] = (slice->d_nnz[i / bs] - 1) * (slice->ofill ? slice->ofill->i[i % bs + 1] - slice->ofill->i[i % bs] : bs) + (slice->dfill ? slice->dfill->i[i % bs + 1] - slice->dfill->i[i % bs] : bs);
        if (so_nnz) so_nnz[i] = slice->o_nnz[i / bs] * (slice->ofill ? slice->ofill->i[i % bs + 1] - slice->ofill->i[i % bs] : bs);
      }
      PetscCall(MatSeqAIJSetPreallocation(*J, slice->d_nz * bs, sd_nnz));
      PetscCall(MatMPIAIJSetPreallocation(*J, slice->d_nz * bs, sd_nnz, slice->o_nz * bs, so_nnz));
      PetscCall(PetscFree2(sd_nnz, so_nnz));
    }
  }

  /* Set up the local to global map.  For the scalar map, we have to translate to entry-wise indexing instead of block-wise. */
  PetscCall(PetscMalloc1(slice->n + slice->Nghosts, &globals));
  PetscCall(MatGetOwnershipRange(*J, &rstart, NULL));
  for (i = 0; i < slice->n; i++) globals[i] = rstart / bs + i;

  for (i = 0; i < slice->Nghosts; i++) globals[slice->n + i] = slice->ghosts[i];
  PetscCall(ISLocalToGlobalMappingCreate(PETSC_COMM_SELF, bs, slice->n + slice->Nghosts, globals, PETSC_OWN_POINTER, &lmap));
  PetscCall(MatSetLocalToGlobalMapping(*J, lmap, lmap));
  PetscCall(ISLocalToGlobalMappingDestroy(&lmap));
  PetscCall(MatSetDM(*J, dm));
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@C
    DMSlicedSetGhosts - Sets the global indices of other processes elements that will
      be ghosts on this process

    Not Collective

    Input Parameters:
+    slice - the `DMSLICED` object
.    bs - block size
.    nlocal - number of local (owned, non-ghost) blocks
.    Nghosts - number of ghost blocks on this process
-    ghosts - global indices of each ghost block

    Level: advanced

.seealso `DMDestroy()`, `DMCreateGlobalVector()`
@*/
PetscErrorCode DMSlicedSetGhosts(DM dm, PetscInt bs, PetscInt nlocal, PetscInt Nghosts, const PetscInt ghosts[])
{
  DM_Sliced *slice = (DM_Sliced *)dm->data;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscCall(PetscFree(slice->ghosts));
  PetscCall(PetscMalloc1(Nghosts, &slice->ghosts));
  PetscCall(PetscArraycpy(slice->ghosts, ghosts, Nghosts));
  slice->bs      = bs;
  slice->n       = nlocal;
  slice->Nghosts = Nghosts;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@C
    DMSlicedSetPreallocation - sets the matrix memory preallocation for matrices computed by DMSliced

    Not Collective

    Input Parameters:
+    slice - the `DM` object
.    d_nz  - number of block nonzeros per block row in diagonal portion of local
           submatrix  (same for all local rows)
.    d_nnz - array containing the number of block nonzeros in the various block rows
           of the in diagonal portion of the local (possibly different for each block
           row) or NULL.
.    o_nz  - number of block nonzeros per block row in the off-diagonal portion of local
           submatrix (same for all local rows).
-    o_nnz - array containing the number of nonzeros in the various block rows of the
           off-diagonal portion of the local submatrix (possibly different for
           each block row) or `NULL`.

    Level: advanced

    Note:
    See `MatMPIBAIJSetPreallocation()` for more details on preallocation.  If a scalar matrix (AIJ) is
    obtained with `DMSlicedGetMatrix()`, the correct preallocation will be set, respecting `DMSlicedSetBlockFills()`.

.seealso `DMDestroy()`, `DMCreateGlobalVector()`, `MatMPIAIJSetPreallocation()`,
         `MatMPIBAIJSetPreallocation()`, `DMSlicedGetMatrix()`, `DMSlicedSetBlockFills()`
@*/
PetscErrorCode DMSlicedSetPreallocation(DM dm, PetscInt d_nz, const PetscInt d_nnz[], PetscInt o_nz, const PetscInt o_nnz[])
{
  DM_Sliced *slice = (DM_Sliced *)dm->data;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  slice->d_nz  = d_nz;
  slice->d_nnz = (PetscInt *)d_nnz;
  slice->o_nz  = o_nz;
  slice->o_nnz = (PetscInt *)o_nnz;
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode DMSlicedSetBlockFills_Private(PetscInt bs, const PetscInt *fill, DMSlicedBlockFills **inf)
{
  PetscInt            i, j, nz, *fi, *fj;
  DMSlicedBlockFills *f;

  PetscFunctionBegin;
  PetscValidPointer(inf, 3);
  if (*inf) PetscCall(PetscFree3(*inf, (*inf)->i, (*inf)->j));
  if (!fill) PetscFunctionReturn(PETSC_SUCCESS);
  for (i = 0, nz = 0; i < bs * bs; i++)
    if (fill[i]) nz++;
  PetscCall(PetscMalloc3(1, &f, bs + 1, &fi, nz, &fj));
  f->bs = bs;
  f->nz = nz;
  f->i  = fi;
  f->j  = fj;
  for (i = 0, nz = 0; i < bs; i++) {
    fi[i] = nz;
    for (j = 0; j < bs; j++)
      if (fill[i * bs + j]) fj[nz++] = j;
  }
  fi[i] = nz;
  *inf  = f;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@C
    DMSlicedSetBlockFills - Sets the fill pattern in each block for a multi-component problem
    of the matrix returned by `DMSlicedGetMatrix()`.

    Logically Collective

    Input Parameters:
+   sliced - the `DM` object
.   dfill - the fill pattern in the diagonal block (may be `NULL`, means use dense block)
-   ofill - the fill pattern in the off-diagonal blocks

    Level: advanced

    Note:
    This only makes sense for multicomponent problems using scalar matrix formats (AIJ).
    See `DMDASetBlockFills()` for example usage.

.seealso `DMSlicedGetMatrix()`, `DMDASetBlockFills()`
@*/
PetscErrorCode DMSlicedSetBlockFills(DM dm, const PetscInt *dfill, const PetscInt *ofill)
{
  DM_Sliced *slice = (DM_Sliced *)dm->data;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscCall(DMSlicedSetBlockFills_Private(slice->bs, dfill, &slice->dfill));
  PetscCall(DMSlicedSetBlockFills_Private(slice->bs, ofill, &slice->ofill));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode DMDestroy_Sliced(DM dm)
{
  DM_Sliced *slice = (DM_Sliced *)dm->data;

  PetscFunctionBegin;
  PetscCall(PetscFree(slice->ghosts));
  if (slice->dfill) PetscCall(PetscFree3(slice->dfill, slice->dfill->i, slice->dfill->j));
  if (slice->ofill) PetscCall(PetscFree3(slice->ofill, slice->ofill->i, slice->ofill->j));
  /* This was originally freed in DMDestroy(), but that prevents reference counting of backend objects */
  PetscCall(PetscFree(slice));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode DMCreateGlobalVector_Sliced(DM dm, Vec *gvec)
{
  DM_Sliced *slice = (DM_Sliced *)dm->data;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
  PetscValidPointer(gvec, 2);
  *gvec = NULL;
  PetscCall(VecCreateGhostBlock(PetscObjectComm((PetscObject)dm), slice->bs, slice->n * slice->bs, PETSC_DETERMINE, slice->Nghosts, slice->ghosts, gvec));
  PetscCall(VecSetDM(*gvec, dm));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode DMGlobalToLocalBegin_Sliced(DM da, Vec g, InsertMode mode, Vec l)
{
  PetscBool flg;

  PetscFunctionBegin;
  PetscCall(VecGhostIsLocalForm(g, l, &flg));
  PetscCheck(flg, PetscObjectComm((PetscObject)da), PETSC_ERR_ARG_WRONG, "Local vector is not local form of global vector");
  PetscCall(VecGhostUpdateEnd(g, mode, SCATTER_FORWARD));
  PetscCall(VecGhostUpdateBegin(g, mode, SCATTER_FORWARD));
  PetscFunctionReturn(PETSC_SUCCESS);
}

static PetscErrorCode DMGlobalToLocalEnd_Sliced(DM da, Vec g, InsertMode mode, Vec l)
{
  PetscBool flg;

  PetscFunctionBegin;
  PetscCall(VecGhostIsLocalForm(g, l, &flg));
  PetscCheck(flg, PetscObjectComm((PetscObject)da), PETSC_ERR_ARG_WRONG, "Local vector is not local form of global vector");
  PetscCall(VecGhostUpdateEnd(g, mode, SCATTER_FORWARD));
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*MC
   DMSLICED = "sliced" - A DM object that is used to manage data for a general graph. Uses VecCreateGhost() ghosted vectors for storing the fields

   See DMCreateSliced() for details.

  Level: intermediate

.seealso: `DMType`, `DMCOMPOSITE`, `DMCreateSliced()`, `DMCreate()`
M*/

PETSC_EXTERN PetscErrorCode DMCreate_Sliced(DM p)
{
  DM_Sliced *slice;

  PetscFunctionBegin;
  PetscCall(PetscNew(&slice));
  p->data = slice;

  p->ops->createglobalvector = DMCreateGlobalVector_Sliced;
  p->ops->creatematrix       = DMCreateMatrix_Sliced;
  p->ops->globaltolocalbegin = DMGlobalToLocalBegin_Sliced;
  p->ops->globaltolocalend   = DMGlobalToLocalEnd_Sliced;
  p->ops->destroy            = DMDestroy_Sliced;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@C
    DMSlicedCreate - Creates a `DM` object, used to manage data for a unstructured problem

    Collective

    Input Parameters:
+   comm - the processors that will share the global vector
.   bs - the block size
.   nlocal - number of vector entries on this process
.   Nghosts - number of ghost points needed on this process
.   ghosts - global indices of all ghost points for this process
.   d_nnz - matrix preallocation information representing coupling within this process
-   o_nnz - matrix preallocation information representing coupling between this process and other processes

    Output Parameters:
.   slice - the slice object

    Level: advanced

    Notes:
        This `DM` does not support `DMCreateLocalVector()`, `DMGlobalToLocalBegin()`, and `DMGlobalToLocalEnd()` instead one directly uses
        `VecGhostGetLocalForm()` and `VecGhostRestoreLocalForm()` to access the local representation and `VecGhostUpdateBegin()` and `VecGhostUpdateEnd()` to update
        the ghost points.

        One can use `DMGlobalToLocalBegin()`, and `DMGlobalToLocalEnd()` instead of `VecGhostUpdateBegin()` and `VecGhostUpdateEnd()`.

.seealso `DMDestroy()`, `DMCreateGlobalVector()`, `DMSetType()`, `DMSLICED`, `DMSlicedSetGhosts()`, `DMSlicedSetPreallocation()`, `VecGhostUpdateBegin()`, `VecGhostUpdateEnd()`,
         `VecGhostGetLocalForm()`, `VecGhostRestoreLocalForm()`

@*/
PetscErrorCode DMSlicedCreate(MPI_Comm comm, PetscInt bs, PetscInt nlocal, PetscInt Nghosts, const PetscInt ghosts[], const PetscInt d_nnz[], const PetscInt o_nnz[], DM *dm)
{
  PetscFunctionBegin;
  PetscValidPointer(dm, 8);
  PetscCall(DMCreate(comm, dm));
  PetscCall(DMSetType(*dm, DMSLICED));
  PetscCall(DMSlicedSetGhosts(*dm, bs, nlocal, Nghosts, ghosts));
  if (d_nnz) PetscCall(DMSlicedSetPreallocation(*dm, 0, d_nnz, 0, o_nnz));
  PetscFunctionReturn(PETSC_SUCCESS);
}