lib/Transforms/Utils/CodeExtractor.cpp

   1 //===- CodeExtractor.cpp - Pull code region into a new function -----------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file was developed by the LLVM research group and is distributed under
   6 // the University of Illinois Open Source License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This file implements the interface to tear out a code region, such as an
  11 // individual loop or a parallel section, into a new function, replacing it with
  12 // a call to the new function.
  13 //
  14 //===----------------------------------------------------------------------===//
  15
  16 #include "llvm/Transforms/Utils/FunctionUtils.h"
  17 #include "llvm/Constants.h"
  18 #include "llvm/DerivedTypes.h"
  19 #include "llvm/Instructions.h"
  20 #include "llvm/Module.h"
  21 #include "llvm/Pass.h"
  22 #include "llvm/Analysis/LoopInfo.h"
  23 #include "llvm/Analysis/Verifier.h"
  24 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
  25 #include "Support/Debug.h"
  26 #include "Support/StringExtras.h"
  27 #include <algorithm>
  28 #include <set>
  29 using namespace llvm;
  30
  31 namespace {
  32
  33   /// getFunctionArg - Return a pointer to F's ARGNOth argument.
  34   ///
  35   Argument *getFunctionArg(Function *F, unsigned argno) {
  36     Function::aiterator I = F->abegin();
  37     std::advance(I, argno);
  38     return I;
  39   }
  40
  41   struct CodeExtractor {
  42     typedef std::vector<Value*> Values;
  43     typedef std::vector<std::pair<unsigned, unsigned> > PhiValChangesTy;
  44     typedef std::map<PHINode*, PhiValChangesTy> PhiVal2ArgTy;
  45     PhiVal2ArgTy PhiVal2Arg;
  46     std::set<BasicBlock*> BlocksToExtract;
  47   public:
  48     Function *ExtractCodeRegion(const std::vector<BasicBlock*> &code);
  49
  50   private:
  51     void findInputsOutputs(Values &inputs, Values &outputs,
  52                            BasicBlock *newHeader,
  53                            BasicBlock *newRootNode);
  54
  55     void processPhiNodeInputs(PHINode *Phi,
  56                               Values &inputs,
  57                               BasicBlock *newHeader,
  58                               BasicBlock *newRootNode);
  59
  60     void rewritePhiNodes(Function *F, BasicBlock *newFuncRoot);
  61
  62     Function *constructFunction(const Values &inputs,
  63                                 const Values &outputs,
  64                                 BasicBlock *newRootNode, BasicBlock *newHeader,
  65                                 Function *oldFunction, Module *M);
  66
  67     void moveCodeToFunction(Function *newFunction);
  68
  69     void emitCallAndSwitchStatement(Function *newFunction,
  70                                     BasicBlock *newHeader,
  71                                     Values &inputs,
  72                                     Values &outputs);
  73
  74   };
  75 }
  76
  77 void CodeExtractor::processPhiNodeInputs(PHINode *Phi,
  78                                          Values &inputs,
  79                                          BasicBlock *codeReplacer,
  80                                          BasicBlock *newFuncRoot) {
  81   // Separate incoming values and BasicBlocks as internal/external. We ignore
  82   // the case where both the value and BasicBlock are internal, because we don't
  83   // need to do a thing.
  84   std::vector<unsigned> EValEBB;
  85   std::vector<unsigned> EValIBB;
  86   std::vector<unsigned> IValEBB;
  87
  88   for (unsigned i = 0, e = Phi->getNumIncomingValues(); i != e; ++i) {
  89     Value *phiVal = Phi->getIncomingValue(i);
  90     if (Instruction *Inst = dyn_cast<Instruction>(phiVal)) {
  91       if (BlocksToExtract.count(Inst->getParent())) {
  92         if (!BlocksToExtract.count(Phi->getIncomingBlock(i)))
  93           IValEBB.push_back(i);
  94       } else {
  95         if (BlocksToExtract.count(Phi->getIncomingBlock(i)))
  96           EValIBB.push_back(i);
  97         else
  98           EValEBB.push_back(i);
  99       }
 100     } else if (Argument *Arg = dyn_cast<Argument>(phiVal)) {
 101       // arguments are external
 102       if (BlocksToExtract.count(Phi->getIncomingBlock(i)))
 103         EValIBB.push_back(i);
 104       else
 105         EValEBB.push_back(i);
 106     } else {
 107       // Globals/Constants are internal, but considered `external' if they are
 108       // coming from an external block.
 109       if (!BlocksToExtract.count(Phi->getIncomingBlock(i)))
 110         EValEBB.push_back(i);
 111     }
 112   }
 113
 114   // Both value and block are external. Need to group all of these, have an
 115   // external phi, pass the result as an argument, and have THIS phi use that
 116   // result.
 117   if (EValEBB.size() > 0) {
 118     if (EValEBB.size() == 1) {
 119       // Now if it's coming from the newFuncRoot, it's that funky input
 120       unsigned phiIdx = EValEBB[0];
 121       if (!isa<Constant>(Phi->getIncomingValue(phiIdx))) {
 122         PhiVal2Arg[Phi].push_back(std::make_pair(phiIdx, inputs.size()));
 123         // We can just pass this value in as argument
 124         inputs.push_back(Phi->getIncomingValue(phiIdx));
 125       }
 126       Phi->setIncomingBlock(phiIdx, newFuncRoot);
 127     } else {
 128       PHINode *externalPhi = new PHINode(Phi->getType(), "extPhi");
 129       codeReplacer->getInstList().insert(codeReplacer->begin(), externalPhi);
 130       for (std::vector<unsigned>::iterator i = EValEBB.begin(),
 131              e = EValEBB.end(); i != e; ++i) {
 132         externalPhi->addIncoming(Phi->getIncomingValue(*i),
 133                                  Phi->getIncomingBlock(*i));
 134
 135         // We make these values invalid instead of deleting them because that
 136         // would shift the indices of other values... The fixPhiNodes should
 137         // clean these phi nodes up later.
 138         Phi->setIncomingValue(*i, 0);
 139         Phi->setIncomingBlock(*i, 0);
 140       }
 141       PhiVal2Arg[Phi].push_back(std::make_pair(Phi->getNumIncomingValues(),
 142                                                inputs.size()));
 143       // We can just pass this value in as argument
 144       inputs.push_back(externalPhi);
 145     }
 146   }
 147
 148   // When the value is external, but block internal...  just pass it in as
 149   // argument, no change to phi node
 150   for (std::vector<unsigned>::iterator i = EValIBB.begin(),
 151          e = EValIBB.end(); i != e; ++i) {
 152     // rewrite the phi input node to be an argument
 153     PhiVal2Arg[Phi].push_back(std::make_pair(*i, inputs.size()));
 154     inputs.push_back(Phi->getIncomingValue(*i));
 155   }
 156
 157   // Value internal, block external this can happen if we are extracting a part
 158   // of a loop.
 159   for (std::vector<unsigned>::iterator i = IValEBB.begin(),
 160          e = IValEBB.end(); i != e; ++i) {
 161     assert(0 && "Cannot (YET) handle internal values via external blocks");
 162   }
 163 }
 164
 165
 166 void CodeExtractor::findInputsOutputs(Values &inputs, Values &outputs,
 167                                       BasicBlock *newHeader,
 168                                       BasicBlock *newRootNode) {
 169   for (std::set<BasicBlock*>::const_iterator ci = BlocksToExtract.begin(),
 170        ce = BlocksToExtract.end(); ci != ce; ++ci) {
 171     BasicBlock *BB = *ci;
 172     for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I) {
 173       // If a used value is defined outside the region, it's an input.  If an
 174       // instruction is used outside the region, it's an output.
 175       if (PHINode *Phi = dyn_cast<PHINode>(I)) {
 176         processPhiNodeInputs(Phi, inputs, newHeader, newRootNode);
 177       } else {
 178         // All other instructions go through the generic input finder
 179         // Loop over the operands of each instruction (inputs)
 180         for (User::op_iterator op = I->op_begin(), opE = I->op_end();
 181              op != opE; ++op)
 182           if (Instruction *opI = dyn_cast<Instruction>(*op)) {
 183             // Check if definition of this operand is within the loop
 184             if (!BlocksToExtract.count(opI->getParent()))
 185               inputs.push_back(opI);
 186           } else if (isa<Argument>(*op)) {
 187             inputs.push_back(*op);
 188           }
 189       }
 190
 191       // Consider uses of this instruction (outputs)
 192       for (Value::use_iterator UI = I->use_begin(), E = I->use_end();
 193            UI != E; ++UI)
 194         if (!BlocksToExtract.count(cast<Instruction>(*UI)->getParent()))
 195           outputs.push_back(*UI);
 196     } // for: insts
 197   } // for: basic blocks
 198 }
 199
 200 void CodeExtractor::rewritePhiNodes(Function *F,
 201                                     BasicBlock *newFuncRoot) {
 202   // Write any changes that were saved before: use function arguments as inputs
 203   for (PhiVal2ArgTy::iterator i = PhiVal2Arg.begin(), e = PhiVal2Arg.end();
 204        i != e; ++i) {
 205     PHINode *phi = i->first;
 206     PhiValChangesTy &values = i->second;
 207     for (unsigned cIdx = 0, ce = values.size(); cIdx != ce; ++cIdx)
 208     {
 209       unsigned phiValueIdx = values[cIdx].first, argNum = values[cIdx].second;
 210       if (phiValueIdx < phi->getNumIncomingValues())
 211         phi->setIncomingValue(phiValueIdx, getFunctionArg(F, argNum));
 212       else
 213         phi->addIncoming(getFunctionArg(F, argNum), newFuncRoot);
 214     }
 215   }
 216
 217   // Delete any invalid Phi node inputs that were marked as NULL previously
 218   for (PhiVal2ArgTy::iterator i = PhiVal2Arg.begin(), e = PhiVal2Arg.end();
 219        i != e; ++i) {
 220     PHINode *phi = i->first;
 221     for (unsigned idx = 0, end = phi->getNumIncomingValues(); idx != end; ++idx)
 222     {
 223       if (phi->getIncomingValue(idx) == 0 && phi->getIncomingBlock(idx) == 0) {
 224         phi->removeIncomingValue(idx);
 225         --idx;
 226         --end;
 227       }
 228     }
 229   }
 230
 231   // We are done with the saved values
 232   PhiVal2Arg.clear();
 233 }
 234
 235
 236 /// constructFunction - make a function based on inputs and outputs, as follows:
 237 /// f(in0, ..., inN, out0, ..., outN)
 238 ///
 239 Function *CodeExtractor::constructFunction(const Values &inputs,
 240                                            const Values &outputs,
 241                                            BasicBlock *newRootNode,
 242                                            BasicBlock *newHeader,
 243                                            Function *oldFunction, Module *M) {
 244   DEBUG(std::cerr << "inputs: " << inputs.size() << "\n");
 245   DEBUG(std::cerr << "outputs: " << outputs.size() << "\n");
 246   BasicBlock *header = *BlocksToExtract.begin();
 247
 248   // This function returns unsigned, outputs will go back by reference.
 249   Type *retTy = Type::UShortTy;
 250   std::vector<const Type*> paramTy;
 251
 252   // Add the types of the input values to the function's argument list
 253   for (Values::const_iterator i = inputs.begin(),
 254          e = inputs.end(); i != e; ++i) {
 255     const Value *value = *i;
 256     DEBUG(std::cerr << "value used in func: " << value << "\n");
 257     paramTy.push_back(value->getType());
 258   }
 259
 260   // Add the types of the output values to the function's argument list, but
 261   // make them pointer types for scalars
 262   for (Values::const_iterator i = outputs.begin(),
 263          e = outputs.end(); i != e; ++i) {
 264     const Value *value = *i;
 265     DEBUG(std::cerr << "instr used in func: " << value << "\n");
 266     const Type *valueType = value->getType();
 267     // Convert scalar types into a pointer of that type
 268     if (valueType->isPrimitiveType()) {
 269       valueType = PointerType::get(valueType);
 270     }
 271     paramTy.push_back(valueType);
 272   }
 273
 274   DEBUG(std::cerr << "Function type: " << retTy << " f(");
 275   for (std::vector<const Type*>::iterator i = paramTy.begin(),
 276          e = paramTy.end(); i != e; ++i)
 277     DEBUG(std::cerr << *i << ", ");
 278   DEBUG(std::cerr << ")\n");
 279
 280   const FunctionType *funcType = FunctionType::get(retTy, paramTy, false);
 281
 282   // Create the new function
 283   Function *newFunction = new Function(funcType,
 284                                        GlobalValue::InternalLinkage,
 285                                        oldFunction->getName() + "_code", M);
 286   newFunction->getBasicBlockList().push_back(newRootNode);
 287
 288   for (unsigned i = 0, e = inputs.size(); i != e; ++i) {
 289     std::vector<User*> Users(inputs[i]->use_begin(), inputs[i]->use_end());
 290     for (std::vector<User*>::iterator use = Users.begin(), useE = Users.end();
 291          use != useE; ++use)
 292       if (Instruction* inst = dyn_cast<Instruction>(*use))
 293         if (BlocksToExtract.count(inst->getParent()))
 294           inst->replaceUsesOfWith(inputs[i], getFunctionArg(newFunction, i));
 295   }
 296
 297   // Rewrite branches to basic blocks outside of the loop to new dummy blocks
 298   // within the new function. This must be done before we lose track of which
 299   // blocks were originally in the code region.
 300   std::vector<User*> Users(header->use_begin(), header->use_end());
 301   for (std::vector<User*>::iterator i = Users.begin(), e = Users.end();
 302        i != e; ++i) {
 303     if (BranchInst *inst = dyn_cast<BranchInst>(*i)) {
 304       BasicBlock *BB = inst->getParent();
 305       if (!BlocksToExtract.count(BB) && BB->getParent() == oldFunction) {
 306         // The BasicBlock which contains the branch is not in the region
 307         // modify the branch target to a new block
 308         inst->replaceUsesOfWith(header, newHeader);
 309       }
 310     }
 311   }
 312
 313   return newFunction;
 314 }
 315
 316 void CodeExtractor::moveCodeToFunction(Function *newFunction) {
 317   Function *oldFunc = (*BlocksToExtract.begin())->getParent();
 318   Function::BasicBlockListType &oldBlocks = oldFunc->getBasicBlockList();
 319   Function::BasicBlockListType &newBlocks = newFunction->getBasicBlockList();
 320
 321   for (std::set<BasicBlock*>::const_iterator i = BlocksToExtract.begin(),
 322          e = BlocksToExtract.end(); i != e; ++i) {
 323     // Delete the basic block from the old function, and the list of blocks
 324     oldBlocks.remove(*i);
 325
 326     // Insert this basic block into the new function
 327     newBlocks.push_back(*i);
 328   }
 329 }
 330
 331 void
 332 CodeExtractor::emitCallAndSwitchStatement(Function *newFunction,
 333                                           BasicBlock *codeReplacer,
 334                                           Values &inputs,
 335                                           Values &outputs)
 336 {
 337   // Emit a call to the new function, passing allocated memory for outputs and
 338   // just plain inputs for non-scalars
 339   std::vector<Value*> params(inputs);
 340
 341   for (Values::const_iterator i = outputs.begin(), e = outputs.end(); i != e;
 342        ++i) {
 343     Value *Output = *i;
 344     // Create allocas for scalar outputs
 345     if (Output->getType()->isPrimitiveType()) {
 346       AllocaInst *alloca =
 347         new AllocaInst((*i)->getType(), 0, Output->getName()+".loc",
 348                        codeReplacer->getParent()->begin()->begin());
 349       params.push_back(alloca);
 350
 351       LoadInst *load = new LoadInst(alloca, Output->getName()+".reload");
 352       codeReplacer->getInstList().push_back(load);
 353       std::vector<User*> Users((*i)->use_begin(), (*i)->use_end());
 354       for (std::vector<User*>::iterator use = Users.begin(), useE =Users.end();
 355            use != useE; ++use) {
 356         if (Instruction* inst = dyn_cast<Instruction>(*use)) {
 357           if (!BlocksToExtract.count(inst->getParent()))
 358             inst->replaceUsesOfWith(*i, load);
 359         }
 360       }
 361     } else {
 362       params.push_back(*i);
 363     }
 364   }
 365
 366   CallInst *call = new CallInst(newFunction, params, "targetBlock");
 367   codeReplacer->getInstList().push_front(call);
 368
 369   // Now we can emit a switch statement using the call as a value.
 370   SwitchInst *TheSwitch = new SwitchInst(call, codeReplacer, codeReplacer);
 371
 372   // Since there may be multiple exits from the original region, make the new
 373   // function return an unsigned, switch on that number.  This loop iterates
 374   // over all of the blocks in the extracted region, updating any terminator
 375   // instructions in the to-be-extracted region that branch to blocks that are
 376   // not in the region to be extracted.
 377   std::map<BasicBlock*, BasicBlock*> ExitBlockMap;
 378
 379   unsigned switchVal = 0;
 380   for (std::set<BasicBlock*>::const_iterator i = BlocksToExtract.begin(),
 381          e = BlocksToExtract.end(); i != e; ++i) {
 382     TerminatorInst *TI = (*i)->getTerminator();
 383     for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
 384       if (!BlocksToExtract.count(TI->getSuccessor(i))) {
 385         BasicBlock *OldTarget = TI->getSuccessor(i);
 386         // add a new basic block which returns the appropriate value
 387         BasicBlock *&NewTarget = ExitBlockMap[OldTarget];
 388         if (!NewTarget) {
 389           // If we don't already have an exit stub for this non-extracted
 390           // destination, create one now!
 391           NewTarget = new BasicBlock(OldTarget->getName() + ".exitStub",
 392                                      newFunction);
 393
 394           ConstantUInt *brVal = ConstantUInt::get(Type::UShortTy, switchVal++);
 395           ReturnInst *NTRet = new ReturnInst(brVal, NewTarget);
 396
 397           // Update the switch instruction.
 398           TheSwitch->addCase(brVal, OldTarget);
 399
 400           // Restore values just before we exit
 401           // FIXME: Use a GetElementPtr to bunch the outputs in a struct
 402           for (unsigned out = 0, e = outputs.size(); out != e; ++out)
 403             new StoreInst(outputs[out], getFunctionArg(newFunction, out),NTRet);
 404         }
 405
 406         // rewrite the original branch instruction with this new target
 407         TI->setSuccessor(i, NewTarget);
 408       }
 409   }
 410
 411   // Now that we've done the deed, make the default destination of the switch
 412   // instruction be one of the exit blocks of the region.
 413   if (TheSwitch->getNumSuccessors() > 1) {
 414     // FIXME: this is broken w.r.t. PHI nodes, but the old code was more broken.
 415     // This edge is not traversable.
 416     TheSwitch->setSuccessor(0, TheSwitch->getSuccessor(1));
 417   }
 418 }
 419
 420
 421 /// ExtractRegion - Removes a loop from a function, replaces it with a call to
 422 /// new function. Returns pointer to the new function.
 423 ///
 424 /// algorithm:
 425 ///
 426 /// find inputs and outputs for the region
 427 ///
 428 /// for inputs: add to function as args, map input instr* to arg#
 429 /// for outputs: add allocas for scalars,
 430 ///             add to func as args, map output instr* to arg#
 431 ///
 432 /// rewrite func to use argument #s instead of instr*
 433 ///
 434 /// for each scalar output in the function: at every exit, store intermediate
 435 /// computed result back into memory.
 436 ///
 437 Function *CodeExtractor::ExtractCodeRegion(const std::vector<BasicBlock*> &code)
 438 {
 439   // 1) Find inputs, outputs
 440   // 2) Construct new function
 441   //  * Add allocas for defs, pass as args by reference
 442   //  * Pass in uses as args
 443   // 3) Move code region, add call instr to func
 444   //
 445   BlocksToExtract.insert(code.begin(), code.end());
 446
 447   Values inputs, outputs;
 448
 449   // Assumption: this is a single-entry code region, and the header is the first
 450   // block in the region.
 451   BasicBlock *header = code[0];
 452   for (unsigned i = 1, e = code.size(); i != e; ++i)
 453     for (pred_iterator PI = pred_begin(code[i]), E = pred_end(code[i]);
 454          PI != E; ++PI)
 455       assert(BlocksToExtract.count(*PI) &&
 456              "No blocks in this region may have entries from outside the region"
 457              " except for the first block!");
 458
 459   Function *oldFunction = header->getParent();
 460
 461   // This takes place of the original loop
 462   BasicBlock *codeReplacer = new BasicBlock("codeRepl", oldFunction);
 463
 464   // The new function needs a root node because other nodes can branch to the
 465   // head of the loop, and the root cannot have predecessors
 466   BasicBlock *newFuncRoot = new BasicBlock("newFuncRoot");
 467   newFuncRoot->getInstList().push_back(new BranchInst(header));
 468
 469   // Find inputs to, outputs from the code region
 470   //
 471   // If one of the inputs is coming from a different basic block and it's in a
 472   // phi node, we need to rewrite the phi node:
 473   //
 474   // * All the inputs which involve basic blocks OUTSIDE of this region go into
 475   //   a NEW phi node that takes care of finding which value really came in.
 476   //   The result of this phi is passed to the function as an argument.
 477   //
 478   // * All the other phi values stay.
 479   //
 480   // FIXME: PHI nodes' incoming blocks aren't being rewritten to accomodate for
 481   // blocks moving to a new function.
 482   // SOLUTION: move Phi nodes out of the loop header into the codeReplacer, pass
 483   // the values as parameters to the function
 484   findInputsOutputs(inputs, outputs, codeReplacer, newFuncRoot);
 485
 486   // Step 2: Construct new function based on inputs/outputs,
 487   // Add allocas for all defs
 488   Function *newFunction = constructFunction(inputs, outputs, newFuncRoot,
 489                                             codeReplacer, oldFunction,
 490                                             oldFunction->getParent());
 491
 492   rewritePhiNodes(newFunction, newFuncRoot);
 493
 494   emitCallAndSwitchStatement(newFunction, codeReplacer, inputs, outputs);
 495
 496   moveCodeToFunction(newFunction);
 497
 498   DEBUG(if (verifyFunction(*newFunction)) abort());
 499   return newFunction;
 500 }
 501
 502 /// ExtractCodeRegion - slurp a sequence of basic blocks into a brand new
 503 /// function
 504 ///
 505 Function* llvm::ExtractCodeRegion(const std::vector<BasicBlock*> &code) {
 506   return CodeExtractor().ExtractCodeRegion(code);
 507 }
 508
 509 /// ExtractBasicBlock - slurp a natural loop into a brand new function
 510 ///
 511 Function* llvm::ExtractLoop(Loop *L) {
 512   return CodeExtractor().ExtractCodeRegion(L->getBlocks());
 513 }
 514
 515 /// ExtractBasicBlock - slurp a basic block into a brand new function
 516 ///
 517 Function* llvm::ExtractBasicBlock(BasicBlock *BB) {
 518   std::vector<BasicBlock*> Blocks;
 519   Blocks.push_back(BB);
 520   return CodeExtractor().ExtractCodeRegion(Blocks);
 521 }