[flang-commits] [flang] [flang][cuda] Adding runtime call to CUFRegisterVariable (PR #113952)

[Renaud Kauffmann via flang-commits]

https://github.com/Renaud-K created https://github.com/llvm/llvm-project/pull/113952

None

>From 9aba1d2855eb546be484218873e126abd19bad3e Mon Sep 17 00:00:00 2001
From: Renaud-K <rkauffmann at nvidia.com>
Date: Mon, 28 Oct 2024 12:23:59 -0700
Subject: [PATCH] [flang][cuda] Adding runtime call to CUFRegisterVariable

---
 flang/include/flang/Runtime/CUDA/registration.h | 5 +++++
 flang/runtime/CUDA/registration.cpp             | 8 ++++++++
 2 files changed, 13 insertions(+)

diff --git a/flang/include/flang/Runtime/CUDA/registration.h b/flang/include/flang/Runtime/CUDA/registration.h
index 009715613e29f7..5237069a4c739c 100644
--- a/flang/include/flang/Runtime/CUDA/registration.h
+++ b/flang/include/flang/Runtime/CUDA/registration.h
@@ -11,6 +11,7 @@
 
 #include "flang/Runtime/entry-names.h"
 #include <cstddef>
+#include <cstdint>
 
 namespace Fortran::runtime::cuda {
 
@@ -23,6 +24,10 @@ void *RTDECL(CUFRegisterModule)(void *data);
 void RTDECL(CUFRegisterFunction)(
     void **module, const char *fctSym, char *fctName);
 
+/// Register a device variable.
+void RTDECL(CUFRegisterVariable)(
+    void **module, char *varSym, const char *varName, int64_t size);
+
 } // extern "C"
 
 } // namespace Fortran::runtime::cuda
diff --git a/flang/runtime/CUDA/registration.cpp b/flang/runtime/CUDA/registration.cpp
index 20d274c4d8d1c2..b7b6ef389bffba 100644
--- a/flang/runtime/CUDA/registration.cpp
+++ b/flang/runtime/CUDA/registration.cpp
@@ -21,6 +21,9 @@ extern void __cudaRegisterFatBinaryEnd(void *);
 extern void __cudaRegisterFunction(void **fatCubinHandle, const char *hostFun,
     char *deviceFun, const char *deviceName, int thread_limit, uint3 *tid,
     uint3 *bid, dim3 *bDim, dim3 *gDim, int *wSize);
+extern void __cudaRegisterVar(void **fatCubinHandle, char *hostVar,
+    const char *deviceAddress, const char *deviceName, int ext, size_t size,
+    int constant, int global);
 
 void *RTDECL(CUFRegisterModule)(void *data) {
   void **fatHandle{__cudaRegisterFatBinary(data)};
@@ -34,6 +37,11 @@ void RTDEF(CUFRegisterFunction)(
       (uint3 *)0, (dim3 *)0, (dim3 *)0, (int *)0);
 }
 
+void RTDEF(CUFRegisterVariable)(
+    void **module, char *varSym, const char *varName, int64_t size) {
+  __cudaRegisterVar(module, varSym, varName, varName, 0, size, 0, 0);
+}
+
 } // extern "C"
 
 } // namespace Fortran::runtime::cuda