From 09d71cada28de22b5b4e0c817a62b47b0bcc4b24 Mon Sep 17 00:00:00 2001 From: Peter Goodman Date: Wed, 3 Feb 2016 12:24:34 -0500 Subject: [PATCH] Improved flags analysis. Move the undefined optimization out into the plugin. Still have an issue where dead-store elimination is not happening :-( --- mcsema/Arch/X86/RegisterAnalysis.cpp | 58 +++++++++++++------- mcsema/Arch/X86/RegisterAnalysis.h | 3 ++ mcsema/BC/IntrinsicTable.cpp | 2 +- mcsema/BC/Translator.cpp | 39 +------------- mcsema/BC/Translator.h | 3 -- mcsema/BC/Util.cpp | 3 +- mcsema/BC/Util.h | 3 +- mcsema/Optimize.cpp | 80 +++++++++++++++++++++++----- scripts/optimize_bitcode.sh | 2 +- tests/X86/Run.cpp | 20 +++++++ 10 files changed, 134 insertions(+), 79 deletions(-) diff --git a/mcsema/Arch/X86/RegisterAnalysis.cpp b/mcsema/Arch/X86/RegisterAnalysis.cpp index d8144bad..2d0588ad 100644 --- a/mcsema/Arch/X86/RegisterAnalysis.cpp +++ b/mcsema/Arch/X86/RegisterAnalysis.cpp @@ -84,6 +84,7 @@ void RegisterAnalysis::AddFunction(const cfg::Function &func) { void RegisterAnalysis::AddBlock(const cfg::Block &block) { auto bb = new BasicBlockRegs; + bb->live_anywhere.flat = 0U; bb->live_exit.flat = 0U; bb->live_entry.flat = 0U; bb->keep_alive.flat = 0U; @@ -96,6 +97,13 @@ void RegisterAnalysis::AddBlock(const cfg::Block &block) { const auto &instr = *it; const auto xedd = DecodeInstruction(instr, arch_name); + if (const auto rflags = xed_decoded_inst_get_rflags_info(&xedd)) { + bb->live_anywhere.flat |= rflags->read.flat; + bb->keep_alive.flat |= rflags->written.flat; + bb->keep_alive.flat |= rflags->undefined.flat; + bb->keep_alive.flat &= ~(rflags->read.flat); + } + if (it == it_begin) { const auto next_pc = instr.address() + instr.size(); bb->flow = FindSuccessors( @@ -124,26 +132,11 @@ void RegisterAnalysis::AddBlock(const cfg::Block &block) { ret_blocks[block.address()] = next_pc; } } - - if (const auto rflags = xed_decoded_inst_get_rflags_info(&xedd)) { - bb->keep_alive.flat |= rflags->written.flat; - bb->keep_alive.flat |= rflags->undefined.flat; - bb->keep_alive.flat &= ~(rflags->read.flat); - } } bb->keep_alive.flat = ~bb->keep_alive.flat; bb->live_entry.flat = bb->live_exit.flat & bb->keep_alive.flat; - if (FLAGS_aggressive_dataflow_analysis) { - if (kFlowSysCall == bb->flow || - kFlowIndirectCall == bb->flow || - kFlowUnknown == bb->flow) { - bb->live_entry.flat = 0U; - bb->keep_alive.flat = 0U; - } - } - blocks[block.address()] = bb; } @@ -249,12 +242,29 @@ void RegisterAnalysis::InitWorkList(AnalysisWorkList &work_list) { } } + // Do some ahead-of-time work to kill any flags that are never read anywhere + // in the program. + xed_flag_set_t live_anywhere; + live_anywhere.flat = 0U; + for (auto b : blocks) { + if (auto block = b.second) { + live_anywhere.flat |= block->live_anywhere.flat; + } + } + // Initialize the worklist for dead flags and register analysis across the // control-flow graph. for (auto b : blocks) { auto block_pc = b.first; - auto block = b.second; - work_list.insert({block->predecessors.size(), block_pc}); + if (auto block = b.second) { + + // Try to kill globally unused flags. + block->live_entry.flat &= ~live_anywhere.flat; + block->live_exit.flat &= ~live_anywhere.flat; + block->keep_alive.flat &= ~live_anywhere.flat; + + work_list.insert({block->predecessors.size(), block_pc}); + } } } @@ -267,8 +277,16 @@ void RegisterAnalysis::AnalyzeBlock(AnalysisWorkItem item, // kills all the flags. auto incoming_live = FLAGS_aggressive_dataflow_analysis ? 0U : ~0U; - for (auto succ_pc : bb->successors) { - incoming_live |= LiveFlags(succ_pc); + // Try to collect flags across system calls. + if (kFlowSysCall == bb->flow) { + if (auto succ_pc = ret_blocks[item.pc]) { + incoming_live |= LiveFlags(succ_pc); + } + + } else { + for (auto succ_pc : bb->successors) { + incoming_live |= LiveFlags(succ_pc); + } } auto live_entry = incoming_live & bb->keep_alive.flat; @@ -290,7 +308,7 @@ uint32_t RegisterAnalysis::LiveFlags(uint64_t pc) { // This is basically an error case anyway, i.e. we have a direct flow that // we can't resolve. This will be warned about during the lifting phase. } else { - return ~0U; + return FLAGS_aggressive_dataflow_analysis ? 0U : ~0U; } } diff --git a/mcsema/Arch/X86/RegisterAnalysis.h b/mcsema/Arch/X86/RegisterAnalysis.h index 2e75d511..184a453d 100644 --- a/mcsema/Arch/X86/RegisterAnalysis.h +++ b/mcsema/Arch/X86/RegisterAnalysis.h @@ -26,6 +26,9 @@ struct BasicBlockRegs { uint64_t address; FlowType flow; + // Flags that are live anywhere in this block. + xed_flag_set_t live_anywhere; + // Flags that are live on entry, after factoring in those flags that are // live from the successors. xed_flag_set_t live_entry; diff --git a/mcsema/BC/IntrinsicTable.cpp b/mcsema/BC/IntrinsicTable.cpp index 9516fc83..722ae15b 100644 --- a/mcsema/BC/IntrinsicTable.cpp +++ b/mcsema/BC/IntrinsicTable.cpp @@ -97,7 +97,7 @@ IntrinsicTable::IntrinsicTable(const llvm::Module *M) atomic_end(FindIntrinsic(M, "__mcsema_atomic_end")), // Optimization guides. - defer_inlining(FindIntrinsic(M, "__mcsema_defer_inlining")), + defer_inlining(FindPureIntrinsic(M, "__mcsema_defer_inlining")), // Optimization enablers. undefined_bool(FindPureIntrinsic(M, "__mcsema_undefined_bool")), diff --git a/mcsema/BC/Translator.cpp b/mcsema/BC/Translator.cpp index 340e364c..689f79a3 100644 --- a/mcsema/BC/Translator.cpp +++ b/mcsema/BC/Translator.cpp @@ -393,7 +393,7 @@ void Translator::LiftCFG(const cfg::Module *cfg) { LinkFunctionsToBlocks(cfg); AnalyzeCFG(cfg); LiftBlocks(cfg); - OptimizeModule(); + //OptimizeModule(); } // Run an architecture-specific data-flow analysis on the module. @@ -431,41 +431,4 @@ llvm::Function *Translator::GetLiftedBlockForPC(uintptr_t pc) const { return F; } -namespace { - -// Replace all uses of a specific intrinsic with an undefined value. -static void ReplaceIntrinsic(llvm::Function *F, unsigned N) { - if (!F) return; - - std::vector Cs; - for (auto U : F->users()) { - if (auto C = llvm::dyn_cast(U)) { - Cs.push_back(C); - } - } - - auto Undef = llvm::UndefValue::get(llvm::Type::getIntNTy(F->getContext(), N)); - for (auto C : Cs) { - C->replaceAllUsesWith(Undef); - C->removeFromParent(); - delete C; - } -} - -} // namespace - -// Remove calls to the undefined intrinsics. The goal here is to improve dead -// store elimination by peppering the instruction semantics with assignments -// to the return values of special `__mcsema_undefined_*` intrinsics. It's hard -// to reliably produce an `undef` LLVM value from C/C++, so we use our trick -// of declaring (but never defining) a special "intrinsic" and then we replace -// all such uses with `undef` values. -void Translator::OptimizeModule(void) { - ReplaceIntrinsic(intrinsics->undefined_bool, 1); - ReplaceIntrinsic(intrinsics->undefined_8, 8); - ReplaceIntrinsic(intrinsics->undefined_16, 16); - ReplaceIntrinsic(intrinsics->undefined_32, 32); - ReplaceIntrinsic(intrinsics->undefined_64, 64); -} - } // namespace mcsema diff --git a/mcsema/BC/Translator.h b/mcsema/BC/Translator.h index 510f95d9..e3ea662e 100644 --- a/mcsema/BC/Translator.h +++ b/mcsema/BC/Translator.h @@ -81,9 +81,6 @@ class Translator { // Run an architecture-specific data-flow analysis on the module. void AnalyzeCFG(const cfg::Module *cfg); - // Optimize the lifted module. - void OptimizeModule(void); - // Architecture of the code contained within the CFG being lifted. const Arch * const arch; diff --git a/mcsema/BC/Util.cpp b/mcsema/BC/Util.cpp index 2cd241c7..380d6b25 100644 --- a/mcsema/BC/Util.cpp +++ b/mcsema/BC/Util.cpp @@ -87,7 +87,8 @@ void AddTerminatingTailCall(llvm::BasicBlock *B, llvm::Function *To) { // Find a local variable defined in the entry block of the function. We use // this to find register variables. -llvm::Value *FindVarInFunction(llvm::Function *F, std::string name) { +llvm::Value *FindVarInFunction(llvm::Function *F, std::string name, + bool allow_failure) { for (auto &I : F->getEntryBlock()) { if (I.getName() == name) { return &I; diff --git a/mcsema/BC/Util.h b/mcsema/BC/Util.h index 118c154e..7cb6ef5b 100644 --- a/mcsema/BC/Util.h +++ b/mcsema/BC/Util.h @@ -33,7 +33,8 @@ void AddTerminatingTailCall(llvm::BasicBlock *From, llvm::Function *To); // Find a local variable defined in the entry block of the function. We use // this to find register variables. -llvm::Value *FindVarInFunction(llvm::Function *F, std::string name); +llvm::Value *FindVarInFunction(llvm::Function *F, std::string name, + bool allow_failure=false); // Find the machine state pointer. The machine state pointer is, by convention, // passed as the first argument to every lifted function. diff --git a/mcsema/Optimize.cpp b/mcsema/Optimize.cpp index 37460407..aa8d72a9 100644 --- a/mcsema/Optimize.cpp +++ b/mcsema/Optimize.cpp @@ -1,9 +1,10 @@ /* Copyright 2015 Peter Goodman (peter@trailofbits.com), all rights reserved. */ -#define DEBUG_TYPE "McSema2Optimizer" +#define DEBUG_TYPE "IntrinsicOptimizer" #include #include +#include #include #include @@ -12,16 +13,65 @@ #include namespace mcsema { +namespace { + +// Looks for a function by name. If we can't find it, try to find an underscore +// prefixed version, just in case this is Mac or Windows. +static llvm::Function *GetFunction(llvm::Module &M, const char *name) { + if (auto F = M.getFunction(name)) { + return F; + } else { + std::stringstream ss; + ss << "_" << name; + return M.getFunction(ss.str()); + } +} + +// Replace all uses of a specific intrinsic with an undefined value. +static void ReplaceIntrinsic(llvm::Module &M, const char *name, unsigned N) { + if (auto F = GetFunction(M, name)) { + std::vector Cs; + for (auto U : F->users()) { + if (auto C = llvm::dyn_cast(U)) { + Cs.push_back(C); + } + } + + auto Undef = llvm::UndefValue::get( + llvm::Type::getIntNTy(F->getContext(), N)); + for (auto C : Cs) { + C->replaceAllUsesWith(Undef); + C->removeFromParent(); + delete C; + } + } +} + +// Remove calls to the undefined intrinsics. The goal here is to improve dead +// store elimination by peppering the instruction semantics with assignments +// to the return values of special `__mcsema_undefined_*` intrinsics. It's hard +// to reliably produce an `undef` LLVM value from C/C++, so we use our trick +// of declaring (but never defining) a special "intrinsic" and then we replace +// all such uses with `undef` values. +void RemoveUndefinedIntrinsics(llvm::Module &M) { + ReplaceIntrinsic(M, "__mcsema_undefined_bool", 1); + ReplaceIntrinsic(M, "__mcsema_undefined_8", 8); + ReplaceIntrinsic(M, "__mcsema_undefined_16", 16); + ReplaceIntrinsic(M, "__mcsema_undefined_32", 32); + ReplaceIntrinsic(M, "__mcsema_undefined_64", 64); +} + +} // namespace // Implements the deferred inlining optimization. McSema2 uses a special // `__mcsema_defer_inlining` intrinsic to mark functions as needing to be // "late" inlined. The idea is that we want some functions to be optimized // away (flag computation functions), but the ones that stick around should // then be inlined into their callers for further optimization. -class DeferredInlineOptimizer : public llvm::ModulePass { +class IntrinsicOptimizer : public llvm::ModulePass { public: - DeferredInlineOptimizer(void); - ~DeferredInlineOptimizer(void); + IntrinsicOptimizer(void); + ~IntrinsicOptimizer(void); virtual const char *getPassName(void) const override; virtual bool runOnModule(llvm::Module &M) override; @@ -31,17 +81,19 @@ class DeferredInlineOptimizer : public llvm::ModulePass { private: }; -DeferredInlineOptimizer::DeferredInlineOptimizer(void) +IntrinsicOptimizer::IntrinsicOptimizer(void) : llvm::ModulePass(ID) {} -DeferredInlineOptimizer::~DeferredInlineOptimizer(void) {} +IntrinsicOptimizer::~IntrinsicOptimizer(void) {} -const char *DeferredInlineOptimizer::getPassName(void) const { - return "DeferredInlineOptimizer"; +const char *IntrinsicOptimizer::getPassName(void) const { + return "IntrinsicOptimizer"; } -bool DeferredInlineOptimizer::runOnModule(llvm::Module &M) { - auto F = M.getFunction("__mcsema_defer_inlining"); +bool IntrinsicOptimizer::runOnModule(llvm::Module &M) { + RemoveUndefinedIntrinsics(M); + + auto F = GetFunction(M, "__mcsema_defer_inlining"); if (!F) { return false; } @@ -89,11 +141,11 @@ bool DeferredInlineOptimizer::runOnModule(llvm::Module &M) { return changed; } -char DeferredInlineOptimizer::ID = 0; +char IntrinsicOptimizer::ID = 0; -static llvm::RegisterPass X( - "deferred_inliner", - "Optimizes `__mcsema_defer_inlining` intrinsics.", +static llvm::RegisterPass X( + "intrinsic_optimizer", + "Removes `__mcsema_defer_inlining` and `__mcsema_undefined_*` intrinsics.", false, // Only looks at CFG. false); // Analysis Pass. diff --git a/scripts/optimize_bitcode.sh b/scripts/optimize_bitcode.sh index 6763996f..95d450bb 100755 --- a/scripts/optimize_bitcode.sh +++ b/scripts/optimize_bitcode.sh @@ -30,7 +30,7 @@ $DIR/third_party/bin/opt -O3 -o=$BIN.opt0.bc $1 || { } $DIR/third_party/bin/opt \ - -load $DIR/build/libOptimize.$DYLIB_SUFFIX -deferred_inliner \ + -load $DIR/build/libOptimize.$DYLIB_SUFFIX -intrinsic_optimizer \ -o=$BIN.opt1.bc $BIN.opt0.bc || { printf "${RED}Could not optimize $BIN.opt0.bc${RESET}\n" > /dev/stderr exit 1 diff --git a/tests/X86/Run.cpp b/tests/X86/Run.cpp index 2ef2669a..949820ea 100644 --- a/tests/X86/Run.cpp +++ b/tests/X86/Run.cpp @@ -184,6 +184,26 @@ void __mcsema_interrupt_return(State &) { __builtin_unreachable(); } +bool __mcsema_undefined_bool(void) { + return false; +} + +uint8_t __mcsema_undefined_8(void) { + return 0; +} + +uint16_t __mcsema_undefined_16(void) { + return 0; +} + +uint32_t __mcsema_undefined_32(void) { + return 0; +} + +uint64_t __mcsema_undefined_64(void) { + return 0; +} + } // extern C // The `State` structure maintains two versions of the `XMM` registers. One