[pull] master from tinygrad:master

docs/abstractions2.py

@@ -84,8 +84,6 @@
 b = UOp.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
 a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
 b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
-a = a.buf_uop_view()
-b = b.buf_uop_view()
 
 # describe the computation
 out = a.alu(Ops.ADD, b)
@@ -104,7 +102,7 @@
 run_schedule(sched)
 
 # check the data out
-assert out.realized is not None and out.realized.as_buffer().cast('I')[0] == 5
+assert out.is_realized and out.buffer.as_buffer().cast('I')[0] == 5
 
 
 print("******** fourth, the Tensor ***********")

test/external/speed_v_theoretical.py

@@ -91,7 +91,7 @@ def test_conv_3x3_256_32_32_256_256(self): self._test_conv_3x3(256, 32, 32, 256,
 
   # theoretical is nv_tflops=165, amd_tflops=123
   def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=115, amd_tflops=80)
-  def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=130, amd_tflops=73)
+  def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=125, amd_tflops=73)
 
   # theoretical is nv_gbs=1008, amd_gbs=960
   def test_gemv_16384_4096(self): self._test_matmul(16384, 4096, 1, nv_gbs=840, amd_gbs=750)

test/imported/test_indexing.py

@@ -21,6 +21,7 @@ def consec(shape, start=1):
 
 # creates strided tensor with base set to reference tensor's base, equivalent to torch.set_()
 def set_(reference: Tensor, shape, strides, offset):
+  raise NotImplementedError("need to implement without calling lazydata.view")
   if reference.lazydata.base.realized is None: reference.realize()
   assert reference.lazydata.base.realized, "base has to be realized before setting it to strided's base"
   strided = Tensor(reference.lazydata.view(ShapeTracker((View.create(shape=shape, strides=strides, offset=offset),))))
@@ -220,28 +221,29 @@ def validate_setting(x):
     # Tensor with stride != 1
     # strided is [1, 3, 5, 7]
 
-    reference = consec((10,))
-    strided = set_(reference, (4,), (2,), 0)
+    # # TODO: set stride
+    # reference = consec((10,))
+    # strided = set_(reference, (4,), (2,), 0)
 
-    numpy_testing_assert_equal_helper(strided[[0]], np.array([1]))
-    numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([1]))
-    numpy_testing_assert_equal_helper(strided[ri([3]), ], np.array([7]))
-    numpy_testing_assert_equal_helper(strided[[1, 2]], np.array([3, 5]))
-    numpy_testing_assert_equal_helper(strided[ri([1, 2]), ], np.array([3, 5]))
-    numpy_testing_assert_equal_helper(strided[ri([[2, 1], [0, 3]]), ],
-                      np.array([[5, 3], [1, 7]]))
+    # numpy_testing_assert_equal_helper(strided[[0]], np.array([1]))
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([1]))
+    # numpy_testing_assert_equal_helper(strided[ri([3]), ], np.array([7]))
+    # numpy_testing_assert_equal_helper(strided[[1, 2]], np.array([3, 5]))
+    # numpy_testing_assert_equal_helper(strided[ri([1, 2]), ], np.array([3, 5]))
+    # numpy_testing_assert_equal_helper(strided[ri([[2, 1], [0, 3]]), ],
+    #                   np.array([[5, 3], [1, 7]]))
 
     # stride is [4, 8]
 
-    strided = set_(reference, (2,), (4,), offset=4)
+    # strided = set_(reference, (2,), (4,), offset=4)
 
-    numpy_testing_assert_equal_helper(strided[[0]], np.array([5]))
-    numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([5]))
-    numpy_testing_assert_equal_helper(strided[ri([1]), ], np.array([9]))
-    numpy_testing_assert_equal_helper(strided[[0, 1]], np.array([5, 9]))
-    numpy_testing_assert_equal_helper(strided[ri([0, 1]), ], np.array([5, 9]))
-    numpy_testing_assert_equal_helper(strided[ri([[0, 1], [1, 0]]), ],
-                      np.array([[5, 9], [9, 5]]))
+    # numpy_testing_assert_equal_helper(strided[[0]], np.array([5]))
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([5]))
+    # numpy_testing_assert_equal_helper(strided[ri([1]), ], np.array([9]))
+    # numpy_testing_assert_equal_helper(strided[[0, 1]], np.array([5, 9]))
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ], np.array([5, 9]))
+    # numpy_testing_assert_equal_helper(strided[ri([[0, 1], [1, 0]]), ],
+    #                   np.array([[5, 9], [9, 5]]))
 
     # reference is 1 2
     #              3 4
