add failure nodes to some checks

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/cos_sin_cache.py

@@ -103,7 +103,7 @@ def check(
         # TODO(rama): handle redundant reshape/expand
         if self._const_freqs:
             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.")
+                return check_result.fail("freqs is not a constant or not 3D.", freqs)
             else:
                 return check_result
         if (
@@ -113,14 +113,14 @@ def check(
         ):
             pass
         else:
-            return check_result.fail("position_ids is not a 1D or 2D tensor.")
+            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 check_result.fail("inv_freq is not 3D.")
+            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 check_result.fail("inv_freq is not a constant.")
+            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].")
+            return check_result.fail("inv_freq is not of shape [1, ., 1].", inv_freq)
         return check_result
 
     def rewrite(

onnxscript/rewriter/ort_fusions/mha.py

@@ -172,36 +172,44 @@ def no_match(val: ir.Value, dims: Sequence[str]) -> bool:
 
         if no_match(query_BSD, ["B", "S", "D"]):
             return check_result.fail(
-                f"Shape mismatch: {query_BSD} does not match expected dimensions ['B', 'S', 'D']"
+                f"Shape mismatch: {query_BSD} does not match expected dimensions ['B', 'S', 'D']",
+                query_BSD,
             )
         if no_match(key_BSD, ["B", "Skv", "D"]):
             return check_result.fail(
-                f"Shape mismatch: {key_BSD} does not match expected dimensions ['B', 'Skv', 'D']"
+                f"Shape mismatch: {key_BSD} does not match expected dimensions ['B', 'Skv', 'D']",
+                query_BSD,
             )
         if no_match(value_BSD, ["B", "Skv", "D"]):
             return check_result.fail(
-                f"Shape mismatch: {value_BSD} does not match expected dimensions ['B', 'Skv', 'D']"
+                f"Shape mismatch: {value_BSD} does not match expected dimensions ['B', 'Skv', 'D']",
+                value_BSD,
             )
 
         if no_match(past_key, ["B", "H", "Spast", "Dh"]):
             return check_result.fail(
-                f"Shape mismatch: {past_key} does not match expected dimensions ['B', 'H', 'Spast', 'Dh']"
+                f"Shape mismatch: {past_key} does not match expected dimensions ['B', 'H', 'Spast', 'Dh']",
+                past_key,
             )
         if no_match(past_value, ["B", "H", "Spast", "Dv"]):
             return check_result.fail(
-                f"Shape mismatch: {past_value} does not match expected dimensions ['B', 'H', 'Spast', 'Dv']"
+                f"Shape mismatch: {past_value} does not match expected dimensions ['B', 'H', 'Spast', 'Dv']",
+                past_value,
             )
         if no_match(query_BSHDh, ["B", "S", "H", "Dh"]):
             return check_result.fail(
-                f"Shape mismatch: {query_BSHDh} does not match expected dimensions ['B', 'S', 'H', 'Dh']"
+                f"Shape mismatch: {query_BSHDh} does not match expected dimensions ['B', 'S', 'H', 'Dh']",
+                query_BSHDh,
             )
         if no_match(key_BSHDh, ["B", "S", "H", "Dh"]):
             return check_result.fail(
-                f"Shape mismatch: {key_BSHDh} does not match expected dimensions ['B', 'S', 'H', 'Dh']"
+                f"Shape mismatch: {key_BSHDh} does not match expected dimensions ['B', 'S', 'H', 'Dh']",
+                query_BSHDh,
             )
         if no_match(value_BSHDh, ["B", "S", "H", "Dh"]):
             return check_result.fail(
-                f"Shape mismatch: {value_BSHDh} does not match expected dimensions ['B', 'S', 'H', 'Dh']"
+                f"Shape mismatch: {value_BSHDh} does not match expected dimensions ['B', 'S', 'H', 'Dh']",
+                query_BSHDh,
             )
         # TODO: mask shape check: ideally, it should be (1 or B, 1 or H, S, St)
         # But this also, unforunately, depends on ORT version.

onnxscript/rewriter/ort_fusions/rms_normalization.py

@@ -58,12 +58,12 @@ def check(self, op, x, scale, epsilon, compute_dtype, target_dtype) -> pattern.M
         # epsilon must be a scalar
         epsilon_value = _ir_utils.get_singleton_value(epsilon)
         if not isinstance(epsilon_value, float):  # TODO: support other types
-            return check_result.fail("Epsilon is not a float value.")
+            return check_result.fail("Epsilon is not a float value.", epsilon)
         # input and output must be same dtype
         if x.dtype not in float_types:
-            return check_result.fail("Input is not a float type.")
+            return check_result.fail("Input is not a float type.", x)
         if scale.dtype not in float_types:
-            return check_result.fail("Scale is not a float type.")
+            return check_result.fail("Scale is not a float type.", scale)
         stash_dtype = compute_dtype.value if self._cast_input else x.dtype
         if stash_dtype not in fp_float_types:
             return check_result.fail("Normalization precision is not a float or double type.")

onnxscript/rewriter/ort_fusions/rotary_embedding.py

@@ -34,12 +34,12 @@ def check(self, op, x, start1, end1, start2, end2, **_) -> pattern.MatchResult:
         check_result = pattern.MatchResult()
         # x needs to be a 4D tensor with known last dimension size (== head_size) and known second dimension (num_heads)
         if x is None or x.shape is None or len(x.shape) != 4:
-            return check_result.fail("Input is not a 4D tensor.")
+            return check_result.fail("Input is not a 4D tensor.", x)
         if not isinstance(x.shape[1], int):
-            return check_result.fail("Input dimension 1 is not an integer.")
+            return check_result.fail("Input dimension 1 is not an integer.", x)
         head_size = x.shape[3]
         if not isinstance(head_size, int):
-            return check_result.fail("Head size is not an integer.")
+            return check_result.fail("Head size is not an integer.", x)
         half_head_size = head_size // 2
 
         # Check that x is being split into two equal halves of size half_head_size