From 916e6ad7d0608a1de6b570483a82509640b210b8 Mon Sep 17 00:00:00 2001
From: Francis Visoiu Mistrih <890283+francisvm@users.noreply.github.com>
Date: Wed, 2 Oct 2024 13:37:25 -0700
Subject: [PATCH] [CodeGen] Fix InstructionCount remarks for MI bundles
 (#107621)

For MI bundles, the instruction count remark doesn't count the
instructions inside the bundle.
---
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp    | 40 ++++++++--
 .../RISCV/instruction-count-remark.mir        | 75 +++++++++++++++++++
 2 files changed, 108 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..3a8cde7330efc0 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1703,6 +1703,20 @@ static bool needFuncLabels(const MachineFunction &MF, const AsmPrinter &Asm) {
       classifyEHPersonality(MF.getFunction().getPersonalityFn()));
 }
 
+// Return the mnemonic of a MachineInstr if available, or the MachineInstr
+// opcode name otherwise.
+static StringRef getMIMnemonic(const MachineInstr &MI, MCStreamer &Streamer) {
+  const TargetInstrInfo *TII =
+      MI.getParent()->getParent()->getSubtarget().getInstrInfo();
+  MCInst MCI;
+  MCI.setOpcode(MI.getOpcode());
+  if (StringRef Name = Streamer.getMnemonic(MCI); !Name.empty())
+    return Name;
+  StringRef Name = TII->getName(MI.getOpcode());
+  assert(!Name.empty() && "Missing mnemonic and name for opcode");
+  return Name;
+}
+
 /// EmitFunctionBody - This method emits the body and trailer for a
 /// function.
 void AsmPrinter::emitFunctionBody() {
@@ -1746,7 +1760,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.
@@ -1845,12 +1858,25 @@ void AsmPrinter::emitFunctionBody() {
         break;
       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 = [&](const MachineInstr &MI) {
+          // Skip Meta instructions inside bundles.
+          if (MI.isMetaInstruction())
+            return;
+          ++NumInstsInFunction;
+          if (CanDoExtraAnalysis) {
+            StringRef Name = getMIMnemonic(MI, *OutStreamer);
+            ++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
+...