[Openmp-commits] [openmp] [OpenMP] Fix issue of indirect function call in `__kmpc_fork_call_if` (PR #65436)

[Shilei Tian via Openmp-commits]

https://github.com/shiltian created https://github.com/llvm/llvm-project/pull/65436:

The outlined function is typically invoked by using `__kmp_invoke_microtask`,
which is written in asm. D138495 introduces a new interface function for parallel
region for OpenMPIRBuilder, where the outlined function is called via the function
pointer. For some reason, it works perfectly well on x86 and x86-64 system, but
doesn't work on aarch64. The 3rd argument in the callee is always `nullptr`, even
if it is not in caller. It appears `x2` always contains `0x0`. This patch adopts
the typical method to invoke the function pointer. It works on my M2 Ultra Mac.

Fix #63194.


>From 4335e96a00f9d4c0f8e4ef32b86f30415e4daff0 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Wed, 6 Sep 2023 00:04:20 -0400
Subject: [PATCH] [OpenMP] Fix issue of indirect function call in
 `__kmpc_fork_call_if`

The outlined function is typically invoked by using `__kmp_invoke_microtask`,
which is written in asm. D138495 introduces a new interface function for parallel
region for OpenMPIRBuilder, where the outlined function is called via the function
pointer. For some reason, it works perfectly well on x86 and x86-64 system, but
doesn't work on aarch64. The 3rd argument in the callee is always `nullptr`, even
if it is not in caller. It appears `x2` always contains `0x0`. This patch adopts
the typical method to invoke the function pointer. It works on my M2 Ultra Mac.

Fix #63194.
---
 openmp/runtime/src/kmp_csupport.cpp | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/openmp/runtime/src/kmp_csupport.cpp b/openmp/runtime/src/kmp_csupport.cpp
index 8bd0e89a7dacd95..fb79adb936412cd 100644
--- a/openmp/runtime/src/kmp_csupport.cpp
+++ b/openmp/runtime/src/kmp_csupport.cpp
@@ -343,7 +343,6 @@ Perform a fork only if the condition is true.
 void __kmpc_fork_call_if(ident_t *loc, kmp_int32 argc, kmpc_micro microtask,
                          kmp_int32 cond, void *args) {
   int gtid = __kmp_entry_gtid();
-  int zero = 0;
   if (cond) {
     if (args)
       __kmpc_fork_call(loc, argc, microtask, args);
@@ -352,10 +351,29 @@ void __kmpc_fork_call_if(ident_t *loc, kmp_int32 argc, kmpc_micro microtask,
   } else {
     __kmpc_serialized_parallel(loc, gtid);
 
+#ifdef OMPT_SUPPORT
+    void *exit_frame_ptr;
+#endif
+
     if (args)
-      microtask(&gtid, &zero, args);
+      __kmp_invoke_microtask(VOLATILE_CAST(microtask_t) microtask, gtid,
+                             /*npr=*/0,
+                             /*argc=*/1, &args
+#ifdef OMPT_SUPPORT
+                             ,
+                             &exit_frame_ptr
+#endif
+      );
     else
-      microtask(&gtid, &zero);
+      __kmp_invoke_microtask(VOLATILE_CAST(microtask_t) microtask, gtid,
+                             /*npr=*/0,
+                             /*argc=*/0,
+                             /*args=*/nullptr
+#ifdef OMPT_SUPPORT
+                             ,
+                             &exit_frame_ptr
+#endif
+      );
 
     __kmpc_end_serialized_parallel(loc, gtid);
   }