diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 5022ee3..5148075 100644 --- a/.github/workflows/push.yml +++ b/.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') }} diff --git a/cpu_id.h b/cpu_id.h index 9c0d035..b6d64a8 100644 --- a/cpu_id.h +++ b/cpu_id.h @@ -1,11 +1,20 @@ #pragma once -#ifdef _MSC_VER - #include - #define CPUID(info, x) __cpuidex(reinterpret_cast(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 - #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 + #define CPUID(info, x) __cpuidex(reinterpret_cast(info), x, 0) + #else + #include + #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; } diff --git a/lap.h b/lap.h index e4ea0a6..aca6663 100644 --- a/lap.h +++ b/lap.h @@ -3,7 +3,11 @@ #include #include +#include "cpu_id.h" + +#if LAPJV_X86 #include +#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 always_inline std::tuple 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 parseable. +template +always_inline std::tuple +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 always_inline std::tuple find_umins( diff --git a/setup.py b/setup.py index 86e403d..b47258c 100644 --- a/setup.py +++ b/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"], }