Actual source code: ex311.c

  1: static const char help[] = "Test MatDiagonalScale() on dense matrices with scaling Vecs of any type\n\n";

  3: // Contributed by: Steven Dargaville

  5: #include <petscmat.h>

  7: /* The scaling Vecs take their type from -vec_type, independently of the Mat type set
  8:    with -mat_type. Scaling is verified against a duplicate matrix scaled with VECSTANDARD
  9:    Vecs holding the same values: for a device matrix both scalings execute the same kernel
 10:    on identical data, so MatEqual() is exact, and no cross-type Vec or Mat copies are
 11:    needed */
 12: int main(int argc, char **args)
 13: {
 14:   Mat           A, B;
 15:   Vec           l, r, lstd, rstd;
 16:   PetscInt      m = 5, n = 4, mloc, nloc, rstart, rend;
 17:   PetscBool     equal = PETSC_FALSE, check_copies = PETSC_FALSE;
 18:   PetscLogEvent event;

 20:   PetscFunctionBeginUser;
 21:   PetscCall(PetscInitialize(&argc, &args, NULL, help));
 22:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-m", &m, NULL));
 23:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &n, NULL));
 24:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-check_copies", &check_copies, NULL));
 25:   if (check_copies) PetscCall(PetscLogDefaultBegin());
 26:   PetscCall(PetscLogEventRegister("ScaleCheck", MAT_CLASSID, &event));

 28:   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
 29:   PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, m, n));
 30:   PetscCall(MatSetType(A, MATDENSE));
 31:   PetscCall(MatSetFromOptions(A));
 32:   PetscCall(MatSetUp(A));
 33:   PetscCall(MatSetRandom(A, NULL));
 34:   PetscCall(MatDuplicate(A, MAT_COPY_VALUES, &B));
 35:   PetscCall(MatGetLocalSize(A, &mloc, &nloc));

 37:   // l and lstd match A's row layout, r and rstd its column layout
 38:   PetscCall(VecCreate(PETSC_COMM_WORLD, &l));
 39:   PetscCall(VecSetSizes(l, mloc, m));
 40:   PetscCall(VecSetFromOptions(l));
 41:   PetscCall(VecCreate(PETSC_COMM_WORLD, &lstd));
 42:   PetscCall(VecSetSizes(lstd, mloc, m));
 43:   PetscCall(VecSetType(lstd, VECSTANDARD));
 44:   PetscCall(VecCreate(PETSC_COMM_WORLD, &r));
 45:   PetscCall(VecSetSizes(r, nloc, n));
 46:   PetscCall(VecSetFromOptions(r));
 47:   PetscCall(VecCreate(PETSC_COMM_WORLD, &rstd));
 48:   PetscCall(VecSetSizes(rstd, nloc, n));
 49:   PetscCall(VecSetType(rstd, VECSTANDARD));

 51:   PetscCall(VecGetOwnershipRange(l, &rstart, &rend));
 52:   for (PetscInt i = rstart; i < rend; i++) {
 53:     PetscCall(VecSetValue(l, i, (PetscScalar)(i + 2), INSERT_VALUES));
 54:     PetscCall(VecSetValue(lstd, i, (PetscScalar)(i + 2), INSERT_VALUES));
 55:   }
 56:   PetscCall(VecGetOwnershipRange(r, &rstart, &rend));
 57:   for (PetscInt j = rstart; j < rend; j++) {
 58:     PetscCall(VecSetValue(r, j, (PetscScalar)(j + 3), INSERT_VALUES));
 59:     PetscCall(VecSetValue(rstd, j, (PetscScalar)(j + 3), INSERT_VALUES));
 60:   }
 61:   PetscCall(VecAssemblyBegin(l));
 62:   PetscCall(VecAssemblyEnd(l));
 63:   PetscCall(VecAssemblyBegin(lstd));
 64:   PetscCall(VecAssemblyEnd(lstd));
 65:   PetscCall(VecAssemblyBegin(r));
 66:   PetscCall(VecAssemblyEnd(r));
 67:   PetscCall(VecAssemblyBegin(rstd));
 68:   PetscCall(VecAssemblyEnd(rstd));

 70:   // left scaling only
 71:   PetscCall(MatDiagonalScale(A, l, NULL));
 72:   PetscCall(MatDiagonalScale(B, lstd, NULL));
 73:   PetscCall(MatEqual(A, B, &equal));
 74:   PetscCheck(equal, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Left scaling gives the wrong result");

 76:   // right scaling on top of the left scaling
 77:   PetscCall(MatDiagonalScale(A, NULL, r));
 78:   PetscCall(MatDiagonalScale(B, NULL, rstd));
 79:   PetscCall(MatEqual(A, B, &equal));
 80:   PetscCheck(equal, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Right scaling gives the wrong result");

 82:   // both sides in one call; l and r were already consumed by the calls above, so any copy logged here is the matrix
 83:   PetscCall(PetscLogEventBegin(event, 0, 0, 0, 0));
 84:   PetscCall(MatDiagonalScale(A, l, r));
 85:   PetscCall(PetscLogEventEnd(event, 0, 0, 0, 0));
 86:   PetscCall(MatDiagonalScale(B, lstd, rstd));
 87:   PetscCall(MatEqual(A, B, &equal));
 88:   PetscCheck(equal, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Two-sided scaling gives the wrong result");

 90: #if PetscDefined(HAVE_DEVICE)
 91:   /* a device matrix must never come back to the host to be scaled, and a device-resident scaling Vec must be
 92:      consumed in place; a host scaling Vec is copied to the device by design, so only the first check applies then */
 93:   if (check_copies) {
 94:     PetscEventPerfInfo info;
 95:     const PetscScalar *array;
 96:     PetscMemType       mtype;

 98:     PetscCall(VecGetArrayReadAndMemType(l, &array, &mtype));
 99:     PetscCall(VecRestoreArrayReadAndMemType(l, &array));
100:     PetscCall(PetscLogEventGetPerfInfo(PETSC_DETERMINE, event, &info));
101:     PetscCheck(info.GpuToCpuCount == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "%g unexpected GPU to CPU copies (%g bytes) in MatDiagonalScale()", info.GpuToCpuCount, info.GpuToCpuSize);
102:     PetscCheck(!PetscMemTypeDevice(mtype) || info.CpuToGpuCount == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "%g unexpected CPU to GPU copies (%g bytes) in MatDiagonalScale() with device scaling Vecs", info.CpuToGpuCount, info.CpuToGpuSize);
103:   }
104: #endif

106:   PetscCall(VecDestroy(&l));
107:   PetscCall(VecDestroy(&lstd));
108:   PetscCall(VecDestroy(&r));
109:   PetscCall(VecDestroy(&rstd));
110:   PetscCall(MatDestroy(&A));
111:   PetscCall(MatDestroy(&B));
112:   PetscCall(PetscFinalize());
113:   return 0;
114: }

116: /*TEST

118:   test:
119:     suffix: cpu
120:     nsize: {{1 2}}
121:     output_file: output/empty.out

123:   test:
124:     suffix: kokkos
125:     nsize: {{1 2}}
126:     requires: kokkos_kernels
127:     args: -vec_type kokkos
128:     output_file: output/empty.out

130:   # -check_copies requires GPU-aware MPI: without it PetscSF stages the device buffers of the
131:   # right scaling through the host, and those copies are logged inside MatDiagonalScale()

133:   testset:
134:     nsize: {{1 2}}
135:     requires: cuda
136:     args: -mat_type densecuda -vec_type {{cuda standard}}
137:     output_file: output/empty.out

139:     test:
140:       suffix: cuda

142:     test:
143:       suffix: cuda_copies
144:       requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
145:       args: -check_copies

147:   testset:
148:     nsize: {{1 2}}
149:     requires: cuda kokkos_kernels
150:     args: -mat_type densecuda -vec_type kokkos
151:     output_file: output/empty.out

153:     test:
154:       suffix: densecuda_vec_kokkos

156:     test:
157:       suffix: densecuda_vec_kokkos_copies
158:       requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
159:       args: -check_copies

161:   testset:
162:     nsize: {{1 2}}
163:     requires: hip
164:     args: -mat_type densehip -vec_type {{hip standard}}
165:     output_file: output/empty.out

167:     test:
168:       suffix: hip

170:     test:
171:       suffix: hip_copies
172:       requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
173:       args: -check_copies

175:   testset:
176:     nsize: {{1 2}}
177:     requires: hip kokkos_kernels
178:     args: -mat_type densehip -vec_type kokkos
179:     output_file: output/empty.out

181:     test:
182:       suffix: densehip_vec_kokkos

184:     test:
185:       suffix: densehip_vec_kokkos_copies
186:       requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
187:       args: -check_copies

189: TEST*/