Added HWY_WARN and fallback instead of exiting

Author: bbhattar

ops/brgemm-inl.h

@@ -54,12 +54,25 @@ static bool MakeBrgemm(dnnl::ukernel::brgemm& brg, int64_t m, int64_t n,
   try {
     brg = dnnl::ukernel::brgemm(m, n, k, batch, lda, ldb, ldc, a_dt, b_dt,
                                  c_dt, true);
-    if (!brg) return false;
+    if (!brg) {
+      HWY_WARN("BRGeMM: kernel creation failed m=%lld n=%lld k=%lld.",
+               static_cast<long long>(m), static_cast<long long>(n),
+               static_cast<long long>(k));
+      return false;
+    }
     brg.set_add_C(add_C);
-    if (!brg.finalize()) return false;
+    if (!brg.finalize()) {
+      HWY_WARN("BRGeMM: kernel finalize failed m=%lld n=%lld k=%lld.",
+               static_cast<long long>(m), static_cast<long long>(n),
+               static_cast<long long>(k));
+      return false;
+    }
     brg.generate();
     return true;
   } catch (...) {
+    HWY_WARN("BRGeMM: kernel JIT exception m=%lld n=%lld k=%lld.",
+             static_cast<long long>(m), static_cast<long long>(n),
+             static_cast<long long>(k));
     return false;
   }
 }
@@ -295,7 +308,7 @@ static HWY_NOINLINE bool InitBRGeMMKernels(
 }
 
 template <typename TA, typename TB, typename TC>
-static HWY_NOINLINE void DoMatMul_BRGeMM(
+static HWY_NOINLINE bool DoMatMul_BRGeMM(
     const MatPtrT<TA>& A, const MatPtrT<TB>& B, RowPtrs<TC> C, size_t M,
     size_t K, size_t N, float scale, const float* HWY_RESTRICT add,
     const BRGeMMConfig& cfg, ThreadingContext& ctx, size_t cluster_idx) {
@@ -310,7 +323,7 @@ static HWY_NOINLINE void DoMatMul_BRGeMM(
   if (kern_it == kern_cache.end()) {
     BRGeMMKernelEntry ke;
     if (!InitBRGeMMKernels(cfg, M, K, N, A.Stride(), B.Stride(), ke)) {
-      return;
+      return false;
     }
     kern_it = kern_cache.emplace(kern_key, std::move(ke)).first;
   }
@@ -344,7 +357,10 @@ static HWY_NOINLINE void DoMatMul_BRGeMM(
 
       pe.B_packed_buf.Resize(total_packed);
       uint8_t* B_packed = pe.B_packed_buf.data();
-      if (!B_packed) return;
+      if (!B_packed) {
+        HWY_WARN("BRGeMM: packed B allocation failed.");
+        return false;
+      }
 
       for (size_t nt = 0; nt < ke.N_total_tiles; ++nt) {
         const int ni = (nt < ke.N_full_tiles) ? 0 : 1;
@@ -366,7 +382,8 @@ static HWY_NOINLINE void DoMatMul_BRGeMM(
                                         B_packed + pe.B_ktail_offset[nt]);
           }
         } catch (...) {
-          return;
+          HWY_WARN("BRGeMM: B-packing execution failed.");
+          return false;
         }
       }
     }
@@ -548,6 +565,7 @@ static HWY_NOINLINE void DoMatMul_BRGeMM(
   dnnl::ukernel::brgemm::release_hw_context();
   auto& main_bufs = GetBRGeMMThreadBufs();
   main_bufs.hw_ctx_kernel = nullptr;
+  return true;
 }
 
 #endif  // GEMMA_ONEDNN_BRGEMM

ops/matmul-inl.h

@@ -1089,9 +1089,11 @@ HWY_NOINLINE MMPerKey* MatMul(const MatPtrT<TA>& A, const MatPtrT<TB>& B,
     MMAutoTune<BRGeMMConfig>& brg_tuner = per_key.brgemm_autotune;
 
     if (HWY_LIKELY(brg_tuner.Best())) {
-      DoMatMul_BRGeMM(A, B, C_rows, M, K, N, scale, add, *brg_tuner.Best(),
-                       env.ctx, cluster_idx);
-      return &per_key;
+      if (DoMatMul_BRGeMM(A, B, C_rows, M, K, N, scale, add,
+                           *brg_tuner.Best(), env.ctx, cluster_idx)) {
+        return &per_key;
+      }
+      // BRGeMM failed; fall through to standard matmul.
     }
 
     if (HWY_UNLIKELY(!brg_tuner.HasCandidates())) {
@@ -1100,21 +1102,23 @@ HWY_NOINLINE MMPerKey* MatMul(const MatPtrT<TA>& A, const MatPtrT<TB>& B,
 
     const BRGeMMConfig& cfg = brg_tuner.NextConfig();
     const uint64_t t0 = hwy::timer::Start();
-    DoMatMul_BRGeMM(A, B, C_rows, M, K, N, scale, add, cfg, env.ctx,
-                     cluster_idx);
-    const uint64_t t1 =
-        env.have_timer_stop ? hwy::timer::Stop() : hwy::timer::Start();
-    brg_tuner.NotifyTicks(t1 - t0);
-
-    if (HWY_UNLIKELY(env.print_best && brg_tuner.Best())) {
-      const BRGeMMConfig& best = *brg_tuner.Best();
-      fprintf(stderr,
-              "BRGeMM best: %zux%zux%zu M_blk=%zu N_blk=%zu K_blk=%zu "
-              "batch=%zu\n",
-              M, K, N, best.M_blk, best.N_blk, best.K_blk,
-              best.batch_size);
+    if (DoMatMul_BRGeMM(A, B, C_rows, M, K, N, scale, add, cfg, env.ctx,
+                         cluster_idx)) {
+      const uint64_t t1 =
+          env.have_timer_stop ? hwy::timer::Stop() : hwy::timer::Start();
+      brg_tuner.NotifyTicks(t1 - t0);
+
+      if (HWY_UNLIKELY(env.print_best && brg_tuner.Best())) {
+        const BRGeMMConfig& best = *brg_tuner.Best();
+        fprintf(stderr,
+                "BRGeMM best: %zux%zux%zu M_blk=%zu N_blk=%zu K_blk=%zu "
+                "batch=%zu\n",
+                M, K, N, best.M_blk, best.N_blk, best.K_blk,
+                best.batch_size);
+      }
+      return &per_key;
     }
-    return &per_key;
+    // BRGeMM failed; fall through to standard matmul.
   }
   }  // if constexpr BF16/float
 #endif  // GEMMA_ONEDNN_BRGEMM