Workaround for fbgemm::FindMinMax

cyyever · cyyever · commit fe87087aebee · 2025-12-08T16:46:36.000+08:00
Signed-off-by: Yuanyuan Chen &lt;cyyever@outlook.com&gt;
diff --git a/aten/src/ATen/native/QuantizedLinear.cpp b/aten/src/ATen/native/QuantizedLinear.cpp
@@ -11,6 +11,7 @@
 #include <ATen/Functions.h>
 #include <ATen/NativeFunctions.h>
 #else
+#include <ATen/ops/aminmax.h>
 #include <ATen/ops/empty.h>
 #include <ATen/ops/empty_like_native.h>
 #include <ATen/ops/fbgemm_linear_fp16_weight_fp32_activation_native.h>
@@ -81,11 +82,19 @@ Tensor fbgemm_linear_int8_weight_fp32_activation(
   // Calculate statistics for quantization of the input Tensor
   float x_min = std::numeric_limits<float>::quiet_NaN();
   float x_max = std::numeric_limits<float>::quiet_NaN();
+#if defined(__AVX__)
   fbgemm::FindMinMax(
       /*m=*/input_ptr,
       /*min=*/&x_min,
       /*max=*/&x_max,
       /*len=*/input.numel());
+#else
+  if (input.numel() > 0) {
+    auto [t_min, t_max] = at::aminmax(input);
+    x_max = t_max.item<float>();
+    x_min = t_min.item<float>();
+  }
+#endif
 
   // Input tensor is quantized as 8-bit unsigned values
   constexpr int kPrecision = 8;
@@ -237,11 +246,19 @@ std::tuple<Tensor, Tensor, double, int64_t> fbgemm_linear_quantize_weight(
   // Calculate weight statistics
   float w_min = std::numeric_limits<float>::quiet_NaN();
   float w_max = std::numeric_limits<float>::quiet_NaN();
+#if defined(__AVX__)
   fbgemm::FindMinMax(
       /*m=*/weight_contig.data_ptr<float>(),
       /*min=*/&w_min,
       /*max=*/&w_max,
       /*len=*/weight_contig.numel());
+#else
+  if (weight_contig.numel() > 0) {
+    auto [t_min, t_max] = at::aminmax(weight_contig);
+    w_max = t_max.item<float>();
+    w_min = t_min.item<float>();
+  }
+#endif
 
   // Choose parameters for quantizing the weight as 8-bit signed integer
   constexpr bool kIsSigned = true;
diff --git a/aten/src/ATen/native/quantized/cpu/qconv_dynamic.cpp b/aten/src/ATen/native/quantized/cpu/qconv_dynamic.cpp
@@ -16,6 +16,7 @@
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
 #else
+#include <ATen/ops/aminmax.h>
 #include <ATen/ops/dequantize.h>                           // for dequantize
 #include <ATen/ops/quantize_per_tensor.h>
 #endif
@@ -29,12 +30,20 @@ at::Tensor PackedConvWeight<kSpatialDim>::apply_dynamic(
   TORCH_CHECK(
       fbgemm::fbgemmSupportedCPU(), "Your CPU does not support FBGEMM.");
 
-  float x_min, x_max;
+  float x_min = std::numeric_limits<float>::quiet_NaN(), x_max = std::numeric_limits<float>::quiet_NaN();
+#if defined(__AVX__)
   fbgemm::FindMinMax(
       /*m=*/input.data_ptr<float>(),
       /*min=*/&x_min,
       /*max=*/&x_max,
       /*len=*/input.numel());
+#else
+  if (input.numel() > 0) {
+    auto [t_min, t_max] = at::aminmax(input);
+    x_max = t_max.item<float>();
+    x_min = t_min.item<float>();
+  }
+#endif
 
   // Input tensor is quantized as 8-bit unsigned values
   static constexpr int precision = 8;
diff --git a/aten/src/ATen/native/quantized/cpu/qlinear_dynamic.cpp b/aten/src/ATen/native/quantized/cpu/qlinear_dynamic.cpp
@@ -69,11 +69,19 @@ at::Tensor PackedLinearWeight::apply_dynamic_impl(
 
   // Calculate statistics for quantization of the input Tensor
   float x_min = std::numeric_limits<float>::quiet_NaN(), x_max = std::numeric_limits<float>::quiet_NaN();
+#if defined(__AVX__)
   fbgemm::FindMinMax(
       /*m=*/input_ptr,
       /*min=*/&x_min,
       /*max=*/&x_max,
       /*len=*/input.numel());
+#else
+  if (input_contig.numel() > 0) {
+    auto [t_min, t_max] = at::aminmax(input_contig);
+    x_max = t_max.item<float>();
+    x_min = t_min.item<float>();
+  }
+#endif
 
   // Input tensor is quantized as 8-bit unsigned values
   static constexpr int precision = 8;
@@ -512,7 +520,7 @@ at::Tensor PackedLinearWeightsOnednn::apply_dynamic_impl(
   x.init(input_desc, input_contig.data_ptr());
   // Find quantization parameters
   float x_max = 0, x_min = 0;
-#ifdef USE_FBGEMM
+#if defined(USE_FBGEMM) && defined(__AVX__)
   // Use FBGEMM's FindMinMax if available since it's faster
   fbgemm::FindMinMax(
       /*m=*/input_contig.data_ptr<float>(),
@@ -738,7 +746,7 @@ at::Tensor PackedLinearWeightsACL::apply_dynamic_impl(
     // Find quantization parameters
     float x_max = 0, x_min = 0;
 
-#ifdef USE_FBGEMM
+#if defined(USE_FBGEMM) && defined(__AVX__)
     // Use FBGEMM's FindMinMax if available since it's faster
     fbgemm::FindMinMax(
         /*m=*/input_contig.data_ptr<float>(),
