mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-02-18 18:35:12 -05:00
cac8bcf8b5
* decrease bert python time [pr]
* order copies
* Revert "order copies"
This reverts commit 3f62c8693b.
* rewrite count
* Ops.REDUCE
* acc first in the add chain
* Fix tensor core acc
* arange patterns look good
* fix multireduce gate
* reduce rewrite rule
* bump that to 15 minutes
* multiwmma isn't fusing
* gep through wmma is gep pushing
* bump that timeout too, it's all env setup
* add failing test
29 lines
936 B
Python
29 lines
936 B
Python
import unittest
|
|
from tinygrad import Tensor, Context, Device
|
|
from tinygrad.codegen.kernel import Kernel, Opt, OptOps
|
|
|
|
class TestLinearizerRewrite(unittest.TestCase):
|
|
def test_reduction(self):
|
|
t = Tensor.ones((64,64), device="NULL").contiguous().realize()
|
|
out = (t*2).sum(axis=1)
|
|
with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0):
|
|
si = out.schedule()[-1]
|
|
k = Kernel(si.ast, Device["CPU"].renderer)
|
|
k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
|
|
k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
|
|
prg = k.to_program()
|
|
print(prg.src)
|
|
|
|
def test_arange(self):
|
|
out = Tensor.arange(32, device="NULL")
|
|
with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0):
|
|
si = out.schedule()[-1]
|
|
k = Kernel(si.ast, Device["CPU"].renderer)
|
|
k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
|
|
k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
|
|
prg = k.to_program()
|
|
print(prg.src)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|