Actual source code: cupminterface.hpp

  1: #pragma once

  3: #include <petscdevice_cupm.h>

  5: #include <petsc/private/cpputil.hpp>
  6: #include <petsc/private/petscadvancedmacros.h>

  8: #include <petsc/private/cpp/array.hpp>

 10: namespace Petsc
 11: {

 13: namespace device
 14: {

 16: namespace cupm
 17: {

 19: // enum describing available cupm devices, this is used as the template parameter to any
 20: // class subclassing the Interface or using it as a member variable
 21: enum class DeviceType : int {
 22:   CUDA,
 23:   HIP
 24: };

 26: // clang-format off
 27: static constexpr std::array<const char *const, 5> DeviceTypes = {
 28:   "cuda",
 29:   "hip",
 30:   "Petsc::device::cupm::DeviceType",
 31:   "Petsc::device::cupm::DeviceType::",
 32:   nullptr
 33: };
 34: // clang-format on

 36: namespace impl
 37: {

 39: #define PetscCallCUPM_(__abort_fn__, __comm__, ...) \
 40:   do { \
 41:     PetscStackUpdateLine; \
 42:     const cupmError_t cerr_p_ = __VA_ARGS__; \
 43:     __abort_fn__(cerr_p_ == cupmSuccess, __comm__, PETSC_ERR_GPU, "%s error %d (%s) : %s", cupmName(), static_cast<PetscErrorCode>(cerr_p_), cupmGetErrorName(cerr_p_), cupmGetErrorString(cerr_p_)); \
 44:   } while (0)

 46: // A backend agnostic PetscCallCUPM() function, this will only work inside the member
 47: // functions of a class inheriting from CUPM::Interface. Thanks to __VA_ARGS__ templated
 48: // functions can also be wrapped inline:
 49: //
 50: // PetscCallCUPM(foo());
 51: #define PetscCallCUPM(...)             PetscCallCUPM_(PetscCheck, PETSC_COMM_SELF, __VA_ARGS__)
 52: #define PetscCallCUPMAbort(comm_, ...) PetscCallCUPM_(PetscCheckAbort, comm_, __VA_ARGS__)

 54: // PETSC_CUPM_ALIAS_FUNCTION() - declaration to alias a cuda/hip function
 55: //
 56: // input params:
 57: // our_name   - the name of the alias
 58: // their_name - the name of the function being aliased
 59: //
 60: // notes:
 61: // see PETSC_ALIAS_FUNCTION() for the exact nature of the expansion
 62: //
 63: // example usage:
 64: // PETSC_CUPM_ALIAS_FUNCTION(cupmMalloc, cudaMalloc) ->
 65: // template 
 66: // static constexpr auto cupmMalloc(T&&... args) *noexcept and trailing return type deduction*
 67: // {
 68: //   return cudaMalloc(std::forward(args)...);
 69: // }
 70: //
 71: // PETSC_CUPM_ALIAS_FUNCTION(cupmMalloc, hipMalloc) ->
 72: // template 
 73: // static constexpr auto cupmMalloc(T&&... args) *noexcept and trailing return type deduction*
 74: // {
 75: //   return hipMalloc(std::forward(args)...);
 76: // }
 77: #define PETSC_CUPM_ALIAS_FUNCTION(our_name, their_name) PETSC_ALIAS_FUNCTION(static our_name, their_name)

 79: // PETSC_CUPM_ALIAS_FUNCTION_GOBBLE() - declaration to alias a cuda/hip function but
 80: // discard the last N arguments
 81: //
 82: // input params:
 83: // our_name   - the name of the alias
 84: // their_name - the name of the function being aliased
 85: // N          - integer constant [0, INT_MAX) dictating how many arguments to chop off the end
 86: //
 87: // notes:
 88: // see PETSC_ALIAS_FUNCTION_GOBBLE_NTH_LAST_ARGS() for the exact nature of the expansion
 89: //
 90: // example use:
 91: // PETSC_CUPM_ALIAS_FUNCTION_GOBBLE_COMMON(cupmMallocAsync, cudaMalloc, 1) ->
 92: // template 
 93: // static constexpr auto cupmMallocAsync(T&&... args, Tend argend) *noexcept and trailing
 94: // return type deduction*
 95: // {
 96: //   (void)argend;
 97: //   return cudaMalloc(std::forward(args)...);
 98: // }
 99: #define PETSC_CUPM_ALIAS_FUNCTION_GOBBLE(our_name, their_name, N) PETSC_ALIAS_FUNCTION_GOBBLE_NTH_LAST_ARGS(static our_name, their_name, N)

101: // Base class that holds functions and variables that don't require CUDA or HIP to be present
102: // on the system
103: template <DeviceType T>
104: struct InterfaceBase {
105:   static const DeviceType type = T;

107:   PETSC_NODISCARD static constexpr const char *cupmName() noexcept
108:   {
109:     static_assert(util::to_underlying(DeviceType::CUDA) == 0, "");
110:     static_assert(util::to_underlying(DeviceType::HIP) == 1, "");
111:     return std::get<util::to_underlying(T)>(DeviceTypes);
112:   }

114:   PETSC_NODISCARD static constexpr const char *cupmNAME() noexcept { return T == DeviceType::CUDA ? "CUDA" : "HIP"; }

116:   PETSC_NODISCARD static constexpr PetscDeviceType PETSC_DEVICE_CUPM() noexcept { return T == DeviceType::CUDA ? PETSC_DEVICE_CUDA : PETSC_DEVICE_HIP; }

118:   PETSC_NODISCARD static constexpr PetscMemType PETSC_MEMTYPE_CUPM() noexcept { return T == DeviceType::CUDA ? PETSC_MEMTYPE_CUDA : PETSC_MEMTYPE_HIP; }
119: };

121: // declare the base class static member variables
122: template <DeviceType T>
123: const DeviceType InterfaceBase<T>::type;

125: #define PETSC_CUPM_BASE_CLASS_HEADER(T) \
126:   using ::Petsc::device::cupm::impl::InterfaceBase<T>::type; \
127:   using ::Petsc::device::cupm::impl::InterfaceBase<T>::cupmName; \
128:   using ::Petsc::device::cupm::impl::InterfaceBase<T>::cupmNAME; \
129:   using ::Petsc::device::cupm::impl::InterfaceBase<T>::PETSC_DEVICE_CUPM; \
130:   using ::Petsc::device::cupm::impl::InterfaceBase<T>::PETSC_MEMTYPE_CUPM

132: // A templated C++ struct that defines the entire CUPM interface. Use of templating vs
133: // preprocessor macros allows us to use both interfaces simultaneously as well as easily
134: // import them into classes.
135: template <DeviceType>
136: struct InterfaceImpl;

138: #if PetscDefined(HAVE_CUDA)
139: template <>
140: struct PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL InterfaceImpl<DeviceType::CUDA> : InterfaceBase<DeviceType::CUDA> {
141:   PETSC_CUPM_BASE_CLASS_HEADER(DeviceType::CUDA);

143:   // typedefs
144:   using cupmError_t             = cudaError_t;
145:   using cupmEvent_t             = cudaEvent_t;
146:   using cupmStream_t            = cudaStream_t;
147:   using cupmDeviceProp_t        = cudaDeviceProp;
148:   using cupmMemcpyKind_t        = cudaMemcpyKind;
149:   using cupmDeviceAttr_t        = cudaDeviceAttr;
150:   using cupmComplex_t           = util::conditional_t<PetscDefined(USE_REAL_SINGLE), cuComplex, cuDoubleComplex>;
151:   using cupmPointerAttributes_t = cudaPointerAttributes;
152:   using cupmMemoryType_t        = enum cudaMemoryType;
153:   using cupmDim3                = dim3;
154:   using cupmHostFn_t            = cudaHostFn_t;
155:   #if PETSC_PKG_CUDA_VERSION_GE(11, 2, 0)
156:   using cupmMemPool_t   = cudaMemPool_t;
157:   using cupmMemPoolAttr = cudaMemPoolAttr;
158:   #else
159:   using cupmMemPool_t   = void *;
160:   using cupmMemPoolAttr = unsigned int;
161:   #endif

163:   // values
164:   static const auto cupmSuccess                 = cudaSuccess;
165:   static const auto cupmErrorNotReady           = cudaErrorNotReady;
166:   static const auto cupmErrorDeviceAlreadyInUse = cudaErrorDeviceAlreadyInUse;
167:   static const auto cupmErrorSetOnActiveProcess = cudaErrorSetOnActiveProcess;
168:   static const auto cupmErrorStubLibrary =
169:   #if PETSC_PKG_CUDA_VERSION_GE(11, 1, 0)
170:     cudaErrorStubLibrary;
171:   #else
172:     cudaErrorInsufficientDriver;
173:   #endif

175:   static const auto cupmErrorNoDevice          = cudaErrorNoDevice;
176:   static const auto cupmStreamDefault          = cudaStreamDefault;
177:   static const auto cupmStreamNonBlocking      = cudaStreamNonBlocking;
178:   static const auto cupmDeviceMapHost          = cudaDeviceMapHost;
179:   static const auto cupmMemcpyHostToDevice     = cudaMemcpyHostToDevice;
180:   static const auto cupmMemcpyDeviceToHost     = cudaMemcpyDeviceToHost;
181:   static const auto cupmMemcpyDeviceToDevice   = cudaMemcpyDeviceToDevice;
182:   static const auto cupmMemcpyHostToHost       = cudaMemcpyHostToHost;
183:   static const auto cupmMemcpyDefault          = cudaMemcpyDefault;
184:   static const auto cupmMemoryTypeHost         = cudaMemoryTypeHost;
185:   static const auto cupmMemoryTypeDevice       = cudaMemoryTypeDevice;
186:   static const auto cupmMemoryTypeManaged      = cudaMemoryTypeManaged;
187:   static const auto cupmEventDisableTiming     = cudaEventDisableTiming;
188:   static const auto cupmHostAllocDefault       = cudaHostAllocDefault;
189:   static const auto cupmHostAllocWriteCombined = cudaHostAllocWriteCombined;
190:   static const auto cupmMemPoolAttrReleaseThreshold =
191:   #if PETSC_PKG_CUDA_VERSION_GE(11, 2, 0)
192:     cudaMemPoolAttrReleaseThreshold;
193:   #else
194:     cupmMemPoolAttr{0};
195:   #endif
196:   static const auto cupmDevAttrClockRate       = cudaDevAttrClockRate;
197:   static const auto cupmDevAttrMemoryClockRate = cudaDevAttrMemoryClockRate;

199:   // error functions
200:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetErrorName, cudaGetErrorName)
201:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetErrorString, cudaGetErrorString)
202:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetLastError, cudaGetLastError)

204:   // device management
205:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetDeviceCount, cudaGetDeviceCount)
206:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetDeviceProperties, cudaGetDeviceProperties)
207:   PETSC_CUPM_ALIAS_FUNCTION(cupmDeviceGetAttribute, cudaDeviceGetAttribute)
208:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetDevice, cudaGetDevice)
209:   PETSC_CUPM_ALIAS_FUNCTION(cupmSetDevice, cudaSetDevice)
210:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetDeviceFlags, cudaGetDeviceFlags)
211:   PETSC_CUPM_ALIAS_FUNCTION(cupmSetDeviceFlags, cudaSetDeviceFlags)
212:   PETSC_CUPM_ALIAS_FUNCTION(cupmPointerGetAttributes, cudaPointerGetAttributes)

