From 2ca2fa1df66fceb4c323601132d20d5d10e7e9cf Mon Sep 17 00:00:00 2001
From: Mathieu Poumeyrol <kali@zoy.org>
Date: Fri, 30 Aug 2024 17:12:32 +0200
Subject: [PATCH] better kernel selection heuristic

---
 core/src/ops/einsum/codegen.rs | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/core/src/ops/einsum/codegen.rs b/core/src/ops/einsum/codegen.rs
index 6c0e6f0d87..ae77f12b54 100644
--- a/core/src/ops/einsum/codegen.rs
+++ b/core/src/ops/einsum/codegen.rs
@@ -329,8 +329,7 @@ fn select_kernel_and_packing(
                     .iter()
                     .min_by_key(|a| {
                         ((m as usize).divceil(a.0.mr()) * (n as usize).divceil(a.0.nr()))
-                            * a.0.mr()
-                            * a.0.nr()
+                            * (a.0.mr() * a.0.nr() + 100)
                     })
                     .unwrap()
             } else {