From 88dcdae4858d921187a5685a70a8ead743366584 Mon Sep 17 00:00:00 2001
From: Francis Lam <flam@alum.mit.edu>
Date: Tue, 2 Apr 2024 16:52:05 -0700
Subject: [PATCH] search: fix counting of upcasts to ignore TC upcasts (#4045)

TC upcasts don't impact the size or complexity of the kernel code
---
 tinygrad/features/search.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tinygrad/features/search.py b/tinygrad/features/search.py
index a11448b77a..3bf3cce6f2 100644
--- a/tinygrad/features/search.py
+++ b/tinygrad/features/search.py
@@ -87,11 +87,11 @@ def get_linearizer_actions(lin:Linearizer, include_0=True) -> Dict[int, Lineariz
     lin2 = lin.copy()
     try:
       lin2.apply_opt(a)
-      up, lcl = 1, 1
+      up, lcl, tc_up = 1, 1, prod(tc.dims)//prod([x[1] for x in tc.threads]) if (tc:=lin2.tensor_core) else 1
       for s,c in zip(lin2.full_shape, lin2.colors()):
         if c in {"magenta", "yellow"}: up *= s
         elif c in {"cyan", "green", "white"}: lcl *= s
-      if up > max_up or lcl > max_lcl: continue
+      if up//tc_up > max_up or lcl > max_lcl: continue
       acted_lins[i+1] = lin2
     except KernelOptError: pass
   return acted_lins