Actual source code: ex319.c

  1: static char help[] = "Tests MatMatMult() of an MPIAIJ-family matrix with an MPIDENSE-family matrix, including the batched path, against a host reference\n\n";

  3: // Contributed by: Steven Dargaville

  5: #include <petscmat.h>
  6: #include <petsc/private/petscimpl.h>

  8: /* The product of the matrices whose types are set with -A_mat_type and -B_mat_type is compared against
  9:    the product of host MATAIJ and MATDENSE copies holding the same values */
 10: static PetscErrorCode CheckProduct(Mat C, Mat Ch, const char *label)
 11: {
 12:   Mat       Ct;
 13:   PetscReal nrm, err;

 15:   PetscFunctionBeginUser;
 16:   /* a host matrix is duplicated, a device one is copied back to the host */
 17:   PetscCall(MatConvert(C, MATDENSE, MAT_INITIAL_MATRIX, &Ct));
 18:   PetscCall(MatNorm(Ch, NORM_INFINITY, &nrm));
 19:   PetscCall(MatAXPY(Ct, -1.0, Ch, SAME_NONZERO_PATTERN));
 20:   PetscCall(MatNorm(Ct, NORM_INFINITY, &err));
 21:   PetscCheck(err <= 100.0 * PETSC_SMALL * PetscMax(nrm, 1.0), PETSC_COMM_WORLD, PETSC_ERR_PLIB, "%s: error %g relative to %g", label, (double)err, (double)nrm);
 22:   PetscCall(MatDestroy(&Ct));
 23:   PetscFunctionReturn(PETSC_SUCCESS);
 24: }

 26: int main(int argc, char **args)
 27: {
 28:   Mat           A, B, C, Ah, Bh, Ch;
 29:   char          atype[256], btype[256];
 30:   PetscScalar   vals[4];
 31:   PetscInt      m = 20, n = 7, rstart, rend, ncols, cols[4];
 32:   PetscBool     block_diagonal = PETSC_FALSE, check_copies = PETSC_FALSE, tridiagonal = PETSC_FALSE;
 33:   PetscLogEvent event;
 34:   PetscMPIInt   size;

 36:   PetscFunctionBeginUser;
 37:   PetscCall(PetscInitialize(&argc, &args, NULL, help));
 38:   PetscCall(PetscStrncpy(atype, MATAIJ, sizeof(atype)));
 39:   PetscCall(PetscStrncpy(btype, MATDENSE, sizeof(btype)));
 40:   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "MatMatMult() test options", "Mat");
 41:   PetscCall(PetscOptionsInt("-m", "Number of rows and columns of A", NULL, m, &m, NULL));
 42:   PetscCall(PetscOptionsInt("-n", "Number of columns of B", NULL, n, &n, NULL));
 43:   PetscCall(PetscOptionsFList("-A_mat_type", "Type of A", "MatSetType", MatList, atype, atype, sizeof(atype), NULL));
 44:   PetscCall(PetscOptionsFList("-B_mat_type", "Type of B", "MatSetType", MatList, btype, btype, sizeof(btype), NULL));
 45:   PetscCall(PetscOptionsBool("-block_diagonal", "Only keep the columns of A owned by this process, so the off-diagonal blocks have no columns", NULL, block_diagonal, &block_diagonal, NULL));
 46:   PetscCall(PetscOptionsBool("-tridiagonal", "Make A tridiagonal, so that only the first and last row of each off-diagonal block are nonempty and the block is stored in compressed-row format", NULL, tridiagonal, &tridiagonal, NULL));
 47:   PetscCall(PetscOptionsBool("-check_copies", "Check that a product of device matrices does not copy between the host and the device", NULL, check_copies, &check_copies, NULL));
 48:   PetscOptionsEnd();
 49:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
 50:   /* PetscSF stages device buffers through the host when GPU-aware MPI is disabled. */
 51:   if (check_copies && size > 1 && !use_gpu_aware_mpi) {
 52:     PetscCall(PetscInfo(NULL, "Skipping copy checks because GPU-aware MPI is disabled; numerical checks remain enabled\n"));
 53:     check_copies = PETSC_FALSE;
 54:   }
 55:   if (check_copies) PetscCall(PetscLogDefaultBegin());
 56:   PetscCall(PetscLogEventRegister("ProductCheck", MAT_CLASSID, &event));

 58:   /* host reference matrices, deliberately not calling MatSetFromOptions() so they keep their host types */
 59:   PetscCall(MatCreate(PETSC_COMM_WORLD, &Ah));
 60:   PetscCall(MatSetSizes(Ah, PETSC_DECIDE, PETSC_DECIDE, m, m));
 61:   PetscCall(MatSetType(Ah, MATAIJ));
 62:   PetscCall(MatSeqAIJSetPreallocation(Ah, 4, NULL));
 63:   PetscCall(MatMPIAIJSetPreallocation(Ah, 4, NULL, 4, NULL));
 64:   PetscCall(MatGetOwnershipRange(Ah, &rstart, &rend));
 65:   for (PetscInt i = rstart; i < rend; i++) {
 66:     if (tridiagonal) {
 67:       ncols = 0;
 68:       if (i > 0) {
 69:         cols[ncols] = i - 1;
 70:         vals[ncols] = -1.0;
 71:         ncols++;
 72:       }
 73:       cols[ncols] = i;
 74:       vals[ncols] = 2.0;
 75:       ncols++;
 76:       if (i < m - 1) {
 77:         cols[ncols] = i + 1;
 78:         vals[ncols] = -1.0;
 79:         ncols++;
 80:       }
 81:     } else {
 82:       ncols   = 4;
 83:       cols[0] = i;
 84:       cols[1] = (i + 1) % m;
 85:       cols[2] = (i + m / 2) % m;
 86:       cols[3] = (i * 7 + 3) % m;
 87:       for (PetscInt k = 0; k < ncols; k++) vals[k] = 1.0 + 0.1 * (PetscReal)(i + cols[k]);
 88:     }
 89:     /* duplicate columns accumulate with ADD_VALUES */
 90:     for (PetscInt k = 0; k < ncols; k++) {
 91:       if (block_diagonal && (cols[k] < rstart || cols[k] >= rend)) continue;
 92:       PetscCall(MatSetValue(Ah, i, cols[k], vals[k], ADD_VALUES));
 93:     }
 94:   }
 95:   PetscCall(MatAssemblyBegin(Ah, MAT_FINAL_ASSEMBLY));
 96:   PetscCall(MatAssemblyEnd(Ah, MAT_FINAL_ASSEMBLY));

 98:   PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, m, n, NULL, &Bh));
 99:   PetscCall(MatGetOwnershipRange(Bh, &rstart, &rend));
