Macro compression for LLVM

MacroCompression.cpp

/*

  Macro compression transform for LLVM
  ====================================

  Description
  -----------
  This pass identifies duplicated code sequences in the whole module and replaces
  them with equivalent "macros". To keep things simple, macros are defined as
  functions made up of two instructions.
  The process is made up of two steps: in the first one, a list of instruction pairs
  is generated; in the second one we iteratively select the most frequent
  instruction pair, create a macro containing the pair and replace all occurrences
  of the pair with a call to the macro. The approach will likely result in very deep
  "macro trees"

  Example
  -------
  The following two instructions:

  %Ar = op_A %A1, %A2, ..., %An
  %Br = op_B %B1, %B2, ..., %Bn

  will be replaced by:

  {%Ar, %Br} = call fastcc @macro_1( %A1, %A2, ..., %An, %B1, %B2, ..., %Bn )

  and the two original instructions (op_A and op_B) will be placed in a function 
  called @macro_1

*/

typedef std::pair< Instruction*, Instruction* > InstructionPair;
typedef std::list< Instruction* > InstructionList;
typedef std::list< InstructionPair > InstructionPairList;
typedef std::list< Macro* > MacroList;

class Macro {
  InstructionPairList l;
  Function *F;
  
  Macro(Instruction *i, Instruction *i2) {
    add(i, i2);
    F = null;
  }
  
  bool compare(Instruction *i, Instruction *i2) {
    assert(refs() > 0 && "comparing to what?");
    InstructionPair p = l.first();
    return p.first->isSameOperationAs(i) && p.end->isSameOperationAs(i2);
  }
  
  bool remove(Instruction *i, Instruction *i2) {
    for (InstructionPairList::iterator j = l.first(), je = l.end(); j != je; )
      if (j->first == i || j->second == i || j->first == i2 || j->second == i2)
        j = l.erase(j);
      else
        j++;
  }
  
  void add(Instruction *i, Instruction *i2) {
    l.push_back(InstructionPair(i, i2));
  }
  
  unsigned refs() {
    return l.size();
  }
  
  void replaceInstructions(Instruction *i, Instruction *i2) {
    if (F == null)
      createFunctionMacro();
    moveIstructions(i, i2);
  }
  
  void moveInstructions(Instruction *i, Instruction *i2) {
    std::set< Value* > in;
    std::list< Value* > args;
    for (User::op_iterator j = i->op_begin(), je = i->op_end(); j != je; ++j) {
      if (!in.contains(j->get())) {
        in.insert(j->get());
        args.push_back(j->get());
      }
    }
    for (User::op_iterator j = i2->op_begin(), je = i2->op_end(); j != je; ++j) {
      if (!in.contains(j->get())) {
        in.insert(j->get());
        args.push_back(j->get());
      }
    }
    CallInst *call = CallInst::Create(F, args, "", i);
    i->replaceInstWithInst(ExtractValueInst::Create(call, 0));
    i2->replaceInstWithInst(ExtractValueInst::Create(call, 1));
  }
  
  void createFunctionMacro() {
    InstructionPair p = l.first();
    std::map< Value*, Value* > in;
    // get (distinct) input values to the two instructions
    for (User::op_iterator i = p.first->op_begin(), e = p.first->op_end(); i != e; ++i)
      in[i->get()] = NULL;
    for (User::op_iterator i = p.second->op_begin(), e = p.second->op_end(); i != e; ++i)
      in[i->get()] = NULL;
    std::list< Type* > inTy;
    for (std::map< Value*, Value* >::iterator i = in.begin(), ie = in.end(); i != ie; ++i)
      inTy.push_back( i->first->getType() );
    // create the return type
    Type *retTy = StructType::get( p.first, p.second );
    // create the function
    FunctionType *funTy = FunctionType::get( retTy, inTy, false );
    Function *fun = Function::Create( funTy, LinkageTypes::PrivateLinkage, "macro", module );
    fun->setCallingConv(CallingConv::FastCC);
    // map the operands of the original instructions to the arguments of the extracted macro
    Function::arg_iterator j = fun->arg_begin();
    for (std::map< Value*, Value* >::iterator i = in.begin(), ie = in.end(); i != ie; ++i, ++j)
      in[i->first] = *j;
    // copy the original instructions into the macro and map their operands to the macro arguments
    BasicBlock *bb = BasicBlock::Create(fun);
    Instruction *i = p.first->clone();
    i->insertBefore(bb->last());
    Instruction *i2 = p.second->clone();
    i2->insertBefore(bb->last());
    for (User::op_iterator j = i->op_begin(), je = i->op_end(); j != je; ++j)
      j->set(in[j->get()]);
    for (User::op_iterator j = i2->op_begin(), je = i2->op_end(); j != je; ++j)
      j->set(in[j->get()]);
    InsertValueInst *t1 = InsertValueInst::Create(Undef::get(retTy), i, 0, "", bb);
    InsertValueInst *t2 = InsertValueInst::Create(t1, i2, 0, "", bb);
    ReturnInst *ret = ReturnInst::Create(t2, bb);
  }
}

class MacroRegistry {
  MacroList l;
  
  void add(Instruction *I, Instruction *I2) {
    for (MacroList::iterator m = l.begin(), me = l.end(); m != me; m++) {
      if (m->compare(I, I2)) {
        m->add(I, I2);
        return;
      }
    }
    l.push_back(new Macro(I, I2));
  }
  
  void remove(Instruction *I, Instruction *I2) {
    for (MacroList::iterator m = l.begin(), me = l.end(); m != me; m++) {
      m->remove(I, I2);
      if (m->refs() > 0) {
        m++;
      } else {
        delete *m;
        m = l.erase(m);
      }
    }
  }
  
  Macro* get() {
    Macro *M = null;
    for (MacroList::iterator m = l.begin(), me = l.end(); m != me; m++)
      if (M == null || M->refs() < (*m)->refs())
        M = *m;
    return M;
  }
  
}

static bool interesting(Instruction *i) {
  return !i->isTerminator() && !is_a<PHINode*>(i);
}

InstructionList findCompanionInstruction(Instruction *i) {
  InstructionList l;
  if (interesting(i)) {
    for (BasicBlock::iterator j(i)++, je = i->getParent()->end(); j != je; j++)
      l.push_back(*j);
    for (InstructionList::iterator j = l.begin(); j != l.end(); ) {
      for (Value::use_iterator k = (*j)->use_begin(), ke = (*j)->use_end(); k != ke; k++)
        if (*j != *k)
          l.remove(*k);
      if (!interesting(*j))
        j = l.erase(j);
      else
        j++;
    }
  }
  return l;
}

bool scan(Module *M) {
  MacroRegistry R;
  for (inst_iterator I = inst_begin(M), Ie = inst_end(M); I != Ie; I++) {
    InstructionList Ic = findCompanionInstruction(I);
    for (InstructionList::iterator i = Ic.begin(), ie = Ic.end(); i != ie; i++)
      R.add(*I, *Ic);
  }
  while (Macro *m = R.get() && m->refs() > 1) {
    m->createFunctionMacro();
    for (InstructionPairList::iterator i = m.first(), ie = m.end(); i != ie; i++) {
      replaceInstructionsWithFunction(i->first, i->second, F);
      R.remove(i->first, i->second);
    }
  }
}