[AARCH64] Fall back to GEMM if mkldnn_matmul fails (#115936) (#116666)

huydhn · malfet · web-flow · commit 7e97e4b4b6df · 2024-01-02T20:38:28.000-08:00
- Add call to `at::globalContext().userEnabledMkldnn()` to `apply_mkldnn_matmul_heur` - Surround calls to `mkldnn_matmul` with `try {} catch {}` - Print warning and fall back to BLAS (by calling `at::globalContext().setUserEnabledMkldnn()`) if `mkldnn_matmul()` fails Test plan: On Linux arm run: ```shell $ sudo chmod 400 /sys; python -c "import torch;m=torch.nn.Linear(1, 32);print(torch.__version__);print(m(torch.rand(32, 1)))" Error in cpuinfo: failed to parse the list of possible processors in /sys/devices/system/cpu/possible Error in cpuinfo: failed to parse the list of present processors in /sys/devices/system/cpu/present Error in cpuinfo: failed to parse both lists of possible and present processors 2.3.0.dev20231215 bad err=11 in Xbyak::Error bad err=11 in Xbyak::Error /home/ubuntu/miniconda3/envs/py311/lib/python3.11/site-packages/torch/nn/modules/linear.py:116: UserWarning: mkldnn_matmul failed, switching to BLAS gemm:internal error (Triggered internally at /pytorch/aten/src/ATen/native/LinearAlgebra.cpp:1509.) return F.linear(input, self.weight, self.bias) tensor([[-0.5183, 0.2279, -0.4035, ..., -0.3446, 0.0938, -0.2113], [-0.5111, 0.2362, -0.3821, ..., -0.3536, 0.1011, -0.2159], [-0.6387, 0.0894, -0.7619, ..., -0.1939, -0.0282, -0.1344], ..., [-0.6352, 0.0934, -0.7516, ..., -0.1983, -0.0247, -0.1366], [-0.4790, 0.2733, -0.2862, ..., -0.3939, 0.1338, -0.2365], [-0.5702, 0.1682, -0.5580, ..., -0.2796, 0.0412, -0.1782]], grad_fn=<AddmmBackward0>) ``` Fixes #114750 Pull Request resolved: #115936 Approved by: https://github.com/lezcano Co-authored-by: Nikita Shulga <nshulga@meta.com>
diff --git a/aten/src/ATen/native/LinearAlgebra.cpp b/aten/src/ATen/native/LinearAlgebra.cpp
@@ -1375,7 +1375,7 @@ static inline int64_t get_mkldnn_matmul_min_size() {
 static inline bool apply_mkldnn_matmul_heur(int64_t m, int64_t k, int64_t n) {
   const int64_t min_dim = get_mkldnn_matmul_min_dim();
   const int64_t min_size = get_mkldnn_matmul_min_size();
-  return m > min_dim && k > min_dim && n > min_dim && m * k * n > min_size;
+  return at::globalContext().userEnabledMkldnn() && m > min_dim && k > min_dim && n > min_dim && m * k * n > min_size;
 }
 
 
@@ -1500,10 +1500,15 @@ static void addmm_impl_cpu_(
   if (transpose_c) {
     bool apply_heur = apply_mkldnn_matmul_heur(b.sizes()[0], b.sizes()[1], a.sizes()[1]);
     if (apply_heur && transpose_a && !transpose_b && result.scalar_type() == at::ScalarType::Float) {
+      try {
         mkldnn_matmul(b, a, c, beta.to<float>(), alpha.to<float>());
         // We have dispatched to ACL GEMM for single precision float
         // so do not need to dispatch to BLAS GEMM below
         dispatched = true;
+      } catch (const std::exception& e) {
+        TORCH_WARN("mkldnn_matmul failed, switching to BLAS gemm:", e.what());
+        at::globalContext().setUserEnabledMkldnn(false);
+      }
     }
   }
 #endif
@@ -1752,8 +1757,13 @@ static inline void bmm_out_or_baddbmm_(const Tensor& self_or_result_, const Tens
 
   bool apply_heur = apply_mkldnn_matmul_heur(batch1.sizes()[1], batch1.sizes()[2], batch2.sizes()[2]);
   if (apply_heur && use_mkldnn_lower_precision_matmul(batch1, batch2, self_or_result)) {
+    try {
       mkldnn_matmul(batch1, batch2, self_or_result, beta.to<float>(), alpha.to<float>());
       return;
+    } catch (const std::exception& e) {
+      TORCH_WARN("mkldnn_matmul failed, switching to baddbmm:", e.what());
+      at::globalContext().setUserEnabledMkldnn(false);
+    }
   }
 
   if (contraction_size * res_rows * res_cols < 400) {
