1: static char help[] = "Tests MatDenseUpdateColumnLayout().\n\n";

  3: #include <petscmat.h>

  5: static PetscErrorCode SetLinearValues(Vec x)
  6: {
  7:   PetscInt rstart, rend;

  9:   PetscFunctionBeginUser;
 10:   PetscCall(VecGetOwnershipRange(x, &rstart, &rend));
 11:   for (PetscInt i = rstart; i < rend; i++) PetscCall(VecSetValue(x, i, (PetscScalar)(i + 1), INSERT_VALUES));
 12:   PetscCall(VecAssemblyBegin(x));
 13:   PetscCall(VecAssemblyEnd(x));
 14:   PetscFunctionReturn(PETSC_SUCCESS);
 15: }

 17: int main(int argc, char **args)
 18: {
 19:   Mat         A, S;
 20:   Vec         x, y, yref;
 21:   PetscLayout clayout, cmap, scmap;
 22:   PetscReal   norm;
 23:   PetscInt    M, N, nloc, rstart, rend, cstart, cend, xstart, xend;
 24:   PetscBool   same;
 25:   PetscMPIInt rank, size;

 27:   PetscFunctionBeginUser;
 28:   PetscCall(PetscInitialize(&argc, &args, NULL, help));
 29:   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
 30:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
 31:   M    = 2 * size + 1;
 32:   N    = size * (size + 1) / 2;
 33:   nloc = rank + 1;

 35:   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
 36:   PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, M, N));
 37:   PetscCall(MatSetType(A, MATDENSE));
 38:   PetscCall(MatSetFromOptions(A));
 39:   PetscCall(MatSetUp(A));
 40:   PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
 41:   for (PetscInt i = rstart; i < rend; i++) {
 42:     for (PetscInt j = 0; j < N; j++) PetscCall(MatSetValue(A, i, j, (PetscScalar)(1 + i + 2 * j), INSERT_VALUES));
 43:   }
 44:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
 45:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));

 47:   /* Cache a full-column submatrix using the original column layout. */
 48:   PetscCall(MatDenseGetSubMatrix(A, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, &S));
 49:   PetscCall(MatDenseRestoreSubMatrix(A, &S));

 51:   /* Build the MPI dense matrix-vector scatter for the original column layout. */
 52:   PetscCall(MatCreateVecs(A, &x, &y));
 53:   PetscCall(SetLinearValues(x));
 54:   PetscCall(MatMult(A, x, y));
 55:   PetscCall(VecDuplicate(y, &yref));
 56:   PetscCall(VecCopy(y, yref));
 57:   PetscCall(VecDestroy(&x));

 59:   PetscCall(PetscLayoutCreateFromSizes(PETSC_COMM_WORLD, nloc, N, 1, &clayout));
 60:   PetscCall(PetscLayoutGetRange(clayout, &cstart, &cend));
 61:   PetscCall(MatDenseUpdateColumnLayout(A, clayout));
 62:   PetscCall(MatGetLayouts(A, NULL, &cmap));
 63:   PetscCheck(cmap == clayout, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Matrix did not adopt the new column layout");
 64:   PetscCall(PetscLayoutDestroy(&clayout));

 66:   PetscCall(MatCreateVecs(A, &x, NULL));
 67:   PetscCall(VecGetOwnershipRange(x, &xstart, &xend));
 68:   PetscCheck(xstart == cstart && xend == cend, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Right vector ownership range [%" PetscInt_FMT ",%" PetscInt_FMT ") does not match the new column layout [%" PetscInt_FMT ",%" PetscInt_FMT ")", xstart, xend, cstart, cend);
 69:   PetscCall(SetLinearValues(x));

 71:   PetscCall(MatDenseGetSubMatrix(A, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, &S));
 72:   PetscCall(MatGetLayouts(S, NULL, &scmap));
 73:   PetscCall(PetscLayoutCompare(scmap, cmap, &same));
 74:   PetscCheck(same, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Submatrix did not adopt the new column layout");
 75:   PetscCall(MatMult(S, x, y));
 76:   PetscCall(MatDenseRestoreSubMatrix(A, &S));
 77:   PetscCall(VecAXPY(y, -1.0, yref));
 78:   PetscCall(VecNorm(y, NORM_INFINITY, &norm));
 79:   PetscCheck(norm <= 100.0 * PETSC_MACHINE_EPSILON, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Submatrix MatMult() changed after updating the column layout, error %g", (double)norm);

 81:   PetscCall(MatMult(A, x, y));
 82:   PetscCall(VecAXPY(y, -1.0, yref));
 83:   PetscCall(VecNorm(y, NORM_INFINITY, &norm));
 84:   PetscCheck(norm <= 100.0 * PETSC_MACHINE_EPSILON, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "MatMult() changed after updating the column layout, error %g", (double)norm);

 86:   PetscCall(VecDestroy(&x));
 87:   PetscCall(VecDestroy(&y));
 88:   PetscCall(VecDestroy(&yref));
 89:   PetscCall(MatDestroy(&A));
 90:   PetscCall(PetscFinalize());
 91:   return 0;
 92: }

 94: /*TEST

 96:   testset:
 97:     nsize: {{1 2}}
 98:     output_file: output/empty.out

100:     test:
101:       suffix: cpu

103:     test:
104:       requires: cuda
105:       suffix: cuda
106:       args: -mat_type densecuda

108:     test:
109:       requires: hip
110:       suffix: hip
111:       args: -mat_type densehip

113: TEST*/