Actual source code: ex17.c
1: static const char help[] = "Test CUDA point-block Jacobi on nonblocking device contexts.\n";
3: #include <petscksp.h>
4: #include <petscdevice_cuda.h>
6: static PetscErrorCode TestApply(PC pc, PC reference, Vec x, Vec y, Vec hx, Vec hy, PetscDeviceContext dctx, PetscBool barrier)
7: {
8: const PetscScalar *actual, *expected;
9: PetscInt n;
10: cublasHandle_t handle;
11: cublasPointerMode_t savedMode, actualMode;
12: const cublasPointerMode_t modes[] = {CUBLAS_POINTER_MODE_HOST, CUBLAS_POINTER_MODE_DEVICE};
14: PetscFunctionBeginUser;
15: PetscCall(VecGetLocalSize(x, &n));
16: PetscCall(PetscCUBLASGetHandle(&handle));
17: PetscCallCUBLAS(cublasGetPointerMode(handle, &savedMode));
18: for (size_t mode = 0; mode < PETSC_STATIC_ARRAY_LENGTH(modes); ++mode) {
19: for (PetscInt transpose = 0; transpose < 2; ++transpose) {
20: PetscCall(VecSet(x, 1));
21: PetscCall(VecSet(y, 0));
22: if (transpose) PetscCall(PCApplyTranspose(pc, x, y));
23: else PetscCall(PCApply(pc, x, y));
24: PetscCall(PetscDeviceContextSynchronize(dctx));
26: PetscCall(VecSet(hx, 2));
27: if (transpose) PetscCall(PCApplyTranspose(reference, hx, hy));
28: else PetscCall(PCApply(reference, hx, hy));
29: PetscCall(VecScale(hy, 3));
31: PetscCall(PetscDeviceContextDelay(dctx, 0.05));
32: PetscCall(VecScale(x, 2));
33: PetscCallCUBLAS(cublasSetPointerMode(handle, modes[mode]));
34: // A separate delay checks completion even when VecScale() honors the barrier.
35: if (barrier) PetscCall(PetscDeviceContextDelay(dctx, 0.05));
36: if (transpose) PetscCall(PCApplyTranspose(pc, x, y));
37: else PetscCall(PCApply(pc, x, y));
38: PetscCallCUBLAS(cublasGetPointerMode(handle, &actualMode));
39: PetscCheck(actualMode == modes[mode], PETSC_COMM_SELF, PETSC_ERR_PLIB, "PCApply() changed the cuBLAS pointer mode");
40: PetscCallCUBLAS(cublasSetPointerMode(handle, savedMode));
41: if (barrier) {
42: PetscBool idle;
44: PetscCall(PetscDeviceContextQueryIdle(dctx, &idle));
45: PetscCheck(idle, PETSC_COMM_SELF, PETSC_ERR_PLIB, "PCApply() returned with work pending on a barrier context");
46: }
47: PetscCall(VecScale(y, 3));
48: PetscCall(PetscDeviceContextSynchronize(dctx));
49: PetscCall(VecGetArrayRead(y, &actual));
50: PetscCall(VecGetArrayRead(hy, &expected));
51: for (PetscInt i = 0; i < n; ++i) PetscCheck(PetscAbsScalar(actual[i] - expected[i]) < 100 * PETSC_MACHINE_EPSILON, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Incorrect point-block Jacobi result at local entry %" PetscInt_FMT, i);
52: PetscCall(VecRestoreArrayRead(hy, &expected));
53: PetscCall(VecRestoreArrayRead(y, &actual));
54: }
55: }
56: PetscFunctionReturn(PETSC_SUCCESS);
57: }
59: int main(int argc, char **argv)
60: {
61: Mat A, host;
62: PC pc, reference;
63: Vec x, y, hx, hy;
64: PetscDeviceContext saved, dctx;
65: const PetscStreamType streams[] = {PETSC_STREAM_NONBLOCKING, PETSC_STREAM_NONBLOCKING_WITH_BARRIER};
66: PetscInt start;
67: const PetscInt blocks[] = {2, 4};
68: const PetscScalar entries[6][6] = {
69: {2, 1, 0, 0, 0, 0 },
70: {0, 3, 0, 0, 0, 0 },
71: {0, 0, 4, 1, 0, 0 },
72: {0, 0, 2, 5, 0, 0 },
73: {0, 0, 0, 0, 6, -1},
74: {0, 0, 0, 0, 2, 7 }
75: };
76: PCType type;
78: PetscFunctionBeginUser;
79: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
80: PetscCall(PetscDeviceContextGetCurrentContext(&saved));
81: PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, 6, 6, PETSC_DECIDE, PETSC_DECIDE, 2, NULL, 0, NULL, &host));
82: PetscCall(MatGetOwnershipRange(host, &start, NULL));
83: for (PetscInt i = 0; i < 6; ++i)
84: for (PetscInt j = 0; j < 6; ++j)
85: if (entries[i][j] != 0) PetscCall(MatSetValue(host, start + i, start + j, entries[i][j], INSERT_VALUES));
86: PetscCall(MatAssemblyBegin(host, MAT_FINAL_ASSEMBLY));
87: PetscCall(MatAssemblyEnd(host, MAT_FINAL_ASSEMBLY));
88: PetscCall(MatSetBlockSize(host, 2));
89: PetscCall(MatSetVariableBlockSizes(host, 2, blocks));
90: PetscCall(MatConvert(host, MATAIJCUSPARSE, MAT_INITIAL_MATRIX, &A));
91: PetscCall(MatSetVariableBlockSizes(A, 2, blocks));
92: PetscCall(MatCreateVecs(A, &x, &y));
93: PetscCall(MatCreateVecs(host, &hx, &hy));
94: PetscCall(PCCreate(PETSC_COMM_WORLD, &pc));
95: PetscCall(PCSetType(pc, PCVPBJACOBI));
96: PetscCall(PCSetOperators(pc, A, A));
97: PetscCall(PCSetFromOptions(pc));
98: PetscCall(PCGetType(pc, &type));
99: PetscCall(PCCreate(PETSC_COMM_WORLD, &reference));
100: PetscCall(PCSetType(reference, type));
101: PetscCall(PCSetOperators(reference, host, host));
102: PetscCall(PetscDeviceContextSynchronize(saved));
103: for (size_t k = 0; k < PETSC_STATIC_ARRAY_LENGTH(streams); ++k) {
104: PetscCall(PetscDeviceContextDuplicate(saved, &dctx));
105: PetscCall(PetscDeviceContextSetStreamType(dctx, streams[k]));
106: PetscCall(PetscDeviceContextSetUp(dctx));
107: PetscCall(PetscDeviceContextSetCurrentContext(dctx));
108: for (PetscInt update = 0; update < 2; ++update) {
109: PetscCall(MatScale(host, 2));
110: PetscCall(MatScale(A, 2));
111: PetscCall(PCSetUp(pc));
112: PetscCall(PCSetUp(reference));
113: PetscCall(TestApply(pc, reference, x, y, hx, hy, dctx, (PetscBool)(streams[k] == PETSC_STREAM_NONBLOCKING_WITH_BARRIER)));
114: }
115: PetscCall(PetscDeviceContextSynchronize(dctx));
116: PetscCall(PetscDeviceContextSetCurrentContext(saved));
117: PetscCall(PetscDeviceContextDestroy(&dctx));
118: }
119: PetscCall(PCDestroy(&reference));
120: PetscCall(PCDestroy(&pc));
121: PetscCall(VecDestroy(&hy));
122: PetscCall(VecDestroy(&hx));
123: PetscCall(VecDestroy(&y));
124: PetscCall(VecDestroy(&x));
125: PetscCall(MatDestroy(&A));
126: PetscCall(MatDestroy(&host));
127: PetscCall(PetscFinalize());
128: return 0;
129: }
131: /*TEST
133: build:
134: requires: cuda
136: test:
137: suffix: cuda
138: args: -pc_type {{pbjacobi vpbjacobi}}
139: nsize: {{1 2}}
140: output_file: output/empty.out
142: TEST*/