[flang-commits] [flang] fa627d9 - [flang][cuda] Add entry point for alloc/free and simple copy (#109867)

via flang-commits flang-commits at lists.llvm.org
Tue Sep 24 20:00:16 PDT 2024


Author: Valentin Clement (バレンタイン クレメン)
Date: 2024-09-24T20:00:11-07:00
New Revision: fa627d98e87504b6f6d621a7dab5d140340ed760

URL: https://github.com/llvm/llvm-project/commit/fa627d98e87504b6f6d621a7dab5d140340ed760
DIFF: https://github.com/llvm/llvm-project/commit/fa627d98e87504b6f6d621a7dab5d140340ed760.diff

LOG: [flang][cuda] Add entry point for alloc/free and simple copy (#109867)

These will be used to translate simple cuf.alloc/cuf.free and
cuf.data_transfer on scalar and constant size arrays.

Added: 
    flang/unittests/Runtime/CUDA/Memory.cpp

Modified: 
    flang/include/flang/Runtime/CUDA/common.h
    flang/include/flang/Runtime/CUDA/memory.h
    flang/runtime/CUDA/memory.cpp
    flang/unittests/Runtime/CUDA/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Runtime/CUDA/common.h b/flang/include/flang/Runtime/CUDA/common.h
index cb8681da161f0d..b73bc390ea8c9e 100644
--- a/flang/include/flang/Runtime/CUDA/common.h
+++ b/flang/include/flang/Runtime/CUDA/common.h
@@ -12,6 +12,13 @@
 #include "flang/Runtime/descriptor.h"
 #include "flang/Runtime/entry-names.h"
 
+/// Type of memory for allocation/deallocation
+static constexpr unsigned kMemTypeDevice = 0;
+static constexpr unsigned kMemTypeManaged = 1;
+static constexpr unsigned kMemTypeUnified = 2;
+static constexpr unsigned kMemTypePinned = 3;
+
+/// Data transfer kinds.
 static constexpr unsigned kHostToDevice = 0;
 static constexpr unsigned kDeviceToHost = 1;
 static constexpr unsigned kDeviceToDevice = 2;

diff  --git a/flang/include/flang/Runtime/CUDA/memory.h b/flang/include/flang/Runtime/CUDA/memory.h
index 33947248dc4831..3c3ae73d4ad7a1 100644
--- a/flang/include/flang/Runtime/CUDA/memory.h
+++ b/flang/include/flang/Runtime/CUDA/memory.h
@@ -17,12 +17,24 @@ namespace Fortran::runtime::cuda {
 
 extern "C" {
 
+/// Allocate memory on the device.
+void *RTDECL(CUFMemAlloc)(std::size_t bytes, unsigned type,
+    const char *sourceFile = nullptr, int sourceLine = 0);
+
+/// Free memory allocated on the device.
+void RTDECL(CUFMemFree)(void *devicePtr, unsigned type,
+    const char *sourceFile = nullptr, int sourceLine = 0);
+
 /// Set value to the data hold by a descriptor. The \p value pointer must be
 /// addressable to the same amount of bytes specified by the element size of
 /// the descriptor \p desc.
 void RTDECL(CUFMemsetDescriptor)(const Descriptor &desc, void *value,
     const char *sourceFile = nullptr, int sourceLine = 0);
 
+/// Data transfer from a pointer to a pointer.
+void RTDECL(CUFDataTransferPtrPtr)(void *dst, void *src, std::size_t bytes,
+    unsigned mode, const char *sourceFile = nullptr, int sourceLine = 0);
+
 /// Data transfer from a pointer to a descriptor.
 void RTDECL(CUFDataTransferDescPtr)(const Descriptor &dst, void *src,
     std::size_t bytes, unsigned mode, const char *sourceFile = nullptr,

diff  --git a/flang/runtime/CUDA/memory.cpp b/flang/runtime/CUDA/memory.cpp
index a287fa14a48789..fc48b4343eea9d 100644
--- a/flang/runtime/CUDA/memory.cpp
+++ b/flang/runtime/CUDA/memory.cpp
@@ -8,12 +8,47 @@
 
 #include "flang/Runtime/CUDA/memory.h"
 #include "../terminator.h"
+#include "flang/Runtime/CUDA/common.h"
 
 #include "cuda_runtime.h"
 
 namespace Fortran::runtime::cuda {
 extern "C" {
 
+void *RTDEF(CUFMemAlloc)(
+    std::size_t bytes, unsigned type, const char *sourceFile, int sourceLine) {
+  void *ptr = nullptr;
+  if (bytes != 0) {
+    if (type == kMemTypeDevice) {
+      CUDA_REPORT_IF_ERROR(cudaMalloc((void **)&ptr, bytes));
+    } else if (type == kMemTypeManaged || type == kMemTypeUnified) {
+      CUDA_REPORT_IF_ERROR(
+          cudaMallocManaged((void **)&ptr, bytes, cudaMemAttachGlobal));
+    } else if (type == kMemTypePinned) {
+      CUDA_REPORT_IF_ERROR(cudaMallocHost((void **)&ptr, bytes));
+    } else {
+      Terminator terminator{sourceFile, sourceLine};
+      terminator.Crash("unsupported memory type");
+    }
+  }
+  return ptr;
+}
+
+void RTDEF(CUFMemFree)(
+    void *ptr, unsigned type, const char *sourceFile, int sourceLine) {
+  if (!ptr)
+    return;
+  if (type == kMemTypeDevice || type == kMemTypeManaged ||
+      type == kMemTypeUnified) {
+    CUDA_REPORT_IF_ERROR(cudaFree(ptr));
+  } else if (type == kMemTypePinned) {
+    CUDA_REPORT_IF_ERROR(cudaFreeHost(ptr));
+  } else {
+    Terminator terminator{sourceFile, sourceLine};
+    terminator.Crash("unsupported memory type");
+  }
+}
+
 void RTDEF(CUFMemsetDescriptor)(const Descriptor &desc, void *value,
     const char *sourceFile, int sourceLine) {
   Terminator terminator{sourceFile, sourceLine};
@@ -21,6 +56,23 @@ void RTDEF(CUFMemsetDescriptor)(const Descriptor &desc, void *value,
                    "value to a descriptor");
 }
 
+void RTDEF(CUFDataTransferPtrPtr)(void *dst, void *src, std::size_t bytes,
+    unsigned mode, const char *sourceFile, int sourceLine) {
+  cudaMemcpyKind kind;
+  if (mode == kHostToDevice) {
+    kind = cudaMemcpyHostToDevice;
+  } else if (mode == kDeviceToHost) {
+    kind = cudaMemcpyDeviceToHost;
+  } else if (mode == kDeviceToDevice) {
+    kind = cudaMemcpyDeviceToDevice;
+  } else {
+    Terminator terminator{sourceFile, sourceLine};
+    terminator.Crash("host to host copy not supported");
+  }
+  // TODO: Use cudaMemcpyAsync when we have support for stream.
+  CUDA_REPORT_IF_ERROR(cudaMemcpy(dst, src, bytes, kind));
+}
+
 void RTDEF(CUFDataTransferDescPtr)(const Descriptor &desc, void *addr,
     std::size_t bytes, unsigned mode, const char *sourceFile, int sourceLine) {
   Terminator terminator{sourceFile, sourceLine};

diff  --git a/flang/unittests/Runtime/CUDA/CMakeLists.txt b/flang/unittests/Runtime/CUDA/CMakeLists.txt
index 30fb8c220233c0..a7fe604d687bdd 100644
--- a/flang/unittests/Runtime/CUDA/CMakeLists.txt
+++ b/flang/unittests/Runtime/CUDA/CMakeLists.txt
@@ -3,6 +3,7 @@ if (FLANG_CUF_RUNTIME)
 add_flang_unittest(FlangCufRuntimeTests
   Allocatable.cpp
   AllocatorCUF.cpp
+  Memory.cpp
 )
 
 if (BUILD_SHARED_LIBS)

diff  --git a/flang/unittests/Runtime/CUDA/Memory.cpp b/flang/unittests/Runtime/CUDA/Memory.cpp
new file mode 100644
index 00000000000000..157d3cdb531def
--- /dev/null
+++ b/flang/unittests/Runtime/CUDA/Memory.cpp
@@ -0,0 +1,31 @@
+//===-- flang/unittests/Runtime/Memory.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/memory.h"
+#include "gtest/gtest.h"
+#include "../../../runtime/terminator.h"
+#include "flang/Common/Fortran.h"
+#include "flang/Runtime/CUDA/common.h"
+
+#include "cuda_runtime.h"
+
+using namespace Fortran::runtime::cuda;
+
+TEST(MemoryCUFTest, SimpleAllocTramsferFree) {
+  int *dev = (int *)RTNAME(CUFMemAlloc)(
+      sizeof(int), kMemTypeDevice, __FILE__, __LINE__);
+  EXPECT_TRUE(dev != 0);
+  int host = 42;
+  RTNAME(CUFDataTransferPtrPtr)
+  ((void *)dev, (void *)&host, sizeof(int), kHostToDevice, __FILE__, __LINE__);
+  host = 0;
+  RTNAME(CUFDataTransferPtrPtr)
+  ((void *)&host, (void *)dev, sizeof(int), kDeviceToHost, __FILE__, __LINE__);
+  EXPECT_EQ(42, host);
+  RTNAME(CUFMemFree)((void *)dev, kMemTypeDevice, __FILE__, __LINE__);
+}


        


More information about the flang-commits mailing list