RFC: [openmp] Provide an assembly implementation of __kmp_invoke_microtask on ARM

[mstorsjo]

Summary

This fixes passing an arbitrarily large number of arguments to
microtasks, fixing the misc_bugs/many-microtask-args.c testcase on
ARM.

I'm not entirely sure if this works on older arm architectures though;
it's been tested on armv7 (Linux and Windows, with both thumb and arm
modes).

[DavidSpickett]

Not trying to end up with more comments than assembly, but it might be the easiest way to do it :)

openmp/runtime/src/z_Linux_asm.S

1425

You might want to add "for when we call pkfn below".

Since AAPCS32 says align to 4, unless at a public interface where it should be 8. Which includes calling a function.
(which evidently you knew but I did not)

Also worth saying which is the extra register.

1427

What does this load? I think p_argv and exit_frame.

1448

r3 is the arg count. So I think you are rounding up for the case where the count is an odd number, which would leave sp at 4 byte alignment? But if that's the case why do you then multiply by 8 not 4.

Because you divided by 2 in the first place? Making the calculation produce a number of 8 byte slots not 4?

Tell me where I'm going wrong there.

1461

This section you doing:

while (argc != 0) {
   copy and argument to the stack
   argc--
}

call pkfn

Correct?

1463

What does this ldr r2 and the ldr r3 below setup?

2069

And this is an elf/coff difference? No need for a equivalent coff line?

[mstorsjo]

Just BTW - the assembly here isn't entirely written from scratch, but I took the existing aarch64 implementation and rewrote it, so most of the style/structure matches that.

I'll try to add comments to clarify the details that you pointed out. Thanks for the review!

openmp/runtime/src/z_Linux_asm.S

1425

Sure, can do.

1427

Yes, that's correct - I'll add a comment.

1448

Exactly - we want to adjust the stack down 4*argc bytes, but rounded up to an even increment of 8 bytes.

1461

Yes, pretty much.

1463

When calling a function, the first 4 arguments go in r0-r3 - r0/r1 are filled with &gtid and &tid, while r2/r3 should get the first arguments among the regular ones. If argc > 2, then we copy the rest of them onto the stack. In the case of aarch64, we'd copy incoming arguments into x2-x7 and the rest on the stack - here we only need to handle r2/r3 separately.

2069

.size is a ELF-specific assembly directive (and object file feature) - COFF doesn't store per-symbol sizes.

[mstorsjo]

Added more comments on the assembly, simplified a few bits about the stack allocation calculation and setting up of the frame pointer.

[mstorsjo]

Updated to use the assembly version on Windows on ARM too.

[DavidSpickett]

openmp/runtime/src/z_Linux_asm.S

1447

What is the -2 in argc-2? (or rather what would it be)

1452

This 1 here, is it part of the rounding or is it the extra 8 bytes for gtid and tid?

[mstorsjo]

openmp/runtime/src/z_Linux_asm.S

1447

We pass 2 arguments out of argc in registers (while the other 2 register parameters are taken by tid and gtid). In the case of aarch64 with 8 register parameters, they'd need space for argc-6 arguments on the stack.

1452

This is the extra 8 bytes for gtid and tid.

[DavidSpickett]

openmp/runtime/src/z_Linux_asm.S

1452

Got it, that's why you say above that we're over allocating. Because of that we're ok with rounding down (rounding up would need some +1 somewhere).

[DavidSpickett]

Assembly looks good to me.

Now we need some way to tell if earlier Arm needs to be supported. Is there other existing Arm assembly for OpenMP? That would give a clue, might have to go to the git blame to find the intent when it was first added. I'm no expert on the very old architectures but I can help check it for things that look "new" or "old" (for lack of better words).

openmp/runtime/src/z_Linux_asm.S

1447

How about:

// We strictly need 4*(argc-2) bytes (2 arguments are passed in registers)
// but allocate 4*argc for....

[mstorsjo]

In D138704#3956913, @DavidSpickett wrote:

Assembly looks good to me.

Thanks for reviewing!

Now we need some way to tell if earlier Arm needs to be supported. Is there other existing Arm assembly for OpenMP? That would give a clue, might have to go to the git blame to find the intent when it was first added. I'm no expert on the very old architectures but I can help check it for things that look "new" or "old" (for lack of better words).

