diff --git a/compiler_opt/memtrace_costmodel/basic_block_trace.cc b/compiler_opt/memtrace_costmodel/basic_block_trace.cc new file mode 100644 index 00000000..2ca92530 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/basic_block_trace.cc @@ -0,0 +1,1000 @@ +#include "compiler_opt/memtrace_costmodel/basic_block_trace.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "compiler_opt/memtrace_costmodel/elf_metadata_parser.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" +#include "compiler_opt/memtrace_costmodel/serialized_mbbs.proto.h" +#include "drmemtrace/analyzer.h" +#undef X86 +#undef X86_64 +#include "absl/algorithm/container.h" +#include "absl/container/btree_map.h" +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/functional/function_ref.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/strings/escaping.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" +#include "compiler_opt/memtrace_costmodel/status_macros.h" +#include "drmemtrace/memref.h" +#include "drmemtrace/trace_entry.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallSet.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/MC/MCContext.h" +#include "llvm/MC/MCDisassembler/MCDisassembler.h" +#include "llvm/MC/MCInstrDesc.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Object/Binary.h" +#include "llvm/Object/ELFObjectFile.h" +#include "llvm/Object/ELFTypes.h" +#include "llvm/Object/ObjectFile.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Error.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/TargetParser/Triple.h" +#include "riegeli/bytes/file_writer.h" +#include "riegeli/records/record_writer.h" + +namespace mlgo { +namespace latency_model { +namespace { +constexpr int kBbTraceCompressionLevel = 1 << 16; +constexpr absl::string_view kBbTraceFileName = "bb_trace.pb"; +constexpr absl::string_view kFunctionIndexFileName = "function_index.pb"; + +template +bool LlvmExpectedSucceeded(llvm::Expected& expected) { + return static_cast(expected); +} + +std::vector GetSectionsWithBbInfoFromBinary( + absl::string_view binary_path, + const absl::flat_hash_map& + bb_addresses_to_ids_and_functions) { + llvm::Expected> + application_binary = llvm::object::createBinary(binary_path); + + QCHECK(LlvmExpectedSucceeded(application_binary)); + + llvm::object::ELFObjectFileBase* object_file = + llvm::cast( + application_binary->getBinary()); + + std::vector sections_with_bb_info; + + for (const llvm::object::SectionRef& section : object_file->sections()) { + if (!section.isText()) { + continue; + } + + uint64_t section_address = section.getAddress(); + uint64_t section_size = section.getSize(); + + bool contains_bb_address = false; + for (const auto& [bb_address, _] : bb_addresses_to_ids_and_functions) { + if (bb_address >= section_address && + bb_address < section_address + section_size) { + contains_bb_address = true; + break; + } + } + + if (contains_bb_address) { + sections_with_bb_info.push_back( + {.address = section_address, .size = section_size}); + } + } + + return sections_with_bb_info; +} + +bool InsideTextSectionWithBbAddrMap(std::vector& section_info, + uint64_t current_address) { + for (const SectionInfo& current_section : section_info) { + if (current_section.address <= current_address && + current_section.address + current_section.size > current_address) { + return true; + } + } + return false; +} + +llvm::MCInst getInstructionFromBytes(llvm::ArrayRef instruction_data, + llvm::MCDisassembler& disassembler, + uint64_t& instruction_size, + uint64_t instruction_address = 0) { + llvm::MCInst instruction; + std::string disassembler_output_buffer; + llvm::raw_string_ostream output_stream(disassembler_output_buffer); + + const llvm::MCDisassembler::DecodeStatus status = disassembler.getInstruction( + instruction, instruction_size, instruction_data, instruction_address, + output_stream); + output_stream.flush(); + if (status != llvm::MCDisassembler::DecodeStatus::Success) { + LOG(WARNING) << "Failed to disassemble instruction at " + << instruction_address << " status: " << status + << " output: " << disassembler_output_buffer; + return llvm::MCInst(); + } + return instruction; +} +} // namespace + +uint64_t GetElfOffset( + const dynamorio::drmemtrace::module_mapper_t* module_mapper, + uint64_t runtime_pc) { + const auto& modules = + const_cast(module_mapper) + ->get_loaded_modules(); + for (const auto& m : modules) { + if (runtime_pc >= reinterpret_cast(m.orig_seg_base) && + runtime_pc < reinterpret_cast(m.orig_seg_base) + m.seg_size) { + return runtime_pc - reinterpret_cast(m.orig_seg_base) + + m.seg_offs; + } + } + return 0; // not found +} + +void* BBMemtraceProcessor::parallel_shard_init(int shard_index, + void* worker_data) { + PerShardData* current_shard_data = new PerShardData(); + LOG(INFO) << "Shard init " << shard_index << " data: " << current_shard_data; + return reinterpret_cast(current_shard_data); +} + +bool BBMemtraceProcessor::parallel_shard_exit(void* shard_data) { + PerShardData* current_shard_data = + reinterpret_cast(shard_data); + + { + absl::MutexLock lock(¤t_shard_data->mutex); + if (!current_shard_data->current_trace_data.mbbs().empty() || + !current_shard_data->current_trace_data.shared_object_traces() + .empty()) { + absl::MutexLock lock(&bb_trace_processor_mutex_); + bb_trace_processor_(current_shard_data->current_trace_data); + } + + { + absl::MutexLock function_name_to_id_lock(&function_name_to_id_mutex_); + function_name_to_id_.mutable_function_ids()->insert( + current_shard_data->function_name_to_id.begin(), + current_shard_data->function_name_to_id.end()); + } + } + + delete current_shard_data; + return true; +} + +bool BBMemtraceProcessor::process_memref( + const dynamorio::drmemtrace::memref_t& memref) { + LOG(QFATAL) << "Intentionally not implemented"; + return false; +} + +bool BBMemtraceProcessor::parallel_shard_memref( + void* shard_data, const dynamorio::drmemtrace::memref_t& memref) { + PerShardData* current_shard_data = + reinterpret_cast(shard_data); + absl::MutexLock lock(¤t_shard_data->mutex); + + if (!dynamorio::drmemtrace::type_is_instr(memref.instr.type)) return true; + + // Size-based segment splitting + if (!current_shard_data->inside_shared_object && + current_shard_data->current_trace_data.mbbs_size() >= + max_blocks_per_segment_) { + { + absl::MutexLock lock(bb_trace_processor_mutex_); + bb_trace_processor_(current_shard_data->current_trace_data); + } + current_shard_data->current_trace_data.Clear(); + } + + // Resolve ELF address using custom module mapping + uint64_t elf_pc = GetElfOffset(module_mapper_, memref.instr.addr); + if (elf_pc == 0) return true; // Ignore unmapped code + + // Handle shared object boundary crossing + if (!current_shard_data->inside_shared_object && + !InsideTextSectionWithBbAddrMap(sections_with_bb_info_, elf_pc)) { + current_shard_data->inside_shared_object = true; + + // We have just entered the shared object. Create a new shared object trace. + current_shard_data->current_trace_data.add_shared_object_traces(); + + { + absl::MutexLock current_trace_id_lock(current_trace_id_mutex_); + current_shard_data->current_trace_data.mutable_shared_object_traces() + ->rbegin() + ->set_trace_id(current_trace_id_); + ++current_trace_id_; + } + } + + if (current_shard_data->inside_shared_object && + InsideTextSectionWithBbAddrMap(sections_with_bb_info_, elf_pc)) { + current_shard_data->inside_shared_object = false; + + // Left shared object, create a block referencing the trace + MachineBbId* trace_bb_id = + current_shard_data->current_trace_data.add_mbbs(); + trace_bb_id->set_basic_block_id( + current_shard_data->current_trace_data.shared_object_traces() + .rbegin() + ->trace_id()); + } + + if (current_shard_data->inside_shared_object) { + std::string current_instruction_data( + reinterpret_cast(memref.instr.encoding), + memref.instr.size); + + current_shard_data->current_trace_data.mutable_shared_object_traces() + ->rbegin() + ->add_instruction_data(std::move(current_instruction_data)); + } + + auto function_address_and_bb_id = + bb_addresses_to_ids_and_functions_.find(elf_pc); + if (function_address_and_bb_id != bb_addresses_to_ids_and_functions_.end()) { + MachineBbId* next_mbb = current_shard_data->current_trace_data.add_mbbs(); + (*next_mbb) = function_address_and_bb_id->second; + + current_shard_data->function_name_to_id.try_emplace( + function_id_to_name_[next_mbb->function_id()], next_mbb->function_id()); + } + + return true; +} + +absl::Status GetBasicBlockTracesFromDirectory( + absl::string_view trace_dir, absl::string_view binary_path, + absl::FunctionRef bb_trace_processor, + absl::Span symbols_of_interest_addresses, + FunctionMapping& function_name_to_id, bool split_on_segment, + int64_t max_blocks_per_segment) { + std::vector function_id_to_name; + + ASSIGN_OR_RETURN(auto bb_addresses_to_ids_and_functions, + GetBbAddressesToIdsMap(binary_path, function_id_to_name)); + + std::vector section_info = GetSectionsWithBbInfoFromBinary( + binary_path, bb_addresses_to_ids_and_functions); + + ASSIGN_OR_RETURN( + const auto unstripped_binary_processor, + mlgo::latency_model::UnstrippedBinaryProcessor::Create(binary_path)); + + // Read modules.log file contents + std::string modules_log_path = std::string(trace_dir) + "/modules.log"; + std::ifstream modules_file(modules_log_path, + std::ios::binary | std::ios::ate); + if (!modules_file) { + return absl::InternalError( + absl::StrCat("Failed to open modules.log at ", modules_log_path)); + } + std::streamsize size = modules_file.tellg(); + modules_file.seekg(0, std::ios::beg); + std::vector modules_buffer(size + 1, 0); + if (!modules_file.read(modules_buffer.data(), size)) { + return absl::InternalError("Failed to read modules.log"); + } + + // Create module mapper + std::unique_ptr module_mapper = + dynamorio::drmemtrace::module_mapper_t::create(modules_buffer.data()); + if (!module_mapper || !module_mapper->get_last_error().empty()) { + return absl::InternalError( + absl::StrCat("Failed to create module mapper: ", + module_mapper ? module_mapper->get_last_error() : "")); + } + + std::vector> tools; + tools.push_back(std::make_unique( + bb_addresses_to_ids_and_functions, function_id_to_name, + module_mapper.get(), symbols_of_interest_addresses, bb_trace_processor, + function_name_to_id, split_on_segment, max_blocks_per_segment, + std::move(section_info), unstripped_binary_processor->GetLinkerBuildID(), + std::string(binary_path))); + + std::vector tool_ptrs; + tool_ptrs.reserve(tools.size()); + for (const auto& t : tools) { + tool_ptrs.push_back(t.get()); + } + + // Run standard DynamoRIO trace analyzer + dynamorio::drmemtrace::analyzer_t analyzer( + std::string(trace_dir), tool_ptrs.data(), tool_ptrs.size()); + if (!analyzer) { + return absl::InternalError("Failed to initialize trace analyzer"); + } + if (!analyzer.run()) { + return absl::InternalError("Failed to run trace analyzer"); + } + + return absl::OkStatus(); +} + +absl::Status GetBasicBlockTracesFromDirectory( + absl::string_view trace_dir, absl::string_view binary_path, + absl::FunctionRef bb_trace_processor, + absl::Span symbols_of_interest, + FunctionMapping& function_name_to_id, bool split_on_segment, + int64_t max_blocks_per_segment) { + ASSIGN_OR_RETURN( + const auto unstripped_binary_processor, + mlgo::latency_model::UnstrippedBinaryProcessor::Create(binary_path)); + std::vector entrypoint_addresses; + RETURN_IF_ERROR(unstripped_binary_processor->ProcessBBAddrMap( + [&symbols_of_interest, &entrypoint_addresses]( + const UnstrippedBinaryProcessor::FunctionBBInfo& function_bb_info) { + // We simply iterate over the symbols as there will always be very + // few (n<5), and moving the symbols to a set and hashing would likely + // be more expensive. + for (absl::string_view symbol_of_interest : symbols_of_interest) { + if (symbol_of_interest != function_bb_info.function_name) { + continue; + } + entrypoint_addresses.push_back(function_bb_info.function_address); + } + })); + CHECK_EQ(symbols_of_interest.size(), entrypoint_addresses.size()) + << "Expected to find only one address per entrypoint."; + + return GetBasicBlockTracesFromDirectory( + trace_dir, binary_path, bb_trace_processor, entrypoint_addresses, + function_name_to_id, split_on_segment, max_blocks_per_segment); +} + +absl::StatusOr> +GetBbAddressesToIdsMap(absl::string_view binary_path, + std::vector& function_id_to_name) { + ASSIGN_OR_RETURN( + const auto unstripped_binary_processor, + mlgo::latency_model::UnstrippedBinaryProcessor::Create(binary_path)); + + BinaryApplicationToBbDisassembler bb_disassembler("x86_64", + std::string(binary_path)); + + absl::flat_hash_map bb_addresses_to_ids_and_functions; + absl::flat_hash_map function_name_to_id_map; + + RETURN_IF_ERROR(unstripped_binary_processor->ProcessBBAddrMap( + [&bb_addresses_to_ids_and_functions, &function_name_to_id_map, + &function_id_to_name, &bb_disassembler]( + const UnstrippedBinaryProcessor::FunctionBBInfo& function_bb_info) { + for (uint32_t i = 0; i < function_bb_info.bb_infos.size(); ++i) { + if (function_bb_info.bb_infos[i].address == 0) continue; + + // Skip empty basic blocks as if there is another block at exactly + // the same address that has a non-zero size, we might end up picking + // up the zero sized BB which would cause downstream consumers to + // use the wrong instructions. + if (function_bb_info.bb_infos[i].size == 0) continue; + + // TODO: For now, assert if we find .cold functions as + // we need to ensure that we can handle them. + QCHECK(!function_bb_info.function_name.ends_with(".cold")); + + const auto function_id_it = + function_name_to_id_map.find(function_bb_info.function_name); + uint32_t function_id = 0; + if (function_id_it == function_name_to_id_map.end()) { + function_id = function_name_to_id_map.size(); + function_name_to_id_map.emplace(function_bb_info.function_name, + function_id); + + function_id_to_name.push_back( + std::string(function_bb_info.function_name)); + QCHECK(function_id_to_name.size() == + function_name_to_id_map.size()); + } else { + function_id = function_id_it->second; + } + + uint32_t current_entry = 0; + + bb_disassembler.ProcessAllEntriesInBlock( + function_bb_info.bb_infos[i].address, + function_bb_info.bb_infos[i].size, + [&](uint64_t address_offset, + std::vector entry_instructions, + llvm::ArrayRef entry_contents) -> void { + MachineBbId function_bb_id; + function_bb_id.set_function_id(function_id); + function_bb_id.set_basic_block_id(i); + function_bb_id.set_entry_id(current_entry++); + auto [_, inserted_bb] = + bb_addresses_to_ids_and_functions.emplace( + function_bb_info.bb_infos[i].address + address_offset, + function_bb_id); + + // Assert that if a BB is not inserted, the basic block IDs + // match up. If a BB is not inserted, that means there is + // another BB already present at that address, which leaves it + // ambiguous which one should be used for modelling. We cannot + // assert that the function name matches as there are multiple + // symbol names that point to the same definition in some cases, + // like whole object and base object constructors and + // destructors. + if (!inserted_bb) { + QCHECK_EQ(i, bb_addresses_to_ids_and_functions + [function_bb_info.bb_infos[i].address] + .basic_block_id()); + } + }); + } + })); + + CHECK_GT(bb_addresses_to_ids_and_functions.size(), 0); + return bb_addresses_to_ids_and_functions; +} + +absl::Status WriteBasicBlockTraces( + absl::string_view trace_dir, absl::string_view binary_path, + absl::string_view output_folder, + absl::Span symbols_of_interest, bool split_on_segment, + int64_t max_blocks_per_segment) { + riegeli::RecordWriter output_writer( + riegeli::Maker(std::string(output_folder) + "/" + + std::string(kBbTraceFileName)), + riegeli::RecordWriterBase::Options().set_zstd()); + + auto bb_trace_processor = + [&output_writer](const mlgo::latency_model::MbbTrace& mbb_trace_segment) { + LOG(INFO) << "Writing a trace with " << mbb_trace_segment.mbbs_size() + << " BBs"; + output_writer.WriteRecord(mbb_trace_segment); + }; + + mlgo::latency_model::FunctionMapping function_name_to_id; + RETURN_IF_ERROR(GetBasicBlockTracesFromDirectory( + trace_dir, binary_path, bb_trace_processor, symbols_of_interest, + function_name_to_id, split_on_segment, max_blocks_per_segment)); + QCHECK(output_writer.Close()) << output_writer.status(); + + if (function_name_to_id.function_ids().empty()) { + LOG(WARNING) + << "No traces were written to " + << std::string(output_folder) + "/" + std::string(kBbTraceFileName) + << ". Check if symbols of interest are present in the profile."; + } + + riegeli::RecordWriter function_index_writer( + riegeli::Maker(std::string(output_folder) + "/" + + std::string(kFunctionIndexFileName)), + riegeli::RecordWriterBase::Options().set_zstd()); + function_index_writer.WriteRecord(function_name_to_id); + QCHECK(function_index_writer.Close()) << function_index_writer.status(); + return absl::OkStatus(); +} + +ApplicationToBbDisassembler::~ApplicationToBbDisassembler() = default; + +void ApplicationToBbDisassembler::PopulateLlvmHelpers( + const std::string& target_triple) { + std::string possible_lookup_error; + llvm::Triple triple(target_triple); + const llvm::Target* const llvm_target = + llvm::TargetRegistry::lookupTarget(triple, possible_lookup_error); + QCHECK(llvm_target); + + llvm::TargetOptions llvm_target_options; + + llvm_target_machine_.reset(llvm_target->createTargetMachine( + triple, /*CPU*/ "", /*Features*/ "", llvm_target_options, std::nullopt)); + QCHECK(llvm_target_machine_); + + llvm_mc_context_ = std::make_unique( + llvm_target_machine_->getTargetTriple(), + llvm_target_machine_->getMCAsmInfo(), + llvm_target_machine_->getMCRegisterInfo(), + llvm_target_machine_->getMCSubtargetInfo()); + QCHECK(llvm_mc_context_); + + llvm_mc_disassembler_.reset(llvm_target->createMCDisassembler( + llvm_target_machine_->getMCSubtargetInfo(), *llvm_mc_context_)); + QCHECK(llvm_mc_disassembler_); + + llvm_mc_instr_info_.reset(llvm_target->createMCInstrInfo()); + QCHECK(llvm_mc_instr_info_); +} + +ApplicationToBbDisassembler::ApplicationToBbDisassembler( + const std::string& target_triple) { + PopulateLlvmHelpers(target_triple); +} + +llvm::ArrayRef +ApplicationToBbDisassembler::GetDisassembledInstructions( + MachineBbId function_basic_block_id) { + // TODO: This hurts performance. We should remove it once we + // can statically know about all the BB traces at the beginning and add + // them to the cache at that point. + absl::ReaderMutexLock disassembling_lock(disassembling_instructions_mutex_); + + const auto disassembled_basic_block = + disassembled_instructions_.find(function_basic_block_id); + + if (disassembled_basic_block == disassembled_instructions_.end() && + function_basic_block_id.entry_id() > 0) { + return {}; + } + + QCHECK(disassembled_basic_block != disassembled_instructions_.end()) + << function_basic_block_id.has_function_id() << ":" + << function_basic_block_id.function_id() << ":" + << function_basic_block_id.basic_block_id() << ":" + << function_basic_block_id.entry_id() << "\n"; + + return disassembled_basic_block->second; +} + +void ApplicationToBbDisassembler::ProcessAllEntriesFromBlockContents( + uint64_t block_start_address, + absl::FunctionRef&&, + llvm::ArrayRef)> + entry_processor, + llvm::ArrayRef block_contents, + uint64_t expected_block_size) const { + if (block_contents.empty()) { + entry_processor(0, {}, {}); + return; + } + + size_t previous_offset = 0; + size_t current_offset = 0; + std::vector bb_instructions; + while (current_offset < expected_block_size && + current_offset < block_contents.size()) { + // while (current_offset < block_contents.size()) { + uint64_t current_instruction_size = 0; + + const uint8_t* instruction_bytes_address = + &block_contents.data()[current_offset]; + llvm::ArrayRef instruction_data( + instruction_bytes_address, block_contents.size() - current_offset); + llvm::MCInst current_instruction = getInstructionFromBytes( + instruction_data, *llvm_mc_disassembler_, current_instruction_size); + + const llvm::MCInstrDesc& current_instruction_description = + llvm_mc_instr_info_->get(current_instruction.getOpcode()); + + bb_instructions.push_back({.instruction = std::move(current_instruction), + .address = block_start_address + current_offset, + .size = current_instruction_size}); + current_offset += current_instruction_size; + + QCHECK_NE(current_instruction_size, 0); + + // We should always be before the end of the block or exactly at the end + // of the block. + QCHECK(current_offset <= block_contents.size()); + + // We cannot have an entry at the end of the block because that is + // actually just the first entry for the next block, so if we run into + // this case, just return here. + if (current_offset == block_contents.size()) { + llvm::ArrayRef entry_contents( + block_contents.data() + previous_offset, + current_offset - previous_offset); + entry_processor(previous_offset, std::move(bb_instructions), + entry_contents); + return; + } + + // Split around calls as they represent a change in control flow not + // captured by the compiler's definition of a basic block. + // Additionally split around any terminator instructions to handle cases + // like inline assembly where a terminator instruction such as a jump + // might be placed in the middle of a block. + if (current_instruction_description.isCall() || + current_instruction_description.isTerminator()) { + llvm::ArrayRef entry_contents( + block_contents.data() + previous_offset, + current_offset - previous_offset); + entry_processor(previous_offset, std::move(bb_instructions), + entry_contents); + bb_instructions.clear(); + previous_offset = current_offset; + } + } +} + +void ApplicationToBbDisassembler::LoadSharedObjectTraces( + const MbbTrace& mbb_trace) { + absl::MutexLock disassembling_lock(disassembling_instructions_mutex_); + + for (const SharedObjectTrace& shared_object_trace : + mbb_trace.shared_object_traces()) { + std::vector trace_instructions; + trace_instructions.reserve(shared_object_trace.instruction_data_size()); + + for (const std::string& instruction_encoding : + shared_object_trace.instruction_data()) { + uint64_t total_encoding_size = 0; + while (total_encoding_size < instruction_encoding.size()) { + uint64_t instruction_size; + llvm::ArrayRef instruction_data( + reinterpret_cast(instruction_encoding.data() + + total_encoding_size), + instruction_encoding.size()); + + trace_instructions.push_back( + {.instruction = getInstructionFromBytes( + instruction_data, *llvm_mc_disassembler_, instruction_size), + .address = 0, + .size = instruction_size}); + + total_encoding_size += instruction_size; + } + } + + MachineBbId mbb_to_insert; + mbb_to_insert.set_basic_block_id(shared_object_trace.trace_id()); + disassembled_instructions_[mbb_to_insert] = std::move(trace_instructions); + } +} + +void ApplicationToBbDisassembler::LoadSerializedBbs( + const SerializedMbbs& serialized_mbbs) { + absl::MutexLock disassembling_lock(disassembling_instructions_mutex_); + + for (const auto& [mbb_id, mbb_data] : + llvm::zip(serialized_mbbs.mbb_ids(), serialized_mbbs.mbb_bytes())) { + uint64_t current_offset = 0; + std::vector current_block_instructions; + + while (current_offset < mbb_data.size()) { + llvm::ArrayRef instruction_data( + reinterpret_cast(mbb_data.data() + current_offset), + mbb_data.size() - current_offset); + + uint64_t instruction_size = 0; + current_block_instructions.push_back( + {.instruction = getInstructionFromBytes( + instruction_data, *llvm_mc_disassembler_, instruction_size), + .address = 0, + .size = instruction_size}); + + current_offset += instruction_size; + } + + disassembled_instructions_.emplace(mbb_id, + std::move(current_block_instructions)); + } +} + +void BinaryApplicationToBbDisassembler::ProcessAllEntriesInBlock( + uint64_t block_address, uint32_t block_size, + absl::FunctionRef, + llvm::ArrayRef)> + entry_processor) const { + llvm::ArrayRef block_contents = + GetBlockContentsFromAddress(block_address, block_size + 15); + + ProcessAllEntriesFromBlockContents(block_address, entry_processor, + block_contents, block_size); +} + +llvm::ArrayRef +BinaryApplicationToBbDisassembler::GetBlockContentsFromAddress( + uint64_t block_address, uint32_t block_size) const { + auto bb_section = absl::c_lower_bound( + address_to_section_, block_address, + [](const std::pair section_info, + uint64_t function_address) { + return std::get<0>(section_info) <= function_address; + }); + + QCHECK(bb_section != address_to_section_.begin()); + bb_section--; + + llvm::Expected section_contents = + bb_section->second.getContents(); + QCHECK(LlvmExpectedSucceeded(section_contents)); + + QCHECK_GE(block_address, bb_section->first); + size_t bb_start_index = block_address - bb_section->first; + const uint8_t* bb_start_address = reinterpret_cast( + §ion_contents->data()[bb_start_index]); + + size_t section_size = section_contents->size(); + size_t available_size = section_size - bb_start_index; + size_t actual_size = + std::min(static_cast(block_size), available_size); + llvm::ArrayRef block_contents(bb_start_address, actual_size); + // llvm::ArrayRef block_contents(bb_start_address, block_size); + + return block_contents; +} + +BinaryApplicationToBbDisassembler::BinaryApplicationToBbDisassembler( + const std::string& target_triple, const std::string& binary_path) + : ApplicationToBbDisassembler(target_triple), binary_path_(binary_path) { + llvm::Expected> + application_binary = llvm::object::createBinary(binary_path); + + QCHECK(LlvmExpectedSucceeded(application_binary)); + binary_ = std::move(*application_binary); + + llvm::object::ObjectFile* object_file = + llvm::cast(binary_.getBinary()); + QCHECK(object_file); + + // Populate the address to section map. + for (const llvm::object::SectionRef& section : object_file->sections()) { + // Skip all non-text sections as we only need to pull instructions from + // these sections later, which are guaranteed to be in a text section. + if (!section.isText()) continue; + + address_to_section_.emplace(section.getAddress(), section); + } +} + +void ApplicationToBbDisassembler::LoadBasicBlocks( + const FunctionMapping& function_name_to_id) { + LOG(QFATAL) + << "LoadBasicBlocks is not implemented for the current implementation."; +} + +void BinaryApplicationToBbDisassembler::LoadBasicBlocks( + const FunctionMapping& function_name_to_id) { + // Populate the basic block to address map. + const absl::StatusOr< + std::unique_ptr> + unstripped_binary_processor = + mlgo::latency_model::UnstrippedBinaryProcessor::Create(binary_path_); + QCHECK_OK(unstripped_binary_processor); + + absl::Status possible_bbaddrmap_processing_error = + (*unstripped_binary_processor) + ->ProcessBBAddrMap( + [this, &function_name_to_id]( + const mlgo::latency_model::UnstrippedBinaryProcessor:: + FunctionBBInfo& function_bb_info) { + // TODO: For now, assert if we find .cold functions + // as we need to ensure that we can handle them. + QCHECK(!function_bb_info.function_name.ends_with(".cold")); + + const auto function_id_it = + function_name_to_id.function_ids().find( + function_bb_info.function_name); + // Skip the function if we cannot find the name to function id + // mapping as it should imply that no basic blocks from this + // function are included in the traces. + if (function_id_it == + function_name_to_id.function_ids().end()) { + VLOG(1) << "Failed to find a function ID for " + << function_bb_info.function_name << "\n"; + return; + } + + for (uint32_t i = 0; i < function_bb_info.bb_infos.size(); + ++i) { + uint32_t current_entry = 0; + + if (function_bb_info.bb_infos[i].address == 0) { + continue; + } + + ProcessAllEntriesInBlock( + function_bb_info.bb_infos[i].address, + function_bb_info.bb_infos[i].size, + [¤t_entry, this, &function_id_it, i]( + uint64_t address_offset, + std::vector entry_instructions, + llvm::ArrayRef entry_contents) -> void { + MachineBbId function_address_and_bb_id; + function_address_and_bb_id.set_function_id( + function_id_it->second); + function_address_and_bb_id.set_basic_block_id(i); + function_address_and_bb_id.set_entry_id(current_entry); + + ++current_entry; + + disassembled_instructions_.emplace( + function_address_and_bb_id, + std::move(entry_instructions)); + }); + } + }); + QCHECK_OK(possible_bbaddrmap_processing_error); +} + +CorpusApplicationToBbDisassembler::CorpusApplicationToBbDisassembler( + const std::string& target_triple, + const std::vector& module_paths, bool store_block_contents) + : ApplicationToBbDisassembler(target_triple), + module_paths_(module_paths), + store_block_contents_(store_block_contents) {} + +void CorpusApplicationToBbDisassembler::LoadBasicBlocks( + const FunctionMapping& function_name_to_id) { + // Load all of the object files. + size_t file_index = 0; + for (const std::string& module_path : module_paths_) { + llvm::Expected> + module_binary = llvm::object::createBinary(module_path); + QCHECK(LlvmExpectedSucceeded(module_binary)); + + llvm::object::ObjectFile* object_file = + llvm::cast(module_binary->getBinary()); + QCHECK_NE(object_file, nullptr); + + // Populate an address to name map to use later for associating addresses + // with the appropriate function. + absl::flat_hash_map, + llvm::SmallSet> + section_index_and_offset_to_name; + + for (const llvm::object::SymbolRef& symbol : object_file->symbols()) { + llvm::Expected symbol_type = + symbol.getType(); + QCHECK(LlvmExpectedSucceeded(symbol_type)); + if (*symbol_type != llvm::object::SymbolRef::ST_Function) continue; + + llvm::Expected offset = symbol.getAddress(); + QCHECK(LlvmExpectedSucceeded(offset)); + + llvm::Expected symbol_section = + symbol.getSection(); + QCHECK(LlvmExpectedSucceeded(symbol_section)); + uint64_t section_index = symbol_section.get()->getIndex(); + + llvm::Expected symbol_name = symbol.getName(); + QCHECK(LlvmExpectedSucceeded(symbol_name)); + + section_index_and_offset_to_name[std::make_pair(section_index, *offset)] + .insert(*symbol_name); + } + + for (const llvm::object::SectionRef& section : object_file->sections()) { + if (!section.isText()) continue; + + // Get the basic block address map. + const auto* elf_object_file = + llvm::dyn_cast(object_file); + QCHECK(elf_object_file); + llvm::Expected> bb_addr_maps = + elf_object_file->readBBAddrMap(section.getIndex()); + QCHECK(LlvmExpectedSucceeded(bb_addr_maps)); + + for (const auto& bb_addr_map : *bb_addr_maps) { + for (const llvm::object::BBAddrMap::BBEntry& entry : + bb_addr_map.getBBEntries()) { + uint64_t bb_offset = entry.Offset + bb_addr_map.getFunctionAddress(); + + auto name_iter = section_index_and_offset_to_name.find(std::make_pair( + section.getIndex(), bb_addr_map.getFunctionAddress())); + if (name_iter == section_index_and_offset_to_name.end()) { + LOG(WARNING) << "Found a BB map entry without a symbol name."; + continue; + } + + const llvm::SmallSet& names = name_iter->second; + for (const auto& name : names) { + // TODO: For now, assert if we find .cold functions as + // we need to ensure that we can handle them. + QCHECK(!name.ends_with(".cold")); + + const auto function_id_it = + function_name_to_id.function_ids().find(name); + // Skip the function if we cannot find the name to function id + // mapping as it should imply that no basic blocks from this + // function are included in the traces. + if (function_id_it == function_name_to_id.function_ids().end()) { + VLOG(1) << "Failed to find a function ID for " << name.str() + << "\n"; + continue; + } + + uint32_t current_entry = 0; + + ProcessAllEntriesInBlock( + bb_offset, entry.Size, section, + [¤t_entry, this, &function_id_it, &entry]( + uint64_t address_offset, + std::vector partial_block_instructions, + llvm::ArrayRef partial_block_contents) { + MachineBbId function_name_and_bb_id; + function_name_and_bb_id.set_function_id( + function_id_it->second); + function_name_and_bb_id.set_basic_block_id(entry.ID); + function_name_and_bb_id.set_entry_id(current_entry); + + ++current_entry; + + disassembled_instructions_.emplace( + function_name_and_bb_id, + std::move(partial_block_instructions)); + + if (store_block_contents_) { + entry_contents_.emplace(function_name_and_bb_id, + partial_block_contents); + } + }); + } + } + } + } + + LOG(INFO) << "Finished loading " << module_path << " - " << file_index; + + ++file_index; + } +} + +void CorpusApplicationToBbDisassembler::ProcessAllEntriesInBlock( + uint64_t bb_offset, uint32_t bb_size, llvm::object::SectionRef bb_section, + absl::FunctionRef, + llvm::ArrayRef)> + entry_processor) const { + llvm::ArrayRef block_contents = + GetBlockContentsFromSectionOffset(bb_offset, bb_size + 15, bb_section); + + ProcessAllEntriesFromBlockContents(bb_offset, entry_processor, block_contents, + bb_size); +} + +llvm::ArrayRef +CorpusApplicationToBbDisassembler::GetBlockContentsFromSectionOffset( + uint64_t bb_offset, uint32_t bb_size, + llvm::object::SectionRef bb_section) const { + llvm::Expected section_contents = bb_section.getContents(); + QCHECK(LlvmExpectedSucceeded(section_contents)); + + size_t bb_start_index = bb_offset; + const uint8_t* bb_start_offset = reinterpret_cast( + §ion_contents->data()[bb_start_index]); + + // llvm::ArrayRef block_contents(bb_start_offset, bb_size); + size_t section_size = section_contents->size(); + size_t available_size = section_size - bb_start_index; + size_t actual_size = std::min(static_cast(bb_size), available_size); + llvm::ArrayRef block_contents(bb_start_offset, actual_size); + + return block_contents; +} + +llvm::ArrayRef CorpusApplicationToBbDisassembler::GetEntryContents( + MachineBbId function_basic_block_id) { + absl::ReaderMutexLock disassembling_lock(disassembling_instructions_mutex_); + + const auto basic_block_contents = + entry_contents_.find(function_basic_block_id); + + if (basic_block_contents == entry_contents_.end() && + function_basic_block_id.entry_id() > 0) { + return {}; + } + + QCHECK(basic_block_contents != entry_contents_.end()) + << function_basic_block_id.has_function_id() << ":" + << function_basic_block_id.function_id() << ":" + << function_basic_block_id.basic_block_id() << ":" + << function_basic_block_id.entry_id() << "\n"; + + return basic_block_contents->second; +} + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/basic_block_trace.h b/compiler_opt/memtrace_costmodel/basic_block_trace.h new file mode 100644 index 00000000..799ed83f --- /dev/null +++ b/compiler_opt/memtrace_costmodel/basic_block_trace.h @@ -0,0 +1,282 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_BASIC_BLOCK_TRACE_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_BASIC_BLOCK_TRACE_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/container/btree_map.h" +#include "absl/container/flat_hash_map.h" +#include "absl/functional/function_ref.h" +#include "absl/hash/hash.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" +#include "compiler_opt/memtrace_costmodel/serialized_mbbs.proto.h" +#include "drmemtrace/analysis_tool.h" +#include "drmemtrace/memref.h" +#include "drmemtrace/raw2trace_shared.h" +#include "llvm/MC/MCDisassembler/MCDisassembler.h" +#include "llvm/MC/MCInst.h" +#include "llvm/MC/MCInstrInfo.h" +#include "llvm/Object/Binary.h" +#include "llvm/Object/ObjectFile.h" +#include "llvm/Target/TargetMachine.h" + +namespace mlgo { +namespace latency_model { + +// We include has_function_id() in the hash and equality checks as we use +// MachineBbIds without function IDs to represent direct traces so they do +// not overlap with a potential function that has ID zero. + +struct MachineBBIDKeyHash { + size_t operator()(const MachineBbId& machine_bbid) const { + return absl::HashOf(machine_bbid.has_function_id(), + machine_bbid.function_id(), + machine_bbid.basic_block_id(), machine_bbid.entry_id()); + } +}; + +struct MachineBBIDKeyEqual { + bool operator()(const MachineBbId& lhs, const MachineBbId& rhs) const { + return lhs.has_function_id() == rhs.has_function_id() && + lhs.function_id() == rhs.function_id() && + lhs.basic_block_id() == rhs.basic_block_id() && + lhs.entry_id() == rhs.entry_id(); + } +}; + +struct SectionInfo { + uint64_t address = 0; + uint64_t size = 0; +}; + +struct InstructionInfo { + llvm::MCInst instruction; + uint64_t address = 0; + uint64_t size = 0; +}; + +class BBMemtraceProcessor : public dynamorio::drmemtrace::analysis_tool_t { + public: + // The backing array for symbols_of_interest_addresses needs to outlive the + // created instance of this class. + explicit BBMemtraceProcessor( + const absl::flat_hash_map& + bb_addresses_to_ids_and_functions, + const std::vector& function_id_to_name, + const dynamorio::drmemtrace::module_mapper_t* module_mapper, + absl::Span symbols_of_interest_addresses, + absl::FunctionRef& bb_trace_processor, + FunctionMapping& function_name_to_id, bool split_on_segment, + int64_t max_blocks_per_segment, + std::vector&& sections_with_bb_info, + std::string expected_build_id, std::string binary_path) + : bb_addresses_to_ids_and_functions_(bb_addresses_to_ids_and_functions), + function_id_to_name_(function_id_to_name), + module_mapper_(module_mapper), + symbols_of_interest_addresses_(symbols_of_interest_addresses), + bb_trace_processor_(bb_trace_processor), + function_name_to_id_(function_name_to_id), + split_on_segment_(split_on_segment), + max_blocks_per_segment_(max_blocks_per_segment), + sections_with_bb_info_(std::move(sections_with_bb_info)), + expected_build_id_(std::move(expected_build_id)), + binary_path_(std::move(binary_path)) {}; + + private: + struct PerShardData { + PerShardData() = default; + + absl::Mutex mutex; + MbbTrace current_trace_data ABSL_GUARDED_BY(mutex); + std::vector entrypoint_segments ABSL_GUARDED_BY(mutex); + bool previously_under_segment ABSL_GUARDED_BY(mutex) = false; + bool inside_shared_object ABSL_GUARDED_BY(mutex) = false; + absl::flat_hash_map function_name_to_id + ABSL_GUARDED_BY(mutex); + }; + + void* parallel_shard_init(int shard_index, void* worker_data) override; + bool parallel_shard_exit(void* shard_data) override; + + bool process_memref(const dynamorio::drmemtrace::memref_t& memref) override; + + bool parallel_shard_supported() override { return true; } + + bool parallel_shard_memref( + void* shard_data, const dynamorio::drmemtrace::memref_t& memref) override; + + bool print_results() override { return true; } + + const absl::flat_hash_map& + bb_addresses_to_ids_and_functions_; + const std::vector& function_id_to_name_; + const dynamorio::drmemtrace::module_mapper_t* module_mapper_; + const absl::Span symbols_of_interest_addresses_; + absl::FunctionRef& bb_trace_processor_; + absl::Mutex bb_trace_processor_mutex_; + absl::Mutex function_name_to_id_mutex_; + FunctionMapping& function_name_to_id_ + ABSL_GUARDED_BY(function_name_to_id_mutex_); + bool split_on_segment_ = false; + const int64_t max_blocks_per_segment_ = 0; + std::vector sections_with_bb_info_; + std::string expected_build_id_; + std::string binary_path_; + + uint32_t current_trace_id_ ABSL_GUARDED_BY(current_trace_id_mutex_) = 0; + absl::Mutex current_trace_id_mutex_; +}; + +absl::Status GetBasicBlockTracesFromDirectory( + absl::string_view trace_dir, absl::string_view binary_path, + absl::FunctionRef bb_trace_processor, + absl::Span symbols_of_interest_addresses, + FunctionMapping& function_name_to_id, bool split_on_segment, + int64_t max_blocks_per_segment = std::numeric_limits::max()); + +absl::Status GetBasicBlockTracesFromDirectory( + absl::string_view trace_dir, absl::string_view binary_path, + absl::FunctionRef bb_trace_processor, + absl::Span symbols_of_interest, + FunctionMapping& function_name_to_id, bool split_on_segment, + int64_t max_blocks_per_segment = std::numeric_limits::max()); + +absl::Status WriteBasicBlockTraces( + absl::string_view trace_dir, absl::string_view binary_path, + absl::string_view output_folder, + absl::Span symbols_of_interest, bool split_on_segment, + int64_t max_blocks_per_segment = std::numeric_limits::max()); + +absl::StatusOr> +GetBbAddressesToIdsMap(absl::string_view binary_path, + std::vector& function_id_to_name); + +// Takes a application in an implementation-defined form (like a binary or +// corpus) and disassembles individual basic blocks within the application on +// demand with caching. +class ApplicationToBbDisassembler { + public: + llvm::ArrayRef GetDisassembledInstructions( + MachineBbId function_basic_block_id); + + // Calls entry_processor with each block entry found in the basic block + // at block_address with size block_size. An entry is defined as an + // instruction where a block can be entered, such as when returning + // from a call within that block. + void ProcessAllEntriesInBlock( + uint64_t block_address, uint32_t block_size, + absl::FunctionRef entry_processor) const; + + virtual void LoadBasicBlocks(const FunctionMapping& function_name_to_id); + + void LoadSharedObjectTraces(const MbbTrace& mbb_trace); + + // Loads serialized machine basic blocks into internal state such that they + // can then be queried using GetDisassembledInstructions. + void LoadSerializedBbs(const SerializedMbbs& serialized_mbbs); + + virtual ~ApplicationToBbDisassembler(); + + protected: + explicit ApplicationToBbDisassembler(const std::string& target_triple); + std::unique_ptr llvm_mc_disassembler_; + std::unique_ptr llvm_mc_instr_info_; + + void ProcessAllEntriesFromBlockContents( + uint64_t block_start_address, + absl::FunctionRef&&, + llvm::ArrayRef)> + entry_processor, + llvm::ArrayRef block_contents, + uint64_t expected_block_size) const; + + absl::flat_hash_map, + MachineBBIDKeyHash, MachineBBIDKeyEqual> + disassembled_instructions_; + + absl::Mutex disassembling_instructions_mutex_; + + private: + void PopulateLlvmHelpers(const std::string& target_triple); + + std::unique_ptr llvm_target_machine_; + std::unique_ptr llvm_mc_context_; +}; + +class BinaryApplicationToBbDisassembler : public ApplicationToBbDisassembler { + public: + explicit BinaryApplicationToBbDisassembler(const std::string& target_triple, + const std::string& binary_path); + + void ProcessAllEntriesInBlock( + uint64_t block_address, uint32_t block_size, + absl::FunctionRef, + llvm::ArrayRef)> + entry_processor) const; + + void LoadBasicBlocks(const FunctionMapping& function_name_to_id) override; + + private: + // Gets the contents of a basic block identified by its address and + // size. This is intended for processing through all the basic blocks + // in a binary so that maps can be set up appropriately for partial + // basic blocks. + llvm::ArrayRef GetBlockContentsFromAddress( + uint64_t block_address, uint32_t block_size) const; + + llvm::object::OwningBinary binary_; + + absl::btree_map address_to_section_; + + std::string binary_path_; +}; + +// Takes an application in the form of a corpus and implements functions +// to find where specific blocks are within the corpus and to grab the bytes +// for those blocks so they can be disassembled and used for cost modelling. +class CorpusApplicationToBbDisassembler : public ApplicationToBbDisassembler { + public: + explicit CorpusApplicationToBbDisassembler( + const std::string& target_triple, + const std::vector& module_paths, + bool store_block_contents = false); + + void ProcessAllEntriesInBlock( + uint64_t bb_offset, uint32_t bb_size, llvm::object::SectionRef bb_section, + absl::FunctionRef, + llvm::ArrayRef)> + entry_processor) const; + + void LoadBasicBlocks(const FunctionMapping& function_name_to_id) override; + + llvm::ArrayRef GetEntryContents(MachineBbId function_basic_block_id); + + private: + llvm::ArrayRef GetBlockContentsFromSectionOffset( + uint64_t bb_offset, uint32_t bb_size, + llvm::object::SectionRef bb_section) const; + + const std::vector& module_paths_; + const bool store_block_contents_ = false; + + absl::flat_hash_map, MachineBBIDKeyHash, + MachineBBIDKeyEqual> + entry_contents_; +}; + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_BASIC_BLOCK_TRACE_H_ diff --git a/compiler_opt/memtrace_costmodel/basic_block_trace_extract.cc b/compiler_opt/memtrace_costmodel/basic_block_trace_extract.cc new file mode 100644 index 00000000..df4a28ec --- /dev/null +++ b/compiler_opt/memtrace_costmodel/basic_block_trace_extract.cc @@ -0,0 +1,57 @@ +#include +#include +#include + +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "absl/log/check.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" +#include "absl/types/span.h" +#include "compiler_opt/memtrace_costmodel/basic_block_trace.h" +#include "llvm-c/Target.h" + +ABSL_FLAG(std::string, memtrace_path, "", + "The path to the memtrace to process."); +ABSL_FLAG(std::vector, symbol_names, {}, + "The names of the entrypoint symbol."); +ABSL_FLAG(std::string, binary_path, "", "The path to the binary."); +ABSL_FLAG(std::string, output_folder, "", "The path to the output folder."); +ABSL_FLAG(bool, split_on_segment, false, + "Whether or not to split entrypoints by segment"); +ABSL_FLAG(int64_t, max_blocks_per_segment, 1 << 23, + "The maximum number of blocks that can be in an individual segment."); + +int main(int argc, char** argv) { + LLVMInitializeX86Target(); + LLVMInitializeX86TargetInfo(); + LLVMInitializeX86TargetMC(); + LLVMInitializeX86Disassembler(); + + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + + if (absl::GetFlag(FLAGS_memtrace_path).empty()) { + LOG(QFATAL) << "--memtrace_path was not specified."; + } + + if (absl::GetFlag(FLAGS_symbol_names).empty()) { + LOG(QFATAL) << "--symbol_names must be specified."; + } + + if (absl::GetFlag(FLAGS_binary_path).empty()) { + LOG(QFATAL) << "--binary_path was not specified."; + } + + if (absl::GetFlag(FLAGS_output_folder).empty()) { + LOG(QFATAL) << "--output_folder was not specified."; + } + + QCHECK_OK(mlgo::latency_model::WriteBasicBlockTraces( + absl::GetFlag(FLAGS_memtrace_path), absl::GetFlag(FLAGS_binary_path), + absl::GetFlag(FLAGS_output_folder), absl::GetFlag(FLAGS_symbol_names), + absl::GetFlag(FLAGS_split_on_segment), + absl::GetFlag(FLAGS_max_blocks_per_segment))); + + return 0; +} diff --git a/compiler_opt/memtrace_costmodel/basic_block_trace_model.cc b/compiler_opt/memtrace_costmodel/basic_block_trace_model.cc new file mode 100644 index 00000000..43018102 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/basic_block_trace_model.cc @@ -0,0 +1,186 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/flags/parse.h" +#include "compiler_opt/memtrace_costmodel/basic_block_trace.h" +#undef X86 +#undef X86_64 +#include "absl/flags/flag.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/costmodel.h" +#include "compiler_opt/memtrace_costmodel/costmodel_factory.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" +#include "compiler_opt/memtrace_costmodel/serialized_mbbs.proto.h" +#include "llvm-c/Target.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/MC/MCInst.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/Target/TargetOptions.h" +#include "llvm/lib/Target/X86/MCTargetDesc/X86MCTargetDesc.h" +#include "nlohmann/json.hpp" +#include "riegeli/bytes/file_reader.h" +#include "riegeli/records/record_reader.h" + +ABSL_FLAG(std::string, bb_trace_path, "", "The path to the basic block trace."); +ABSL_FLAG(std::string, binary_path, "", "The path to the binary."); +ABSL_FLAG(std::string, corpus_path, "", + "The path to the corpus description JSON."); +ABSL_FLAG(std::string, target_triple, "x86_64", + "The target triple of the binary."); +ABSL_FLAG(std::string, cpu_name, "skylake", "The CPU name to model."); +ABSL_FLAG(std::string, function_index_path, "", + "The path to the function name to ID mapping."); +ABSL_FLAG(std::string, serialized_bbs_path, "", + "The path to the serialized basic blocks to load."); + +// Model specific flags. +ABSL_FLAG(mlgo::latency_model::CostModelType, model_type, + mlgo::latency_model::CostModelType::Mca, + "The type of cost model to use. (\"mca\" | \"print\" " + "| \"instruction_counting\")"); +ABSL_FLAG(std::string, print_output_file, "", + "The output file if the print cost model is selected."); + +int main(int argc, char** argv) { + LLVMInitializeX86Target(); + LLVMInitializeX86TargetInfo(); + LLVMInitializeX86TargetMC(); + LLVMInitializeX86Disassembler(); + LLVMInitializeX86TargetMCA(); + + absl::ParseCommandLine(argc, argv); + + if (absl::GetFlag(FLAGS_bb_trace_path).empty()) { + LOG(QFATAL) << "--bb_trace_path was not specified."; + } + + if (absl::GetFlag(FLAGS_binary_path).empty() && + absl::GetFlag(FLAGS_corpus_path).empty()) { + LOG(QFATAL) << "--binary_path or --corpus_path needs to be specified."; + } + + if (!absl::GetFlag(FLAGS_binary_path).empty() && + !absl::GetFlag(FLAGS_corpus_path).empty()) { + LOG(QFATAL) << "--binary_path and --corpus_path cannot both be specified " + "at the same " + "time."; + } + + if (absl::GetFlag(FLAGS_function_index_path).empty()) { + LOG(QFATAL) << "--function_index_path was not specified."; + } + + mlgo::latency_model::ValidateCostModelFlags( + absl::GetFlag(FLAGS_model_type), "", + absl::GetFlag(FLAGS_print_output_file)); + + riegeli::RecordReader function_index_reader( + riegeli::Maker( + absl::GetFlag(FLAGS_function_index_path))); + mlgo::latency_model::FunctionMapping function_name_to_id; + QCHECK(function_index_reader.ReadRecord(function_name_to_id)); + QCHECK(function_index_reader.Close()) << function_index_reader.status(); + + std::unique_ptr + application_disassembler; + if (!absl::GetFlag(FLAGS_binary_path).empty()) { + application_disassembler = std::make_unique< + mlgo::latency_model::BinaryApplicationToBbDisassembler>( + absl::GetFlag(FLAGS_target_triple), absl::GetFlag(FLAGS_binary_path)); + application_disassembler->LoadBasicBlocks(function_name_to_id); + } else { + // Load module list from corpus JSON file using standard C++ streams. + std::string corpus_description_path = absl::GetFlag(FLAGS_corpus_path); + std::ifstream corpus_file(corpus_description_path); + if (!corpus_file) { + LOG(QFATAL) << "Failed to open corpus path: " << corpus_description_path; + } + std::stringstream buffer; + buffer << corpus_file.rdbuf(); + std::string corpus_description_contents = buffer.str(); + nlohmann::json corpus_description = + nlohmann::json::parse(corpus_description_contents); + + QCHECK(corpus_description.contains("modules")); + QCHECK(corpus_description["modules"].is_array()); + + std::vector module_full_paths; + module_full_paths.reserve(corpus_description["modules"].size()); + + std::string corpus_path = absl::GetFlag(FLAGS_corpus_path); + std::string corpus_dir_path = + std::string(std::filesystem::path(corpus_path).parent_path()); + + for (const std::string relative_module_path : + corpus_description["modules"]) { + module_full_paths.push_back(corpus_dir_path + "/" + relative_module_path + + ".bc.o"); + } + + application_disassembler = std::make_unique< + mlgo::latency_model::CorpusApplicationToBbDisassembler>( + absl::GetFlag(FLAGS_target_triple), module_full_paths); + application_disassembler->LoadBasicBlocks(function_name_to_id); + } + + if (!absl::GetFlag(FLAGS_serialized_bbs_path).empty()) { + riegeli::RecordReader serialized_bbs_reader( + riegeli::Maker( + absl::GetFlag(FLAGS_serialized_bbs_path))); + mlgo::latency_model::SerializedMbbs serialized_mbbs; + QCHECK(serialized_bbs_reader.ReadRecord(serialized_mbbs)); + QCHECK(serialized_bbs_reader.Close()) << serialized_bbs_reader.status(); + + application_disassembler->LoadSerializedBbs(serialized_mbbs); + } + + const std::string& bb_trace_path = absl::GetFlag(FLAGS_bb_trace_path); + + riegeli::RecordReader trace_reader( + riegeli::Maker(bb_trace_path)); + + std::string possible_lookup_error; + llvm::Triple target_triple(absl::GetFlag(FLAGS_target_triple)); + const llvm::Target* target = + llvm::TargetRegistry::lookupTarget(target_triple, possible_lookup_error); + QCHECK_NE(target, nullptr) << possible_lookup_error; + + auto target_machine = std::unique_ptr( + target->createTargetMachine(target_triple, absl::GetFlag(FLAGS_cpu_name), + "", llvm::TargetOptions(), std::nullopt)); + + auto cost_model_factory = mlgo::latency_model::CreateCostModelFactory( + absl::GetFlag(FLAGS_model_type), target_machine.get(), + absl::GetFlag(FLAGS_cpu_name), absl::GetFlag(FLAGS_target_triple), "", 0, + absl::GetFlag(FLAGS_print_output_file)); + + mlgo::latency_model::MbbTrace mbb_trace; + while (trace_reader.ReadRecord(mbb_trace)) { + application_disassembler->LoadSharedObjectTraces(mbb_trace); + auto cost_model = cost_model_factory(); + + for (const mlgo::latency_model::MachineBbId& basic_block : + mbb_trace.mbbs()) { + llvm::ArrayRef bb_instructions = + application_disassembler->GetDisassembledInstructions(basic_block); + for (const mlgo::latency_model::InstructionInfo& instruction : + bb_instructions) { + cost_model->AddInstruction(instruction.instruction); + } + } + + std::cout << "Segment Cost: " << cost_model->GetCost() << "\n"; + } + + QCHECK(trace_reader.Close()) << trace_reader.status(); + return 0; +} diff --git a/compiler_opt/memtrace_costmodel/compare_binary_cfgs.cc b/compiler_opt/memtrace_costmodel/compare_binary_cfgs.cc new file mode 100644 index 00000000..4857e684 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/compare_binary_cfgs.cc @@ -0,0 +1,133 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "absl/log/check.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/compare_binary_cfgs_lib.h" +#include "nlohmann/json.hpp" + +ABSL_FLAG(std::string, binary_path_a, "", "The path to binary A."); +ABSL_FLAG(std::string, corpus_path_a, "", "The path to corpus A."); +ABSL_FLAG(std::string, binary_path_b, "", "The path to binary B."); +ABSL_FLAG(std::string, corpus_path_b, "", "The path to corpus B."); +ABSL_FLAG(bool, continue_on_diff, false, + "Whether or not to continue running if a difference is found."); + +absl::flat_hash_map> +GetCfgFromBinaryOrCorpus(absl::string_view binary_path, + absl::string_view corpus_path) { + if (!binary_path.empty()) { + absl::StatusOr>> + binary_cfgs = mlgo::latency_model::LoadBinaryCfgs(binary_path); + QCHECK_OK(binary_cfgs); + return *binary_cfgs; + } + + std::vector module_full_paths; + + std::string corpus_json_path = + (std::filesystem::path(std::string(corpus_path)) / + "corpus_description.json") + .string(); + std::ifstream corpus_file(corpus_json_path); + QCHECK(corpus_file) << "Failed to open corpus JSON path: " + << corpus_json_path; + std::stringstream buffer; + buffer << corpus_file.rdbuf(); + std::string corpus_json_contents = buffer.str(); + + nlohmann::json corpus_description = + nlohmann::json::parse(corpus_json_contents); + + QCHECK(corpus_description.contains("modules")); + QCHECK(corpus_description["modules"].is_array()); + + for (const std::string module_path : corpus_description["modules"]) { + std::string module_full_path = + (std::filesystem::path(std::string(corpus_path)) / module_path) + .string() + + ".bc.o"; + module_full_paths.push_back(std::move(module_full_path)); + } + + absl::StatusOr>> + corpus_cfgs = mlgo::latency_model::LoadCorpusCfgs(module_full_paths); + QCHECK_OK(corpus_cfgs); + + return *corpus_cfgs; +} + +int main(int argc, char** argv) { + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + + if (absl::GetFlag(FLAGS_binary_path_a).empty() && + absl::GetFlag(FLAGS_corpus_path_a).empty()) { + LOG(QFATAL) << "--binary_path_a or --corpus_path_a needs to be specified."; + } + + if (!absl::GetFlag(FLAGS_binary_path_a).empty() && + !absl::GetFlag(FLAGS_corpus_path_a).empty()) { + LOG(QFATAL) << "--binary_path_a and --corpus_path_a cannot both be set at " + "the same time."; + } + + if (absl::GetFlag(FLAGS_binary_path_b).empty() && + absl::GetFlag(FLAGS_corpus_path_b).empty()) { + LOG(QFATAL) << "--binary_path_b or --corpus_path_b needs to be specified."; + } + + if (!absl::GetFlag(FLAGS_binary_path_b).empty() && + !absl::GetFlag(FLAGS_corpus_path_b).empty()) { + LOG(QFATAL) << "--binary_path_b and --corpus_path_b cannot both be set at " + "the same time."; + } + + absl::flat_hash_map> + a_cfgs = GetCfgFromBinaryOrCorpus(absl::GetFlag(FLAGS_binary_path_a), + absl::GetFlag(FLAGS_corpus_path_a)); + absl::flat_hash_map> + b_cfgs = GetCfgFromBinaryOrCorpus(absl::GetFlag(FLAGS_binary_path_b), + absl::GetFlag(FLAGS_corpus_path_b)); + + bool continue_on_diff = absl::GetFlag(FLAGS_continue_on_diff); + + if (a_cfgs.size() != b_cfgs.size()) { + std::cout << "Binaries have a different function set.\n"; + if (!continue_on_diff) return 0; + } + + for (const auto& function_info : a_cfgs) { + const auto function_b_info_it = b_cfgs.find(function_info.first); + if (function_b_info_it == b_cfgs.end()) { + std::cout << "Failed to find function " << function_info.first + << " from binary A in binary B.\n"; + if (!continue_on_diff) return 0; + } + + bool are_cfgs_different = mlgo::latency_model::AreCfgsDifferent( + function_info.second, function_b_info_it->second); + if (are_cfgs_different) { + std::cout << "Function " << function_info.first + << " has a different CFG between the two versions.\n"; + if (!continue_on_diff) return 0; + } + } + + return 0; +} diff --git a/compiler_opt/memtrace_costmodel/compare_binary_cfgs_lib.cc b/compiler_opt/memtrace_costmodel/compare_binary_cfgs_lib.cc new file mode 100644 index 00000000..5a1a7ad7 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/compare_binary_cfgs_lib.cc @@ -0,0 +1,189 @@ +#include "compiler_opt/memtrace_costmodel/compare_binary_cfgs_lib.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" +#include "absl/hash/hash.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "compiler_opt/memtrace_costmodel/elf_metadata_parser.h" +#include "compiler_opt/memtrace_costmodel/status_macros.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/Object/Binary.h" +#include "llvm/Object/ELFObjectFile.h" +#include "llvm/Object/ELFTypes.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/raw_ostream.h" + +namespace mlgo { +namespace latency_model { +namespace { + +// Helper utilities for converting LLVM Errors and Expecteds into +// absl equivalents. +absl::Status LlvmErrorToStatus(llvm::Error error) { + if (!error) return absl::OkStatus(); + std::string error_string; + llvm::raw_string_ostream error_string_stream(error_string); + error_string_stream << error; + return absl::InternalError(error_string); +} + +template +absl::StatusOr LlvmExpectedToStatusOr(llvm::Expected expected) { + if (expected) return std::move(*expected); + return LlvmErrorToStatus(expected.takeError()); +} + +} // namespace + +struct CfgEdgeComparator { + bool operator()(const ControlFlowEdge& lhs, + const ControlFlowEdge& rhs) const { + return std::forward_as_tuple(lhs.from_block, lhs.to_block) < + std::forward_as_tuple(rhs.from_block, rhs.to_block); + } +}; + +struct FunctionID { + uint64_t function_address; + uint64_t section_index; +}; + +struct FunctionIDKeyHash { + size_t operator()(const FunctionID& function_id) const { + return absl::HashOf(function_id.function_address, + function_id.section_index); + } +}; + +struct FunctionIDKeyEqual { + bool operator()(const FunctionID& lhs, const FunctionID& rhs) const { + return lhs.function_address == rhs.function_address && + lhs.section_index == rhs.section_index; + } +}; + +absl::StatusOr< + absl::flat_hash_map>> +LoadBinaryCfgs(absl::string_view binary_path) { + // Get the mapping of addresses to function names. + absl::flat_hash_map + function_id_to_name; + + ASSIGN_OR_RETURN( + auto unstripped_binary_processor, + mlgo::latency_model::UnstrippedBinaryProcessor::Create(binary_path)); + + RETURN_IF_ERROR(unstripped_binary_processor->ProcessBBAddrMap( + [&function_id_to_name]( + const mlgo::latency_model::UnstrippedBinaryProcessor::FunctionBBInfo& + bb_info) { + FunctionID function_id = {.function_address = bb_info.function_address, + .section_index = bb_info.section_index}; + function_id_to_name.emplace(function_id, bb_info.function_name); + })); + + // Load the executable through the LLVM APIs. + ASSIGN_OR_RETURN( + llvm::object::OwningBinary object_binary, + LlvmExpectedToStatusOr(llvm::object::createBinary(binary_path))); + llvm::object::ELFObjectFileBase* elf_object = + llvm::cast(object_binary.getBinary()); + + // Load the CFG into the map. + absl::flat_hash_map> + binary_cfgs; + + for (const auto& section : elf_object->sections()) { + if (!section.isText()) continue; + + std::vector pgo_analysis_maps; + ASSIGN_OR_RETURN(std::vector bb_addr_maps, + LlvmExpectedToStatusOr(elf_object->readBBAddrMap( + section.getIndex(), &pgo_analysis_maps))); + + for (const auto& [bb_addr_map, pgo_analysis_map] : + llvm::zip(bb_addr_maps, pgo_analysis_maps)) { + FunctionID function_id = { + .function_address = bb_addr_map.getFunctionAddress(), + .section_index = section.getIndex()}; + const auto function_name_it = function_id_to_name.find(function_id); + if (function_name_it == function_id_to_name.end()) { + // LOG(WARNING) + // << "Failed to find a symbol name for function at address 0x" + // << absl::Hex(bb_addr_map.getFunctionAddress()) << " in section " + // << section.getIndex() << ", skipping."; + continue; + } + absl::string_view function_name = function_name_it->second; + + std::vector function_edges; + + for (const auto& [bb_entry, pgo_bb_entry] : + llvm::zip(bb_addr_map.getBBEntries(), pgo_analysis_map.BBEntries)) { + for (const auto& bb_successor_entry : pgo_bb_entry.Successors) { + ControlFlowEdge new_cfg_edge = {.from_block = bb_entry.ID, + .to_block = bb_successor_entry.ID}; + function_edges.push_back(new_cfg_edge); + } + } + + absl::c_sort(function_edges, CfgEdgeComparator()); + + binary_cfgs.emplace(function_name, std::move(function_edges)); + } + } + + return binary_cfgs; +} + +absl::StatusOr< + absl::flat_hash_map>> +LoadCorpusCfgs(absl::Span module_paths) { + absl::flat_hash_map> + binary_cfgs; + + for (const std::string& module_path : module_paths) { + ASSIGN_OR_RETURN(auto module_cfgs, LoadBinaryCfgs(module_path)); + + binary_cfgs.insert(module_cfgs.begin(), module_cfgs.end()); + } + + return binary_cfgs; +} + +bool AreCfgsDifferent(absl::Span function_a_cfg, + absl::Span function_b_cfg) { + if (function_a_cfg.size() != function_b_cfg.size()) { + return true; + } + + for (size_t i = 0; i < function_a_cfg.size(); ++i) { + if (function_a_cfg[i].from_block != function_b_cfg[i].from_block || + function_a_cfg[i].to_block != function_b_cfg[i].to_block) { + LOG(INFO) << "Expected CFG edge (" << function_a_cfg[i].from_block << "," + << function_a_cfg[i].to_block << ") and (" + << function_b_cfg[i].from_block << "," + << function_b_cfg[i].to_block << ") to be the same."; + return true; + } + } + + return false; +} + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/compare_binary_cfgs_lib.h b/compiler_opt/memtrace_costmodel/compare_binary_cfgs_lib.h new file mode 100644 index 00000000..a545cb83 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/compare_binary_cfgs_lib.h @@ -0,0 +1,37 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_COMPARE_BINARY_CFGS_LIB_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_COMPARE_BINARY_CFGS_LIB_H_ + +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" + +namespace mlgo { +namespace latency_model { + +struct ControlFlowEdge { + uint32_t from_block = 0; + uint32_t to_block = 0; +}; + +absl::StatusOr< + absl::flat_hash_map>> +LoadBinaryCfgs(absl::string_view binary_path); + +absl::StatusOr< + absl::flat_hash_map>> +LoadCorpusCfgs(absl::Span module_paths); + +// Detects if the CFGs are different by directly comparing the CFG Edge +// vectors. They are assumed to be sorted. +bool AreCfgsDifferent(absl::Span function_a_cfg, + absl::Span function_b_cfg); + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_COMPARE_BINARY_CFGS_LIB_H_ diff --git a/compiler_opt/memtrace_costmodel/costmodel.h b/compiler_opt/memtrace_costmodel/costmodel.h new file mode 100644 index 00000000..0881fb0d --- /dev/null +++ b/compiler_opt/memtrace_costmodel/costmodel.h @@ -0,0 +1,31 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_COSTMODEL_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_COSTMODEL_H_ + +#include "llvm/MC/MCInst.h" + +namespace mlgo { +namespace latency_model { + +// This class is an abstraction around a cost model used to estimate the +// performance characteristics of a given sequence of instructions. +// It serves as a wrapper around the specific cost model implementation, +// providing a common interface for all cost models. +class CostModel { + public: + virtual ~CostModel() = default; + + // This function should be called for each instruction in a stream of + // execution whose cost we are interested in modeling, such as individual + // basic blocks or longer trace segments. + virtual void AddInstruction(const llvm::MCInst& new_instruction) = 0; + + // This function should be called once all `AddInstruction` has been called + // for all instructions in the stream being modeled. Returns the cost of the + // instruction stream. + virtual double GetCost() = 0; +}; + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_COSTMODEL_H_ diff --git a/compiler_opt/memtrace_costmodel/costmodel_factory.cc b/compiler_opt/memtrace_costmodel/costmodel_factory.cc new file mode 100644 index 00000000..49155d5f --- /dev/null +++ b/compiler_opt/memtrace_costmodel/costmodel_factory.cc @@ -0,0 +1,94 @@ +#include "compiler_opt/memtrace_costmodel/costmodel_factory.h" + +#include +#include +#include + +#include "absl/log/log.h" +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/costmodel.h" +#include "compiler_opt/memtrace_costmodel/instruction_counting_costmodel.h" +#include "compiler_opt/memtrace_costmodel/print_costmodel.h" +#include "compiler_opt/memtrace_costmodel/trace_segment_mca.h" +#include "llvm/Target/TargetMachine.h" + +namespace mlgo { +namespace latency_model { + +std::string AbslUnparseFlag(CostModelType cost_model_type) { + switch (cost_model_type) { + case (CostModelType::Mca): + return "mca"; + case (CostModelType::Print): + return "print"; + case (CostModelType::InstructionCounting): + return "instruction_counting"; + default: + LOG(QFATAL) << "Cannot unparse cost model type " + << static_cast(cost_model_type); + } +} + +bool AbslParseFlag(absl::string_view text, CostModelType* cost_model_type, + std::string* error) { + if (text == "mca") { + *cost_model_type = CostModelType::Mca; + return true; + } else if (text == "print") { + *cost_model_type = CostModelType::Print; + return true; + } else if (text == "instruction_counting") { + *cost_model_type = CostModelType::InstructionCounting; + return true; + } else { + *error = "There is no cost model with the given name,"; + return false; + } +} + +void ValidateCostModelFlags(const CostModelType model_type, + absl::string_view gematria_model_path, + absl::string_view print_output_file) { + if (model_type == CostModelType::Print && print_output_file.empty()) { + LOG(QFATAL) << "--print_output_file must be specified when --model_type is " + "\"print\"."; + } +} + +std::function()> CreateCostModelFactory( + const CostModelType model_type, llvm::TargetMachine* target_machine, + absl::string_view llvm_cpu_name, absl::string_view target_triple, + absl::string_view gematria_model_path, const int gematria_task_index, + absl::string_view print_output_file) { + if (model_type == CostModelType::Mca) { + return [llvm_cpu_name = std::string(llvm_cpu_name), + target_triple = + std::string(target_triple)]() -> std::unique_ptr { + std::unique_ptr cost_model = + std::make_unique(target_triple, + llvm_cpu_name); + return cost_model; + }; + } else if (model_type == CostModelType::Print) { + return [target_triple = std::string(target_triple), + llvm_cpu_name = std::string(llvm_cpu_name), + print_output_file = std::string( + print_output_file)]() -> std::unique_ptr { + std::unique_ptr cost_model = + std::make_unique( + target_triple, llvm_cpu_name, print_output_file); + return cost_model; + }; + } else if (model_type == CostModelType::InstructionCounting) { + return []() -> std::unique_ptr { + std::unique_ptr cost_model = + std::make_unique(); + return cost_model; + }; + } else { + LOG(QFATAL) << "Unknown model type: " << static_cast(model_type); + } +} + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/costmodel_factory.h b/compiler_opt/memtrace_costmodel/costmodel_factory.h new file mode 100644 index 00000000..a3109e55 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/costmodel_factory.h @@ -0,0 +1,44 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_COSTMODEL_FACTORY_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_COSTMODEL_FACTORY_H_ + +#include +#include +#include + +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/costmodel.h" +#include "llvm/Target/TargetMachine.h" + +namespace mlgo { +namespace latency_model { + +enum class CostModelType { + Invalid, + Mca, + Print, + InstructionCounting, + CacheLines +}; + +std::string AbslUnparseFlag(CostModelType cost_model_type); +bool AbslParseFlag(absl::string_view text, CostModelType* cost_model_type, + std::string* error); + +// Validates all cost-model related flags, ensuring that the necessary +// parameters are set for the cost model that is requested. +void ValidateCostModelFlags(CostModelType model_type, + absl::string_view gematria_model_path, + absl::string_view print_output_file); + +// Creates a factory function that will create the requested cost model +// with the provided flags. +std::function()> CreateCostModelFactory( + CostModelType model_type, llvm::TargetMachine* target_machine, + absl::string_view llvm_cpu_name, absl::string_view target_triple, + absl::string_view gematria_model_path, int gematria_task_index, + absl::string_view print_output_file); + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_COSTMODEL_FACTORY_H_ diff --git a/compiler_opt/memtrace_costmodel/elf_metadata_parser.cc b/compiler_opt/memtrace_costmodel/elf_metadata_parser.cc new file mode 100644 index 00000000..1df5e4be --- /dev/null +++ b/compiler_opt/memtrace_costmodel/elf_metadata_parser.cc @@ -0,0 +1,188 @@ +#include "compiler_opt/memtrace_costmodel/elf_metadata_parser.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/functional/any_invocable.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/memory/memory.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "compiler_opt/memtrace_costmodel/status_macros.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallSet.h" +#include "llvm/ADT/StringExtras.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Object/Binary.h" +#include "llvm/Object/BuildID.h" +#include "llvm/Object/ELFObjectFile.h" +#include "llvm/Object/ELFTypes.h" +#include "llvm/Object/ObjectFile.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Error.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/MemoryBufferRef.h" +#include "llvm/Support/raw_ostream.h" + +namespace { + +template +absl::StatusOr AsStatusOr(llvm::Expected&& expected) { + if (expected) { + return std::move(*expected); + } + std::string ret; + llvm::raw_string_ostream OS(ret); + OS << expected.takeError(); + return absl::InternalError(ret); +} + +template +void InsertIntoVectorAtPos(std::vector* vector, int index, V value) { + if (vector->size() < index + 1) vector->resize(index + 1); + (*vector)[index] = value; +} + +} // namespace + +namespace mlgo { +namespace latency_model { + +UnstrippedBinaryProcessor::UnstrippedBinaryProcessor( + std::unique_ptr buffer, + std::unique_ptr elfobj) + : buffer_(std::move(buffer)), elfobj_(std::move(elfobj)) {} + +UnstrippedBinaryProcessor::~UnstrippedBinaryProcessor() = default; + +absl::StatusOr> +UnstrippedBinaryProcessor::Create(absl::string_view unstripped_binary_path) { + llvm::ErrorOr> buffer_or_err = + llvm::MemoryBuffer::getFile(llvm::StringRef( + unstripped_binary_path.data(), unstripped_binary_path.size())); + if (std::error_code ec = buffer_or_err.getError()) { + return absl::InternalError( + absl::StrCat("Failed to open binary: ", ec.message())); + } + std::unique_ptr buffer = std::move(*buffer_or_err); + + llvm::Expected> obj_binary_or_err = + llvm::object::createBinary(buffer->getMemBufferRef()); + if (!obj_binary_or_err) { + std::string err_str; + llvm::raw_string_ostream OS(err_str); + OS << obj_binary_or_err.takeError(); + return absl::InternalError(err_str); + } + std::unique_ptr obj_binary = + std::move(*obj_binary_or_err); + + if (llvm::isa(obj_binary.get())) { + auto elf_obj = std::unique_ptr( + llvm::cast(obj_binary.release())); + return absl::WrapUnique( + new UnstrippedBinaryProcessor(std::move(buffer), std::move(elf_obj))); + } + return absl::InternalError( + absl::StrCat("Non-elf binary: ", unstripped_binary_path)); +} + +std::string UnstrippedBinaryProcessor::GetLinkerBuildID() const { + return ::llvm::toHex(::llvm::object::getBuildID(elfobj_.get()), + /*lowercase=*/true); +} + +// NameInfo contains both a function name and the section that it is contained +// in so that can provide the user section information later on to disambiguate +// between functions in relocatable object files compiled with +// -ffunction-sections. +struct NameInfo { + ::llvm::StringRef function_name; + uint64_t function_section = 0; + + bool operator==(const NameInfo& other) const = default; +}; + +// We defer to sorting the function names as they are guaranteed to be unique +// and in many cases the vast majority of functions will all have the same +// section index. +bool operator<(const NameInfo& lhs, const NameInfo& rhs) { + return lhs.function_name < rhs.function_name; +} + +absl::Status UnstrippedBinaryProcessor::ProcessBBAddrMap( + absl::AnyInvocable + record_processor) const { + // Because of aliasing, an address would potentially be mapped to more than + // one name. We'll report all that back as separate entries with different + // names but same mapping. One of them should match what the compiler has. + absl::flat_hash_map> address_to_name; + for (const ::llvm::object::ELFSymbolRef& symbol : elfobj_->symbols()) { + if (symbol.getSize() == 0) { + continue; + } + + ASSIGN_OR_RETURN(const auto symbol_type, AsStatusOr(symbol.getType())); + + if (symbol_type != ::llvm::object::SymbolRef::ST_Function) continue; + ASSIGN_OR_RETURN(const auto address, AsStatusOr(symbol.getAddress())); + + // We want to skip symbols at address 0, but only in non-relocatable + // binaries as real function symbols can exist at address 0 in relocatable + // object files. + if (address == 0 && !elfobj_->isRelocatableObject()) continue; + ASSIGN_OR_RETURN(::llvm::StringRef name, AsStatusOr(symbol.getName())); + + ASSIGN_OR_RETURN(::llvm::object::section_iterator section_it, + AsStatusOr(symbol.getSection())); + address_to_name[address].insert( + {.function_name = name, .function_section = section_it->getIndex()}); + } + std::vector pgo_data; + ASSIGN_OR_RETURN(const auto bb_addr_map_list, + AsStatusOr(elfobj_->readBBAddrMap(std::nullopt, &pgo_data))); + CHECK_EQ(pgo_data.size(), bb_addr_map_list.size()); + + // Try to avoid churning through allocating/deallocating this buffer, + // and instead, clear it whenever transitioning to a new function. + std::vector bb_infos; + for (const auto& [bb_addr_map, pgo] : llvm::zip(bb_addr_map_list, pgo_data)) { + CHECK_EQ(bb_addr_map.getBBEntries().size(), pgo.BBEntries.size()); + bb_infos.clear(); + auto name_iter = address_to_name.find(bb_addr_map.getFunctionAddress()); + if (name_iter == address_to_name.end()) { + LOG(WARNING) << "Found a BB map entry without a symbol name: 0x" + << absl::Hex(bb_addr_map.getFunctionAddress()); + continue; + } + const auto& names = name_iter->second; + for (const auto& [entry, freq] : + llvm::zip(bb_addr_map.getBBEntries(), pgo.BBEntries)) { + InsertIntoVectorAtPos( + &bb_infos, entry.ID, + {.address = entry.Offset + bb_addr_map.getFunctionAddress(), + .size = entry.Size, + .frequency = freq.BlockFreq.getFrequency()}); + } + for (const auto& name : names) { + FunctionBBInfo current_function_info( + name.function_name, bb_addr_map.getFunctionAddress(), + pgo.FuncEntryCount, name.function_section, bb_infos); + record_processor(current_function_info); + } + } + return absl::OkStatus(); +} + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/elf_metadata_parser.h b/compiler_opt/memtrace_costmodel/elf_metadata_parser.h new file mode 100644 index 00000000..ce9ae4d1 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/elf_metadata_parser.h @@ -0,0 +1,84 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_ELF_METADATA_PARSER_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_ELF_METADATA_PARSER_H_ + +#include +#include +#include + +#include "absl/functional/any_invocable.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "llvm/Object/ELFObjectFile.h" +#include "llvm/Support/MemoryBuffer.h" + +namespace mlgo { +namespace latency_model { + +// Utility that loads sections from an unstripped binary: linker build id, the +// symbols table and the bb address map (assuming -fbasic-block-address-map) +// Together, the latter 2 are used to capture the mbb address map for each +// function, identified by name (we also assume -fno-split-machine-functions). +// Because a lot can go wrong in loading the binary, we expose a factory method +// and hide the implementation. +class UnstrippedBinaryProcessor { + public: + std::string GetLinkerBuildID() const; + + // Information about a machine basic block (BB). + struct BBInfo { + // The address is the start address of the BB. + uint64_t address = 0; + // The size is the size of the BB. + uint32_t size = 0; + // The frequency is the number of times the BB was executed (profile + // information). + uint64_t frequency = 0; + }; + + struct FunctionBBInfo { + explicit FunctionBBInfo(absl::string_view func_name, uint64_t func_address, + uint64_t func_entrycount, uint64_t section_index, + absl::Span bb_infos) + : function_name(func_name), + function_address(func_address), + func_entrycount(func_entrycount), + section_index(section_index), + bb_infos(bb_infos) {} + + FunctionBBInfo() = delete; + + absl::string_view function_name; + uint64_t function_address = 0; + uint64_t func_entrycount = 0; + uint64_t section_index = 0; + absl::Span bb_infos; + }; + + // Call `record_processor` once, passing a vector containing, at index `i`, + // the start binary address for machine basic block with ID `i`; and a vector + // containing, at position 0, the function entrycount, and then at position + // `i` the BB frequency of BB with ID `i - 1`. `record_processor` will be + // called for each alias name separately. + absl::Status ProcessBBAddrMap( + absl::AnyInvocable record_processor) const; + + ~UnstrippedBinaryProcessor(); + + static absl::StatusOr> Create( + absl::string_view unstripped_binary_path); + + private: + std::unique_ptr buffer_; + std::unique_ptr elfobj_; + + UnstrippedBinaryProcessor( + std::unique_ptr buffer, + std::unique_ptr elfobj); +}; + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_ELF_METADATA_PARSER_H_ diff --git a/compiler_opt/memtrace_costmodel/extract_corpus_subset.cc b/compiler_opt/memtrace_costmodel/extract_corpus_subset.cc new file mode 100644 index 00000000..65faae7b --- /dev/null +++ b/compiler_opt/memtrace_costmodel/extract_corpus_subset.cc @@ -0,0 +1,108 @@ +#include +#include +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "absl/log/check.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/extract_corpus_subset_lib.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" +#include "nlohmann/json.hpp" +#include "riegeli/bytes/file_reader.h" +#include "riegeli/records/record_reader.h" + +ABSL_FLAG(std::string, corpus_json_path, "", "The corpus to process."); +ABSL_FLAG(std::string, bb_trace_path, "", "The path to the input trace."); +ABSL_FLAG(std::string, output_path, "", + "The output path to put the corpus subset in."); +ABSL_FLAG(std::string, function_index_path, "", + "The path to the function name to ID mapping."); + +int main(int argc, char** argv) { + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + + if (absl::GetFlag(FLAGS_corpus_json_path).empty()) { + LOG(QFATAL) << "--corpus_path must be set.\n"; + } + + if (absl::GetFlag(FLAGS_bb_trace_path).empty()) { + LOG(QFATAL) << "--bb_trace_path must be set.\n"; + } + + if (absl::GetFlag(FLAGS_output_path).empty()) { + LOG(QFATAL) << "--output_path must be set.\n"; + } + + if (absl::GetFlag(FLAGS_function_index_path).empty()) { + LOG(QFATAL) << "--function_index_path was not specified."; + } + + riegeli::RecordReader function_index_reader( + riegeli::Maker( + absl::GetFlag(FLAGS_function_index_path))); + mlgo::latency_model::FunctionMapping function_name_to_id; + QCHECK(function_index_reader.ReadRecord(function_name_to_id)); + QCHECK(function_index_reader.Close()) << function_index_reader.status(); + + std::string corpus_path = absl::GetFlag(FLAGS_corpus_json_path); + + std::ifstream corpus_file(corpus_path); + QCHECK(corpus_file) << "Failed to open corpus JSON path: " << corpus_path; + std::stringstream buffer; + buffer << corpus_file.rdbuf(); + std::string corpus_json_contents = buffer.str(); + + std::cout << "Loaded corpus JSON\n"; + + nlohmann::json corpus_description = + nlohmann::json::parse(corpus_json_contents); + + std::string corpus_base_path = + std::filesystem::path(corpus_path).parent_path().string(); + + // Check that the corpus has a modules field. + QCHECK(corpus_description.contains("modules")); + QCHECK(corpus_description["modules"].is_array()); + + riegeli::RecordReader trace_reader( + riegeli::Maker(absl::GetFlag(FLAGS_bb_trace_path))); + + mlgo::latency_model::MbbTrace current_trace; + absl::flat_hash_set included_files; + + mlgo::latency_model::FunctionMapping processed_function_mapping = + mlgo::latency_model::ProcessFunctionMappingForModifiedFunctions( + function_name_to_id); + + absl::flat_hash_map function_to_module_name = + mlgo::latency_model::GetFunctionToModuleMapping( + corpus_description["modules"], corpus_base_path, + processed_function_mapping); + + while (trace_reader.ReadRecord(current_trace)) { + absl::flat_hash_set current_trace_files = + mlgo::latency_model::GetIncludedFilesList(function_to_module_name, + current_trace); + included_files.insert(current_trace_files.begin(), + current_trace_files.end()); + } + QCHECK(trace_reader.Close()) << trace_reader.status(); + + std::string output_corpus_base_path = absl::GetFlag(FLAGS_output_path); + + // The first BB in the trace should be from the entrypoint function. + QCHECK_OK(mlgo::latency_model::CopySubsetCorpus( + corpus_base_path, output_corpus_base_path, included_files, + corpus_description)); + + return 0; +} diff --git a/compiler_opt/memtrace_costmodel/extract_corpus_subset_lib.cc b/compiler_opt/memtrace_costmodel/extract_corpus_subset_lib.cc new file mode 100644 index 00000000..47888f8e --- /dev/null +++ b/compiler_opt/memtrace_costmodel/extract_corpus_subset_lib.cc @@ -0,0 +1,374 @@ +#include "compiler_opt/memtrace_costmodel/extract_corpus_subset_lib.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "compiler_opt/memtrace_costmodel/status_macros.h" +#include "llvm/ADT/StableHashing.h" +#include "llvm/IR/Attributes.h" +#include "llvm/IR/GlobalAlias.h" +#include "llvm/IR/GlobalValue.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/StructuralHash.h" +#include "llvm/IRReader/IRReader.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/MemoryBufferRef.h" +#include "llvm/Support/SourceMgr.h" + +namespace mlgo { +namespace latency_model { +namespace { + +std::string GetFullModulePath(absl::string_view corpus_path, + absl::string_view module_relative_path) { + return (std::filesystem::path(std::string(corpus_path)) / + std::string(module_relative_path)) + .string() + + ".bc"; +} + +// Contains information on an individual function relevant to corpus +// extraction. Designed to be relatively small so that it can efficiently +// be sent between threads after a bitcode module has been processed into +// a series of this struct. +struct FunctionInfo { + uint32_t function_id = 0; + size_t module_index = 0; + llvm::stable_hash function_hash = 0; + llvm::GlobalValue::LinkageTypes linkage_type = + llvm::GlobalValue::LinkageTypes::ExternalLinkage; +}; + +// Checks if there are differing function definitions present, updating the +// function hashes mapping if necessary. +void CheckDifferingDefinitionsAndUpdateHashes( + absl::flat_hash_map& function_hashes, + const FunctionInfo& current_function) { + auto [function_hashes_it, inserted] = function_hashes.emplace( + current_function.function_id, current_function.function_hash); + if (!inserted) { + if (function_hashes_it->second != current_function.function_hash) { + // We can ignore weak/linkonce linkage types here as lld will always + // resolve to a non-weak symbol, or the first weak symbol that is + // found. So these symbols should never end up being used in a link + // anyways, and thus we can safely skip them. + if (current_function.linkage_type == + llvm::GlobalValue::LinkageTypes::LinkOnceODRLinkage || + current_function.linkage_type == + llvm::GlobalValue::LinkageTypes::WeakAnyLinkage) { + LOG(WARNING) << "Found differing function definitions. Ignoring " + "due to weak linkage.\n"; + return; + } + } + + QCHECK(function_hashes_it->second == current_function.function_hash) + << "Function " << current_function.function_id + << " has multiple differing definitions."; + } +} + +std::vector GetFunctionInfoForModule( + const std::string& module_path, size_t module_index, + const FunctionMapping& function_name_to_id) { + llvm::LLVMContext llvm_context; + + auto memory_buffer_or = llvm::MemoryBuffer::getFile(module_path); + QCHECK(memory_buffer_or) << "Failed to open module file: " << module_path; + std::unique_ptr memory_buffer = + std::move(*memory_buffer_or); + llvm::MemoryBufferRef module_memory_buffer = memory_buffer->getMemBufferRef(); + + llvm::SMDiagnostic parse_error; + std::unique_ptr current_module = + llvm::parseIR(module_memory_buffer, parse_error, llvm_context); + QCHECK(current_module) << "Failed to parse IR for module: " << module_path + << " Error: " << parse_error.getMessage().str(); + + std::vector function_info; + function_info.reserve(current_module->size()); + + for (const llvm::Function& current_function : *current_module) { + // Skip all declarations as we only care about the function definitions. + if (current_function.isDeclaration()) { + continue; + } + + // If the function has the available_externally attribute, it will + // not get emitted into the object file and thus we cannot put it + // in the map as the BB trace modelling tooling works over object + if (current_function.hasAvailableExternallyLinkage()) { + continue; + } + + const auto function_id_it = + function_name_to_id.function_ids().find(current_function.getName()); + // The corpus will have more functions than the original binary, so if we + // fail to find one in the name to ID map, we should simply skip it. + // Missing functions will be caught later. + if (function_id_it == function_name_to_id.function_ids().end()) { + continue; + } + + // TODO: The checks below use the standard StructuralHash + // rather than the detailed StructuralHash due to issues with function + // names being slightly different between definitions in different + // modules. Eventually we should switch over to detailed StructuralHash + // here. + llvm::stable_hash function_hash = + llvm::StructuralHash(current_function, false); + + FunctionInfo current_function_info = { + .function_id = function_id_it->second, + .module_index = module_index, + .function_hash = function_hash, + .linkage_type = current_function.getLinkage()}; + + function_info.push_back(std::move(current_function_info)); + } + + // Additionally capture all global aliases as otherwise we might miss + // functions where multiple symbols point to the same address, but we + // only put one of them into function_to_module_name, leading to missing + // functions. This is particularly common with C++ constructors and + // destructors. + for (const llvm::GlobalAlias& current_alias : current_module->aliases()) { + if (const auto* aliasee_function = + llvm::dyn_cast(current_alias.getAliaseeObject())) { + QCHECK(!aliasee_function->isDeclaration()); + + const auto function_id_it = + function_name_to_id.function_ids().find(current_alias.getName()); + // The corpus will have more functions than the original binary, so if + // we fail to find one in the name to ID map, we should simply skip it. + // Missing functions will be caught later. + if (function_id_it == function_name_to_id.function_ids().end()) { + continue; + } + + // We can just check for differing definitions for the alias name in the + // same way as for functions. If the alias name does not have a function + // ID, then it is not part of the trace and we only care that functions + // with that exact alias name are identical because that is all we will + // resolve. If it does have a function ID and there are multiple differing + // definitions, that will be caught when iterating through the functions. + + // TODO: The checks below use the standard StructuralHash + // rather than the detailed StructuralHash due to issues with function + // names being slightly different between definitions in different + // modules. Eventually we should switch over to detailed StructuralHash + // here. + llvm::stable_hash function_hash = + llvm::StructuralHash(*aliasee_function, false); + FunctionInfo current_function_info = { + .function_id = function_id_it->second, + .module_index = module_index, + .function_hash = function_hash, + .linkage_type = aliasee_function->getLinkage()}; + function_info.push_back(std::move(current_function_info)); + } + } + + return function_info; +} + +} // end anonymous namespace + +absl::flat_hash_map GetFunctionToModuleMapping( + const std::vector& module_paths, absl::string_view corpus_path, + const FunctionMapping& function_name_to_id) { + absl::flat_hash_map function_to_module_names; + absl::flat_hash_map function_hashes; + + for (size_t i = 0; i < module_paths.size(); ++i) { + std::string module_full_path = + GetFullModulePath(corpus_path, module_paths[i]); + + QCHECK(std::filesystem::exists(module_full_path)) + << "Module path does not exist: " << module_full_path; + + std::vector current_module_info = + GetFunctionInfoForModule(module_full_path, i, function_name_to_id); + + for (const FunctionInfo& current_function_info : current_module_info) { + function_to_module_names.emplace( + current_function_info.function_id, + module_paths[current_function_info.module_index]); + + CheckDifferingDefinitionsAndUpdateHashes(function_hashes, + current_function_info); + } + + LOG(INFO) << "Finished processing module " << (i + 1) << "/" + << module_paths.size() << "\n"; + } + + return function_to_module_names; +} + +absl::flat_hash_set GetIncludedFilesList( + const absl::flat_hash_map& function_to_module_names, + const MbbTrace& mbb_trace) { + absl::flat_hash_set included_modules; + + for (const MachineBbId& bb_info : mbb_trace.mbbs()) { + // Skip blocks that do not have function IDs, as they are from shared + // object invocations and recorded and thus do not have a function name + // associated with them. + if (!bb_info.has_function_id()) { + continue; + } + + const auto function_to_module_it = + function_to_module_names.find(bb_info.function_id()); + QCHECK(function_to_module_it != function_to_module_names.end()) + << "Failed to find a function, ID " << bb_info.function_id() + << " in the corpus."; + + included_modules.emplace(function_to_module_it->second); + } + + return included_modules; +} + +absl::Status CopySubsetCorpus( + absl::string_view input_corpus_base_path, + absl::string_view output_corpus_base_path, + const absl::flat_hash_set& included_modules, + nlohmann::json& corpus_description) { + // See if the output path exists and if it does not, create it. + std::error_code ec; + std::filesystem::create_directories(std::string(output_corpus_base_path), ec); + if (ec) { + return absl::InternalError( + absl::StrCat("Failed to create directory: ", output_corpus_base_path, + " Error: ", ec.message())); + } + + // Copy the modules included in the corpus subset over to the output + // directory. + for (const std::string& included_module : included_modules) { + std::string input_module_full_path = + (std::filesystem::path(std::string(input_corpus_base_path)) / + included_module) + .string(); + std::string output_module_full_path = + (std::filesystem::path(std::string(output_corpus_base_path)) / + included_module) + .string(); + + // Create the output subdirectory if it does not exist. + std::filesystem::create_directories( + std::filesystem::path(output_module_full_path).parent_path(), ec); + if (ec) { + return absl::InternalError(absl::StrCat( + "Failed to create directory: ", + std::filesystem::path(output_module_full_path).parent_path().string(), + " Error: ", ec.message())); + } + + std::filesystem::copy_file( + input_module_full_path + ".bc", output_module_full_path + ".bc", + std::filesystem::copy_options::overwrite_existing, ec); + if (ec) { + return absl::InternalError( + absl::StrCat("Failed to copy file: ", input_module_full_path + ".bc", + " Error: ", ec.message())); + } + + std::filesystem::copy_file( + input_module_full_path + ".cmd", output_module_full_path + ".cmd", + std::filesystem::copy_options::overwrite_existing, ec); + if (ec) { + return absl::InternalError( + absl::StrCat("Failed to copy file: ", input_module_full_path + ".cmd", + " Error: ", ec.message())); + } + + // If we have a distributed ThinLTO corpus, we need to copy over the + // ThinLTO index files (.thinlto.bc) to the output directory. + if (corpus_description["has_thinlto"]) { + std::filesystem::copy_file( + input_module_full_path + ".thinlto.bc", + output_module_full_path + ".thinlto.bc", + std::filesystem::copy_options::overwrite_existing, ec); + if (ec) { + return absl::InternalError(absl::StrCat( + "Failed to copy file: ", input_module_full_path + ".thinlto.bc", + " Error: ", ec.message())); + } + } + } + + // Modify the corpus description to only include the necessary files and + // then copy them over to the output directory. + corpus_description["modules"].clear(); + + for (const std::string& included_module : included_modules) { + corpus_description["modules"].push_back(included_module); + } + + // Sort the output modules to ensure deterministic output of the corpus + // description. This is mainly useful for testing purposes. + std::sort(corpus_description["modules"].begin(), + corpus_description["modules"].end()); + + std::string output_corpus_description_path = + (std::filesystem::path(std::string(output_corpus_base_path)) / + "corpus_description.json") + .string(); + + // Write out the corpus description. Use an indent of two so the + // corpus description is appropriately indented for manual inspection + // and/or modification. + std::ofstream output_file(output_corpus_description_path); + if (!output_file) { + return absl::InternalError( + absl::StrCat("Failed to open output corpus description path: ", + output_corpus_description_path)); + } + output_file << corpus_description.dump(2); + if (!output_file.good()) { + return absl::InternalError( + absl::StrCat("Failed to write corpus description to path: ", + output_corpus_description_path)); + } + + return absl::OkStatus(); +} + +constexpr absl::string_view kFunctionSpecializationSentinel = ".specialized."; + +FunctionMapping ProcessFunctionMappingForModifiedFunctions( + const FunctionMapping& function_mapping) { + FunctionMapping processed_function_mapping; + + for (const auto& [function_name, function_id] : + function_mapping.function_ids()) { + if (absl::StrContains(function_name, kFunctionSpecializationSentinel)) { + absl::string_view standard_function_name = + *absl::StrSplit(function_name, kFunctionSpecializationSentinel) + .begin(); + processed_function_mapping.mutable_function_ids()->try_emplace( + standard_function_name, function_id); + } + + processed_function_mapping.mutable_function_ids()->try_emplace( + function_name, function_id); + } + + return processed_function_mapping; +} + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/extract_corpus_subset_lib.h b/compiler_opt/memtrace_costmodel/extract_corpus_subset_lib.h new file mode 100644 index 00000000..d72f92c1 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/extract_corpus_subset_lib.h @@ -0,0 +1,50 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_EXTRACT_CORPUS_SUBSET_LIB_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_EXTRACT_CORPUS_SUBSET_LIB_H_ + +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/status/status.h" +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" +#include "nlohmann/json.hpp" + +namespace mlgo { +namespace latency_model { + +// Returns a mapping of functions to the modules that they are +// contained within. +absl::flat_hash_map GetFunctionToModuleMapping( + const std::vector& module_paths, absl::string_view corpus_path, + const FunctionMapping& function_name_to_id); + +// Gets the minimal set of modules that need to be included in the corpus +// to completely cover all the functions in the MBB trace. +absl::flat_hash_set GetIncludedFilesList( + const absl::flat_hash_map& function_to_module_names, + const MbbTrace& mbb_trace); + +// Copies the corpus subset containing the relevant modules in the trace +// to the output corpus path in addition to modifying the corpus +// description to fit the module set in the subset. +// TODO: We ideally want to avoid the JSON dependency here +// and pass in a better representation of the corpus description. +absl::Status CopySubsetCorpus( + absl::string_view input_corpus_base_path, + absl::string_view output_corpus_base_path, + const absl::flat_hash_set& included_modules, + nlohmann::json& corpus_description); + +// Processes a function mapping and adds additional entries pointing at +// existing function IDs to handle cases like function specialization +// and hot cold splitting that occur within PGO optimized binaries. +FunctionMapping ProcessFunctionMappingForModifiedFunctions( + const FunctionMapping& function_mapping); + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_EXTRACT_CORPUS_SUBSET_LIB_H_ diff --git a/compiler_opt/memtrace_costmodel/instruction_counting_costmodel.cc b/compiler_opt/memtrace_costmodel/instruction_counting_costmodel.cc new file mode 100644 index 00000000..9931b276 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/instruction_counting_costmodel.cc @@ -0,0 +1,16 @@ +#include "compiler_opt/memtrace_costmodel/instruction_counting_costmodel.h" + +#include "llvm/MC/MCInst.h" + +namespace mlgo { +namespace latency_model { + +void InstructionCountingCostModel::AddInstruction( + const llvm::MCInst& new_instruction) { + total_cost_ += 1.0; +} + +double InstructionCountingCostModel::GetCost() { return total_cost_; } + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/instruction_counting_costmodel.h b/compiler_opt/memtrace_costmodel/instruction_counting_costmodel.h new file mode 100644 index 00000000..ce262477 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/instruction_counting_costmodel.h @@ -0,0 +1,26 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_INSTRUCTION_COUNTING_COSTMODEL_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_INSTRUCTION_COUNTING_COSTMODEL_H_ + +#include "compiler_opt/memtrace_costmodel/costmodel.h" +#include "llvm/MC/MCInst.h" + +namespace mlgo { +namespace latency_model { + +class InstructionCountingCostModel : public CostModel { + public: + void AddInstruction(const llvm::MCInst& new_instruction) override; + + // This function does not actually get the cost of the instruction trace, + // rather just serializing the instructions to a text file where they can + // be inspected later. + double GetCost() override; + + private: + double total_cost_ = 0.0; +}; + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_INSTRUCTION_COUNTING_COSTMODEL_H_ diff --git a/compiler_opt/memtrace_costmodel/mbb_trace.proto b/compiler_opt/memtrace_costmodel/mbb_trace.proto new file mode 100644 index 00000000..78973d2a --- /dev/null +++ b/compiler_opt/memtrace_costmodel/mbb_trace.proto @@ -0,0 +1,27 @@ +edition = "2024"; + +package mlgo.latency_model; + +option features.utf8_validation = NONE; + +message MachineBbId { + uint32 function_id = 3; + uint32 basic_block_id = 2; + uint32 entry_id = 4; + + reserved 1; +} + +message SharedObjectTrace { + uint32 trace_id = 1; + repeated bytes instruction_data = 2; +} + +message MbbTrace { + repeated MachineBbId mbbs = 1; + repeated SharedObjectTrace shared_object_traces = 2; +} + +message FunctionMapping { + map function_ids = 1; +} diff --git a/compiler_opt/memtrace_costmodel/memtrace_costmodel.cc b/compiler_opt/memtrace_costmodel/memtrace_costmodel.cc new file mode 100644 index 00000000..c1f280c1 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/memtrace_costmodel.cc @@ -0,0 +1,315 @@ +#include "compiler_opt/memtrace_costmodel/memtrace_costmodel.h" + +#include +#include +#include +#include + +#include "absl/functional/function_ref.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" +#include "compiler_opt/memtrace_costmodel/costmodel.h" + +// Standard DynamoRIO inclusions +#include "drmemtrace/analysis_tool.h" +#include "drmemtrace/analyzer.h" +#include "drmemtrace/memref.h" +#include "drmemtrace/raw2trace_shared.h" + +// Undefine conflicting DynamoRIO macro before including LLVM headers +#undef X86 +#undef X86_64 + +#include "llvm/MC/MCContext.h" +#include "llvm/MC/MCDisassembler/MCDisassembler.h" +#include "llvm/MC/MCInst.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/Target/TargetOptions.h" +#include "llvm/lib/Target/X86/MCTargetDesc/X86MCTargetDesc.h" + +namespace mlgo { +namespace latency_model { + +using ::dynamorio::drmemtrace::memref_t; + +absl::Status GetMemtraceCost( + std::string file_name, + absl::Span symbols_of_interest_addresses, + absl::FunctionRef)> entrypoint_cost_processor, + absl::FunctionRef()> cost_model_factory, + bool split_on_segment, std::string target_triple, std::string cpu_name) { + for (uint64_t symbol_of_interest_address : symbols_of_interest_addresses) { + CHECK_NE(symbol_of_interest_address, 0); + } + + // The module mapper is used to read the module list and map instruction + // addresses. For offline traces, it can be created from the trace directory + // metadata. + std::unique_ptr module_mapper = + dynamorio::drmemtrace::module_mapper_t::create(nullptr); + + auto memtrace_processor = std::make_unique( + module_mapper.get(), symbols_of_interest_addresses, + entrypoint_cost_processor, cost_model_factory, target_triple, cpu_name, + split_on_segment); + + std::vector tools; + tools.push_back(memtrace_processor.get()); + + // Reimplemented using standard, portable DynamoRIO analyzer_t. + dynamorio::drmemtrace::analyzer_t analyzer(file_name, &tools[0], 1); + if (!analyzer.run()) { + return absl::InternalError("Failed to run memtrace analyzer"); + } + + return absl::OkStatus(); +} + +void CostModelCallstackContext::PushToLeafstack(uint64_t addr) { + if (IsPcDirectlyUnderEntrypoint(addr)) { + ++entrypoint_count_; + } + leafstack_.push_back(addr); +} + +void CostModelCallstackContext::PopLeafstack() { + if (leafstack_.empty()) return; + if (IsPcDirectlyUnderEntrypoint(leafstack_.back())) { + --entrypoint_count_; + if (entrypoint_count_ == 0) { + just_left_point_ = true; + } + } + leafstack_.pop_back(); +} + +uint64_t CostModelCallstackContext::ElfAddressFromPc(uint64_t current_pc) { + // Standard DynamoRIO traces are post-processed by raw2trace, which maps all + // runtime PCs back to binary segment offsets (ELF PCs), resolving ASLR. + // Thus, trace PCs are already identical to ELF addresses. + // If further mapping is needed, the module_mapper_t can be used to map trace + // addresses. + if (module_mapper_ != nullptr) { + app_pc mapped_pc = module_mapper_->find_mapped_trace_address( + reinterpret_cast(current_pc)); + if (mapped_pc != nullptr) { + return reinterpret_cast(mapped_pc); + } + } + return current_pc; +} + +bool CostModelCallstackContext::IsPcDirectlyUnderEntrypoint( + uint64_t current_pc) { + uint64_t elf_address = ElfAddressFromPc(current_pc); + for (uint64_t symbol_of_interest_address : symbols_of_interest_addresses_) { + if (elf_address == symbol_of_interest_address) { + return true; + } + } + return false; +} + +llvm::MCInst getInstructionFromBytes(llvm::ArrayRef instruction_data, + llvm::MCDisassembler& disassembler, + uint64_t& instruction_size, + uint64_t instruction_address) { + llvm::MCInst instruction; + std::string disassembler_output_buffer; + llvm::raw_string_ostream output_stream(disassembler_output_buffer); + + const llvm::MCDisassembler::DecodeStatus status = disassembler.getInstruction( + instruction, instruction_size, instruction_data, instruction_address, + output_stream); + output_stream.flush(); + QCHECK_EQ(status, llvm::MCDisassembler::DecodeStatus::Success) + << disassembler_output_buffer; + + return instruction; +} + +CostMemtraceProcessor::CostMemtraceProcessor( + dynamorio::drmemtrace::module_mapper_t* module_mapper, + absl::Span symbols_of_interest_addresses, + absl::FunctionRef)> entrypoint_cost_processor, + absl::FunctionRef()> cost_model_factory, + std::string target_triple, std::string cpu_name, bool split_on_segment) + : module_mapper_(module_mapper), + symbols_of_interest_addresses_(symbols_of_interest_addresses), + entrypoint_cost_processor_(entrypoint_cost_processor), + cost_model_factory_(cost_model_factory), + target_triple_(target_triple), + cpu_name_(cpu_name), + split_on_segment_(split_on_segment) { + // Initialize LLVM. + std::string possible_lookup_error; + llvm::Triple triple(target_triple_); + const llvm::Target* const llvm_target = + llvm::TargetRegistry::lookupTarget(triple, possible_lookup_error); + QCHECK(llvm_target) << possible_lookup_error; + + llvm::TargetOptions llvm_target_options; + + llvm_target_machine_.reset(llvm_target->createTargetMachine( + triple, /*CPU*/ "", /*Features*/ "", llvm_target_options, std::nullopt)); + QCHECK(llvm_target_machine_); + llvm_mc_context_ = std::make_unique( + llvm_target_machine_->getTargetTriple(), + llvm_target_machine_->getMCAsmInfo(), + llvm_target_machine_->getMCRegisterInfo(), + llvm_target_machine_->getMCSubtargetInfo()); + QCHECK(llvm_mc_context_); +} + +void* CostMemtraceProcessor::parallel_shard_init(int shard_index, + void* worker_data) { + std::string possible_lookup_error; + const llvm::Target* const llvm_target = llvm::TargetRegistry::lookupTarget( + llvm::Triple(target_triple_), possible_lookup_error); + QCHECK(llvm_target); + + std::unique_ptr cost_model = cost_model_factory_(); + PerThreadData* thread_data = new PerThreadData( + [this]() { + return std::make_unique( + module_mapper_, symbols_of_interest_addresses_); + }, + std::move(cost_model), + std::unique_ptr(llvm_target->createMCDisassembler( + llvm_target_machine_->getMCSubtargetInfo(), *llvm_mc_context_))); + return reinterpret_cast(thread_data); +} + +bool CostMemtraceProcessor::parallel_shard_exit(void* shard_data) { + PerThreadData* current_shard_data = + reinterpret_cast(shard_data); + delete current_shard_data; + return true; +} + +bool CostMemtraceProcessor::process_memref(const memref_t& memref) { + LOG(QFATAL) << "Intentionally not implemented"; + return false; +} + +void CostMemtraceProcessor::PushCostAndResetModel( + PerThreadData& current_shard_data) { + current_shard_data.entrypoint_segment_costs.push_back( + current_shard_data.segment_cost_model->GetCost()); + current_shard_data.segment_cost_model = cost_model_factory_(); +} + +bool CostMemtraceProcessor::parallel_shard_memref(void* shard_data, + const memref_t& memref) { + PerThreadData* current_shard_data = + reinterpret_cast(shard_data); + + // Context updating logic is performed manually or using a portable tracker. + // We update the callstack tracker depending on the instruction memref type. + if (dynamorio::drmemtrace::type_is_instr(memref.instr.type)) { + // Simple tracking of calls and returns: + if (memref.instr.type == + dynamorio::drmemtrace::TRACE_TYPE_INSTR_DIRECT_CALL || + memref.instr.type == + dynamorio::drmemtrace::TRACE_TYPE_INSTR_INDIRECT_CALL) { + current_shard_data->callstack.Current().PushToLeafstack( + memref.instr.addr); + } else if (memref.instr.type == + dynamorio::drmemtrace::TRACE_TYPE_INSTR_RETURN) { + current_shard_data->callstack.Current().PopLeafstack(); + } + } + + // Skip all non-instruction memrefs. + if (!dynamorio::drmemtrace::type_is_instr(memref.instr.type)) return true; + + if (current_shard_data->previously_under_segment && + !current_shard_data->callstack.Current().UnderEntrypoint() && + split_on_segment_) { + // We have just left an entrypoint segment. We need to add the current + // cost and recreate the cost model, but not return the vector to the user + // yet. + PushCostAndResetModel(*current_shard_data); + } + + if (current_shard_data->callstack.Current().JustLeftEntrypoint()) { + current_shard_data->callstack.Current().ResetJustLeftEntrypoint(); + + // If we are not splitting by segment, then we need to capture the + // cost for the entire entrypoint here. + if (!split_on_segment_) { + PushCostAndResetModel(*current_shard_data); + } + + { + absl::MutexLock lock(&cost_processor_mutex_); + entrypoint_cost_processor_( + std::move(current_shard_data->entrypoint_segment_costs)); + } + } + + current_shard_data->previously_under_segment = + current_shard_data->callstack.Current().UnderEntrypoint(); + + // Skip instructions that are not under the entrypoint of interest. + if (!current_shard_data->callstack.Current().UnderEntrypoint()) return true; + + // Process the instruction into LLVM MCInsts so that we can utilize LLVM + // tooling for downstream processing. We loop through until the total + // instruction size is equal to the size of the encoding captured in the + // memref as we might have multiple MCInsts in a single instruction, for + // example with a lock prefix before a cmpxchg instruction. + uint64_t total_instruction_size = 0; + while (total_instruction_size < memref.instr.size) { + llvm::ArrayRef instruction_data( + reinterpret_cast(memref.instr.encoding + + total_instruction_size), + memref.instr.size); + uint64_t instruction_size = 0; + llvm::MCInst current_instruction = getInstructionFromBytes( + instruction_data, *current_shard_data->llvm_mc_disassembler, + instruction_size, memref.instr.addr); + + ProcessMachineInstruction(shard_data, current_instruction); + + total_instruction_size += instruction_size; + } + + return true; +} + +bool isInstructionNoop(const llvm::MCInst& machine_instruction) { + if (machine_instruction.getOpcode() == llvm::X86::NOOP || + machine_instruction.getOpcode() == llvm::X86::NOOPL || + machine_instruction.getOpcode() == llvm::X86::NOOPLr || + machine_instruction.getOpcode() == llvm::X86::NOOPQ || + machine_instruction.getOpcode() == llvm::X86::NOOPQr || + machine_instruction.getOpcode() == llvm::X86::NOOPW || + machine_instruction.getOpcode() == llvm::X86::NOOPWr) { + return true; + } + + return false; +} + +void CostMemtraceProcessor::ProcessMachineInstruction( + void* shard_data, const llvm::MCInst& machine_instruction) { + // Do not process no-op instructions as we do not currently (at least for + // MCA) model the (pre)decoder, which means they do not impact the model + // at all. They can also cause differences between the basic block trace + // and the raw disassembled memtrace. + if (isInstructionNoop(machine_instruction)) { + return; + } + + PerThreadData* current_shard_data = + reinterpret_cast(shard_data); + current_shard_data->segment_cost_model->AddInstruction(machine_instruction); +} + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/memtrace_costmodel.h b/compiler_opt/memtrace_costmodel/memtrace_costmodel.h new file mode 100644 index 00000000..07909626 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/memtrace_costmodel.h @@ -0,0 +1,153 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_MEMTRACE_COSTMODEL_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_MEMTRACE_COSTMODEL_H_ + +#include +#include +#include +#include + +#include "absl/functional/function_ref.h" +#include "absl/status/status.h" +#include "absl/types/span.h" +#include "compiler_opt/memtrace_costmodel/costmodel.h" +#include "llvm/MC/MCInst.h" + +// Standard DynamoRIO inclusions +#include "absl/synchronization/mutex.h" +#include "drmemtrace/analysis_tool.h" +#include "drmemtrace/memref.h" +#include "drmemtrace/raw2trace_shared.h" +#include "llvm/MC/MCContext.h" +#include "llvm/MC/MCDisassembler/MCDisassembler.h" +#include "llvm/Target/TargetMachine.h" + +namespace mlgo { +namespace latency_model { + +// Simple portable callstack context tracker stub. +template +class CallstackTracker { + public: + CallstackTracker(std::function()> factory) + : factory_(factory), context_(factory_()) {} + T& Current() { return *context_; } + + private: + std::function()> factory_; + std::unique_ptr context_; +}; + +class CostModelCallstackContext { + public: + CostModelCallstackContext( + dynamorio::drmemtrace::module_mapper_t* module_mapper, + absl::Span symbols_of_interest_addresses) + : module_mapper_(module_mapper), + symbols_of_interest_addresses_(symbols_of_interest_addresses) {}; + + void PushToLeafstack(uint64_t addr); + void PopLeafstack(); + bool UnderEntrypoint() const { return entrypoint_count_ > 0; }; + bool JustLeftEntrypoint() const { return just_left_point_; }; + void ResetJustLeftEntrypoint() { just_left_point_ = false; }; + const std::vector& Leafstack() const { return leafstack_; } + + private: + uint64_t ElfAddressFromPc(uint64_t current_pc); + bool IsPcDirectlyUnderEntrypoint(uint64_t current_pc); + + int64_t entrypoint_count_ = 0; + bool just_left_point_ = false; + std::vector leafstack_; + dynamorio::drmemtrace::module_mapper_t* module_mapper_; + const absl::Span symbols_of_interest_addresses_; +}; + +using CostModelCallstackTracker = CallstackTracker; + +llvm::MCInst getInstructionFromBytes(llvm::ArrayRef instruction_data, + llvm::MCDisassembler& disassembler, + uint64_t& instruction_size, + uint64_t instruction_address = 0); + +bool isInstructionNoop(const llvm::MCInst& machine_instruction); + +// CostMemtraceProcessor inherits from standard DynamoRIO analysis_tool_t +// to keep the trace analysis logic functional and open-source. +class CostMemtraceProcessor : public dynamorio::drmemtrace::analysis_tool_t { + public: + explicit CostMemtraceProcessor( + dynamorio::drmemtrace::module_mapper_t* module_mapper, + absl::Span symbols_of_interest_addresses, + absl::FunctionRef)> entrypoint_cost_processor, + absl::FunctionRef()> cost_model_factory, + std::string target_triple, std::string cpu_name, bool split_on_segment); + + // Implement pure virtual methods of analysis_tool_t + bool process_memref(const dynamorio::drmemtrace::memref_t& entry) override; + bool print_results() override { return true; } + + private: + struct PerThreadData { + explicit PerThreadData( + std::function()> + callstack_factory_, + std::unique_ptr cost_model, + std::unique_ptr mc_disassembler) + : callstack(std::move(callstack_factory_)), + segment_cost_model(std::move(cost_model)), + llvm_mc_disassembler(std::move(mc_disassembler)) {}; + + CostModelCallstackTracker callstack; + std::unique_ptr segment_cost_model; + bool previously_under_segment = false; + std::unique_ptr llvm_mc_disassembler; + std::vector entrypoint_segment_costs = {}; + }; + + void* parallel_shard_init(int shard_index, void* worker_data) override; + bool parallel_shard_exit(void* shard_data) override; + + bool parallel_shard_supported() override { return true; } + + void PushCostAndResetModel(PerThreadData& current_shard_data); + + bool parallel_shard_memref( + void* shard_data, const dynamorio::drmemtrace::memref_t& memref) override; + + // Processes individual instructions, eventually handing them off to the + // underlying cost model. + void ProcessMachineInstruction(void* shard_data, + const llvm::MCInst& machine_instruction); + + std::unique_ptr llvm_target_machine_; + std::unique_ptr llvm_mc_context_; + + dynamorio::drmemtrace::module_mapper_t* module_mapper_; + absl::Span symbols_of_interest_addresses_; + absl::FunctionRef)> entrypoint_cost_processor_; + absl::Mutex cost_processor_mutex_; + + absl::FunctionRef()> cost_model_factory_; + + std::string target_triple_; + std::string cpu_name_; + + bool split_on_segment_ = false; +}; + +// Re-enabled standard GetMemtraceCost using DynamoRIO structures. +// It might not fully compile if some test setups require internal analyzer_t, +// but the core functionality is preserved and adapted to standard DynamoRIO. +absl::Status GetMemtraceCost( + std::string file_name, + absl::Span symbols_of_interest_addresses, + absl::FunctionRef)> entrypoint_cost_processor, + absl::FunctionRef()> cost_model_factory, + bool split_on_segment, std::string target_triple = "x86_64", + std::string cpu_name = "skylake"); + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_MEMTRACE_COSTMODEL_H_ diff --git a/compiler_opt/memtrace_costmodel/memtrace_costmodel_runner.cc b/compiler_opt/memtrace_costmodel/memtrace_costmodel_runner.cc new file mode 100644 index 00000000..fb3e7b98 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/memtrace_costmodel_runner.cc @@ -0,0 +1,123 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "absl/log/check.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/costmodel.h" +#include "compiler_opt/memtrace_costmodel/costmodel_factory.h" +#include "compiler_opt/memtrace_costmodel/memtrace_costmodel.h" +#undef X86 +#undef X86_64 +#include "llvm-c/Target.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/TargetParser/Triple.h" + +struct TargetTripleString { + std::string target_triple = ""; +}; + +std::string AbslUnparseFlag(TargetTripleString target_triple_string) { + return target_triple_string.target_triple; +} + +bool AbslParseFlag(absl::string_view text, + TargetTripleString* target_triple_string, + std::string* error) { + const llvm::Target* const llvm_target = llvm::TargetRegistry::lookupTarget( + llvm::Triple(llvm::StringRef(text)), *error); + if (!llvm_target) return false; + target_triple_string->target_triple = text; + return true; +} + +ABSL_FLAG(std::string, memtrace_path, "", + "The path to the memtrace to process."); +// TODO: Lookup the symbol address and size automatically from +// a symbol name rather than requiring the user specify them manually. +ABSL_FLAG(uint64_t, symbol_address, 0, "The address of the entrypoint symbol"); +ABSL_FLAG(TargetTripleString, target_triple, {"x86_64"}, + "The target triple of the platform to model"); +ABSL_FLAG(std::string, llvm_cpu_name, "skylake", + "The name of the CPU microarchitecture of the platform to model"); +ABSL_FLAG(bool, split_on_segment, false, + "Whether or not to split entrypoints by segment"); + +// Model specific flags. +ABSL_FLAG(mlgo::latency_model::CostModelType, model_type, + mlgo::latency_model::CostModelType::Mca, + "The type of cost model to use. (\"mca\" | \"print\" " + "| \"instruction_counting\")"); +ABSL_FLAG(std::string, print_output_file, "", + "The output file if the print cost model is selected."); + +int main(int argc, char** argv) { +#ifdef __x86_64__ + LLVMInitializeX86Target(); + LLVMInitializeX86TargetInfo(); + LLVMInitializeX86TargetMC(); + LLVMInitializeX86Disassembler(); + LLVMInitializeX86TargetMCA(); +#else +#error memtrace_costmodel_runner is only supported on X86_64. +#endif // __x86_64__ + + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + + if (absl::GetFlag(FLAGS_memtrace_path).empty()) { + LOG(ERROR) << "--memtrace_path was not specified."; + return 1; + } + + if (absl::GetFlag(FLAGS_symbol_address) == 0) { + LOG(ERROR) << "--symbol_address must specify a symbol address."; + return 1; + } + + mlgo::latency_model::ValidateCostModelFlags( + absl::GetFlag(FLAGS_model_type), "", + absl::GetFlag(FLAGS_print_output_file)); + + std::string possible_lookup_error; + llvm::Triple triple(absl::GetFlag(FLAGS_target_triple).target_triple); + const llvm::Target* target = + llvm::TargetRegistry::lookupTarget(triple, possible_lookup_error); + auto target_machine = std::unique_ptr( + target->createTargetMachine(triple, absl::GetFlag(FLAGS_llvm_cpu_name), + "", llvm::TargetOptions(), std::nullopt)); + + std::function()> + cost_model_factory = mlgo::latency_model::CreateCostModelFactory( + absl::GetFlag(FLAGS_model_type), target_machine.get(), + absl::GetFlag(FLAGS_llvm_cpu_name), + absl::GetFlag(FLAGS_target_triple).target_triple, "", 0, + absl::GetFlag(FLAGS_print_output_file)); + + QCHECK_OK(mlgo::latency_model::GetMemtraceCost( + FLAGS_memtrace_path.CurrentValue(), {absl::GetFlag(FLAGS_symbol_address)}, + [](std::vector entrypoint_segment_costs) { + std::cout << "BEGIN ENTRYPOINT\n"; + for (const double entrypoint_segment_cost : entrypoint_segment_costs) { + std::cout << entrypoint_segment_cost << "\n"; + } + std::cout << "END ENTRYPOINT\n"; + }, + std::move(cost_model_factory), absl::GetFlag(FLAGS_split_on_segment), + absl::GetFlag(FLAGS_target_triple).target_triple, + absl::GetFlag(FLAGS_llvm_cpu_name))); + + LOG(INFO) << "Finished processing memtrace."; + + return 0; +} diff --git a/compiler_opt/memtrace_costmodel/print_costmodel.cc b/compiler_opt/memtrace_costmodel/print_costmodel.cc new file mode 100644 index 00000000..8ad50c0c --- /dev/null +++ b/compiler_opt/memtrace_costmodel/print_costmodel.cc @@ -0,0 +1,78 @@ +#include "compiler_opt/memtrace_costmodel/print_costmodel.h" + +#include +#include +#include +#include +#include + +#include "absl/log/check.h" +#include "absl/memory/memory.h" +#include "llvm/IR/InlineAsm.h" +#include "llvm/MC/MCInst.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/TargetParser/Triple.h" + +namespace mlgo { +namespace latency_model { + +PrintCostModel::PrintCostModel(std::string_view target_triple, + std::string_view cpu_name, + std::string_view output_file_path) + : output_file_path_(output_file_path) { + std::string possible_lookup_error; + llvm::Triple triple(target_triple); + const llvm::Target* const llvm_target = + llvm::TargetRegistry::lookupTarget(triple, possible_lookup_error); + QCHECK(llvm_target); + + mc_instruction_info_ = absl::WrapUnique(llvm_target->createMCInstrInfo()); + QCHECK(mc_instruction_info_); + + mc_register_info_ = absl::WrapUnique(llvm_target->createMCRegInfo(triple)); + QCHECK(mc_register_info_); + + llvm::MCTargetOptions target_options; + mc_asm_info_ = absl::WrapUnique( + llvm_target->createMCAsmInfo(*mc_register_info_, triple, target_options)); + QCHECK(mc_asm_info_); + + mc_instruction_printer_ = absl::WrapUnique(llvm_target->createMCInstPrinter( + llvm::Triple(target_triple), llvm::InlineAsm::AD_ATT, *mc_asm_info_, + *mc_instruction_info_, *mc_register_info_)); + + mc_subtarget_info_ = absl::WrapUnique( + llvm_target->createMCSubtargetInfo(triple, cpu_name, "")); +} + +void PrintCostModel::AddInstruction(const llvm::MCInst& new_instruction) { + std::string output_buffer; + llvm::raw_string_ostream output_stream(output_buffer); + mc_instruction_printer_->printInst(&new_instruction, 0, "", + *mc_subtarget_info_, output_stream); + instructions_.push_back(std::move(output_buffer)); +} + +double PrintCostModel::GetCost() { + if (!instructions_.empty()) { + std::string output_buffer; + llvm::raw_string_ostream output_stream(output_buffer); + + for (const std::string_view instruction : instructions_) { + output_stream << instruction << "\n"; + } + + std::ofstream output_file(output_file_path_); + QCHECK(output_file) << "Failed to open print costmodel output file: " + << output_file_path_; + output_file << output_buffer; + QCHECK(output_file.good()) + << "Failed to write to output file: " << output_file_path_; + } + + return 0; +} + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/print_costmodel.h b/compiler_opt/memtrace_costmodel/print_costmodel.h new file mode 100644 index 00000000..d2bb5f8e --- /dev/null +++ b/compiler_opt/memtrace_costmodel/print_costmodel.h @@ -0,0 +1,47 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_PRINT_COSTMODEL_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_PRINT_COSTMODEL_H_ + +#include +#include +#include +#include + +#include "compiler_opt/memtrace_costmodel/costmodel.h" +#include "llvm/MC/MCAsmInfo.h" +#include "llvm/MC/MCInst.h" +#include "llvm/MC/MCInstPrinter.h" +#include "llvm/MC/MCInstrInfo.h" +#include "llvm/MC/MCRegisterInfo.h" +#include "llvm/MC/MCSubtargetInfo.h" + +namespace mlgo { +namespace latency_model { + +class PrintCostModel : public CostModel { + public: + explicit PrintCostModel(std::string_view target_triple, + std::string_view cpu_name, + std::string_view output_file_path); + + void AddInstruction(const llvm::MCInst& new_instruction) override; + + // This function does not actually get the cost of the instruction trace, + // rather just serializing the instructions to a text file where they can + // be inspected later. + double GetCost() override; + + private: + std::unique_ptr mc_instruction_info_; + std::unique_ptr mc_register_info_; + std::unique_ptr mc_asm_info_; + std::unique_ptr mc_instruction_printer_; + std::unique_ptr mc_subtarget_info_; + + std::vector instructions_; + std::string output_file_path_; +}; + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_PRINT_COSTMODEL_H_ diff --git a/compiler_opt/memtrace_costmodel/serialize_trace_functions.cc b/compiler_opt/memtrace_costmodel/serialize_trace_functions.cc new file mode 100644 index 00000000..4e101957 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/serialize_trace_functions.cc @@ -0,0 +1,143 @@ +#include +#include +#include +#include +#include +#include + +#include "absl/container/flat_hash_set.h" +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "absl/log/check.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/basic_block_trace.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" +#include "compiler_opt/memtrace_costmodel/serialize_trace_functions_lib.h" +#include "compiler_opt/memtrace_costmodel/serialized_mbbs.proto.h" +#include "llvm-c/Target.h" +#include "nlohmann/json.hpp" +#include "riegeli/bytes/file_reader.h" +#include "riegeli/bytes/file_writer.h" +#include "riegeli/records/record_reader.h" +#include "riegeli/records/record_writer.h" + +ABSL_FLAG(std::string, trace_path, "", "Path to the trace recordio."); +ABSL_FLAG(std::string, function_index_path, "", + "The path to the function index."); +ABSL_FLAG(std::string, output_path, "", "The path to the output proto."); +ABSL_FLAG(std::string, exclude_functions_list, "", + "The path to the file listing function names to exclude."); +ABSL_FLAG(std::string, corpus_path, "", + "The path to the corpus description JSON to pull blocks from."); +ABSL_FLAG(std::string, target_triple, "x86_64", "The target triple to use."); + +namespace { +std::vector GetObjectPaths(absl::string_view corpus_path) { + std::string corpus_path_str(corpus_path); + std::ifstream corpus_file(corpus_path_str); + QCHECK(corpus_file) << "Failed to open corpus path: " << corpus_path; + std::stringstream buffer; + buffer << corpus_file.rdbuf(); + std::string corpus_description_contents = buffer.str(); + + nlohmann::json corpus_description = + nlohmann::json::parse(corpus_description_contents); + + auto modules_it = corpus_description.find("modules"); + CHECK(modules_it != corpus_description.end()); + + std::vector object_file_paths; + object_file_paths.reserve(modules_it->size()); + + std::string corpus_dir_path = + std::filesystem::path(std::string(corpus_path)).parent_path().string(); + + for (const std::string relative_module_path : *modules_it) { + object_file_paths.push_back( + (std::filesystem::path(corpus_dir_path) / relative_module_path) + .string() + + ".bc.o"); + } + + return object_file_paths; +} +} // namespace + +int main(int argc, char** argv) { + LLVMInitializeX86Target(); + LLVMInitializeX86TargetInfo(); + LLVMInitializeX86TargetMC(); + LLVMInitializeX86Disassembler(); + + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + + const std::string trace_path = absl::GetFlag(FLAGS_trace_path); + const std::string function_index_path = + absl::GetFlag(FLAGS_function_index_path); + const std::string output_path = absl::GetFlag(FLAGS_output_path); + const std::string exclude_functions_list = + absl::GetFlag(FLAGS_exclude_functions_list); + const std::string corpus_path = absl::GetFlag(FLAGS_corpus_path); + + if (trace_path.empty()) { + LOG(QFATAL) << "--trace_path was not specified."; + } + + if (function_index_path.empty()) { + LOG(QFATAL) << "--function_index_path was not specified."; + } + + if (output_path.empty()) { + LOG(QFATAL) << "--output_path was not specified."; + } + + if (exclude_functions_list.empty()) { + LOG(QFATAL) << "--exclude_functions_list was not specified."; + } + + if (corpus_path.empty()) { + LOG(QFATAL) << "--corpus_path was not specified."; + } + + std::vector excluded_functions = + mlgo::latency_model::LoadExcludedFunctions(exclude_functions_list); + + riegeli::RecordReader function_index_reader( + riegeli::Maker(function_index_path)); + mlgo::latency_model::FunctionMapping function_name_to_id; + CHECK(function_index_reader.ReadRecord(function_name_to_id)); + QCHECK(function_index_reader.Close()) << function_index_reader.status(); + + std::vector mbb_traces; + riegeli::RecordReader trace_reader( + riegeli::Maker(trace_path)); + mlgo::latency_model::MbbTrace current_mbb_trace; + while (trace_reader.ReadRecord(current_mbb_trace)) { + mbb_traces.push_back(std::move(current_mbb_trace)); + } + QCHECK(trace_reader.Close()) << trace_reader.status(); + + absl::flat_hash_set + required_mbbs = mlgo::latency_model::GetRequiredMbbs( + mbb_traces, excluded_functions, function_name_to_id); + + std::vector object_file_paths = GetObjectPaths(corpus_path); + + mlgo::latency_model::SerializedMbbs serialized_mbbs = + mlgo::latency_model::GetSerializedMbbs( + object_file_paths, absl::GetFlag(FLAGS_target_triple), + function_name_to_id, required_mbbs); + + riegeli::RecordWriter output_writer( + riegeli::Maker(output_path), + riegeli::RecordWriterBase::Options().set_zstd()); + output_writer.WriteRecord(serialized_mbbs); + QCHECK(output_writer.Close()) << output_writer.status(); + + return 0; +} diff --git a/compiler_opt/memtrace_costmodel/serialize_trace_functions_lib.cc b/compiler_opt/memtrace_costmodel/serialize_trace_functions_lib.cc new file mode 100644 index 00000000..c6e8d8c1 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/serialize_trace_functions_lib.cc @@ -0,0 +1,105 @@ +#include "compiler_opt/memtrace_costmodel/serialize_trace_functions_lib.h" + +#include +#include +#include +#include + +#include "absl/container/flat_hash_set.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "compiler_opt/memtrace_costmodel/basic_block_trace.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" +#include "compiler_opt/memtrace_costmodel/serialized_mbbs.proto.h" +#include "llvm/ADT/ArrayRef.h" + +namespace mlgo { +namespace latency_model { + +absl::flat_hash_set +GetRequiredMbbs(absl::Span mbb_traces, + absl::Span functions_to_exclude, + const FunctionMapping& function_names_to_id) { + absl::flat_hash_set + unique_bbs; + absl::flat_hash_set functions_ids_to_exclude; + + for (const std::string& function_to_exclude : functions_to_exclude) { + const auto function_id_it = + function_names_to_id.function_ids().find(function_to_exclude); + if (function_id_it == function_names_to_id.function_ids().end()) { + LOG(WARNING) << "Failed to find an ID for function " + << function_to_exclude << " in the function mapping"; + continue; + } + + functions_ids_to_exclude.emplace(function_id_it->second); + } + + for (const MbbTrace& mbb_trace : mbb_traces) { + for (const MachineBbId& basic_block : mbb_trace.mbbs()) { + if (functions_ids_to_exclude.contains(basic_block.function_id())) { + continue; + } + + // Blocks without function IDs are shared object traces. We want to skip + // them because they are held within the trace. + if (!basic_block.has_function_id()) { + continue; + } + + unique_bbs.emplace(basic_block); + } + } + + return unique_bbs; +} + +SerializedMbbs GetSerializedMbbs( + const std::vector& bitcode_paths, + const std::string& target_triple, + const FunctionMapping& function_names_to_id, + const absl::flat_hash_set& mbbs_to_include) { + SerializedMbbs mbbs_to_return; + + CorpusApplicationToBbDisassembler corpus_to_bb_contents( + target_triple, bitcode_paths, /*store_block_contents=*/true); + corpus_to_bb_contents.LoadBasicBlocks(function_names_to_id); + + for (const MachineBbId& entry_id : mbbs_to_include) { + llvm::ArrayRef entry_contents = + corpus_to_bb_contents.GetEntryContents(entry_id); + + *mbbs_to_return.add_mbb_ids() = entry_id; + + absl::string_view current_instruction_data( + reinterpret_cast(entry_contents.data()), + entry_contents.size()); + mbbs_to_return.add_mbb_bytes(current_instruction_data); + } + + return mbbs_to_return; +} + +std::vector LoadExcludedFunctions( + absl::string_view excluded_functions_list) { + std::vector excluded_functions; + std::ifstream file((std::string(excluded_functions_list))); + QCHECK(file) << "Failed to open excluded functions list: " + << excluded_functions_list; + std::string line; + while (std::getline(file, line)) { + // Remove carriage return if present (similar to REMOVE_LINEFEED) + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + excluded_functions.push_back(line); + } + return excluded_functions; +} + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/serialize_trace_functions_lib.h b/compiler_opt/memtrace_costmodel/serialize_trace_functions_lib.h new file mode 100644 index 00000000..f8182a38 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/serialize_trace_functions_lib.h @@ -0,0 +1,35 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_SERIALIZE_TRACE_FUNCTIONS_LIB_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_SERIALIZE_TRACE_FUNCTIONS_LIB_H_ + +#include +#include + +#include "absl/container/flat_hash_set.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "compiler_opt/memtrace_costmodel/basic_block_trace.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" +#include "compiler_opt/memtrace_costmodel/serialized_mbbs.proto.h" + +namespace mlgo { +namespace latency_model { + +absl::flat_hash_set +GetRequiredMbbs(absl::Span mbb_traces, + absl::Span functions_to_exclude, + const FunctionMapping& function_names_to_id); + +SerializedMbbs GetSerializedMbbs( + const std::vector& bitcode_paths, + const std::string& target_triple, + const FunctionMapping& function_names_to_id, + const absl::flat_hash_set& mbbs_to_include); + +std::vector LoadExcludedFunctions( + absl::string_view excluded_functions_list); + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_SERIALIZE_TRACE_FUNCTIONS_LIB_H_ diff --git a/compiler_opt/memtrace_costmodel/serialized_mbbs.proto b/compiler_opt/memtrace_costmodel/serialized_mbbs.proto new file mode 100644 index 00000000..d58703a8 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/serialized_mbbs.proto @@ -0,0 +1,12 @@ +edition = "2024"; + +package mlgo.latency_model; + +import "mbb_trace.proto"; + +option features.utf8_validation = NONE; + +message SerializedMbbs { + repeated MachineBbId mbb_ids = 1; + repeated bytes mbb_bytes = 2; +} diff --git a/compiler_opt/memtrace_costmodel/sort_bb_traces.cc b/compiler_opt/memtrace_costmodel/sort_bb_traces.cc new file mode 100644 index 00000000..38fcb334 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/sort_bb_traces.cc @@ -0,0 +1,50 @@ +#include +#include +#include + +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "absl/log/check.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" +#include "compiler_opt/memtrace_costmodel/sort_bb_traces_lib.h" + +ABSL_FLAG(std::string, input_trace, "", "The trace to take as input."); +ABSL_FLAG(std::string, output_path, "", "The output path pattern."); +ABSL_FLAG(int, shard_count, 1, "The shard count."); +ABSL_FLAG( + int, compression_level, 1 << 16, + "The compression level to use while writing the basic block trace protos."); + +int main(int argc, char** argv) { + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + + if (absl::GetFlag(FLAGS_input_trace).empty()) { + LOG(QFATAL) << "--input_trace must be set."; + } + if (absl::GetFlag(FLAGS_output_path).empty()) { + LOG(QFATAL) << "--output_path must be set."; + } + + const std::string input_trace = absl::GetFlag(FLAGS_input_trace); + const std::string output_path = absl::GetFlag(FLAGS_output_path); + const int shard_count = absl::GetFlag(FLAGS_shard_count); + const int compression_level = absl::GetFlag(FLAGS_compression_level); + + std::vector trace_info = + mlgo::latency_model::LoadTraceInfo(input_trace); + LOG(INFO) << "Finished loading trace info."; + + std::vector trace_index_to_shard_map = + mlgo::latency_model::GetTraceIndexToShardMap(trace_info, shard_count); + LOG(INFO) << "Finished planning shards."; + + mlgo::latency_model::WriteOutShards(input_trace, shard_count, + trace_index_to_shard_map, output_path, + compression_level); + + return 0; +} diff --git a/compiler_opt/memtrace_costmodel/sort_bb_traces_lib.cc b/compiler_opt/memtrace_costmodel/sort_bb_traces_lib.cc new file mode 100644 index 00000000..7264f9f1 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/sort_bb_traces_lib.cc @@ -0,0 +1,86 @@ +#include "compiler_opt/memtrace_costmodel/sort_bb_traces_lib.h" + +#include +#include +#include +#include + +#include "absl/log/check.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" +#include "riegeli/bytes/file_reader.h" +#include "riegeli/bytes/file_writer.h" +#include "riegeli/records/record_reader.h" +#include "riegeli/records/record_writer.h" + +namespace mlgo { +namespace latency_model { + +std::vector LoadTraceInfo(absl::string_view trace_file_path) { + riegeli::RecordReader trace_reader( + riegeli::Maker(std::string(trace_file_path))); + mlgo::latency_model::MbbTrace mbb_trace; + std::vector trace_info; + size_t index = 0; + while (trace_reader.ReadRecord(mbb_trace)) { + trace_info.push_back( + TraceInfo{.index = index, .size = mbb_trace.mbbs_size()}); + ++index; + } + QCHECK(trace_reader.Close()) << trace_reader.status(); + return trace_info; +} + +std::vector GetTraceIndexToShardMap(std::vector& traces_info, + size_t shard_count) { + int trace_count = traces_info.back().index + 1; + std::sort( + traces_info.begin(), traces_info.end(), + [](const TraceInfo& a, const TraceInfo& b) { return a.size > b.size; }); + + std::vector index_to_shard_map(trace_count); + int current_shard_index = 0; + for (int i = 0; i < shard_count; ++i) { + std::vector current_shard; + for (int j = i; j < traces_info.size(); j += shard_count) { + index_to_shard_map[traces_info[j].index] = current_shard_index; + } + ++current_shard_index; + } + + return index_to_shard_map; +} + +void WriteOutShards(absl::string_view trace_file_path, int shard_count, + const std::vector& trace_shards, + absl::string_view output_path_template, + int compression_level) { + std::vector>>> + shard_writers; + for (int i = 0; i < shard_count; ++i) { + shard_writers.push_back( + std::make_unique>>( + riegeli::Maker>( + (absl::StrCat(output_path_template, i, ".pb"))), + riegeli::RecordWriterBase::Options().set_zstd())); + } + + int index = 0; + mlgo::latency_model::MbbTrace mbb_trace; + riegeli::RecordReader trace_reader( + riegeli::Maker(std::string(trace_file_path))); + while (trace_reader.ReadRecord(mbb_trace)) { + std::unique_ptr>>& + shard_writer = shard_writers[trace_shards[index]]; + shard_writer->WriteRecord(mbb_trace); + ++index; + } + QCHECK(trace_reader.Close()) << trace_reader.status(); + for (int i = 0; i < shard_writers.size(); ++i) { + QCHECK(shard_writers[i]->Close()) << shard_writers[i]->status(); + } +} + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/sort_bb_traces_lib.h b/compiler_opt/memtrace_costmodel/sort_bb_traces_lib.h new file mode 100644 index 00000000..a9cda8ab --- /dev/null +++ b/compiler_opt/memtrace_costmodel/sort_bb_traces_lib.h @@ -0,0 +1,33 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_SORT_BB_TRACES_LIB_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_SORT_BB_TRACES_LIB_H_ + +#include +#include + +#include "absl/strings/string_view.h" +#include "compiler_opt/memtrace_costmodel/mbb_trace.proto.h" + +namespace mlgo { +namespace latency_model { + +struct TraceInfo { + size_t index; + int size; +}; + +std::vector LoadTraceInfo(absl::string_view trace_file_path); + +// The traces are expected to be passed in the order in which they are parsed +// (i.e., with the highest index trace at the back). +std::vector GetTraceIndexToShardMap(std::vector& traces_info, + size_t shard_count); + +void WriteOutShards(absl::string_view trace_file_path, int shard_count, + const std::vector& trace_shards, + absl::string_view output_path_template, + int compression_level = 1 << 16); + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_SORT_BB_TRACES_LIB_H_ diff --git a/compiler_opt/memtrace_costmodel/status_macros.h b/compiler_opt/memtrace_costmodel/status_macros.h new file mode 100644 index 00000000..a9705bdc --- /dev/null +++ b/compiler_opt/memtrace_costmodel/status_macros.h @@ -0,0 +1,21 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_STATUS_MACROS_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_STATUS_MACROS_H_ + +#include "absl/status/status.h" +#include "absl/status/statusor.h" + +#define CONCAT_IMPL(x, y) x##y +#define CONCAT(x, y) CONCAT_IMPL(x, y) + +#define RETURN_IF_ERROR(expr) \ + if (auto _status = (expr); !_status.ok()) return _status; + +#define ASSIGN_OR_RETURN_IMPL(tmp, lhs, rexpr) \ + auto tmp = (rexpr); \ + if (!tmp.ok()) return tmp.status(); \ + lhs = std::move(*tmp); + +#define ASSIGN_OR_RETURN(lhs, rexpr) \ + ASSIGN_OR_RETURN_IMPL(CONCAT(_status_or_, __LINE__), lhs, rexpr) + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_STATUS_MACROS_H_ diff --git a/compiler_opt/memtrace_costmodel/trace_segment_mca.cc b/compiler_opt/memtrace_costmodel/trace_segment_mca.cc new file mode 100644 index 00000000..7a8f74a9 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/trace_segment_mca.cc @@ -0,0 +1,177 @@ +#include "compiler_opt/memtrace_costmodel/trace_segment_mca.h" + +#include +#include +#include +#include + +#include "absl/log/check.h" +#include "absl/memory/memory.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/MC/MCInst.h" +#include "llvm/MC/MCInstrDesc.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/MCA/Context.h" +#include "llvm/MCA/CustomBehaviour.h" +#include "llvm/MCA/InstrBuilder.h" +#include "llvm/MCA/Instruction.h" +#include "llvm/MCA/Stages/Stage.h" +#include "llvm/Support/Error.h" +#include "llvm/TargetParser/Triple.h" + +constexpr const int kCallLatency = 100; + +namespace mlgo { +namespace latency_model { + +TraceSegmentMca::TraceSegmentMca(std::string_view target_triple, + std::string_view cpu_name, int batch_size) + : recycle_freed_instruction_( + [this](llvm::mca::Instruction* freed_instruction) { + recycled_mca_instructions_[&freed_instruction->getDesc()].insert( + freed_instruction); + }), + get_recycled_instruction_( + [this](const llvm::mca::InstrDesc& instruction_description) + -> llvm::mca::Instruction* { + auto recycled_instruction_iterator = + recycled_mca_instructions_.find(&instruction_description); + if (recycled_instruction_iterator != + recycled_mca_instructions_.end()) { + llvm::SmallPtrSetImpl& + recycled_instructions = recycled_instruction_iterator->second; + if (!recycled_instructions.empty()) { + llvm::mca::Instruction* recycled_instruction = + *recycled_instructions.begin(); + recycled_instructions.erase(recycled_instruction); + return recycled_instruction; + } + } + return nullptr; + }), + batch_size_(batch_size) { + std::string possible_lookup_error; + llvm::Triple triple((llvm::StringRef(target_triple))); + const llvm::Target* const llvm_target = + llvm::TargetRegistry::lookupTarget(triple, possible_lookup_error); + QCHECK(llvm_target) << possible_lookup_error; + mc_subtarget_info_ = absl::WrapUnique( + llvm_target->createMCSubtargetInfo(triple, cpu_name, "")); + QCHECK(mc_subtarget_info_); + mc_register_info_ = absl::WrapUnique(llvm_target->createMCRegInfo(triple)); + QCHECK(mc_register_info_); + mc_instruction_info_ = absl::WrapUnique(llvm_target->createMCInstrInfo()); + QCHECK(mc_instruction_info_); + mc_instruction_analysis_ = absl::WrapUnique( + llvm_target->createMCInstrAnalysis(mc_instruction_info_.get())); + QCHECK(mc_instruction_analysis_); + mca_context_ = std::make_unique(*mc_register_info_, + *mc_subtarget_info_); + QCHECK(mca_context_); + + CreateMcaPipeline(); + + mca_instrument_manager_ = std::make_unique( + *mc_subtarget_info_, *mc_instruction_info_); + QCHECK(mca_instrument_manager_); + mca_instruction_builder_ = std::make_unique( + *mc_subtarget_info_, *mc_instruction_info_, *mc_register_info_, + mc_instruction_analysis_.get(), *mca_instrument_manager_, kCallLatency); + QCHECK(mca_instruction_builder_); + mca_instruction_processor_ = + absl::WrapUnique(llvm_target->createInstrPostProcess( + *mc_subtarget_info_, *mc_instruction_info_)); + // If there is not a target specific instruction post processor available, + // create a generic one. + if (!mca_instruction_processor_) { + mca_instruction_processor_ = + absl::WrapUnique(llvm_target->createInstrPostProcess( + *mc_subtarget_info_, *mc_instruction_info_)); + } + QCHECK(mca_instruction_processor_); + + source_manager_.setOnInstFreedCallback(recycle_freed_instruction_); + mca_instruction_builder_->setInstRecycleCallback(get_recycled_instruction_); +} + +void TraceSegmentMca::CreateMcaPipeline() { + // Setting these values to zero/their defaults makes MCA use the values + // provided by the scheduling model. + llvm::mca::PipelineOptions mca_pipeline_options( + /*UOPQSize=*/0, /*DecThr=*/0, + /*DW=*/0, + /*RFS=*/0, + /*LQS=*/0, /*SQS=*/0, + /*NoAlias=*/true, + /*ShouldEnableBottleneckAnalysis=*/false); + mca_custombehavior_ = std::make_unique( + *mc_subtarget_info_, source_manager_, *mc_instruction_info_); + QCHECK(mca_custombehavior_); + mca_pipeline_ = mca_context_->createDefaultPipeline( + mca_pipeline_options, source_manager_, *mca_custombehavior_); + QCHECK(mca_pipeline_); +} + +void TraceSegmentMca::AddInstruction(const llvm::MCInst& new_instruction) { + const llvm::MCInstrDesc& instruction_description = + mc_instruction_info_->get(new_instruction.getOpcode()); + + // Skip call and return instructions. These are not modeled properly by MCA + // so omitting them will have little impact on modeling the performance + // characteristics we are interested in. Passing in return instructions also + // currently results in use after frees within MCA. + if (instruction_description.isCall() || instruction_description.isReturn()) + return; + + llvm::Expected> new_mca_instruction = + mca_instruction_builder_->createInstruction(new_instruction, /*IVec=*/{}); + + if (!new_mca_instruction) { + llvm::mca::Instruction* recycled_instruction = nullptr; + llvm::Error leftover_error = llvm::handleErrors( + new_mca_instruction.takeError(), + [&recycled_instruction]( + const llvm::mca::RecycledInstErr& recycling_error) { + recycled_instruction = recycling_error.getInst(); + }); + QCHECK(!leftover_error); + QCHECK(recycled_instruction); + llvm::consumeError(std::move(leftover_error)); + mca_instruction_processor_->postProcessInstruction(*recycled_instruction, + new_instruction); + source_manager_.addRecycledInst(recycled_instruction); + } else { + // We only need to recycle instructions if they are new (i.e., not + // recycled). Recycled instructions will already contain the relevant + // modifications. The x86 instruction post processor is stateless, so this + // is safe. + mca_instruction_processor_->postProcessInstruction( + *new_mca_instruction.get(), new_instruction); + source_manager_.addInst(std::move(new_mca_instruction.get())); + } + + ++instruction_count_; + + if (instruction_count_ % batch_size_ == 0) { + llvm::Expected cycles = mca_pipeline_->run(); + + // We expect to get an InstStreamPause error rather than a cycles + // value as we have not called endOfStream yet. Check that this is + // the case. Consume the error as we are expecting it, and we are + // not done processing. + QCHECK(!cycles); + QCHECK(cycles.errorIsA()); + llvm::consumeError(cycles.takeError()); + } +} + +double TraceSegmentMca::GetCost() { + source_manager_.endOfStream(); + llvm::Expected cycles = mca_pipeline_->run(); + QCHECK(cycles); + return cycles.get(); +} + +} // namespace latency_model +} // namespace mlgo diff --git a/compiler_opt/memtrace_costmodel/trace_segment_mca.h b/compiler_opt/memtrace_costmodel/trace_segment_mca.h new file mode 100644 index 00000000..bb828377 --- /dev/null +++ b/compiler_opt/memtrace_costmodel/trace_segment_mca.h @@ -0,0 +1,81 @@ +#ifndef COMPILER_OPT_MEMTRACE_COSTMODEL_TRACE_SEGMENT_MCA_H_ +#define COMPILER_OPT_MEMTRACE_COSTMODEL_TRACE_SEGMENT_MCA_H_ + +#include +#include +#include +#include +#include + +#include "compiler_opt/memtrace_costmodel/costmodel.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/MC/MCInst.h" +#include "llvm/MC/MCInstrAnalysis.h" +#include "llvm/MC/MCInstrInfo.h" +#include "llvm/MC/MCRegisterInfo.h" +#include "llvm/MC/MCSubtargetInfo.h" +#include "llvm/MCA/Context.h" +#include "llvm/MCA/CustomBehaviour.h" +#include "llvm/MCA/IncrementalSourceMgr.h" +#include "llvm/MCA/InstrBuilder.h" +#include "llvm/MCA/Instruction.h" +#include "llvm/MCA/Pipeline.h" + +namespace mlgo { +namespace latency_model { + +// This class is an abstraction around LLVM MCA (machine code analyzer), aimed +// specifically at supporting the memtrace case where instructions are +// processed one by one. This class automatically handles setting up the +// relevant LLVM State and keeping the pipeline state so users only have to +// add instructions one by one and then query for the cost after a trace +// segment is complete (i.e., upon entrypoint exit). +class TraceSegmentMca : public CostModel { + public: + explicit TraceSegmentMca(std::string_view target_triple, + std::string_view cpu_name, int batch_size = 1); + + // This function should be called for each instruction encountered during + // memtrace processing within a context where we are interested in modeling + // the cost (i.e., running under an entrypoint). + void AddInstruction(const llvm::MCInst& new_instruction) override; + + // This function should only be called after all relevant instructions have + // been added and the user is ready to get a final cost for a specific + // instruction stream (i.e., upon entrypoint exit). + double GetCost() override; + + private: + void CreateMcaPipeline(); + + std::unique_ptr mca_context_; + std::unique_ptr mc_subtarget_info_; + std::unique_ptr mc_register_info_; + std::unique_ptr mc_instruction_info_; + std::unique_ptr mc_instruction_analysis_; + + llvm::mca::IncrementalSourceMgr source_manager_; + + std::unique_ptr mca_custombehavior_; + std::unique_ptr mca_pipeline_; + std::unique_ptr mca_instrument_manager_; + std::unique_ptr mca_instruction_builder_; + std::unique_ptr mca_instruction_processor_; + + size_t instruction_count_ = 0; + + std::function recycle_freed_instruction_; + std::function + get_recycled_instruction_; + + std::unordered_map> + recycled_mca_instructions_; + + const int batch_size_; +}; + +} // namespace latency_model +} // namespace mlgo + +#endif // COMPILER_OPT_MEMTRACE_COSTMODEL_TRACE_SEGMENT_MCA_H_