[llvm] [flang][cuda] Use the provided stream in kernel launch (PR #135267)

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

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

None

>From b2752c98b9cf4aa43f7227eec762ccc3df93884a Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 10 Apr 2025 14:56:03 -0700
Subject: [PATCH] [flang][cuda] Use the provided stream in kernel launch

---
 flang-rt/lib/cuda/kernel.cpp | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/flang-rt/lib/cuda/kernel.cpp b/flang-rt/lib/cuda/kernel.cpp
index 6b60b72630a15..bc2f2d49447df 100644
--- a/flang-rt/lib/cuda/kernel.cpp
+++ b/flang-rt/lib/cuda/kernel.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "flang/Runtime/CUDA/kernel.h"
+#include "flang-rt/runtime/descriptor.h"
 #include "flang-rt/runtime/terminator.h"
 #include "flang/Runtime/CUDA/common.h"
 
@@ -74,9 +75,9 @@ void RTDEF(CUFLaunchKernel)(const void *kernel, intptr_t gridX, intptr_t gridY,
     Fortran::runtime::Terminator terminator{__FILE__, __LINE__};
     terminator.Crash("Too many invalid grid dimensions");
   }
-  cudaStream_t cuStream = 0; // TODO stream managment
-  CUDA_REPORT_IF_ERROR(
-      cudaLaunchKernel(kernel, gridDim, blockDim, params, smem, cuStream));
+  cudaStream_t deafultStream = 0;
+  CUDA_REPORT_IF_ERROR(cudaLaunchKernel(kernel, gridDim, blockDim, params, smem,
+      stream != kNoAsyncId ? (cudaStream_t)stream : deafultStream));
 }
 
 void RTDEF(CUFLaunchClusterKernel)(const void *kernel, intptr_t clusterX,
@@ -140,7 +141,11 @@ void RTDEF(CUFLaunchClusterKernel)(const void *kernel, intptr_t clusterX,
     terminator.Crash("Too many invalid grid dimensions");
   }
   config.dynamicSmemBytes = smem;
-  config.stream = 0; // TODO stream managment
+  if (stream != kNoAsyncId) {
+    config.stream = (cudaStream_t)stream;
+  } else {
+    config.stream = 0;
+  }
   cudaLaunchAttribute launchAttr[1];
   launchAttr[0].id = cudaLaunchAttributeClusterDimension;
   launchAttr[0].val.clusterDim.x = clusterX;
@@ -212,9 +217,10 @@ void RTDEF(CUFLaunchCooperativeKernel)(const void *kernel, intptr_t gridX,
     Fortran::runtime::Terminator terminator{__FILE__, __LINE__};
     terminator.Crash("Too many invalid grid dimensions");
   }
-  cudaStream_t cuStream = 0; // TODO stream managment
-  CUDA_REPORT_IF_ERROR(cudaLaunchCooperativeKernel(
-      kernel, gridDim, blockDim, params, smem, cuStream));
+  cudaStream_t deafultStream = 0;
+  CUDA_REPORT_IF_ERROR(
+      cudaLaunchCooperativeKernel(kernel, gridDim, blockDim, params, smem,
+          stream != kNoAsyncId ? (cudaStream_t)stream : deafultStream));
 }
 
 } // extern "C"