[flang-commits] [flang] 19de5dd - [flang][cuda] Add CUFSetAssociatedStream entry point (#181313)

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2026-02-13T19:15:44Z
New Revision: 19de5dd722a4919cf29e67aee22af006033f7f20

URL: https://github.com/llvm/llvm-project/commit/19de5dd722a4919cf29e67aee22af006033f7f20
DIFF: https://github.com/llvm/llvm-project/commit/19de5dd722a4919cf29e67aee22af006033f7f20.diff

LOG: [flang][cuda] Add CUFSetAssociatedStream entry point (#181313)

Added: 
    

Modified: 
    flang-rt/lib/cuda/allocator.cpp
    flang-rt/unittests/Runtime/CUDA/Allocatable.cpp
    flang/include/flang/Runtime/CUDA/allocator.h

Removed: 
    


################################################################################
diff  --git a/flang-rt/lib/cuda/allocator.cpp b/flang-rt/lib/cuda/allocator.cpp
index dc3ce0ee1b590..917b279b38f3c 100644
--- a/flang-rt/lib/cuda/allocator.cpp
+++ b/flang-rt/lib/cuda/allocator.cpp
@@ -140,6 +140,21 @@ cudaStream_t RTDECL(CUFGetAssociatedStream)(void *p) {
   }
   return nullptr;
 }
+
+int RTDECL(CUFSetAssociatedStream)(void *p, cudaStream_t stream, bool hasStat,
+    const Descriptor *errMsg, const char *sourceFile, int sourceLine) {
+  Terminator terminator{sourceFile, sourceLine};
+  if (p == nullptr) {
+    return ReturnError(terminator, StatBaseNull, errMsg, hasStat);
+  }
+  int pos = findAllocation(p);
+  if (pos >= 0) {
+    deviceAllocations[pos].stream = stream;
+  } else {
+    insertAllocation(p, 0, stream);
+  }
+  return StatOk;
+}
 }
 
 void *CUFAllocPinned(

diff  --git a/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp b/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp
index f061c082cc614..e308e8c8bdadb 100644
--- a/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp
+++ b/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp
@@ -11,6 +11,7 @@
 #include "gtest/gtest.h"
 #include "flang-rt/runtime/allocator-registry.h"
 #include "flang-rt/runtime/descriptor.h"
+#include "flang-rt/runtime/stat.h"
 #include "flang-rt/runtime/terminator.h"
 #include "flang/Runtime/CUDA/allocator.h"
 #include "flang/Runtime/CUDA/common.h"
@@ -172,3 +173,39 @@ 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);
+
+  int stat1 = RTDECL(CUFSetAssociatedStream)(a->raw().base_addr, stream);
+  EXPECT_EQ(stat1, StatOk);
+  s = RTDECL(CUFGetAssociatedStream)(a->raw().base_addr);
+  EXPECT_EQ(s, stream);
+
+  // REAL(4), DEVICE, ALLOCATABLE :: b(:) - unallocated, base_addr is null
+  auto b{createAllocatable(TypeCategory::Real, 4)};
+  int stat2 = RTDECL(CUFSetAssociatedStream)(
+      b->raw().base_addr, stream, true, nullptr, __FILE__, __LINE__);
+  EXPECT_EQ(stat2, StatBaseNull);
+}

diff  --git a/flang/include/flang/Runtime/CUDA/allocator.h b/flang/include/flang/Runtime/CUDA/allocator.h
index 56176360296a9..6a64bdeccbc2c 100644
--- a/flang/include/flang/Runtime/CUDA/allocator.h
+++ b/flang/include/flang/Runtime/CUDA/allocator.h
@@ -21,6 +21,9 @@ extern "C" {
 
 void RTDECL(CUFRegisterAllocator)();
 cudaStream_t RTDECL(CUFGetAssociatedStream)(void *);
+int RTDECL(CUFSetAssociatedStream)(void *, cudaStream_t, bool hasStat = false,
+    const Descriptor *errMsg = nullptr, const char *sourceFile = nullptr,
+    int sourceLine = 0);
 }
 
 void *CUFAllocPinned(std::size_t, std::int64_t *);