Add SequenceType Eq comparison (#482)

Author: oulgen

helion/_compiler/type_propagation.py

@@ -1654,6 +1654,39 @@ def _compare(self, op: ast.cmpop, left: TypeInfo, right: TypeInfo) -> TypeInfo:
                     raise
                 except Exception as e:
                     raise exc.TorchOpTracingError(e) from e
+        if (
+            isinstance(left, SequenceType)
+            and isinstance(right, SequenceType)
+            and isinstance(op, ast.Eq)
+        ):
+            if len(left.element_types) != len(right.element_types):
+                return LiteralType(origin=self.origin(), value=False)
+
+            can_determine_statically = True
+            all_elements_equal = True
+            for left_elem, right_elem in zip(
+                left.element_types, right.element_types, strict=False
+            ):
+                if isinstance(left_elem, LiteralType) and isinstance(
+                    right_elem, LiteralType
+                ):
+                    if left_elem.value != right_elem.value:
+                        all_elements_equal = False
+                        break
+                elif isinstance(left_elem, (NumericType, LiteralType)) and isinstance(
+                    right_elem, (NumericType, LiteralType)
+                ):
+                    if NumericType.known_equal(left_elem.value, right_elem.value):
+                        continue
+                    can_determine_statically = False
+                    break
+                else:
+                    can_determine_statically = False
+                    break
+
+            if can_determine_statically:
+                return LiteralType(origin=self.origin(), value=all_elements_equal)
+            return SymBoolType.new_unbacked(self.origin())
         raise exc.TypeInferenceError(
             f"{type(op).__name__} not supported on {left!s} and {right!s}"
         )

requirements.txt

@@ -2,3 +2,5 @@ pytest
 typing_extensions
 pre-commit
 filecheck
+expecttest
+numpy

test/test_misc.expected

@@ -91,6 +91,59 @@ def kernel_with_scalar_item(x: torch.Tensor, scalar_tensor: torch.Tensor, *, _la
     _launcher(_kernel_with_scalar_item_kernel, (triton.cdiv(x.size(0), _BLOCK_SIZE_0),), x, result, x.size(0), result.stride(0), x.stride(0), scalar_val, _BLOCK_SIZE_0, num_warps=4, num_stages=3)
     return result
 
+--- assertExpectedJournal(TestMisc.test_sequence_assert_static_shapes_False)
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from helion.runtime import default_launcher as _default_launcher
+
+@triton.jit
+def _kernel_kernel(a, out, a_size_0, a_stride_0, out_stride_0, _BLOCK_SIZE_0: tl.constexpr):
+    num_blocks_0 = tl.cdiv(a_size_0, _BLOCK_SIZE_0)
+    pid_0 = tl.program_id(0) % num_blocks_0
+    offset_0 = pid_0 * _BLOCK_SIZE_0
+    indices_0 = (offset_0 + tl.arange(0, _BLOCK_SIZE_0)).to(tl.int32)
+    mask_0 = indices_0 < a_size_0
+    load = tl.load(a + indices_0[:, None] * a_stride_0, mask_0[:, None], other=0)
+    load_1 = tl.load(a + indices_0[:, None] * a_stride_0, mask_0[:, None], other=0)
+    v_0 = load + load_1
+    tl.store(out + indices_0[:, None] * out_stride_0, v_0, mask_0[:, None])
+
+def kernel(a: torch.Tensor, b: torch.Tensor, *, _launcher=_default_launcher):
+    assert a.size() == b.size()
+    out = torch.empty_like(a)
+    _BLOCK_SIZE_0 = 16
+    _launcher(_kernel_kernel, (triton.cdiv(a.size(0), _BLOCK_SIZE_0) * 1,), a, out, a.size(0), a.stride(0), out.stride(0), _BLOCK_SIZE_0, num_warps=4, num_stages=3)
+    return out
+
+--- assertExpectedJournal(TestMisc.test_sequence_assert_static_shapes_True)
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from helion.runtime import default_launcher as _default_launcher
+
+@triton.jit
+def _kernel_kernel(a, b, out, _BLOCK_SIZE_0: tl.constexpr):
+    num_blocks_0 = tl.cdiv(16, _BLOCK_SIZE_0)
+    pid_0 = tl.program_id(0) % num_blocks_0
+    offset_0 = pid_0 * _BLOCK_SIZE_0
+    indices_0 = (offset_0 + tl.arange(0, _BLOCK_SIZE_0)).to(tl.int32)
+    load = tl.load(a + indices_0[:, None] * 1, None)
+    load_1 = tl.load(b + indices_0[:, None] * 1, None)
+    v_0 = load + load_1
+    tl.store(out + indices_0[:, None] * 1, v_0, None)
+
+def kernel(a: torch.Tensor, b: torch.Tensor, *, _launcher=_default_launcher):
+    assert a.size() == b.size()
+    out = torch.empty_like(a)
+    _BLOCK_SIZE_0 = 16
+    _launcher(_kernel_kernel, (triton.cdiv(16, _BLOCK_SIZE_0) * 1,), a, b, out, _BLOCK_SIZE_0, num_warps=4, num_stages=3)
+    return out
+
 --- assertExpectedJournal(TestMisc.test_tile_block_size_constexpr_fix)
 from __future__ import annotations
 

test/test_misc.py

@@ -7,6 +7,8 @@
 from packaging import version
 import pytest
 import torch
+from torch.testing._internal.common_utils import instantiate_parametrized_tests
+from torch.testing._internal.common_utils import parametrize
 
 import helion
 from helion._compat import supports_tensor_descriptor
@@ -415,6 +417,25 @@ def copy_kernel(a: torch.Tensor) -> torch.Tensor:
         torch.testing.assert_close(result, args[0])
         self.assertExpectedJournal(code)
 
+    @parametrize("static_shapes", (True, False))
+    def test_sequence_assert(self, static_shapes):
+        @helion.kernel(static_shapes=static_shapes)
+        def kernel(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+            assert a.size() == b.size()
+            out = torch.empty_like(a)
+
+            for tile in hl.tile(a.size()):
+                out[tile] = a[tile] + b[tile]
+            return out
+
+        a = torch.randn(16, 1, device=DEVICE)
+        code, result = code_and_output(kernel, (a, a))
+        torch.testing.assert_close(result, a + a)
+        self.assertExpectedJournal(code)
+
+
+instantiate_parametrized_tests(TestMisc)
+
 
 if __name__ == "__main__":
     unittest.main()