[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.

Author: vpykhtin

llvm/test/tools/UpdateTestChecks/update_llc_test_checks/Inputs/x86_asm_mir_mixed.ll

@@ -0,0 +1,17 @@
+; RUN: llc -mtriple=x86_64 < %s | FileCheck %s --check-prefix=ASM
+; RUN: llc -mtriple=x86_64 -stop-after=finalize-isel < %s | FileCheck %s --check-prefix=MIR
+
+define i64 @test1(i64 %i) nounwind readnone {
+  %loc = alloca i64
+  %j = load i64, i64 * %loc
+  %r = add i64 %i, %j
+  ret i64 %r
+}
+
+define i64 @test2(i32 %i) nounwind readnone {
+  %loc = alloca i32
+  %j = load i32, i32 * %loc
+  %r = add i32 %i, %j
+  %ext = zext i32 %r to i64
+  ret i64 %ext
+}

llvm/test/tools/UpdateTestChecks/update_llc_test_checks/Inputs/x86_asm_mir_mixed.ll.expected

@@ -0,0 +1,47 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=x86_64 < %s | FileCheck %s --check-prefix=ASM
+; RUN: llc -mtriple=x86_64 -stop-after=finalize-isel < %s | FileCheck %s --check-prefix=MIR
+
+define i64 @test1(i64 %i) nounwind readnone {
+; ASM-LABEL: test1:
+; ASM:       # %bb.0:
+; ASM-NEXT:    movq %rdi, %rax
+; ASM-NEXT:    addq -{{[0-9]+}}(%rsp), %rax
+; ASM-NEXT:    retq
+  ; MIR-LABEL: name: test1
+  ; MIR: bb.0 (%ir-block.0):
+  ; MIR-NEXT:   liveins: $rdi
+  ; MIR-NEXT: {{  $}}
+  ; MIR-NEXT:   [[COPY:%[0-9]+]]:gr64 = COPY $rdi
+  ; MIR-NEXT:   [[ADD64rm:%[0-9]+]]:gr64 = ADD64rm [[COPY]], %stack.0.loc, 1, $noreg, 0, $noreg, implicit-def dead $eflags :: (dereferenceable load (s64) from %ir.loc)
+  ; MIR-NEXT:   $rax = COPY [[ADD64rm]]
+  ; MIR-NEXT:   RET 0, $rax
+  %loc = alloca i64
+  %j = load i64, i64 * %loc
+  %r = add i64 %i, %j
+  ret i64 %r
+}
+
+define i64 @test2(i32 %i) nounwind readnone {
+; ASM-LABEL: test2:
+; ASM:       # %bb.0:
+; ASM-NEXT:    movl %edi, %eax
+; ASM-NEXT:    addl -{{[0-9]+}}(%rsp), %eax
+; ASM-NEXT:    retq
+  ; MIR-LABEL: name: test2
+  ; MIR: bb.0 (%ir-block.0):
+  ; MIR-NEXT:   liveins: $edi
+  ; MIR-NEXT: {{  $}}
+  ; MIR-NEXT:   [[COPY:%[0-9]+]]:gr32 = COPY $edi
+  ; MIR-NEXT:   [[ADD32rm:%[0-9]+]]:gr32 = ADD32rm [[COPY]], %stack.0.loc, 1, $noreg, 0, $noreg, implicit-def dead $eflags :: (dereferenceable load (s32) from %ir.loc)
+  ; MIR-NEXT:   [[SUBREG_TO_REG:%[0-9]+]]:gr64 = SUBREG_TO_REG 0, killed [[ADD32rm]], %subreg.sub_32bit
+  ; MIR-NEXT:   $rax = COPY [[SUBREG_TO_REG]]
+  ; MIR-NEXT:   RET 0, $rax
+  %loc = alloca i32
+  %j = load i32, i32 * %loc
+  %r = add i32 %i, %j
+  %ext = zext i32 %r to i64
+  ret i64 %ext
+}
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; MIR: {{.*}}

llvm/test/tools/UpdateTestChecks/update_llc_test_checks/x86-asm-mir-mixed.test

@@ -0,0 +1,5 @@
+# REQUIRES: x86-registered-target
+## Test checking that update_llc_test_checks.py can generate both ASM and MIR checks in the same file
+
+# RUN: cp -f %S/Inputs/x86_asm_mir_mixed.ll %t.ll && %update_llc_test_checks %t.ll
+# RUN: diff -u %S/Inputs/x86_asm_mir_mixed.ll.expected %t.ll

llvm/utils/UpdateTestChecks/common.py

@@ -604,6 +604,7 @@ def invoke_tool(exe, cmd_args, ir, preprocess_cmd=None, verbose=False):
 TRIPLE_ARG_RE = re.compile(r"-m?triple[= ]([^ ]+)")
 MARCH_ARG_RE = re.compile(r"-march[= ]([^ ]+)")
 DEBUG_ONLY_ARG_RE = re.compile(r"-debug-only[= ]([^ ]+)")
+STOP_PASS_RE = re.compile(r"-stop-(before|after)=(\w+)")
 
 IS_DEBUG_RECORD_RE = re.compile(r"^(\s+)#dbg_")
 

llvm/utils/update_llc_test_checks.py

@@ -16,6 +16,7 @@
 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.
@@ -119,6 +120,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 +146,37 @@ 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
+        # Check if we are in MIR output mode. If -debug-only is present assume
+        # the debug output is the main point of interest.
+        if common.STOP_PASS_RE.search(llc_args) and not common.DEBUG_ONLY_ARG_RE.search(
+            llc_args
         ):
-            common.warn(
-                "Couldn't match any function. Possibly the wrong target triple has been provided"
+            common.debug("Detected MIR output mode for prefixes:", str(prefixes))
+            for prefix in prefixes:
+                if prefix not in mir_func_dict:
+                    mir_func_dict[prefix] = {}
+
+            update_mir_test_checks.build_function_info_dictionary(
+                ti.path,
+                raw_tool_output,
+                triple,
+                prefixes,
+                mir_func_dict,
+                ti.args.verbose,
+            )
+
+            mir_run_list.append(
+                (prefixes, llc_tool, llc_args, triple_in_cmd, march_in_cmd)
             )
-        builder.processed_prefixes(prefixes)
+        else:
+            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 +249,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: