[Flang] Add Assign_omp fortran-rt function

This commit is c-p from ivanradanov commit
d7e4499

Author: skc7

flang-rt/lib/runtime/CMakeLists.txt

@@ -21,6 +21,7 @@ set(supported_sources
   allocatable.cpp
   array-constructor.cpp
   assign.cpp
+  assign_omp.cpp
   buffer.cpp
   character.cpp
   connection.cpp
@@ -97,6 +98,7 @@ set(gpu_sources
   allocatable.cpp
   array-constructor.cpp
   assign.cpp
+  assign_omp.cpp
   buffer.cpp
   character.cpp
   connection.cpp

flang-rt/lib/runtime/assign_omp.cpp

@@ -0,0 +1,72 @@
+//===-- lib/runtime/assign_omp.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/assign.h"
+#include "flang-rt/runtime/assign-impl.h"
+#include "flang-rt/runtime/derived.h"
+#include "flang-rt/runtime/descriptor.h"
+#include "flang-rt/runtime/stat.h"
+#include "flang-rt/runtime/terminator.h"
+#include "flang-rt/runtime/tools.h"
+#include "flang-rt/runtime/type-info.h"
+
+#include <omp.h>
+
+namespace Fortran::runtime {
+
+RT_API_ATTRS static void Assign(Descriptor &to, const Descriptor &from,
+    Terminator &terminator, int flags, int32_t omp_device) {
+  std::size_t toElementBytes{to.ElementBytes()};
+  std::size_t fromElementBytes{from.ElementBytes()};
+  std::size_t toElements{to.Elements()};
+  std::size_t fromElements{from.Elements()};
+
+  if (toElementBytes != fromElementBytes)
+    terminator.Crash("Assign: toElementBytes != fromElementBytes");
+  if (toElements != fromElements)
+    terminator.Crash("Assign: toElements != fromElements");
+
+  void *host_to_ptr = to.raw().base_addr;
+  void *host_from_ptr = from.raw().base_addr;
+  size_t length = toElements * toElementBytes;
+
+  printf("assign length: %zu\n", length);
+
+  if (!omp_target_is_present(host_to_ptr, omp_device))
+    terminator.Crash("Assign: !omp_target_is_present(host_to_ptr, omp_device)");
+  if (!omp_target_is_present(host_from_ptr, omp_device))
+    terminator.Crash(
+        "Assign: !omp_target_is_present(host_from_ptr, omp_device)");
+
+  printf("host_to_ptr: %p\n", host_to_ptr);
+#pragma omp target data use_device_ptr(host_to_ptr, host_from_ptr) device(omp_device)
+  {
+    printf("device_to_ptr: %p\n", host_to_ptr);
+    // TODO do we need to handle overlapping memory? does this function do that?
+    omp_target_memcpy(host_to_ptr, host_from_ptr, length, /*dst_offset*/ 0,
+        /*src_offset*/ 0, /*dst*/ omp_device, /*src*/ omp_device);
+  }
+
+  return;
+}
+
+extern "C" {
+RT_EXT_API_GROUP_BEGIN
+void RTDEF(Assign_omp)(Descriptor &to, const Descriptor &from,
+    const char *sourceFile, int sourceLine, int32_t omp_device) {
+  Terminator terminator{sourceFile, sourceLine};
+  // All top-level defined assignments can be recognized in semantics and
+  // will have been already been converted to calls, so don't check for
+  // defined assignment apart from components.
+  Assign(to, from, terminator,
+      MaybeReallocate | NeedFinalization | ComponentCanBeDefinedAssignment,
+      omp_device);
+}
+} // extern "C"
+
+}

flang/include/flang/Runtime/assign.h

@@ -55,6 +55,8 @@ extern "C" {
 // API for lowering assignment
 void RTDECL(Assign)(Descriptor &to, const Descriptor &from,
     const char *sourceFile = nullptr, int sourceLine = 0);
+void RTDECL(Assign_omp)(Descriptor &to, const Descriptor &from,
+    const char *sourceFile = nullptr, int sourceLine = 0, int32_t omp_device = 0);
 // This variant has no finalization, defined assignment, or allocatable
 // reallocation.
 void RTDECL(AssignTemporary)(Descriptor &to, const Descriptor &from,

flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp

@@ -809,7 +809,7 @@ static void moveToHost(omp::TargetOp targetOp, RewriterBase &rewriter) {
         auto module = runtimeCall->getParentOfType<ModuleOp>();
         auto callee = cast<func::FuncOp>(
             module.lookupSymbol(runtimeCall.getCalleeAttr()));
-        std::string newCalleeName = (callee.getName()).str();
+        std::string newCalleeName = (callee.getName() + "_omp").str();
         mlir::OpBuilder moduleBuilder(module.getBodyRegion());
         func::FuncOp newCallee =
             cast_or_null<func::FuncOp>(module.lookupSymbol(newCalleeName));