@@ -335,79 +337,73 @@ def validate_setting(x):
     # strided is [[1 3 5 7],
     #             [9 11 13 15]]
 
-    reference = Tensor.arange(0., 24).reshape(3, 8)
-    strided = set_(reference, (2,4), (8,2), 1)
-
-    numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([0])],
-                      np.array([1, 9]))
-    numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1])],
-                      np.array([3, 11]))
-    numpy_testing_assert_equal_helper(strided[ri([0]), ri([0])],
-                      np.array([1]))
-    numpy_testing_assert_equal_helper(strided[ri([1]), ri([3])],
-                      np.array([15]))
-    numpy_testing_assert_equal_helper(strided[[ri([0, 0]), ri([0, 3])]],
-                      np.array([1, 7]))
-    numpy_testing_assert_equal_helper(strided[[ri([1]), ri([0, 1, 1, 0, 3])]],
-                      np.array([9, 11, 11, 9, 15]))
-    numpy_testing_assert_equal_helper(strided[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]],
-                      np.array([1, 3, 9, 9]))
-
-    rows = ri([[0, 0],
-                [1, 1]])
-    columns = [0],
-    numpy_testing_assert_equal_helper(strided[rows, columns],
-                      np.array([[1, 1], [9, 9]]))
-
-    rows = ri([[0, 1],
-                [1, 0]])
-    columns = ri([1, 2])
-    numpy_testing_assert_equal_helper(strided[rows, columns],
-                      np.array([[3, 13], [11, 5]]))
-    rows = ri([[0, 0],
-                [1, 1]])
-    columns = ri([[0, 1],
-                  [1, 2]])
-    numpy_testing_assert_equal_helper(strided[rows, columns],
-                      np.array([[1, 3], [11, 13]]))
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).reshape(3, 8)
+    # strided = set_(reference, (2,4), (8,2), 1)
+
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([0])], np.array([1, 9]))
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1])], np.array([3, 11]))
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ri([0])], np.array([1]))
+    # numpy_testing_assert_equal_helper(strided[ri([1]), ri([3])], np.array([15]))
+    # numpy_testing_assert_equal_helper(strided[[ri([0, 0]), ri([0, 3])]], np.array([1, 7]))
+    # numpy_testing_assert_equal_helper(strided[[ri([1]), ri([0, 1, 1, 0, 3])]], np.array([9, 11, 11, 9, 15]))
+    # numpy_testing_assert_equal_helper(strided[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 3, 9, 9]))
+
+    # rows = ri([[0, 0],
+    #            [1, 1]])
+    # columns = [0],
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 1], [9, 9]]))
+
+    # rows = ri([[0, 1],
+    #            [1, 0]])
+    # columns = ri([1, 2])
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[3, 13], [11, 5]]))
+    # rows = ri([[0, 0],
+    #            [1, 1]])
+    # columns = ri([[0, 1],
+    #               [1, 2]])
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 3], [11, 13]]))
 
     # setting values
 
     # strided is [[10, 11],
     #             [17, 18]]
 
-    reference = Tensor.arange(0., 24).reshape(3, 8)
-    strided = set_(reference, (2,2), (7,1), 10)
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).reshape(3, 8)
+    # strided = set_(reference, (2,2), (7,1), 10)
+
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])], np.array([11]))
 
-    numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])],
-                      np.array([11]))
     # TODO non contiguous setitem
     '''
     strided[ri([0]), ri([1])] = -1
     numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])],
                       Tensor([-1]))
     '''
-    reference = Tensor.arange(0., 24).reshape(3, 8)
-    strided = set_(reference, (2,2), (7,1), 10)
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).reshape(3, 8)
+    # strided = set_(reference, (2,2), (7,1), 10)
+
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])], np.array([11, 17]))
 
-    numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])],
-                      np.array([11, 17]))
     # TODO non contiguous setitem
     '''
     strided[ri([0, 1]), ri([1, 0])] = Tensor([-1, 2])
     numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])],
                       Tensor([-1, 2]))
     '''
 