100:   for (PetscInt i = rstart; i < rend; i++) {
101:     for (PetscInt j = 0; j < n; j++) PetscCall(MatSetValue(Bh, i, j, 0.5 + PetscSinReal((PetscReal)(i + 3 * j)), INSERT_VALUES));
102:   }
103:   PetscCall(MatAssemblyBegin(Bh, MAT_FINAL_ASSEMBLY));
104:   PetscCall(MatAssemblyEnd(Bh, MAT_FINAL_ASSEMBLY));

106:   /* the matrices under test hold the same values with the requested types */
107:   PetscCall(MatDuplicate(Ah, MAT_COPY_VALUES, &A));
108:   PetscCall(MatConvert(A, atype, MAT_INPLACE_MATRIX, &A));
109:   PetscCall(MatDuplicate(Bh, MAT_COPY_VALUES, &B));
110:   PetscCall(MatConvert(B, btype, MAT_INPLACE_MATRIX, &B));

112:   PetscCall(MatMatMult(A, B, MAT_INITIAL_MATRIX, PETSC_DETERMINE, &C));
113:   PetscCall(MatMatMult(Ah, Bh, MAT_INITIAL_MATRIX, PETSC_DETERMINE, &Ch));
114:   PetscCall(CheckProduct(C, Ch, "MAT_INITIAL_MATRIX"));

116:   /* change the values of both pairs identically and reuse the products */
117:   PetscCall(MatScale(A, 2.0));
118:   PetscCall(MatShift(A, 1.0));
119:   PetscCall(MatScale(B, -0.5));
120:   PetscCall(MatScale(Ah, 2.0));
121:   PetscCall(MatShift(Ah, 1.0));
122:   PetscCall(MatScale(Bh, -0.5));

124:   PetscCall(MatMatMult(A, B, MAT_REUSE_MATRIX, PETSC_DETERMINE, &C));
125:   PetscCall(MatMatMult(Ah, Bh, MAT_REUSE_MATRIX, PETSC_DETERMINE, &Ch));
126:   PetscCall(CheckProduct(C, Ch, "MAT_REUSE_MATRIX"));

128:   /* the reuse above brings the modified values to the device (MatShift() may run on the host), so any copy logged
129:      in a further reuse is made by the product itself */
130:   PetscCall(PetscLogEventBegin(event, 0, 0, 0, 0));
131:   PetscCall(MatMatMult(A, B, MAT_REUSE_MATRIX, PETSC_DETERMINE, &C));
132:   PetscCall(PetscLogEventEnd(event, 0, 0, 0, 0));
133:   PetscCall(CheckProduct(C, Ch, "MAT_REUSE_MATRIX again"));

