Make the return type of rewrite check functions a MatchResult object (#2138)

- Check function returns a MatchResult object instead of bool
- This allows propagating the failure reason to the tracer to help in
debugging

Author: shubhambhokare1

onnxscript/rewriter/llama_rule_sets.py

@@ -26,9 +26,12 @@ def pattern(self, op, x):
     def rewrite(self, op, x: ir.Value):
         return op.Identity(x)
 
-    def check(self, context, x) -> bool:
+    def check(self, context, x) -> orp.MatchResult:
         del context  # Unused
-        return ir_utils.has_rank(x, 1)
+        check_result = orp.MatchResult()
+        if not ir_utils.has_rank(x, 1):
+            return check_result.fail("Input is not 1D")
+        return check_result
 
 
 class CastIdentity(orp.RewriteRuleAsClass):
@@ -43,8 +46,11 @@ def rewrite(cls, op, x: ir.Value, to: ir.Attr):
         return op.Identity(x)
 
     @classmethod
-    def check(cls, context, x, to) -> bool:
-        return x.dtype == to.value
+    def check(cls, context, x, to) -> orp.MatchResult:
+        check_result = orp.MatchResult()
+        if x.dtype != to.value:
+            return check_result.fail("Input and output types are not the same")
+        return check_result
 
 
 class CastCast(orp.RewriteRuleAsClass):
@@ -62,11 +68,13 @@ def pattern(cls, op, x, to, to_ignored):
         return op.Cast(op.Cast(x, to=to_ignored), to=to)
 
     @classmethod
-    def check(cls, context, x: ir.Value, to: ir.Attr, to_ignored: ir.Attr) -> bool:
-        return (
-            to.value in cls._allowed_tensor_types
-            and to_ignored.value in cls._allowed_tensor_types
-        )
+    def check(cls, context, x: ir.Value, to: ir.Attr, to_ignored: ir.Attr) -> orp.MatchResult:
+        check_result = orp.MatchResult()
+        if to.value not in cls._allowed_tensor_types:
+            return check_result.fail(f"Output type {to.value} is not allowed")
+        if to_ignored.value not in cls._allowed_tensor_types:
+            return check_result.fail(f"Ignored type {to_ignored.value} is not allowed")
+        return check_result
 
     @classmethod
     def rewrite(cls, op, x: ir.Value, to: ir.Attr, to_ignored: ir.Attr):
@@ -85,14 +93,19 @@ def rewrite(cls, op, x: ir.Value, shape: ir.Value):
         return op.Identity(x)
 
     @classmethod
-    def check(cls, context, x, shape) -> bool:
+    def check(cls, context, x, shape) -> orp.MatchResult:
+        check_result = orp.MatchResult()
         if shape.const_value is None:
             # Shape is not a constant and cannot be guessed.
-            return False
+            return check_result.fail("Shape is not a constant and cannot be guessed.")
         if (x_shape := x.shape) is None:
             # We don't know the shape of the input
-            return False
-        return x_shape.dims == tuple(shape.const_value.numpy().tolist())
+            return check_result.fail("Input shape is not known.")
+        if x_shape.dims != tuple(shape.const_value.numpy().tolist()):
+            return check_result.fail(
+                f"Input shape {x_shape.dims} does not match the shape {shape.const_value.numpy().tolist()}."
+            )
+        return check_result
 
 
 class ReshapeReshape(orp.RewriteRuleAsClass):
@@ -110,12 +123,15 @@ def rewrite(cls, op, x: ir.Value, shape_ignored: ir.Value, shape: ir.Value):
         return op.Reshape(x, shape)
 
     @classmethod
-    def check(cls, context, x, shape_ignored, shape) -> bool:
-        if shape_ignored.const_value is None or shape.const_value is None:
-            return False
+    def check(cls, context, x, shape_ignored, shape) -> orp.MatchResult:
+        check_result = orp.MatchResult()
+        if shape_ignored.const_value is None:
+            return check_result.fail("Shape ignored is not a constant.")
+        if shape.const_value is None:
+            return check_result.fail("Shape is not a constant.")
         if shape.const_value.numpy().min() <= 0:
-            return False
-        return True
+            return check_result.fail("Shape has non-positive values.")
+        return check_result
 
 
 class SlicesSplit(orp.RewriteRuleAsClass):
@@ -128,49 +144,50 @@ def pattern(cls, op, x, begin0, end0, axes0, begin1, end1, axes1):
         return op.Slice(x, begin0, end0, axes0), op.Slice(x, begin1, end1, axes1)
 
     @classmethod
-    def check(cls, context, x, begin0, end0, axes0, begin1, end1, axes1) -> bool:
+    def check(cls, context, x, begin0, end0, axes0, begin1, end1, axes1) -> orp.MatchResult:
+        check_result = orp.MatchResult()
         if (
             axes0.const_value is None
             or axes1.const_value is None
             or axes0.const_value.numpy().tolist() != axes1.const_value.numpy().tolist()
         ):
-            return False
+            return check_result.fail("Axes are not equal or not constant.")
         axes = axes0.const_value.numpy().tolist()
         if len(axes) != 1:
-            return False
+            return check_result.fail("Axes has more than one dimension.")
         if x.shape:
             rk = len(x.shape)
         else:
             rk = x.rank
         if axes[0] != -1 and axes[0] != rk - 1:
-            return False
+            return check_result.fail("Axes is not -1 or last dimension.")
         if (
             begin0.const_value is None
             or end0.const_value is None
             or begin1.const_value is None
             or end1.const_value is None
         ):
-            return False
+            return check_result.fail("Begin or end are not constant values.")
         if begin0.const_value.numpy().tolist() != [0]:
-            return False
+            return check_result.fail("First begin value is not 0.")
         e0, b1, e1 = (
             end0.const_value.numpy().tolist(),
             begin1.const_value.numpy().tolist(),
             end1.const_value.numpy().tolist(),
         )
         if e0[0] != b1[0]:
-            return False
+            return check_result.fail("End0 is not equal to Begin1.")
         shape = x.shape
         if shape is None:
-            return False
+            return check_result.fail("Shape is not known.")
         last_dim = shape[-1]
         if not isinstance(last_dim, int):
-            return False
+            return check_result.fail("Last dimension is not known.")
         if last_dim != e1[0]:
-            return False
+            return check_result.fail("Last dimension is not equal to End1.")
         if last_dim // 2 != b1[0]:
-            return False
-        return True
+            return check_result.fail("Last dimension is not equal to Begin1.")
+        return check_result
 
     @classmethod
     def rewrite(cls, op, x, begin0, end0, axes0, begin1, end1, axes1):
@@ -187,13 +204,14 @@ def pattern(cls, op, x, perm):
         return op.Transpose(x, perm=perm)
 
     @classmethod
-    def check(cls, context, x: ir.Value, perm: ir.Attr) -> bool:
+    def check(cls, context, x: ir.Value, perm: ir.Attr) -> orp.MatchResult:
+        check_result = orp.MatchResult()
         if isinstance(perm, ir.RefAttr):
-            return False
+            return check_result.fail("Permutation is a reference attribute.")
         if perm.type == ir.AttributeType.INTS:
             if perm.value == list(range(len(perm.value))):
-                return True
-        return False
+                return check_result
+        return check_result.fail("Permutation is not identity.")
 
     @classmethod
     def rewrite(cls, op, x: ir.Value, perm: ir.Attr):
@@ -210,10 +228,11 @@ def pattern(cls, op, x, perm1, perm2):
         return op.Transpose(op.Transpose(x, perm=perm1), perm=perm2)
 
     @classmethod
-    def check(cls, context, x: ir.Value, perm1: ir.Attr, perm2: ir.Attr) -> bool:
+    def check(cls, context, x: ir.Value, perm1: ir.Attr, perm2: ir.Attr) -> orp.MatchResult:
+        check_result = orp.MatchResult()
         if isinstance(perm1, ir.RefAttr) or isinstance(perm2, ir.RefAttr):
-            return False
-        return True
+            return check_result.fail("Permutation is a reference attribute.")
+        return check_result
 
     @classmethod
     def _apply_transpose(cls, perm: tuple[int, ...], on: list[int]) -> list[int]:
@@ -257,17 +276,18 @@ def rewrite(cls, op, x: ir.Value, axes1: ir.Value, axes2: ir.Value):
         return op.Unsqueeze(x, op.Constant(value=ir.tensor(axes, dtype=ir.DataType.INT64)))
 
     @classmethod
-    def check(cls, context, x, axes1, axes2) -> bool:
+    def check(cls, context, x, axes1, axes2) -> orp.MatchResult:
+        check_result = orp.MatchResult()
         del context  # Unused
         del x  # Unused
         # Currently restricted to single element positive axis
         v1 = ir_utils.get_singleton_value(axes1)
         v2 = ir_utils.get_singleton_value(axes2)
         if v1 is None or v2 is None:
-            return False
+            return check_result.fail("Axes are not constant.")
         if (v1 < 0) or (v2 < 0):
-            return False
-        return True
+            return check_result.fail("Axes are negative.")
+        return check_result
 
 
 cast_cast_rule = orp.make_rewrite_rule_from_class(CastCast)

onnxscript/rewriter/ort_fusions/cos_sin_cache.py

@@ -96,24 +96,32 @@ def pattern(
             _domain="ai.onnxruntime.fusion",
         )
 
-    def check(self, context, inv_freq, position_ids, freqs, extra_dims, **_):
+    def check(
+        self, context, inv_freq, position_ids, freqs, extra_dims, **_
+    ) -> pattern.MatchResult:  # type: ignore[name-defined]
+        check_result = pattern.MatchResult()
         # TODO(rama): handle redundant reshape/expand
         if self._const_freqs:
-            return (freqs.const_value is not None) and _ir_utils.has_rank(freqs, 3)
+            if (freqs.const_value is None) or not _ir_utils.has_rank(freqs, 3):
+                return check_result.fail("freqs is not a constant or not 3D.", freqs)
+            else:
+                return check_result
         if (
             _ir_utils.has_rank(position_ids, 2) and _ir_utils.is_singleton_value(extra_dims, 1)
         ) or (
             _ir_utils.has_rank(position_ids, 1) and _ir_utils.is_1d_value(extra_dims, [0, 1])
         ):
             pass
         else:
-            return False
+            return check_result.fail("position_ids is not a 1D or 2D tensor.", position_ids)
         if not _ir_utils.has_rank(inv_freq, 3):
-            return False
+            return check_result.fail("inv_freq is not 3D.", inv_freq)
         inv_freq_shape = inv_freq.shape
         if inv_freq.const_value is None:  # TODO: should this be inv_freq_shape?
-            return False
-        return inv_freq_shape[0] == 1 and inv_freq_shape[2] == 1
+            return check_result.fail("inv_freq is not a constant.", inv_freq)
+        if inv_freq_shape[0] != 1 or inv_freq_shape[2] != 1:
+            return check_result.fail("inv_freq is not of shape [1, ., 1].", inv_freq)
+        return check_result
 
     def rewrite(
         self, op, x, inv_freq, position_ids, interleaved, num_heads, freqs, dtype, **_

onnxscript/rewriter/ort_fusions/fused_matmul_rule_sets.py

@@ -15,13 +15,14 @@ def pattern(cls, op, x, y, cst):
         return op.Div(op.MatMul(x, y), cst)
 
     @classmethod
-    def check(cls, context, x, y, cst) -> bool:
+    def check(cls, context, x, y, cst) -> orp.MatchResult:
+        check_result = orp.MatchResult()
         if cst.const_value is None:
-            return False
+            return check_result.fail("Divisor is not a constant value.")
         value = cst.const_value.numpy()
         if value.size > 1:
-            return False
-        return True
+            return check_result.fail("Divisor is not a scalar value.")
+        return check_result
 
     @classmethod
     def rewrite(cls, op, x, y, cst):
@@ -38,12 +39,13 @@ def pattern(cls, op, x, y, cst):
         return op.Div(op.FusedMatMul(x, y, _domain="com.microsoft"), cst)
 
     @classmethod
-    def check(cls, context, x, y, cst) -> bool:
+    def check(cls, context, x, y, cst) -> orp.MatchResult:
+        check_result = orp.MatchResult()
         if cst.const_value is None:
-            return False
+            return check_result.fail("Divisor is not a constant value.")
         if cst.const_value.numpy().size > 1:
-            return False
-        return True
+            return check_result.fail("Divisor is not a scalar value.")
+        return check_result
 
     @classmethod
     def rewrite(cls, op, x, y, cst):
@@ -65,11 +67,14 @@ class _TransposeMatMulBase(orp.RewriteRuleAsClass):
     _pos: ClassVar = 1
 
     @classmethod
-    def check(cls, context, x, y) -> bool:
+    def check(cls, context, x, y) -> orp.MatchResult:
+        check_result = orp.MatchResult()
         perm = list((x if cls._pos == 1 else y).uses())[0][0].attributes["perm"].value  # noqa: RUF015
         expected_perm = list(range(len(perm)))
         expected_perm[-2], expected_perm[-1] = expected_perm[-1], expected_perm[-2]
-        return perm == expected_perm
+        if perm != expected_perm:
+            return check_result.fail("Permutation values for Transpose are not correct.")
+        return check_result
 
     @classmethod
     def rewrite(cls, op, x, y):
@@ -126,13 +131,16 @@ def pattern(cls, op, x, y):
         return op.Transpose(op.MatMul(x, y))
 
     @classmethod
-    def check(cls, context, x, y) -> bool:
+    def check(cls, context, x, y) -> orp.MatchResult:
+        check_result = orp.MatchResult()
         matmul = list(x.uses())[0][0]  # noqa: RUF015
         transpose = list(matmul.outputs[0].uses())[0][0]  # noqa: RUF015
         perm = transpose.attributes["perm"].value
         expected_perm = list(range(len(perm)))
         expected_perm[-2], expected_perm[-1] = expected_perm[-1], expected_perm[-2]
-        return perm == expected_perm
+        if perm != expected_perm:
+            return check_result.fail("Permutation values for Transpose are not correct.")
+        return check_result
 
     @classmethod
     def rewrite(cls, op, x, y):

onnxscript/rewriter/ort_fusions/gqa.py

@@ -94,7 +94,8 @@ def check(
         # key_transposed,
         # attention_reshaped,
         **_,
-    ):
+    ) -> pattern.MatchResult:  # type: ignore[name-defined]
+        check_result = pattern.MatchResult()
         # bindings: dict[str, int] = {}
         # status = (
         #     _check_shape(bindings, query_mm_reshaped, ["B", "S", "H", "d_h"])
@@ -110,7 +111,7 @@ def check(
         #     return False
         # if bindings["H"] * bindings["d_h"] != bindings["H*d_h"]:
         #     return False
-        return True
+        return check_result
 
     def rewrite(
         self,