[flang][cuda] Add entry point for set/get default stream (#181440)

flang-rt/lib/cuda/allocator.cpp

@@ -21,6 +21,8 @@
 
 namespace Fortran::runtime::cuda {
 
+static thread_local cudaStream_t defaultStream{nullptr};
+
 struct DeviceAllocation {
   void *ptr;
   std::size_t size;
@@ -153,6 +155,12 @@ int RTDECL(CUFSetAssociatedStream)(void *p, cudaStream_t stream) {
   }
   return StatOk;
 }
+
+void RTDECL(CUFSetDefaultStream)(cudaStream_t stream) {
+  defaultStream = stream;
+}
+
+cudaStream_t RTDECL(CUFGetDefaultStream)() { return defaultStream; }
 }
 
 void *CUFAllocPinned(

flang-rt/unittests/Runtime/CUDA/CMakeLists.txt

@@ -9,6 +9,7 @@
 add_flangrt_unittest(FlangCufRuntimeTests
   Allocatable.cpp
   AllocatorCUF.cpp
+  DefaultStream.cpp
   Memory.cpp
 )
 

flang-rt/unittests/Runtime/CUDA/DefaultStream.cpp

@@ -0,0 +1,27 @@
+//===-- unittests/Runtime/CUDA/DefaultStream.cpp ----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "cuda_runtime.h"
+#include "gtest/gtest.h"
+#include "flang/Runtime/CUDA/allocator.h"
+
+using namespace Fortran::runtime;
+using namespace Fortran::runtime::cuda;
+
+TEST(DefaultStreamTest, GetAndSetTest) {
+  using Fortran::common::TypeCategory;
+  cudaStream_t defaultStream = RTDECL(CUFGetDefaultStream)();
+  EXPECT_EQ(defaultStream, nullptr);
+
+  cudaStream_t stream;
+  cudaStreamCreate(&stream);
+  EXPECT_EQ(cudaSuccess, cudaGetLastError());
+  RTDECL(CUFSetDefaultStream)(stream);
+  cudaStream_t outStream = RTDECL(CUFGetDefaultStream)();
+  EXPECT_EQ(outStream, stream);
+}

flang/include/flang/Runtime/CUDA/allocator.h

@@ -22,6 +22,8 @@ extern "C" {
 void RTDECL(CUFRegisterAllocator)();
 cudaStream_t RTDECL(CUFGetAssociatedStream)(void *);
 int RTDECL(CUFSetAssociatedStream)(void *, cudaStream_t);
+void RTDECL(CUFSetDefaultStream)(cudaStream_t);
+cudaStream_t RTDECL(CUFGetDefaultStream)();
 }
 
 void *CUFAllocPinned(std::size_t, std::int64_t *);