[llvm-branch-commits] [llvm] [lit] Add lit.llvm.decorate: --param decorate=PROG prepends PROG to a pipeline stage (PR #202299)

[via llvm-branch-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-testing-tools

Author: jofrn

<details>
<summary>Changes</summary>

[lit] Add lit.llvm.decorate: --param decorate=PROG prepends PROG to a pipeline stage

---
Full diff: https://github.com/llvm/llvm-project/pull/202299.diff


8 Files Affected:

- (modified) llvm/test/lit.cfg.py (+7) 
- (added) llvm/utils/lit/lit/llvm/decorate.py (+30) 
- (added) llvm/utils/lit/tests/Inputs/decorate/lit.cfg (+10) 
- (added) llvm/utils/lit/tests/Inputs/decorate/sample.ll (+1) 
- (added) llvm/utils/lit/tests/Inputs/narrow-and-decorate/lit.cfg (+16) 
- (added) llvm/utils/lit/tests/Inputs/narrow-and-decorate/sample.ll (+21) 
- (added) llvm/utils/lit/tests/decorate.py (+30) 
- (added) llvm/utils/lit/tests/narrow-and-decorate.py (+18) 


``````````diff
diff --git a/llvm/test/lit.cfg.py b/llvm/test/lit.cfg.py
index f4b9e8aaed6ef..81d582dff64f3 100644
--- a/llvm/test/lit.cfg.py
+++ b/llvm/test/lit.cfg.py
@@ -39,6 +39,13 @@
 )
 config.test_format = lit.formats.ShTest(not use_lit_shell, extra_substitutions)
 
+# --param decorate=PROG prepends PROG to a pipeline stage (debugger, profiler,
+# etc.). Registered before fn_param so subsequent fn_extract substitutions
+# correctly retarget the now-decorated stage.
+from lit.llvm import decorate
+
+decorate.install(config, lit_config)
+
 # --param fn=NAMES narrows compilation to the named functions and their
 # transitive dependencies. See lit.llvm.fn_param.install for dispatch details.
 from lit.llvm import fn_param
diff --git a/llvm/utils/lit/lit/llvm/decorate.py b/llvm/utils/lit/lit/llvm/decorate.py
new file mode 100644
index 0000000000000..3f0e8c1b6a9c4
--- /dev/null
+++ b/llvm/utils/lit/lit/llvm/decorate.py
@@ -0,0 +1,30 @@
+"""`--param decorate=PROG[@N]` prepends PROG to pipeline stage N (0-indexed,
+default 0) of every RUN line, so the downstream tool can be run under a
+debugger, profiler, etc. Examples:
+
+    --param 'decorate=gdb --args'   # run stage 0 under gdb
+    --param decorate=perf            # profile stage 0
+    --param decorate=time at 1          # time the second stage
+
+A stage index past the end of a pipeline is a no-op for that RUN line."""
+
+import re
+
+from lit.llvm.fn_param import add_capture_sub
+
+_TRAILING_INDEX = re.compile(r"^(.+)@(\d+)$")
+
+
+def install(config, lit_config):
+    prog = lit_config.params.get("decorate")
+    if not prog:
+        return
+    m = _TRAILING_INDEX.match(prog)
+    if m:
+        prog, stage = m.group(1), m.group(2)
+    else:
+        stage = "0"
+    # Match (optional %dbg marker)(stage many `|`-delimited stages and their
+    # bars) and insert PROG after, preserving everything captured.
+    pattern = r"^(%dbg\([^)]*\)\s*)?((?:[^|]*\|\s*){" + stage + r"})"
+    add_capture_sub(config, pattern, r"\1\2" + prog + " ")
diff --git a/llvm/utils/lit/tests/Inputs/decorate/lit.cfg b/llvm/utils/lit/tests/Inputs/decorate/lit.cfg
new file mode 100644
index 0000000000000..6d4913e168243
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/decorate/lit.cfg
@@ -0,0 +1,10 @@
+import lit.formats
+from lit.llvm import decorate
+
+config.name = "decorate"
+config.suffixes = [".ll"]
+config.test_format = lit.formats.ShTest()
+config.test_source_root = None
+config.test_exec_root = None
+
+decorate.install(config, lit_config)
diff --git a/llvm/utils/lit/tests/Inputs/decorate/sample.ll b/llvm/utils/lit/tests/Inputs/decorate/sample.ll
new file mode 100644
index 0000000000000..b17fb677c75fd
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/decorate/sample.ll
@@ -0,0 +1 @@
+; RUN: echo opt %s -S -passes='instcombine' | echo FileCheck %s
diff --git a/llvm/utils/lit/tests/Inputs/narrow-and-decorate/lit.cfg b/llvm/utils/lit/tests/Inputs/narrow-and-decorate/lit.cfg
new file mode 100644
index 0000000000000..54079be2ece4c
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/narrow-and-decorate/lit.cfg
@@ -0,0 +1,16 @@
+import lit.formats
+import shutil
+from lit.llvm import decorate, fn_param
+
+config.name = "narrow-and-decorate"
+config.suffixes = [".ll"]
+config.test_format = lit.formats.ShTest()
+config.test_source_root = None
+config.test_exec_root = None
+
+for tool in ("opt", "llvm-extract", "FileCheck"):
+    if shutil.which(tool):
+        config.available_features.add(tool)
+
+decorate.install(config, lit_config)
+fn_param.install(config, lit_config)
diff --git a/llvm/utils/lit/tests/Inputs/narrow-and-decorate/sample.ll b/llvm/utils/lit/tests/Inputs/narrow-and-decorate/sample.ll
new file mode 100644
index 0000000000000..7f08b045715b5
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/narrow-and-decorate/sample.ll
@@ -0,0 +1,21 @@
+; REQUIRES: opt
+;
+; RUN: opt -S < %s -passes=instcombine | FileCheck %s --check-prefix=ALL
+; RUN: opt -S < %s -passes=instcombine | FileCheck %s --check-prefix=ONLYFOO
+
+define i32 @foo(i32 %x) {
+  ret i32 %x
+}
+
+define i32 @bar(i32 %x) {
+  ret i32 %x
+}
+
+; ALL-LABEL: @foo
+; ALL: ret i32 %x
+
+; ALL-LABEL: @bar
+; ALL: ret i32 FILTER_AWAY
+
+; ONLYFOO-LABEL: @foo
+; ONLYFOO: ret i32 %x
diff --git a/llvm/utils/lit/tests/decorate.py b/llvm/utils/lit/tests/decorate.py
new file mode 100644
index 0000000000000..909c92f984431
--- /dev/null
+++ b/llvm/utils/lit/tests/decorate.py
@@ -0,0 +1,30 @@
+# Verify --param decorate=PROG[@N] prepends PROG to pipeline stage N (default 0).
+
+# --- Default stage (0): PROG prepended to the first stage ---
+# RUN: %{lit} -a --param decorate=time %{inputs}/decorate/sample.ll \
+# RUN:   | FileCheck --check-prefix=STAGE0 %s
+#
+# STAGE0: time echo opt {{.*}}sample.ll -S -passes='instcombine' | echo FileCheck {{.*}}sample.ll
+
+# --- Explicit @0 is the same as default ---
+# RUN: %{lit} -a --param decorate=time at 0 %{inputs}/decorate/sample.ll \
+# RUN:   | FileCheck --check-prefix=STAGE0 %s
+
+# --- @1: PROG prepended to the second stage ---
+# RUN: %{lit} -a --param decorate=time at 1 %{inputs}/decorate/sample.ll \
+# RUN:   | FileCheck --check-prefix=STAGE1 %s
+#
+# STAGE1: echo opt {{.*}}sample.ll -S -passes='instcombine' | time echo FileCheck {{.*}}sample.ll
+
+# --- @N past the end: no-op, pipeline unchanged ---
+# RUN: %{lit} -a --param decorate=time at 9 %{inputs}/decorate/sample.ll \
+# RUN:   | FileCheck --check-prefix=OOB %s
+#
+# OOB-NOT: time
+# OOB: echo opt {{.*}}sample.ll -S -passes='instcombine' | echo FileCheck {{.*}}sample.ll
+
+# --- No --param decorate: pipeline unchanged ---
+# RUN: %{lit} -a %{inputs}/decorate/sample.ll | FileCheck --check-prefix=NONE %s
+#
+# NONE-NOT: time
+# NONE: echo opt {{.*}}sample.ll -S -passes='instcombine' | echo FileCheck {{.*}}sample.ll
diff --git a/llvm/utils/lit/tests/narrow-and-decorate.py b/llvm/utils/lit/tests/narrow-and-decorate.py
new file mode 100644
index 0000000000000..d80022f8793ce
--- /dev/null
+++ b/llvm/utils/lit/tests/narrow-and-decorate.py
@@ -0,0 +1,18 @@
+# End-to-end check that --check, --param fn=, and --param decorate= compose
+# correctly on a real opt+FileCheck pipeline. The inner sample has two RUN
+# lines and a deliberately-wrong CHECK in @bar's section.
+
+# --- No params: RUN 0 fails on @bar's FILTER_AWAY check ---
+# RUN: not %{lit} %{inputs}/narrow-and-decorate/sample.ll
+
+# --- --check=1: skip RUN 0, run only the ONLYFOO prefix RUN, which passes ---
+# RUN: %{lit} --check=1 %{inputs}/narrow-and-decorate/sample.ll
+
+# --- --param fn=foo --check=0: extract drops @bar; filter drops its CHECK ---
+# RUN: %{lit} --param fn=foo --check=0 %{inputs}/narrow-and-decorate/sample.ll
+
+# --- All three: --check=0 + --param fn=foo + --param decorate=time ---
+# RUN: %{lit} --param fn=foo --param decorate=time --check=0 %{inputs}/narrow-and-decorate/sample.ll
+
+# --- --param decorate=time alone: pipeline still has the wrong CHECK, fails ---
+# RUN: not %{lit} --param decorate=time %{inputs}/narrow-and-decorate/sample.ll

``````````

</details>


https://github.com/llvm/llvm-project/pull/202299