214:   PETSC_NODISCARD static cupmMemoryType_t cupmPointerGetMemType(const cupmPointerAttributes_t &attr, PetscBool *managed) noexcept
215:   {
216:   #if CUDART_VERSION < 10000
217:     if (managed) *managed = static_cast<PetscBool>(attr.isManaged);
218:     return attr.memoryType;
219:   #else
220:     if (managed) *managed = static_cast<PetscBool>(attr.type == cupmMemoryTypeManaged);
221:     return attr.type;
222:   #endif
223:   }

225:   #if PETSC_PKG_CUDA_VERSION_GE(11, 2, 0)
226:   PETSC_CUPM_ALIAS_FUNCTION(cupmDeviceGetMemPool, cudaDeviceGetMemPool)
227:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemPoolSetAttribute, cudaMemPoolSetAttribute)
228:   #else
229:   PETSC_NODISCARD static cupmError_t cupmDeviceGetMemPool(cupmMemPool_t *pool, int) noexcept
230:   {
231:     *pool = nullptr;
232:     return cupmSuccess;
233:   }

235:   PETSC_NODISCARD static cupmError_t cupmMemPoolSetAttribute(cupmMemPool_t, cupmMemPoolAttr, void *) noexcept { return cupmSuccess; }
236:   #endif
237:   // CUDA has no cudaInit() to match hipInit()
238:   PETSC_NODISCARD static cupmError_t cupmInit(unsigned int) noexcept { return cudaFree(nullptr); }

240:   // stream management
241:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventCreate, cudaEventCreate)
242:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventCreateWithFlags, cudaEventCreateWithFlags)
243:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventDestroy, cudaEventDestroy)
244:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventRecord, cudaEventRecord)
245:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventSynchronize, cudaEventSynchronize)
246:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventElapsedTime, cudaEventElapsedTime)
247:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventQuery, cudaEventQuery)
248:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamCreate, cudaStreamCreate)
249:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamCreateWithFlags, cudaStreamCreateWithFlags)
250:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamGetFlags, cudaStreamGetFlags)
251:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamDestroy, cudaStreamDestroy)
252:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamWaitEvent, cudaStreamWaitEvent)
253:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamQuery, cudaStreamQuery)
254:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamSynchronize, cudaStreamSynchronize)
255:   PETSC_CUPM_ALIAS_FUNCTION(cupmDeviceSynchronize, cudaDeviceSynchronize)
256:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetSymbolAddress, cudaGetSymbolAddress)

