test_mismatch_reduce

.github/workflows/test.yml

@@ -375,8 +375,8 @@ jobs:
         PYTHONPATH=. python extra/optimization/extract_dataset.py
         gzip -c /tmp/sops > extra/datasets/sops.gz
         DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
-    - name: Repo line count < 16000 lines
-      run: MAX_LINE_COUNT=16000 python sz.py
+    - name: Repo line count < 15500 lines
+      run: MAX_LINE_COUNT=15500 python sz.py
 
   fuzzing:
     name: Fuzzing

test/test_linearizer.py

@@ -214,7 +214,9 @@ class TestLinearizer(unittest.TestCase):
     # these are of size 3 to avoid float4 coalesce
     r = a[:-1] + a[1:]
 
-    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    k = Kernel(r.schedule()[-1].ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
     assert num_loads <= 4, "more load uops than needed"
     assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
@@ -225,7 +227,9 @@ class TestLinearizer(unittest.TestCase):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = a.expand([2]) + b.expand([2])
 
-    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    k = Kernel(r.schedule()[-1].ast)
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops <= 1, "more alu uops than needed"
 

test/test_opt_gemm.py

@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import get_single_element
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 
 class TestOptGemm(unittest.TestCase):
@@ -17,7 +17,9 @@ class TestOptGemm(unittest.TestCase):
     t = self.a.T @ self.b.T
     # TODO: this should be a generic test helper
     si = get_single_element(t.schedule())
-    run = CompiledRunner(get_program(si.ast, opts=opts))
+    k = Kernel(si.ast)
+    k.apply_opts(opts)
+    run = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))
     ExecItem(run, si.bufs).run()
     test = si.bufs[0].numpy().reshape(self.res.shape)
     np.testing.assert_allclose(self.res, test, atol=1e-4)

test/test_uops_stats.py

@@ -173,7 +173,8 @@ class TestStatsOptimized(unittest.TestCase):
     self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats
 
   def test_gemm(self):
-    p = get_program(self.ast_gemm, opts=[])
+    k = Kernel(self.ast_gemm)
+    p = get_program(k.get_optimized_ast(), k.opts)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)
 
@@ -188,30 +189,42 @@ class TestStatsOptimized(unittest.TestCase):
   # this is a good lesson about why UPCASTing is a good idea
 
   def test_gemm_one_upcasted(self):
-    p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4)])
+    k = Kernel(self.ast_gemm)
+    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
+    p = get_program(k.get_optimized_ast(), k.opts)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted(self):
-    p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
+    k = Kernel(self.ast_gemm)
+    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
+    k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
+    k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
+    p = get_program(k.get_optimized_ast(), k.opts)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted_locals(self):
+    k = Kernel(self.ast_gemm)
+    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
+    k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
     try:
-      p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4),
-                                           Opt(OptOps.LOCAL, 0, 4),  Opt(OptOps.LOCAL, 1, 4)])
+      k.apply_opt(Opt(OptOps.LOCAL, 0, 5))
+      k.apply_opt(Opt(OptOps.LOCAL, 1, 5))
     except KernelOptError:
       raise unittest.SkipTest("no locals")
+    p = get_program(k.get_optimized_ast(), k.opts)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_group(self):
+    k = Kernel(self.ast_gemm)
     try:
-      p = get_program(self.ast_gemm, opts=[Opt(OptOps.GROUP, 0, 4)])
+      k.apply_opt(Opt(OptOps.GROUP, 0, 4))
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     SZ = N*N*4
+    p = get_program(k.get_optimized_ast(), k.opts)
     # NOTE: these are sort of wrong. they aren't honoring the IF statement
     self.check_gemm(p, extra_flops=SZ*4)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)

tinygrad/engine/realize.py

@@ -3,17 +3,16 @@ import time, pprint
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer, graph_rewrite, print_uops, track_rewrites
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
 from tinygrad.engine.schedule import ScheduleItem
 from tinygrad.codegen import full_rewrite
-from tinygrad.opt.kernel import Opt
 
 # **************** Program Creation ****************
 
 @track_rewrites(name=lambda _ast,_renderer,ret: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret))
