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*/