[lit] Add --fn to prepend llvm-extract for function-level test narrowing Add a --fn=name1,name2 flag to llvm-lit that prepends llvm-extract --func=<name> to the first pipeline stage of each RUN line whose first stage references %s. This lets users narrow IR tests to specific functions and their dependencies without modifying test files.
diff --git a/llvm/utils/lit/lit/LitConfig.py b/llvm/utils/lit/lit/LitConfig.py index 4be2a0f..acc8c52 100644 --- a/llvm/utils/lit/lit/LitConfig.py +++ b/llvm/utils/lit/lit/LitConfig.py
@@ -42,6 +42,7 @@ per_test_coverage=False, gtest_sharding=True, update_tests=False, + fnSelection=None, ): # The name of the test runner. self.progname = progname @@ -96,6 +97,7 @@ self.gtest_sharding = bool(gtest_sharding) self.update_tests = update_tests self.test_updaters = [diff_test_updater] + self.fnSelection = fnSelection @property def maxIndividualTestTime(self):
diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py index 82852f1..7254acc 100644 --- a/llvm/utils/lit/lit/TestRunner.py +++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1921,6 +1921,34 @@ commandLine = "%s && test -e %s" % (commandLine, filePath) return commandLine + +def _applyFnSelection(script, fn_names): + """Prepend llvm-extract to each command so the pipeline receives reduced IR.""" + if not fn_names: + return script + func_args = " ".join("--func=" + n for n in fn_names) + extract = "llvm-extract " + func_args + " %s -o -" + for directive in script: + if not isinstance(directive, CommandDirective): + continue + cmd = directive.command + pipe_idx = cmd.find("|") + first, rest = (cmd[:pipe_idx], cmd[pipe_idx:]) if pipe_idx != -1 else (cmd, "") + if "%s" not in first: + continue + # Preserve %dbg(...) marker at the start of the command + dbg = re.match(r"(%dbg\([^)]*\))\s*", first) + prefix = dbg.group(1) + " " if dbg else "" + if dbg: + first = first[dbg.end():] + # Strip '< %s' stdin redirect so the tool reads from the pipe instead + first = re.sub(r"<\s*%s(?=\s|$)", "", first) + # Strip '%s' positional arg (only when whitespace-delimited) + first = re.sub(r"(?<!\S)%s(?=\s|$)", "", first) + directive.command = prefix + extract + " | " + first.strip() + " " + rest + return script + + def executeShTest( test, litConfig, useExternalSh, extra_substitutions=[], preamble_commands=[] ): @@ -1938,6 +1966,8 @@ if litConfig.noExecute: return lit.Test.Result(Test.PASS) + script = _applyFnSelection(script, litConfig.fnSelection) + tmpDir, tmpBase = getTempPaths(test) substitutions = list(extra_substitutions) substitutions += getDefaultSubstitutions(
diff --git a/llvm/utils/lit/lit/cl_arguments.py b/llvm/utils/lit/lit/cl_arguments.py index bebde4b..926b1bb 100644 --- a/llvm/utils/lit/lit/cl_arguments.py +++ b/llvm/utils/lit/lit/cl_arguments.py
@@ -481,6 +481,16 @@ default=os.environ.get("LIT_RUN_SHARD"), ) + selection_group.add_argument( + "--fn", + dest="fnSelection", + metavar="NAMES", + type=_comma_list, + default=None, + help="Pipe IR output through llvm-extract to keep only the named " + "functions (comma-separated) and their dependencies", + ) + debug_group = parser.add_argument_group("Debug and Experimental Options") debug_group.add_argument( "--debug", help="Enable debugging (for 'lit' development)", action="store_true" @@ -587,6 +597,13 @@ return arg.split(";") +def _comma_list(arg): + names = [n.strip() for n in arg.split(",") if n.strip()] + if not names: + raise _error("empty function name list") + return names + + def _error(desc, *args): msg = desc.format(*args) return argparse.ArgumentTypeError(msg)
diff --git a/llvm/utils/lit/lit/main.py b/llvm/utils/lit/lit/main.py index 77b23bf..2dcd0dc 100755 --- a/llvm/utils/lit/lit/main.py +++ b/llvm/utils/lit/lit/main.py
@@ -44,6 +44,7 @@ gtest_sharding=opts.gtest_sharding, maxRetriesPerTest=opts.maxRetriesPerTest, update_tests=opts.update_tests, + fnSelection=opts.fnSelection, ) discovered_tests = lit.discovery.find_tests_for_inputs(
diff --git a/llvm/utils/lit/tests/Inputs/fn-selection/lit.cfg b/llvm/utils/lit/tests/Inputs/fn-selection/lit.cfg new file mode 100644 index 0000000..23a704a --- /dev/null +++ b/llvm/utils/lit/tests/Inputs/fn-selection/lit.cfg
@@ -0,0 +1,11 @@ +import lit.formats +config.name = "fn-selection" +config.suffixes = [".ll"] +config.test_format = lit.formats.ShTest() +config.test_source_root = os.path.dirname(__file__) +config.test_exec_root = config.test_source_root +config.environment["PATH"] = ( + os.path.join(config.test_source_root, "mock-bin") + + os.pathsep + + config.environment.get("PATH", "") +)
diff --git a/llvm/utils/lit/tests/Inputs/fn-selection/mock-bin/llvm-extract b/llvm/utils/lit/tests/Inputs/fn-selection/mock-bin/llvm-extract new file mode 100755 index 0000000..039e4d0 --- /dev/null +++ b/llvm/utils/lit/tests/Inputs/fn-selection/mock-bin/llvm-extract
@@ -0,0 +1,2 @@ +#!/bin/sh +exit 0
diff --git a/llvm/utils/lit/tests/Inputs/fn-selection/sample.ll b/llvm/utils/lit/tests/Inputs/fn-selection/sample.ll new file mode 100644 index 0000000..d1c6ec4 --- /dev/null +++ b/llvm/utils/lit/tests/Inputs/fn-selection/sample.ll
@@ -0,0 +1,2 @@ +; RUN: echo opt %s -S -passes='instcombine' | echo FileCheck %s +; RUN: echo opt < %s -S -passes="instcombine" | echo FileCheck %s
diff --git a/llvm/utils/lit/tests/fn-selection.py b/llvm/utils/lit/tests/fn-selection.py new file mode 100644 index 0000000..0364592 --- /dev/null +++ b/llvm/utils/lit/tests/fn-selection.py
@@ -0,0 +1,25 @@ +# Verify lit's --fn flag prepends llvm-extract to pipelines. + +# --- --fn=foo: single function --- +# RUN: %{lit} -a --fn=foo %{inputs}/fn-selection/sample.ll \ +# RUN: | FileCheck --check-prefix=SINGLE %s +# +# Positional %s form: +# SINGLE: llvm-extract --func=foo {{.*}}sample.ll -o - | echo opt -S -passes='instcombine' | echo FileCheck {{.*}}sample.ll +# Redirect < %s form: +# SINGLE: llvm-extract --func=foo {{.*}}sample.ll -o - | echo opt -S -passes="instcombine" | echo FileCheck {{.*}}sample.ll + +# --- --fn=foo,bar: multiple functions --- +# RUN: %{lit} -a --fn=foo,bar %{inputs}/fn-selection/sample.ll \ +# RUN: | FileCheck --check-prefix=MULTI %s +# +# MULTI: llvm-extract --func=foo --func=bar {{.*}}sample.ll -o - | echo opt -S -passes='instcombine' | echo FileCheck {{.*}}sample.ll +# MULTI: llvm-extract --func=foo --func=bar {{.*}}sample.ll -o - | echo opt -S -passes="instcombine" | echo FileCheck {{.*}}sample.ll + +# --- No --fn: passes unchanged --- +# RUN: %{lit} -a %{inputs}/fn-selection/sample.ll \ +# RUN: | FileCheck --check-prefix=NONE %s +# +# NONE-NOT: llvm-extract +# NONE: echo opt {{.*}}sample.ll -S -passes='instcombine' +# NONE: echo opt < {{.*}}sample.ll -S -passes="instcombine"