[flang-commits] [flang] [llvm] [flang][cuda] Add CUFSetAssociatedStream entry point (PR #181313)

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/181313

None

>From 03deb53cb502ac5e022bb25858cf722596047225 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 12 Feb 2026 22:15:20 -0800
Subject: [PATCH] [flang][cuda] Add CUFSetAssociatedStream entry point

---
 flang-rt/lib/cuda/allocator.cpp               | 10 +++++++
 .../unittests/Runtime/CUDA/Allocatable.cpp    | 29 +++++++++++++++++++
 flang/include/flang/Runtime/CUDA/allocator.h  |  1 +
 3 files changed, 40 insertions(+)

diff --git a/flang-rt/lib/cuda/allocator.cpp b/flang-rt/lib/cuda/allocator.cpp
index dc3ce0ee1b590..7aeabc2e12e1c 100644
--- a/flang-rt/lib/cuda/allocator.cpp
+++ b/flang-rt/lib/cuda/allocator.cpp
@@ -140,6 +140,16 @@ cudaStream_t RTDECL(CUFGetAssociatedStream)(void *p) {
   }
   return nullptr;
 }
+
+void RTDECL(CUFSetAssociatedStream)(void *p, cudaStream_t stream) {
+  int pos = findAllocation(p);
+  if (pos >= 0) {
+    deviceAllocations[pos].stream = stream;
+  } else {
+    insertAllocation(p, 0, stream);
+  }
+}
+
 }
 
 void *CUFAllocPinned(
diff --git a/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp b/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp
index f061c082cc614..38d215fd0c0c5 100644
--- a/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp
+++ b/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp
@@ -172,3 +172,32 @@ TEST(AllocatableAsyncTest, StreamDeviceAllocatable) {
   cudaStream_t empty = RTDECL(CUFGetAssociatedStream)(a->raw().base_addr);
   EXPECT_EQ(empty, nullptr);
 }
+
+TEST(AllocatableAsyncTest, SetStreamTest) {
+    using Fortran::common::TypeCategory;
+    RTNAME(CUFRegisterAllocator)();
+    // REAL(4), DEVICE, ALLOCATABLE :: a(:)
+    auto a{createAllocatable(TypeCategory::Real, 4)};
+    a->SetAllocIdx(kDeviceAllocatorPos);
+    EXPECT_EQ((int)kDeviceAllocatorPos, a->GetAllocIdx());
+    EXPECT_FALSE(a->HasAddendum());
+    RTNAME(AllocatableSetBounds)(*a, 0, 1, 10);
+  
+    cudaStream_t stream;
+    cudaStreamCreate(&stream);
+    EXPECT_EQ(cudaSuccess, cudaGetLastError());
+  
+    RTNAME(AllocatableAllocate)
+    (*a, /*asyncObject=*/nullptr, /*hasStat=*/false,
+        /*errMsg=*/nullptr, __FILE__, __LINE__);
+    EXPECT_TRUE(a->IsAllocated());
+    cudaDeviceSynchronize();
+    EXPECT_EQ(cudaSuccess, cudaGetLastError());
+    cudaStream_t defaultStream = 0;
+    cudaStream_t s = RTDECL(CUFGetAssociatedStream)(a->raw().base_addr);
+    EXPECT_EQ(s, defaultStream);
+
+    RTDECL(CUFSetAssociatedStream)(a->raw().base_addr, stream);
+    s = RTDECL(CUFGetAssociatedStream)(a->raw().base_addr);
+    EXPECT_EQ(s, stream);
+  }
diff --git a/flang/include/flang/Runtime/CUDA/allocator.h b/flang/include/flang/Runtime/CUDA/allocator.h
index 56176360296a9..e87fe2799af1e 100644
--- a/flang/include/flang/Runtime/CUDA/allocator.h
+++ b/flang/include/flang/Runtime/CUDA/allocator.h
@@ -21,6 +21,7 @@ extern "C" {
 
 void RTDECL(CUFRegisterAllocator)();
 cudaStream_t RTDECL(CUFGetAssociatedStream)(void *);
+void RTDECL(CUFSetAssociatedStream)(void *, cudaStream_t);
 }
 
 void *CUFAllocPinned(std::size_t, std::int64_t *);