code/day55-mallocasync/mempool.cuThis is the source used by the lesson and its recorded evidence. Compile commands and expected output live in the directory README.
// SPDX-License-Identifier: MIT
//
// Day 55: stream-ordered allocation and memory pools.
//
// One loop, three ways. Each iteration needs a 16 MiB scratch buffer, runs
// two kernels through it, and lets it go. The first loop calls cudaMalloc
// and cudaFree every iteration, the second hoists one buffer out of the
// loop, the third uses cudaMallocAsync and cudaFreeAsync against the
// device's default memory pool. All three loops produce the same
// accumulator, and the program fails if they do not.
//
// After the timed loops it reads the pool's reserved size twice: once at
// the default release threshold of zero, once after raising the threshold
// to UINT64_MAX, to show when the pool gives its memory back to the OS.
//
// Build: nvcc -std=c++17 -O3 -arch=sm_75 -lineinfo -o mempool mempool.cu
// Run: ./mempool
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <vector>
#include <cuda_runtime.h>
#include <nvtx3/nvToolsExt.h>
// The one error macro. This file is standalone, the way a Compiler Explorer
// embed is, so it carries its own verbatim copy.
#define CUDA_CHECK(call) \
do { \
cudaError_t err_ = (call); \
if (err_ != cudaSuccess) { \
std::fprintf(stderr, "CUDA error %s:%d: %s: %s\n", __FILE__, \
__LINE__, #call, cudaGetErrorString(err_)); \
std::exit(EXIT_FAILURE); \
} \
} while (0)
// kElems is deliberately not a multiple of the block size, so every kernel's
// bounds check runs on every launch. 4 Mi floats is 16 MiB, big enough that
// an allocator that re-maps it from the OS each iteration has real work to
// do, small enough that one hundred iterations finish in well under a
// second of kernel time on a Tesla T4.
constexpr size_t kElems = 4ull * 1024ull * 1024ull + 611ull;
constexpr size_t kScratchBytes = kElems * sizeof(float);
constexpr int kThreadsPerBlock = 256; // 8 warps
constexpr int kIters = 100;
constexpr int kPoolDemoIters = 10;
constexpr float kRelTolerance = 1e-5f;
static_assert(kThreadsPerBlock % 32 == 0,
"block size must be a whole number of warps");
// out[i] = 2 * in[i]. One thread owns one element.
//
// Memory: consecutive threads take consecutive elements, so one warp's 32
// addresses cover 128 contiguous bytes on the load and on the store.
//
// Launch assumption: gridDim.x * blockDim.x >= n. The kernel is deliberately
// trivial: the loop around it, not the kernel, is what this program measures.
__global__ void writeScaled(const float* __restrict__ in,
float* __restrict__ out, size_t n) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i < n) {
out[i] = 2.0f * in[i];
}
}
// acc[i] += in[i]. One thread owns one element.
//
// Memory: same contiguous pattern as writeScaled, plus a read of acc.
// Updates acc in place, which is safe here because every timed phase runs
// exactly once over a freshly zeroed accumulator; nothing launches this
// kernel repeatedly over the same input inside one measurement.
__global__ void accumulate(const float* __restrict__ in,
float* __restrict__ acc, size_t n) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i < n) {
acc[i] += in[i];
}
}
// snippet: per-iter
// The loop everyone writes first. cudaMalloc and cudaFree are not stream
// operations: cudaFree cannot hand back memory the queued kernels still
// use, so it waits for the device to drain, once per iteration.
static void loopMallocPerIter(const float* d_in, float* d_scratch, float* d_acc,
cudaStream_t stream, int blocks) {
(void)d_scratch; // this variant allocates its own
nvtxRangePushA("malloc-per-iter");
for (int iter = 0; iter < kIters; ++iter) {
float* d_tmp = nullptr;
CUDA_CHECK(cudaMalloc(&d_tmp, kScratchBytes));
writeScaled<<<blocks, kThreadsPerBlock, 0, stream>>>(d_in, d_tmp,
kElems);
accumulate<<<blocks, kThreadsPerBlock, 0, stream>>>(d_tmp, d_acc,
kElems);
CUDA_CHECK(cudaFree(d_tmp));
}
CUDA_CHECK(cudaGetLastError());
nvtxRangePop();
}
// end snippet
// The classic fix: hoist the buffer out of the loop. This is the baseline
// the pool has to approach, and the version day 54's pipeline already uses.
static void loopHoisted(const float* d_in, float* d_scratch, float* d_acc,
cudaStream_t stream, int blocks) {
nvtxRangePushA("hoisted");
for (int iter = 0; iter < kIters; ++iter) {
writeScaled<<<blocks, kThreadsPerBlock, 0, stream>>>(d_in, d_scratch,
kElems);
accumulate<<<blocks, kThreadsPerBlock, 0, stream>>>(d_scratch, d_acc,
kElems);
}
CUDA_CHECK(cudaGetLastError());
nvtxRangePop();
}
// snippet: async
// Same shape as the first loop, but allocation and free are operations in
// stream order against the device's default pool. The host queues them and
// moves on; the free returns the block to the pool, and the next iteration's
// cudaMallocAsync reuses it without touching the OS.
static void loopMallocAsync(const float* d_in, float* d_scratch, float* d_acc,
cudaStream_t stream, int blocks) {
(void)d_scratch; // this variant allocates its own
nvtxRangePushA("malloc-async");
for (int iter = 0; iter < kIters; ++iter) {
float* d_tmp = nullptr;
CUDA_CHECK(cudaMallocAsync(&d_tmp, kScratchBytes, stream));
writeScaled<<<blocks, kThreadsPerBlock, 0, stream>>>(d_in, d_tmp,
kElems);
accumulate<<<blocks, kThreadsPerBlock, 0, stream>>>(d_tmp, d_acc,
kElems);
CUDA_CHECK(cudaFreeAsync(d_tmp, stream));
}
CUDA_CHECK(cudaGetLastError());
nvtxRangePop();
}
// end snippet
// Times one phase with events recorded into the stream. No warm-up inside,
// unlike day 9's timeKernel: a phase mutates the accumulator, so warming up
// in here would corrupt the correctness check. main() warms up every kernel
// and both allocators once, explicitly, before any phase is timed.
//
// The events bracket the whole loop on the device clock, so time the device
// spends idle while the host sits inside cudaMalloc or cudaFree is counted,
// which is the point.
static float timePhase(void (*phase)(const float*, float*, float*, cudaStream_t,
int),
const float* d_in, float* d_scratch, float* d_acc,
cudaStream_t stream, int blocks) {
CUDA_CHECK(cudaMemsetAsync(d_acc, 0, kScratchBytes, stream));
cudaEvent_t start;
cudaEvent_t stop;
CUDA_CHECK(cudaEventCreate(&start));
CUDA_CHECK(cudaEventCreate(&stop));
CUDA_CHECK(cudaEventRecord(start, stream));
phase(d_in, d_scratch, d_acc, stream, blocks);
CUDA_CHECK(cudaEventRecord(stop, stream));
CUDA_CHECK(cudaEventSynchronize(stop));
CUDA_CHECK(cudaGetLastError());
float ms = 0.0f;
CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
CUDA_CHECK(cudaEventDestroy(start));
CUDA_CHECK(cudaEventDestroy(stop));
return ms;
}
// Reads how much backing memory the default pool currently holds. The
// documented attribute type is cuuint64_t; uint64_t is the same 64 bits
// without pulling in the driver API header.
static uint64_t poolReservedBytes(cudaMemPool_t pool) {
uint64_t reserved = 0;
CUDA_CHECK(cudaMemPoolGetAttribute(pool, cudaMemPoolAttrReservedMemCurrent,
&reserved));
return reserved;
}
// snippet: threshold
// The release threshold is the amount of memory the pool may keep across a
// synchronisation. At the default of zero, any stream or event synchronise
// hands everything back to the OS; at UINT64_MAX the pool keeps what it has.
static void raiseReleaseThreshold(cudaMemPool_t pool) {
uint64_t threshold = UINT64_MAX;
CUDA_CHECK(cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold,
&threshold));
}
// end snippet
int main() {
const int device = 0;
CUDA_CHECK(cudaSetDevice(device));
cudaDeviceProp prop;
CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
std::printf("GPU: %s (compute capability %d.%d)\n", prop.name, prop.major,
prop.minor);
int poolsSupported = 0;
CUDA_CHECK(cudaDeviceGetAttribute(&poolsSupported,
cudaDevAttrMemoryPoolsSupported, device));
if (poolsSupported == 0) {
std::fprintf(stderr, "this device does not support cudaMallocAsync\n");
return EXIT_FAILURE;
}
const int blocks =
static_cast<int>((kElems + kThreadsPerBlock - 1) / kThreadsPerBlock);
std::printf("n = %zu (%zu bytes scratch), %d iterations per loop\n", kElems,
kScratchBytes, kIters);
// Small whole numbers, so one hundred accumulations of 2 * in[i] stay
// exact in float and a mismatch below can only be a real bug.
std::vector<float> h_in(kElems);
std::vector<float> h_want(kElems);
for (size_t i = 0; i < kElems; ++i) {
h_in[i] = static_cast<float>(i % 33);
h_want[i] = static_cast<float>(2 * kIters) * h_in[i];
}
float* d_in = nullptr;
float* d_scratch = nullptr;
float* d_acc = nullptr;
CUDA_CHECK(cudaMalloc(&d_in, kScratchBytes));
CUDA_CHECK(cudaMalloc(&d_scratch, kScratchBytes));
CUDA_CHECK(cudaMalloc(&d_acc, kScratchBytes));
CUDA_CHECK(
cudaMemcpy(d_in, h_in.data(), kScratchBytes, cudaMemcpyHostToDevice));
cudaStream_t stream;
CUDA_CHECK(cudaStreamCreate(&stream));
// Warm-up covering everything the phases time: both kernels (lazy module
// loading makes each kernel's first launch pay its own load), one
// cudaMalloc/cudaFree pair, and one cudaMallocAsync/cudaFreeAsync pair,
// which also gives the default pool its first block. The accumulator is
// scratch here; timePhase zeroes it before each measured phase.
{
float* d_tmp = nullptr;
CUDA_CHECK(cudaMalloc(&d_tmp, kScratchBytes));
writeScaled<<<blocks, kThreadsPerBlock, 0, stream>>>(d_in, d_tmp,
kElems);
accumulate<<<blocks, kThreadsPerBlock, 0, stream>>>(d_tmp, d_acc,
kElems);
CUDA_CHECK(cudaFree(d_tmp));
float* d_tmp2 = nullptr;
CUDA_CHECK(cudaMallocAsync(&d_tmp2, kScratchBytes, stream));
writeScaled<<<blocks, kThreadsPerBlock, 0, stream>>>(d_in, d_tmp2,
kElems);
CUDA_CHECK(cudaFreeAsync(d_tmp2, stream));
CUDA_CHECK(cudaStreamSynchronize(stream));
CUDA_CHECK(cudaGetLastError());
}
// The three phases. Each zeroes the accumulator, runs its loop under one
// NVTX range, and its result is copied out before the next phase runs.
std::vector<float> h_perIter(kElems);
std::vector<float> h_hoisted(kElems);
std::vector<float> h_async(kElems);
const float msPerIterMalloc =
timePhase(loopMallocPerIter, d_in, d_scratch, d_acc, stream, blocks);
CUDA_CHECK(cudaMemcpy(h_perIter.data(), d_acc, kScratchBytes,
cudaMemcpyDeviceToHost));
const float msHoisted =
timePhase(loopHoisted, d_in, d_scratch, d_acc, stream, blocks);
CUDA_CHECK(cudaMemcpy(h_hoisted.data(), d_acc, kScratchBytes,
cudaMemcpyDeviceToHost));
const float msAsync =
timePhase(loopMallocAsync, d_in, d_scratch, d_acc, stream, blocks);
CUDA_CHECK(cudaMemcpy(h_async.data(), d_acc, kScratchBytes,
cudaMemcpyDeviceToHost));
std::printf("\n%-18s %10s %14s %10s\n", "loop", "total ms", "us per iter",
"vs hoisted");
std::printf("%-18s %10.3f %14.1f %9.2fx\n", "malloc-per-iter",
msPerIterMalloc, 1000.0f * msPerIterMalloc / kIters,
msPerIterMalloc / msHoisted);
std::printf("%-18s %10.3f %14.1f %9.2fx\n", "hoisted", msHoisted,
1000.0f * msHoisted / kIters, 1.0f);
std::printf("%-18s %10.3f %14.1f %9.2fx\n", "malloc-async", msAsync,
1000.0f * msAsync / kIters, msAsync / msHoisted);
// The pool demo. The event synchronise that ended the malloc-async
// phase already counted as a synchronisation point, and the threshold
// is still at its default of zero, so the pool should be empty now.
nvtxRangePushA("pool-demo");
cudaMemPool_t pool; // owned by the device; nothing to create or destroy
CUDA_CHECK(cudaDeviceGetDefaultMemPool(&pool, device));
std::printf("\npool reserved after sync, threshold 0: %llu bytes\n",
static_cast<unsigned long long>(poolReservedBytes(pool)));
raiseReleaseThreshold(pool);
for (int iter = 0; iter < kPoolDemoIters; ++iter) {
float* d_tmp = nullptr;
CUDA_CHECK(cudaMallocAsync(&d_tmp, kScratchBytes, stream));
writeScaled<<<blocks, kThreadsPerBlock, 0, stream>>>(d_in, d_tmp,
kElems);
CUDA_CHECK(cudaFreeAsync(d_tmp, stream));
}
CUDA_CHECK(cudaStreamSynchronize(stream));
CUDA_CHECK(cudaGetLastError());
std::printf("pool reserved after sync, threshold max: %llu bytes\n",
static_cast<unsigned long long>(poolReservedBytes(pool)));
nvtxRangePop();
CUDA_CHECK(cudaFree(d_in));
CUDA_CHECK(cudaFree(d_scratch));
CUDA_CHECK(cudaFree(d_acc));
CUDA_CHECK(cudaStreamDestroy(stream));
// All three loops must produce the identical accumulator. A failure in
// the async loop here would mean a kernel read a block the pool had
// already handed to another allocation, which is exactly the class of
// bug stream ordering exists to prevent.
struct Named {
const char* name;
const std::vector<float>* acc;
};
const Named results[] = {{"malloc-per-iter", &h_perIter},
{"hoisted", &h_hoisted},
{"malloc-async", &h_async}};
for (const Named& r : results) {
for (size_t i = 0; i < kElems; ++i) {
const float want = h_want[i];
const float scale = (want == 0.0f) ? 1.0f : std::fabs(want);
if (std::fabs((*r.acc)[i] - want) > kRelTolerance * scale) {
std::fprintf(stderr, "%s wrong at %zu: got %.9g, want %.9g\n",
r.name, i, (*r.acc)[i], want);
return EXIT_FAILURE;
}
}
}
std::printf("\nall three loops agree with the reference\n");
return EXIT_SUCCESS;
}