[flang][cuda][NFC] Move set/get default stream to its own file (#181927)

flang-rt/lib/cuda/CMakeLists.txt

@@ -16,6 +16,7 @@ add_flangrt_library(flang_rt.cuda STATIC SHARED
   memory.cpp
   pointer.cpp
   registration.cpp
+  stream.cpp
 
   TARGET_PROPERTIES
     # libflang_rt.runtime depends on a certain version of CUDA. To be able to have

flang-rt/lib/cuda/allocator.cpp

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

flang-rt/lib/cuda/stream.cpp

@@ -0,0 +1,34 @@
+//===-- lib/cuda/stream.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 "flang/Runtime/CUDA/stream.h"
+#include "flang-rt/runtime/allocator-registry.h"
+#include "flang-rt/runtime/derived.h"
+#include "flang-rt/runtime/descriptor.h"
+#include "flang-rt/runtime/environment.h"
+#include "flang-rt/runtime/lock.h"
+#include "flang-rt/runtime/stat.h"
+#include "flang-rt/runtime/terminator.h"
+#include "flang/Runtime/CUDA/common.h"
+#include "flang/Support/Fortran.h"
+
+namespace Fortran::runtime::cuda {
+
+static thread_local cudaStream_t defaultStream{nullptr};
+
+extern "C" {
+
+int RTDECL(CUFSetDefaultStream)(cudaStream_t stream) {
+  defaultStream = stream;
+  return StatOk;
+}
+
+cudaStream_t RTDECL(CUFGetDefaultStream)() { return defaultStream; }
+}
+
+} // namespace Fortran::runtime::cuda

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

@@ -16,6 +16,7 @@
 #include "flang/Runtime/CUDA/allocator.h"
 #include "flang/Runtime/CUDA/common.h"
 #include "flang/Runtime/CUDA/descriptor.h"
+#include "flang/Runtime/CUDA/stream.h"
 #include "flang/Support/Fortran.h"
 
 using namespace Fortran::runtime;

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

@@ -9,6 +9,7 @@
 #include "cuda_runtime.h"
 #include "gtest/gtest.h"
 #include "flang/Runtime/CUDA/allocator.h"
+#include "flang/Runtime/CUDA/stream.h"
 
 using namespace Fortran::runtime;
 using namespace Fortran::runtime::cuda;

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

@@ -18,12 +18,9 @@
 namespace Fortran::runtime::cuda {
 
 extern "C" {
-
-void RTDECL(CUFRegisterAllocator)();
 cudaStream_t RTDECL(CUFGetAssociatedStream)(void *);
 int RTDECL(CUFSetAssociatedStream)(void *, cudaStream_t);
-int RTDECL(CUFSetDefaultStream)(cudaStream_t);
-cudaStream_t RTDECL(CUFGetDefaultStream)();
+void RTDECL(CUFRegisterAllocator)();
 }
 
 void *CUFAllocPinned(std::size_t, std::int64_t *);

flang/include/flang/Runtime/CUDA/stream.h

@@ -0,0 +1,27 @@
+//===-- include/flang/Runtime/CUDA/stream.h ---------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_RUNTIME_CUDA_STREAM_H_
+#define FORTRAN_RUNTIME_CUDA_STREAM_H_
+
+#include "common.h"
+#include "flang/Runtime/descriptor-consts.h"
+#include "flang/Runtime/entry-names.h"
+
+#include "cuda_runtime.h"
+
+namespace Fortran::runtime::cuda {
+
+extern "C" {
+
+int RTDECL(CUFSetDefaultStream)(cudaStream_t);
+cudaStream_t RTDECL(CUFGetDefaultStream)();
+}
+
+} // namespace Fortran::runtime::cuda
+#endif // FORTRAN_RUNTIME_CUDA_STREAM_H_