Add extra checks to select the right mapping and filter samples within it.

gematria/datasets/annotating_importer.cc

@@ -73,17 +73,26 @@ absl::StatusOr<const quipper::PerfDataProto *> AnnotatingImporter::LoadPerfData(
   return &perf_reader_.proto();
 }
 
+namespace {
+
+std::string GetFileNameFromPath(const std::string &path) {
+  int idx = path.find_last_of('/');
+  if (idx == std::string::npos) {
+    return path;
+  }
+  return path.substr(idx + 1);
+}
+
+}  // namespace
+
 absl::StatusOr<const quipper::PerfDataProto_MMapEvent *>
 AnnotatingImporter::GetMainMapping(
     const llvm::object::ELFObjectFileBase *elf_object,
     const quipper::PerfDataProto *perf_data) {
-  // Find the relevant mapping.
-  // TODO(virajbshah): Make sure the mapping was found. (Use num_mmap_events)
+  std::string file_name = GetFileNameFromPath(elf_object->getFileName().str());
   for (const auto &event : perf_data->events()) {
-    // TODO(virajbshah): Not sure if this always works, i.e. does the main
-    // binary always correspond to the first MMapEvent. Implement BuildID or
-    // name based checking.
     if (event.has_mmap_event() &&
+        GetFileNameFromPath(event.mmap_event().filename()) == file_name &&
         event.mmap_event().prot() & 0b001 /* PROT_READ */ &&
         event.mmap_event().prot() & 0b100 /* PROT_EXEC */) {
       return &event.mmap_event();
@@ -199,8 +208,8 @@ AnnotatingImporter::GetBlocksFromELF(
     for (const llvm::object::BBAddrMap::BBRangeEntry &bb_range :
          map.getBBRanges()) {
       for (const llvm::object::BBAddrMap::BBEntry &bb : bb_range.BBEntries) {
-        uint64_t begin_idx = function_addr + bb.Offset,
-                 end_idx = begin_idx + bb.Size;
+        uint64_t begin_idx = function_addr + bb.Offset;
+        uint64_t end_idx = begin_idx + bb.Size;
         if (begin_idx == end_idx) {
           continue;  // Skip any empty basic blocks.
         }
@@ -258,6 +267,9 @@ AnnotatingImporter::GetSamples(
     }
 
     // Filter out sample events from outside the profiled binary.
+    if (!event.sample_event().has_pid() ||
+        !(event.sample_event().pid() == mapping->pid()))
+      continue;
     uint64_t sample_ip = event.sample_event().ip();
     if (sample_ip < mmap_begin_addr || sample_ip >= mmap_end_addr) continue;
 
@@ -320,6 +332,13 @@ AnnotatingImporter::GetLBRBlocksWithLatency(
         !event.sample_event().branch_stack_size()) {
       continue;
     }
+
+    // Check if the sample PID matches that of the relevant mapping.
+    if (!event.sample_event().has_pid() ||
+        !(event.sample_event().pid() == mapping->pid())) {
+      continue;
+    }
+
     const auto &branch_stack = event.sample_event().branch_stack();
     for (int branch_idx = branch_stack.size() - 2; branch_idx >= 0;
          --branch_idx) {
@@ -336,9 +355,6 @@ AnnotatingImporter::GetLBRBlocksWithLatency(
 
       // Remove blocks not belonging to the binary we are importing from.
       if (block_begin < mmap_begin_addr || mmap_end_addr < block_end) continue;
-      if (block_begin < main_header->p_offset ||
-          main_header->p_offset + main_header->p_filesz < block_end)
-        continue;
 
       uint32_t block_latency = next_branch_entry.cycles();
 

gematria/datasets/annotating_importer.h

@@ -66,6 +66,11 @@ class AnnotatingImporter {
 
   // Searches all MMap events for the one that most likely corresponds to the
   // executable load segment of the given object.
+  // This requires that the ELF object's filename has not changed from when it
+  // was profiled, since we check its name against the filenames from the
+  // recorded MMap events. Note the object file can still be moved, since we
+  // check only the name and not the path.
+  // TODO(virajbshah): Find a better way to identify the relevant mapping.
   absl::StatusOr<const quipper::PerfDataProto_MMapEvent*> GetMainMapping(
       const llvm::object::ELFObjectFileBase* elf_object,
       const quipper::PerfDataProto* perf_data);

gematria/datasets/python/import_annotated_basic_blocks.py

@@ -93,6 +93,7 @@ def main(argv: Sequence[str]) -> None:
   with tf.io.TFRecordWriter(_OUTPUT_TFRECORD_FILE.value) as writer:
     for proto in protos:
       writer.write(proto.SerializeToString())
+  print(f'Wrote {len(protos)} (pseudo-)basic block(s).')
 
 
 if __name__ == '__main__':