[MLGO] Change MBB Profile Dump from using MBB numbers to MBB IDs

[aidengrossman]

Summary

Currenty, setting the -mbb-profile-dump dumps a CSV file with blocks
inside an individual function identified by their MBB numbers. This
patch changes the MBBs to be identified by their ID which is set at MBB
creation and not changed afterwards, making it inherently stable
throughout the backend. This alleviates concerns with the MBB IDs
changing between the profile dump and what ends up in the final object
file. The MBBs inside the SHT_LLVM_BB_ADDR_MAP sections are also
identified using their MBB ID rather than number, so if we want to match
them up we need to identify the MBBs here by number.

[mtrofin]

lgtm, but wait for @rahmanl to comment, he's more familiar with these aspects.

[aidengrossman]

@rahmanl Bumping this for when you have a chance to take a look. Thanks!

[rahmanl]

Sorry for the delay.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

1939

Currently, these IDs are only generated with -basic-block-sections=labels|list, but IIUC -mbb-profile-dump does not require these flags. We can either have these IDs always generated, or add sanity checks to make sure -mbb-profile-dump is always used with -basic-block-sections.

llvm/test/CodeGen/MLRegalloc/bb-profile-dump.ll

6

It's interesting that this works. Maybe just use -basic-block-sections=labels to be safe.

[aidengrossman]

Address reviewer feedback.

[aidengrossman]

All good on the timing. Thanks for the review!

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

1939

We are requiring -basic-block-sections=labels according to the help string, but we aren't checking if both are enabled at the same time. I've added a check (and test) to ensure that the MachineFunction we're going over has BB labels. The check emits an error otherwise.

[rahmanl]

llvm/test/CodeGen/MLRegalloc/bb-profile-dump.ll

39

Remove this redundant line.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

 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 "
+             "performed with -basic-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";
 const char DbgTimerDescription[] = "Debug Info Emission";
 const char EHTimerName[] = "write_exception";
 const char EHTimerDescription[] = "DWARF Exception Writer";
void AsmPrinter::emitFunctionBody() {
 
   emitPatchableFunctionEntries();
 
   if (isVerbose())
     OutStreamer->getCommentOS() << "-- End function\n";
 
   OutStreamer->addBlankLine();
 
-  // Output MBB numbers, function names, and frequencies if the flag to dump
+  // Output MBB ids, function names, and frequencies if the flag to dump
   // MBB profile information has been set
   if (MBBProfileDumpFileOutput) {
+    if (!MF->hasBBLabels())
+      MF->getContext().reportError(
+          SMLoc(),
+          "Unable to find BB labels for MBB profile dump. -mbb-profile-dump "
+          "must be called with -basic-block-sections=labels");
     MachineBlockFrequencyInfo &MBFI =
         getAnalysis<LazyMachineBlockFrequencyInfoPass>().getBFI();
     for (const auto &MBB : *MF) {
       *MBBProfileDumpFileOutput.get()
-          << MF->getName() << "," << MBB.getNumber() << ","
+          << MF->getName() << "," << MBB.getBBID() << ","
           << MBFI.getBlockFreqRelativeToEntryBlock(&MBB) << "\n";
     }
   }
 }
 
 /// Compute the number of Global Variables that uses a Constant.
 static unsigned getNumGlobalVariableUses(const Constant *C) {
   if (!C)

llvm/test/CodeGen/MLRegalloc/bb-profile-dump.ll

 ; REQUIRES: x86-registered-target
 ;
 ; Check that the basic block profile dump outputs data and in the correct
 ; format.
 ;
-; RUN: llc -mtriple=x86_64-linux-unknown -o /dev/null -mbb-profile-dump=- %s | FileCheck %s
+; RUN: llc -mtriple=x86_64-linux-unknown -o /dev/null -basic-block-sections=labels -mbb-profile-dump=- %s | FileCheck %s
 
 ; bb profile dump is not supported on NVPTX
 ; UNSUPPORTED: target=nvptx{{.*}}
 
+; Check that given a simple case, we can return the default MBFI
+
 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)
     %isEqual = icmp eq i64 %sum, 4
     br i1 %isEqual, label %ifEqual, label %ifNotEqual
 ifEqual:
     ret i64 0
 ifNotEqual:
     ret i64 %sum
 }
 
 ; CHECK: f2,0,1.000000e+00
 ; CHECK-NEXT: f1,0,1.000000e+00
 ; CHECK-NEXT: f1,1,5.000000e-01
 ; CHECK-NEXT: f1,2,1.000000e+00
+
+; Check that if we pass -mbb-profile-dump but don't set -basic-block-sections,
+; we get an appropriate error message
+
+; RUN: not llc -mtriple=x86_64-linux-unknown -o /dev/null -mbb-profile-dump=- %s 2>&1 | FileCheck --check-prefix=NO-SECTIONS %s
+
+; NO-SECTIONS: <unknown>:0: error: Unable to find BB labels for MBB profile dump. -mbb-profile-dump must be called with -basic-block-sections=labels
+