[flang-commits] [flang] [llvm] [flang][cuda] Fix return value for CUFSetDefaultStream (PR #181884)

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

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

The interface return an integer value but the entry point and lowering were missing it. 

>From 924523ec6d5af86e9a6470c8a122e51ec9ee17b4 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Tue, 17 Feb 2026 11:15:38 -0800
Subject: [PATCH] [flang][cuda] Fix return value for CUFSetDefaultStream

---
 flang-rt/lib/cuda/allocator.cpp                  |  3 ++-
 flang/include/flang/Runtime/CUDA/allocator.h     |  2 +-
 .../lib/Optimizer/Builder/CUDAIntrinsicCall.cpp  |  2 +-
 flang/test/Lower/CUDA/cuda-default-stream.cuf    | 16 ++++++++++++++++
 4 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/flang-rt/lib/cuda/allocator.cpp b/flang-rt/lib/cuda/allocator.cpp
index 795b26b641720..347195926def3 100644
--- a/flang-rt/lib/cuda/allocator.cpp
+++ b/flang-rt/lib/cuda/allocator.cpp
@@ -156,8 +156,9 @@ int RTDECL(CUFSetAssociatedStream)(void *p, cudaStream_t stream) {
   return StatOk;
 }
 
-void RTDECL(CUFSetDefaultStream)(cudaStream_t stream) {
+int RTDECL(CUFSetDefaultStream)(cudaStream_t stream) {
   defaultStream = stream;
+  return StatOk;
 }
 
 cudaStream_t RTDECL(CUFGetDefaultStream)() { return defaultStream; }
diff --git a/flang/include/flang/Runtime/CUDA/allocator.h b/flang/include/flang/Runtime/CUDA/allocator.h
index d21f7d7c421bb..c45b97a6df4fe 100644
--- a/flang/include/flang/Runtime/CUDA/allocator.h
+++ b/flang/include/flang/Runtime/CUDA/allocator.h
@@ -22,7 +22,7 @@ extern "C" {
 void RTDECL(CUFRegisterAllocator)();
 cudaStream_t RTDECL(CUFGetAssociatedStream)(void *);
 int RTDECL(CUFSetAssociatedStream)(void *, cudaStream_t);
-void RTDECL(CUFSetDefaultStream)(cudaStream_t);
+int RTDECL(CUFSetDefaultStream)(cudaStream_t);
 cudaStream_t RTDECL(CUFGetDefaultStream)();
 }
 
diff --git a/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp b/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp
index 4986c57048081..bbc353634cd42 100644
--- a/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp
+++ b/flang/lib/Optimizer/Builder/CUDAIntrinsicCall.cpp
@@ -1131,7 +1131,7 @@ fir::ExtendedValue CUDAIntrinsicLibrary::genCUDASetDefaultStream(
   mlir::Value stream = fir::getBase(args[0]);
   mlir::Type i64Ty = builder.getI64Type();
   auto ctx = builder.getContext();
-  mlir::FunctionType ftype = mlir::FunctionType::get(ctx, {i64Ty}, {});
+  mlir::FunctionType ftype = mlir::FunctionType::get(ctx, {i64Ty}, {resTy});
   auto funcOp =
       builder.createFunction(loc, RTNAME_STRING(CUFSetDefaultStream), ftype);
   auto call = fir::CallOp::create(builder, loc, funcOp, {stream});
diff --git a/flang/test/Lower/CUDA/cuda-default-stream.cuf b/flang/test/Lower/CUDA/cuda-default-stream.cuf
index 59c6bc6b70612..beacb409f44f6 100644
--- a/flang/test/Lower/CUDA/cuda-default-stream.cuf
+++ b/flang/test/Lower/CUDA/cuda-default-stream.cuf
@@ -22,3 +22,19 @@ end subroutine
 ! CHECK: %[[VOIDPTR:.*]] = fir.convert %[[ADDR]] : (!fir.heap<!fir.array<?xi32>>) -> !fir.llvm_ptr<i8>
 ! CHECK: %[[STREAM:.*]] = fir.call @_FortranACUFGetAssociatedStream(%[[VOIDPTR]]) fastmath<contract> : (!fir.llvm_ptr<i8>) -> i64
 ! CHECK: hlfir.assign %[[STREAM]] to %{{.*}}#0 : i64, !fir.ref<i64>
+
+subroutine default_stream
+  use cuda_runtime_api
+  
+  integer(kind=cuda_stream_kind) :: strm, strm2
+  integer :: istat
+  istat = cudaStreamCreate(strm2)
+  istat = cudaforSetDefaultStream(strm2)
+  strm = cudaforGetDefaultStream()
+  istat = cudaStreamSynchronize(cudaforGetDefaultStream())
+end subroutine
+
+! CHECK-LABEL: func.func @_QPdefault_stream()
+! CHECK: %{{.*}} = fir.call @_FortranACUFSetDefaultStream(%{{.*}}) fastmath<contract> : (i64) -> i32
+! CHECK: %{{.*}} = fir.call @_FortranACUFGetDefaultStream() fastmath<contract> : () -> i64
+! CHECK: %{{.*}} = fir.call @_FortranACUFGetDefaultStream() fastmath<contract> : () -> i64