add test

Signed-off-by: Yaoyao Ding <dingyaoyao.cs@gmail.com>

Author: yaoyaoding

tests/instructions/test_reduce.py

@@ -15,7 +15,7 @@
 import pytest
 import tilus
 import torch
-from tilus import int32
+from tilus import boolean, int32
 from tilus.ir.layout import RegisterLayout, register_layout
 from tilus.ir.layout.ops import spatial
 
@@ -41,6 +41,19 @@ def __call__(self, out_ptr: ~int32):
         self.store_global(g_out, b, offsets=[0, 0], dims=[0, 1])
 
 
+class AnyAllInstExample(tilus.Script):
+    def __call__(self, x_ptr: ~int32, y_ptr: ~boolean):
+        self.attrs.blocks = 1
+        self.attrs.warps = 1
+
+        g_x = self.global_view(ptr=x_ptr, dtype=int32, shape=(32, 32))
+        g_y = self.global_view(ptr=y_ptr, dtype=boolean, shape=[2])
+        r_x = self.load_global(g_x, offsets=[0, 0], shape=[32, 32])
+
+        self.store_global(g_y, src=self.any(r_x != 0), offsets=[0], dims=[])
+        self.store_global(g_y, src=self.all(r_x != 0), offsets=[1], dims=[])
+
+
 @pytest.mark.parametrize("dim", [0, 1])
 @pytest.mark.parametrize(
     "layout",
@@ -69,3 +82,19 @@ def test_reduce_instruction(dim: int, layout: RegisterLayout):
     demo = ReduceKernelExample(layout, dim=dim)
     demo(actual)
     assert torch.allclose(actual, expected), f"Failed for layout {layout} and dim {dim}"
+
+
+def test_any_all_reduce_instruction():
+    kernel = AnyAllInstExample()
+    x0 = torch.zeros((32, 32), dtype=torch.int32).cuda()
+    y0 = torch.asarray([False, False], dtype=torch.bool).cuda()
+    x1 = torch.ones((32, 32), dtype=torch.int32).cuda()
+    y1 = torch.asarray([True, True], dtype=torch.bool).cuda()
+    x2 = torch.randint(0, 2, size=(32, 32), dtype=torch.int32).cuda()
+    x2[0, 0] = 1
+    x2[0, 1] = 0
+    y2 = torch.asarray([True, False], dtype=torch.bool).cuda()
+    for x, y in zip([x0, x1, x2], [y0, y1, y2]):
+        y_actual = torch.empty_like(y)
+        kernel(x, y_actual)
+        assert torch.allclose(y_actual, y), f"Failed for x={x} and y={y}, y_actual={y_actual}"