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

[Francis Visoiu Mistrih via llvm-commits]

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

For MI bundles, the instruction count remark doesn't count the instructions inside the bundle.

>From 6ba6dff0f5cd16a895dfe6ec060a1124f0ffff85 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    | 22 ++++++++----
 .../RISCV/instruction-count-remark.mir        | 35 +++++++++++++++++++
 2 files changed, 51 insertions(+), 6 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 60cb26973ead41..d160f08c2f90e7 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1724,7 +1724,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.
@@ -1816,11 +1815,22 @@ void AsmPrinter::emitFunctionBody() {
       default:
         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 = [&](unsigned Opcode) {
+            MCInst MCI;
+            MCI.setOpcode(Opcode);
+            auto Name = OutStreamer->getMnemonic(MCI);
+            ++MnemonicCounts[Name];
+            ++NumInstsInFunction;
+          };
+          if (!MI.isBundle()) {
+            CountInstruction(MI.getOpcode());
+            break;
+          }
+          // Separately count all the instructions in a bundle.
+          for (auto It = std::next(MI.getIterator());
+               It != MBB.end() && It->isInsideBundle(); ++It) {
+            CountInstruction(It->getOpcode());
+          }
         }
         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..539eb3aff3b888
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/instruction-count-remark.mir
@@ -0,0 +1,35 @@
+# 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
+...