[llvm-branch-commits] [llvm] [lit] Add lit.llvm.fn_selection: opt-in select-function pass via --param fn-pass (PR #200352)

[via llvm-branch-commits]

https://github.com/jofrn updated https://github.com/llvm/llvm-project/pull/200352

>From 66126dc09fc82cb47a5f09655cd959b3ed3a0847 Mon Sep 17 00:00:00 2001
From: jofrn <165626406+jofrn at users.noreply.github.com>
Date: Fri, 29 May 2026 01:18:00 -0700
Subject: [PATCH] [lit] Add lit.llvm.fn_selection: opt-in select-function pass
 via --param fn-pass

---
 llvm/utils/lit/lit/llvm/fn_param.py           |  6 ++---
 llvm/utils/lit/lit/llvm/fn_selection.py       | 16 ++++++++++++
 .../lit/tests/Inputs/fn-selection/lit.cfg     | 10 ++++++++
 .../lit/tests/Inputs/fn-selection/sample.ll   |  2 ++
 llvm/utils/lit/tests/fn-selection.py          | 25 +++++++++++++++++++
 5 files changed, 56 insertions(+), 3 deletions(-)
 create mode 100644 llvm/utils/lit/lit/llvm/fn_selection.py
 create mode 100644 llvm/utils/lit/tests/Inputs/fn-selection/lit.cfg
 create mode 100644 llvm/utils/lit/tests/Inputs/fn-selection/sample.ll
 create mode 100644 llvm/utils/lit/tests/fn-selection.py

diff --git a/llvm/utils/lit/lit/llvm/fn_param.py b/llvm/utils/lit/lit/llvm/fn_param.py
index ebab824674dab..38ba56d59a250 100644
--- a/llvm/utils/lit/lit/llvm/fn_param.py
+++ b/llvm/utils/lit/lit/llvm/fn_param.py
@@ -28,9 +28,9 @@ def install(config, lit_config):
     pass, opt-only); otherwise `lit.llvm.fn_extract` is used (prepends
     llvm-extract, tool-agnostic)."""
     if lit_config.params.get("fn-pass"):
-        # from lit.llvm import fn_selection
-        # fn_selection.install(config, lit_config)
-        pass
+        from lit.llvm import fn_selection
+
+        fn_selection.install(config, lit_config)
     else:
         from lit.llvm import fn_extract
 
diff --git a/llvm/utils/lit/lit/llvm/fn_selection.py b/llvm/utils/lit/lit/llvm/fn_selection.py
new file mode 100644
index 0000000000000..3addb9be732fc
--- /dev/null
+++ b/llvm/utils/lit/lit/llvm/fn_selection.py
@@ -0,0 +1,16 @@
+"""Splice a `select-function<fn=...>` pass at the head of every `-passes=`
+pipeline so only the named functions (and their transitive dependencies) are
+compiled by `opt`. Driven by `--param fn=NAMES`."""
+
+from lit.llvm.fn_param import add_capture_sub, parse_fn_names
+
+
+def install(config, lit_config):
+    names = parse_fn_names(lit_config)
+    if not names:
+        return
+    sel = "select-function<" + ";".join("fn=" + n for n in names) + ">"
+    # -passes='...' / -passes="..." — splice select-function after the quote
+    add_capture_sub(config, r"""-passes=(['"])""", r"-passes=\1" + sel + ",")
+    # -passes=word (unquoted) — wrap to protect angle brackets
+    add_capture_sub(config, r"-passes=([^'\"\s]\S*)", r"-passes='" + sel + r",\1'")
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 0000000000000..43cb0c32d7dc4
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/fn-selection/lit.cfg
@@ -0,0 +1,10 @@
+import lit.formats
+from lit.llvm import fn_param
+
+config.name = "fn-selection"
+config.suffixes = [".ll"]
+config.test_format = lit.formats.ShTest()
+config.test_source_root = None
+config.test_exec_root = None
+
+fn_param.install(config, lit_config)
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 0000000000000..30c01a9469750
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/fn-selection/sample.ll
@@ -0,0 +1,2 @@
+; RUN: echo -passes='instcombine,mem2reg'
+; RUN: echo -passes="instcombine,mem2reg"
diff --git a/llvm/utils/lit/tests/fn-selection.py b/llvm/utils/lit/tests/fn-selection.py
new file mode 100644
index 0000000000000..5a1e7c6f6056d
--- /dev/null
+++ b/llvm/utils/lit/tests/fn-selection.py
@@ -0,0 +1,25 @@
+# Verify --param fn=NAMES + --param fn-pass=1 splices select-function into
+# -passes= pipelines via lit.llvm.fn_selection (same dispatch as
+# llvm/test/lit.cfg.py).
+
+# --- --param fn=foo: single function ---
+# RUN: %{lit} -a --param fn=foo --param fn-pass=1 %{inputs}/fn-selection/sample.ll \
+# RUN:   | FileCheck --check-prefix=SINGLE %s
+#
+# SINGLE: -passes='select-function<fn=foo>,instcombine,mem2reg'
+# SINGLE: -passes="select-function<fn=foo>,instcombine,mem2reg"
+
+# --- --param fn=foo,bar: multiple functions ---
+# RUN: %{lit} -a --param fn=foo,bar --param fn-pass=1 %{inputs}/fn-selection/sample.ll \
+# RUN:   | FileCheck --check-prefix=MULTI %s
+#
+# MULTI: -passes='select-function<fn=foo;fn=bar>,instcombine,mem2reg'
+# MULTI: -passes="select-function<fn=foo;fn=bar>,instcombine,mem2reg"
+
+# --- No --param: passes unchanged ---
+# RUN: %{lit} -a %{inputs}/fn-selection/sample.ll \
+# RUN:   | FileCheck --check-prefix=NONE %s
+#
+# NONE-NOT: select-function
+# NONE: -passes='instcombine,mem2reg'
+# NONE: -passes="instcombine,mem2reg"