diff --git a/src/numpy-stubs/_core/_simd.pyi b/src/numpy-stubs/_core/_simd.pyi index 87b2aaa5..efa905b2 100644 --- a/src/numpy-stubs/_core/_simd.pyi +++ b/src/numpy-stubs/_core/_simd.pyi @@ -1,28 +1,43 @@ -import sys from types import ModuleType -from typing import TypedDict, type_check_only +from typing import Required, TypedDict, type_check_only -# NOTE: these 6 are only defined on systems with an intel processor +# NOTE: these are only defined on (some) systems with an intel processor SSE42: ModuleType | None = ... FMA3: ModuleType | None = ... AVX2: ModuleType | None = ... AVX512F: ModuleType | None = ... AVX512_SKX: ModuleType | None = ... -if sys.platform == "linux": - X86_V4: ModuleType | None = ... +X86_V3: ModuleType | None = ... +X86_V4: ModuleType | None = ... + +# NOTE: these 2 are only defined on systems with an arm processor +ASIMD: ModuleType | None = ... +NEON: ModuleType | None = ... + +# NOTE: This is only defined on systems with an riscv64 processor. +RVV: ModuleType | None = ... baseline: ModuleType | None = ... @type_check_only -class SimdTargets(TypedDict): +class SimdTargets(TypedDict, total=False): + # x86 SSE42: ModuleType | None AVX2: ModuleType | None FMA3: ModuleType | None AVX512F: ModuleType | None AVX512_SKX: ModuleType | None + X86_V3: ModuleType | None X86_V4: ModuleType | None - baseline: ModuleType | None + # ARM + ASIMD: ModuleType | None + NEON: ModuleType | None + + # RISC-V + RVV: ModuleType | None + + baseline: Required[ModuleType | None] targets: SimdTargets = ... diff --git a/tool/allowlists/common.txt b/tool/allowlists/common.txt index 7e077d5e..f1a35463 100644 --- a/tool/allowlists/common.txt +++ b/tool/allowlists/common.txt @@ -13,6 +13,9 @@ numpy(\.\w+)?\.tests.* numpy(\..+)?\.float(96|128) numpy(\..+)?\.complex(192|256) +# system-dependent private SIMD submodules +numpy\._core\._simd\.(SSE42|FMA3|AVX(2|512F|512_SKX)|X86_V(3|4)|ASIMD|NEON|RVV) + # these are always either missing float96/complex192 or float128/complex256 numpy\.__all__ numpy\._?core\.__all__ diff --git a/tool/allowlists/simd.txt b/tool/allowlists/simd.txt deleted file mode 100644 index 6f65e6fa..00000000 --- a/tool/allowlists/simd.txt +++ /dev/null @@ -1,7 +0,0 @@ -# only defined on systems with an intel processor -numpy\._core\._simd\.SSE42 -numpy\._core\._simd\.FMA3 -numpy\._core\._simd\.AVX2 -numpy\._core\._simd\.AVX512F -numpy\._core\._simd\.AVX512_SKX -# numpy\._core\._simd\.X86_V4 diff --git a/tool/stubtest.py b/tool/stubtest.py index c4eacab9..48a8b238 100644 --- a/tool/stubtest.py +++ b/tool/stubtest.py @@ -46,19 +46,6 @@ ALLOWLISTS = ["common.txt", f"py3{sys.version_info.minor}.txt"] -def __check_simd() -> None: - try: - from numpy._core import _simd # noqa: PLC0415, PLC2701 - except ImportError: - return - - if not hasattr(_simd, "AVX512F"): - ALLOWLISTS.append("simd.txt") - - -__check_simd() - - def _allowlists() -> list[str]: relpath = (TOOL_DIR / "allowlists").relative_to(CWD) return [str(relpath / fname) for fname in ALLOWLISTS]