From 21f1c4cc09d81f63f62f970dd0cb01d058d5ba3d Mon Sep 17 00:00:00 2001 From: Ignacio Sica Date: Wed, 25 Jun 2025 16:37:17 -0300 Subject: [PATCH] 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 | 1 - test/test_linearizer.py | 21 ++++++++++++--------- test/test_uops.py | 12 +++++++----- test/test_winograd.py | 1 - test/unit/test_uop_spec.py | 14 +++++++------- 5 files changed, 26 insertions(+), 23 deletions(-) diff --git a/test/external/speed_compare_cuda_ptx.py b/test/external/speed_compare_cuda_ptx.py index 068b7619ba..f5ed26f755 100644 --- a/test/external/speed_compare_cuda_ptx.py +++ b/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 diff --git a/test/test_linearizer.py b/test/test_linearizer.py index 04652c1e4c..44d7e85920 100644 --- a/test/test_linearizer.py +++ b/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): diff --git a/test/test_uops.py b/test/test_uops.py index 9c794de8c7..953f18af62 100644 --- a/test/test_uops.py +++ b/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): diff --git a/test/test_winograd.py b/test/test_winograd.py index fa34146a8d..515694c6c2 100644 --- a/test/test_winograd.py +++ b/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: diff --git a/test/unit/test_uop_spec.py b/test/unit/test_uop_spec.py index f5227e5894..5fafb6d1c3 100644 --- a/test/unit/test_uop_spec.py +++ b/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):