[flang-commits] [flang] [flang][cuda] Add entry point to launch cuda fortran kernel (PR #113490)

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

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

None

>From a7f34fb38d150e5dd0985e45849ddc96ec6d270d Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Wed, 23 Oct 2024 11:44:14 -0700
Subject: [PATCH] [flang][cuda] Add entry point to launch cuda fortran kernel

---
 flang/include/flang/Runtime/CUDA/kernel.h | 29 ++++++++++++++++++
 flang/runtime/CUDA/CMakeLists.txt         |  1 +
 flang/runtime/CUDA/kernel.cpp             | 37 +++++++++++++++++++++++
 3 files changed, 67 insertions(+)
 create mode 100644 flang/include/flang/Runtime/CUDA/kernel.h
 create mode 100644 flang/runtime/CUDA/kernel.cpp

diff --git a/flang/include/flang/Runtime/CUDA/kernel.h b/flang/include/flang/Runtime/CUDA/kernel.h
new file mode 100644
index 00000000000000..f08cea8c1e4a7d
--- /dev/null
+++ b/flang/include/flang/Runtime/CUDA/kernel.h
@@ -0,0 +1,29 @@
+//===-- include/flang/Runtime/CUDA/kernel.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_KERNEL_H_
+#define FORTRAN_RUNTIME_CUDA_KERNEL_H_
+
+#include "flang/Runtime/entry-names.h"
+#include <cstddef>
+#include <stdint.h>
+
+namespace Fortran::runtime::cuda {
+
+extern "C" {
+
+// This function uses intptr_t instead of CUDA's unsigned int to match
+// the type of MLIR's index type. This avoids the need for casts in the
+// generated MLIR code.
+void RTDEF(CUFLaunchKernel)(const void *kernelName, intptr_t gridX,
+    intptr_t gridY, intptr_t gridZ, intptr_t blockX, intptr_t blockY,
+    intptr_t blockZ, int32_t smem, void **params, void **extra);
+
+} // extern "C"
+
+} // namespace Fortran::runtime::cuda
diff --git a/flang/runtime/CUDA/CMakeLists.txt b/flang/runtime/CUDA/CMakeLists.txt
index 86523b419f8711..ce87f3efdc3632 100644
--- a/flang/runtime/CUDA/CMakeLists.txt
+++ b/flang/runtime/CUDA/CMakeLists.txt
@@ -17,6 +17,7 @@ add_flang_library(${CUFRT_LIBNAME}
   allocator.cpp
   allocatable.cpp
   descriptor.cpp
+  kernel.cpp
   memory.cpp
   registration.cpp
 )
diff --git a/flang/runtime/CUDA/kernel.cpp b/flang/runtime/CUDA/kernel.cpp
new file mode 100644
index 00000000000000..844aefd59462dc
--- /dev/null
+++ b/flang/runtime/CUDA/kernel.cpp
@@ -0,0 +1,37 @@
+//===-- runtime/CUDA/kernel.cpp -------------------------------------------===//
+//
+// 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/kernel.h"
+#include "../terminator.h"
+#include "flang/Runtime/CUDA/common.h"
+
+#include "cuda_runtime.h"
+
+namespace Fortran::runtime::cuda {
+
+extern "C" {
+
+void RTDEF(CUFLaunchKernel)(const void *kernel, intptr_t gridX, intptr_t gridY,
+    intptr_t gridZ, intptr_t blockX, intptr_t blockY, intptr_t blockZ,
+    int32_t smem, void **params, void **extra) {
+  dim3 gridDim;
+  gridDim.x = gridX;
+  gridDim.y = gridY;
+  gridDim.z = gridZ;
+  dim3 blockDim;
+  blockDim.x = blockX;
+  blockDim.y = blockY;
+  blockDim.z = blockZ;
+  cudaStream_t stream = 0;
+  CUDA_REPORT_IF_ERROR(
+      cudaLaunchKernel(kernel, gridDim, blockDim, params, smem, stream));
+}
+
+} // extern "C"
+
+} // namespace Fortran::runtime::cuda