[openmp] Use z_Linux_asm.S to provide __kmp_invoke_microtask with Clang for Windows/aarch64

[mstorsjo]

Summary

When building for Windows aarch64, and not using the actual MSVC,
we can assemble gnu assembly files just fine, and the existing
correct implementation of __kmp_invoke_microtask is fully usable.

(This current patch only uses it when building in mingw mode, but
when building with clang-cl, we could also just as well build and
rely on the assembly file - but that would require a bit more
preprocessor and cmake checks to distinguish that case.)

The C implementation of __kmp_invoke_microtask in
z_Windows_NT-586_util.cpp relies on unguaranteed assumptions about
the compiler behaviour - it might work on MSVC, but doesn't necessarily
on other compilers. That function uses an alloca to pass parameters
on the stack to the called functions.

However, there's no guarantee that the buffer allocated by alloca is
exactly at the bottom of the stack when doing the call; the compiler
might have left space for extra things to save on the stack there.

Additionally, when compiled with Clang with optimization, Clang
optimizes out the alloca and memcpy entirely. On the C language
level, they don't have any visible effect outside of the function
and thus can be omitted entirely.

[mstorsjo]

openmp/runtime/src/CMakeLists.txt

98

I guess we could make this AND (NOT MSVC OR CMAKE_C_COMPILER_ID STREQUAL "Clang") - so we'd use the assembly form with clang in MSVC mode too. Clang's optimizations break the current C implementation, but Clang should be able to handle the gas assembly even in msvc mode.

Then we'd make the condition in the source file KMP_ARCH_AARCH64 && defined(_MSC_VER) && !defined(__clang__).

[mstorsjo]

Use the assembly implementation for Clang-cl builds for arm64 too (as the C implementation gets misoptimized by Clang but not by MSVC).

(Building openmp with clang-cl for arm64 requires one other unrelated patch, which I can send when the others I have outstanding have settled.)

[mstorsjo]

Ping

[DavidSpickett]

I have zero idea about OpenMP, just one comment on the cmake.

openmp/runtime/src/CMakeLists.txt

104

The implicit else here is else use the .cpp file ? This would happen when you are on AArch64 and using MSVC.

Worth adding that as a comment?

[mstorsjo]

In D137827#3949038, @DavidSpickett wrote:

I have zero idea about OpenMP, just one comment on the cmake.

Yeah, it seems hard to get someone to approve this... I presume the asm changes look straightforward/ok to you too? I'm wrapping the .comm directive since the alignment in that directive is interpreted differently between ELF and COFF.

openmp/runtime/src/CMakeLists.txt

104

Sure, it never hurts to be more explicit with things like this.

[mstorsjo]

Clarified the aarch64+msvc case in the cmake file.

[DavidSpickett]

The C implementation of __kmp_invoke_microtask in z_Windows_NT-586_util.cpp relies on unguaranteed assumptions about the compiler behaviour

Meaning no reason to use the C version, that part is clear to me.

I'm wrapping the .comm directive since the alignment in that directive is interpreted differently between ELF and COFF.

I didn't know this was a thing. Can you cite some documentation for the .comm behaviour?

I assume this passes whatever OpenMP testing we have. I'd be ok approving it on that basis.

openmp/runtime/src/z_Linux_asm.S

153

I guess this is 2 byte function start alignment? I'd have guessed 4 so maybe I'm wrong.

[mstorsjo]

In D137827#3949184, @DavidSpickett wrote:

I'm wrapping the .comm directive since the alignment in that directive is interpreted differently between ELF and COFF.

I didn't know this was a thing. Can you cite some documentation for the .comm behaviour?

https://sourceware.org/binutils/docs/as/Comm.html - which says:

When using ELF or (as a GNU extension) PE, the .comm directive takes an optional third argument. This is the desired alignment of the symbol, specified for ELF as a byte boundary (for example, an alignment of 16 means that the least significant 4 bits of the address should be zero), and for PE as a power of two (for example, an alignment of 5 means aligned to a 32-byte boundary).

I assume this passes whatever OpenMP testing we have. I'd be ok approving it on that basis.

Yes - this, together with D137748 and D137772, makes all tests pass in the mingw/aarch64 build configuration.

openmp/runtime/src/z_Linux_asm.S

153

No, the ALIGN macro takes a power of two - so this sets 4 byte alignment.

