[llvm] [CodeGen] Fix InstructionCount remarks for MI bundles (PR #107621)

[Francis Visoiu Mistrih via llvm-commits]

https://github.com/francisvm updated https://github.com/llvm/llvm-project/pull/107621

>From 3caf9d480b5419505117b3ee2bc34d1b0d7340b7 Mon Sep 17 00:00:00 2001
From: Francis Visoiu Mistrih <francisvm at apple.com>
Date: Fri, 6 Sep 2024 08:54:53 -0700
Subject: [PATCH] [CodeGen] Fix InstructionCount remarks for MI bundles

For MI bundles, the instruction count remark doesn't count the
instructions inside the bundle.
---
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp    | 31 ++++++--
 .../RISCV/instruction-count-remark.mir        | 75 +++++++++++++++++++
 2 files changed, 99 insertions(+), 7 deletions(-)
 create mode 100644 llvm/test/CodeGen/RISCV/instruction-count-remark.mir

diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 317278911b28f6..64b01c0189073f 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1746,7 +1746,6 @@ void AsmPrinter::emitFunctionBody() {
       if (!MI.isPosition() && !MI.isImplicitDef() && !MI.isKill() &&
           !MI.isDebugInstr()) {
         HasAnyRealCode = true;
-        ++NumInstsInFunction;
       }
 
       // If there is a pre-instruction symbol, emit a label for it here.
@@ -1844,13 +1843,31 @@ void AsmPrinter::emitFunctionBody() {
         // actual initialization is needed.
         break;
       default:
+        assert(!MI.isMetaInstruction() &&
+               "Meta instructions should be already handled here");
         emitInstruction(&MI);
-        if (CanDoExtraAnalysis) {
-          MCInst MCI;
-          MCI.setOpcode(MI.getOpcode());
-          auto Name = OutStreamer->getMnemonic(MCI);
-          auto I = MnemonicCounts.insert({Name, 0u});
-          I.first->second++;
+
+        auto CountInstruction = [&](const MachineInstr &MI) {
+          // Skip Meta instructions inside bundles.
+          if (MI.isMetaInstruction())
+            return;
+          ++NumInstsInFunction;
+          if (CanDoExtraAnalysis) {
+            MCInst MCI;
+            MCI.setOpcode(MI.getOpcode());
+            auto Name = OutStreamer->getMnemonic(MCI);
+            assert(!Name.empty() && "Missing mnemonic for opcode");
+            ++MnemonicCounts[Name];
+          }
+        };
+        if (!MI.isBundle()) {
+          CountInstruction(MI);
+          break;
+        }
+        // Separately count all the instructions in a bundle.
+        for (auto It = std::next(MI.getIterator());
+             It != MBB.end() && It->isInsideBundle(); ++It) {
+          CountInstruction(*It);
         }
         break;
       }
diff --git a/llvm/test/CodeGen/RISCV/instruction-count-remark.mir b/llvm/test/CodeGen/RISCV/instruction-count-remark.mir
new file mode 100644
index 00000000000000..4f429ab5274e87
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/instruction-count-remark.mir
@@ -0,0 +1,75 @@
+# RUN: llc -mtriple=riscv32 -verify-machineinstrs -start-before=riscv-expand-pseudo -simplify-mir -o /dev/null -pass-remarks-analysis=asm-printer %s 2>&1 | FileCheck %s
+---
+name: instrs
+tracksRegLiveness: true
+body: |
+  bb.0:
+    $x0 = ADDI $x0, 0
+    $x0 = ADDI $x0, 0
+    $x0 = ADDI $x0, 0
+    $x0 = LW $x0, 0
+    $x0 = LW $x0, 0
+    $x0 = XORI $x0, 0
+    ; CHECK: addi : 3
+    ; CHECK-NEXT: lw : 2
+    ; CHECK-NEXT: xori : 1
+    ; CHECK: 6 instructions in function
+...
+---
+name: bundles
+tracksRegLiveness: true
+body: |
+  bb.0:
+    $x0 = ADDI $x0, 0
+    BUNDLE {
+    $x0 = ADDI $x0, 0
+    $x0 = ADDI $x0, 0
+    $x0 = LW $x0, 0
+    }
+    $x0 = LW $x0, 0
+    $x0 = XORI $x0, 0
+    ; CHECK: addi : 3
+    ; CHECK-NEXT: lw : 2
+    ; CHECK-NEXT: xori : 1
+    ; CHECK: 6 instructions in function
+...
+---
+name: metainstrs
+tracksRegLiveness: true
+body: |
+  bb.0:
+    $x0 = ADDI $x0, 0
+    $x0 = ADDI $x0, 0
+    $x0 = ADDI $x0, 0
+    $x0 = IMPLICIT_DEF
+    $x0 = LW $x0, 0
+    $x0 = LW $x0, 0
+    CFI_INSTRUCTION adjust_cfa_offset 4
+    $x0 = XORI $x0, 0
+    DBG_VALUE $x0, 0
+    ; CHECK: addi : 3
+    ; CHECK-NEXT: lw : 2
+    ; CHECK-NEXT: xori : 1
+    ; CHECK: 6 instructions in function
+...
+---
+name: metabundles
+tracksRegLiveness: true
+body: |
+  bb.0:
+    $x0 = ADDI $x0, 0
+    BUNDLE {
+    CFI_INSTRUCTION adjust_cfa_offset 4
+    $x0 = ADDI $x0, 0
+    $x0 = ADDI $x0, 0
+    DBG_VALUE $x0, 0
+    $x0 = LW $x0, 0
+    }
+    $x0 = LW $x0, 0
+    $x0 = IMPLICIT_DEF
+    $x0 = XORI $x0, 0
+    ; CHECK: addi : 3
+    ; CHECK-NEXT: lw : 2
+    ; CHECK-NEXT: xori : 1
+    ; CHECK: 6 instructions in function
+...