mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
ced5c89b6f
This patch is part of half float performance optimization on CPU: * add specification for dtype `Half` in `Vectorized<>` under both avx256 and avx512. * add specification for dtype `Half` in functional utils, e.g. `vec::map_reduce<>()`, which uses float32 as accumulate type. Also add a helper struct `vec_hold_type<scalar_t>`, since Vectorized<Half>::value_type is pointing to its underlying storage type which is `uint16_t`, leading to error if the kernel uses `Vec::value_type`. Half uses the same logic as BFloat16 in the Vectorized<>, each half vector is mapped to 2x float vectors for computation. Notice that this patch modified the cmake files by adding **-mf16c** on AVX2 build, from https://gcc.gnu.org/onlinedocs/gcc/x86-Options.html, we can see that all the hardware platforms that support **avx2** already have **f16c** Pull Request resolved: https://github.com/pytorch/pytorch/pull/96076 Approved by: https://github.com/malfet
83 lines
2.4 KiB
CMake
83 lines
2.4 KiB
CMake
INCLUDE(CheckCSourceRuns)
|
|
INCLUDE(CheckCXXSourceRuns)
|
|
|
|
SET(AVX_CODE "
|
|
#include <immintrin.h>
|
|
|
|
int main()
|
|
{
|
|
__m256 a;
|
|
a = _mm256_set1_ps(0);
|
|
return 0;
|
|
}
|
|
")
|
|
|
|
SET(AVX512_CODE "
|
|
#include <immintrin.h>
|
|
|
|
int main()
|
|
{
|
|
__m512i a = _mm512_set_epi8(0, 0, 0, 0, 0, 0, 0, 0,
|
|
0, 0, 0, 0, 0, 0, 0, 0,
|
|
0, 0, 0, 0, 0, 0, 0, 0,
|
|
0, 0, 0, 0, 0, 0, 0, 0,
|
|
0, 0, 0, 0, 0, 0, 0, 0,
|
|
0, 0, 0, 0, 0, 0, 0, 0,
|
|
0, 0, 0, 0, 0, 0, 0, 0,
|
|
0, 0, 0, 0, 0, 0, 0, 0);
|
|
__m512i b = a;
|
|
__mmask64 equality_mask = _mm512_cmp_epi8_mask(a, b, _MM_CMPINT_EQ);
|
|
return 0;
|
|
}
|
|
")
|
|
|
|
SET(AVX2_CODE "
|
|
#include <immintrin.h>
|
|
|
|
int main()
|
|
{
|
|
__m256i a = {0};
|
|
a = _mm256_abs_epi16(a);
|
|
__m256i x;
|
|
_mm256_extract_epi64(x, 0); // we rely on this in our AVX2 code
|
|
return 0;
|
|
}
|
|
")
|
|
|
|
MACRO(CHECK_SSE lang type flags)
|
|
SET(__FLAG_I 1)
|
|
SET(CMAKE_REQUIRED_FLAGS_SAVE ${CMAKE_REQUIRED_FLAGS})
|
|
FOREACH(__FLAG ${flags})
|
|
IF(NOT ${lang}_${type}_FOUND)
|
|
SET(CMAKE_REQUIRED_FLAGS ${__FLAG})
|
|
IF(lang STREQUAL "CXX")
|
|
CHECK_CXX_SOURCE_RUNS("${${type}_CODE}" ${lang}_HAS_${type}_${__FLAG_I})
|
|
ELSE()
|
|
CHECK_C_SOURCE_RUNS("${${type}_CODE}" ${lang}_HAS_${type}_${__FLAG_I})
|
|
ENDIF()
|
|
IF(${lang}_HAS_${type}_${__FLAG_I})
|
|
SET(${lang}_${type}_FOUND TRUE CACHE BOOL "${lang} ${type} support")
|
|
SET(${lang}_${type}_FLAGS "${__FLAG}" CACHE STRING "${lang} ${type} flags")
|
|
ENDIF()
|
|
MATH(EXPR __FLAG_I "${__FLAG_I}+1")
|
|
ENDIF()
|
|
ENDFOREACH()
|
|
SET(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_SAVE})
|
|
|
|
IF(NOT ${lang}_${type}_FOUND)
|
|
SET(${lang}_${type}_FOUND FALSE CACHE BOOL "${lang} ${type} support")
|
|
SET(${lang}_${type}_FLAGS "" CACHE STRING "${lang} ${type} flags")
|
|
ENDIF()
|
|
|
|
MARK_AS_ADVANCED(${lang}_${type}_FOUND ${lang}_${type}_FLAGS)
|
|
|
|
ENDMACRO()
|
|
|
|
CHECK_SSE(C "AVX" " ;-mavx;/arch:AVX")
|
|
CHECK_SSE(C "AVX2" " ;-mavx2 -mfma -mf16c;/arch:AVX2")
|
|
CHECK_SSE(C "AVX512" " ;-mavx512f -mavx512dq -mavx512vl -mavx512bw -mfma;/arch:AVX512")
|
|
|
|
CHECK_SSE(CXX "AVX" " ;-mavx;/arch:AVX")
|
|
CHECK_SSE(CXX "AVX2" " ;-mavx2 -mfma -mf16c;/arch:AVX2")
|
|
CHECK_SSE(CXX "AVX512" " ;-mavx512f -mavx512dq -mavx512vl -mavx512bw -mfma;/arch:AVX512")
|