From 17cec8d64584f24eb826b037cef13cf8f62aae4a Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Wed, 24 Sep 2025 23:42:32 -0400
Subject: [PATCH] RANGEIFY winograd test (#12297)

speed seems fine
---
 .github/workflows/test.yml |  2 ++
 test/unit/test_winograd.py | 19 +++++++++++--------
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 8d128fa414..9254c27bd4 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -578,6 +578,8 @@ jobs:
       with:
         key: metal
         deps: testing
+    - name: some unit tests
+      run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/unit/test_winograd.py --durations=20
     - name: Test METAL=1 RANGEIFY=1
       run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/test_ops.py --durations=20
     - name: Run process replay tests
diff --git a/test/unit/test_winograd.py b/test/unit/test_winograd.py
index d1ddfc8bae..e9cbd822bc 100644
--- a/test/unit/test_winograd.py
+++ b/test/unit/test_winograd.py
@@ -1,7 +1,7 @@
 import unittest, sys
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, Context, nn
-from tinygrad.helpers import CI, Profiling, WINO
+from tinygrad.helpers import CI, Profiling, WINO, RANGEIFY
 
 @unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
 class TestWinogradClose(unittest.TestCase):
@@ -35,32 +35,35 @@ class TestWinograd(unittest.TestCase):
   def test_forward_kernels(self):
     x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
     out = Tensor.conv2d(x,w)
-    self.assertEqual(len(out.schedule()), 4)
+    self.assertEqual(len(out.schedule()), 2 if RANGEIFY else 4)
 
   def test_backward_kernels(self):
     x,w = Tensor.empty(1,4,9,9,requires_grad=True).realize(), Tensor.empty(4,4,3,3,requires_grad=True).realize()
     out = Tensor.conv2d(x,w, padding=1)
     out.mean().backward()
     backward_schedule = Tensor.schedule(x.grad, w.grad)
-    self.assertEqual(len(backward_schedule), 9)
+    self.assertEqual(len(backward_schedule), 6 if RANGEIFY else 9)
 
   def test_counters(self):
     IC, OC, X, Y = 4,4,9,9
     #OC, IC, X, Y = 512, 256, 8, 8
     x,w = Tensor.rand(1,IC,Y,X).realize(), Tensor.rand(OC,IC,3,3).realize()
     GlobalCounters.reset()
-    Tensor.conv2d(x,w).realize()
+    with Context(WINO=1):
+      Tensor.conv2d(x,w).realize()
     ops_wino, mem_wino = GlobalCounters.global_ops, GlobalCounters.global_mem
-    WINO.value = 0
     GlobalCounters.reset()
-    Tensor.conv2d(x,w).realize()
+    with Context(WINO=0):
+      Tensor.conv2d(x,w).realize()
     ops_normal, mem_normal = GlobalCounters.global_ops, GlobalCounters.global_mem
 
     ops_ratio, mem_ratio = ops_wino/ops_normal, mem_wino/mem_normal
     print(f"ops: normal {ops_normal:9d} wino {ops_wino:9d} ratio {ops_ratio:.2f}")
     print(f"mem: normal {mem_normal:9d} wino {mem_wino:9d} ratio {mem_ratio:.2f}")
-    self.assertLess(ops_ratio, 2.6)  # TODO: there's issues with factorization now
-    self.assertLess(mem_ratio, 10)
+
+    if not RANGEIFY:
+      self.assertLess(ops_ratio, 2.6)  # TODO: there's issues with factorization now
+      self.assertLess(mem_ratio, 10)
 
   def test_dtype(self):
     IC, OC, X, Y = 4,4,9,9