Merge "simpleperf: Support LBR data in branch_list.proto" into main am: 39913fe999tmp_amf_315507370

Original change: https://android-review.googlesource.com/c/platform/system/extras/+/2864660

Change-Id: I715e8f9d41cfeef292fa8bdabb5a3409d3e24e08
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

6 files changed, 372 insertions, 124 deletions

diff --git a/simpleperf/BranchListFile.cpp b/simpleperf/BranchListFile.cpp
index e8f6b22c..253dcdbc 100644
--- a/simpleperf/BranchListFile.cpp
+++ b/simpleperf/BranchListFile.cpp
@@ -132,32 +132,8 @@ static UnorderedETMBranchMap BuildUnorderedETMBranchMap(const proto::ETMBinary&
 }
 
 bool StringToETMBinaryMap(const std::string& s, ETMBinaryMap& binary_map) {
-  proto::BranchList branch_list_proto;
-  if (!branch_list_proto.ParseFromString(s)) {
-    PLOG(ERROR) << "failed to read ETMBranchList msg";
-    return false;
-  }
-  if (branch_list_proto.magic() != ETM_BRANCH_LIST_PROTO_MAGIC) {
-    PLOG(ERROR) << "not in etm branch list format in branch_list.proto";
-    return false;
-  }
-
-  for (size_t i = 0; i < branch_list_proto.etm_data_size(); i++) {
-    const auto& binary_proto = branch_list_proto.etm_data(i);
-    BinaryKey key(binary_proto.path(), BuildId(binary_proto.build_id()));
-    if (binary_proto.has_kernel_info()) {
-      key.kernel_start_addr = binary_proto.kernel_info().kernel_start_addr();
-    }
-    ETMBinary& binary = binary_map[key];
-    auto dso_type = ToDsoType(binary_proto.type());
-    if (!dso_type) {
-      LOG(ERROR) << "invalid binary type " << binary_proto.type();
-      return false;
-    }
-    binary.dso_type = dso_type.value();
-    binary.branch_map = BuildUnorderedETMBranchMap(binary_proto);
-  }
-  return true;
+  LBRData lbr_data;
+  return ParseBranchListData(s, binary_map, lbr_data);
 }
 
 class ETMThreadTreeWhenRecording : public ETMThreadTree {
@@ -382,4 +358,83 @@ std::unique_ptr<ETMBranchListGenerator> ETMBranchListGenerator::Create(bool dump
 
 ETMBranchListGenerator::~ETMBranchListGenerator() {}
 
+bool LBRDataToString(const LBRData& data, std::string& s) {
+  proto::BranchList branch_list_proto;
+  branch_list_proto.set_magic(ETM_BRANCH_LIST_PROTO_MAGIC);
+  auto lbr_proto = branch_list_proto.mutable_lbr_data();
+  for (const LBRSample& sample : data.samples) {
+    auto sample_proto = lbr_proto->add_samples();
+    sample_proto->set_binary_id(sample.binary_id);
+    sample_proto->set_vaddr_in_file(sample.vaddr_in_file);
+    for (const LBRBranch& branch : sample.branches) {
+      auto branch_proto = sample_proto->add_branches();
+      branch_proto->set_from_binary_id(branch.from_binary_id);
+      branch_proto->set_to_binary_id(branch.to_binary_id);
+      branch_proto->set_from_vaddr_in_file(branch.from_vaddr_in_file);
+      branch_proto->set_to_vaddr_in_file(branch.to_vaddr_in_file);
+    }
+  }
+  for (const BinaryKey& binary : data.binaries) {
+    auto binary_proto = lbr_proto->add_binaries();
+    binary_proto->set_path(binary.path);
+    binary_proto->set_build_id(binary.build_id.ToString().substr(2));
+  }
+  if (!branch_list_proto.SerializeToString(&s)) {
+    LOG(ERROR) << "failed to serialize lbr data";
+    return false;
+  }
+  return true;
+}
+
+bool ParseBranchListData(const std::string& s, ETMBinaryMap& etm_data, LBRData& lbr_data) {
+  proto::BranchList branch_list_proto;
+  if (!branch_list_proto.ParseFromString(s)) {
+    PLOG(ERROR) << "failed to read ETMBranchList msg";
+    return false;
+  }
+  if (branch_list_proto.magic() != ETM_BRANCH_LIST_PROTO_MAGIC) {
+    PLOG(ERROR) << "not in etm branch list format in branch_list.proto";
+    return false;
+  }
+  for (size_t i = 0; i < branch_list_proto.etm_data_size(); i++) {
+    const auto& binary_proto = branch_list_proto.etm_data(i);
+    BinaryKey key(binary_proto.path(), BuildId(binary_proto.build_id()));
+    if (binary_proto.has_kernel_info()) {
+      key.kernel_start_addr = binary_proto.kernel_info().kernel_start_addr();
+    }
+    ETMBinary& binary = etm_data[key];
+    auto dso_type = ToDsoType(binary_proto.type());
+    if (!dso_type) {
+      LOG(ERROR) << "invalid binary type " << binary_proto.type();
+      return false;
+    }
+    binary.dso_type = dso_type.value();
+    binary.branch_map = BuildUnorderedETMBranchMap(binary_proto);
+  }
+  if (branch_list_proto.has_lbr_data()) {
+    const auto& lbr_data_proto = branch_list_proto.lbr_data();
+    lbr_data.samples.resize(lbr_data_proto.samples_size());
+    for (size_t i = 0; i < lbr_data_proto.samples_size(); ++i) {
+      const auto& sample_proto = lbr_data_proto.samples(i);
+      LBRSample& sample = lbr_data.samples[i];
+      sample.binary_id = sample_proto.binary_id();
+      sample.vaddr_in_file = sample_proto.vaddr_in_file();
+      sample.branches.resize(sample_proto.branches_size());
+      for (size_t j = 0; j < sample_proto.branches_size(); ++j) {
+        const auto& branch_proto = sample_proto.branches(j);
+        LBRBranch& branch = sample.branches[j];
+        branch.from_binary_id = branch_proto.from_binary_id();
+        branch.to_binary_id = branch_proto.to_binary_id();
+        branch.from_vaddr_in_file = branch_proto.from_vaddr_in_file();
+        branch.to_vaddr_in_file = branch_proto.to_vaddr_in_file();
+      }
+    }
+    for (size_t i = 0; i < lbr_data_proto.binaries_size(); ++i) {
+      const auto& binary_proto = lbr_data_proto.binaries(i);
+      lbr_data.binaries.emplace_back(binary_proto.path(), BuildId(binary_proto.build_id()));
+    }
+  }
+  return true;
+}
+
 }  // namespace simpleperf
diff --git a/simpleperf/BranchListFile.h b/simpleperf/BranchListFile.h
index 031a026f..de97ef88 100644
--- a/simpleperf/BranchListFile.h
+++ b/simpleperf/BranchListFile.h
@@ -129,7 +129,6 @@ struct ETMBinary {
 };
 
 using ETMBinaryMap = std::unordered_map<BinaryKey, ETMBinary, BinaryKeyHash>;
-
 bool ETMBinaryMapToString(const ETMBinaryMap& binary_map, std::string& s);
 bool StringToETMBinaryMap(const std::string& s, ETMBinaryMap& binary_map);
 
@@ -167,6 +166,9 @@ struct LBRData {
   std::vector<BinaryKey> binaries;
 };
 
+bool LBRDataToString(const LBRData& data, std::string& s);
+bool ParseBranchListData(const std::string& s, ETMBinaryMap& etm_data, LBRData& lbr_data);
+
 // for testing
 std::string ETMBranchToProtoString(const std::vector<bool>& branch);
 std::vector<bool> ProtoStringToETMBranch(const std::string& s, size_t bit_size);
diff --git a/simpleperf/branch_list.proto b/simpleperf/branch_list.proto
index 29247933..6d6bfb17 100644
--- a/simpleperf/branch_list.proto
+++ b/simpleperf/branch_list.proto
@@ -26,6 +26,7 @@ message BranchList {
   // Should always be "simpleperf:EtmBranchList".
   string magic = 1;
   repeated ETMBinary etm_data = 2;
+  LBRData lbr_data = 3;
 }
 
 message ETMBinary {
@@ -65,3 +66,31 @@ message ETMBinary {
 
   KernelBinaryInfo kernel_info = 5;
 }
+
+message LBRData {
+  repeated Sample samples = 1;
+  repeated Binary binaries = 2;
+
+  message Sample {
+    // If binary_id >= 1, it refers to LBRData.binaries[binary_id - 1]. Otherwise, it's invalid.
+    uint32 binary_id = 1;
+    uint64 vaddr_in_file = 2;
+    repeated Branch branches = 3;
+
+    message Branch {
+      // If from_binary_id >= 1, it refers to LBRData.binaries[from_binary_id - 1]. Otherwise, it's
+      // invalid.
+      uint32 from_binary_id = 1;
+      // If to_binary_id >= 1, it refers to LBRData.binaries[to_binary_id - 1]. Otherwise, it's
+      // invalid.
+      uint32 to_binary_id = 2;
+      uint64 from_vaddr_in_file = 3;
+      uint64 to_vaddr_in_file = 4;
+    }
+  }
+
+  message Binary {
+    string path = 1;
+    string build_id = 2;
+  }
+}
diff --git a/simpleperf/cmd_inject.cpp b/simpleperf/cmd_inject.cpp
index be8f4173..182f4fd2 100644
--- a/simpleperf/cmd_inject.cpp
+++ b/simpleperf/cmd_inject.cpp
@@ -146,9 +146,9 @@ class PerfDataReader {
 
   std::string GetDataType() const { return GetDataType(*reader_); }
 
-  void SetCallback(const AutoFDOBinaryCallback& callback) { autofdo_callback_ = callback; }
-  virtual void SetCallback(const ETMBinaryCallback&) {}
-  virtual void SetCallback(const LBRDataCallback&) {}
+  void AddCallback(const AutoFDOBinaryCallback& callback) { autofdo_callback_ = callback; }
+  void AddCallback(const ETMBinaryCallback& callback) { etm_binary_callback_ = callback; }
+  void AddCallback(const LBRDataCallback& callback) { lbr_data_callback_ = callback; }
 
   virtual bool Read() {
     if (exclude_perf_) {
@@ -196,6 +196,8 @@ class PerfDataReader {
   std::optional<int> exclude_pid_;
   ThreadTree thread_tree_;
   AutoFDOBinaryCallback autofdo_callback_;
+  ETMBinaryCallback etm_binary_callback_;
+  LBRDataCallback lbr_data_callback_;
   // Store results for AutoFDO.
   std::unordered_map<const Dso*, AutoFDOBinaryInfo> autofdo_binary_map_;
 };
@@ -231,8 +233,6 @@ class ETMPerfDataReader : public PerfDataReader {
         etm_dump_option_(etm_dump_option),
         etm_thread_tree_(thread_tree_, exclude_pid_) {}
 
-  void SetCallback(const ETMBinaryCallback& callback) override { etm_binary_callback_ = callback; }
-
   bool Read() override {
     if (reader_->HasFeature(PerfFileFormat::FEAT_ETM_BRANCH_LIST)) {
       return ProcessETMBranchListFeature();
@@ -363,7 +363,6 @@ class ETMPerfDataReader : public PerfDataReader {
   }
 
   ETMDumpOption etm_dump_option_;
-  ETMBinaryCallback etm_binary_callback_;
   ETMThreadTreeWithFilter etm_thread_tree_;
   std::vector<uint8_t> aux_data_buffer_;
   std::unique_ptr<ETMDecoder> etm_decoder_;
@@ -372,12 +371,49 @@ class ETMPerfDataReader : public PerfDataReader {
   std::unordered_map<Dso*, ETMBinary> etm_binary_map_;
 };
 
+static std::optional<std::vector<AutoFDOBinaryInfo>> ConvertLBRDataToAutoFDO(
+    const LBRData& lbr_data) {
+  std::vector<AutoFDOBinaryInfo> binaries(lbr_data.binaries.size());
+  for (const LBRSample& sample : lbr_data.samples) {
+    if (sample.binary_id != 0) {
+      if (sample.binary_id > binaries.size()) {
+        LOG(ERROR) << "binary_id out of range";
+        return std::nullopt;
+      }
+      binaries[sample.binary_id - 1].AddAddress(sample.vaddr_in_file);
+    }
+    for (size_t i = 0; i < sample.branches.size(); ++i) {
+      const LBRBranch& branch = sample.branches[i];
+      if (branch.from_binary_id == 0) {
+        continue;
+      }
+      if (branch.from_binary_id > binaries.size()) {
+        LOG(ERROR) << "binary_id out of range";
+        return std::nullopt;
+      }
+      if (branch.from_binary_id == branch.to_binary_id) {
+        binaries[branch.from_binary_id - 1].AddBranch(branch.from_vaddr_in_file,
+                                                      branch.to_vaddr_in_file);
+      }
+      if (i > 0 && branch.from_binary_id == sample.branches[i - 1].to_binary_id) {
+        uint64_t begin = sample.branches[i - 1].to_vaddr_in_file;
+        uint64_t end = branch.from_vaddr_in_file;
+        // Use the same logic to skip bogus LBR data as AutoFDO.
+        if (end < begin || end - begin > (1 << 20)) {
+          continue;
+        }
+        binaries[branch.from_binary_id - 1].AddRange(begin, end);
+      }
+    }
+  }
+  return binaries;
+}
+
 class LBRPerfDataReader : public PerfDataReader {
  public:
   LBRPerfDataReader(std::unique_ptr<RecordFileReader> reader, bool exclude_perf,
                     const RegEx* binary_name_regex)
       : PerfDataReader(std::move(reader), exclude_perf, binary_name_regex) {}
-  void SetCallback(const LBRDataCallback& callback) override { lbr_data_callback_ = callback; }
 
  private:
   bool ProcessRecord(Record& r) override {
@@ -419,8 +455,16 @@ class LBRPerfDataReader : public PerfDataReader {
 
   bool PostProcess() override {
     if (autofdo_callback_) {
-      ConvertLBRDataToAutoFDO();
+      std::optional<std::vector<AutoFDOBinaryInfo>> binaries = ConvertLBRDataToAutoFDO(lbr_data_);
+      if (!binaries) {
+        return false;
+      }
+      for (const auto& [dso, binary_id] : dso_map_) {
+        autofdo_binary_map_[dso] = std::move(binaries.value()[binary_id - 1]);
+      }
       ProcessAutoFDOBinaryInfo();
+    } else if (lbr_data_callback_) {
+      lbr_data_callback_(lbr_data_);
     }
     return true;
   }
@@ -440,55 +484,25 @@ class LBRPerfDataReader : public PerfDataReader {
     if (auto it = dso_map_.find(dso); it != dso_map_.end()) {
       return it->second;
     }
-    uint32_t binary_id = static_cast<uint32_t>(lbr_data_.binaries.size() + 1);
+    lbr_data_.binaries.emplace_back(dso, 0);
+    uint32_t binary_id = static_cast<uint32_t>(lbr_data_.binaries.size());
     dso_map_[dso] = binary_id;
     return binary_id;
   }
 
-  void ConvertLBRDataToAutoFDO() {
-    std::vector<AutoFDOBinaryInfo> binaries(dso_map_.size());
-    for (const LBRSample& sample : lbr_data_.samples) {
-      if (sample.binary_id != 0) {
-        binaries[sample.binary_id - 1].AddAddress(sample.vaddr_in_file);
-      }
-      for (size_t i = 0; i < sample.branches.size(); ++i) {
-        const LBRBranch& branch = sample.branches[i];
-        if (branch.from_binary_id == 0) {
-          continue;
-        }
-        if (branch.from_binary_id == branch.to_binary_id) {
-          binaries[branch.from_binary_id - 1].AddBranch(branch.from_vaddr_in_file,
-                                                        branch.to_vaddr_in_file);
-        }
-        if (i > 0 && branch.from_binary_id == sample.branches[i - 1].to_binary_id) {
-          uint64_t begin = sample.branches[i - 1].to_vaddr_in_file;
-          uint64_t end = branch.from_vaddr_in_file;
-          // Use the same logic to skip bogus LBR data as AutoFDO.
-          if (end < begin || end - begin > (1 << 20)) {
-            continue;
-          }
-          binaries[branch.from_binary_id - 1].AddRange(begin, end);
-        }
-      }
-    }
-    for (const auto& [dso, binary_id] : dso_map_) {
-      autofdo_binary_map_[dso] = std::move(binaries[binary_id - 1]);
-    }
-  }
-
-  LBRDataCallback lbr_data_callback_;
   LBRData lbr_data_;
   // Map from dso to binary_id in lbr_data_.
   std::unordered_map<const Dso*, uint32_t> dso_map_;
 };
 
-// Read a protobuf file specified by etm_branch_list.proto, and generate ETMBinary.
-class ETMBranchListReader {
+// Read a protobuf file specified by branch_list.proto.
+class BranchListReader {
  public:
-  ETMBranchListReader(const std::string& filename, const RegEx* binary_name_regex)
+  BranchListReader(const std::string& filename, const RegEx* binary_name_regex)
       : filename_(filename), binary_filter_(binary_name_regex) {}
 
-  void SetCallback(const ETMBinaryCallback& callback) { callback_ = callback; }
+  void AddCallback(const ETMBinaryCallback& callback) { etm_binary_callback_ = callback; }
+  void AddCallback(const LBRDataCallback& callback) { lbr_data_callback_ = callback; }
 
   bool Read() {
     std::string s;
@@ -496,24 +510,84 @@ class ETMBranchListReader {
       PLOG(ERROR) << "failed to read " << filename_;
       return false;
     }
-    ETMBinaryMap binary_map;
-    if (!StringToETMBinaryMap(s, binary_map)) {
+    ETMBinaryMap etm_data;
+    LBRData lbr_data;
+    if (!ParseBranchListData(s, etm_data, lbr_data)) {
       PLOG(ERROR) << "file is in wrong format: " << filename_;
       return false;
     }
-    for (auto& [key, binary] : binary_map) {
+    if (etm_binary_callback_ && !etm_data.empty()) {
+      ProcessETMData(etm_data);
+    }
+    if (lbr_data_callback_ && !lbr_data.samples.empty()) {
+      ProcessLBRData(lbr_data);
+    }
+    return true;
+  }
+
+ private:
+  void ProcessETMData(ETMBinaryMap& etm_data) {
+    for (auto& [key, binary] : etm_data) {
       if (!binary_filter_.Filter(key.path)) {
         continue;
       }
-      callback_(key, binary);
+      etm_binary_callback_(key, binary);
     }
-    return true;
   }
 
- private:
+  void ProcessLBRData(LBRData& lbr_data) {
+    // 1. Check if we need to remove binaries.
+    std::vector<uint32_t> new_ids(lbr_data.binaries.size());
+    uint32_t next_id = 1;
+
+    for (size_t i = 0; i < lbr_data.binaries.size(); ++i) {
+      if (!binary_filter_.Filter(lbr_data.binaries[i].path)) {
+        new_ids[i] = 0;
+      } else {
+        new_ids[i] = next_id++;
+      }
+    }
+
+    if (next_id <= lbr_data.binaries.size()) {
+      // 2. Modify lbr_data.binaries.
+      for (size_t i = 0; i < lbr_data.binaries.size(); ++i) {
+        if (new_ids[i] != 0) {
+          size_t new_pos = new_ids[i] - 1;
+          lbr_data.binaries[new_pos] = lbr_data.binaries[i];
+        }
+      }
+      lbr_data.binaries.resize(next_id - 1);
+
+      // 3. Modify lbr_data.samples.
+      auto convert_id = [&](uint32_t& binary_id) {
+        if (binary_id != 0) {
+          binary_id = (binary_id <= new_ids.size()) ? new_ids[binary_id - 1] : 0;
+        }
+      };
+      std::vector<LBRSample> new_samples;
+      for (LBRSample& sample : lbr_data.samples) {
+        convert_id(sample.binary_id);
+        bool has_valid_binary_id = sample.binary_id != 0;
+        for (LBRBranch& branch : sample.branches) {
+          convert_id(branch.from_binary_id);
+          convert_id(branch.to_binary_id);
+          if (branch.from_binary_id != 0 || branch.to_binary_id != 0) {
+            has_valid_binary_id = true;
+          }
+        }
+        if (has_valid_binary_id) {
+          new_samples.emplace_back(std::move(sample));
+        }
+      }
+      lbr_data.samples = std::move(new_samples);
+    }
+    lbr_data_callback_(lbr_data);
+  }
+
   const std::string filename_;
   BinaryFilter binary_filter_;
-  ETMBinaryCallback callback_;
+  ETMBinaryCallback etm_binary_callback_;
+  LBRDataCallback lbr_data_callback_;
 };
 
 // Convert ETMBinary into AutoFDOBinaryInfo.
@@ -662,42 +736,82 @@ class AutoFDOWriter {
   std::unordered_map<BinaryKey, AutoFDOBinaryInfo, BinaryKeyHash> binary_map_;
 };
 
-// Merge ETMBinary.
-struct ETMBranchListMerger {
+// Merge branch list data.
+struct BranchListMerger {
   void AddETMBinary(const BinaryKey& key, ETMBinary& binary) {
-    auto it = binary_map.find(key);
-    if (it == binary_map.end()) {
-      binary_map[key] = std::move(binary);
-    } else {
+    if (auto it = etm_data_.find(key); it != etm_data_.end()) {
       it->second.Merge(binary);
+    } else {
+      etm_data_[key] = std::move(binary);
     }
   }
 
-  ETMBinaryMap binary_map;
-};
+  void AddLBRData(LBRData& lbr_data) {
+    // 1. Merge binaries.
+    std::vector<uint32_t> new_ids(lbr_data.binaries.size());
+    for (size_t i = 0; i < lbr_data.binaries.size(); i++) {
+      const BinaryKey& key = lbr_data.binaries[i];
+      if (auto it = lbr_binary_id_map_.find(key); it != lbr_binary_id_map_.end()) {
+        new_ids[i] = it->second;
+      } else {
+        uint32_t next_id = static_cast<uint32_t>(lbr_binary_id_map_.size()) + 1;
+        new_ids[i] = next_id;
+        lbr_binary_id_map_[key] = next_id;
+        lbr_data_.binaries.emplace_back(key);
+      }
+    }
 
-// Write branch lists to a protobuf file specified by etm_branch_list.proto.
-class ETMBranchListWriter {
- public:
-  bool Write(const std::string& output_filename, const ETMBinaryMap& binary_map) {
-    // Don't produce empty output file.
-    if (binary_map.empty()) {
-      LOG(INFO) << "Skip empty output file.";
-      unlink(output_filename.c_str());
-      return true;
+    // 2. Merge samples.
+    auto convert_id = [&](uint32_t& binary_id) {
+      if (binary_id != 0) {
+        binary_id = (binary_id <= new_ids.size()) ? new_ids[binary_id - 1] : 0;
+      }
+    };
+
+    for (LBRSample& sample : lbr_data.samples) {
+      convert_id(sample.binary_id);
+      for (LBRBranch& branch : sample.branches) {
+        convert_id(branch.from_binary_id);
+        convert_id(branch.to_binary_id);
+      }
+      lbr_data_.samples.emplace_back(std::move(sample));
     }
-    std::string s;
-    if (!ETMBinaryMapToString(binary_map, s)) {
-      LOG(ERROR) << "invalid ETMBinaryMap";
+  }
+
+  ETMBinaryMap& GetETMData() { return etm_data_; }
+
+  LBRData& GetLBRData() { return lbr_data_; }
+
+ private:
+  ETMBinaryMap etm_data_;
+  LBRData lbr_data_;
+  std::unordered_map<BinaryKey, uint32_t, BinaryKeyHash> lbr_binary_id_map_;
+};
+
+// Write branch lists to a protobuf file specified by branch_list.proto.
+static bool WriteBranchListFile(const std::string& output_filename, const ETMBinaryMap& etm_data,
+                                const LBRData& lbr_data) {
+  std::string s;
+  if (!etm_data.empty()) {
+    if (!ETMBinaryMapToString(etm_data, s)) {
       return false;
     }
-    if (!android::base::WriteStringToFile(s, output_filename)) {
-      PLOG(ERROR) << "failed to write to " << output_filename;
+  } else if (!lbr_data.samples.empty()) {
+    if (!LBRDataToString(lbr_data, s)) {
       return false;
     }
+  } else {
+    // Don't produce empty output file.
+    LOG(INFO) << "Skip empty output file.";
+    unlink(output_filename.c_str());
     return true;
   }
-};
+  if (!android::base::WriteStringToFile(s, output_filename)) {
+    PLOG(ERROR) << "failed to write to " << output_filename;
+    return false;
+  }
+  return true;
+}
 
 class InjectCommand : public Command {
  public:
@@ -877,7 +991,7 @@ class InjectCommand : public Command {
     auto afdo_callback = [&](const BinaryKey& key, AutoFDOBinaryInfo& binary) {
       autofdo_writer.AddAutoFDOBinary(key, binary);
     };
-    auto reader_callback = [&](PerfDataReader& reader) { reader.SetCallback(afdo_callback); };
+    auto reader_callback = [&](PerfDataReader& reader) { reader.AddCallback(afdo_callback); };
     if (!ReadPerfDataFiles(reader_callback)) {
       return false;
     }
@@ -885,36 +999,42 @@ class InjectCommand : public Command {
   }
 
   bool ConvertPerfDataToBranchList() {
-    ETMBranchListMerger branch_list_merger;
+    BranchListMerger merger;
     auto etm_callback = [&](const BinaryKey& key, ETMBinary& binary) {
-      branch_list_merger.AddETMBinary(key, binary);
+      merger.AddETMBinary(key, binary);
+    };
+    auto lbr_callback = [&](LBRData& lbr_data) { merger.AddLBRData(lbr_data); };
+
+    auto reader_callback = [&](PerfDataReader& reader) {
+      reader.AddCallback(etm_callback);
+      reader.AddCallback(lbr_callback);
     };
-    auto reader_callback = [&](PerfDataReader& reader) { reader.SetCallback(etm_callback); };
     if (!ReadPerfDataFiles(reader_callback)) {
       return false;
     }
-    ETMBranchListWriter branch_list_writer;
-    return branch_list_writer.Write(output_filename_, branch_list_merger.binary_map);
+    return WriteBranchListFile(output_filename_, merger.GetETMData(), merger.GetLBRData());
   }
 
   bool ConvertBranchListToAutoFDO() {
     // Step1 : Merge branch lists from all input files.
-    ETMBranchListMerger branch_list_merger;
-    auto callback = [&](const BinaryKey& key, ETMBinary& binary) {
-      branch_list_merger.AddETMBinary(key, binary);
+    BranchListMerger merger;
+    auto etm_callback = [&](const BinaryKey& key, ETMBinary& binary) {
+      merger.AddETMBinary(key, binary);
     };
+    auto lbr_callback = [&](LBRData& lbr_data) { merger.AddLBRData(lbr_data); };
     for (const auto& input_filename : input_filenames_) {
-      ETMBranchListReader reader(input_filename, binary_name_regex_.get());
-      reader.SetCallback(callback);
+      BranchListReader reader(input_filename, binary_name_regex_.get());
+      reader.AddCallback(etm_callback);
+      reader.AddCallback(lbr_callback);
       if (!reader.Read()) {
         return false;
       }
     }
 
-    // Step2: Convert ETMBinary to AutoFDOBinaryInfo.
+    // Step2: Convert ETMBinary and LBRData to AutoFDOBinaryInfo.
     AutoFDOWriter autofdo_writer;
     ETMBranchListToAutoFDOConverter converter;
-    for (auto& p : branch_list_merger.binary_map) {
+    for (auto& p : merger.GetETMData()) {
       const BinaryKey& key = p.first;
       ETMBinary& binary = p.second;
       std::unique_ptr<AutoFDOBinaryInfo> autofdo_binary = converter.Convert(key, binary);
@@ -924,6 +1044,16 @@ class InjectCommand : public Command {
         autofdo_writer.AddAutoFDOBinary(BinaryKey(key.path, key.build_id), *autofdo_binary);
       }
     }
+    if (!merger.GetLBRData().samples.empty()) {
+      LBRData& lbr_data = merger.GetLBRData();
+      std::optional<std::vector<AutoFDOBinaryInfo>> binaries = ConvertLBRDataToAutoFDO(lbr_data);
+      if (!binaries) {
+        return false;
+      }
+      for (size_t i = 0; i < binaries.value().size(); ++i) {
+        autofdo_writer.AddAutoFDOBinary(lbr_data.binaries[i], binaries.value()[i]);
+      }
+    }
 
     // Step3: Write AutoFDOBinaryInfo.
     return autofdo_writer.Write(output_filename_);
@@ -931,20 +1061,21 @@ class InjectCommand : public Command {
 
   bool ConvertBranchListToBranchList() {
     // Step1 : Merge branch lists from all input files.
-    ETMBranchListMerger branch_list_merger;
-    auto callback = [&](const BinaryKey& key, ETMBinary& binary) {
-      branch_list_merger.AddETMBinary(key, binary);
+    BranchListMerger merger;
+    auto etm_callback = [&](const BinaryKey& key, ETMBinary& binary) {
+      merger.AddETMBinary(key, binary);
     };
+    auto lbr_callback = [&](LBRData& lbr_data) { merger.AddLBRData(lbr_data); };
     for (const auto& input_filename : input_filenames_) {
-      ETMBranchListReader reader(input_filename, binary_name_regex_.get());
-      reader.SetCallback(callback);
+      BranchListReader reader(input_filename, binary_name_regex_.get());
+      reader.AddCallback(etm_callback);
+      reader.AddCallback(lbr_callback);
       if (!reader.Read()) {
         return false;
       }
     }
     // Step2: Write ETMBinary.
-    ETMBranchListWriter branch_list_writer;
-    return branch_list_writer.Write(output_filename_, branch_list_merger.binary_map);
+    return WriteBranchListFile(output_filename_, merger.GetETMData(), merger.GetLBRData());
   }
 
   std::unique_ptr<RegEx> binary_name_regex_;
diff --git a/simpleperf/cmd_inject_test.cpp b/simpleperf/cmd_inject_test.cpp
index 3b25553b..bdc2820f 100644
--- a/simpleperf/cmd_inject_test.cpp
+++ b/simpleperf/cmd_inject_test.cpp
@@ -226,12 +226,40 @@ TEST(cmd_inject, accept_missing_aux_data) {
 }
 
 TEST(cmd_inject, read_lbr_data) {
+  // Convert perf.data to AutoFDO text format.
+  std::string perf_data_path = GetTestData("lbr/perf_lbr.data");
   std::string data;
-  ASSERT_TRUE(RunInjectCmd({"-i", GetTestData("lbr/perf_lbr.data")}, &data));
+  ASSERT_TRUE(RunInjectCmd({"-i", perf_data_path}, &data));
   data.erase(std::remove(data.begin(), data.end(), '\r'), data.end());
 
   std::string expected_data;
   ASSERT_TRUE(android::base::ReadFileToString(
       GetTestData(std::string("lbr") + OS_PATH_SEPARATOR + "inject_lbr.data"), &expected_data));
-  ASSERT_NE(data.find(expected_data), data.npos);
+  ASSERT_EQ(data, expected_data);
+
+  // Convert perf.data to branch_list.proto format.
+  // Then convert branch_list.proto format to AutoFDO text format.
+  TemporaryFile branch_list_file;
+  close(branch_list_file.release());
+  ASSERT_TRUE(
+      RunInjectCmd({"-i", perf_data_path, "--output", "branch-list", "-o", branch_list_file.path}));
+  ASSERT_TRUE(RunInjectCmd({"-i", branch_list_file.path}, &data));
+  ASSERT_EQ(data, expected_data);
+
+  // Test binary filter on LBR data.
+  ASSERT_TRUE(RunInjectCmd({"-i", perf_data_path, "--binary", "no_lbr_test_loop"}, &data));
+  ASSERT_EQ(data.find("lbr_test_loop"), data.npos);
+
+  // Test binary filter on branch list file.
+  ASSERT_TRUE(RunInjectCmd({"-i", branch_list_file.path, "--binary", "no_lbr_test_loop"}, &data));
+  ASSERT_EQ(data.find("lbr_test_loop"), data.npos);
+
+  // Test multiple input files.
+  ASSERT_TRUE(RunInjectCmd(
+      {
+          "-i",
+          std::string(branch_list_file.path) + "," + branch_list_file.path,
+      },
+      &data));
+  ASSERT_NE(data.find("194d->1940:706"), data.npos);
 }
diff --git a/simpleperf/testdata/lbr/inject_lbr.data b/simpleperf/testdata/lbr/inject_lbr.data
index ffc9b5b5..6f141648 100644
--- a/simpleperf/testdata/lbr/inject_lbr.data
+++ b/simpleperf/testdata/lbr/inject_lbr.data
@@ -9,3 +9,6 @@
 2
 191b->1910:32
 194d->1940:353
+// build_id: 0x0000000000000000000000000000000000000000
+// /home/yabinc/lbr_test_loop
+