-def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> ProgramSpec:
+def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec:
   """
   Transform an AST into a ProgramSpec. May trigger BEAM search.
 
@@ -28,10 +27,6 @@ def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None)
   if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
 
   # linearize
-  if renderer is None: renderer = Device.default.renderer
-  if opts is not None:
-    assert ast.arg is None, "can't apply opts if sink has an arg"
-    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
   try:
     uops = full_rewrite(ast, renderer)
   except RuntimeError:

tinygrad/opt/kernel.py

@@ -451,8 +451,7 @@ class Kernel:
         return ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
       if op.op is Ops.SINK:
         # NOTE: should group_for_reduces be added to the local_dims?
-        # TODO: arg.name should be able to be None
-        kernel_name = ret.arg.name if ret.arg is not None and ret.arg.name != "test" else self.name if name_override is None else name_override
+        kernel_name = ret.arg.name if ret.arg is not None else self.name if name_override is None else name_override
         return ret.replace(arg=KernelInfo(kernel_name, tuple(self.axis_types), self.dont_use_locals, tuple(self.applied_opts)))
       if op.op is Ops.REDUCE_AXIS:
         reduce_idx = len(self.bufs) + self.reduceops.index(op) * 2

tinygrad/opt/swizzler.py

@@ -15,6 +15,9 @@ merge_views = PatternMatcher([
    lambda x,view: x if x.st is not None and x.op not in GroupOp.Defines and view.st.contiguous and view.shape == x.shape else None),
   (UPat(GroupOp.All-{Ops.DEFINE_GLOBAL}).view(name="view"),
    lambda view: view.const_like(0) if (mask:=view.st.views[-1].mask) is not None and any((x[1]-x[0]) == 0 for x in mask) else None),
+  # only unmaksed VIEW on CONST replaces the ShapeTracker
+  (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="x"),), name="view"),
+   lambda x,view: x.replace(src=(x.src[0].replace(arg=x.st+view.st),)) if all(v.mask is None for v in (x.st+view.st).views) else None),
 ])
 
 def reduce_push_add_ones(src:UOp, r:UOp, view:UOp):
@@ -104,7 +107,7 @@ view_right = merge_views+PatternMatcher([
   (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.REDUCE_AXIS, name="r1"),), name="r2"),
    lambda r1,r2: r1.replace(arg=(r1.arg[0], r2.arg[1]+r1.arg[1])) if r1.arg[0] is r2.arg[0] else None),
   # remove view from sink
-  (UPat(Ops.VIEW, name="v").sink(name="sink"), lambda v,sink: v.src[0].sink(arg=sink.arg)),
+  (UPat(Ops.VIEW, name="v").sink(), lambda v: v.src[0].sink()),
 ])
 
 def check_load_st(glbl:UOp, view:UOp):
@@ -121,9 +124,6 @@ fix_kernel_ops = view_left_through_load+PatternMatcher([
   # add view to LOAD and STORE
   (UPat(Ops.DEFINE_GLOBAL, name="g").load(), lambda g: g.view(g.st).load()),
   (UPat(Ops.DEFINE_GLOBAL, name="g").store(UPat.var('x')), lambda g,x: g.view(g.st).store(x)),
-  # only unmaksed VIEW on CONST replaces the ShapeTracker
-  (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="x"),), name="view"),
-   lambda x,view: x.replace(src=(UOp(Ops.VIEW, arg=x.st+view.st),)) if all(v.mask is None for v in (x.st+view.st).views) else None),
   # VALID
   (UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"),
    lambda self: UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)),

tinygrad/schedule/kernelize.py

@@ -160,8 +160,7 @@ replace_buffers = PatternMatcher([
   (UPat(Ops.SINK, src=(UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Meta, name="x"),),))), lambda x:x),
   # STORE (except for meta ops)
   (UPat(Ops.SINK, src=UPat(GroupOp.All-{Ops.STORE}), name="sink"), lambda ctx,sink:
-   UOp.sink(*[UOp.store(UOp(Ops.DEFINE_GLOBAL, (s:=x.base).dtype.ptr(ctx[i].size), (), i).view(s.st), s) for i,x in enumerate(sink.src)],
-            arg=sink.arg)),
+   UOp.sink(*[UOp.store(UOp(Ops.DEFINE_GLOBAL, (s:=x.base).dtype.ptr(ctx[i].size), (), i).view(s.st), s) for i,x in enumerate(sink.src)])),
   # remove CONTIGUOUS/DEVICE from kernel AST
   (UPat((Ops.CONTIGUOUS, Ops.MSELECT), src=(UPat.var("x"),)), lambda x: x),
   (UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),), name="view"), lambda view: view.replace(src=())),

tinygrad/uop/ops.py

@@ -255,8 +255,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b
     if isinstance(b, tuple) and all_same(b): b = b[0]  # doesn't have to be a VCONST if they are all the same
     ret = UOp(Ops.VCONST if isinstance(b, tuple) else Ops.CONST, dtype, arg=dtypes.as_const(b, dtype))
-    if device is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),))
-    if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape)
+    if shape is not None:
+      from tinygrad.shape.shapetracker import ShapeTracker
+      ret = ret.replace(src=(UOp(Ops.VIEW, dtypes.void, (), ShapeTracker.from_shape(shape, (0,)*len(shape))),))
+    if device is not None:
+      ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device).view(unwrap(ret.st)),))
     return ret
   @staticmethod
   def range(dtype:DType, end:sint, idx:int): return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end),), arg=idx)