fold stencil.cast op with equal input-output types

xdsl/dialects/stencil.py

@@ -690,6 +690,16 @@ class AllocOp(IRDLOperation):
     traits = traits_def(AllocOpEffect())
 
 
+class CastOpHasCanonicalizationPatternsTrait(HasCanonicalizationPatternsTrait):
+    @classmethod
+    def get_canonicalization_patterns(cls) -> tuple[RewritePattern, ...]:
+        from xdsl.transforms.canonicalization_patterns.stencil import (
+            RemoveCastWithNoEffect,
+        )
+
+        return (RemoveCastWithNoEffect(),)
+
+
 @irdl_op_definition
 class CastOp(IRDLOperation):
     """
@@ -722,7 +732,7 @@ class CastOp(IRDLOperation):
         "$field attr-dict-with-keyword `:` type($field) `->` type($result)"
     )
 
-    traits = traits_def(NoMemoryEffect())
+    traits = traits_def(NoMemoryEffect(), CastOpHasCanonicalizationPatternsTrait())
 
     @staticmethod
     def get(

xdsl/transforms/canonicalization_patterns/stencil.py

@@ -109,3 +109,14 @@ def match_and_rewrite(self, op: stencil.ApplyOp, rewriter: PatternRewriter) -> N
 
         rewriter.replace_op(old_return, stencil.ReturnOp.get(return_args))
         rewriter.replace_matched_op(new, replace_results)
+
+
+class RemoveCastWithNoEffect(RewritePattern):
+    """
+    Remove `stencil.cast` where input and output types are equal.
+    """
+
+    @op_type_rewrite_pattern
+    def match_and_rewrite(self, op: stencil.CastOp, rewriter: PatternRewriter) -> None:
+        if op.result.type == op.field.type:
+            rewriter.replace_matched_op([], new_results=[op.field])