(The ALIGN macro is probably used because the .align directive is inconsistent across platforms - that's why there are the .p2align and .balign directives too. But changing the existing code to use them instead would be a different change.)

[DavidSpickett]

LGTM

openmp/runtime/src/z_Linux_asm.S

153

Of course yes. Every time I have to read these things I get number of bytes mixed up with power of 2.

openmp/runtime/src/CMakeLists.txt

include_directories(
     ${LLVM_MAIN_INCLUDE_DIR}
     ${LLVM_INCLUDE_DIR}
   )
 endif()
 
 # Getting correct source files to build library
 set(LIBOMP_CXXFILES)
 set(LIBOMP_ASMFILES)
+set(LIBOMP_GNUASMFILES)
 if(STUBS_LIBRARY)
   set(LIBOMP_CXXFILES kmp_stub.cpp)
 else()
   # Get C++ files
   set(LIBOMP_CXXFILES
     kmp_alloc.cpp
     kmp_atomic.cpp
     kmp_csupport.cpp
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.
+      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)
     libomp_append(LIBOMP_CXXFILES kmp_gsupport.cpp)
-    libomp_append(LIBOMP_ASMFILES z_Linux_asm.S) # Unix assembly file
+    libomp_append(LIBOMP_GNUASMFILES z_Linux_asm.S) # Unix assembly file
   endif()
   libomp_append(LIBOMP_CXXFILES thirdparty/ittnotify/ittnotify_static.cpp LIBOMP_USE_ITT_NOTIFY)
   libomp_append(LIBOMP_CXXFILES kmp_debugger.cpp LIBOMP_USE_DEBUGGER)
   libomp_append(LIBOMP_CXXFILES kmp_stats.cpp LIBOMP_STATS)
   libomp_append(LIBOMP_CXXFILES kmp_stats_timing.cpp LIBOMP_STATS)
   libomp_append(LIBOMP_CXXFILES kmp_taskdeps.cpp)
   libomp_append(LIBOMP_CXXFILES kmp_cancel.cpp)
 endif()
 # Files common to stubs and normal library
 libomp_append(LIBOMP_CXXFILES kmp_ftn_cdecl.cpp)
 libomp_append(LIBOMP_CXXFILES kmp_ftn_extra.cpp)
 libomp_append(LIBOMP_CXXFILES kmp_version.cpp)
 libomp_append(LIBOMP_CXXFILES ompt-general.cpp IF_TRUE LIBOMP_OMPT_SUPPORT)
 libomp_append(LIBOMP_CXXFILES ompd-specific.cpp IF_TRUE LIBOMP_OMPD_SUPPORT)
 
-set(LIBOMP_SOURCE_FILES ${LIBOMP_CXXFILES} ${LIBOMP_ASMFILES})
+set(LIBOMP_SOURCE_FILES ${LIBOMP_CXXFILES} ${LIBOMP_ASMFILES} ${LIBOMP_GNUASMFILES})
 # For Windows, there is a resource file (.rc -> .res) that is also compiled
 libomp_append(LIBOMP_SOURCE_FILES libomp.rc WIN32)
 
 # Get compiler and assembler flags
 libomp_get_cxxflags(LIBOMP_CONFIGURED_CXXFLAGS)
 libomp_get_asmflags(LIBOMP_CONFIGURED_ASMFLAGS)
 # Set the compiler flags for each type of source
 set_source_files_properties(${LIBOMP_CXXFILES} PROPERTIES COMPILE_FLAGS "${LIBOMP_CONFIGURED_CXXFLAGS}")
-set_source_files_properties(${LIBOMP_ASMFILES} PROPERTIES COMPILE_FLAGS "${LIBOMP_CONFIGURED_ASMFLAGS}")
-# Let the compiler handle the assembly files on Unix-like systems
-if(NOT WIN32)
-  set_source_files_properties(${LIBOMP_ASMFILES} PROPERTIES LANGUAGE C)
-endif()
+set_source_files_properties(${LIBOMP_ASMFILES} ${LIBOMP_GNUASMFILES} PROPERTIES COMPILE_FLAGS "${LIBOMP_CONFIGURED_ASMFLAGS}")
+# Let the compiler handle the GNU assembly files
+set_source_files_properties(${LIBOMP_GNUASMFILES} PROPERTIES LANGUAGE C)
 
 # Remove any cmake-automatic linking of the standard C++ library.
 # We neither need (nor want) the standard C++ library dependency even though we compile c++ files.
 if(NOT ${LIBOMP_USE_STDCPPLIB})
   set(LIBOMP_LINKER_LANGUAGE C)
   set(CMAKE_CXX_IMPLICIT_LINK_LIBRARIES)
 else()
   set(LIBOMP_LINKER_LANGUAGE CXX)

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_ARCH_AARCH64
+#if (KMP_OS_LINUX || KMP_OS_DARWIN || KMP_OS_WINDOWS) && KMP_ARCH_AARCH64
 
 # 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 COMMON name, size, align_power
+  .comm \name, \size
+.endm
 
 .macro DEBUG_INFO
 /* Not sure what .size does in icc, not sure if we need to do something
    similar for OS X*.
 */
 .endmacro
 
 .macro PROC
   ALIGN  4
   .globl KMP_PREFIX_UNDERSCORE($0)
 KMP_PREFIX_UNDERSCORE($0):
 .endmacro
-# else // KMP_OS_DARWIN
+# elif KMP_OS_WINDOWS
+#  define KMP_PREFIX_UNDERSCORE(x) x  // no extra underscore for Windows/ARM64 symbols
+// Format labels so that they don't override function names in gdb's backtraces
+#  define KMP_LABEL(x) .L_##x         // local label hidden from backtraces
+
+.macro ALIGN size
+  .align 1<<(\size)
+.endm
+.macro COMMON name, size, align_power
+  .comm \name, \size, \align_power
+.endm
+
+.macro DEBUG_INFO proc
+  ALIGN 2
+.endm
+
+.macro PROC proc
+  ALIGN 2
+  .globl KMP_PREFIX_UNDERSCORE(\proc)
+KMP_PREFIX_UNDERSCORE(\proc):
+.endm
+# else // KMP_OS_DARWIN || KMP_OS_WINDOWS
 #  define KMP_PREFIX_UNDERSCORE(x) x  // no extra underscore for Linux* OS symbols
 // Format labels so that they don't override function names in gdb's backtraces
 #  define KMP_LABEL(x) .L_##x         // local label hidden from backtraces
 
 .macro ALIGN size
   .align 1<<(\size)
 .endm
+.macro COMMON name, size, align_power
+  .comm \name, \size, (1<<(\align_power))
+.endm
 
 .macro DEBUG_INFO proc
   .cfi_endproc
 // Not sure why we need .type and .size for the functions
   ALIGN 2
   .type  \proc,@function
   .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_ARCH_AARCH64
+#endif // (KMP_OS_LINUX || KMP_OS_DARWIN || KMP_OS_WINDOWS) && KMP_ARCH_AARCH64
 
 // -----------------------------------------------------------------------
 // data
 // -----------------------------------------------------------------------
 
 #ifdef KMP_GOMP_COMPAT
 
 // Support for unnamed common blocks.
// __kmp_bsr32( int );
         ret
 
         DEBUG_INFO __kmp_bsr32
 
 // -----------------------------------------------------------------------
 #endif /* KMP_ARCH_X86_64 */
 
 // '
-#if (KMP_OS_LINUX || KMP_OS_DARWIN) && KMP_ARCH_AARCH64
+#if (KMP_OS_LINUX || KMP_OS_DARWIN || KMP_OS_WINDOWS) && KMP_ARCH_AARCH64
 
 //------------------------------------------------------------------------
 // int
 // __kmp_invoke_microtask( void (*pkfn) (int gtid, int tid, ...),
 //                         int gtid, int tid,
 //                         int argc, void *p_argv[]
 // #if OMPT_SUPPORT
 //                         ,
# if OMPT_SUPPORT
   ldp x19, x20, [sp], #16
 # endif
   ldp x29, x30, [sp], #16
   ret
 
   DEBUG_INFO __kmp_invoke_microtask
 // -- End  __kmp_invoke_microtask
 
-#endif /* (KMP_OS_LINUX || KMP_OS_DARWIN) && KMP_ARCH_AARCH64 */
+#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[]
.Lfunc_end0:
   .cfi_endproc
 
 // -- End  __kmp_invoke_microtask
 
 #endif /* KMP_ARCH_LOONGARCH64 */
 
 #if KMP_ARCH_ARM || KMP_ARCH_MIPS
     .data
-    .comm .gomp_critical_user_,32,8
+    COMMON .gomp_critical_user_, 32, 3
     .data
     .align 4
     .global __kmp_unnamed_critical_addr
 __kmp_unnamed_critical_addr:
     .4byte .gomp_critical_user_
     .size __kmp_unnamed_critical_addr,4
 #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
-    .comm .gomp_critical_user_,32,8
+    COMMON .gomp_critical_user_, 32, 3
     .data
     .align 8
     .global KMP_PREFIX_UNDERSCORE(__kmp_unnamed_critical_addr)
 KMP_PREFIX_UNDERSCORE(__kmp_unnamed_critical_addr):
     .8byte .gomp_critical_user_
 #ifdef __ELF__
     .size KMP_PREFIX_UNDERSCORE(__kmp_unnamed_critical_addr),8
 #endif

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

while (!KMP_COMPARE_AND_STORE_REL64((volatile kmp_int64 *)p, old_value,
     KMP_CPU_PAUSE();
     old_value = TCR_8(*p);
     new_value = old_value & d;
   }
 
   return old_value;
 }
 
-#if KMP_ARCH_AARCH64
+#if KMP_ARCH_AARCH64 && KMP_COMPILER_MSVC
+// For !KMP_COMPILER_MSVC, this function is provided in assembly form
+// by z_Linux_asm.S.
 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