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

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

Author: mr-c

mypy_self_check.ini

@@ -8,7 +8,7 @@ pretty = True
 always_false = MYPYC
 plugins = mypy.plugins.proper_plugin
 python_version = 3.9
-exclude = mypy/typeshed/|mypyc/test-data/|mypyc/lib-rt/
+exclude = mypy/typeshed/|mypyc/test-data/
 enable_error_code = ignore-without-code,redundant-expr
 enable_incomplete_feature = PreciseTupleTypes
 show_error_code_links = True

mypyc/build.py

@@ -36,7 +36,7 @@
 from mypy.util import write_junit_xml
 from mypyc.annotate import generate_annotated_html
 from mypyc.codegen import emitmodule
-from mypyc.common import IS_FREE_THREADED, RUNTIME_C_FILES, X86_64, shared_lib_name
+from mypyc.common import IS_FREE_THREADED, RUNTIME_C_FILES, shared_lib_name
 from mypyc.errors import Errors
 from mypyc.ir.pprint import format_modules
 from mypyc.namegen import exported_name
@@ -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,52 @@ 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",
+    },
+}
+
+__spawn = ccompiler.CCompiler.spawn
+
+
+def spawn(self, cmd, **kwargs):  # type: ignore[no-untyped-def]
+    compiler_type: str = self.compiler_type
+    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):
+                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
+    __spawn(self, new_cmd, **kwargs)
+
+
+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 +711,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 +739,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/common.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import platform
 import sys
 import sysconfig
 from typing import Any, Final
@@ -45,8 +44,6 @@
 
 IS_32_BIT_PLATFORM: Final = int(SIZEOF_SIZE_T) == 4
 
-X86_64: Final = platform.machine() in ("x86_64", "AMD64", "amd64")
-
 PLATFORM_SIZE = 4 if IS_32_BIT_PLATFORM else 8
 
 # Maximum value for a short tagged integer.

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

@@ -6,9 +6,9 @@
 from __future__ import annotations
 
 import os
-import platform
 import subprocess
 import sys
+from collections.abc import Iterable
 from distutils import ccompiler, sysconfig
 from typing import Any
 
@@ -25,7 +25,51 @@
     "pythonsupport.c",
 ]
 
-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",
+    },
+}
+
+__spawn = ccompiler.CCompiler.spawn
+
+
+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
+    __spawn(self, new_cmd, **kwargs)
+
+
+ccompiler.CCompiler.spawn = spawn  # type: ignore[method-assign]
 
 
 class BuildExtGtest(build_ext):
@@ -80,14 +124,10 @@ def run(self) -> None:
     compiler = ccompiler.new_compiler()
     sysconfig.customize_compiler(compiler)
     cflags: list[str] = []
-    if compiler.compiler_type == "unix":
+    if compiler.compiler_type == "unix":  # type: ignore[attr-defined]
         cflags += ["-O3"]
-        if X86_64:
-            cflags.append("-msse4.2")  # Enable SIMD (see also mypyc/build.py)
-    elif compiler.compiler_type == "msvc":
+    elif compiler.compiler_type == "msvc":  # type: ignore[attr-defined]
         cflags += ["/O2"]
-        if X86_64:
-            cflags.append("/arch:SSE4.2")  # Enable SIMD (see also mypyc/build.py)
 
     setup(
         ext_modules=[