From bb1f376ae631e352fce11fe3131cb82fe60dbbf4 Mon Sep 17 00:00:00 2001
From: Sieds Lykles <93992551+S-Lykles@users.noreply.github.com>
Date: Fri, 19 Sep 2025 22:52:06 +0200
Subject: [PATCH] profile z3 (#12248)

---
 tinygrad/uop/spec.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py
index c1d1384103..79f4769bea 100644
--- a/tinygrad/uop/spec.py
+++ b/tinygrad/uop/spec.py
@@ -1,7 +1,7 @@
 from typing import cast, Callable
 from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid
-from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context
+from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context, cpu_profile
 from tinygrad.shape.shapetracker import ShapeTracker
 try:
   import z3
@@ -138,11 +138,12 @@ def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)):
   solver = z3.Solver(ctx=z3.Context())
   z3_idx, z3_mask = uops_to_z3(solver, idx.src[1], mask)
   solver.add(z3_mask)
-  if solver.check((z3_idx<0)|(sz<=z3_idx)) == z3.sat:
-    print(f"idx={idx.src[1].render(simplify=False)}")
-    print(f"mask & gate={mask.render(simplify=False)}")
-    print(f"# OUT OF BOUNDS ACCESS: at {solver.model()} INDEX not in 0 - {sz}\nconstraints = {solver}")
-    return False
+  with cpu_profile("validate index with z3", "TINY"):
+    if solver.check((z3_idx<0)|(sz<=z3_idx)) == z3.sat:
+      print(f"idx={idx.src[1].render(simplify=False)}")
+      print(f"mask & gate={mask.render(simplify=False)}")
+      print(f"# OUT OF BOUNDS ACCESS: at {solver.model()} INDEX not in 0 - {sz}\nconstraints = {solver}")
+      return False
   return True
 
 def validate_store(idx:UOp, val:UOp, gate:UOp=UOp.const(dtypes.bool, True)):