[llvm] [flang][cuda] Handle zero sized allocation correctly (PR #160929)

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

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

Like on the host allocate 1 byte when zero size is requested. 

>From 627953b8ded3edd0594f6f9ed3e6888186df5973 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Fri, 26 Sep 2025 11:07:18 -0700
Subject: [PATCH] [flang][cuda] Handle zero sized allocation correctly

---
 flang-rt/lib/cuda/memory.cpp               | 25 +++++++++++-----------
 flang-rt/unittests/Runtime/CUDA/Memory.cpp |  6 ++++++
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/flang-rt/lib/cuda/memory.cpp b/flang-rt/lib/cuda/memory.cpp
index d830580e6a066..78270fef07c36 100644
--- a/flang-rt/lib/cuda/memory.cpp
+++ b/flang-rt/lib/cuda/memory.cpp
@@ -25,23 +25,22 @@ 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) {
-      if (Fortran::runtime::executionEnvironment.cudaDeviceIsManaged) {
-        CUDA_REPORT_IF_ERROR(
-            cudaMallocManaged((void **)&ptr, bytes, cudaMemAttachGlobal));
-      } else {
-        CUDA_REPORT_IF_ERROR(cudaMalloc((void **)&ptr, bytes));
-      }
-    } else if (type == kMemTypeManaged || type == kMemTypeUnified) {
+  bytes = bytes ? bytes : 1;
+  if (type == kMemTypeDevice) {
+    if (Fortran::runtime::executionEnvironment.cudaDeviceIsManaged) {
       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");
+      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;
 }
diff --git a/flang-rt/unittests/Runtime/CUDA/Memory.cpp b/flang-rt/unittests/Runtime/CUDA/Memory.cpp
index f2e17870f7999..c84c54a1376e5 100644
--- a/flang-rt/unittests/Runtime/CUDA/Memory.cpp
+++ b/flang-rt/unittests/Runtime/CUDA/Memory.cpp
@@ -35,6 +35,12 @@ TEST(MemoryCUFTest, SimpleAllocTramsferFree) {
   RTNAME(CUFMemFree)((void *)dev, kMemTypeDevice, __FILE__, __LINE__);
 }
 
+TEST(MemoryCUFTest, AllocZero) {
+  int *dev = (int *)RTNAME(CUFMemAlloc)(0, kMemTypeDevice, __FILE__, __LINE__);
+  EXPECT_TRUE(dev != 0);
+  RTNAME(CUFMemFree)((void *)dev, kMemTypeDevice, __FILE__, __LINE__);
+}
+
 static OwningPtr<Descriptor> createAllocatable(
     Fortran::common::TypeCategory tc, int kind, int rank = 1) {
   return Descriptor::Create(TypeCode{tc, kind}, kind, nullptr, rank, nullptr,