Reland '[flang][runtime] Allow different memmov function in assign' (#114587)

Reland #114301

Author: clementval

flang/include/flang/Runtime/assign.h

@@ -24,11 +24,40 @@
 #define FORTRAN_RUNTIME_ASSIGN_H_
 
 #include "flang/Runtime/entry-names.h"
+#include "flang/Runtime/freestanding-tools.h"
 
 namespace Fortran::runtime {
 class Descriptor;
+class Terminator;
+
+enum AssignFlags {
+  NoAssignFlags = 0,
+  MaybeReallocate = 1 << 0,
+  NeedFinalization = 1 << 1,
+  CanBeDefinedAssignment = 1 << 2,
+  ComponentCanBeDefinedAssignment = 1 << 3,
+  ExplicitLengthCharacterLHS = 1 << 4,
+  PolymorphicLHS = 1 << 5,
+  DeallocateLHS = 1 << 6
+};
+
+using MemmoveFct = void *(*)(void *, const void *, std::size_t);
+
+#ifdef RT_DEVICE_COMPILATION
+static RT_API_ATTRS void *MemmoveWrapper(
+    void *dest, const void *src, std::size_t count) {
+  return Fortran::runtime::memmove(dest, src, count);
+}
+RT_API_ATTRS void Assign(Descriptor &to, const Descriptor &from,
+    Terminator &terminator, int flags, MemmoveFct memmoveFct = &MemmoveWrapper);
+#else
+RT_API_ATTRS void Assign(Descriptor &to, const Descriptor &from,
+    Terminator &terminator, int flags,
+    MemmoveFct memmoveFct = &Fortran::runtime::memmove);
+#endif
 
 extern "C" {
+
 // API for lowering assignment
 void RTDECL(Assign)(Descriptor &to, const Descriptor &from,
     const char *sourceFile = nullptr, int sourceLine = 0);

flang/include/flang/Runtime/freestanding-tools.h

@@ -82,7 +82,7 @@ using std::fill_n;
 #if STD_MEMMOVE_UNSUPPORTED
 // Provides alternative implementation for std::memmove(), if
 // it is not supported.
-static inline RT_API_ATTRS void memmove(
+static inline RT_API_ATTRS void *memmove(
     void *dest, const void *src, std::size_t count) {
   char *to{reinterpret_cast<char *>(dest)};
   const char *from{reinterpret_cast<const char *>(src)};
@@ -103,6 +103,7 @@ static inline RT_API_ATTRS void memmove(
       *--to = *--from;
     }
   }
+  return dest;
 }
 #else // !STD_MEMMOVE_UNSUPPORTED
 using std::memmove;

flang/runtime/assign.cpp

@@ -17,17 +17,6 @@
 
 namespace Fortran::runtime {
 
-enum AssignFlags {
-  NoAssignFlags = 0,
-  MaybeReallocate = 1 << 0,
-  NeedFinalization = 1 << 1,
-  CanBeDefinedAssignment = 1 << 2,
-  ComponentCanBeDefinedAssignment = 1 << 3,
-  ExplicitLengthCharacterLHS = 1 << 4,
-  PolymorphicLHS = 1 << 5,
-  DeallocateLHS = 1 << 6
-};
-
 // Predicate: is the left-hand side of an assignment an allocated allocatable
 // that must be deallocated?
 static inline RT_API_ATTRS bool MustDeallocateLHS(
@@ -250,8 +239,8 @@ static RT_API_ATTRS void BlankPadCharacterAssignment(Descriptor &to,
 // of elements, but their shape need not to conform (the assignment is done in
 // element sequence order). This facilitates some internal usages, like when
 // dealing with array constructors.
-RT_API_ATTRS static void Assign(
-    Descriptor &to, const Descriptor &from, Terminator &terminator, int flags) {
+RT_API_ATTRS void Assign(Descriptor &to, const Descriptor &from,
+    Terminator &terminator, int flags, MemmoveFct memmoveFct) {
   bool mustDeallocateLHS{(flags & DeallocateLHS) ||
       MustDeallocateLHS(to, from, terminator, flags)};
   DescriptorAddendum *toAddendum{to.Addendum()};
@@ -423,14 +412,14 @@ RT_API_ATTRS static void Assign(
             Assign(toCompDesc, fromCompDesc, terminator, nestedFlags);
           } else { // Component has intrinsic type; simply copy raw bytes
             std::size_t componentByteSize{comp.SizeInBytes(to)};
-            Fortran::runtime::memmove(to.Element<char>(toAt) + comp.offset(),
+            memmoveFct(to.Element<char>(toAt) + comp.offset(),
                 from.Element<const char>(fromAt) + comp.offset(),
                 componentByteSize);
           }
           break;
         case typeInfo::Component::Genre::Pointer: {
           std::size_t componentByteSize{comp.SizeInBytes(to)};
-          Fortran::runtime::memmove(to.Element<char>(toAt) + comp.offset(),
+          memmoveFct(to.Element<char>(toAt) + comp.offset(),
               from.Element<const char>(fromAt) + comp.offset(),
               componentByteSize);
         } break;
@@ -476,14 +465,14 @@ RT_API_ATTRS static void Assign(
         const auto &procPtr{
             *procPtrDesc.ZeroBasedIndexedElement<typeInfo::ProcPtrComponent>(
                 k)};
-        Fortran::runtime::memmove(to.Element<char>(toAt) + procPtr.offset,
+        memmoveFct(to.Element<char>(toAt) + procPtr.offset,
             from.Element<const char>(fromAt) + procPtr.offset,
             sizeof(typeInfo::ProcedurePointer));
       }
     }
   } else { // intrinsic type, intrinsic assignment
     if (isSimpleMemmove()) {
-      Fortran::runtime::memmove(to.raw().base_addr, from.raw().base_addr,
+      memmoveFct(to.raw().base_addr, from.raw().base_addr,
           toElements * toElementBytes);
     } else if (toElementBytes > fromElementBytes) { // blank padding
       switch (to.type().raw()) {
@@ -507,8 +496,8 @@ RT_API_ATTRS static void Assign(
     } else { // elemental copies, possibly with character truncation
       for (std::size_t n{toElements}; n-- > 0;
            to.IncrementSubscripts(toAt), from.IncrementSubscripts(fromAt)) {
-        Fortran::runtime::memmove(to.Element<char>(toAt),
-            from.Element<const char>(fromAt), toElementBytes);
+        memmoveFct(to.Element<char>(toAt), from.Element<const char>(fromAt),
+            toElementBytes);
       }
     }
   }