258:   // memory management
259:   PETSC_CUPM_ALIAS_FUNCTION(cupmFree, cudaFree)
260:   PETSC_CUPM_ALIAS_FUNCTION(cupmMalloc, cudaMalloc)
261:   #if PETSC_PKG_CUDA_VERSION_GE(11, 2, 0)
262:   PETSC_CUPM_ALIAS_FUNCTION(cupmFreeAsync, cudaFreeAsync)
263:   PETSC_CUPM_ALIAS_FUNCTION(cupmMallocAsync, cudaMallocAsync)
264:   #else
265:   PETSC_CUPM_ALIAS_FUNCTION_GOBBLE(cupmFreeAsync, cudaFree, 1)
266:   PETSC_CUPM_ALIAS_FUNCTION_GOBBLE(cupmMallocAsync, cudaMalloc, 1)
267:   #endif
268:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemcpy, cudaMemcpy)
269:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemcpyAsync, cudaMemcpyAsync)
270:   PETSC_CUPM_ALIAS_FUNCTION(cupmMallocHost, cudaMallocHost)
271:   PETSC_CUPM_ALIAS_FUNCTION(cupmFreeHost, cudaFreeHost)
272:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemset, cudaMemset)
273:   #if PETSC_PKG_CUDA_VERSION_GE(11, 2, 0)
274:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemsetAsync, cudaMemsetAsync)
275:   #else
276:   PETSC_CUPM_ALIAS_FUNCTION_GOBBLE(cupmMemsetAsync, cudaMemset, 1)
277:   #endif
278:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemcpy2D, cudaMemcpy2D)
279:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemcpy2DAsync, cudaMemcpy2DAsync)
280:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemset2D, cudaMemset2D)
281:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemset2DAsync, cudaMemset2DAsync)

283:   // launch control
284:   PETSC_CUPM_ALIAS_FUNCTION(cupmLaunchHostFunc, cudaLaunchHostFunc)
285:   template <typename FunctionT, typename... KernelArgsT>
286:   PETSC_NODISCARD static cudaError_t cupmLaunchKernel(FunctionT &&func, dim3 gridDim, dim3 blockDim, std::size_t sharedMem, cudaStream_t stream, KernelArgsT &&...kernelArgs) noexcept
287:   {
288:     static_assert(!std::is_pointer<FunctionT>::value, "kernel function must not be passed by pointer");
289:     void *args[] = {(void *)std::addressof(kernelArgs)...};

291:     return cudaLaunchKernel<util::remove_reference_t<FunctionT>>(std::addressof(func), std::move(gridDim), std::move(blockDim), args, sharedMem, std::move(stream));
292:   }
293: };
294: #endif // PetscDefined(HAVE_CUDA)

296: #if PetscDefined(HAVE_HIP)
297: template <>
298: struct PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL InterfaceImpl<DeviceType::HIP> : InterfaceBase<DeviceType::HIP> {
299:   PETSC_CUPM_BASE_CLASS_HEADER(DeviceType::HIP);

301:   // typedefs
302:   using cupmError_t             = hipError_t;
303:   using cupmEvent_t             = hipEvent_t;
304:   using cupmStream_t            = hipStream_t;
305:   using cupmDeviceProp_t        = hipDeviceProp_t;
306:   using cupmMemcpyKind_t        = hipMemcpyKind;
307:   using cupmDeviceAttr_t        = hipDeviceAttribute_t;
308:   using cupmComplex_t           = util::conditional_t<PetscDefined(USE_REAL_SINGLE), hipComplex, hipDoubleComplex>;
309:   using cupmPointerAttributes_t = hipPointerAttribute_t;
310:   using cupmMemoryType_t        = enum hipMemoryType;
311:   using cupmDim3                = dim3;
312:   #if PETSC_PKG_HIP_VERSION_GE(5, 2, 0)
313:   using cupmHostFn_t    = hipHostFn_t;
314:   using cupmMemPool_t   = hipMemPool_t;
315:   using cupmMemPoolAttr = hipMemPoolAttr;
316:   #else
317:   using cupmHostFn_t    = void (*)(void *);
318:   using cupmMemPool_t   = void *;
319:   using cupmMemPoolAttr = unsigned int;
320:   #endif

322:   // values
323:   static const auto cupmSuccess       = hipSuccess;
324:   static const auto cupmErrorNotReady = hipErrorNotReady;
325:   // see https://github.com/ROCm-Developer-Tools/HIP/blob/develop/bin/hipify-perl
326:   static const auto cupmErrorDeviceAlreadyInUse = hipErrorContextAlreadyInUse;
327:   static const auto cupmErrorSetOnActiveProcess = hipErrorSetOnActiveProcess;
328:   // as of HIP v4.2 cudaErrorStubLibrary has no HIP equivalent
329:   static const auto cupmErrorStubLibrary     = hipErrorInsufficientDriver;
330:   static const auto cupmErrorNoDevice        = hipErrorNoDevice;
331:   static const auto cupmStreamDefault        = hipStreamDefault;
332:   static const auto cupmStreamNonBlocking    = hipStreamNonBlocking;
333:   static const auto cupmDeviceMapHost        = hipDeviceMapHost;
334:   static const auto cupmMemcpyHostToDevice   = hipMemcpyHostToDevice;
335:   static const auto cupmMemcpyDeviceToHost   = hipMemcpyDeviceToHost;
336:   static const auto cupmMemcpyDeviceToDevice = hipMemcpyDeviceToDevice;
337:   static const auto cupmMemcpyHostToHost     = hipMemcpyHostToHost;
338:   static const auto cupmMemcpyDefault        = hipMemcpyDefault;
339:   static const auto cupmMemoryTypeHost       = hipMemoryTypeHost;
340:   static const auto cupmMemoryTypeDevice     = hipMemoryTypeDevice;
341:   // see
342:   // https://github.com/ROCm-Developer-Tools/HIP/blob/develop/include/hip/hip_runtime_api.h#L156
343:   static const auto cupmMemoryTypeManaged      = hipMemoryTypeUnified;
344:   static const auto cupmEventDisableTiming     = hipEventDisableTiming;
345:   static const auto cupmHostAllocDefault       = hipHostMallocDefault;
346:   static const auto cupmHostAllocWriteCombined = hipHostMallocWriteCombined;
347:   static const auto cupmMemPoolAttrReleaseThreshold =
348:   #if PETSC_PKG_HIP_VERSION_GE(5, 2, 0)
349:     hipMemPoolAttrReleaseThreshold;
350:   #else
351:     cupmMemPoolAttr{0};
352:   #endif
353:   static const auto cupmDevAttrClockRate       = hipDeviceAttributeClockRate;
354:   static const auto cupmDevAttrMemoryClockRate = hipDeviceAttributeMemoryClockRate;

356:   // error functions
357:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetErrorName, hipGetErrorName)
358:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetErrorString, hipGetErrorString)
359:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetLastError, hipGetLastError)

