[BOLT] Add number of basic blocks to BAT

YAML profile reader checks the number of basic blocks in regular,
no-stale-matching mode. Add it to BAT.

This increases the size of BAT section to:
- large binary: 39583080 bytes (1.02x of the original),
- medium binary: 3816492 bytes (0.64x),
- small binary: 920 bytes (0.64x, no change due to alignment).

Test Plan: Updated bolt-address-translation-yaml.test

Reviewers: rafaelauler, ayermolo, maksfb, dcci

Reviewed By: rafaelauler

Pull Request: #86045

Author: aaupov

bolt/docs/BAT.md

@@ -79,6 +79,7 @@ Hot indices are delta encoded, implicitly starting at zero.
 | `Address` | Continuous, Delta, ULEB128 | Function address in the output binary |
 | `HotIndex` | Delta, ULEB128 | Cold functions only: index of corresponding hot function in hot functions table |
 | `FuncHash` | 8b | Hot functions only: function hash for input function |
+| `NumBlocks` | ULEB128 | Hot functions only: number of basic blocks in the original function |
 | `NumEntries` | ULEB128 | Number of address translation entries for a function |
 | `EqualElems` | ULEB128 | Hot functions only: number of equal offsets in the beginning of a function |
 | `BranchEntries` | Bitmask, `alignTo(EqualElems, 8)` bits | Hot functions only: if `EqualElems` is non-zero, bitmask denoting entries with `BRANCHENTRY` bit |

bolt/include/bolt/Profile/BoltAddressTranslation.h

@@ -162,6 +162,9 @@ class BoltAddressTranslation {
   using BBHashMap = std::unordered_map<uint32_t, std::pair<unsigned, size_t>>;
   std::unordered_map<uint64_t, std::pair<size_t, BBHashMap>> FuncHashes;
 
+  /// Map a function to its basic blocks count
+  std::unordered_map<uint64_t, size_t> NumBasicBlocksMap;
+
   /// Links outlined cold bocks to their original function
   std::map<uint64_t, uint64_t> ColdPartSource;
 

bolt/lib/Profile/BoltAddressTranslation.cpp

@@ -196,6 +196,10 @@ void BoltAddressTranslation::writeMaps(std::map<uint64_t, MapTy> &Maps,
       // Function hash
       LLVM_DEBUG(dbgs() << "Hash: " << formatv("{0:x}\n", FuncHashPair.first));
       OS.write(reinterpret_cast<char *>(&FuncHashPair.first), 8);
+      // Number of basic blocks
+      size_t NumBasicBlocks = FuncHashPair.second.size();
+      LLVM_DEBUG(dbgs() << "Basic blocks: " << NumBasicBlocks << '\n');
+      encodeULEB128(NumBasicBlocks, OS);
     }
     encodeULEB128(NumEntries, OS);
     // For hot fragments only: encode the number of equal offsets
@@ -293,6 +297,12 @@ void BoltAddressTranslation::parseMaps(std::vector<uint64_t> &HotFuncs,
       const size_t FuncHash = DE.getU64(&Offset, &Err);
       FuncHashes[Address].first = FuncHash;
       LLVM_DEBUG(dbgs() << formatv("{0:x}: hash {1:x}\n", Address, FuncHash));
+      // Number of basic blocks
+      const size_t NumBasicBlocks = DE.getULEB128(&Offset, &Err);
+      NumBasicBlocksMap.emplace(Address, NumBasicBlocks);
+      LLVM_DEBUG(dbgs() << formatv("{0:x}: #bbs {1}, {2} bytes\n", Address,
+                                   NumBasicBlocks,
+                                   getULEB128Size(NumBasicBlocks)));
     }
     const uint32_t NumEntries = DE.getULEB128(&Offset, &Err);
     // Equal offsets, hot fragments only.

bolt/test/X86/bolt-address-translation-yaml.test

@@ -18,7 +18,7 @@ RUN:   | FileCheck --check-prefix CHECK-BOLT-YAML %s
 
 WRITE-BAT-CHECK: BOLT-INFO: Wrote 5 BAT maps
 WRITE-BAT-CHECK: BOLT-INFO: Wrote 4 function and 22 basic block hashes
-WRITE-BAT-CHECK: BOLT-INFO: BAT section size (bytes): 376
+WRITE-BAT-CHECK: BOLT-INFO: BAT section size (bytes): 380
 
 READ-BAT-CHECK-NOT: BOLT-ERROR: unable to save profile in YAML format for input file processed by BOLT
 READ-BAT-CHECK: BOLT-INFO: Parsed 5 BAT entries