static char help[] = "Tests MatMPIBAIJSetPreallocationCSR()\n\n";

/*
  Include "petscmat.h" so that we can use matrices.  Note that this file
  automatically includes:
     petscsys.h       - base PETSc routines   petscvec.h - vectors
     petscmat.h - matrices
     petscis.h     - index sets
     petscviewer.h - viewers
*/
#include <petscmat.h>

int main(int argc, char **args)
{
  Mat         A;
  PetscInt   *ia, *ja, bs = 2;
  PetscInt    N = 9, n;
  PetscInt    rstart, rend, row, col;
  PetscInt    i;
  PetscMPIInt rank, size;
  Vec         v;

  PetscFunctionBeginUser;
  PetscCall(PetscInitialize(&argc, &args, NULL, help));
  PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
  PetscCheck(size < 5, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Can only use at most 4 processors.");
  PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));

  /* Get a partition range based on the vector size */
  PetscCall(VecCreateFromOptions(PETSC_COMM_WORLD, NULL, 1, PETSC_DECIDE, N, &v));
  PetscCall(VecGetLocalSize(v, &n));
  PetscCall(VecGetOwnershipRange(v, &rstart, &rend));
  PetscCall(VecDestroy(&v));

  PetscCall(PetscMalloc1(n + 1, &ia));
  PetscCall(PetscMalloc1(3 * n, &ja));

  /* Construct a tri-diagonal CSR indexing */
  i     = 1;
  ia[0] = 0;
  for (row = rstart; row < rend; row++) {
    ia[i] = ia[i - 1];

    /* diagonal */
    col = row;
    {
      ja[ia[i]] = col;
      ia[i]++;
    }

    /* lower diagonal */
    col = row - 1;
    if (col >= 0) {
      ja[ia[i]] = col;
      ia[i]++;
    }

    /* upper diagonal */
    col = row + 1;
    if (col < N) {
      ja[ia[i]] = col;
      ia[i]++;
    }
    i++;
  }

  PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
  PetscCall(MatSetSizes(A, n, n, PETSC_DETERMINE, PETSC_DETERMINE));
  PetscCall(MatSetType(A, MATMPIAIJ));
  PetscCall(MatMPIAIJSetPreallocationCSR(A, ia, ja, NULL));
  PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
  PetscCall(MatDestroy(&A));

  PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
  PetscCall(MatSetSizes(A, bs * n, bs * n, PETSC_DETERMINE, PETSC_DETERMINE));
  PetscCall(MatSetType(A, MATMPIBAIJ));
  PetscCall(MatMPIBAIJSetPreallocationCSR(A, bs, ia, ja, NULL));
  PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
  PetscCall(MatDestroy(&A));

  PetscCall(PetscFree(ia));
  PetscCall(PetscFree(ja));
  PetscCall(PetscFinalize());
  return 0;
}

/*TEST

    test:
      nsize: 4

TEST*/