No, as far as I can see, there's no other existing assembly for ARM here.

AFAIK this implementation should be interworking-safe on armv5 and above, but wouldn't necessarily work as-is for armv4t. It probably won't assemble in thumb1 mode either.

openmp/runtime/src/z_Linux_asm.S

1447

Thanks, that sounds good to me, I'll amend the patch accordingly.

[DavidSpickett]

The oldest architecture distro I know of for Linux is Debian (https://wiki.debian.org/ArmEabiPort) which has now moved to v5. The last supported LTS for armv4t ran out June 2022 (https://www.debian.org/releases/stretch/).

https://clang.llvm.org/docs/OpenMPSupport.html doesn't mention Arm at all. Perhaps it should.

Landing this seems like a low risk to me.

Anyone from the OpenMP side want to check my logic there? @jdoerfert perhaps?

[mstorsjo]

Updated the comment as suggested.

Rebased on top of other recent related pushed commits; use the assembly instead of the C implementation on Windows (just like on aarch64), remove the XFAIL from the testcase that gets fixed by the assembly.

[jdoerfert]

In D138704#3957350, @DavidSpickett wrote:

https://clang.llvm.org/docs/OpenMPSupport.html doesn't mention Arm at all. Perhaps it should.

It probably should, feel free to change that. Also openmp.llvm.org if you want to provide more information.

Landing this seems like a low risk to me.

Agreed.

Anyone from the OpenMP side want to check my logic there? @jdoerfert perhaps?

I did not check "the logic" as I am not familiar with arm but given that the test works and the assembly was proof read I think this is good to go.

[tianshilei1992]

openmp/runtime/src/z_Linux_asm.S

1429

It is reported (https://github.com/llvm/llvm-project/issues/60370) compile error encountered when compiling this. Can you take a look at it?

openmp/runtime/src/CMakeLists.txt

set(LIBOMP_CXXFILES
     kmp_sched.cpp
   )
   if(WIN32)
     # Windows specific files
     libomp_append(LIBOMP_CXXFILES z_Windows_NT_util.cpp)
     libomp_append(LIBOMP_CXXFILES z_Windows_NT-586_util.cpp)
     if(${LIBOMP_ARCH} STREQUAL "i386" OR ${LIBOMP_ARCH} STREQUAL "x86_64")
       libomp_append(LIBOMP_ASMFILES z_Windows_NT-586_asm.asm) # Windows assembly file
-    elseif(${LIBOMP_ARCH} STREQUAL "aarch64" AND (NOT MSVC OR CMAKE_C_COMPILER_ID STREQUAL "Clang"))
-      # z_Linux_asm.S works for AArch64 Windows too.
+    elseif((${LIBOMP_ARCH} STREQUAL "aarch64" OR ${LIBOMP_ARCH} STREQUAL "arm") AND (NOT MSVC OR CMAKE_C_COMPILER_ID STREQUAL "Clang"))
+      # z_Linux_asm.S works for AArch64 and ARM Windows too.
       libomp_append(LIBOMP_GNUASMFILES z_Linux_asm.S)
     else()
       # AArch64 with MSVC gets implementations of the functions from
       # ifdeffed sections in z_Windows_NT-586_util.cpp.
     endif()
   else()
     # Unix specific files
     libomp_append(LIBOMP_CXXFILES z_Linux_util.cpp)

openmp/runtime/src/z_Linux_asm.S

.macro KMP_CFI_REGISTER reg
   .cfi_def_cfa_register \reg
 .endm
 .macro KMP_CFI_DEF reg, sz
   .cfi_def_cfa  \reg,\sz
 .endm
 # endif // KMP_OS_DARWIN
 #endif // KMP_ARCH_X86 || KMP_ARCH_x86_64
 
-#if (KMP_OS_LINUX || KMP_OS_DARWIN || KMP_OS_WINDOWS) && KMP_ARCH_AARCH64
+#if (KMP_OS_LINUX || KMP_OS_DARWIN || KMP_OS_WINDOWS) && (KMP_ARCH_AARCH64 || KMP_ARCH_ARM)
 
 # if KMP_OS_DARWIN
 #  define KMP_PREFIX_UNDERSCORE(x) _##x  // extra underscore for OS X* symbols
 #  define KMP_LABEL(x) L_##x             // form the name of label
 
 .macro ALIGN
   .align $0
 .endmacro

 .macro ALIGN size
   .align 1<<(\size)
 .endm
 
 .macro DEBUG_INFO proc
   .cfi_endproc
 // Not sure why we need .type and .size for the functions
   ALIGN 2
+#if KMP_ARCH_ARM
+  .type  \proc,%function
+#else
   .type  \proc,@function
+#endif
   .size  \proc,.-\proc
 .endm
 
 .macro PROC proc
   ALIGN 2
   .globl KMP_PREFIX_UNDERSCORE(\proc)
 KMP_PREFIX_UNDERSCORE(\proc):
   .cfi_startproc
 .endm
 # endif // KMP_OS_DARWIN
 
-#endif // (KMP_OS_LINUX || KMP_OS_DARWIN || KMP_OS_WINDOWS) && KMP_ARCH_AARCH64
+#endif // (KMP_OS_LINUX || KMP_OS_DARWIN || KMP_OS_WINDOWS) && (KMP_ARCH_AARCH64 || KMP_ARCH_ARM)
 
 .macro COMMON name, size, align_power
 #if KMP_OS_DARWIN
   .comm \name, \size
 #elif KMP_OS_WINDOWS
   .comm \name, \size, \align_power
 #else // !KMP_OS_DARWIN && !KMP_OS_WINDOWS
   .comm \name, \size, (1<<(\align_power))
# endif
   ldp x29, x30, [sp], #16
   ret
 
   DEBUG_INFO __kmp_invoke_microtask
 // -- End  __kmp_invoke_microtask
 
 #endif /* (KMP_OS_LINUX || KMP_OS_DARWIN || KMP_OS_WINDOWS) && KMP_ARCH_AARCH64 */
 
+#if (KMP_OS_LINUX || KMP_OS_DARWIN || KMP_OS_WINDOWS) && KMP_ARCH_ARM
+
+//------------------------------------------------------------------------
+// int
+// __kmp_invoke_microtask( void (*pkfn) (int gtid, int tid, ...),
+//                         int gtid, int tid,
+//                         int argc, void *p_argv[]
+// #if OMPT_SUPPORT
+//                         ,
+//                         void **exit_frame_ptr
+// #endif
+//                       ) {
+// #if OMPT_SUPPORT
+//   *exit_frame_ptr = OMPT_GET_FRAME_ADDRESS(0);
+// #endif
+//
+//   (*pkfn)( & gtid, & tid, argv[0], ... );
+//
+// // FIXME: This is done at call-site and can be removed here.
+// #if OMPT_SUPPORT
+//   *exit_frame_ptr = 0;
+// #endif
+//
+//   return 1;
+// }
+//
+// parameters:
+//  r0: pkfn
+//  r1: gtid
+//  r2: tid
+//  r3: argc
+//  r4(stack):  p_argv
+//  r5(stack):  &exit_frame
+//
+// locals:
+//  __gtid: gtid parm pushed on stack so can pass &gtid to pkfn
+//  __tid:  tid parm pushed on stack so can pass &tid to pkfn
+//
+// reg temps:
+//   r4:  used to hold pkfn address
+//   r5:  used as temporary for number of pkfn parms
+//   r6:  used to traverse p_argv array
+//   r7:  frame pointer (in some configurations)
+//   r8:  used as temporary for stack placement calculation
+//    and as pointer to base of callee saved area
+//   r9:  used as temporary for stack parameters
+//  r10:  used to preserve exit_frame_ptr, callee-save
+//  r11:  frame pointer (in some configurations)
+//
+// return:  r0  (always 1/TRUE)
+//
+
+__gtid = 4
+__tid = 8
+
+// -- Begin __kmp_invoke_microtask
+// mark_begin;
+  .text
+  PROC __kmp_invoke_microtask
+
+  // Pushing one extra register (r3) to keep the stack aligned
+  // for when we call pkfn below
+  push  {r3-r11,lr}
+  // Load p_argv and &exit_frame
+  ldrd  r4, r5, [sp, #10*4]
+
+# if KMP_OS_DARWIN || (defined(__thumb__) && !KMP_OS_WINDOWS)
+# define FP r7
+# define FPOFF 4*4
+#else
+# define FP r11
+# define FPOFF 8*4
+#endif
+  add FP, sp, #FPOFF
+# if OMPT_SUPPORT
+  mov r10, r5
+  str FP, [r10]
+# endif
+  mov r8, sp
+
+  // Calculate how much stack to allocate, in increments of 8 bytes.
+  // We strictly need 4*(argc-2) bytes, but allocate 4*argc for
+  // simplicity (to avoid needing to handle the argc<2 cases).
+  // We align the number of bytes allocated to 8 bytes, to keep the
+  // stack aligned. (Since we already allocate more than enough, it's ok
+  // to round down instead of up for the alignment.) We allocate another
+  // extra 8 bytes for gtid and tid.
+  mov r5, #1
+  add r5, r5, r3, lsr #1
+  sub sp, sp, r5, lsl #3
+
+  str r1, [r8, #-__gtid]
+  str r2, [r8, #-__tid]
+  mov r5, r3
+  mov r6, r4
+  mov r4, r0
+
+  // Prepare the first 2 parameters to pkfn - pointers to gtid and tid
+  // in our stack frame.
+  sub r0, r8, #__gtid
+  sub r1, r8, #__tid
+
+  mov r8, sp
+
+  // Load p_argv[0] and p_argv[1] into r2 and r3, if argc >= 1/2
+  cmp r5, #0
+  beq KMP_LABEL(kmp_1)
+  ldr r2, [r6]
+
+  subs  r5, r5, #1
+  beq KMP_LABEL(kmp_1)
+  ldr r3, [r6, #4]!
+
+  // Loop, loading the rest of p_argv and writing the elements on the
+  // stack.
+KMP_LABEL(kmp_0):
+  subs  r5, r5, #1
+  beq KMP_LABEL(kmp_1)
+  ldr r12, [r6, #4]!
+  str r12, [r8], #4
+  b KMP_LABEL(kmp_0)
+KMP_LABEL(kmp_1):
+  blx r4
+  mov r0, #1
+
+  sub r4, FP, #FPOFF
+  mov sp, r4
+# undef FP
+# undef FPOFF
+
+# if OMPT_SUPPORT
+  mov r1, #0
+  str r1, [r10]
+# endif
+  pop {r3-r11,pc}
+
+  DEBUG_INFO __kmp_invoke_microtask
+// -- End  __kmp_invoke_microtask
+
+#endif /* (KMP_OS_LINUX || KMP_OS_DARWIN || KMP_OS_WINDOWS) && KMP_ARCH_AARCH64 */
+
 #if KMP_ARCH_PPC64
 
 //------------------------------------------------------------------------
 // int
 // __kmp_invoke_microtask( void (*pkfn) (int gtid, int tid, ...),
 //                         int gtid, int tid,
 //                         int argc, void *p_argv[]
 // #if OMPT_SUPPORT

 #if KMP_ARCH_ARM || KMP_ARCH_MIPS
     .data
     COMMON .gomp_critical_user_, 32, 3
     .data
     .align 4
     .global __kmp_unnamed_critical_addr
 __kmp_unnamed_critical_addr:
     .4byte .gomp_critical_user_
+#ifdef __ELF__
     .size __kmp_unnamed_critical_addr,4
+#endif
 #endif /* KMP_ARCH_ARM */
 
 #if KMP_ARCH_PPC64 || KMP_ARCH_AARCH64 || KMP_ARCH_MIPS64 || KMP_ARCH_RISCV64 || KMP_ARCH_LOONGARCH64
 #ifndef KMP_PREFIX_UNDERSCORE
 # define KMP_PREFIX_UNDERSCORE(x) x
 #endif
     .data
     COMMON .gomp_critical_user_, 32, 3

openmp/runtime/src/z_Linux_util.cpp

 } // __kmp_get_load_balance
 
 #endif // KMP_OS_DARWIN
 
 #endif // USE_LOAD_BALANCE
 
 #if !(KMP_ARCH_X86 || KMP_ARCH_X86_64 || KMP_MIC ||                            \
       ((KMP_OS_LINUX || KMP_OS_DARWIN) && KMP_ARCH_AARCH64) ||                 \
-      KMP_ARCH_PPC64 || KMP_ARCH_RISCV64 || KMP_ARCH_LOONGARCH64)
+      KMP_ARCH_PPC64 || KMP_ARCH_RISCV64 || KMP_ARCH_LOONGARCH64 ||            \
+      KMP_ARCH_ARM)
 
 // we really only need the case with 1 argument, because CLANG always build
 // a struct of pointers to shared variables referenced in the outlined function
 int __kmp_invoke_microtask(microtask_t pkfn, int gtid, int tid, int argc,
                            void *p_argv[]
 #if OMPT_SUPPORT
                            ,
                            void **exit_frame_ptr

openmp/runtime/src/z_Windows_NT-586_util.cpp

 #if OMPT_SUPPORT
   *exit_frame_ptr = 0;
 #endif
 
   return 1;
 }
 #endif
 
-#if KMP_ARCH_ARM
-// This matches the generic fallback implementation of __kmp_invoke_microtask
-// from z_Linux_util.cpp, which is used on Linux on ARM.
-int __kmp_invoke_microtask(microtask_t pkfn, int gtid, int tid, int argc,
-                           void *p_argv[]
-#if OMPT_SUPPORT
-                           ,
-                           void **exit_frame_ptr
-#endif
-) {
-#if OMPT_SUPPORT
-  *exit_frame_ptr = OMPT_GET_FRAME_ADDRESS(0);
-#endif
-
-  switch (argc) {
-  default:
-    fprintf(stderr, "Too many args to microtask: %d!\n", argc);
-    fflush(stderr);
-    exit(-1);
-  case 0:
-    (*pkfn)(&gtid, &tid);
-    break;
-  case 1:
-    (*pkfn)(&gtid, &tid, p_argv[0]);
-    break;
-  case 2:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1]);
-    break;
-  case 3:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2]);
-    break;
-  case 4:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3]);
-    break;
-  case 5:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4]);
-    break;
-  case 6:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4],
-            p_argv[5]);
-    break;
-  case 7:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4],
-            p_argv[5], p_argv[6]);
-    break;
-  case 8:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4],
-            p_argv[5], p_argv[6], p_argv[7]);
-    break;
-  case 9:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4],
-            p_argv[5], p_argv[6], p_argv[7], p_argv[8]);
-    break;
-  case 10:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4],
-            p_argv[5], p_argv[6], p_argv[7], p_argv[8], p_argv[9]);
-    break;
-  case 11:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4],
-            p_argv[5], p_argv[6], p_argv[7], p_argv[8], p_argv[9], p_argv[10]);
-    break;
-  case 12:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4],
-            p_argv[5], p_argv[6], p_argv[7], p_argv[8], p_argv[9], p_argv[10],
-            p_argv[11]);
-    break;
-  case 13:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4],
-            p_argv[5], p_argv[6], p_argv[7], p_argv[8], p_argv[9], p_argv[10],
-            p_argv[11], p_argv[12]);
-    break;
-  case 14:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4],
-            p_argv[5], p_argv[6], p_argv[7], p_argv[8], p_argv[9], p_argv[10],
-            p_argv[11], p_argv[12], p_argv[13]);
-    break;
-  case 15:
-    (*pkfn)(&gtid, &tid, p_argv[0], p_argv[1], p_argv[2], p_argv[3], p_argv[4],
-            p_argv[5], p_argv[6], p_argv[7], p_argv[8], p_argv[9], p_argv[10],
-            p_argv[11], p_argv[12], p_argv[13], p_argv[14]);
-    break;
-  }
-
-#if OMPT_SUPPORT
-  *exit_frame_ptr = 0;
-#endif
-
-  return 1;
-}
-#endif
-
 #endif /* KMP_ARCH_X86 || KMP_ARCH_X86_64 || KMP_ARCH_AARCH64 || KMP_ARCH_ARM */