mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 00:53:45 +08:00
Daniel Koepping
64 lines
2.7 KiB
Python
64 lines
2.7 KiB
Python
import unittest
|
|
from tinygrad import Tensor, UOp, dtypes
|
|
from tinygrad.helpers import Context
|
|
from tinygrad.uop.ops import Ops
|
|
from test.helpers import KernelCountException
|
|
from tinygrad.engine.realize import run_linear
|
|
|
|
class TestRingAllReduce(unittest.TestCase):
|
|
def test_hierarchy(self):
|
|
ds = tuple(f"CPU:{i}" for i in range(4))
|
|
with Context(ALL2ALL=1, ALLREDUCE_NODE_NDEVS=2): # two nodes of 2
|
|
for size in (1, 17):
|
|
x = (Tensor.arange(4 * size, dtype=dtypes.int32).reshape(4, size) % 13).realize()
|
|
self.assertEqual(x.shard(ds, axis=0).sum(0).tolist(), x.sum(0).tolist())
|
|
|
|
def test_schedule_all2all(self):
|
|
with Context(ALL2ALL=2):
|
|
N = 4
|
|
M = N*100
|
|
ds = tuple(f"CPU:{i}" for i in range(N))
|
|
x = Tensor.arange(N*M, dtype=dtypes.float).reshape(N, M)
|
|
t = (x*x).clone().shard(ds, axis=0).realize()
|
|
out = t.sum(0).mul(2.).contiguous()
|
|
linear, var_vals = out.linear_with_vars()
|
|
copies = [si for si in linear.src if si.src[0].op is Ops.STORE]
|
|
sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
|
|
# N*(N-1) copies for input and output
|
|
copy_count = N*(N-1)*2
|
|
if len(copies) != copy_count: raise KernelCountException(copy_count, len(copies))
|
|
# N*(N-1) shrinks from other devices becoming contigs, N ALU, N extra contig, reassembly (cat), and mul
|
|
sink_count = (N*(N-1))+(N)+(N)+(1)+(1)
|
|
if len(sinks) != sink_count: raise KernelCountException(sink_count, len(sinks))
|
|
# correctness
|
|
run_linear(linear, var_vals)
|
|
expected = [2*sum((d*M+i)**2 for d in range(N)) for i in range(M)]
|
|
dev_nums = Tensor.arange(1, N+1, dtype=dtypes.float).reshape(N, 1).expand(N, M).shard(ds, axis=0)
|
|
shards = out.reshape(1, M).expand(N, M)+dev_nums
|
|
self.assertListEqual(shards.tolist(), [[x+d+1 for x in expected] for d in range(N)])
|
|
|
|
def test_symbolic_shape(self):
|
|
rows = UOp.variable("rows", 1, 4).bind(3)
|
|
t = Tensor.ones(4, 4).shard(("CPU:0", "CPU:1"), axis=1).realize()
|
|
out = t[:rows].sum(1).realize()
|
|
self.assertEqual(out.shape, (rows,))
|
|
self.assertTrue((out == 4).all().item())
|
|
|
|
def test_symbolic_shape_2d(self):
|
|
rows = UOp.variable("rows", 1, 4).bind(3)
|
|
t = Tensor.ones(4, 4, 2).shard(("CPU:0", "CPU:1"), axis=1).realize()
|
|
out = t[:rows].sum(1).realize()
|
|
self.assertEqual(out.shape, (rows, 2))
|
|
self.assertTrue((out == 4).all().item())
|
|
|
|
def test_correct_ring(self):
|
|
with Context(RING=2):
|
|
N = 4
|
|
ds = tuple(f"CPU:{i}" for i in range(N))
|
|
t = Tensor.ones(N, N*100).contiguous().shard(ds, axis=0).realize()
|
|
out = t.sum(0)
|
|
self.assertListEqual(out.tolist(), [4]*N*100)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|