1: static const char help[] = "Test CUPM two-dimensional memory operations on a delayed device context.\n";

  3: #include <petsc/private/cupminterface.hpp>

  5: using Petsc::device::cupm::DeviceType;

  7: template <DeviceType T>
  8: struct TestMemory2D : Petsc::device::cupm::impl::Interface<T> {
  9:   PETSC_CUPM_INHERIT_INTERFACE_TYPEDEFS_USING(T);

 11:   static PetscErrorCode Run(PetscDeviceContext dctx)
 12:   {
 13:     const PetscInt width = 3, height = 2, src_pitch = 7, dst_pitch = 5;
 14:     const auto     dtod = cupmMemcpyDeviceToDevice;
 15:     unsigned char *src, *dst, *result;
 16:     void          *handle;
 17:     cupmStream_t   stream;

 19:     PetscFunctionBeginUser;
 20:     PetscCall(PetscDeviceContextGetStreamHandle(dctx, &handle));
 21:     stream = *static_cast<cupmStream_t *>(handle);
 22:     PetscCall(PetscDeviceMalloc(dctx, PETSC_MEMTYPE_DEVICE, src_pitch * height, PETSC_DECIDE, &src));
 23:     PetscCall(PetscDeviceMalloc(dctx, PETSC_MEMTYPE_DEVICE, dst_pitch * height, PETSC_DECIDE, &dst));
 24:     PetscCall(PetscDeviceMalloc(dctx, PETSC_MEMTYPE_HOST, dst_pitch * height, PETSC_DECIDE, &result));
 25:     for (PetscInt op = 0; op < 2; ++op) {
 26:       PetscCall(PetscDeviceMemset(dctx, src, 1, src_pitch * height));
 27:       PetscCall(PetscDeviceMemset(dctx, dst, 7, dst_pitch * height));
 28:       PetscCall(PetscDeviceContextSynchronize(dctx));
 29:       PetscCall(PetscDeviceContextDelay(dctx, 0.1));
 30:       if (op) {
 31:         PetscCall(PetscDeviceMemset(dctx, dst, 2, dst_pitch * height));
 32:         PetscCallCUPM(cupmMemset2DAsync(dst, dst_pitch, 0, width, height, stream));
 33:       } else {
 34:         PetscCall(PetscDeviceMemset(dctx, src, 2, src_pitch * height));
 35:         PetscCallCUPM(cupmMemcpy2DAsync(dst, dst_pitch, src, src_pitch, width, height, dtod, stream));
 36:       }
 37:       // Drain misplaced default-stream work before the delayed producer can overwrite its result.
 38:       PetscCall(PetscDeviceContextSynchronize(PetscDeviceContextDefault));
 39:       PetscCall(PetscDeviceContextSynchronize(dctx));
 40:       PetscCall(PetscDeviceArrayCopy(dctx, result, dst, dst_pitch * height));
 41:       PetscCall(PetscDeviceContextSynchronize(dctx));
 42:       for (PetscInt j = 0; j < height; ++j) {
 43:         for (PetscInt i = 0; i < dst_pitch; ++i) {
 44:           const unsigned char expected = i < width ? (op ? 0 : 2) : (op ? 2 : 7);

 46:           PetscCheck(result[j * dst_pitch + i] == expected, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Incorrect %s result at row %" PetscInt_FMT ", column %" PetscInt_FMT ": got %u, expected %u", op ? "memset" : "copy", j, i, (unsigned int)result[j * dst_pitch + i], (unsigned int)expected);
 47:         }
 48:       }
 49:     }
 50:     PetscCall(PetscDeviceFree(dctx, src));
 51:     PetscCall(PetscDeviceFree(dctx, dst));
 52:     PetscCall(PetscDeviceFree(dctx, result));
 53:     PetscCall(PetscDeviceContextSynchronize(dctx));
 54:     PetscFunctionReturn(PETSC_SUCCESS);
 55:   }
 56: };

 58: int main(int argc, char **argv)
 59: {
 60:   PetscDeviceContext current, dctx;
 61:   PetscDeviceType    dtype;

 63:   PetscFunctionBeginUser;
 64:   PetscCall(PetscInitialize(&argc, &argv, nullptr, help));
 65:   PetscCall(PetscDeviceContextGetCurrentContext(&current));
 66:   PetscCall(PetscDeviceContextDuplicate(current, &dctx));
 67:   PetscCall(PetscDeviceContextSetStreamType(dctx, PETSC_STREAM_NONBLOCKING));
 68:   PetscCall(PetscDeviceContextSetUp(dctx));
 69:   PetscCall(PetscDeviceContextGetDeviceType(dctx, &dtype));
 70:   PetscCall(PetscDeviceContextSynchronize(PetscDeviceContextDefault));
 71:   switch (dtype) {
 72: #if PetscDefined(HAVE_CUDA)
 73:   case PETSC_DEVICE_CUDA:
 74:     PetscCall(TestMemory2D<DeviceType::CUDA>::Run(dctx));
 75:     break;
 76: #endif
 77: #if PetscDefined(HAVE_HIP)
 78:   case PETSC_DEVICE_HIP:
 79:     PetscCall(TestMemory2D<DeviceType::HIP>::Run(dctx));
 80:     break;
 81: #endif
 82:   default:
 83:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "This test requires CUDA or HIP");
 84:   }
 85:   PetscCall(PetscDeviceContextDestroy(&dctx));
 86:   PetscCall(PetscFinalize());
 87:   return 0;
 88: }

 90: /*TEST

 92:   build:
 93:     requires: cxx

 95:   testset:
 96:     output_file: output/empty.out
 97:     test:
 98:       suffix: cuda
 99:       requires: cuda
100:       args: -default_device_type cuda
101:     test:
102:       suffix: hip
103:       requires: hip
104:       args: -default_device_type hip

106: TEST*/