Remove AVX compilation on aarch64 (#5065)

Nicoshev · meta-codesync[bot] · commit 9d38b8df4bad · 2025-10-28T14:07:51.000-07:00
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2072 Pull Request resolved: #5065 We are removing AVX2/AVX512 compilation of fbgemm for aarch64. These modules are no longer used, thus removing them should decrease build time and code size. Reviewed By: mcfi, YifanYuan3 Differential Revision: D85603930 fbshipit-source-id: 42ed869c5ba406bdc1430d8e94b3726440eeb904
diff --git a/include/fbgemm/Fbgemm.h b/include/fbgemm/Fbgemm.h
@@ -616,7 +616,7 @@ class FBGEMM_API PackWeightsForConv {
     return W_im2col_packed_;
   }
 
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
   std::shared_ptr<PackedDepthWiseConvMatrix> getPackedWForDepthwise() {
     return W_dw_packed_;
   }
@@ -672,7 +672,7 @@ class FBGEMM_API PackWeightsForConv {
   const conv_param_t<SPATIAL_DIM> conv_param_;
   // Packed weights if we use im2col based convolution implementation
   std::shared_ptr<PackBMatrix<T, accT>> W_im2col_packed_;
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
   // Packed weights if we use depthwise convolution implementation
   std::shared_ptr<PackedDepthWiseConvMatrix> W_dw_packed_;
 #endif // __aarch64__
diff --git a/include/fbgemm/FbgemmConvert.h b/include/fbgemm/FbgemmConvert.h
@@ -47,6 +47,7 @@ FloatToBfloat16_simd(const float* src, bfloat16* dst, size_t size);
 FBGEMM_API void
 Bfloat16ToFloat_simd(const bfloat16* src, float* dst, size_t size);
 
+#if !defined(__aarch64__)
 /**
  * @brief AVX2 implementation to convert fp32 numbers to bf16 numbers.
  *
@@ -58,10 +59,8 @@ FloatToBfloat16_avx2(const float* src, bfloat16* dst, size_t size);
  * @brief AVX512 implementation to convert fp32 numbers to bf16 numbers.
  *
  */
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
 FBGEMM_API void
 FloatToBfloat16_avx512(const float* src, bfloat16* dst, size_t size);
-#endif
 
 /**
  * @brief AVX2 implementation to convert bf16 numbers to fp32 numbers.
@@ -74,7 +73,6 @@ Bfloat16ToFloat_avx2(const bfloat16* src, float* dst, size_t size);
  * @brief AVX512 implementation to convert bf16 numbers to fp32 numbers.
  *
  */
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
 FBGEMM_API void
 Bfloat16ToFloat_avx512(const bfloat16* src, float* dst, size_t size);
 #endif
@@ -124,6 +122,7 @@ Float16ToFloat_simd(const float16* src, float* dst, size_t size);
  * @brief AVX2 implementation to convert fp32 numbers to fp16 numbers.
  *
  */
+#if !defined(__aarch64__)
 FBGEMM_API void FloatToFloat16_avx2(
     const float* src,
     float16* dst,
@@ -134,7 +133,6 @@ FBGEMM_API void FloatToFloat16_avx2(
  * @brief AVX512 implementation to convert fp32 numbers to fp16 numbers.
  *
  */
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
 FBGEMM_API void FloatToFloat16_avx512(
     const float* src,
     float16* dst,
@@ -152,6 +150,7 @@ FBGEMM_API void FloatToFloat16_sve2(
     size_t size,
     bool do_clip = false);
 
+#if !defined(__aarch64__)
 /**
  * @brief AVX2 implementation to convert fp16 numbers to fp32 numbers.
  *
@@ -163,7 +162,6 @@ Float16ToFloat_avx2(const float16* src, float* dst, size_t size);
  * @brief AVX512 implementation to convert fp16 numbers to fp32 numbers.
  *
  */
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
 FBGEMM_API void
 Float16ToFloat_avx512(const float16* src, float* dst, size_t size);
 #endif
diff --git a/include/fbgemm/FbgemmEmbedding.h b/include/fbgemm/FbgemmEmbedding.h
@@ -349,7 +349,7 @@ FBGEMM_API bool EmbeddingSpMDMBlockSize1_(
     bool use_offsets = true,
     bool is_bf16 = false);
 
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
 template <typename IndexType, bool HAS_WEIGHTS>
 void compressed_indices_remap_avx512(
     std::int32_t offsets_numel,
diff --git a/include/fbgemm/FbgemmI8DepthwiseAvx2.h b/include/fbgemm/FbgemmI8DepthwiseAvx2.h
@@ -8,6 +8,8 @@
 
 #pragma once
 
+#if !defined(__aarch64__)
+
 #include <cstdint>
 #include "fbgemm/ConvUtils.h"
 #include "fbgemm/FbgemmBuild.h"
@@ -110,3 +112,5 @@ FBGEMM_API void depthwise_3d_same_pad(
     int num_threads = 1);
 
 } // namespace fbgemm
+
+#endif // !defined(__aarch64__)
diff --git a/include/fbgemm/FbgemmSparse.h b/include/fbgemm/FbgemmSparse.h
@@ -166,7 +166,7 @@ void SparseDenseMMAvx2(
     int ldc,
     bool accum = false);
 
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
 void SparseDenseMMAvx512(
     int M,
     int N,
diff --git a/include/fbgemm/OutputProcessing-inl.h b/include/fbgemm/OutputProcessing-inl.h
@@ -125,7 +125,7 @@ ReQuantizeOutput<FUSE_RELU, Q_GRAN, BIAS_TYPE, outT, inT, nextOPType>::f(
       }
     }
 
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
 
   } else if constexpr (
       instSet == inst_set_t::avx2 || instSet == inst_set_t::avx512) {
@@ -249,7 +249,7 @@ inline int ReQuantizeForFloat<FUSE_RELU, Q_GRAN, outT, inT, nextOPType>::f(
       }
     }
 
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
   } else if constexpr (
       instSet == inst_set_t::avx2 || instSet == inst_set_t::avx512) {
     bool b_symmetric =
diff --git a/include/fbgemm/QuantUtilsAvx2.h b/include/fbgemm/QuantUtilsAvx2.h
@@ -52,6 +52,13 @@ struct FBGEMM_API RequantizationParams {
   TensorQuantizationParams target_qparams;
 };
 
+/// @ingroup fbgemm-quant-utils-avx2
+///
+/// @brief Find the min and max value in a float matrix.
+void FBGEMM_API FindMinMax(const float* m, float* min, float* max, int64_t len);
+
+#if !defined(__aarch64__)
+
 ////////////////////////////////////////////////////////////////////////////////
 // Utility functions
 ////////////////////////////////////////////////////////////////////////////////
@@ -77,11 +84,6 @@ void FusedQuantizeDequantizeAvx2(
 /// <a href="https://www.jstatsoft.org/v08/i14/paper">this paper</a>.
 uint32_t FBGEMM_API Xor128();
 
-/// @ingroup fbgemm-quant-utils-avx2
-///
-/// @brief Find the min and max value in a float matrix.
-void FBGEMM_API FindMinMax(const float* m, float* min, float* max, int64_t len);
-
 void RequantizeFixedPointAvx2(
     const std::int32_t* src,
     std::uint8_t* dst,
@@ -176,4 +178,6 @@ void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfAvx2(
     int input_columns,
     OutputType* output);
 
+#endif // !defined(__aarch64__)
+
 } // namespace fbgemm
diff --git a/include/fbgemm/QuantUtilsAvx512.h b/include/fbgemm/QuantUtilsAvx512.h
@@ -9,7 +9,7 @@
 #pragma once
 
 #include "Types.h"
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
 
 #include <cstdint>
 #include "./FbgemmBuild.h" // @manual
diff --git a/src/FbgemmBfloat16Convert.cc b/src/FbgemmBfloat16Convert.cc
@@ -29,7 +29,7 @@ namespace fbgemm {
 void FloatToBfloat16_simd(const float* src, bfloat16* dst, size_t size) {
   // Run time CPU detection
   if (cpuinfo_initialize()) {
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
     if (fbgemmHasAvx512Support()) {
       FloatToBfloat16_avx512(src, dst, size);
     } else if (fbgemmHasAvx2Support()) {
@@ -48,7 +48,7 @@ void FloatToBfloat16_simd(const float* src, bfloat16* dst, size_t size) {
 void Bfloat16ToFloat_simd(const bfloat16* src, float* dst, size_t size) {
   // Run time CPU detection
   if (cpuinfo_initialize()) {
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
     if (fbgemmHasAvx512Support()) {
       Bfloat16ToFloat_avx512(src, dst, size);
     } else if (fbgemmHasAvx2Support()) {
diff --git a/src/FbgemmFP16.cc b/src/FbgemmFP16.cc
@@ -35,7 +35,7 @@ namespace {
 // the restrictions of ymm register numbers (16).
 constexpr kernel_array_t<float16> kernel_fp16_avx2 = {
     nullptr,
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
     gemmkernel_1x2_Avx2_fp16_fA0fB0fC0,
     gemmkernel_2x2_Avx2_fp16_fA0fB0fC0,
     gemmkernel_3x2_Avx2_fp16_fA0fB0fC0,
@@ -79,7 +79,7 @@ constexpr kernel_array_t<float16> kernel_fp16_neon = {
 
 constexpr kernel_array_t<float16> kernel_fp16_avx512_256 = {
     nullptr,
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
     gemmkernel_1x2_Avx2_fp16_fA0fB0fC0,
     gemmkernel_2x2_Avx2_fp16_fA0fB0fC0,
     gemmkernel_3x2_Avx2_fp16_fA0fB0fC0,
diff --git a/src/FbgemmFP16UKernelsAvx2.h b/src/FbgemmFP16UKernelsAvx2.h
@@ -16,11 +16,15 @@ namespace fbgemm {
 
 using GemmParamsFP16 = GemmParams<float16>;
 
+#if !defined(__aarch64__)
+
 void NOINLINE gemmkernel_1x2_Avx2_fp16_fA0fB0fC0(GemmParamsFP16* gp);
 void NOINLINE gemmkernel_2x2_Avx2_fp16_fA0fB0fC0(GemmParamsFP16* gp);
 void NOINLINE gemmkernel_3x2_Avx2_fp16_fA0fB0fC0(GemmParamsFP16* gp);
 void NOINLINE gemmkernel_4x2_Avx2_fp16_fA0fB0fC0(GemmParamsFP16* gp);
 void NOINLINE gemmkernel_5x2_Avx2_fp16_fA0fB0fC0(GemmParamsFP16* gp);
 void NOINLINE gemmkernel_6x2_Avx2_fp16_fA0fB0fC0(GemmParamsFP16* gp);
 
+#endif // !defined(__aarch64__)
+
 } // namespace fbgemm
diff --git a/src/FbgemmFloat16Convert.cc b/src/FbgemmFloat16Convert.cc
@@ -23,7 +23,7 @@ void FloatToFloat16_simd(
     bool do_clip) {
   // Run time CPU detection
   if (cpuinfo_initialize()) {
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
     if (fbgemmHasAvx512Support()) {
       FloatToFloat16_avx512(src, dst, size, do_clip);
     } else if (fbgemmHasAvx2Support()) {
@@ -42,7 +42,7 @@ void FloatToFloat16_simd(
 void Float16ToFloat_simd(const float16* src, float* dst, size_t size) {
   // Run time CPU detection
   if (cpuinfo_initialize()) {
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
     if (fbgemmHasAvx512Support()) {
       Float16ToFloat_avx512(src, dst, size);
     } else if (fbgemmHasAvx2Support()) {
diff --git a/src/FbgemmSparseDense.cc b/src/FbgemmSparseDense.cc
@@ -193,7 +193,7 @@ void SparseDenseMM(
     float* C,
     int ldc,
     bool accum) {
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
   // Run time CPU detection
   static const auto iset = fbgemmInstructionSet();
 
@@ -229,7 +229,7 @@ FBGEMM_API void fbgemmSparseDenseInt8MM(
     return;
   }
 
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
   // Run time CPU detection
   static const auto iset = fbgemmInstructionSet();
 
diff --git a/src/GroupwiseConv.cc b/src/GroupwiseConv.cc
@@ -121,7 +121,7 @@ static jit_conv_kernel_fp getOrCreateConvKernel(
       accum);
 
   if (cpuinfo_initialize()) {
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
     if (fbgemmHasAvx512VnniSupport()) {
       return GenConvKernel<SPATIAL_DIM, inst_set_t::avx512_vnni>::codeCache_
           .getOrCreate(kernelSig, [&]() {
@@ -954,7 +954,7 @@ static void dispatchOutputProcessing(
   }
 
   if (cpuinfo_initialize()) {
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
     if (fbgemmHasAvx512Support() || fbgemmHasAvx512VnniSupport()) {
       REQUANTIZE_C_PER_G(Avx512);
     } else if (fbgemmHasAvx2Support() || fbgemmHasArmNeonSupport()) {
diff --git a/src/PackWeightsForConv.cc b/src/PackWeightsForConv.cc
@@ -25,7 +25,7 @@ PackWeightsForConv<SPATIAL_DIM, T, accT>::PackWeightsForConv(
   // FbgemmConv.cc
   switch (ConvFastPath<SPATIAL_DIM, accT>(conv_p)) {
     case optimized_conv_t::depthwise: {
-#if !defined(FBGEMM_FBCODE) && defined(__aarch64__)
+#if defined(__aarch64__)
       throw std::runtime_error(
           "PackWeightsForConv<SPATIAL_DIM, T, accT>::PackWeightsForConv(): No fallback available for aarch64");
 #else
@@ -98,7 +98,7 @@ PackWeightsForConv<SPATIAL_DIM, T, accT>::PackWeightsForConv(
 
 template <int SPATIAL_DIM, typename T, typename accT>
 void PackWeightsForConv<SPATIAL_DIM, T, accT>::unpack(T* origin_buf) {
-#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#if !defined(__aarch64__)
   if (W_dw_packed_) {
     W_dw_packed_->unpack(origin_buf);
   } else
diff --git a/src/PackWeightsForDirectConv.cc b/src/PackWeightsForDirectConv.cc
@@ -239,7 +239,7 @@ void fbgemmDirectConv(
     return;
   }
 
-#if !defined(FBGEMM_FBCODE) && defined(__aarch64__)
+#if defined(__aarch64__)
   throw std::runtime_error(
       "fbgemmDirectConv<SPATIAL_DIM, Q_GRAN, FUSE_RELU, BIAS_TYPE>(): No fallback available for aarch64");
 #else
@@ -459,7 +459,7 @@ void fbgemmDirectConv(
     }
   } // else SPATIAL_DIM
 
-#endif // defined(FBGEMM_FBCODE) || !defined(__aarch64__)
+#endif // !defined(__aarch64__)
 }
 
 #define INSTANTIATE_REQUANTIZE_SPATIAL_DIM(                        \
diff --git a/src/QuantUtilsNeon.cc b/src/QuantUtilsNeon.cc
@@ -6,13 +6,15 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include "fbgemm/Utils.h"
+#if defined(__aarch64__)
 
-#if HAVE_SVE
+#include "fbgemm/Utils.h"
 
 #define FBGEMM_EXPORTS
 #include <arm_neon.h> // @manual
+#if HAVE_SVE
 #include <arm_sve.h> // @manual
+#endif
 
 #include <arm_neon_sve_bridge.h> // @manual
 #include <algorithm> //for std::min/std::max
@@ -31,6 +33,50 @@ using namespace std;
 ////////////////////////////////////////////////////////////////////////////////
 // Utility functions
 
+void FindMinMax(const float* m, float* min, float* max, int64_t len) {
+  if (__builtin_expect(len <= 0, 0)) {
+    *min = 0.0f;
+    *max = 0.0f;
+    return;
+  }
+
+  float first = *m;
+
+  float32x4_t temp_min_0 = vdupq_n_f32(first);
+  float32x4_t temp_min_1 = vdupq_n_f32(first);
+  float32x4_t temp_max_0 = vdupq_n_f32(first);
+  float32x4_t temp_max_1 = vdupq_n_f32(first);
+  uint64_t i = 0;
+  uint64_t count = static_cast<uint64_t>(len);
+  uint64_t loopBound = count - (count % 8);
+
+  for (; i < loopBound; i += 8) {
+    float32x4_t v0 = vld1q_f32(m + i);
+    float32x4_t v1 = vld1q_f32(m + i + 4);
+    temp_min_0 = vminq_f32(temp_min_0, v0);
+    temp_min_1 = vminq_f32(temp_min_1, v1);
+    temp_max_0 = vmaxq_f32(temp_max_0, v0);
+    temp_max_1 = vmaxq_f32(temp_max_1, v1);
+  }
+
+  temp_min_0 = vminq_f32(temp_min_0, temp_min_1);
+  temp_max_0 = vmaxq_f32(temp_max_0, temp_max_1);
+
+  float tmp_min_s = vminvq_f32(temp_min_0);
+  float tmp_max_s = vmaxvq_f32(temp_max_0);
+
+  for (; i < count; i++) {
+    float tmp = *m;
+    tmp_min_s = std::min(tmp_min_s, tmp);
+    tmp_max_s = std::max(tmp_max_s, tmp);
+  }
+
+  *min = tmp_min_s;
+  *max = tmp_max_s;
+}
+
+#if HAVE_SVE
+
 template <typename OutputType>
 void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfNeon(
     const std::uint8_t* input,
@@ -141,6 +187,8 @@ INSTANTIATE_QuantizationNeonFunctions8Bits(float16)
 // clang-format on
 #undef INSTANTIATE_QuantizationNeonFunctions8Bits
 
+#endif // HAVE_SVE
+
 } // namespace fbgemm
 
 #endif // __aarch64__
diff --git a/src/TransposeUtils.cc b/src/TransposeUtils.cc

Original file line number	Diff line number	Diff line change
`@@ -125,7 +125,7 @@ ReQuantizeOutput<FUSE_RELU, Q_GRAN, BIAS_TYPE, outT, inT, nextOPType>::f(`
`125`	`125`	`}`
`126`	`126`	`}`
`127`	`127`
`128`		`-#if defined(FBGEMM_FBCODE) \|\| !defined(__aarch64__)`
	`128`	`+#if !defined(__aarch64__)`
`129`	`129`
`130`	`130`	`} else if constexpr (`
`131`	`131`	`instSet == inst_set_t::avx2 \|\| instSet == inst_set_t::avx512) {`
`@@ -249,7 +249,7 @@ inline int ReQuantizeForFloat<FUSE_RELU, Q_GRAN, outT, inT, nextOPType>::f(`
`249`	`249`	`}`
`250`	`250`	`}`
`251`	`251`
`252`		`-#if defined(FBGEMM_FBCODE) \|\| !defined(__aarch64__)`
	`252`	`+#if !defined(__aarch64__)`
`253`	`253`	`} else if constexpr (`
`254`	`254`	`instSet == inst_set_t::avx2 \|\| instSet == inst_set_t::avx512) {`
`255`	`255`	`bool b_symmetric =`