update_test_checks: add new --filter-all-after option

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 llvm#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-all-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-all-after so I've amended the
code to ensure we treat this filter differently.

Author: david-arm

llvm/test/tools/UpdateTestChecks/update_test_checks/Inputs/filter-all-after.ll

@@ -0,0 +1,17 @@
+; RUN: opt < %s -S | FileCheck %s
+
+define i32 @func({i32, i32} %x, i32 %y, i1 %cond) {
+entry:
+  br i1 %cond, label %b1, label %b2
+
+b1:
+  %x.idx0 = extractvalue {i32, i32} %x, 0
+  %add1 = add i32 %y, 1
+  %add2 = add i32 %x.idx0, %add1
+  %mul = mul i32 %add2, 3
+  br label %b2
+
+b2:
+  %res = phi i32 [ -1, %entry ], [ %mul, %b1 ]
+  ret i32 %res
+}

llvm/test/tools/UpdateTestChecks/update_test_checks/Inputs/filter-all-after.ll.expected

@@ -0,0 +1,30 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --filter-all-after "^b2:" --version 5
+; RUN: opt < %s -S | FileCheck %s
+
+define i32 @func({i32, i32} %x, i32 %y, i1 %cond) {
+; CHECK-LABEL: define i32 @func(
+; CHECK-SAME: { i32, i32 } [[X:%.*]], i32 [[Y:%.*]], i1 [[COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br i1 [[COND]], label %[[B1:.*]], label %[[B2:.*]]
+; CHECK:       [[B1]]:
+; CHECK-NEXT:    [[X_IDX0:%.*]] = extractvalue { i32, i32 } [[X]], 0
+; CHECK-NEXT:    [[ADD1:%.*]] = add i32 [[Y]], 1
+; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[X_IDX0]], [[ADD1]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[ADD2]], 3
+; CHECK-NEXT:    br label %[[B2]]
+; CHECK:       [[B2]]:
+;
+entry:
+  br i1 %cond, label %b1, label %b2
+
+b1:
+  %x.idx0 = extractvalue {i32, i32} %x, 0
+  %add1 = add i32 %y, 1
+  %add2 = add i32 %x.idx0, %add1
+  %mul = mul i32 %add2, 3
+  br label %b2
+
+b2:
+  %res = phi i32 [ -1, %entry ], [ %mul, %b1 ]
+  ret i32 %res
+}

llvm/test/tools/UpdateTestChecks/update_test_checks/Inputs/filter-all-after.ll.expected2

@@ -0,0 +1,28 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --filter-all-after "= mul i32" --filter-all-after "^b2:" --version 5
+; RUN: opt < %s -S | FileCheck %s
+
+define i32 @func({i32, i32} %x, i32 %y, i1 %cond) {
+; CHECK-LABEL: define i32 @func(
+; CHECK-SAME: { i32, i32 } [[X:%.*]], i32 [[Y:%.*]], i1 [[COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br i1 [[COND]], label %[[B1:.*]], [[B2:label %.*]]
+; CHECK:       [[B1]]:
+; CHECK-NEXT:    [[X_IDX0:%.*]] = extractvalue { i32, i32 } [[X]], 0
+; CHECK-NEXT:    [[ADD1:%.*]] = add i32 [[Y]], 1
+; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[X_IDX0]], [[ADD1]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[ADD2]], 3
+;
+entry:
+  br i1 %cond, label %b1, label %b2
+
+b1:
+  %x.idx0 = extractvalue {i32, i32} %x, 0
+  %add1 = add i32 %y, 1
+  %add2 = add i32 %x.idx0, %add1
+  %mul = mul i32 %add2, 3
+  br label %b2
+
+b2:
+  %res = phi i32 [ -1, %entry ], [ %mul, %b1 ]
+  ret i32 %res
+}

llvm/test/tools/UpdateTestChecks/update_test_checks/filter_all_after.test

@@ -0,0 +1,23 @@
+## Check that --filter-all-after works properly.
+# RUN: cp -f %S/Inputs/filter-all-after.ll %t.ll && %update_test_checks --version=5 --filter-all-after="^b2:" %t.ll
+# RUN: diff -u %t.ll %S/Inputs/filter-all-after.ll.expected
+
+## Check that running the script again does not change the result:
+# RUN: %update_test_checks --version=5 --filter-all-after="^b2:" %t.ll
+# RUN: diff -u %t.ll %S/Inputs/filter-all-after.ll.expected
+
+## Check that running the script again, without arguments, does not change the result:
+# RUN: %update_test_checks --version=5 %t.ll
+# RUN: diff -u %t.ll %S/Inputs/filter-all-after.ll.expected
+
+## Check that multiple --filter-all-after options work properly.
+# RUN: cp -f %S/Inputs/filter-all-after.ll %t.ll && %update_test_checks --version=5 --filter-all-after="= mul i32" --filter-all-after="^b2:" %t.ll
+# RUN: diff -u %t.ll %S/Inputs/filter-all-after.ll.expected2
+
+## Check that running the script again does not change the result:
+# RUN: %update_test_checks --version=5 --filter-all-after="= mul i32" --filter-all-after="^b2:" %t.ll
+# RUN: diff -u %t.ll %S/Inputs/filter-all-after.ll.expected2
+
+## Check that running the script again, without arguments, does not change the result:
+# RUN: %update_test_checks --version=5 %t.ll
+# RUN: diff -u %t.ll %S/Inputs/filter-all-after.ll.expected2

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_all_after):
         super(Filter, self).__init__(regex)
+        if is_filter_out and is_filter_all_after:
+            raise ValueError("cannot use both --filter-out and --filter-all-after")
         self.is_filter_out = is_filter_out
+        self.is_filter_all_after = is_filter_all_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_all_after = copy.deepcopy(self.is_filter_all_after, memo)
         return result
 
 
@@ -127,7 +131,11 @@ def __call__(self, parser, namespace, values, option_string=None):
 
             is_filter_out = option_string == "--filter-out"
 
-            value_list[-1] = Filter(value_list[-1].regex, is_filter_out)
+            is_filter_all_after = option_string == "--filter-all-after"
+
+            value_list[-1] = Filter(
+                value_list[-1].regex, is_filter_out, is_filter_all_after
+            )
 
             setattr(namespace, self.dest, value_list)
 
@@ -151,6 +159,13 @@ def __call__(self, parser, namespace, values, option_string=None):
         metavar="REGEX",
         help="Exclude lines matching REGEX",
     )
+    filter_group.add_argument(
+        "--filter-all-after",
+        action=FilterAction,
+        dest="filters",
+        metavar="REGEX",
+        help="Exclude all lines 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_all_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_all(filters):
+    for f in filters:
+        if f.is_filter_all_after:
+            return True
+    return False
+
+
+def filter_all_after(body, filters):
+    lines = []
+    for line in body.splitlines():
+        lines.append(line)
+        for f in filters:
+            if f.is_filter_all_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
+    has_filter_after = has_filter_all(filters)
+    lines = []
+    if has_filter_after:
+        lines = filter_all_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 @@ def __init__(self, run_list, flags, scrubber_args, path, ginfo):
             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_all_after,
                     ),
                     flags.filters,
                 )
@@ -831,7 +870,11 @@ def global_var_dict(self):
         return self._global_var_dict
 
     def is_filtered(self):
-        return bool(self._filters)
+        total_filters = 0
+        for f in self._filters:
+            if not f.is_filter_all_after:
+                total_filters += 1
+        return bool(total_filters)
 
     def process_run_line(self, function_re, scrubber, raw_tool_output, prefixes):
         build_global_values_dictionary(
@@ -2526,7 +2569,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_all_after:
+                    opt_name = "filter-all-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: