pmalloc_sycl.cpp File Reference

#include "gmxpre.h"
#include <exception>
#include <optional>
#include "gromacs/gpu_utils/device_context.h"
#include "gromacs/gpu_utils/gmxsycl.h"
#include "gromacs/utility/gmxassert.h"
#include "gromacs/utility/smalloc.h"
#include "pmalloc.h"

Include dependency graph for pmalloc_sycl.cpp:

Description

Pinned memory allocation routines for SYCL.

This module is a direct wrapper around the sycl::malloc_host / sycl::free, except for the management of the default context.

Unlike in CUDA, pinning memory in SYCL requires a context. It can be passed explicitly to pmalloc and pfree, but that is not straightforward in some calling code. Therefore, we use pmallocSetDefaultDeviceContext and pmallocClearDefaultDeviceContext to manage the default context (stored in g_threadDefaultContext).

That puts a constraint on allocation order: we shall free all the pinned memory before resetting the default context. This is easy to achieve in normal use, but hard to guarantee during stack unwinding when handling an exception. Therefore, we introduce g_threadAllocCount to count the number of allocations that are using the default context. If pmallocClearDefaultDeviceContext is called while handling an exception, we check g_threadAllocCount, and, if there are any remaining allocations, set g_threadDelayContextClearing to defer the context resetting until all the allocations are freed. We also use g_threadAllocCount to control the correctness of memory allocations in normal runs.

GROMACS (at least in 2022 and earlier) uses separate contexts for each rank. Since we support threadMPI, the context management is per-thread, and all the static variables are thread_local.

Author

Andrey Alekseenko al42a.nosp@m.nd@g.nosp@m.mail..nosp@m.com

Functions
void	pmalloc (void **h_ptr, size_t nbytes, const DeviceContext *deviceContext)
	Allocates nbytes of host memory. Use pfree to free memory allocated with this function. More...
void	pfree (void *h_ptr, const DeviceContext *deviceContext)
	Frees memory allocated with pmalloc. More...
void	pmallocSetDefaultDeviceContext (const DeviceContext *deviceContext)
void	pmallocClearDefaultDeviceContext ()

Variables
static thread_local std::optional< const sycl::context >	g_threadDefaultContext = std::nullopt
	Default context to use for pinning memory.
static thread_local int	g_threadAllocCount = 0
	Count the number of memory allocations in the default context.
static thread_local bool	g_threadDelayContextClearing = false
	Whether we should delay resetting the default context because there is still memory allocated there.

Function Documentation

void pfree	(	void *	h_ptr,
		const DeviceContext *	deviceContext
	)

Frees memory allocated with pmalloc.

Parameters

[in]	h_ptr	Buffer allocated with pmalloc that needs to be freed.
[in]	deviceContext	SYCL context to use. Will use the default one (see pmallocSetDefaultDeviceContext) if not set.

void pmalloc	(	void **	h_ptr,
		size_t	nbytes,
		const DeviceContext *	deviceContext
	)

Allocates nbytes of host memory. Use pfree to free memory allocated with this function.

Parameters

[in,out]	h_ptr	Pointer where to store the address of the newly allocated buffer.
[in]	nbytes	Size in bytes of the buffer to be allocated.
[in]	deviceContext	SYCL context to use. Will use the default one (see pmallocSetDefaultDeviceContext) if not set.