361:   // device management
362:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetDeviceCount, hipGetDeviceCount)
363:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetDeviceProperties, hipGetDeviceProperties)
364:   PETSC_CUPM_ALIAS_FUNCTION(cupmDeviceGetAttribute, hipDeviceGetAttribute)
365:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetDevice, hipGetDevice)
366:   PETSC_CUPM_ALIAS_FUNCTION(cupmSetDevice, hipSetDevice)
367:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetDeviceFlags, hipGetDeviceFlags)
368:   PETSC_CUPM_ALIAS_FUNCTION(cupmSetDeviceFlags, hipSetDeviceFlags)
369:   PETSC_CUPM_ALIAS_FUNCTION(cupmPointerGetAttributes, hipPointerGetAttributes)

371:   PETSC_NODISCARD static cupmMemoryType_t cupmPointerGetMemType(const cupmPointerAttributes_t &attr, PetscBool *managed) noexcept
372:   {
373:     // HIP used memoryType on both backends before 5.5.
374:   #if PETSC_PKG_HIP_VERSION_LT(5, 5, 0)
375:     if (managed) *managed = static_cast<PetscBool>(attr.isManaged);
376:     return attr.memoryType;
377:   #else
378:     if (managed) *managed = static_cast<PetscBool>(attr.type == cupmMemoryTypeManaged);
379:     return attr.type;
380:   #endif
381:   }

383:   #if PETSC_PKG_HIP_VERSION_GE(5, 2, 0)
384:   PETSC_CUPM_ALIAS_FUNCTION(cupmDeviceGetMemPool, hipDeviceGetMemPool)
385:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemPoolSetAttribute, hipMemPoolSetAttribute)
386:   #else
387:   PETSC_NODISCARD static cupmError_t cupmDeviceGetMemPool(cupmMemPool_t *pool, int) noexcept
388:   {
389:     *pool = nullptr;
390:     return cupmSuccess;
391:   }

393:   PETSC_NODISCARD static cupmError_t cupmMemPoolSetAttribute(cupmMemPool_t, cupmMemPoolAttr, void *) noexcept { return cupmSuccess; }
394:   #endif
395:   PETSC_CUPM_ALIAS_FUNCTION(cupmInit, hipInit)

397:   // stream management
398:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventCreate, hipEventCreate)
399:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventCreateWithFlags, hipEventCreateWithFlags)
400:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventDestroy, hipEventDestroy)
401:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventRecord, hipEventRecord)
402:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventSynchronize, hipEventSynchronize)
403:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventElapsedTime, hipEventElapsedTime)
404:   PETSC_CUPM_ALIAS_FUNCTION(cupmEventQuery, hipEventQuery)
405:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamCreate, hipStreamCreate)
406:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamCreateWithFlags, hipStreamCreateWithFlags)
407:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamGetFlags, hipStreamGetFlags)
408:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamDestroy, hipStreamDestroy)
409:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamWaitEvent, hipStreamWaitEvent)
410:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamQuery, hipStreamQuery)
411:   PETSC_CUPM_ALIAS_FUNCTION(cupmStreamSynchronize, hipStreamSynchronize)
412:   PETSC_CUPM_ALIAS_FUNCTION(cupmDeviceSynchronize, hipDeviceSynchronize)
413:   PETSC_CUPM_ALIAS_FUNCTION(cupmGetSymbolAddress, hipGetSymbolAddress)

415:   // memory management
416:   PETSC_CUPM_ALIAS_FUNCTION(cupmFree, hipFree)
417:   PETSC_CUPM_ALIAS_FUNCTION(cupmMalloc, hipMalloc)
418:   #if PETSC_PKG_HIP_VERSION_GE(5, 2, 0)
419:   PETSC_CUPM_ALIAS_FUNCTION(cupmMallocAsync, hipMallocAsync)
420:   PETSC_CUPM_ALIAS_FUNCTION(cupmFreeAsync, hipFreeAsync)
421:   #else
422:   PETSC_CUPM_ALIAS_FUNCTION_GOBBLE(cupmMallocAsync, hipMalloc, 1)
423:   PETSC_CUPM_ALIAS_FUNCTION_GOBBLE(cupmFreeAsync, hipFree, 1)
424:   #endif
425:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemcpy, hipMemcpy)
426:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemcpyAsync, hipMemcpyAsync)
427:   // hipMallocHost is deprecated
428:   PETSC_CUPM_ALIAS_FUNCTION(cupmMallocHost, hipHostMalloc)
429:   // hipFreeHost is deprecated
430:   PETSC_CUPM_ALIAS_FUNCTION(cupmFreeHost, hipHostFree)
431:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemset, hipMemset)
432:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemsetAsync, hipMemsetAsync)
433:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemcpy2D, hipMemcpy2D)
434:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemcpy2DAsync, hipMemcpy2DAsync)
435:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemset2D, hipMemset2D)
436:   PETSC_CUPM_ALIAS_FUNCTION(cupmMemset2DAsync, hipMemset2DAsync)

438:   // launch control
439:   // HIP appears to only have hipLaunchHostFunc from 5.2.0 onwards
440:   // https://github.com/ROCm-Developer-Tools/HIPIFY/blob/master/doc/markdown/CUDA_Runtime_API_functions_supported_by_HIP.md#7-execution-control=
441:   #if PetscDefined(HAVE_HIPCUDA)
442:   PETSC_NODISCARD static hipError_t cupmLaunchHostFunc(cupmStream_t stream, void (*fn)(void *), void *ctx) noexcept { return hipCUDAErrorTohipError(cudaLaunchHostFunc(stream, fn, ctx)); }
443:   #elif PETSC_PKG_HIP_VERSION_GE(5, 2, 0)
444:   PETSC_CUPM_ALIAS_FUNCTION(cupmLaunchHostFunc, hipLaunchHostFunc)
445:   #else
446:   PETSC_NODISCARD static hipError_t cupmLaunchHostFunc(hipStream_t stream, cupmHostFn_t fn, void *ctx) noexcept
447:   {
448:     // the only correct way to spoof this function is to do it synchronously...
449:     auto herr = hipStreamSynchronize(stream);
450:     if (PetscUnlikely(herr != hipSuccess)) return herr;
451:     fn(ctx);
452:     return herr;
453:   }
454:   #endif

456:   template <typename FunctionT, typename... KernelArgsT>
457:   PETSC_NODISCARD static hipError_t cupmLaunchKernel(FunctionT &&func, dim3 gridDim, dim3 blockDim, std::size_t sharedMem, hipStream_t stream, KernelArgsT &&...kernelArgs) noexcept
458:   {
459:     void *args[] = {(void *)std::addressof(kernelArgs)...};

461:     return hipLaunchKernel((void *)func, std::move(gridDim), std::move(blockDim), args, sharedMem, std::move(stream));
462:   }
463: };
464: #endif // PetscDefined(HAVE_HIP)

