fix CSE (#414)

pointed out by @kaihsin we should actually use a key object with both
hash and eq implemented.

Author: Roger-luo

src/kirin/ir/attrs/abc.py

@@ -50,6 +50,9 @@ class Attribute(ABC, Printable, metaclass=AttributeMeta):
     @abstractmethod
     def __hash__(self) -> int: ...
 
+    @abstractmethod
+    def __eq__(self, value: object) -> bool: ...
+
     @classmethod
     def has_trait(cls, trait_type: type[Trait["Attribute"]]) -> bool:
         """Check if the Statement has a specific trait.

src/kirin/ir/attrs/py.py

@@ -34,14 +34,14 @@ def __init__(self, data: T, pytype: TypeAttribute | None = None):
         else:
             self.type = pytype
 
-    def __hash__(self):
-        # Fix hash(-1) == hash(-2) collision
-        # assume maximum is 8 bytes == 64 bits
-        if isinstance(self.data, int):
-            return hash(self.data.to_bytes(signed=True, byteorder="big", length=8))
-        elif isinstance(self.data, float):
-            return hash(self.data.hex())
-        return hash(self.data) + hash(self.type)
+    def __hash__(self) -> int:
+        return hash((self.type, self.data))
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, PyAttr):
+            return False
+
+        return self.type == value.type and self.data == value.data
 
     def print_impl(self, printer: Printer) -> None:
         printer.plain_print(repr(self.data))

src/kirin/ir/attrs/types.py

@@ -351,14 +351,20 @@ def print_impl(self, printer: Printer) -> None:
 
 
 @typing.final
-@dataclass(eq=False)
+@dataclass
 class Vararg(Attribute):
     name = "Vararg"
     typ: TypeAttribute
 
     def __hash__(self) -> int:
         return hash((Vararg, self.typ))
 
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, Vararg):
+            return False
+
+        return self.typ == value.typ
+
     def print_impl(self, printer: Printer) -> None:
         printer.plain_print("*")
         printer.print(self.typ)

src/kirin/rewrite/cse.py

@@ -1,40 +1,74 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
 
-from kirin.ir import Pure, Block, Statement
+from kirin import ir
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 
+@dataclass
+class Info:
+    """An object to hold the comparison information of a statement."""
+
+    head: type[ir.Statement]
+    args: tuple[ir.SSAValue, ...]
+    attributes: tuple[ir.Attribute, ...]
+    successors: tuple[ir.Block, ...]
+    regions: tuple[ir.Region, ...]
+
+    def __hash__(self) -> int:
+        return hash(
+            (id(self.head),)
+            + tuple(id(ssa) for ssa in self.args)
+            + tuple(hash(attr) for attr in self.attributes)
+            + tuple(id(succ) for succ in self.successors)
+            + tuple(id(region) for region in self.regions)
+        )
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Info):
+            return False
+
+        return (
+            self.head == other.head
+            and self.args == other.args
+            and self.attributes == other.attributes
+            and self.successors == other.successors
+            and self.regions == other.regions
+        )
+
+
 @dataclass
 class CommonSubexpressionElimination(RewriteRule):
 
-    def rewrite_Block(self, node: Block) -> RewriteResult:
-        seen: dict[int, Statement] = {}
+    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
+        seen: dict[Info, ir.Statement] = {}
 
         for stmt in node.stmts:
-            if not stmt.has_trait(Pure):
+            if not stmt.has_trait(ir.Pure):
                 continue
 
             if stmt.regions:
                 continue
 
-            hash_value = hash(
-                (type(stmt),)
-                + tuple(stmt.args)
-                + tuple(stmt.attributes.values())
-                + tuple(stmt.successors)
-                + tuple(stmt.regions)
+            info = Info(
+                head=type(stmt),
+                args=tuple(stmt.args),
+                attributes=tuple(stmt.attributes.values()),
+                successors=tuple(stmt.successors),
+                regions=tuple(stmt.regions),
             )
-            if hash_value in seen:
-                old_stmt = seen[hash_value]
+            if info in seen:
+                old_stmt = seen[info]
                 for result, old_result in zip(stmt._results, old_stmt.results):
                     result.replace_by(old_result)
                 stmt.delete()
                 return RewriteResult(has_done_something=True)
             else:
-                seen[hash_value] = stmt
+                seen[info] = stmt
         return RewriteResult()
 
-    def rewrite_Statement(self, node: Statement) -> RewriteResult:
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         if not node.regions:
             return RewriteResult()