1: static const char help[] = "Test DM section-SF transfers with Kokkos vectors on a different stream.\n";

  3: #include <petscdmshell.h>
  4: #include <petscsection.h>
  5: #include <petscsf.h>
  6: #include <petscdevice.h>

  8: int main(int argc, char **argv)
  9: {
 10:   DM                 dm;
 11:   PetscSection       section;
 12:   PetscSF            sf;
 13:   PetscSFNode        remote[2];
 14:   Vec                g, l;
 15:   PetscDeviceContext saved, other, current;
 16:   const PetscScalar *values;
 17:   PetscInt           local[] = {0, 2};
 18:   PetscMPIInt        rank, size;

 20:   PetscFunctionBeginUser;
 21:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
 22:   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
 23:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
 24:   PetscCall(PetscDeviceContextGetCurrentContext(&saved));
 25:   PetscCall(PetscDeviceContextDuplicate(saved, &other));
 26:   PetscCall(PetscDeviceContextSetStreamType(other, PETSC_STREAM_NONBLOCKING));
 27:   PetscCall(PetscDeviceContextSetUp(other));

 29:   // Initialize Kokkos on a context that stays alive until PetscFinalize().
 30:   PetscCall(VecCreate(PETSC_COMM_WORLD, &g));
 31:   PetscCall(VecSetSizes(g, 4, PETSC_DECIDE));
 32:   PetscCall(VecSetType(g, VECMPIKOKKOS));
 33:   PetscCall(VecCreate(PETSC_COMM_SELF, &l));
 34:   PetscCall(VecSetSizes(l, 4, 4));
 35:   PetscCall(VecSetType(l, VECSEQKOKKOS));
 36:   PetscCall(DMShellCreate(PETSC_COMM_WORLD, &dm));
 37:   PetscCall(PetscSectionCreate(PETSC_COMM_WORLD, &section));
 38:   PetscCall(PetscSectionSetChart(section, 0, 4));
 39:   for (PetscInt i = 0; i < 4; ++i) PetscCall(PetscSectionSetDof(section, i, 1));
 40:   PetscCall(PetscSectionSetUp(section));
 41:   PetscCall(DMSetLocalSection(dm, section));
 42:   PetscCall(PetscSectionDestroy(&section));
 43:   remote[0].rank  = rank;
 44:   remote[0].index = 3;
 45:   remote[1].rank  = (rank + 1) % size;
 46:   remote[1].index = 1;
 47:   PetscCall(PetscSFCreate(PETSC_COMM_WORLD, &sf));
 48:   PetscCall(PetscSFSetGraph(sf, 4, 2, local, PETSC_COPY_VALUES, remote, PETSC_COPY_VALUES));
 49:   PetscCall(PetscSFSetFromOptions(sf));
 50:   PetscCall(DMSetSectionSF(dm, sf));
 51:   PetscCall(PetscSFDestroy(&sf));

 53:   // Warm the SF links on Kokkos' stream before switching to a nonblocking stream.
 54:   for (PetscInt pass = 0; pass < 2; ++pass) {
 55:     current = pass ? other : saved;
 56:     PetscCall(PetscDeviceContextSetCurrentContext(current));
 57:     for (PetscInt reverse = 0; reverse < 2; ++reverse) {
 58:       Vec input = reverse ? l : g, output = reverse ? g : l;

 60:       PetscCall(VecSet(input, 1));
 61:       PetscCall(VecSet(output, 0));
 62:       PetscCall(PetscDeviceContextSynchronize(saved));
 63:       if (pass) PetscCall(PetscDeviceContextDelay(saved, 0.05));
 64:       PetscCall(VecScale(input, 2));
 65:       if (reverse) PetscCall(DMLocalToGlobalBegin(dm, l, ADD_VALUES, g));
 66:       else PetscCall(DMGlobalToLocalBegin(dm, g, INSERT_VALUES, l));
 67:       if (pass) PetscCall(PetscDeviceContextDelay(current, 0.05));
 68:       if (reverse) PetscCall(DMLocalToGlobalEnd(dm, l, ADD_VALUES, g));
 69:       else PetscCall(DMGlobalToLocalEnd(dm, g, INSERT_VALUES, l));
 70:       PetscCall(VecScale(output, 3));
 71:       PetscCall(PetscDeviceContextSynchronize(saved));
 72:       PetscCall(PetscDeviceContextSynchronize(current));
 73:       PetscCall(VecGetArrayRead(output, &values));
 74:       for (PetscInt i = 0; i < 4; ++i) {
 75:         PetscScalar expected = (i == (reverse ? 3 : 0) || i == (reverse ? 1 : 2)) ? 6 : 0;

 77:         PetscCheck(values[i] == expected, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Pass %" PetscInt_FMT ", reverse %" PetscInt_FMT ": incorrect DM transfer at local entry %" PetscInt_FMT, pass, reverse, i);
 78:       }
 79:       PetscCall(VecRestoreArrayRead(output, &values));
 80:     }
 81:   }
 82:   PetscCall(DMDestroy(&dm));
 83:   PetscCall(VecDestroy(&g));
 84:   PetscCall(VecDestroy(&l));
 85:   PetscCall(PetscDeviceContextSetCurrentContext(saved));
 86:   PetscCall(PetscDeviceContextDestroy(&other));
 87:   PetscCall(PetscFinalize());
 88:   return 0;
 89: }

 91: /*TEST

 93:   build:
 94:     requires: kokkos_kernels

 96:   testset:
 97:     nsize: 2
 98:     output_file: output/empty.out
 99:     test:
100:       suffix: host
101:       requires: !cuda !hip !sycl
102:     test:
103:       suffix: cuda
104:       requires: cuda
105:       args: -sf_backend cuda -use_gpu_aware_mpi 0
106:     test:
107:       suffix: hip
108:       requires: hip
109:       args: -sf_backend hip -use_gpu_aware_mpi 0

111: TEST*/