466: // shorthand for bringing all of the typedefs from the base Interface class into your own,
467: // it's annoying that c++ doesn't have a way to do this automatically
468: #define PETSC_CUPM_IMPL_CLASS_HEADER(T) \
469:   PETSC_CUPM_BASE_CLASS_HEADER(T); \
470:   /* types */ \
471:   using cupmError_t             = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmError_t; \
472:   using cupmEvent_t             = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmEvent_t; \
473:   using cupmStream_t            = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmStream_t; \
474:   using cupmDeviceProp_t        = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmDeviceProp_t; \
475:   using cupmMemcpyKind_t        = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemcpyKind_t; \
476:   using cupmDeviceAttr_t        = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmDeviceAttr_t; \
477:   using cupmComplex_t           = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmComplex_t; \
478:   using cupmPointerAttributes_t = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmPointerAttributes_t; \
479:   using cupmMemoryType_t        = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemoryType_t; \
480:   using cupmDim3                = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmDim3; \
481:   using cupmHostFn_t            = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmHostFn_t; \
482:   using cupmMemPool_t           = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemPool_t; \
483:   using cupmMemPoolAttr         = typename ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemPoolAttr; \
484:   /* variables */ \
485:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmSuccess; \
486:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmErrorNotReady; \
487:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmErrorDeviceAlreadyInUse; \
488:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmErrorSetOnActiveProcess; \
489:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmErrorStubLibrary; \
490:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmErrorNoDevice; \
491:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmStreamDefault; \
492:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmStreamNonBlocking; \
493:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmDeviceMapHost; \
494:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemcpyHostToDevice; \
495:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemcpyDeviceToHost; \
496:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemcpyDeviceToDevice; \
497:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemcpyHostToHost; \
498:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemcpyDefault; \
499:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemoryTypeHost; \
500:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemoryTypeDevice; \
501:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemoryTypeManaged; \
502:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmEventDisableTiming; \
503:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmHostAllocDefault; \
504:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmHostAllocWriteCombined; \
505:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemPoolAttrReleaseThreshold; \
506:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmDevAttrClockRate; \
507:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmDevAttrMemoryClockRate; \
508:   /* functions */ \
509:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmGetErrorName; \
510:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmGetErrorString; \
511:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmGetLastError; \
512:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmGetDeviceCount; \
513:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmGetDeviceProperties; \
514:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmDeviceGetAttribute; \
515:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmGetDevice; \
516:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmSetDevice; \
517:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmGetDeviceFlags; \
518:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmSetDeviceFlags; \
519:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmPointerGetAttributes; \
520:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmDeviceGetMemPool; \
521:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemPoolSetAttribute; \
522:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmInit; \
523:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmEventCreate; \
524:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmEventCreateWithFlags; \
525:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmEventDestroy; \
526:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmEventRecord; \
527:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmEventSynchronize; \
528:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmEventElapsedTime; \
529:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmEventQuery; \
530:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmStreamCreate; \
531:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmStreamCreateWithFlags; \
532:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmStreamGetFlags; \
533:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmStreamDestroy; \
534:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmStreamWaitEvent; \
535:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmStreamQuery; \
536:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmStreamSynchronize; \
537:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmDeviceSynchronize; \
538:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmGetSymbolAddress; \
539:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMalloc; \
540:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMallocAsync; \
541:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemcpy; \
542:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemcpyAsync; \
543:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMallocHost; \
544:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemset; \
545:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemsetAsync; \
546:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemcpy2D; \
547:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemcpy2DAsync; \
548:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemset2D; \
549:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmMemset2DAsync; \
550:   using ::Petsc::device::cupm::impl::InterfaceImpl<T>::cupmLaunchHostFunc

552: #if PetscHasAttribute(always_inline)
553:   // https://gcc.gnu.org/bugzilla//show_bug.cgi?id=109464
554:   #define PETSC_GCC_LINKER_UNDEFINED_REFERENCE_BUG_WORKAROUND __attribute__((always_inline))
555: #else
556:   #define PETSC_GCC_LINKER_UNDEFINED_REFERENCE_BUG_WORKAROUND
557: #endif

