Fix build break (#5076)

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2081

Pull Request resolved: #5076

D85603930 removed AVX from aarch64 compilation, and broke Sigrid build.

Proposed changes fix the build break.

There is an fbgemm routine without ref implementation, so we need to implement a NEON port at a later diff.

For now, four AVX2 files are compiled with the NEON package

Reviewed By: YifanYuan3

Differential Revision: D85918535

fbshipit-source-id: 3a9892535046a0edc05c8c1fccbb0dac8ca8de35

Author: Nicoshev

include/fbgemm/Fbgemm.h

@@ -616,7 +616,7 @@ class FBGEMM_API PackWeightsForConv {
     return W_im2col_packed_;
   }
 
-#if !defined(__aarch64__)
+#if defined(FBGEMM_FBCODE) || !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(__aarch64__)
+#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
   // Packed weights if we use depthwise convolution implementation
   std::shared_ptr<PackedDepthWiseConvMatrix> W_dw_packed_;
 #endif // __aarch64__

include/fbgemm/FbgemmI8DepthwiseAvx2.h

@@ -8,8 +8,6 @@
 
 #pragma once
 
-#if !defined(__aarch64__)
-
 #include <cstdint>
 #include "fbgemm/ConvUtils.h"
 #include "fbgemm/FbgemmBuild.h"
@@ -112,5 +110,3 @@ FBGEMM_API void depthwise_3d_same_pad(
     int num_threads = 1);
 
 } // namespace fbgemm
-
-#endif // !defined(__aarch64__)

include/fbgemm/QuantUtilsAvx2.h

@@ -96,6 +96,8 @@ void RequantizeAvx2(
     int len,
     const RequantizationParams& params);
 
+#endif // !defined(__aarch64__)
+
 /// @ingroup fbgemm-quant-utils-avx2
 ///
 /// Requantize with avx2 and bias is fused.
@@ -145,6 +147,8 @@ FBGEMM_API void requantizeForFloatAvx2(
     int ld_in,
     const requantizationForFloatParams_t& r);
 
+#if !defined(__aarch64__)
+
 template <typename InputType, int BIT_RATE>
 void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2(
     const InputType* input,

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(__aarch64__)
+#if !defined(FBGEMM_FBCODE) && defined(__aarch64__)
       throw std::runtime_error(
           "PackWeightsForConv<SPATIAL_DIM, T, accT>::PackWeightsForConv(): No fallback available for aarch64");
 #else
@@ -61,7 +61,7 @@ PackWeightsForConv<SPATIAL_DIM, T, accT>::PackWeightsForConv(
       break;
     }
     case optimized_conv_t::directconv: {
-#if defined(__aarch64__)
+#if !defined(FBGEMM_FBCODE) && 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(__aarch64__)
+#if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
   if (W_dw_packed_) {
     W_dw_packed_->unpack(origin_buf);
   } else

src/PackWeightsForDirectConv.cc

@@ -239,7 +239,7 @@ void fbgemmDirectConv(
     return;
   }
 
-#if defined(__aarch64__)
+#if !defined(FBGEMM_FBCODE) && defined(__aarch64__)
   throw std::runtime_error(
       "fbgemmDirectConv<SPATIAL_DIM, Q_GRAN, FUSE_RELU, BIAS_TYPE>(): No fallback available for aarch64");
 #else

src/QuantUtilsAvx2.cc

@@ -280,6 +280,8 @@ SPECIALIZE_FUSEDDQAVX2(int8_t)
 
 #undef SPECIALIZE_FUSEDDQAVX2
 
+#ifndef __aarch64__
+
 void FindMinMax(const float* m, float* min, float* max, int64_t len) {
   if (len <= 0) {
     *min = 0.0f;
@@ -317,6 +319,8 @@ void FindMinMax(const float* m, float* min, float* max, int64_t len) {
   *max = temp_max;
 }
 
+#endif
+
 ////////////////////////////////////////////////////////////////////////////////
 // Requantization (with floats)