mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-04-07 03:00:26 -04:00
b12b60af20
* fix binop, other tests failure * that was a bad idea * better layernorm * inference kernel count tests * new style reshape pushing * fixup replacement * 199 kernels is okay. fix flops * push reshape through unaryops only * GRAPH=2 draws the phantom ops * found resnet issue * non working test * mul is cheaper than div * OPT inflation * SHUFFLE_PAD_OPS in OPT=2
49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
#!/usr/bin/env python
|
|
import unittest
|
|
from typing import NamedTuple, Tuple
|
|
from tinygrad.ops import LazyOp, BinaryOps, ReduceOps, get_lazyop_info
|
|
from tinygrad.helpers import DType, dtypes
|
|
|
|
class TestBuffer(NamedTuple):
|
|
shape: Tuple[int, ...]
|
|
dtype: DType
|
|
|
|
class TestFlopCounter(unittest.TestCase):
|
|
def setUp(self):
|
|
self.buf0 = TestBuffer(shape=(4,), dtype=dtypes.float32)
|
|
self.buf1 = TestBuffer(shape=(4,), dtype=dtypes.float32)
|
|
|
|
def test_flops_add(self):
|
|
op0 = LazyOp(BinaryOps.ADD, (self.buf0,self.buf1,), None)
|
|
info = get_lazyop_info(op0)
|
|
self.assertEqual(info.flops, 4)
|
|
|
|
def test_flops_add_twice(self):
|
|
op0 = LazyOp(BinaryOps.ADD, (self.buf0,self.buf1,), None)
|
|
op1 = LazyOp(BinaryOps.ADD, (op0,self.buf1,), None)
|
|
info = get_lazyop_info(op1)
|
|
self.assertEqual(info.flops, 8)
|
|
|
|
def test_flops_add_self(self):
|
|
op0 = LazyOp(BinaryOps.ADD, (self.buf0,self.buf1,), None)
|
|
op1 = LazyOp(BinaryOps.ADD, (op0,op0,), None)
|
|
info = get_lazyop_info(op1)
|
|
self.assertEqual(info.flops, 8)
|
|
|
|
def test_flops_add_roundabout_self(self):
|
|
op0 = LazyOp(BinaryOps.ADD, (self.buf0,self.buf1,), None)
|
|
op1 = LazyOp(BinaryOps.ADD, (op0,self.buf1,), None)
|
|
op2 = LazyOp(BinaryOps.ADD, (op0,op1,), None)
|
|
info = get_lazyop_info(op2)
|
|
self.assertEqual(info.flops, 12)
|
|
|
|
def test_flops_red(self):
|
|
op0 = LazyOp(BinaryOps.MUL, (self.buf0,self.buf1,), None)
|
|
op1 = LazyOp(ReduceOps.SUM, (op0,), (1,))
|
|
op2 = LazyOp(BinaryOps.ADD, (op1, op1,), None)
|
|
info = get_lazyop_info(op2)
|
|
self.assertEqual(info.flops, 9)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|