mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
Daniel Koepping
31 lines
1.5 KiB
Python
31 lines
1.5 KiB
Python
# PTX is shared by the NV and CUDA backends; both CI jobs run this suite.
|
|
import unittest
|
|
import numpy as np
|
|
from tinygrad import Tensor, Device, dtypes, UOp
|
|
from tinygrad.renderer.ptx import PTXRenderer
|
|
from tinygrad.uop.ops import Ops, KernelInfo
|
|
from test.runtime.test_renderer_failures import _test_uop_result
|
|
|
|
@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "tests for ptx renderer")
|
|
class TestPTXFailures(unittest.TestCase):
|
|
@unittest.skip("INDEX can only have a gate ALU parent, not an IF")
|
|
def test_gated_store_with_if(self):
|
|
a = UOp.param(0, dtypes.int, 4)
|
|
gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
|
|
val = UOp.const(1).cast(dtypes.int)
|
|
if_uop = UOp(Ops.IF, src=(gate_alu,))
|
|
gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0, if_uop), val))
|
|
sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
|
|
ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
|
|
np.testing.assert_equal(ret, [0, 1, 1, 1])
|
|
|
|
@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
|
|
def test_gated_define_acc_with_half_dtype(self):
|
|
a = Tensor.randn(32, 32, dtype=dtypes.half).realize()
|
|
b = Tensor.randn(34, 32, dtype=dtypes.half).realize()
|
|
result = a.pad((1,1)).matmul(b, dtype=dtypes.half).numpy()
|
|
reference = a.pad((1,1)).matmul(b, dtype=dtypes.float).numpy()
|
|
np.testing.assert_allclose(result, reference, atol=1e-2, rtol=1e-2)
|
|
|
|
if __name__ == "__main__": unittest.main()
|