llvm · vporpo · Apr 14, 2025
diff --git a/llvm/test/tools/UpdateTestChecks/update_llc_test_checks/Inputs/downstream_target.ll b/llvm/test/tools/UpdateTestChecks/update_llc_test_checks/Inputs/downstream_target.ll
@@ -0,0 +1,4 @@
+; RUN: llc < %s -mtriple=i686-- | FileCheck %s
+define void @foo() {
+  ret void
+}
diff --git a/llvm/test/tools/UpdateTestChecks/update_llc_test_checks/Inputs/downstream_target.ll.expected b/llvm/test/tools/UpdateTestChecks/update_llc_test_checks/Inputs/downstream_target.ll.expected
@@ -0,0 +1,8 @@
+; RUN: llc < %s -mtriple=i686-- | FileCheck %s
+define void @foo() {
+; CHECK-LABEL: foo:
+; CHECK:       Test
+; CHECK-NEXT:  Downstream
+; CHECK-NEXT:  Target
+  ret void
+}
diff --git a/llvm/test/tools/UpdateTestChecks/update_llc_test_checks/Inputs/downstream_target.py b/llvm/test/tools/UpdateTestChecks/update_llc_test_checks/Inputs/downstream_target.py
@@ -0,0 +1,65 @@
+# This is an example of a downstream module, which in this case is for i686.
+# The goal is to test the updater's downstream-target-module-path option.
+#
+# This target module implements a simple scrubber that returns a fixed string,
+# a regex that matches the fixed strings generated by the scrubber.
+#
+# This module needs to define:
+# 2. The `downstream_scrub(asm, arg)` scrubber function
+# 3. The `downstream_asm_fn_re` regex for matching func, body, etc.
+
+import re
+import sys
+import os
+from pathlib import Path
+
+common_dir = Path(__file__).parent / "../../../../../utils/"
+sys.path.append(common_dir)
+from UpdateTestChecks import common
+
+
+# The asm scrubber returns the scrubbed asm.
+# In this test we return a fixed string.
+def scrub(asm, args):
+    return "Test\n" "Downstream\n" "Target\n"
+
+
+# The regular expression for matching func, body, etc.
+regex = re.compile(
+    r'^_?(?P<func>[^:]+):[ \t]*#+[ \t]*(@"?(?P=func)"?| -- Begin function (?P=func))\n(?:\s*\.?Lfunc_begin[^:\n]*:\n)?'
+    r"(?:\.L(?P=func)\$local:\n)?"  # drop .L<func>$local:
+    r"(?:\s*\.type\s+\.L(?P=func)\$local,@function\n)?"  # drop .type .L<func>$local
+    r"(?:[ \t]*(?:\.cfi_startproc|\.cfi_personality|\.cfi_lsda|\.seh_proc|\.seh_handler)\b[^\n]*\n)*"  # drop optional cfi
+    r"(?P<body>^##?[ \t]+[^:]+:.*?)\s*"
+    r"^\s*(?:[^:\n]+?:\s*\n\s*\.size|\.cfi_endproc|\.globl|\.comm|\.(?:sub)?section|#+ -- End function)",
+    flags=(re.M | re.S),
+)
+
+
+def get_run_handler(triple):
+    return (scrub, regex)
+
+
+def add_checks(
+    output_lines,
+    comment_marker,
+    prefix_list,
+    func_dict,
+    func_name,
+    ginfo: common.GeneralizerInfo,
+    global_vars_seen_dict,
+    is_filtered,
+):
+    # Label format is based on ASM string.
+    check_label_format = "{} %s-LABEL: %s%s%s%s".format(comment_marker)
+    return common.add_checks(
+        output_lines,
+        comment_marker,
+        prefix_list,
+        func_dict,
+        func_name,
+        check_label_format,
+        ginfo,
+        global_vars_seen_dict,
+        is_filtered=is_filtered,
+    )
diff --git a/llvm/test/tools/UpdateTestChecks/update_llc_test_checks/downstream_target.test b/llvm/test/tools/UpdateTestChecks/update_llc_test_checks/downstream_target.test
@@ -0,0 +1,7 @@
+# REQUIRES: x86-registered-target
+## Tests that the downstream module support works
+
+# RUN: cp -f %S/Inputs/downstream_target.ll %t.ll && %update_llc_test_checks %t.ll --downstream-target-handler-path=%S/Inputs/downstream_target.py
+# RUN: grep -v '^; NOTE:' %t.ll > %t.ll.expected
+# RUN: diff -b -u %S/Inputs/downstream_target.ll.expected %t.ll.expected
+
diff --git a/llvm/utils/update_llc_test_checks.py b/llvm/utils/update_llc_test_checks.py
@@ -11,6 +11,8 @@
 
 import argparse
 import os  # Used to advertise this file's name ("autogenerated_note").
+import sys
+import importlib.util
 
 from UpdateTestChecks import common
 
@@ -21,6 +23,19 @@
 ]
 
 
+# Helper that imports module from `module_path` and returns it.
+def import_module(module_path):
+    if not os.path.exists(module_path):
+        common.error("Path does not exist", module_path)
+        sys.exit(1)
+    module_name_with_ext = os.path.basename(module_path)
+    module_name = os.path.splitext(module_name_with_ext)[0]
+    spec = importlib.util.spec_from_file_location(module_name, module_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
+
 def main():
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
@@ -67,6 +82,12 @@ def main():
         help="Set a default -march for when neither triple nor arch are found in a RUN line",
     )
     parser.add_argument("tests", nargs="+")
+    parser.add_argument(
+        "--downstream-target-handler-path",
+        default=None,
+        dest="downstream_target_handler_path",
+        help="The path of a python module that implements the scrubber and the regex generator and adds it to the targets handlers map for a downstream target",
+    )
     initial_args = common.parse_commandline_args(parser)
 
     script_name = os.path.basename(__file__)
@@ -107,10 +128,17 @@ def main():
                 march_in_cmd = m.groups()[0]
 
             m = common.DEBUG_ONLY_ARG_RE.search(llc_cmd)
-            if m and m.groups()[0] == "isel":
-                from UpdateTestChecks import isel as output_type
+            if initial_args.downstream_target_handler_path:
+                common.warn(
+                    "Downstream handlers provided by '%s'"
+                    % initial_args.downstream_target_handler_path
+                )
+                output_type = import_module(initial_args.downstream_target_handler_path)
             else:
-                from UpdateTestChecks import asm as output_type
+                if m and m.groups()[0] == "isel":
+                    from UpdateTestChecks import isel as output_type
+                else:
+                    from UpdateTestChecks import asm as output_type
 
             common.verify_filecheck_prefixes(filecheck_cmd)