LLVM 指令统计

#include "llvm/ADT/STLExtras.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Type.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/IR/IRBuilder.h"

#include <memory>
#include "iostream"
#include <sstream>
#include <iomanip>

using namespace llvm;

class InstructionCounter {
public:
    /// String constants for instruction count names.
    static const char* TOTAL_INSTS;
    static const char* TOTAL_BLOCKS;
    static const char* TOTAL_FUNCTIONS;
    static const char* TERMINATOR_INSTS;
    static const char* BINARY_INSTS;
    static const char* MEMORY_INSTS;
    static const char* CAST_INSTS;
    static const char* OTHER_INSTS;

    InstructionCounter();

    /// Visits each Function in Module M.
    void visit(const llvm::Module& M);

    /// Increments "Total Functions" InstructionCounter and visits each BasicBlock in F.
    void visit(const llvm::Function &F);

    /// Increments "Total Blocks" InstructionCounter and visits each Instruction in BB.
    void visit(const llvm::BasicBlock &BB);

    /// Increments "Total Instructions" and whichever instruction count I is delegated to in
    /// this functions switch statement.
    void visit(const llvm::Instruction &I);

    /// Prints a single counter described by name and count
    void PrintCounter(const char* name, int count, int max_count_len,
        std::stringstream* stream) const;

    /// Prints all counters
    std::string PrintCounters() const;

    /// Return count of counter described by name
    int GetCount(const char* name);

    /// Set all counts to 0.
    void ResetCount();

private:
    typedef std::map<std::string, int> CounterMap;

    /// Allows for easy visitation of iterators.
    template<class Iterator>
    void visit(Iterator start, Iterator end) {
        while (start != end) {
            visit(*start++);
        }
    }

    /// Increment InstructionCount with name_ equal to name argument.
    void IncrementCount(const char* name);

    /// This maps instruction names to their respective count.
    CounterMap counters_;
};

const char* InstructionCounter::TOTAL_INSTS = "Total Instructions";
const char* InstructionCounter::TOTAL_BLOCKS = "Total Blocks";
const char* InstructionCounter::TOTAL_FUNCTIONS = "Total Functions";
const char* InstructionCounter::TERMINATOR_INSTS = "Terminator Instructions ";
const char* InstructionCounter::BINARY_INSTS = "Binary Instructions ";
const char* InstructionCounter::MEMORY_INSTS = "Memory Instructions ";
const char* InstructionCounter::CAST_INSTS = "Cast Instructions ";
const char* InstructionCounter::OTHER_INSTS = "Other Instructions ";

InstructionCounter::InstructionCounter() {
    counters_.insert(std::make_pair(TOTAL_INSTS, 0));
    counters_.insert(std::make_pair(TOTAL_BLOCKS, 0));
    counters_.insert(std::make_pair(TOTAL_FUNCTIONS, 0));
    counters_.insert(std::make_pair(TERMINATOR_INSTS, 0));
    counters_.insert(std::make_pair(BINARY_INSTS, 0));
    counters_.insert(std::make_pair(MEMORY_INSTS, 0));
    counters_.insert(std::make_pair(CAST_INSTS, 0));
    counters_.insert(std::make_pair(OTHER_INSTS, 0));
}

void InstructionCounter::visit(const llvm::Module& M) {
    visit(M.begin(), M.end());
}

void InstructionCounter::visit(const llvm::Function& F) {
    IncrementCount(TOTAL_FUNCTIONS);
    visit(F.begin(), F.end());
}

void InstructionCounter::visit(const llvm::BasicBlock& BB) {
    IncrementCount(TOTAL_BLOCKS);
    visit(BB.begin(), BB.end());
}

int InstructionCounter::GetCount(const char* name) {
    CounterMap::const_iterator counter = counters_.find(name);
    if (counter == counters_.end())
        return 0;
    return counter->second;
}

