[flang-commits] [flang] [flang][cuda] Force default allocator in device code (PR #102238)

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

https://github.com/clementval updated https://github.com/llvm/llvm-project/pull/102238

>From 89ba16b84a2142da229bc9e887e62c8604d5d743 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Tue, 6 Aug 2024 15:04:53 -0700
Subject: [PATCH 1/3] [flang][cuda] Force default allocator in device code

---
 flang/runtime/descriptor.cpp | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/flang/runtime/descriptor.cpp b/flang/runtime/descriptor.cpp
index 34f7a02ea8c7bd..6d43bacaed697f 100644
--- a/flang/runtime/descriptor.cpp
+++ b/flang/runtime/descriptor.cpp
@@ -162,11 +162,17 @@ RT_API_ATTRS int Descriptor::Allocate() {
     elementBytes = raw_.elem_len = 0;
   }
   std::size_t byteSize{Elements() * elementBytes};
+
+  // Force default allocator in device code.
+#ifdef RT_DEVICE_COMPILATION
+  AllocFct alloc{allocatorRegistry.GetAllocator(kDefaultAllocator)};
+#else
+  AllocFct alloc{allocatorRegistry.GetAllocator(GetAllocIdx())};
+#endif
+
   // Zero size allocation is possible in Fortran and the resulting
   // descriptor must be allocated/associated. Since std::malloc(0)
   // result is implementation defined, always allocate at least one byte.
-
-  AllocFct alloc{allocatorRegistry.GetAllocator(GetAllocIdx())};
   void *p{alloc(byteSize ? byteSize : 1)};
   if (!p) {
     return CFI_ERROR_MEM_ALLOCATION;
@@ -209,7 +215,12 @@ RT_API_ATTRS int Descriptor::Deallocate() {
   if (!descriptor.base_addr) {
     return CFI_ERROR_BASE_ADDR_NULL;
   } else {
+    // Force default deallocator in device code.
+#ifdef RT_DEVICE_COMPILATION
+    FreeFct free{allocatorRegistry.GetDeallocator(kDefaultAllocator)};
+#else
     FreeFct free{allocatorRegistry.GetDeallocator(GetAllocIdx())};
+#endif
     free(descriptor.base_addr);
     descriptor.base_addr = nullptr;
     return CFI_SUCCESS;

>From c5faa72b42500a97b23c10f8af2088631adde59d Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Tue, 6 Aug 2024 17:01:04 -0700
Subject: [PATCH 2/3] Put logic in a function

---
 flang/include/flang/Runtime/descriptor.h |  8 ++++++++
 flang/runtime/descriptor.cpp             | 16 ++--------------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/flang/include/flang/Runtime/descriptor.h b/flang/include/flang/Runtime/descriptor.h
index 043f6931afad93..8f08dfd857e662 100644
--- a/flang/include/flang/Runtime/descriptor.h
+++ b/flang/include/flang/Runtime/descriptor.h
@@ -436,6 +436,14 @@ class Descriptor {
   RT_API_ATTRS inline int GetAllocIdx() const {
     return (raw_.extra & _CFI_ALLOCATOR_IDX_MASK) >> _CFI_ALLOCATOR_IDX_SHIFT;
   }
+  RT_API_ATTRS inline int GetNormalizedAllocIdx() const {
+#ifdef RT_DEVICE_COMPILATION
+    // Force default allocator in device code.
+    return kDefaultAllocator;
+#else
+    return GetAllocIdx();
+#endif
+  }
   RT_API_ATTRS inline void SetAllocIdx(int pos) {
     raw_.extra &= ~_CFI_ALLOCATOR_IDX_MASK; // Clear the allocator index bits.
     raw_.extra |= (pos << _CFI_ALLOCATOR_IDX_SHIFT);
diff --git a/flang/runtime/descriptor.cpp b/flang/runtime/descriptor.cpp
index 6d43bacaed697f..74a3f069700bef 100644
--- a/flang/runtime/descriptor.cpp
+++ b/flang/runtime/descriptor.cpp
@@ -162,14 +162,7 @@ RT_API_ATTRS int Descriptor::Allocate() {
     elementBytes = raw_.elem_len = 0;
   }
   std::size_t byteSize{Elements() * elementBytes};
-
-  // Force default allocator in device code.
-#ifdef RT_DEVICE_COMPILATION
-  AllocFct alloc{allocatorRegistry.GetAllocator(kDefaultAllocator)};
-#else
-  AllocFct alloc{allocatorRegistry.GetAllocator(GetAllocIdx())};
-#endif
-
+  AllocFct alloc{allocatorRegistry.GetAllocator(GetNormalizedAllocIdx())};
   // Zero size allocation is possible in Fortran and the resulting
   // descriptor must be allocated/associated. Since std::malloc(0)
   // result is implementation defined, always allocate at least one byte.
@@ -215,12 +208,7 @@ RT_API_ATTRS int Descriptor::Deallocate() {
   if (!descriptor.base_addr) {
     return CFI_ERROR_BASE_ADDR_NULL;
   } else {
-    // Force default deallocator in device code.
-#ifdef RT_DEVICE_COMPILATION
-    FreeFct free{allocatorRegistry.GetDeallocator(kDefaultAllocator)};
-#else
-    FreeFct free{allocatorRegistry.GetDeallocator(GetAllocIdx())};
-#endif
+    FreeFct free{allocatorRegistry.GetDeallocator(GetNormalizedAllocIdx())};
     free(descriptor.base_addr);
     descriptor.base_addr = nullptr;
     return CFI_SUCCESS;

>From cfa14bf5dca708ef56bfd08b376218e7b7a43a88 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 8 Aug 2024 10:07:18 -0700
Subject: [PATCH 3/3] Add MapAllocIdx

---
 flang/include/flang/Runtime/descriptor.h |  8 --------
 flang/runtime/descriptor.cpp             | 13 +++++++++++--
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/flang/include/flang/Runtime/descriptor.h b/flang/include/flang/Runtime/descriptor.h
index 8f08dfd857e662..043f6931afad93 100644
--- a/flang/include/flang/Runtime/descriptor.h
+++ b/flang/include/flang/Runtime/descriptor.h
@@ -436,14 +436,6 @@ class Descriptor {
   RT_API_ATTRS inline int GetAllocIdx() const {
     return (raw_.extra & _CFI_ALLOCATOR_IDX_MASK) >> _CFI_ALLOCATOR_IDX_SHIFT;
   }
-  RT_API_ATTRS inline int GetNormalizedAllocIdx() const {
-#ifdef RT_DEVICE_COMPILATION
-    // Force default allocator in device code.
-    return kDefaultAllocator;
-#else
-    return GetAllocIdx();
-#endif
-  }
   RT_API_ATTRS inline void SetAllocIdx(int pos) {
     raw_.extra &= ~_CFI_ALLOCATOR_IDX_MASK; // Clear the allocator index bits.
     raw_.extra |= (pos << _CFI_ALLOCATOR_IDX_SHIFT);
diff --git a/flang/runtime/descriptor.cpp b/flang/runtime/descriptor.cpp
index 74a3f069700bef..10919913ebaa88 100644
--- a/flang/runtime/descriptor.cpp
+++ b/flang/runtime/descriptor.cpp
@@ -154,6 +154,15 @@ RT_API_ATTRS std::size_t Descriptor::Elements() const {
   return elements;
 }
 
+RT_API_ATTRS static int MapAllocIdx(const Descriptor &desc) {
+#ifdef RT_DEVICE_COMPILATION
+  // Force default allocator in device code.
+  return kDefaultAllocator;
+#else
+  return desc.GetAllocIdx();
+#endif
+}
+
 RT_API_ATTRS int Descriptor::Allocate() {
   std::size_t elementBytes{ElementBytes()};
   if (static_cast<std::int64_t>(elementBytes) < 0) {
@@ -162,7 +171,7 @@ RT_API_ATTRS int Descriptor::Allocate() {
     elementBytes = raw_.elem_len = 0;
   }
   std::size_t byteSize{Elements() * elementBytes};
-  AllocFct alloc{allocatorRegistry.GetAllocator(GetNormalizedAllocIdx())};
+  AllocFct alloc{allocatorRegistry.GetAllocator(MapAllocIdx(*this))};
   // Zero size allocation is possible in Fortran and the resulting
   // descriptor must be allocated/associated. Since std::malloc(0)
   // result is implementation defined, always allocate at least one byte.
@@ -208,7 +217,7 @@ RT_API_ATTRS int Descriptor::Deallocate() {
   if (!descriptor.base_addr) {
     return CFI_ERROR_BASE_ADDR_NULL;
   } else {
-    FreeFct free{allocatorRegistry.GetDeallocator(GetNormalizedAllocIdx())};
+    FreeFct free{allocatorRegistry.GetDeallocator(MapAllocIdx(*this))};
     free(descriptor.base_addr);
     descriptor.base_addr = nullptr;
     return CFI_SUCCESS;