135: #if PetscDefined(HAVE_DEVICE)
136:   /* a product of device matrices must run entirely on the device, in either direction and with any number of processes */
137:   if (check_copies) {
138:     PetscEventPerfInfo info;

140:     PetscCall(PetscLogEventGetPerfInfo(PETSC_DETERMINE, event, &info));
141:     PetscCheck(info.GpuToCpuCount == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "%g unexpected GPU to CPU copies (%g bytes) in MatMatMult()", info.GpuToCpuCount, info.GpuToCpuSize);
142:     PetscCheck(info.CpuToGpuCount == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "%g unexpected CPU to GPU copies (%g bytes) in MatMatMult()", info.CpuToGpuCount, info.CpuToGpuSize);
143:   }
144: #endif

146:   PetscCall(MatDestroy(&A));
147:   PetscCall(MatDestroy(&B));
148:   PetscCall(MatDestroy(&C));
149:   PetscCall(MatDestroy(&Ah));
150:   PetscCall(MatDestroy(&Bh));
151:   PetscCall(MatDestroy(&Ch));
152:   PetscCall(PetscFinalize());
153:   return 0;
154: }

156: /*TEST

158:   testset:
159:     output_file: output/empty.out

161:     test:
162:       suffix: host
163:       nsize: {{1 3}}
164:       args: -block_diagonal {{0 1}}

166:     test:
167:       suffix: host_batch
168:       nsize: 3
169:       args: -matproduct_batch_size {{2 3 4}} -tridiagonal {{0 1}}

171:     # the local block of B is a host MATSEQDENSE here, so the scatter of the off-process rows
172:     # does not go through device memory and -use_gpu_aware_mpi would only duplicate the runs
173:     test:
174:       suffix: kokkos
175:       requires: kokkos_kernels
176:       nsize: {{1 3}}
177:       args: -A_mat_type aijkokkos

179:     test:
180:       suffix: kokkos_batch
181:       requires: kokkos_kernels
182:       nsize: 3
183:       args: -A_mat_type aijkokkos -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}}

185:   # -check_copies in parallel requires GPU-aware MPI: without it PetscSF stages the device buffers of the
186:   # scatter through the host, and those copies are logged inside MatMatMult(). The plain device tests run
187:   # with that staging, which is valid on every build, and the _copies_par tests cover the GPU-aware path

189:   testset:
190:     requires: cuda kokkos_kernels
191:     args: -A_mat_type aijkokkos -B_mat_type densecuda
192:     output_file: output/empty.out

194:     test:
195:       suffix: kokkos_cuda
196:       nsize: 3
197:       args: -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0

199:     test:
200:       suffix: kokkos_cuda_copies
201:       nsize: 1
202:       requires: defined(PETSC_USE_LOG)
203:       args: -check_copies

205:     test:
206:       suffix: kokkos_cuda_copies_par
207:       nsize: 3
208:       requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
209:       args: -check_copies -matproduct_batch_size {{0 3}}

211:   testset:
212:     requires: cuda
213:     args: -A_mat_type aijcusparse
214:     output_file: output/empty.out

216:     test:
217:       suffix: cuda
218:       nsize: 3
219:       args: -B_mat_type {{dense densecuda}} -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0

221:     test:
222:       suffix: cuda_copies
223:       nsize: 1
224:       requires: defined(PETSC_USE_LOG)
225:       args: -B_mat_type densecuda -check_copies

227:     test:
228:       suffix: cuda_copies_par
229:       nsize: 3
230:       requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
231:       args: -B_mat_type densecuda -check_copies -matproduct_batch_size {{0 3}}

233:   testset:
234:     requires: hip kokkos_kernels
235:     args: -A_mat_type aijkokkos -B_mat_type densehip
236:     output_file: output/empty.out

238:     test:
239:       suffix: kokkos_hip
240:       nsize: 3
241:       args: -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0

243:     test:
244:       suffix: kokkos_hip_copies
245:       nsize: 1
246:       requires: defined(PETSC_USE_LOG)
247:       args: -check_copies

249:     test:
250:       suffix: kokkos_hip_copies_par
251:       nsize: 3
252:       requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
253:       args: -check_copies -matproduct_batch_size {{0 3}}

255:   testset:
256:     requires: hip
257:     args: -A_mat_type aijhipsparse
258:     output_file: output/empty.out

260:     test:
261:       suffix: hip
262:       nsize: 3
263:       args: -B_mat_type {{dense densehip}} -matproduct_batch_size {{0 3}} -tridiagonal {{0 1}} -use_gpu_aware_mpi 0

265:     test:
266:       suffix: hip_copies
267:       nsize: 1
268:       requires: defined(PETSC_USE_LOG)
269:       args: -B_mat_type densehip -check_copies

271:     test:
272:       suffix: hip_copies_par
273:       nsize: 3
274:       requires: defined(PETSC_HAVE_MPI_GPU_AWARE) defined(PETSC_USE_LOG)
275:       args: -B_mat_type densehip -check_copies -matproduct_batch_size {{0 3}}

277: TEST*/