librt base64: use existing SIMD CPU dispatch by customizing build flags

Inspired by https://stackoverflow.com/a/68508804

Author: mr-c

mypyc/build.py

@@ -70,6 +70,10 @@ class ModDesc(NamedTuple):
             "base64/arch/neon64/codec.c",
         ],
         [
+            "base64/arch/avx/enc_loop_asm.c",
+            "base64/arch/avx2/enc_loop_asm.c",
+            "base64/arch/avx2/dec_loop.c",
+            "base64/arch/avx2/dec_reshuffle.c",
             "base64/arch/generic/32/enc_loop.c",
             "base64/arch/generic/64/enc_loop.c",
             "base64/arch/generic/32/dec_loop.c",
@@ -118,6 +122,51 @@ class ModDesc(NamedTuple):
 else:
     from distutils import ccompiler, sysconfig
 
+EXTRA_FLAGS_PER_COMPILER_TYPE_PER_PATH_COMPONENT = {
+        "unix": {
+            "base64/arch/ssse3": "-mssse3",
+            "base64/arch/sse41": "-msse4.1",
+            "base64/arch/sse42": "-msse4.2",
+            "base64/arch/avx2": "-mavx2",
+            "base64/arch/avx": "-mavx",
+            },
+        "msvc": {
+            "base64/arch/sse42": "/arch:SSE4.2",
+            "base64/arch/avx2": "/arch:AVX2",
+            "base64/arch/avx": "/arch:AVX",
+            }
+        }
+
+
+def spawn(self: ccompiler.CCompiler, cmd: Iterable[str], **kwargs: Any) -> None:
+    compiler_type: str = self.compiler_type  # type: ignore[attr-defined]
+    extra_options = EXTRA_FLAGS_PER_COMPILER_TYPE_PER_PATH_COMPONENT[compiler_type]
+    new_cmd = list(cmd)
+    if extra_options is not None:
+        # filenames are closer to the end of command line
+        for argument in reversed(new_cmd):
+            # Check if argument contains a filename. We must check for all
+            # possible extensions; checking for target extension is faster.
+            if self.obj_extension and not str(argument).endswith(self.obj_extension):  # type: ignore[attr-defined]
+                continue
+
+            for path in extra_options.keys():
+                if path in str(argument):
+                    if compiler_type == 'bcpp':
+                        # Borland accepts a source file name at the end,
+                        # insert the options before it
+                        new_cmd[-1:-1] = extra_options[path]
+                    else:
+                        new_cmd.append(extra_options[path])
+
+                    # path component is found, no need to search any further
+                    break
+    self.__spawn(new_cmd, **kwargs)  # type: ignore[attr-defined]
+
+
+ccompiler.CCompiler.__spawn = ccompiler.CCompiler.spawn  # type: ignore[attr-defined]
+ccompiler.CCompiler.spawn = spawn  # type: ignore[method-assign]
+
 
 def get_extension() -> type[Extension]:
     # We can work with either setuptools or distutils, and pick setuptools
@@ -661,9 +710,6 @@ def mypycify(
             # See https://github.com/mypyc/mypyc/issues/956
             "-Wno-cpp",
         ]
-        if X86_64:
-            # Enable SIMD extensions. All CPUs released since ~2010 support SSE4.2.
-            cflags.append("-msse4.2")
         if log_trace:
             cflags.append("-DMYPYC_LOG_TRACE")
         if experimental_features:
@@ -692,10 +738,6 @@ def mypycify(
             # that we actually get the compilation speed and memory
             # use wins that multi-file mode is intended for.
             cflags += ["/GL-", "/wd9025"]  # warning about overriding /GL
-        if X86_64:
-            # Enable SIMD extensions. All CPUs released since ~2010 support SSE4.2.
-            # Also Windows 11 requires SSE4.2 since 24H2.
-            cflags.append("/arch:SSE4.2")
         if log_trace:
             cflags.append("/DMYPYC_LOG_TRACE")
         if experimental_features:

mypyc/lib-rt/base64/arch/avx/codec.c

@@ -24,7 +24,7 @@
 #include "../ssse3/dec_loop.c"
 
 #if BASE64_AVX_USE_ASM
-# include "enc_loop_asm.c"
+# include "./enc_loop_asm.c"
 #else
 # include "../ssse3/enc_translate.c"
 # include "../ssse3/enc_reshuffle.c"

mypyc/lib-rt/base64/arch/avx2/codec.c

@@ -20,15 +20,15 @@
 # endif
 #endif
 
-#include "dec_reshuffle.c"
-#include "dec_loop.c"
+#include "./dec_reshuffle.c"
+#include "./dec_loop.c"
 
 #if BASE64_AVX2_USE_ASM
-# include "enc_loop_asm.c"
+# include "./enc_loop_asm.c"
 #else
-# include "enc_translate.c"
-# include "enc_reshuffle.c"
-# include "enc_loop.c"
+# include "./enc_translate.c"
+# include "./enc_reshuffle.c"
+# include "./enc_loop.c"
 #endif
 
 #endif	// HAVE_AVX2

mypyc/lib-rt/base64/config.h

@@ -1,24 +1,19 @@
 #ifndef BASE64_CONFIG_H
 #define BASE64_CONFIG_H
 
-#define BASE64_WITH_SSSE3 0
+#define BASE64_WITH_SSSE3 1
 #define HAVE_SSSE3 BASE64_WITH_SSSE3
 
-#define BASE64_WITH_SSE41 0
+#define BASE64_WITH_SSE41 1
 #define HAVE_SSE41 BASE64_WITH_SSE41
 
-#if defined(__x86_64__) || defined(_M_X64)
 #define BASE64_WITH_SSE42 1
-#else
-#define BASE64_WITH_SSE42 0
-#endif
-
 #define HAVE_SSE42 BASE64_WITH_SSE42
 
-#define BASE64_WITH_AVX 0
+#define BASE64_WITH_AVX 1
 #define HAVE_AVX BASE64_WITH_AVX
 
-#define BASE64_WITH_AVX2 0
+#define BASE64_WITH_AVX2 1
 #define HAVE_AVX2 BASE64_WITH_AVX2
 
 #define BASE64_WITH_AVX512 0

mypyc/lib-rt/setup.py

@@ -9,6 +9,7 @@
 import platform
 import subprocess
 import sys
+from collections.abc import Iterable
 from distutils import ccompiler, sysconfig
 from typing import Any
 
@@ -27,6 +28,51 @@
 
 X86_64 = platform.machine() in ("x86_64", "AMD64", "amd64")
 
+EXTRA_FLAGS_PER_COMPILER_TYPE_PER_PATH_COMPONENT = {
+        "unix": {
+            "base64/arch/ssse3": "-mssse3",
+            "base64/arch/sse41": "-msse4.1",
+            "base64/arch/sse42": "-msse4.2",
+            "base64/arch/avx2": "-mavx2",
+            "base64/arch/avx": "-mavx",
+            },
+        "msvc": {
+            "base64/arch/sse42": "/arch:SSE4.2",
+            "base64/arch/avx2": "/arch:AVX2",
+            "base64/arch/avx": "/arch:AVX",
+            }
+        }
+
+
+def spawn(self: ccompiler.CCompiler, cmd: Iterable[str], **kwargs: Any) -> None:
+    compiler_type: str = self.compiler_type  # type: ignore[attr-defined]
+    extra_options = EXTRA_FLAGS_PER_COMPILER_TYPE_PER_PATH_COMPONENT[compiler_type]
+    new_cmd = list(cmd)
+    if extra_options is not None:
+        # filenames are closer to the end of command line
+        for argument in reversed(new_cmd):
+            # Check if argument contains a filename. We must check for all
+            # possible extensions; checking for target extension is faster.
+            if self.obj_extension and not str(argument).endswith(self.obj_extension):  # type: ignore[attr-defined]
+                continue
+
+            for path in extra_options.keys():
+                if path in str(argument):
+                    if compiler_type == 'bcpp':
+                        # Borland accepts a source file name at the end,
+                        # insert the options before it
+                        new_cmd[-1:-1] = extra_options[path]
+                    else:
+                        new_cmd.append(extra_options[path])
+
+                    # path component is found, no need to search any further
+                    break
+    self.__spawn(new_cmd, **kwargs)  # type: ignore[attr-defined]
+
+
+ccompiler.CCompiler.__spawn = ccompiler.CCompiler.spawn  # type: ignore[attr-defined]
+ccompiler.CCompiler.spawn = spawn  # type: ignore[method-assign]
+
 
 class BuildExtGtest(build_ext):
     def get_library_names(self) -> list[str]:
@@ -82,12 +128,8 @@ def run(self) -> None:
     cflags: list[str] = []
     if compiler.compiler_type == "unix":
         cflags += ["-O3"]
-        if X86_64:
-            cflags.append("-msse4.2")  # Enable SIMD (see also mypyc/build.py)
     elif compiler.compiler_type == "msvc":
         cflags += ["/O2"]
-        if X86_64:
-            cflags.append("/arch:SSE4.2")  # Enable SIMD (see also mypyc/build.py)
 
     setup(
         ext_modules=[