[flang-commits] [flang] [llvm] [flang][cuda] Expose some function to deal with async device allocation (PR #220770)
Valentin Clement バレンタイン クレメン via flang-commits
flang-commits at lists.llvm.org
Wed Sep 2 16:56:58 PDT 2026
https://github.com/clementval created https://github.com/llvm/llvm-project/pull/220770
None
>From 6e5f749964603caad9af539142e4f5f1067b3a04 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Wed, 2 Sep 2026 16:55:47 -0700
Subject: [PATCH] [flang][cuda] Expose some function to deal with async device
allocation
---
flang-rt/lib/cuda/allocator.cpp | 83 +++++++++++---------
flang/include/flang/Runtime/CUDA/allocator.h | 12 +++
2 files changed, 59 insertions(+), 36 deletions(-)
diff --git a/flang-rt/lib/cuda/allocator.cpp b/flang-rt/lib/cuda/allocator.cpp
index d8287d6f696e3..fc9a6924a1fc9 100644
--- a/flang-rt/lib/cuda/allocator.cpp
+++ b/flang-rt/lib/cuda/allocator.cpp
@@ -97,17 +97,17 @@ int compareDeviceAlloc(const void *a, const void *b) {
}
// Dynamic array for tracking asynchronous allocations.
-static DeviceAllocation *deviceAllocations = nullptr;
-Lock lock;
+static DeviceAllocation *asyncDeviceAllocations = nullptr;
+Lock asyncDeviceAllocationTableLock;
static int maxDeviceAllocations{512}; // Initial size
static int numDeviceAllocations{0};
static constexpr int allocNotFound{-1};
-static void initAllocations() {
- if (!deviceAllocations) {
- deviceAllocations = static_cast<DeviceAllocation *>(
+static void initAsyncDeviceAllocations() {
+ if (!asyncDeviceAllocations) {
+ asyncDeviceAllocations = static_cast<DeviceAllocation *>(
malloc(maxDeviceAllocations * sizeof(DeviceAllocation)));
- if (!deviceAllocations) {
+ if (!asyncDeviceAllocations) {
Terminator terminator{__FILE__, __LINE__};
terminator.Crash("Failed to allocate tracking array");
}
@@ -117,16 +117,16 @@ static void initAllocations() {
static void doubleAllocationArray() {
unsigned newSize = maxDeviceAllocations * 2;
DeviceAllocation *newArray = static_cast<DeviceAllocation *>(
- realloc(deviceAllocations, newSize * sizeof(DeviceAllocation)));
+ realloc(asyncDeviceAllocations, newSize * sizeof(DeviceAllocation)));
if (!newArray) {
Terminator terminator{__FILE__, __LINE__};
terminator.Crash("Failed to reallocate tracking array");
}
- deviceAllocations = newArray;
+ asyncDeviceAllocations = newArray;
maxDeviceAllocations = newSize;
}
-static unsigned findAllocation(void *ptr) {
+int findAsyncDeviceAllocation(void *ptr) {
if (numDeviceAllocations == 0) {
return allocNotFound;
}
@@ -140,10 +140,10 @@ static unsigned findAllocation(void *ptr) {
while (left <= right) {
int mid = left + (right - left) / 2;
- if (deviceAllocations[mid].ptr == ptr) {
+ if (asyncDeviceAllocations[mid].ptr == ptr) {
return mid;
}
- if (deviceAllocations[mid].ptr < ptr) {
+ if (asyncDeviceAllocations[mid].ptr < ptr) {
left = mid + 1;
} else {
right = mid - 1;
@@ -152,34 +152,45 @@ static unsigned findAllocation(void *ptr) {
return allocNotFound;
}
-static void insertAllocation(void *ptr, std::size_t size, cudaStream_t stream) {
- CriticalSection critical{lock};
- initAllocations();
+void insertAsyncDeviceAllocation(
+ void *ptr, std::size_t size, cudaStream_t stream) {
+ CriticalSection critical{asyncDeviceAllocationTableLock};
+ initAsyncDeviceAllocations();
if (numDeviceAllocations >= maxDeviceAllocations) {
doubleAllocationArray();
}
- deviceAllocations[numDeviceAllocations].ptr = ptr;
- deviceAllocations[numDeviceAllocations].size = size;
- deviceAllocations[numDeviceAllocations].stream = stream;
+ asyncDeviceAllocations[numDeviceAllocations].ptr = ptr;
+ asyncDeviceAllocations[numDeviceAllocations].size = size;
+ asyncDeviceAllocations[numDeviceAllocations].stream = stream;
++numDeviceAllocations;
- qsort(deviceAllocations, numDeviceAllocations, sizeof(DeviceAllocation),
+ qsort(asyncDeviceAllocations, numDeviceAllocations, sizeof(DeviceAllocation),
compareDeviceAlloc);
}
-static void eraseAllocation(int pos) {
- deviceAllocations[pos].ptr = nullptr;
- deviceAllocations[pos].size = 0;
- deviceAllocations[pos].stream = (cudaStream_t)0;
- qsort(deviceAllocations, numDeviceAllocations, sizeof(DeviceAllocation),
+cudaStream_t getAsyncDeviceAllocationStream(int pos) {
+ if (pos < 0 || pos >= numDeviceAllocations) {
+ return nullptr;
+ }
+ return asyncDeviceAllocations[pos].stream;
+}
+
+void eraseAsyncDeviceAllocation(int pos) {
+ if (pos < 0 || pos >= numDeviceAllocations) {
+ return;
+ }
+ asyncDeviceAllocations[pos].ptr = nullptr;
+ asyncDeviceAllocations[pos].size = 0;
+ asyncDeviceAllocations[pos].stream = (cudaStream_t)0;
+ qsort(asyncDeviceAllocations, numDeviceAllocations, sizeof(DeviceAllocation),
compareDeviceAlloc);
--numDeviceAllocations;
}
void CUFResetStream(cudaStream_t stream) {
- CriticalSection critical{lock};
+ CriticalSection critical{asyncDeviceAllocationTableLock};
for (int i = 0; i < numDeviceAllocations; ++i) {
- if (deviceAllocations[i].stream == stream) {
- deviceAllocations[i].stream = nullptr;
+ if (asyncDeviceAllocations[i].stream == stream) {
+ asyncDeviceAllocations[i].stream = nullptr;
}
}
}
@@ -200,9 +211,9 @@ void RTDEF(CUFRegisterAllocator)() {
bool RTDEF(CUFDeviceIsActive)() { return !deviceContextTornDown(); }
cudaStream_t RTDECL(CUFGetAssociatedStream)(void *p) {
- int pos = findAllocation(p);
+ int pos = findAsyncDeviceAllocation(p);
if (pos >= 0) {
- cudaStream_t stream = deviceAllocations[pos].stream;
+ cudaStream_t stream = asyncDeviceAllocations[pos].stream;
return stream;
}
return nullptr;
@@ -212,11 +223,11 @@ int RTDECL(CUFSetAssociatedStream)(void *p, cudaStream_t stream) {
if (p == nullptr) {
return StatBaseNull;
}
- int pos = findAllocation(p);
+ int pos = findAsyncDeviceAllocation(p);
if (pos >= 0) {
- deviceAllocations[pos].stream = stream;
+ asyncDeviceAllocations[pos].stream = stream;
} else {
- insertAllocation(p, 0, stream);
+ insertAsyncDeviceAllocation(p, 0, stream);
}
return StatOk;
}
@@ -244,7 +255,7 @@ void *CUFAllocDevice(std::size_t sizeInBytes,
} else {
CUDA_REPORT_IF_ERROR(
cudaMallocAsync(&p, sizeInBytes, (cudaStream_t)*asyncObject));
- insertAllocation(p, sizeInBytes, (cudaStream_t)*asyncObject);
+ insertAsyncDeviceAllocation(p, sizeInBytes, (cudaStream_t)*asyncObject);
}
}
return p;
@@ -253,11 +264,11 @@ void *CUFAllocDevice(std::size_t sizeInBytes,
// Scope-exit cleanup is guarded in lowering; explicit deallocation after a
// reset is unsupported, keeping cudaFreeAsync free of context-query overhead.
void CUFFreeDevice(void *p) {
- CriticalSection critical{lock};
- int pos = findAllocation(p);
+ CriticalSection critical{asyncDeviceAllocationTableLock};
+ int pos = findAsyncDeviceAllocation(p);
if (pos >= 0) {
- cudaStream_t stream = deviceAllocations[pos].stream;
- eraseAllocation(pos);
+ cudaStream_t stream = asyncDeviceAllocations[pos].stream;
+ eraseAsyncDeviceAllocation(pos);
CUDA_REPORT_IF_ERROR(cudaFreeAsync(p, stream));
} else {
CUDA_REPORT_IF_ERROR(cudaFree(p));
diff --git a/flang/include/flang/Runtime/CUDA/allocator.h b/flang/include/flang/Runtime/CUDA/allocator.h
index 8c1e0f0f4f690..4b38e30903809 100644
--- a/flang/include/flang/Runtime/CUDA/allocator.h
+++ b/flang/include/flang/Runtime/CUDA/allocator.h
@@ -15,6 +15,10 @@
#include "cuda_runtime.h"
+namespace Fortran::runtime {
+class Lock;
+}
+
namespace Fortran::runtime::cuda {
extern "C" {
@@ -25,6 +29,14 @@ void RTDECL(CUFRegisterAllocator)();
void CUFResetStream(cudaStream_t stream);
+extern Lock asyncDeviceAllocationTableLock;
+
+int findAsyncDeviceAllocation(void *ptr);
+void insertAsyncDeviceAllocation(
+ void *ptr, std::size_t size, cudaStream_t stream);
+cudaStream_t getAsyncDeviceAllocationStream(int pos);
+void eraseAsyncDeviceAllocation(int pos);
+
void *CUFAllocPinned(std::size_t, std::size_t, std::int64_t *);
void CUFFreePinned(void *);
More information about the flang-commits
mailing list