mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 00:53:45 +08:00
Daniel Koepping
47 lines
2.1 KiB
Python
47 lines
2.1 KiB
Python
import unittest
|
|
from tinygrad.device import Device
|
|
from tinygrad.dtype import dtypes
|
|
from tinygrad.codegen import to_program
|
|
from tinygrad.helpers import Target
|
|
from tinygrad.renderer.cstyle import CStyleLanguage
|
|
from tinygrad.renderer.wgsl import WGSLRenderer
|
|
from test.helpers import check_schedule
|
|
from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
|
|
from tinygrad.tensor import Tensor
|
|
|
|
@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage), "uops are for cstyle")
|
|
class TestCStyleFailures(unittest.TestCase):
|
|
def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
|
|
dtype = "bool" if op in (Ops.OR, Ops.XOR, Ops.AND) else None
|
|
ret = Tensor.empty(1, dtype=dtype)
|
|
for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype))
|
|
linear, _ = check_schedule(ret, 1)
|
|
src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[2].arg
|
|
self.assertEqual("("*5 not in src, should_strip_paren)
|
|
|
|
def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
|
|
|
|
def test_repeat_mul(self): self._test_src_strip_paren(Ops.MUL)
|
|
|
|
def test_repeat_xor(self): self._test_src_strip_paren(Ops.XOR)
|
|
|
|
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "wgsl ends up with '(' * 5")
|
|
def test_repeat_or(self): self._test_src_strip_paren(Ops.OR)
|
|
|
|
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "wgsl ends up with '(' * 5")
|
|
def test_repeat_and(self): self._test_src_strip_paren(Ops.AND)
|
|
|
|
def test_repeat_sub(self): self._test_src_strip_paren(Ops.SUB, should_strip_paren=False)
|
|
|
|
class TestWGSLFailures(unittest.TestCase):
|
|
def test_folded_packed_store(self):
|
|
b = UOp.param(0, dtypes.char, 4)
|
|
word = b.index(UOp.const(0).cast(dtypes.int)).cast(dtypes.uint32)
|
|
store = word.store(word.load() & UOp.const(0xffffff00).cast(dtypes.uint32))
|
|
src = WGSLRenderer(Target("WEBGPU")).render(UOp.sink(store, arg=KernelInfo()).toposort())
|
|
self.assertIn("atomicAnd(&data0_4[0],4294967040u);", src)
|
|
self.assertNotIn("atomicAdd", src)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|