update_test_checks: add new --filter-out-after option (#129739)

Whilst trying to clean up some loop vectoriser IR tests (see
test/Transforms/LoopVectorize/AArch64/partial-reduce-chained.ll
for example) a reviewer on PR #129047 suggested it would be
nice to have an option to stop generating CHECK lines after a
certain point. Typically when performing a transformation with
the loop vectoriser we don't usually care about any CHECK lines
generated for the scalar tail of the loop, since the scalar
loop is kept intact. Previously if you wanted to eliminate such
unwanted CHECK lines you had to run the update script, then
manually delete all the lines corresponding to the scalar loop.
This can be very time consuming if the tests ever need changing.

What I've tried to do here is add a new --filter-out-after
option alongside the existing --filter* options that provides
support for stopping the generation of any CHECK lines beyond
the line that matches the filter. With the existing filter
options we never generate CHECK-NEXT lines, but we still care
about ordering with --filter-out-after so I've amended the
code to ensure we treat this filter differently.

1 files changed, 59 insertions, 12 deletions

diff --git a/llvm/utils/UpdateTestChecks/common.py b/llvm/utils/UpdateTestChecks/common.py
index 0700486..274614c 100644
--- a/llvm/utils/UpdateTestChecks/common.py
+++ b/llvm/utils/UpdateTestChecks/common.py
@@ -74,13 +74,17 @@ class Filter(Regex):
 
     """
 
-    def __init__(self, regex, is_filter_out):
+    def __init__(self, regex, is_filter_out, is_filter_out_after):
         super(Filter, self).__init__(regex)
+        if is_filter_out and is_filter_out_after:
+            raise ValueError("cannot use both --filter-out and --filter-out-after")
         self.is_filter_out = is_filter_out
+        self.is_filter_out_after = is_filter_out_after
 
     def __deepcopy__(self, memo):
         result = copy.deepcopy(super(Filter, self), memo)
         result.is_filter_out = copy.deepcopy(self.is_filter_out, memo)
+        result.is_filter_out_after = copy.deepcopy(self.is_filter_out_after, memo)
         return result
 
 
@@ -127,7 +131,11 @@ def parse_commandline_args(parser):
 
             is_filter_out = option_string == "--filter-out"
 
-            value_list[-1] = Filter(value_list[-1].regex, is_filter_out)
+            is_filter_out_after = option_string == "--filter-out-after"
+
+            value_list[-1] = Filter(
+                value_list[-1].regex, is_filter_out, is_filter_out_after
+            )
 
             setattr(namespace, self.dest, value_list)
 
@@ -151,6 +159,13 @@ def parse_commandline_args(parser):
         metavar="REGEX",
         help="Exclude lines matching REGEX",
     )
+    filter_group.add_argument(
+        "--filter-out-after",
+        action=FilterAction,
+        dest="filters",
+        metavar="REGEX",
+        help="Exclude all lines within a given function after line matching REGEX",
+    )
 
     parser.add_argument(
         "--include-generated-funcs",
@@ -671,6 +686,8 @@ def get_globals_name_prefix(raw_tool_output):
 def apply_filters(line, filters):
     has_filter = False
     for f in filters:
+        if f.is_filter_out_after:
+            continue
         if not f.is_filter_out:
             has_filter = True
         if f.search(line):
@@ -680,14 +697,34 @@ def apply_filters(line, filters):
     return False if has_filter else True
 
 
+def has_filter_out_after(filters):
+    for f in filters:
+        if f.is_filter_out_after:
+            return True
+    return False
+
+
+def filter_out_after(body, filters):
+    lines = []
+    for line in body.splitlines():
+        lines.append(line)
+        for f in filters:
+            if f.is_filter_out_after:
+                if f.search(line):
+                    return lines
+    return lines
+
+
 def do_filter(body, filters):
-    return (
-        body
-        if not filters
-        else "\n".join(
-            filter(lambda line: apply_filters(line, filters), body.splitlines())
-        )
-    )
+    if not filters:
+        return body
+    filter_out_after_flag = has_filter_out_after(filters)
+    lines = []
+    if filter_out_after_flag:
+        lines = filter_out_after(body, filters)
+    else:
+        lines = body.splitlines()
+    return "\n".join(filter(lambda line: apply_filters(line, filters), lines))
 
 
 def scrub_body(body):
@@ -788,7 +825,9 @@ class FunctionTestBuilder:
             list(
                 map(
                     lambda f: Filter(
-                        re.compile(f.pattern().strip('"'), f.flags()), f.is_filter_out
+                        re.compile(f.pattern().strip('"'), f.flags()),
+                        f.is_filter_out,
+                        f.is_filter_out_after,
                     ),
                     flags.filters,
                 )
@@ -831,7 +870,10 @@ class FunctionTestBuilder:
         return self._global_var_dict
 
     def is_filtered(self):
-        return bool(self._filters)
+        for f in self._filters:
+            if not f.is_filter_out_after:
+                return True
+        return False
 
     def process_run_line(self, function_re, scrubber, raw_tool_output, prefixes):
         build_global_values_dictionary(
@@ -2526,7 +2568,12 @@ def get_autogennote_suffix(parser, args):
             # Create a separate option for each filter element.  The value is a list
             # of Filter objects.
             for elem in value:
-                opt_name = "filter-out" if elem.is_filter_out else "filter"
+                if elem.is_filter_out:
+                    opt_name = "filter-out"
+                elif elem.is_filter_out_after:
+                    opt_name = "filter-out-after"
+                else:
+                    opt_name = "filter"
                 opt_value = elem.pattern()
                 new_arg = '--%s "%s" ' % (opt_name, opt_value.strip('"'))
                 if new_arg not in autogenerated_note_args: