From 786b3b4741f944ed32b3a7b4fbd2991c2ad98184 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Valentin=20Clement=20=28=E3=83=90=E3=83=AC=E3=83=B3?= =?UTF-8?q?=E3=82=BF=E3=82=A4=E3=83=B3=20=E3=82=AF=E3=83=AC=E3=83=A1?= =?UTF-8?q?=E3=83=B3=29?= Date: Tue, 17 Feb 2026 14:43:38 -0800 Subject: [PATCH] [flang][cuda][NFC] Move set/get default stream to its own file (#181927) --- flang-rt/lib/cuda/CMakeLists.txt | 1 + flang-rt/lib/cuda/allocator.cpp | 9 ----- flang-rt/lib/cuda/stream.cpp | 34 +++++++++++++++++++ .../unittests/Runtime/CUDA/Allocatable.cpp | 1 + .../unittests/Runtime/CUDA/DefaultStream.cpp | 1 + flang/include/flang/Runtime/CUDA/allocator.h | 5 +-- flang/include/flang/Runtime/CUDA/stream.h | 27 +++++++++++++++ 7 files changed, 65 insertions(+), 13 deletions(-) create mode 100644 flang-rt/lib/cuda/stream.cpp create mode 100644 flang/include/flang/Runtime/CUDA/stream.h diff --git a/flang-rt/lib/cuda/CMakeLists.txt b/flang-rt/lib/cuda/CMakeLists.txt index c8547a1e36cb..118fd20f39e5 100644 --- a/flang-rt/lib/cuda/CMakeLists.txt +++ b/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 diff --git a/flang-rt/lib/cuda/allocator.cpp b/flang-rt/lib/cuda/allocator.cpp index 347195926def..df7e43de00c7 100644 --- a/flang-rt/lib/cuda/allocator.cpp +++ b/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( diff --git a/flang-rt/lib/cuda/stream.cpp b/flang-rt/lib/cuda/stream.cpp new file mode 100644 index 000000000000..af9e63f6a85b --- /dev/null +++ b/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 diff --git a/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp b/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp index 0680c0086ea7..1e98acfd5151 100644 --- a/flang-rt/unittests/Runtime/CUDA/Allocatable.cpp +++ b/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; diff --git a/flang-rt/unittests/Runtime/CUDA/DefaultStream.cpp b/flang-rt/unittests/Runtime/CUDA/DefaultStream.cpp index 7e255905ceb1..4c7886bfdc38 100644 --- a/flang-rt/unittests/Runtime/CUDA/DefaultStream.cpp +++ b/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; diff --git a/flang/include/flang/Runtime/CUDA/allocator.h b/flang/include/flang/Runtime/CUDA/allocator.h index c45b97a6df4f..698b979636da 100644 --- a/flang/include/flang/Runtime/CUDA/allocator.h +++ b/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 *); diff --git a/flang/include/flang/Runtime/CUDA/stream.h b/flang/include/flang/Runtime/CUDA/stream.h new file mode 100644 index 000000000000..ce97d0d06721 --- /dev/null +++ b/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_