559: // The actual interface class
560: template <DeviceType T>
561: struct Interface : InterfaceImpl<T> {
562: private:
563:   using interface_type = InterfaceImpl<T>;

565: public:
566:   PETSC_CUPM_IMPL_CLASS_HEADER(T);

568:   using cupmReal_t   = util::conditional_t<PetscDefined(USE_REAL_SINGLE), float, double>;
569:   using cupmScalar_t = util::conditional_t<PetscDefined(USE_COMPLEX), cupmComplex_t, cupmReal_t>;

571:   PETSC_NODISCARD PETSC_GCC_LINKER_UNDEFINED_REFERENCE_BUG_WORKAROUND static constexpr cupmScalar_t cupmScalarCast(PetscScalar s) noexcept
572:   {
573: #if PetscDefined(USE_COMPLEX)
574:     return cupmComplex_t{PetscRealPart(s), PetscImaginaryPart(s)};
575: #else
576:     return static_cast<cupmScalar_t>(s);
577: #endif
578:   }

580:   PETSC_NODISCARD PETSC_GCC_LINKER_UNDEFINED_REFERENCE_BUG_WORKAROUND static constexpr const cupmScalar_t *cupmScalarPtrCast(const PetscScalar *s) noexcept { return reinterpret_cast<const cupmScalar_t *>(s); }

582:   PETSC_NODISCARD PETSC_GCC_LINKER_UNDEFINED_REFERENCE_BUG_WORKAROUND static constexpr cupmScalar_t *cupmScalarPtrCast(PetscScalar *s) noexcept { return reinterpret_cast<cupmScalar_t *>(s); }

584:   PETSC_NODISCARD PETSC_GCC_LINKER_UNDEFINED_REFERENCE_BUG_WORKAROUND static constexpr const cupmReal_t *cupmRealPtrCast(const PetscReal *s) noexcept { return reinterpret_cast<const cupmReal_t *>(s); }

586:   PETSC_NODISCARD PETSC_GCC_LINKER_UNDEFINED_REFERENCE_BUG_WORKAROUND static constexpr cupmReal_t *cupmRealPtrCast(PetscReal *s) noexcept { return reinterpret_cast<cupmReal_t *>(s); }

588: #if !defined(PETSC_PKG_CUDA_VERSION_GE)
589:   #define PETSC_PKG_CUDA_VERSION_GE(...) 0
590:   #define CUPM_DEFINED_PETSC_PKG_CUDA_VERSION_GE
591: #endif

593:   static PetscErrorCode PetscCUPMGetMemType(const void *data, PetscMemType *type, PetscBool *registered = nullptr, PetscBool *managed = nullptr) noexcept
594:   {
595:     cupmPointerAttributes_t attr;
596:     cupmError_t             cerr;

598:     PetscFunctionBegin;
599:     if (type) PetscAssertPointer(type, 2);
600:     if (registered) {
601:       PetscAssertPointer(registered, 3);
602:       *registered = PETSC_FALSE;
603:     }
604:     if (managed) {
605:       PetscAssertPointer(managed, 4);
606:       *managed = PETSC_FALSE;
607:     }
608:     // Do not check error, instead reset it via GetLastError() since before CUDA 11.0, passing
609:     // a host pointer returns cudaErrorInvalidValue
610:     cerr = cupmPointerGetAttributes(&attr, data);
611:     cerr = cupmGetLastError();
612:     if (PETSC_PKG_CUDA_VERSION_GE(11, 0, 0) && (T == DeviceType::CUDA)) PetscCallCUPM(cerr);
613:     if (type) *type = PETSC_MEMTYPE_HOST;
614:     if (cerr == cupmSuccess) {
615:       const auto mtype = interface_type::cupmPointerGetMemType(attr, managed);

617:       // Managed memory is device accessible, so treat it as device memory.
618:       if (type && (mtype == cupmMemoryTypeDevice || mtype == cupmMemoryTypeManaged)) *type = PETSC_MEMTYPE_CUPM();
619:       if (registered && (mtype == cupmMemoryTypeHost)) *registered = PETSC_TRUE;
620:     }
621:     PetscFunctionReturn(PETSC_SUCCESS);
622:   }
623: #if defined(CUPM_DEFINED_PETSC_PKG_CUDA_VERSION_GE)
624:   #undef PETSC_PKG_CUDA_VERSION_GE
625: #endif

627:   PETSC_NODISCARD static PETSC_CONSTEXPR_14 cupmMemcpyKind_t PetscDeviceCopyModeToCUPMMemcpyKind(PetscDeviceCopyMode mode) noexcept
628:   {
629:     switch (mode) {
630:     case PETSC_DEVICE_COPY_HTOH:
631:       return cupmMemcpyHostToHost;
632:     case PETSC_DEVICE_COPY_HTOD:
633:       return cupmMemcpyHostToDevice;
634:     case PETSC_DEVICE_COPY_DTOD:
635:       return cupmMemcpyDeviceToDevice;
636:     case PETSC_DEVICE_COPY_DTOH:
637:       return cupmMemcpyDeviceToHost;
638:     case PETSC_DEVICE_COPY_AUTO:
639:       return cupmMemcpyDefault;
640:     }
641:     PetscUnreachable();
642:     return cupmMemcpyDefault;
643:   }

645:   // these change what the arguments mean, so need to namespace these
646:   template <typename M>
647:   static PetscErrorCode PetscCUPMMallocAsync(M **ptr, std::size_t n, cupmStream_t stream = nullptr) noexcept
648:   {
649:     static_assert(!std::is_void<M>::value, "");

651:     PetscFunctionBegin;
652:     PetscAssertPointer(ptr, 1);
653:     *ptr = nullptr;
654:     if (n) {
655:       const auto bytes = n * sizeof(M);
656:       // https://developer.nvidia.com/blog/using-cuda-stream-ordered-memory-allocator-part-2/
657:       //
658:       // TLD;DR: cudaMallocAsync() does not work with NVIDIA GPUDirect which OPENMPI uses to
659:       // underpin its cuda-aware MPI implementation, so we cannot just async allocate
660:       // blindly...
661:       if (stream) {
662:         PetscCallCUPM(cupmMallocAsync(reinterpret_cast<void **>(ptr), bytes, stream));
663:       } else {
664:         PetscCallCUPM(cupmMalloc(reinterpret_cast<void **>(ptr), bytes));
665:       }
666:     }
667:     PetscFunctionReturn(PETSC_SUCCESS);
668:   }

670:   template <typename M>
671:   static PetscErrorCode PetscCUPMMalloc(M **ptr, std::size_t n) noexcept
672:   {
673:     PetscFunctionBegin;
674:     PetscCall(PetscCUPMMallocAsync(ptr, n));
675:     PetscFunctionReturn(PETSC_SUCCESS);
676:   }

678:   template <typename M>
679:   static PetscErrorCode PetscCUPMMallocHost(M **ptr, std::size_t n, unsigned int flags = cupmHostAllocDefault) noexcept
680:   {
681:     static_assert(!std::is_void<M>::value, "");

683:     PetscFunctionBegin;
684:     PetscAssertPointer(ptr, 1);
685:     *ptr = nullptr;
686:     if (n) PetscCallCUPM(cupmMallocHost(reinterpret_cast<void **>(ptr), n * sizeof(M), flags));
687:     PetscFunctionReturn(PETSC_SUCCESS);
688:   }

690:   template <typename D>
691:   static PetscErrorCode PetscCUPMMemcpyAsync(D *dest, const util::type_identity_t<D> *src, std::size_t n, cupmMemcpyKind_t kind, cupmStream_t stream = nullptr, bool use_async = false) noexcept
692:   {
693:     static_assert(!std::is_void<D>::value, "");
694:     const auto size = n * sizeof(D);

696:     PetscFunctionBegin;
697:     if (PetscUnlikely(!n)) PetscFunctionReturn(PETSC_SUCCESS);
698:     // cannot dereference (i.e. cannot call PetscAssertPointer() here)
699:     PetscCheck(dest, PETSC_COMM_SELF, PETSC_ERR_POINTER, "Trying to copy to a NULL pointer");
700:     PetscCheck(src, PETSC_COMM_SELF, PETSC_ERR_POINTER, "Trying to copy from a NULL pointer");
701:     if (kind == cupmMemcpyHostToHost) {
702:       // If we are HTOH it is cheaper to check if the stream is idle and do a basic mempcy()
703:       // than it is to just call the vendor functions. This assumes of course that the stream
704:       // accounts for both memory regions being "idle"
705:       if (cupmStreamQuery(stream) == cupmSuccess) {
706:         PetscCall(PetscMemcpy(dest, src, size));
707:         PetscFunctionReturn(PETSC_SUCCESS);
708:       }
709:       // need to clear the potential cupmErrorNotReady generated by query above...
710:       auto cerr = cupmGetLastError();

712:       if (PetscUnlikely(cerr != cupmErrorNotReady)) PetscCallCUPM(cerr);
713:     }
714:     if (use_async || stream || (kind != cupmMemcpyDeviceToHost)) {
715:       PetscCallCUPM(cupmMemcpyAsync(dest, src, size, kind, stream));
716:     } else {
717:       PetscCallCUPM(cupmMemcpy(dest, src, size, kind));
718:     }
719:     PetscCall(PetscLogCUPMMemcpyTransfer(kind, size));
720:     PetscFunctionReturn(PETSC_SUCCESS);
721:   }

723:   template <typename D>
724:   static PetscErrorCode PetscCUPMMemcpy(D *dest, const util::type_identity_t<D> *src, std::size_t n, cupmMemcpyKind_t kind) noexcept
725:   {
726:     PetscFunctionBegin;
727:     PetscCall(PetscCUPMMemcpyAsync(dest, src, n, kind));
728:     PetscFunctionReturn(PETSC_SUCCESS);
729:   }

731:   template <typename D>
732:   static PetscErrorCode PetscCUPMMemcpy2DAsync(D *dest, std::size_t dest_pitch, const util::type_identity_t<D> *src, std::size_t src_pitch, std::size_t width, std::size_t height, cupmMemcpyKind_t kind, cupmStream_t stream = nullptr)
733:   {
734:     static_assert(!std::is_void<D>::value, "");
735:     const auto dest_pitch_bytes = dest_pitch * sizeof(D);
736:     const auto src_pitch_bytes  = src_pitch * sizeof(D);
737:     const auto width_bytes      = width * sizeof(D);
738:     const auto size             = height * width_bytes;

740:     PetscFunctionBegin;
741:     if (PetscUnlikely(!size)) PetscFunctionReturn(PETSC_SUCCESS);
742:     PetscCheck(dest, PETSC_COMM_SELF, PETSC_ERR_POINTER, "Trying to copy to a NULL pointer");
743:     PetscCheck(src, PETSC_COMM_SELF, PETSC_ERR_POINTER, "Trying to copy from a NULL pointer");
744:     if (stream || (kind != cupmMemcpyDeviceToHost)) {
745:       PetscCallCUPM(cupmMemcpy2DAsync(dest, dest_pitch_bytes, src, src_pitch_bytes, width_bytes, height, kind, stream));
746:     } else {
747:       PetscCallCUPM(cupmMemcpy2D(dest, dest_pitch_bytes, src, src_pitch_bytes, width_bytes, height, kind));
748:     }
749:     PetscCall(PetscLogCUPMMemcpyTransfer(kind, size));
750:     PetscFunctionReturn(PETSC_SUCCESS);
751:   }

753:   template <typename D>
754:   static PetscErrorCode PetscCUPMMemcpy2D(D *dest, std::size_t dest_pitch, const util::type_identity_t<D> *src, std::size_t src_pitch, std::size_t width, std::size_t height, cupmMemcpyKind_t kind)
755:   {
756:     PetscFunctionBegin;
757:     PetscCall(PetscCUPMMemcpy2DAsync(dest, dest_pitch, src, src_pitch, width, height, kind));
758:     PetscFunctionReturn(PETSC_SUCCESS);
759:   }

761:   template <typename M>
762:   static PetscErrorCode PetscCUPMMemsetAsync(M *ptr, int value, std::size_t n, cupmStream_t stream = nullptr, bool use_async = false) noexcept
763:   {
764:     static_assert(!std::is_void<M>::value, "");

766:     PetscFunctionBegin;
767:     if (PetscLikely(n)) {
768:       const auto bytes = n * sizeof(M);

770:       PetscCheck(ptr, PETSC_COMM_SELF, PETSC_ERR_POINTER, "Trying to memset a NULL pointer with size %zu != 0", n);
771:       if (stream || use_async) {
772:         PetscCallCUPM(cupmMemsetAsync(ptr, value, bytes, stream));
773:       } else {
774:         PetscCallCUPM(cupmMemset(ptr, value, bytes));
775:       }
776:     }
777:     PetscFunctionReturn(PETSC_SUCCESS);
778:   }

780:   template <typename M>
781:   static PetscErrorCode PetscCUPMMemset(M *ptr, int value, std::size_t n) noexcept
782:   {
783:     PetscFunctionBegin;
784:     PetscCall(PetscCUPMMemsetAsync(ptr, value, n));
785:     PetscFunctionReturn(PETSC_SUCCESS);
786:   }

788:   template <typename D>
789:   static PetscErrorCode PetscCUPMMemset2DAsync(D *ptr, std::size_t pitch, int value, std::size_t width, std::size_t height, cupmStream_t stream = nullptr)
790:   {
791:     static_assert(!std::is_void<D>::value, "");
792:     const auto pitch_bytes = pitch * sizeof(D);
793:     const auto width_bytes = width * sizeof(D);
794:     const auto size        = width_bytes * height;

796:     PetscFunctionBegin;
797:     if (PetscUnlikely(!size)) PetscFunctionReturn(PETSC_SUCCESS);
798:     PetscAssert(ptr, PETSC_COMM_SELF, PETSC_ERR_POINTER, "Trying to memset a NULL pointer with size %zu != 0", size);
799:     if (stream) {
800:       PetscCallCUPM(cupmMemset2DAsync(ptr, pitch_bytes, value, width_bytes, height, stream));
801:     } else {
802:       PetscCallCUPM(cupmMemset2D(ptr, pitch_bytes, value, width_bytes, height));
803:     }
804:     PetscFunctionReturn(PETSC_SUCCESS);
805:   }

807:   // these we can transparently wrap, no need to namespace it to Petsc
808:   template <typename M>
809:   PETSC_NODISCARD static cupmError_t cupmFreeAsync(M &ptr, cupmStream_t stream = nullptr) noexcept
810:   {
811:     static_assert(std::is_pointer<util::decay_t<M>>::value, "");
812:     static_assert(!std::is_const<M>::value, "");

814:     if (ptr) {
815:       auto cerr = interface_type::cupmFreeAsync(std::forward<M>(ptr), stream);

817:       ptr = nullptr;
818:       if (PetscUnlikely(cerr != cupmSuccess)) return cerr;
819:     }
820:     return cupmSuccess;
821:   }

823:   PETSC_NODISCARD static cupmError_t cupmFreeAsync(std::nullptr_t ptr, cupmStream_t stream = nullptr) { return interface_type::cupmFreeAsync(ptr, stream); }

825:   template <typename M>
826:   PETSC_NODISCARD static cupmError_t cupmFree(M &ptr) noexcept
827:   {
828:     return cupmFreeAsync(ptr);
829:   }

831:   PETSC_NODISCARD static cupmError_t cupmFree(std::nullptr_t ptr) { return cupmFreeAsync(ptr); }

833:   template <typename M>
834:   PETSC_NODISCARD static cupmError_t cupmFreeHost(M &ptr) noexcept
835:   {
836:     static_assert(std::is_pointer<util::decay_t<M>>::value, "");
837:     const auto cerr = interface_type::cupmFreeHost(std::forward<M>(ptr));
838:     ptr             = nullptr;
839:     return cerr;
840:   }

842:   PETSC_NODISCARD static cupmError_t cupmFreeHost(std::nullptr_t ptr) { return interface_type::cupmFreeHost(ptr); }

844:   // specific wrapper for device launch function, as the real function is a C routine and
845:   // doesn't have variable arguments. The actual mechanics of this are a bit complicated but
846:   // boils down to the fact that ultimately we pass a
847:   //
848:   // void *args[] = {(void*)&kernel_args...};
849:   //
850:   // to the kernel launcher. Since we pass void* this means implicit conversion does **not**
851:   // happen to the kernel arguments so we must do it ourselves here. This function does this in
852:   // 3 stages:
853:   // 1. Enumerate the kernel arguments (cupmLaunchKernel)
854:   // 2. Deduce the signature of func() and static_cast the kernel arguments to the type
855:   //    expected by func() using the enumeration above (deduceKernelCall)
856:   // 3. Form the void* array with the converted arguments and call cuda/hipLaunchKernel with
857:   //    it. (interface_type::cupmLaunchKernel)
858:   template <typename F, typename... Args>
859:   PETSC_NODISCARD static cupmError_t cupmLaunchKernel(F &&func, cupmDim3 gridDim, cupmDim3 blockDim, std::size_t sharedMem, cupmStream_t stream, Args &&...kernelArgs) noexcept
860:   {
861:     return deduceKernelCall(util::index_sequence_for<Args...>{}, std::forward<F>(func), std::move(gridDim), std::move(blockDim), std::move(sharedMem), std::move(stream), std::forward<Args>(kernelArgs)...);
862:   }

864:   template <std::size_t block_size = 256, std::size_t warp_size = 32, typename F, typename... Args>
865:   static PetscErrorCode PetscCUPMLaunchKernel1D(std::size_t n, std::size_t sharedMem, cupmStream_t stream, F &&func, Args &&...kernelArgs) noexcept
866:   {
867:     static_assert(block_size > 0, "");
868:     static_assert(warp_size > 0, "");
869:     // want block_size to be a multiple of the warp_size
870:     static_assert(block_size % warp_size == 0, "");
871:     const auto nthread = std::min(n, block_size);
872:     const auto nblock  = (n + block_size - 1) / block_size;

874:     PetscFunctionBegin;
875:     // if n = 0 then nthread = 0, which is not allowed. rather than letting the user try to
876:     // decipher cryptic 'cuda/hipErrorLaunchFailure' we explicitly check for zero here
877:     PetscAssert(nthread, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Trying to launch kernel with grid/block size 0");
878:     PetscCallCUPM(cupmLaunchKernel(std::forward<F>(func), (unsigned int)nblock, (unsigned int)nthread, sharedMem, stream, std::forward<Args>(kernelArgs)...));
879:     PetscFunctionReturn(PETSC_SUCCESS);
880:   }

882: private:
883:   template <typename S, typename D, typename = void>
884:   struct is_static_castable : std::false_type { };

886:   template <typename S, typename D>
887:   struct is_static_castable<S, D, util::void_t<decltype(static_cast<D>(std::declval<S>()))>> : std::true_type { };

889:   template <typename D, typename S>
890:   static constexpr util::enable_if_t<is_static_castable<S, D>::value, D> cast_to(S &&src) noexcept
891:   {
892:     return static_cast<D>(std::forward<S>(src));
893:   }

895:   template <typename D, typename S>
896:   static constexpr util::enable_if_t<!is_static_castable<S, D>::value, D> cast_to(S &&src) noexcept
897:   {
898:     return const_cast<D>(std::forward<S>(src));
899:   }

901:   template <typename F, typename... Args, std::size_t... Idx>
902:   PETSC_NODISCARD static cupmError_t deduceKernelCall(util::index_sequence<Idx...>, F &&func, cupmDim3 gridDim, cupmDim3 blockDim, std::size_t sharedMem, cupmStream_t stream, Args &&...kernelArgs) noexcept
903:   {
904:     // clang-format off
905:     return interface_type::cupmLaunchKernel(
906:       std::forward<F>(func),
907:       std::move(gridDim), std::move(blockDim), std::move(sharedMem), std::move(stream),
908:       // can't static_cast() here since the function argument type may be cv-qualified, in
909:       // which case we would need to const_cast(). But you can only const_cast() indirect types
910:       // (pointers, references). So we need a SFINAE monster that is a static_cast() if
911:       // possible, and a const_cast() if not. We could just use a C-style cast which *would*
912:       // work here since it tries the following and uses the first one that succeeds:
913:       //
914:       // 1. const_cast()
915:       // 2. static_cast()
916:       // 3. static_cast() then const_cast()
917:       // 4. reinterpret_cast()...
918:       //
919:       // the issue however is the final reinterpret_cast(). We absolutely cannot get there
920:       // because doing so would silently hide a ton of bugs, for example casting a PetscScalar
921:       // * to double * in complex builds, a PetscInt * to int * in 64idx builds, etc.
922:       cast_to<typename util::func_traits<F>::template arg<Idx>::type>(std::forward<Args>(kernelArgs))...
923:     );
924:     // clang-format on
925:   }

927:   static PetscErrorCode PetscLogCUPMMemcpyTransfer(cupmMemcpyKind_t kind, std::size_t size) noexcept
928:   {
929:     PetscFunctionBegin;
930:     // only the explicit HTOD or DTOH are handled, since we either don't log the other cases
931:     // (yet) or don't know the direction
932:     if (kind == cupmMemcpyDeviceToHost) PetscCall(PetscLogGpuToCpu(static_cast<PetscLogDouble>(size)));
933:     else if (kind == cupmMemcpyHostToDevice) PetscCall(PetscLogCpuToGpu(static_cast<PetscLogDouble>(size)));
934:     else (void)size;
935:     PetscFunctionReturn(PETSC_SUCCESS);
936:   }
937: };

939: #undef PETSC_GCC_LINKER_UNDEFINED_REFERENCE_BUG_WORKAROUND

941: #define PETSC_CUPM_INHERIT_INTERFACE_TYPEDEFS_USING(T) \
942:   PETSC_CUPM_IMPL_CLASS_HEADER(T); \
943:   using cupmReal_t   = typename ::Petsc::device::cupm::impl::Interface<T>::cupmReal_t; \
944:   using cupmScalar_t = typename ::Petsc::device::cupm::impl::Interface<T>::cupmScalar_t; \
945:   using ::Petsc::device::cupm::impl::Interface<T>::cupmScalarCast; \
946:   using ::Petsc::device::cupm::impl::Interface<T>::cupmScalarPtrCast; \
947:   using ::Petsc::device::cupm::impl::Interface<T>::cupmRealPtrCast; \
948:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMGetMemType; \
949:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMMemset; \
950:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMMemsetAsync; \
951:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMMalloc; \
952:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMMallocAsync; \
953:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMMallocHost; \
954:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMMemcpy; \
955:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMMemcpyAsync; \
956:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMMemcpy2D; \
957:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMMemcpy2DAsync; \
958:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMMemset2DAsync; \
959:   using ::Petsc::device::cupm::impl::Interface<T>::cupmFree; \
960:   using ::Petsc::device::cupm::impl::Interface<T>::cupmFreeAsync; \
961:   using ::Petsc::device::cupm::impl::Interface<T>::cupmFreeHost; \
962:   using ::Petsc::device::cupm::impl::Interface<T>::cupmLaunchKernel; \
963:   using ::Petsc::device::cupm::impl::Interface<T>::PetscCUPMLaunchKernel1D; \
964:   using ::Petsc::device::cupm::impl::Interface<T>::PetscDeviceCopyModeToCUPMMemcpyKind

966: #if PetscDefined(HAVE_CUDA)
967: extern template struct PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL Interface<DeviceType::CUDA>;
968: #endif

970: #if PetscDefined(HAVE_HIP)
971: extern template struct PETSC_SINGLE_LIBRARY_VISIBILITY_INTERNAL Interface<DeviceType::HIP>;
972: #endif

974: } // namespace impl

976: } // namespace cupm

978: } // namespace device

980: } // namespace Petsc