1: static const char help[] = "Compare LMVM updates, products, and solves with host results.\n";

  3: #include <petscksp.h>
  4: #include <petscdevice.h>

  6: static PetscErrorCode CheckResult(Vec actual, Vec expected)
  7: {
  8:   const PetscScalar *a, *e;
  9:   PetscInt           n;

 11:   PetscFunctionBeginUser;
 12:   PetscCall(VecGetLocalSize(actual, &n));
 13:   PetscCall(VecGetArrayRead(actual, &a));
 14:   PetscCall(VecGetArrayRead(expected, &e));
 15:   for (PetscInt i = 0; i < n; ++i)
 16:     PetscCheck(PetscAbsScalar(a[i] - e[i]) <= 500 * PETSC_MACHINE_EPSILON * PetscMax(1.0, PetscAbsScalar(e[i])), PETSC_COMM_SELF, PETSC_ERR_PLIB, "LMVM result differs at local entry %" PetscInt_FMT ": error %g", i, (double)PetscAbsScalar(a[i] - e[i]));
 17:   PetscCall(VecRestoreArrayRead(expected, &e));
 18:   PetscCall(VecRestoreArrayRead(actual, &a));
 19:   PetscFunctionReturn(PETSC_SUCCESS);
 20: }

 22: static PetscErrorCode TestLMVM(Vec x, PetscDeviceContext dctx)
 23: {
 24:   Mat          B, reference;
 25:   Vec          f, result, hx, hf, expected;
 26:   PetscInt     n, N, start;
 27:   PetscScalar *xa, *fa;

 29:   PetscFunctionBeginUser;
 30:   PetscCall(VecGetLocalSize(x, &n));
 31:   PetscCall(VecGetSize(x, &N));
 32:   PetscCall(VecGetOwnershipRange(x, &start, NULL));
 33:   PetscCall(VecDuplicate(x, &f));
 34:   PetscCall(VecDuplicate(x, &result));
 35:   PetscCall(VecCreate(PETSC_COMM_WORLD, &hx));
 36:   PetscCall(VecSetSizes(hx, n, N));
 37:   PetscCall(VecSetType(hx, VECSTANDARD));
 38:   PetscCall(VecDuplicate(hx, &hf));
 39:   PetscCall(VecDuplicate(hx, &expected));
 40:   PetscCall(MatCreate(PETSC_COMM_WORLD, &B));
 41:   PetscCall(MatSetSizes(B, n, n, N, N));
 42:   PetscCall(MatSetType(B, MATLMVMDBFGS));
 43:   PetscCall(MatSetFromOptions(B));
 44:   PetscCall(MatLMVMSetHistorySize(B, 3));
 45:   PetscCall(MatLMVMAllocate(B, x, f));
 46:   PetscCall(MatCreate(PETSC_COMM_WORLD, &reference));
 47:   PetscCall(MatSetSizes(reference, n, n, N, N));
 48:   PetscCall(MatSetType(reference, MATLMVMDBFGS));
 49:   PetscCall(MatSetFromOptions(reference));
 50:   PetscCall(MatLMVMSetHistorySize(reference, 3));
 51:   PetscCall(MatLMVMAllocate(reference, hx, hf));
 52:   PetscCall(VecSet(hx, 0));
 53:   for (PetscInt update = 0; update < 8; ++update) {
 54:     PetscCall(VecGetArray(hx, &xa));
 55:     PetscCall(VecGetArray(hf, &fa));
 56:     for (PetscInt i = 0; i < n; ++i) {
 57:       xa[i] += PetscSinReal((start + i + 1) * (update + 1));
 58:       fa[i] = (2.0 + (PetscReal)(start + i) / N) * xa[i];
 59:     }
 60:     PetscCall(VecRestoreArray(hf, &fa));
 61:     PetscCall(VecRestoreArray(hx, &xa));
 62:     PetscCall(VecCopy(hx, x));
 63:     PetscCall(VecCopy(hf, f));
 64:     PetscCall(MatLMVMUpdate(reference, hx, hf));
 65:     PetscCall(MatLMVMUpdate(B, x, f));
 66:     // Repeated applications exercise cached products as well as new history.
 67:     for (PetscInt repeat = 0; repeat < 2; ++repeat) {
 68:       for (PetscInt solve = 0; solve < 2; ++solve) {
 69:         if (solve) PetscCall(MatSolve(reference, hf, expected));
 70:         else PetscCall(MatMult(reference, hf, expected));
 71:         if (solve) PetscCall(MatSolve(B, f, result));
 72:         else PetscCall(MatMult(B, f, result));
 73:         PetscCall(PetscDeviceContextSynchronize(dctx));
 74:         PetscCall(CheckResult(result, expected));
 75:       }
 76:     }
 77:   }
 78:   PetscCall(MatDestroy(&reference));
 79:   PetscCall(MatDestroy(&B));
 80:   PetscCall(VecDestroy(&expected));
 81:   PetscCall(VecDestroy(&hf));
 82:   PetscCall(VecDestroy(&hx));
 83:   PetscCall(VecDestroy(&result));
 84:   PetscCall(VecDestroy(&f));
 85:   PetscFunctionReturn(PETSC_SUCCESS);
 86: }

 88: int main(int argc, char **argv)
 89: {
 90:   Vec                   x;
 91:   PetscDeviceContext    saved, dctx;
 92:   const PetscStreamType streams[] = {PETSC_STREAM_DEFAULT};

 94:   PetscFunctionBeginUser;
 95:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
 96:   PetscCall(KSPInitializePackage());
 97:   PetscCall(VecCreate(PETSC_COMM_WORLD, &x));
 98:   PetscCall(VecSetSizes(x, PETSC_DECIDE, 8));
 99:   PetscCall(VecSetFromOptions(x));
100:   PetscCall(PetscDeviceContextGetCurrentContext(&saved));
101:   PetscCall(PetscDeviceContextSynchronize(saved));
102:   for (size_t k = 0; k < PETSC_STATIC_ARRAY_LENGTH(streams); ++k) {
103:     PetscCall(PetscDeviceContextDuplicate(saved, &dctx));
104:     PetscCall(PetscDeviceContextSetStreamType(dctx, streams[k]));
105:     PetscCall(PetscDeviceContextSetUp(dctx));
106:     PetscCall(PetscDeviceContextSetCurrentContext(dctx));
107:     PetscCall(TestLMVM(x, dctx));
108:     PetscCall(PetscDeviceContextSynchronize(dctx));
109:     PetscCall(PetscDeviceContextSetCurrentContext(saved));
110:     PetscCall(PetscDeviceContextDestroy(&dctx));
111:   }
112:   PetscCall(VecDestroy(&x));
113:   PetscCall(PetscFinalize());
114:   return 0;
115: }

117: /*TEST

119:   testset:
120:     nsize: {{1 2}}
121:     output_file: output/empty.out
122:     args: -mat_type {{lmvmdbfgs lmvmddfp lmvmdqn lmvmbfgs lmvmdfp}}
123:     test:
124:       suffix: cuda
125:       requires: cuda
126:       args: -vec_type cuda
127:     test:
128:       suffix: hip
129:       requires: hip !complex
130:       args: -vec_type hip

132: TEST*/