[llvm] [utils][UpdateLLCTestChecks] Add MIR support to update_llc_test_checks.py. (PR #164965)

[Valery Pykhtin via llvm-commits]

https://github.com/vpykhtin updated https://github.com/llvm/llvm-project/pull/164965

>From 9d05987ccbac42f837220bf1e2feff4ded8ef408 Mon Sep 17 00:00:00 2001
From: Valery Pykhtin <valery.pykhtin at amd.com>
Date: Fri, 24 Oct 2025 11:41:06 +0000
Subject: [PATCH] [utils][UpdateLLCTestChecks] Add MIR support to
 update_llc_test_checks.py

This change enables update_llc_test_checks.py to automatically generate
MIR checks for RUN lines that use -stop-before or -stop-after flags.

This allows tests to verify intermediate compilation stages (e.g., after
instruction selection but before peephole optimizations) alongside the
final assembly output.
---
 llvm/utils/update_llc_test_checks.py | 64 +++++++++++++++++++++++++---
 1 file changed, 57 insertions(+), 7 deletions(-)

diff --git a/llvm/utils/update_llc_test_checks.py b/llvm/utils/update_llc_test_checks.py
index 8c57e75f34f75..ec4c3cefdbe56 100755
--- a/llvm/utils/update_llc_test_checks.py
+++ b/llvm/utils/update_llc_test_checks.py
@@ -13,9 +13,11 @@
 from traceback import print_exc
 import argparse
 import os  # Used to advertise this file's name ("autogenerated_note").
+import re
 import sys
 
 from UpdateTestChecks import common
+import update_mir_test_checks  # Reuse MIR parsing code.
 
 # llc is the only llc-like in the LLVM tree but downstream forks can add
 # additional ones here if they have them.
@@ -24,6 +26,11 @@
 ]
 
 
+def has_stop_pass_flag(llc_args):
+    """Check if llc arguments contain -stop-before or -stop-after flag."""
+    return re.search(r"-stop-(?:before|after)=", llc_args) is not None
+
+
 def update_test(ti: common.TestInfo):
     triple_in_ir = None
     for l in ti.input_lines:
@@ -119,6 +126,10 @@ def update_test(ti: common.TestInfo):
         ginfo=ginfo,
     )
 
+    # Dictionary to store MIR function bodies separately
+    mir_func_dict = {}
+    mir_run_list = []
+
     for (
         prefixes,
         llc_tool,
@@ -141,14 +152,38 @@ def update_test(ti: common.TestInfo):
         if not triple:
             triple = common.get_triple_from_march(march_in_cmd)
 
-        scrubber, function_re = output_type.get_run_handler(triple)
-        if 0 == builder.process_run_line(
-            function_re, scrubber, raw_tool_output, prefixes
-        ):
-            common.warn(
-                "Couldn't match any function. Possibly the wrong target triple has been provided"
+        # Check if this run generates MIR output
+        if has_stop_pass_flag(llc_args):
+            common.debug("Detected MIR output mode for prefixes:", str(prefixes))
+            # Initialize MIR func_dict for these prefixes.
+            for prefix in prefixes:
+                if prefix not in mir_func_dict:
+                    mir_func_dict[prefix] = {}
+
+            # Build MIR function dictionary using imported function.
+            update_mir_test_checks.build_function_info_dictionary(
+                ti.path,
+                raw_tool_output,
+                triple,
+                prefixes,
+                mir_func_dict,
+                ti.args.verbose,
             )
-        builder.processed_prefixes(prefixes)
+
+            # Store this run info for later.
+            mir_run_list.append(
+                (prefixes, llc_tool, llc_args, triple_in_cmd, march_in_cmd)
+            )
+        else:
+            # Regular assembly output.
+            scrubber, function_re = output_type.get_run_handler(triple)
+            if 0 == builder.process_run_line(
+                function_re, scrubber, raw_tool_output, prefixes
+            ):
+                common.warn(
+                    "Couldn't match any function. Possibly the wrong target triple has been provided"
+                )
+            builder.processed_prefixes(prefixes)
 
     func_dict = builder.finish_and_get_func_dict()
     global_vars_seen_dict = {}
@@ -221,6 +256,21 @@ def update_test(ti: common.TestInfo):
                         is_filtered=builder.is_filtered(),
                     )
                 )
+
+                # Also add MIR checks if we have them for this function
+                if mir_run_list and func_name:
+                    common.add_mir_checks_for_function(
+                        ti.path,
+                        output_lines,
+                        mir_run_list,
+                        mir_func_dict,
+                        func_name,
+                        single_bb=False,  # Don't skip basic block labels.
+                        print_fixed_stack=False,  # Don't print fixed stack (ASM tests don't need it).
+                        first_check_is_next=False,  # First check is LABEL, not NEXT.
+                        at_the_function_name=False,  # Use "name:" not "@name".
+                    )
+
                 is_in_function_start = False
 
             if is_in_function: