add support for Linux Aarch64

.github/workflows/push.yml

@@ -21,7 +21,7 @@ jobs:
     - name: actions/checkout
       uses: actions/checkout@v2
     - name: actions/cache
-      uses: actions/cache@v2
+      uses: actions/cache@v4
       with:
         path: ${{ env.PIP_CACHE }}
         key: ubuntu-22.04-pip-lint-${{ hashFiles('lint-requirements.txt') }}
@@ -58,7 +58,7 @@ jobs:
       with:
         python-version: '3.11'
     - name: actions/cache pip
-      uses: actions/cache@v2
+      uses: actions/cache@v4
       with:
         path: ${{ env.PIP_CACHE }}
         key: ubuntu-22.04-pip-main-${{ hashFiles('requirements.txt', 'test-requirements.txt') }}

cpu_id.h

@@ -1,11 +1,20 @@
 #pragma once
 
-#ifdef _MSC_VER
-  #include<intrin.h>
-  #define CPUID(info, x)  __cpuidex(reinterpret_cast<int *>(info), x, 0)
+// CPUID and the AVX2 kernels only exist on x86
+#if defined(__x86_64__) || defined(__i386__) || defined(_M_X64) || defined(_M_IX86)
+  #define LAPJV_X86 1
 #else
-  #include <cpuid.h>
-  #define CPUID(info, x)  __cpuid_count(x, 0, info[0], info[1], info[2], info[3])
+  #define LAPJV_X86 0
+#endif
+
+#if LAPJV_X86
+  #ifdef _MSC_VER
+    #include<intrin.h>
+    #define CPUID(info, x)  __cpuidex(reinterpret_cast<int *>(info), x, 0)
+  #else
+    #include <cpuid.h>
+    #define CPUID(info, x)  __cpuid_count(x, 0, info[0], info[1], info[2], info[3])
+  #endif
 #endif
 
 class SIMDFlags final {
@@ -14,6 +23,7 @@ class SIMDFlags final {
   SIMDFlags(const SIMDFlags &) = delete;
   SIMDFlags &operator=(const SIMDFlags &) = delete;
 
+#if LAPJV_X86
   SIMDFlags() {
     unsigned int cpuInfo[4];
     // CPUID: https://en.wikipedia.org/wiki/CPUID
@@ -34,6 +44,9 @@ class SIMDFlags final {
     CPUID(cpuInfo, 0x80000001);
     simd_flags_ |= cpuInfo[2] & (1 << 16) ? SIMD_FMA4  : SIMD_NONE;
   }
+#else
+  SIMDFlags() = default;
+#endif
 
   inline bool hasSSE()   const { return simd_flags_ & SIMD_SSE;   }
   inline bool hasSSE2()  const { return simd_flags_ & SIMD_SSE2;  }

lap.h

@@ -3,7 +3,11 @@
 #include <limits>
 #include <memory>
 
+#include "cpu_id.h"
+
+#if LAPJV_X86
 #include <immintrin.h>
+#endif
 
 
 #ifdef __GNUC__
@@ -49,6 +53,8 @@ find_umins_regular(
 #define FLOAT_MIN_DIM 64
 #define DOUBLE_MIN_DIM 100000  // 64-bit code is actually always slower
 
+#if LAPJV_X86
+
 template <typename idx>
 always_inline std::tuple<float, float, idx, idx>
 find_umins_avx2(
@@ -207,6 +213,20 @@ find_umins_avx2(
   return std::make_tuple(umin, usubmin, j1, j2);
 }
 
+#else
+
+// No AVX2 on non-x86. The scalar kernel is the only kernel. SIMDFlags never
+// reports AVX2 there either, so this exists to keep find_umins<true> parseable.
+template <typename idx, typename cost>
+always_inline std::tuple<cost, cost, idx, idx>
+find_umins_avx2(
+    idx dim, idx i, const cost *restrict assign_cost,
+    const cost *restrict v) {
+  return find_umins_regular(dim, i, assign_cost, v);
+}
+
+#endif  // LAPJV_X86
+
 template <bool avx2, typename idx, typename cost>
 always_inline std::tuple<cost, cost, idx, idx>
 find_umins(

setup.py

@@ -3,10 +3,14 @@
 
 from setuptools import Extension, setup
 
-UNIX_CXXFLAGS = [
-    "-std=c++17",
+X86_MACHINES = ["x86_64", "amd64", "x86", "i386", "i686", "AMD64"]
+IS_X86 = platform.machine().lower() in X86_MACHINES
+
+UNIX_CXXFLAGS = ["-std=c++17", "-ftree-vectorize"]
+
+# x86-only flags for GCC, not supported by aarch64
+UNIX_X86_CXXFLAGS = [
     "-mavx2",
-    "-ftree-vectorize",
     # GCP N2
     "-march=haswell",
     "-maes",
@@ -17,9 +21,18 @@
     "--param", "l2-cache-size=33792",
 ]
 
+
+def unix_flags(*extra_x86):
+    """Portable flags, plus the Haswell-tuned ones only where they mean something."""
+    flags = [*UNIX_CXXFLAGS]
+    if IS_X86:
+        flags += [*UNIX_X86_CXXFLAGS, *extra_x86]
+    return flags
+
+
 CXX_ARGS = {
     # "Darwin": [*UNIX_CXXFLAGS],  not supported anymore due to M1, PRs welcome
-    "Linux": ["-fopenmp", *UNIX_CXXFLAGS, "-mabm"],
+    "Linux": ["-fopenmp", *unix_flags("-mabm")],
     "Windows": ["/openmp", "/std:c++latest", "/arch:AVX2"],
 }