IQ.Pilot Release Commit @ f2a861c

artifacts/package_sources/tinygrad/extra/optimization/extract_dataset.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+# extract asts from process replay artifacts
+import os
+from test.external.process_replay.process_replay import _pmap
+
+LOGOPS = os.getenv("LOGOPS", "/tmp/sops")
+
+def extract_ast(*args) -> None:
+  open(LOGOPS, "a").write(str(args[1]).replace("\n", "").replace(" ", "")+"\n")
+  return None
+
+if __name__ == "__main__":
+  _pmap({"do_to_program":extract_ast})

artifacts/package_sources/tinygrad/extra/optimization/generate_dataset.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+set -e
+export PAGE_SIZE=1
+export PYTHONPATH=.
+export LOGOPS=/tmp/ops
+export CAPTURE_PROCESS_REPLAY=1
+rm "$LOGOPS" 2>/dev/null || true
+test/external/process_replay/reset.py
+
+CI=1 python3 -m pytest -n=auto test/backend/test_ops.py test/backend/test_nn.py test/unit/test_winograd.py test/null/test_real_world.py --durations=20
+DEV=CL python3 -m pytest test/test_tiny.py
+
+# extract, sort and uniq
+extra/optimization/extract_dataset.py
+sort -u /tmp/ops > /tmp/sops
+ls -lh /tmp/ops /tmp/sops
+gzip -k /tmp/sops
+# mv /tmp/sops.gz extra/datasets/

artifacts/package_sources/tinygrad/extra/optimization/test_beam_search.py

@@ -0,0 +1,105 @@
+import unittest
+import numpy as np
+
+from tinygrad.helpers import BEAM, Timing, prod
+from tinygrad import Variable, Device, Tensor
+from tinygrad.nn import Conv2d
+from tinygrad.uop.ops import AxisType, Ops
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.codegen.opt.postrange import Scheduler
+from tinygrad.codegen.opt.search import get_kernel_actions
+
+def rand(*shape):
+  return Tensor(np.random.rand(*shape).astype(np.float32))
+
+class TestBeamSearch(unittest.TestCase):
+  def setUp(self):
+    self.old_beam = BEAM.value
+    BEAM.value = 2
+  def tearDown(self):
+    BEAM.value = self.old_beam
+
+  def test_variable_ast_beam(self):
+    vi = Variable("a", 1, 10).bind(3)
+    a = rand(10, 3)[:vi]
+    a = (a+1).realize()
+
+  def test_big_prime_number(self):
+    a = rand(367, 367)
+    b = rand(367, 367)
+    c = (a@b).realize()
+    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), atol=1e-4, rtol=1e-4)
+
+  def test_big_prime_number_max(self):
+    a = -rand(367, 367)
+    b = rand(367, 367)
+    # if incorrectly padded 0, the max would be 0 instead of a negative number
+    c = (a*b).max(1)
+    np.testing.assert_allclose(c.numpy(), (a.numpy() * b.numpy()).max(1), atol=1e-4, rtol=1e-4)
+
+  def test_big_prime_number_sum(self):
+    a = rand(367, 367)
+    b = rand(367, 367)
+    # if incorrectly padded 0, the sum would be inf
+    c = (a/b).sum(1).realize()
+    np.testing.assert_allclose(c.numpy(), (a.numpy() / b.numpy()).sum(1), atol=1e-4, rtol=1e-4)
+
+  def test_variable_big_prime_number(self):
+    v = Variable("v", 1, 400).bind(367)
+    a = rand(367, 400)
+    b = rand(400, 367)
+    c = (a[:, :v] @ b[:v, :]).realize()
+    np.testing.assert_allclose(c.numpy(), a[:, :367].numpy() @ b[:367, :].numpy(), atol=1e-4, rtol=1e-4)
+
+  def test_variable_shrink_prime_number(self):
+    v = Variable("v", 1, 400).bind(367)
+    a = rand(400, 367)
+    b = (a.shrink(((0,v), None))+1)[:367,:367].realize()
+    np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)
+
+  def test_no_mutate_rawbuffers(self):
+    a = rand(3, 3).realize()
+    desired = a.numpy() + 1
+    a.assign(a+1)
+    actual = a.numpy()
+    np.testing.assert_allclose(actual, desired)
+
+  @unittest.skip("flaky. CL_OUT_OF_RESOURCES")
+  def test_conv_beam(self):
+    c = Conv2d(3, 16, (3,3))
+    x = rand(1,3,32,32)
+    with Timing():
+      c(x).realize()
+
+  @unittest.skip("flaky, Fatal Python error: Floating point exception")
+  def test_large_ast(self):
+    a = Tensor.rand(3, 3)
+    for _ in range(5):
+      for _ in range(4):
+        a = (a + a) * a
+    a.realize()
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  def test_tc_up(self):
+    tc = Device[Device.DEFAULT].renderer.tensor_cores[0]
+    size = max(tc.dims[0], tc.dims[1]) * 8
+    a, b = Tensor.rand(size, size, dtype=tc.dtype_in), Tensor.rand(size, size, dtype=tc.dtype_in)
+    ast = a.matmul(b, dtype=tc.dtype_out).schedule_linear().src[-1].src[0]
+    s = Scheduler(ast, Device[Device.DEFAULT].renderer)
+    s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1)))
+    up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)])
+    actions = get_kernel_actions(s, include_0=False, max_up=int(up))
+    upcasted = [s for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]
+    assert len(upcasted) > 0, f"expected upcast/unroll actions after TC with max_up={up}, but got none"
+
+  def test_max_up(self):
+    a = Tensor.rand(16, 16)
+    ast = a.schedule_linear().src[-1].src[0]
+    s = Scheduler(ast, Device[Device.DEFAULT].renderer)
+    for max_up in (2, 4):
+      actions = get_kernel_actions(s, include_0=False, max_up=max_up)
+      for up_opts in [s.applied_opts for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]:
+        assert len([opt for opt in up_opts if opt.arg > max_up]) == 0 and len([op for op in up_opts if op.arg <= max_up]) > 0
+
+if __name__ == '__main__':
+  unittest.main()