void InstructionCounter::visit(const llvm::Instruction& I) {
    IncrementCount(TOTAL_INSTS);
    switch (I.getOpcode()) {
    case llvm::Instruction::Ret:
    case llvm::Instruction::Br:
    case llvm::Instruction::Switch:
    case llvm::Instruction::IndirectBr:
    case llvm::Instruction::Invoke:
    case llvm::Instruction::Resume:
    case llvm::Instruction::Unreachable:
        IncrementCount(TERMINATOR_INSTS);
        break;
    case llvm::Instruction::Add:
    case llvm::Instruction::FAdd:
    case llvm::Instruction::Sub:
    case llvm::Instruction::FSub:
    case llvm::Instruction::Mul:
    case llvm::Instruction::FMul:
    case llvm::Instruction::UDiv:
    case llvm::Instruction::SDiv:
    case llvm::Instruction::FDiv:
    case llvm::Instruction::URem:
    case llvm::Instruction::SRem:
    case llvm::Instruction::FRem:
    case llvm::Instruction::Shl:
    case llvm::Instruction::LShr:
    case llvm::Instruction::AShr:
    case llvm::Instruction::And:
    case llvm::Instruction::Or:
    case llvm::Instruction::Xor:
        IncrementCount(BINARY_INSTS);
        break;
    case llvm::Instruction::Alloca:
    case llvm::Instruction::Load:
    case llvm::Instruction::Store:
    case llvm::Instruction::Fence:
    case llvm::Instruction::AtomicCmpXchg:
    case llvm::Instruction::AtomicRMW:
        IncrementCount(MEMORY_INSTS);
        break;
    case llvm::Instruction::Trunc:
    case llvm::Instruction::ZExt:
    case llvm::Instruction::SExt:
    case llvm::Instruction::FPToUI:
    case llvm::Instruction::FPToSI:
    case llvm::Instruction::UIToFP:
    case llvm::Instruction::SIToFP:
    case llvm::Instruction::FPTrunc:
    case llvm::Instruction::FPExt:
    case llvm::Instruction::PtrToInt:
    case llvm::Instruction::IntToPtr:
    case llvm::Instruction::BitCast:
        IncrementCount(CAST_INSTS);
        break;
    case llvm::Instruction::ICmp:
    case llvm::Instruction::FCmp:
    case llvm::Instruction::PHI:
    case llvm::Instruction::Call:
    case llvm::Instruction::Select:
    case llvm::Instruction::UserOp1:
    case llvm::Instruction::UserOp2:
    case llvm::Instruction::VAArg:
    case llvm::Instruction::ExtractElement:
    case llvm::Instruction::InsertElement:
    case llvm::Instruction::ShuffleVector:
    case llvm::Instruction::ExtractValue:
    case llvm::Instruction::InsertValue:
    case llvm::Instruction::LandingPad:
    case llvm::Instruction::GetElementPtr:
        IncrementCount(OTHER_INSTS);
        break;
    default:
        break;
    }
}

void InstructionCounter::ResetCount() {
    for (CounterMap::iterator iter = counters_.begin(); iter != counters_.end(); iter++) {
        iter->second = 0;
    }
}

void InstructionCounter::PrintCounter(const char* name, int count, int max_count_len,
    std::stringstream* stream) const {
    if (count > 0) {
        *stream << std::left << std::setw(max_count_len + 1) << std::setfill(' ') << count << std::setw(35)
            << std::setfill(' ') << name << std::left << std::setw(10 - max_count_len) << std::setfill(' ')
            << "Number of " << name << std::endl;
    }
}

std::string InstructionCounter::PrintCounters() const {
    int max_count_len = 0;
    std::stringstream count_stream;
    for (CounterMap::const_iterator counter = counters_.begin(); counter != counters_.end(); counter++) {
        count_stream << counter->second;
        max_count_len =
            std::max(max_count_len, static_cast<int>(strlen(count_stream.str().c_str())));
        count_stream.str("");
    }
    std::stringstream stream;
    // Print header
    stream << "\n===" << std::string(73, '-') << "===\n\n"
        << "                          ... Instruction Counts ...\n\n"
        << "===" << std::string(73, '-') << "===\n\n";

    for (CounterMap::const_iterator counter = counters_.begin(); counter != counters_.end(); counter++) {
        // Conditional is only used in order to print the top level counters
        // separate from the other counters.
        if (strcmp(counter->first.c_str(), TOTAL_BLOCKS) == 0) {
            stream << "\n===" << std::string(73, '-') << "===\n"
                << "                                ... Totals ...\n"
                << "===" << std::string(73, '-') << "===\n\n";
        }
        PrintCounter(counter->first.c_str(), counter->second, max_count_len, &stream);
    }
    stream << '\n';
    return stream.str();
}

void InstructionCounter::IncrementCount(const char* name) {
    CounterMap::iterator iter = counters_.find(name);
    if (iter == counters_.end())
        return;
    iter->second++;
}

int main() {

    LLVMContext Context;

    // Create some module to put our function into it.
    std::unique_ptr<Module> Owner = make_unique<Module>("test", Context);
    Module *mod = Owner.get();

    // Create the add1 function entry and insert this entry into module M.  The
    // function will have a return type of "int" and take an argument of "int".
    // The '0' terminates the list of argument types.
    Function *Add1F =
        cast<Function>(mod->getOrInsertFunction("add1", Type::getInt32Ty(Context),
        Type::getInt32Ty(Context),
        nullptr));

    // Add a basic block to the function. As before, it automatically inserts
    // because of the last argument.
    BasicBlock *BB = BasicBlock::Create(Context, "EntryBlock", Add1F);

    // Create a basic block builder with default parameters.  The builder will
    // automatically append instructions to the basic block `BB'.
    IRBuilder<> builder(BB);

    // Get pointers to the constant `1'.
    Value *One = builder.getInt32(1);

    // Get pointers to the integer argument of the add1 function...
    assert(Add1F->arg_begin() != Add1F->arg_end()); // Make sure there's an arg
    Argument *ArgX = &*Add1F->arg_begin();          // Get the arg
    ArgX->setName("AnArg");            // Give it a nice symbolic name for fun.

    // Create the add instruction, inserting it into the end of BB.
    Value *Add = builder.CreateAdd(One, ArgX);

    // Create the return instruction and add it to the basic block
    builder.CreateRet(Add);

    InstructionCounter instruction_counter;

    instruction_counter.visit(*mod);
    std::string result = instruction_counter.PrintCounters();
    std::cout << result << std::endl;

    return 0;
}