better kernel selection heuristic

sonos · Aug 30, 2024 · 2ca2fa1 · 2ca2fa1
1 parent 947d6b5
commit 2ca2fa1
Showing 1 changed file with 1 addition and 2 deletions.
diff --git 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 {