mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +08:00
Daniel Koepping
66 lines
3.0 KiB
Python
66 lines
3.0 KiB
Python
#!/usr/bin/env python
|
|
import unittest
|
|
from tinygrad import dtypes, Tensor, GlobalCounters
|
|
from test.helpers import assert_kernel_count
|
|
N = 200 # has to be bigger than the cache to fail
|
|
|
|
class TestAssign(unittest.TestCase):
|
|
def test_assign_contiguous(self):
|
|
b = Tensor.arange(16).reshape(4,4).clone().realize()
|
|
a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1)
|
|
GlobalCounters.reset()
|
|
b.assign(a.contiguous()).realize()
|
|
assert_kernel_count(2)
|
|
|
|
def test_assign_contiguous_permute(self):
|
|
b = Tensor.arange(16).reshape(4,4).clone().realize()
|
|
a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1).permute((1,0))
|
|
GlobalCounters.reset()
|
|
b.assign(a.contiguous()).realize()
|
|
assert_kernel_count(2)
|
|
|
|
# IEEE 754: 1.0f = 0x3f800000, 2.0f = 0x40000000, 3.0f = 0x40400000, 4.0f = 0x40800000
|
|
REVERSED = [0x40800000, 0x40400000, 0x40000000, 0x3f800000]
|
|
|
|
def test_assign_dtype_mismatch(self):
|
|
# assign should not implicitly cast dtypes - this can lose precision
|
|
a = Tensor.zeros(4, dtype=dtypes.float32).contiguous().realize()
|
|
b = Tensor([1, 2, 3, 4], dtype=dtypes.int32)
|
|
with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
|
|
a.assign(b)
|
|
|
|
def test_chained_assign_kernel_count(self):
|
|
"""Chained pending assigns must not produce excessive kernels (tests recursive transitive processing)."""
|
|
D, N = 4, 5
|
|
caches = [Tensor.zeros(8, D).contiguous().realize() for _ in range(N)]
|
|
caches[0][0:1].assign(Tensor.ones(1, D, buffer=False) * 10)
|
|
x = caches[0][:1].sum(0, keepdim=True)
|
|
for i in range(1, N):
|
|
caches[i][0:1].assign(x)
|
|
x = caches[i][:1].sum(0, keepdim=True)
|
|
GlobalCounters.reset()
|
|
x.realize()
|
|
# N assigns (1 kernel each) producing N kernels total
|
|
assert_kernel_count(N)
|
|
|
|
def test_shared_computation_assign_kernel_count(self):
|
|
"""When a .contiguous() is shared between an assign value and the next layer's input (like QKV projection in LLM),
|
|
substitute optimization replaces already-realized sub-graphs in remaining pending assigns, preventing kernel escalation.
|
|
Without substitute, pending assign graphs grow linearly and produce 153 kernels instead of 48."""
|
|
D, N = 16, 16
|
|
caches = [Tensor.zeros(4, D).contiguous().realize() for _ in range(N)]
|
|
W = [Tensor.full((D, D*2), 0.01).contiguous().realize() for _ in range(N)]
|
|
x = Tensor.ones(1, D).contiguous().realize()
|
|
for i in range(N):
|
|
shared = (x @ W[i]).contiguous() # .contiguous() UOp is shared between assign (k) and next layer (q)
|
|
k, q = shared[:, :D], shared[:, D:]
|
|
caches[i][0:1].assign(k) # assign references the CONTIGUOUS
|
|
x = q + caches[i][:1] # next layer also references the same CONTIGUOUS through q
|
|
GlobalCounters.reset()
|
|
caches[-1][:1].contiguous().realize()
|
|
# N matmuls + N assigns + 1 final read = 2*N+1 (AFTER embedding allows full graph scheduling with shared contiguous reuse)
|
|
assert_kernel_count(2*N+1)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|