remove some linearize calls from tests [pr] (#10978)

* remove some linearize calls from tests

speed_compare_cuda_ptx
test_uop_spec
test_linearizer
test_uops
test_winograd

* more clear assert message

test/external/speed_compare_cuda_ptx.py

@@ -33,7 +33,6 @@ if __name__ == "__main__":
     dev.compiler = PTXCompiler(dev.arch)
     lin = ast_str_to_lin(ast, opts=ptx)
     lin.apply_opts(hand_coded_optimizations(lin))
-    lin.linearize()
     ptx_prg = CompiledRunner(lin.to_program())
 
     # warmup

test/test_linearizer.py

@@ -6,12 +6,12 @@ from dataclasses import replace
 from test.helpers import ast_const
 from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel
 from tinygrad.codegen.lowerer import get_grouped_dims
-from tinygrad.uop.ops import UOp, Ops, GroupOp
+from tinygrad.uop.ops import UOp, Ops, GroupOp, KernelInfo
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner
+from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
 from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import prod, Context, getenv, CI, flatten, dedup, AMX
 from tinygrad.dtype import DType, dtypes
@@ -51,17 +51,20 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d
   r = a.matmul(b, dtype=dtype_out)
   sched = r.schedule()
   realized_ast = sched[-1].ast
-  k = Kernel(realized_ast)
-  k.apply_tensor_cores(1, axis=axis, tc_select=tc_select, tc_opt=tc_opt)
-  k.linearize()
-  wmmas = len([uop for uop in k.uops if uop.op is Ops.WMMA])
-  tcs = len([x for x in k.applied_opts if x.op is OptOps.TC])
+  opts_to_apply = [Opt(OptOps.TC, axis, (tc_select, tc_opt, 1))]
+  realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+
   if ensure_triggered:
+    program = get_program(realized_ast, Device[Device.DEFAULT].renderer)
+    wmmas = len([uop for uop in program.uops if uop.op is Ops.WMMA])
+    tcs = len([x for x in program.applied_opts if x.op is OptOps.TC])
     assert wmmas > 0, "tensor core not triggered"
     assert tcs == 1, "tensor core opt not included"
   else:
-    assert wmmas == 0, "tensor core is incorrectly triggered"
-    assert tcs == 0, "tensor core opt is incorrectly included"
+    try:
+      program = get_program(realized_ast, Device[Device.DEFAULT].renderer)
+      assert False, "OptOps.TC triggered, expected KernelOptError"
+    except KernelOptError: pass
 
 class TestLinearizer(unittest.TestCase):
   def test_arg_dedup(self):

test/test_uops.py

@@ -11,11 +11,11 @@ from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401
 from tinygrad.uop.spec import spec
 from tinygrad.renderer import ProgramSpec
 from tinygrad.kernelize.kernelize import fix_kernel_ops
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 def to_uops_list(u:list[UOp], opts=None, skip_check=False) -> list[UOp]: return full_rewrite(UOp.sink(*u), opts)
 
@@ -409,9 +409,11 @@ class TestAssembly(unittest.TestCase):
     a = Tensor.empty(1024)
     b = Tensor.empty(1024)
     c = (a*b).sum()
-    k = Kernel(c.schedule()[-1].ast)
-    k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
-    uops = k.linearize().uops
+    ast = c.schedule()[-1].ast
+    opts_to_apply = [Opt(OptOps.UNROLL, 0, 4)]
+    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+    program = get_program(ast, Device[Device.DEFAULT].renderer)
+    uops = program.uops
     self.assertEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)
 
 class TestUOpMethod(unittest.TestCase):

test/test_winograd.py

@@ -44,7 +44,6 @@ class TestWinograd(unittest.TestCase):
       with Timing(f"linearize {i} with {len(ops):4d} ops: "):
         l = Kernel(s.ast)
         l.apply_opts(hand_coded_optimizations(l))
-        l.linearize()
       assert len(l.sts) <= 256  # just the current value to prevent regression
       if DEBUG >= 2: print(f"{len(l.sts):4d} shapetrackers with max {max(len(x.views) for x in l.sts)} views")
       for st in l.sts:

test/unit/test_uop_spec.py

@@ -2,26 +2,26 @@ from __future__ import annotations
 import unittest
 
 from tinygrad import Tensor
-from tinygrad.opt.kernel import Kernel
 from tinygrad.helpers import DEBUG
 from tinygrad.uop.ops import UOp, Ops, print_uops
 from tinygrad.uop.spec import type_verify, ast_spec, tensor_uop_spec
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad import dtypes
 from tinygrad.shape.view import View
+from tinygrad.engine.realize import get_program
+from tinygrad.device import Device
 
 class InvalidASTException(Exception): pass
-def helper_test_verify_ast(*stores:UOp) -> Kernel:
+def helper_test_verify_ast(*stores:UOp):
   sink = UOp(Ops.SINK, dtypes.void, stores)
   if DEBUG >= 3:
     for op in stores: print(op)
   try: type_verify(list(sink.toposort()), ast_spec)
   except RuntimeError as e: raise InvalidASTException(e.args)
-  k = Kernel(sink)
-  k.linearize()
-  if DEBUG >= 6: print_uops(k.uops)
-  if DEBUG >= 4: print(k.to_program().src)
-  return k
+  program = get_program(sink, Device[Device.DEFAULT].renderer)
+
+  if DEBUG >= 6: print_uops(program.uops)
+  if DEBUG >= 4: print(program.src)
 
 class TestUOpSpec(unittest.TestCase):
   def test_tiny_add(self):