mirror of
https://github.com/pytorch/pytorch.git
synced 2026-10-25 18:53:39 +08:00
Nikita Shulga
4dce2962ca
Neoverse V1 supports SVE256, but Neoverse V2 narrows it back down to SVE128 — still useful to dispatch to. This PR only wires up the dispatch + a small NEON<->SVE bridge. When compiled with `-msve-vector-bits=128`, `svfloat32_t` and `float32x4_t` have identical layout, so the two new helpers `neon_to_sve` / `sve_to_neon` compile down to zero instructions. That lets the SVE128 build call into the existing SLEEF SVE routines from inside what is otherwise the NEON `Vectorized<float>` implementation — currently used by `exp2()` (`Sleef_exp2fx_u10sve`) as a demonstration of the bridge. Note: unlike the X86 AVX hierarchy (AVX512 ⊇ AVX2 ⊇ DEFAULT), SVE dispatch here is **unversioned** — `SVE128` and `SVE256` are sibling capabilities, not ordered. The dispatcher matches with `==` rather than `>=`, so you cannot dispatch a SVE128 kernel on an SVE256 machine, nor vice versa. Each kernel must be registered for the specific vector length it targets. Because of the binary size growth, we have to pass an explicit `--stub-group-size=0x2000000` flag to work around R_AARCH64_CALL26 relocation errors. Also, workaround gcc-13.x bf16 constant-miscompile bug, when it confuses fp16 and bf16 constant which results in vdupq_n_bf16((bfloat16_t)1.0f) being miscomputed into `fmov v?.8h, 1.0e+0`, see https://godbolt.org/z/4sqdxj1GM (gcc-15 uses an interesting numeric trick here) This PR is provisional and may be reverted if the expected performance gains on SVE128 hardware (e.g. Neoverse V2) do not materialize in benchmarks. Pull Request resolved: https://github.com/pytorch/pytorch/pull/176256 Approved by: https://github.com/Skylion007, https://github.com/fadara01
40 lines
1.4 KiB
CMake
40 lines
1.4 KiB
CMake
# Check ARM feature availability for SVE compile-time support.
|
|
IF(CMAKE_SYSTEM_NAME MATCHES "Linux")
|
|
INCLUDE(CheckCXXSourceCompiles)
|
|
|
|
SET(SVE_BF16_CODE "
|
|
#include <arm_sve.h>
|
|
#include <arm_neon.h>
|
|
int main()
|
|
{
|
|
svfloat64_t a;
|
|
a = svdup_n_f64(0);
|
|
float32x4_t b = vdupq_n_f32(0);
|
|
bfloat16x8_t c = vreinterpretq_bf16_f32(b);
|
|
bfloat16x4_t d = vget_low_bf16(c);
|
|
return 0;
|
|
}
|
|
")
|
|
|
|
SET(CMAKE_REQUIRED_FLAGS_SAVE ${CMAKE_REQUIRED_FLAGS})
|
|
SET(CMAKE_REQUIRED_FLAGS "${CMAKE_CXX_FLAGS_INIT} -march=armv8-a+sve+bf16 -msve-vector-bits=256")
|
|
CHECK_CXX_SOURCE_COMPILES("${SVE_BF16_CODE}" CXX_SVE256_FOUND)
|
|
SET(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_SAVE})
|
|
|
|
if(CXX_SVE256_FOUND)
|
|
# Any compiler that supports SVE256 also supports SVE128
|
|
set(CXX_SVE_FOUND TRUE CACHE BOOL "SVE available on host")
|
|
message(STATUS "SVE support detected.")
|
|
else()
|
|
set(CXX_SVE_FOUND FALSE CACHE BOOL "SVE not available on host")
|
|
if(CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64" AND NOT DEFINED ENV{BUILD_IGNORE_SVE_UNAVAILABLE})
|
|
message(FATAL_ERROR "No SVE support on this machine. "
|
|
"Set BUILD_IGNORE_SVE_UNAVAILABLE environment variable to ignore this error.")
|
|
else()
|
|
message(STATUS "No SVE support on this machine.")
|
|
endif()
|
|
endif()
|
|
|
|
mark_as_advanced(CXX_SVE_FOUND CXX_SVE256_FOUND)
|
|
ENDIF()
|