Merge pull request #13204 from vogma/aarch64_build_update

aarch64 op: Enable two‑stage SVE detection in component configuration

Author: bosilca

ompi/mca/op/aarch64/configure.m4

@@ -13,7 +13,7 @@
 #
 
 # MCA_ompi_op_arm_CONFIG([action-if-can-compile],
-#		         [action-if-cant-compile])
+#                        [action-if-cant-compile])
 # ------------------------------------------------
 AC_DEFUN([MCA_ompi_op_aarch64_CONFIG],[
     AC_CONFIG_FILES([ompi/mca/op/aarch64/Makefile])
@@ -74,33 +74,70 @@ AC_DEFUN([MCA_ompi_op_aarch64_CONFIG],[
            #
            # Check for SVE support
            #
-           AC_CACHE_CHECK([for SVE support], op_cv_sve_support,
-                 [AS_IF([test "$op_cv_neon_support" = "yes"],
-                        [
-                          AC_LINK_IFELSE(
-                              [AC_LANG_PROGRAM([[
+          AC_CACHE_CHECK([for SVE support], [op_cv_sve_support], [
+              AC_MSG_RESULT([])      
+              # initialize result variables
+              op_cv_sve_support=no
+              op_cv_sve_add_flags=no
+
+              # first attempt: no extra flags
+              AC_MSG_CHECKING([for SVE support (no additional flags)])
+              AC_LINK_IFELSE(
+                  [AC_LANG_SOURCE([[
 #if defined(__aarch64__) && defined(__ARM_FEATURE_SVE)
 #include <arm_sve.h>
 #else
 #error "No support for __aarch64__ or SVE"
 #endif
-                                             ]],
-                                             [[
-#if defined(__aarch64__) && defined(_ARM_FEATURE_SVE)
-    svfloat32_t vA;
-    vA = svdup_n_f32(0)
+
+int main(void) {
+  svfloat32_t vA;
+  vA = svdup_n_f32(0);
+  return 0;
+}
+               ]])],
+               [ op_cv_sve_support=yes
+                 AC_MSG_RESULT([yes]) ],
+               [ AC_MSG_RESULT([no ]) ]
+             )
+
+            # second attempt: use +sve attribute
+            AS_IF([test "$op_cv_sve_support" = "no"],[
+                AC_MSG_CHECKING([for SVE support (with +sve)])
+                AC_LINK_IFELSE(
+                    [AC_LANG_SOURCE([[
+#if defined(__aarch64__) && defined(__linux__)
+  #include <arm_sve.h>
+#else
+  #error "this feature is only supported on aarch64 + linux platforms"
 #endif
-                                             ]])],
-                      [op_cv_sve_support=yes],
-                      [op_cv_sve_support=no])])])
-          ])
 
+__attribute__((__target__("+sve")))
+int main(void) {
+  svbool_t    pg = svptrue_b32();
+  svuint32_t  a  = svdup_u32(0);
+  svuint32_t  b  = svdup_u32(0);
+  svuint32_t  c  = svadd_u32_m(pg, a, b);
+  return (int)svaddv_u32(pg, c);
+}
+                 ]])],
+                 [ op_cv_sve_support=yes
+                   op_cv_sve_add_flags=yes
+                   AC_MSG_RESULT([yes]) ],
+                 [ AC_MSG_RESULT([no ]) ]
+               )
+             ])
+           ])
+
+           AC_LANG_POP
+])
     AM_CONDITIONAL([MCA_BUILD_ompi_op_has_neon_support],
                    [test "$op_cv_neon_support" = "yes"])
     AM_CONDITIONAL([MCA_BUILD_ompi_op_has_neon_fp_support],
                    [test "$op_cv_neon_fp_support" = "yes"])
     AM_CONDITIONAL([MCA_BUILD_ompi_op_has_sve_support],
                    [test "$op_cv_sve_support" = "yes"])
+
     AC_SUBST(MCA_BUILD_ompi_op_has_neon_support)
     AC_SUBST(MCA_BUILD_ompi_op_has_neon_fp_support)
     AC_SUBST(MCA_BUILD_ompi_op_has_sve_support)
@@ -111,9 +148,12 @@ AC_DEFUN([MCA_ompi_op_aarch64_CONFIG],[
           [AC_DEFINE([OMPI_MCA_OP_HAVE_NEON_FP], [1],[NEON FP supported in the current build])])
     AS_IF([test "$op_cv_sve_support" = "yes"],
           [AC_DEFINE([OMPI_MCA_OP_HAVE_SVE], [1],[SVE supported in the current build])])
+    AS_IF([test "$op_cv_sve_add_flags" = "yes"],
+          [AC_DEFINE([OMPI_MCA_OP_SVE_EXTRA_FLAGS], [1],[SVE supported with additional compile attributes])],
+          [AC_DEFINE([OMPI_MCA_OP_SVE_EXTRA_FLAGS], [0],[SVE not supported])])
 
-    # If we have at least support for Neon
-    AS_IF([test "$op_cv_neon_support" = "yes"],
+    # If we have at least support for Neon or SVE
+    AS_IF([test "$op_cv_neon_support" = "yes" || test "$op_cv_sve_support" = "yes" ],
           [$1],
           [$2])
 ])dnl

ompi/mca/op/aarch64/op_aarch64.h

@@ -24,6 +24,12 @@
 
 BEGIN_C_DECLS
 
+#if OMPI_MCA_OP_SVE_EXTRA_FLAGS
+#define OMPI_SVE_ATTR __attribute__ ((__target__ ("+sve")))
+#else
+#define OMPI_SVE_ATTR
+#endif
+
 /**
  * Derive a struct from the base op component struct, allowing us to
  * cache some component-specific information on our well-known

ompi/mca/op/aarch64/op_aarch64_component.c

@@ -101,7 +101,7 @@ static int mca_op_aarch64_component_close(void)
 /*
  * Register MCA params.
  */
-static int mca_op_aarch64_component_register(void)
+OMPI_SVE_ATTR static int mca_op_aarch64_component_register(void)
 {
 
     mca_op_aarch64_component.hardware_available = 1;  /* Check for Neon */

ompi/mca/op/aarch64/op_aarch64_functions.c

@@ -137,6 +137,7 @@ _Generic((*(out)), \
     }
 #elif defined(GENERATE_SVE_CODE)
 #define OP_AARCH64_FUNC(name, type_name, type_size, type_cnt, type, op)           \
+    OMPI_SVE_ATTR                                                                 \
     static void OP_CONCAT(ompi_op_aarch64_2buff_##name##_##type##type_size##_t, APPEND) \
                             (const void *_in, void *_out, int *count,             \
                              struct ompi_datatype_t **dtype,                      \
@@ -303,6 +304,7 @@ static void OP_CONCAT(ompi_op_aarch64_3buff_##name##_##type##type_size##_t, APPE
 }
 #elif defined(GENERATE_SVE_CODE)
 #define OP_AARCH64_FUNC_3BUFF(name, type_name, type_size, type_cnt, type, op)             \
+OMPI_SVE_ATTR                                                                             \
 static void OP_CONCAT(ompi_op_aarch64_3buff_##name##_##type##type_size##_t, APPEND)       \
                              (const void *_in1, const void *_in2, void *_out, int *count, \
                               struct ompi_datatype_t **dtype,                             \