PETSc Preallocation

#include "petscmat.h"
#include "petscviewer.h"

int main(int argc, char **argv)
{
  const int size = 10;
  const double nz_ratio = 0.8;
  const bool doPrealloc = true;
  printf("nz_ratio = %f\n", nz_ratio);

  PetscInt colNum = 0;
  PetscInt colIdx[size];
  PetscScalar rowVals[size];

  PetscInt d_nnz[size], o_nnz[size];

  PetscInitialize(&argc, &argv, "", NULL);
  PetscErrorCode ierr = 0;

  Mat A;
  ierr = MatCreate(PETSC_COMM_WORLD, &A); CHKERRQ(ierr);
  MatSetType(A, MATAIJ); CHKERRQ(ierr);
  ierr = MatSetSizes(A, size, size, PETSC_DECIDE, PETSC_DECIDE); CHKERRQ(ierr);
  ierr = MatSetUp(A); CHKERRQ(ierr);
  MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);

  PetscInt range_start, range_end;
  MatGetOwnershipRange(A, &range_start, &range_end);

  PetscInt col_range_start, col_range_end;
  MatGetOwnershipRangeColumn(A, &col_range_start, &col_range_end);
  printf("Column Ownership Range = %i / %i\n", col_range_start, col_range_end);

  if (doPrealloc) {
    for (int row = range_start; row < range_end; row++) {
      // printf("Row %i\n", row);
      d_nnz[row] = 0;
      o_nnz[row] = 0;
      int colStart = row - nz_ratio/2 * size;
      int colEnd   = row + nz_ratio/2 * size;
      colStart = (colStart < 0)    ? 0    : colStart;
      colEnd   = (colEnd   > size) ? size : colEnd;
      // printf("colStart = %i, colEnd = %i\n", colStart, colEnd);
      for (int col = 0; col < size; col++) {
        // printf("Column %i ", col);
        if (col >= colStart and col < colEnd) {
          // printf("\n");
          if (col >= col_range_start and col < col_range_end) {
            // printf("  Adding prealloc diagonal\n");
            d_nnz[row]++;
          } else {
            // printf("  Adding prealloc non-diagonal\n");
            o_nnz[row]++;
          }
        }
      }
      // printf("d_nnz[row] = %i, o_nnz[row] = %i\n", d_nnz[row], o_nnz[row]);
    }
    ierr = MatMPIAIJSetPreallocation(A, 0, d_nnz, 0, o_nnz); CHKERRQ(ierr);
  }

  // Printing the prealloc array
  for (int i = 0; i < size; i++) {
    printf("%i o_nnz = %i, d_nnz = %i\n", i, o_nnz[i], d_nnz[i]);
  }

  printf("End preallocation loop\n");

  for (int row = range_start; row < range_end; row++) {
    int colStart = row - nz_ratio/2 * size;
    int colEnd   = row + nz_ratio/2 * size;
    colStart = (colStart < 0)    ? 0      : colStart;
    colEnd   = (colEnd   > size) ? size   : colEnd;
    for (int col = 0; col < size; col++) {
      if (col >= colStart and col < colEnd) {
        colIdx[colNum] = col;
        rowVals[colNum] = row * col;
        // printf("colNum = %i\n", colNum);
        colNum++;
      }
    }
    printf("Inserting %i elements in row %i\n", colNum, row);
    ierr = MatSetValues(A, 1,  &row, colNum, colIdx, rowVals, INSERT_VALUES); CHKERRQ(ierr);
    colNum = 0;
  }

  ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);

  PetscFinalize();
  return 0;
}