-    reference = Tensor.arange(0., 24).realize().reshape(3, 8)
-    strided = set_(reference, (2,2), (7,1), 10)
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).realize().reshape(3, 8)
+    # strided = set_(reference, (2,2), (7,1), 10)
+
+    # rows = ri([[0],
+    #             [1]])
+    # columns = ri([[0, 1],
+    #               [0, 1]])
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[10, 11], [17, 18]]))
 
-    rows = ri([[0],
-                [1]])
-    columns = ri([[0, 1],
-                  [0, 1]])
-    numpy_testing_assert_equal_helper(strided[rows, columns],
-                      np.array([[10, 11], [17, 18]]))
     # TODO non contiguous setitem
     '''
     strided[rows, columns] = Tensor([[4, 6], [2, 3]])

test/test_renderer_failures.py

@@ -8,6 +8,7 @@
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import dedup, flatten, prod
 from tinygrad.renderer.cstyle import CStyleLanguage
+from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.ops import UOp, Ops
 from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Tensor, _to_np_dtype
@@ -41,7 +42,7 @@ def test_inline_const_alu(self):
     ret = _test_uop_result([Tensor([1])], uops)[0]
     self.assertEqual(ret[0], 1)
 
-@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local and Device.DEFAULT == "PTX", "need local")
+@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "tests for ptx renderer")
 class TestPTXFailures(unittest.TestCase):
   def test_gated_store_with_alu(self):
     a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
@@ -63,5 +64,12 @@ def test_gated_store_with_if(self):
     ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
 
+  def test_gated_define_acc_with_half_dtype(self):
+    a = Tensor.randn(32, 32, dtype=dtypes.half).realize()
+    b = Tensor.randn(34, 32, dtype=dtypes.half).realize()
+    result = a.pad((1,1)).matmul(b, acc_dtype=dtypes.half).numpy()
+    reference = a.pad((1,1)).matmul(b, acc_dtype=dtypes.float).numpy()
+    np.testing.assert_allclose(result, reference)
+
 if __name__ == '__main__':
   unittest.main()

test/test_schedule.py

@@ -2226,7 +2226,7 @@ def test_src_masked_const_folding(self):
     self.assertEqual(len(sched), 1)
     run_schedule(sched)
     # add gets assigned to a new buffer
-    self.assertIsNot(add.lazydata.realized, b.lazydata.realized)
+    self.assertIsNot(add.lazydata.base.realized, b.lazydata.base.realized)
     self.assertListEqual(add.tolist(), [4, 2, 2, 2, 2, 4])
 
   # ** part 3: Tensor variable bindings
@@ -2272,6 +2272,15 @@ def test_copy_to_same_device(self):
     b = a.copy_to_device(a.device)
     check_schedule(b, 0, filter_sink=False)
     b = schedule_graph_rewrite(b)
+    # NOTE: Tensor.empty(4) always creates a VIEW(BUFFER) with ShapeTracker((4,)), we simplify this to jsut a BUFFER
+    # in the scheduler because buffer already has shape (4,)
+    self.assertIs(b, a.base)
+
+  def test_copy_to_same_device_alt(self):
+    a = Tensor.empty(4, 4).lazydata
+    b = a.copy_to_device(a.device)
+    check_schedule(b, 0, filter_sink=False)
+    b = schedule_graph_rewrite(b)
     self.assertIs(b, a)
 
   def test_clone(self):
@@ -2455,14 +2464,17 @@ def test_new_buffer(self):
     b = Tensor.empty(4, 4)
     add = a+b
     check_schedule(add, 1)
-    assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER))).match(add.lazydata.base, {})
+    # NOTE: realized base is always a flat buffer
+    assert UPat(Ops.BUFFER).match(add.lazydata.base, {})
+    # the Tensor UOp can optionally stack a VIEW on top of BUFFER
+    assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER),)).match(add.lazydata, {})
 
   def test_new_buffer_view(self):
     a = Tensor.empty(4, 4)
     b = Tensor.empty(4, 4)
     add = (a+b).reshape(8, 2)
     check_schedule(add, 1)
-    assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER))).match(add.lazydata.base, {})
+    assert UPat(Ops.BUFFER).match(add.lazydata.base, {})
     # VIEW is preserverd after the becomes rewrite.
     self.assertEqual(add.lazydata.shape, (8, 2))
     assert add.lazydata is not add.lazydata.base
@@ -2472,7 +2484,7 @@ def test_become_existing_buffer(self):
     b = a*1
     assert UPat(Ops.MUL).match(b.lazydata, {}) # before scheduling it's a mul
     check_schedule(b, 0)
-    assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER))).match(b.lazydata.base, {}) # scheduling replaces the tensor lazydata with a VIEW(BUFFER)
+    assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER))).match(b.lazydata, {}) # scheduling replaces the tensor lazydata with a VIEW(BUFFER)
     self.assertIs(a.lazydata.base.buffer, b.lazydata.base.buffer)
 
   def test_become_const_in_base(self):

test/test_uops.py

@@ -586,7 +586,6 @@ def test_setitem(self):
     assign.realize()
     self.assertEqual(a.tolist(), [1, 0, 1, 1])
 
-  @unittest.expectedFailure
   def test_buffer_st(self):
     a = UOp.new_buffer(Device.DEFAULT, 10, dtypes.float)
     self.assertEqual(a.st, ShapeTracker.from_shape((10,)))

test/unit/test_tensor_uop_representation.py

@@ -2,7 +2,10 @@
 from tinygrad import Tensor
 from tinygrad.ops import UPat, Ops, UOp
 
-realized_pattern = UPat(Ops.VIEW, src=(UPat(Ops.BUFFER),))
+# NOTE: unlike before base for a realized tensor is always a BUFFER
+realized_pattern = UPat(Ops.BUFFER)
+# after realization, tensor uops become VIEW(BUFFER)
+buffer_view_pattern = UPat(Ops.VIEW, src=(UPat(Ops.BUFFER),))
 const_pattern = UPat(Ops.CONST, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),),)))
 def is_pattern_uop(u:UOp, pat:UPat): assert pat.match(u, {}), f"{u}\nis not\n{pat}"
 def is_pattern(ten:Tensor, pat:UPat): is_pattern_uop(ten.lazydata, pat)
@@ -19,9 +22,6 @@ def test_mutate_add(self):
     self.assertIsNot(pa, a.lazydata)
     self.assertIsNot(pb, b.lazydata)
     self.assertIsNot(pr, ret.lazydata)
-    # NOTE: this becomes a VIEW(VIEW(BUFFER)) because UOp.view no longer instantly folds contiguous VIEW of the same shape
-    # this is fine because realized exists on the base.
-    # TODO: we can make this always be a VIEW(BUFFER) once BUFFER has a ShapeTracker of shape=(N,)
     for t in [a,b,ret]: is_pattern_uop(t.lazydata.base, realized_pattern)
 
   def test_reshape_is_same_parent(self):
@@ -32,6 +32,9 @@ def test_reshape_is_same_parent(self):
     d.realize()
     is_pattern_uop(d.lazydata.base, realized_pattern)
     is_pattern_uop(c.lazydata.base, realized_pattern)
+    # NOTE: we keep movement ops on top of the buffer view
+    is_pattern_uop(c.lazydata, buffer_view_pattern)
+    is_pattern_uop(d.lazydata, UPat(Ops.RESHAPE, src=(buffer_view_pattern,)))
 
   def test_reshape_is_same_child(self):
     a = Tensor([1,2,3])
@@ -53,8 +56,8 @@ def test_add_realized(self):
     b = Tensor([4.,5,6]).realize()
     c = a+b
     print(c.lazydata)
-    is_pattern(c, UPat(Ops.ADD, src=(realized_pattern, realized_pattern)))
-    #is_pattern(c, UPat(Ops.ADD, src=(UPat(Ops.VIEW, src=(realized_pattern,)), UPat(Ops.VIEW, src=(realized_pattern,)))))
+    #is_pattern(c, UPat(Ops.ADD, src=(realized_pattern, realized_pattern)))
+    is_pattern(c, UPat(Ops.ADD, src=(UPat(Ops.VIEW, src=(realized_pattern,)), UPat(Ops.VIEW, src=(realized_pattern,)))))
 
   def test_const_pattern(self):
     a = Tensor(1)
@@ -112,8 +115,8 @@ def test_copyin(self):
     c = a.to("TEST")   # NOTE: this isn't checked
     print(c.lazydata)
     # TODO: COPY on a Tensor becomes a VIEW(COPY), this should be done in the scheduler not in ops
-    is_pattern(c, UPat(Ops.COPY, src=(UPat(Ops.DEVICE), realized_pattern,)))
-    #is_pattern(c, UPat(Ops.VIEW, src=(UPat(Ops.COPY, src=(UPat(Ops.DEVICE), realized_pattern,)),)))
+    #is_pattern(c, UPat(Ops.COPY, src=(UPat(Ops.DEVICE), realized_pattern,)))
+    is_pattern(c, UPat(Ops.VIEW, src=(UPat(Ops.COPY, src=(UPat(Ops.DEVICE), realized_pattern,)),)))
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_uop_symbolic.py

@@ -140,6 +140,9 @@ def test_div_1(self):
   def test_mod_1(self):
     self.helper_test_variable(Variable("a", 0, 8)%1, 0, 0, "0")
 
+  def test_max_folds(self):
+    self.helper_test_variable(Variable("a", 0, 20).maximum(10).maximum(11), 11, 20, "max(a, 11)")
+
   def test_add_min_max(self):
     self.helper_test_variable(Variable("a", 0, 8) * 2 + 12, 12, 16+12, "((a*2)+12)")
 

tinygrad/engine/schedule.py

@@ -90,7 +90,8 @@ def add_buffers(buf:UOp, tensor_map:dict[UOp, list[UOp]], ctx:ScheduleContext, c
   # SINK is passthrough
   if buf.op is Ops.SINK: return buf.replace(src=tuple(add_buffers(x, tensor_map, ctx, cache) for x in buf.src))
   # skip creating buffers for CONST/BIND/DEVICE/BUFFER
-  if buf.base.is_realized or buf.base.op in {Ops.CONST, Ops.BIND, Ops.DEVICE}: return buf
+  if buf.base.op in {Ops.CONST, Ops.BIND, Ops.DEVICE}: return buf
+  if buf.base.op is Ops.BUFFER: return buf.view(unwrap(buf.st))
   # VIEW is passthrough
   if buf is not buf.base:
     cache[buf] = ret = add_buffers(buf.base, tensor_map, ctx, cache).view(unwrap(buf.st))
@@ -387,8 +388,8 @@ def replace_contiguous(ctx:dict[UOp, UOp], alu:UOp):
   # remove contiguous if we can just view the buffer
   (UPat(Ops.CONTIGUOUS, name="root", src=(UPat(Ops.VIEW, name="view", src=(UPat(Ops.BUFFER, name="buf"),)),)),
    lambda root,view,buf: view if view.st.contiguous and view.size == buf.size else None),
-  # double contiguous is one contiguous
-  (UPat(Ops.CONTIGUOUS, name="root", src=(UPat(Ops.CONTIGUOUS),)), lambda root: root.src[0]),
+  # contiguous/buffer is already contiguous
+  (UPat(Ops.CONTIGUOUS, name="root", src=(UPat((Ops.CONTIGUOUS, Ops.BUFFER)),)), lambda root: root.src[0]),
   # support for using a contiguous permuted view instead of the parent view if one exists
   (UPat(Ops.CONTIGUOUS, name="contig", src=(UPat(Ops.VIEW, name="src"),)), found_contiguous),
   (UPat(GroupOp.ALU, name="alu"), replace_contiguous),
@@ -480,7 +481,7 @@ def append_uop(ctx:ScheduleContext, view:UOp, buf_uop:UOp) -> None:
 
 remove_movement_ops = merge_views+PatternMatcher([
   # NOTE: movement ops are always applied to base
-  (UPat(GroupOp.Movement, name="mov", src=(UPat.any(UPat.var("x").view(), UPat.var("x")))), lambda x,mov: x.view(unwrap(mov.st))),
+  (UPat(GroupOp.Movement, name="mov", src=(UPat.var("x"),)), lambda x,mov: x.view(unwrap(mov.st))),
   # some masked views can collapse to 0, VIEW(x) -> CONST(VIEW)
   (UPat(Ops.VIEW, name="view"),
    lambda view: view.const_like(0) if (vm:=view.st.views[-1].mask) is not None and any((x[1]-x[0]) == 0 for x in vm) else None),