[llvm] a95aa86 - [MLGO] Add BB Profile Dump in AsmPrinter

[Aiden Grossman via llvm-commits]

Author: Aiden Grossman
Date: 2023-02-08T23:13:42Z
New Revision: a95aa86b538f643cde8502a58e2ce8804e895e96

URL: https://github.com/llvm/llvm-project/commit/a95aa86b538f643cde8502a58e2ce8804e895e96
DIFF: https://github.com/llvm/llvm-project/commit/a95aa86b538f643cde8502a58e2ce8804e895e96.diff

LOG: [MLGO] Add BB Profile Dump in AsmPrinter

This patch adds a basic block profile dump option within the AsmPrinter
and dumps basic block profile information so that cost models can use
the data for downstream analysis.

Differential Revision: https://reviews.llvm.org/D143311

Added: 
    core
    llvm/test/CodeGen/MLRegalloc/bb-profile-dump.ll

Modified: 
    llvm/include/llvm/CodeGen/AsmPrinter.h
    llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

Removed: 
    


################################################################################
diff  --git a/core b/core
new file mode 100644
index 000000000000..22890d321252
Binary files /dev/null and b/core 
diff er

diff  --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h
index 3c20fc06abef..15f76d8de217 100644
--- a/llvm/include/llvm/CodeGen/AsmPrinter.h
+++ b/llvm/include/llvm/CodeGen/AsmPrinter.h
@@ -236,6 +236,10 @@ class AsmPrinter : public MachineFunctionPass {
   /// split stack prologue.
   bool HasNoSplitStack = false;
 
+  /// Raw FDOstream for outputting machine basic block frequncies if the
+  /// --mbb-profile-dump flag is set for downstream cost modelling applications
+  std::unique_ptr<raw_fd_ostream> MBBProfileDumpFileOutput;
+
 protected:
   explicit AsmPrinter(TargetMachine &TM, std::unique_ptr<MCStreamer> Streamer);
 

diff  --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 39e6ffed9aad..4fbe1f4da9b1 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -37,6 +37,7 @@
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/CodeGen/GCMetadata.h"
 #include "llvm/CodeGen/GCMetadataPrinter.h"
+#include "llvm/CodeGen/LazyMachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineConstantPool.h"
 #include "llvm/CodeGen/MachineDominators.h"
@@ -128,6 +129,13 @@ using namespace llvm;
 
 #define DEBUG_TYPE "asm-printer"
 
+static cl::opt<std::string> BasicBlockProfileDump(
+    "mbb-profile-dump", cl::Hidden,
+    cl::desc("Basic block profile dump for external cost modelling. If "
+             "matching up BBs with afterwards, the compilation must be "
+             "performed with -fbasic-block-sections=labels. Enabling this "
+             "flag during in-process ThinLTO is not supported."));
+
 const char DWARFGroupName[] = "dwarf";
 const char DWARFGroupDescription[] = "DWARF Emission";
 const char DbgTimerName[] = "emit";
@@ -414,6 +422,7 @@ void AsmPrinter::getAnalysisUsage(AnalysisUsage &AU) const {
   MachineFunctionPass::getAnalysisUsage(AU);
   AU.addRequired<MachineOptimizationRemarkEmitterPass>();
   AU.addRequired<GCModuleInfo>();
+  AU.addRequired<LazyMachineBlockFrequencyInfoPass>();
 }
 
 bool AsmPrinter::doInitialization(Module &M) {
@@ -585,6 +594,16 @@ bool AsmPrinter::doInitialization(Module &M) {
     HI.Handler->beginModule(&M);
   }
 
+  if (!BasicBlockProfileDump.empty()) {
+    std::error_code PossibleFileError;
+    MBBProfileDumpFileOutput = std::make_unique<raw_fd_ostream>(
+        BasicBlockProfileDump, PossibleFileError);
+    if (PossibleFileError) {
+      M.getContext().emitError("Failed to open file for MBB Profile Dump: " +
+                               PossibleFileError.message() + "\n");
+    }
+  }
+
   return false;
 }
 
@@ -1888,6 +1907,18 @@ void AsmPrinter::emitFunctionBody() {
     OutStreamer->getCommentOS() << "-- End function\n";
 
   OutStreamer->addBlankLine();
+
+  // Output MBB numbers, function names, and frequencies if the flag to dump
+  // MBB profile information has been set
+  if (MBBProfileDumpFileOutput) {
+    MachineBlockFrequencyInfo &MBFI =
+        getAnalysis<LazyMachineBlockFrequencyInfoPass>().getBFI();
+    for (const auto &MBB : *MF) {
+      *MBBProfileDumpFileOutput.get()
+          << MF->getName() << "," << MBB.getNumber() << ","
+          << MBFI.getBlockFreqRelativeToEntryBlock(&MBB) << "\n";
+    }
+  }
 }
 
 /// Compute the number of Global Variables that uses a Constant.

diff  --git a/llvm/test/CodeGen/MLRegalloc/bb-profile-dump.ll b/llvm/test/CodeGen/MLRegalloc/bb-profile-dump.ll
new file mode 100644
index 000000000000..5486dd67fa86
--- /dev/null
+++ b/llvm/test/CodeGen/MLRegalloc/bb-profile-dump.ll
@@ -0,0 +1,21 @@
+; REQUIRES: have_tflite
+; REQUIRES: default_triple
+;
+; Check that the basic block profile dump outputs data and in the correct
+; format.
+;
+; RUN: llc -o /dev/null -mbb-profile-dump=%t %s
+; RUN: FileCheck --input-file %t %s
+
+define i64 @f2(i64 %a, i64 %b) {
+    %sum = add i64 %a, %b
+    ret i64 %sum
+}
+
+define i64 @f1() {
+    %sum = call i64 @f2(i64 2, i64 2)
+    ret i64 %sum
+}
+
+; CHECK: f2,0,1.000000e+00
+; CHECK-NEXT: f1,0,1.000000e+00