Changes in ‘CodeExtractor’ for call-site block

[ashutosh.nema]

Summary

This change is to expose call site block of outlined function in ‘CodeExtractor’.

When code extractor works on blocks inside a loop it extracts properly.
But it didn’t register extracted function call site block to loop.
Because of it loop verification fails with error:
"Loop block has no in-loop successors!".

Call site block registration to loop is very much needed.
I like to expose call site block (‘codeReplacer’) by adding getter functions to ‘CodeExtractor’ class.
So who ever using code extractor they can access that call-site block and register if needed.

Made following changes in CodeExtractor.

1. Made ‘codeReplacer’ as member of ‘CodeExtractor’

2. Added a new method ‘getCodeReplacerIfAvailable’ to ‘CodeExtractor’ class.

This method return ‘codeReplacer’ block.

3. ‘constructFunction’ method works on ‘codeReplacer’ block and accepts as argument.

Now ‘codeReplacer’ block is a member to class so no need of passing it as argument.

4. ‘emitCallAndSwitchStatement’ method works on ‘codeReplacer’ block and accepts as argument.

Now ‘codeReplacer’ block is a member to class so no need of passing it as argument.

5. Updated callsites to ‘constructFunction’ & ‘emitCallAndSwitchStatement’.

6. In ‘CodeExtractor’ constructor setting ‘CodeReplacer’ as nullptr.

Files Modified:
include/llvm/Transforms/Utils/CodeExtractor.h
lib/Transforms/Utils/CodeExtractor.cpp

[chandlerc]

There are a bunch of coding convention problems here. Also, you don't provide a test case that this fixes. It's hard to evaluate what problem you're actually solving without the test case.

[ashutosh.nema]

Thanks Chandler for looking into this.

There are a bunch of coding convention problems here.

I’ll fix these.

Also, you don't provide a test case that this fixes.
It's hard to evaluate what problem you're actually solving without the test case.

I was trying to use code extractor for inner loops.
i.e.
for (i…N){

for(j..N)
{}

}

And like to move inner loop to a function.

CodeExtractor is moving loop to a outlined function properly:

OutlinedFunction(args)
{

for(j..N)
{}

}
for (i…N){

call OutlinedFunction(args) // call-site appears in a block

}

But it’s not registering call-site block to outer loop(if available).

Because of un registered call site block in outer loop, loop verification fails with below error:
"Loop block has no in-loop successors!".

I’m not very sure we can write test case for this because currently no optimization
moves inner loop to a outlined function. I can try if you have any suggestions.

So in this change, I exposed call site block.
If required, optimizations can register call-site block to outer loop.

Hope this explains the problem I'm trying to solve.

Regards,
Ashutosh

include/llvm/Transforms/Utils/CodeExtractor.h

namespace llvm {
   /// 3) Add allocas for any scalar outputs, adding all of the outputs' allocas
   ///    as arguments, and inserting stores to the arguments for any scalars.
   class CodeExtractor {
     typedef SetVector<Value *> ValueSet;
 
     // Various bits of state computed on construction.
     DominatorTree *const DT;
     const bool AggregateArgs;
+    BasicBlock * CodeReplacer;
 
     // Bits of intermediate state computed at various phases of extraction.
     SetVector<BasicBlock *> Blocks;
     unsigned NumExitBlocks;
     Type *RetTy;
 
   public:
     /// \brief Create a code extractor for a single basic block.
     ///
     /// In this formation, we don't require a dominator tree. The given basic
     /// block is set up for extraction.
     CodeExtractor(BasicBlock *BB, bool AggregateArgs = false);
+    BasicBlock * getCodeReplacerIfAvailable() const;
 
     /// \brief Create a code extractor for a sequence of blocks.
     ///
     /// Given a sequence of basic blocks where the first block in the sequence
     /// dominates the rest, prepare a code extractor object for pulling this
     /// sequence out into its new function. When a DominatorTree is also given,
     /// extra checking and transformations are enabled.
     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
namespace llvm {
 
   private:
     void severSplitPHINodes(BasicBlock *&Header);
     void splitReturnBlocks();
 
     Function *constructFunction(const ValueSet &inputs,
                                 const ValueSet &outputs,
                                 BasicBlock *header,
-                                BasicBlock *newRootNode, BasicBlock *newHeader,
+                                BasicBlock *newRootNode,
                                 Function *oldFunction, Module *M);
 
     void moveCodeToFunction(Function *newFunction);
 
     void emitCallAndSwitchStatement(Function *newFunction,
-                                    BasicBlock *newHeader,
                                     ValueSet &inputs,
                                     ValueSet &outputs);
   };
 }
 
 #endif

lib/Transforms/Utils/CodeExtractor.cpp

if (!RN.isSubRegion())
     return buildExtractionBlockSet(RN.getNodeAs<BasicBlock>());
 
   const Region &R = *RN.getNodeAs<Region>();
 
   return buildExtractionBlockSet(R.block_begin(), R.block_end());
 }
 
 CodeExtractor::CodeExtractor(BasicBlock *BB, bool AggregateArgs)
-  : DT(nullptr), AggregateArgs(AggregateArgs||AggregateArgsOpt),
-    Blocks(buildExtractionBlockSet(BB)), NumExitBlocks(~0U) {}
+  : DT(nullptr), AggregateArgs(AggregateArgs||AggregateArgsOpt), 
+    CodeReplacer(nullptr), Blocks(buildExtractionBlockSet(BB)), 
+    NumExitBlocks(~0U) {}
 
 CodeExtractor::CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
                              bool AggregateArgs)
-  : DT(DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
-    Blocks(buildExtractionBlockSet(BBs)), NumExitBlocks(~0U) {}
+  : DT(DT), AggregateArgs(AggregateArgs||AggregateArgsOpt), 
+    CodeReplacer(nullptr), Blocks(buildExtractionBlockSet(BBs)), 
+    NumExitBlocks(~0U) {}
 
 CodeExtractor::CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs)
-  : DT(&DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
-    Blocks(buildExtractionBlockSet(L.getBlocks())), NumExitBlocks(~0U) {}
+  : DT(&DT), AggregateArgs(AggregateArgs||AggregateArgsOpt), 
+    CodeReplacer(nullptr), Blocks(buildExtractionBlockSet(L.getBlocks())), 
+    NumExitBlocks(~0U) {}
 
 CodeExtractor::CodeExtractor(DominatorTree &DT, const RegionNode &RN,
                              bool AggregateArgs)
-  : DT(&DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
-    Blocks(buildExtractionBlockSet(RN)), NumExitBlocks(~0U) {}
+  : DT(&DT), AggregateArgs(AggregateArgs||AggregateArgsOpt), 
+    CodeReplacer(nullptr), Blocks(buildExtractionBlockSet(RN)), 
+    NumExitBlocks(~0U) {}
 
 /// definedInRegion - Return true if the specified value is defined in the
 /// extracted region.
 static bool definedInRegion(const SetVector<BasicBlock *> &Blocks, Value *V) {
   if (Instruction *I = dyn_cast<Instruction>(V))
     if (Blocks.count(I->getParent()))
       return true;
   return false;

 
 /// constructFunction - make a function based on inputs and outputs, as follows:
 /// f(in0, ..., inN, out0, ..., outN)
 ///
 Function *CodeExtractor::constructFunction(const ValueSet &inputs,
                                            const ValueSet &outputs,
                                            BasicBlock *header,
                                            BasicBlock *newRootNode,
-                                           BasicBlock *newHeader,
                                            Function *oldFunction,
                                            Module *M) {
   DEBUG(dbgs() << "inputs: " << inputs.size() << "\n");
   DEBUG(dbgs() << "outputs: " << outputs.size() << "\n");
 
   // This function returns unsigned, outputs will go back by reference.
   switch (NumExitBlocks) {
   case 0:
Function *CodeExtractor::constructFunction(const ValueSet &inputs,
   // blocks were originally in the code region.
   std::vector<User*> Users(header->user_begin(), header->user_end());
   for (unsigned i = 0, e = Users.size(); i != e; ++i)
     // The BasicBlock which contains the branch is not in the region
     // modify the branch target to a new block
     if (TerminatorInst *TI = dyn_cast<TerminatorInst>(Users[i]))
       if (!Blocks.count(TI->getParent()) &&
           TI->getParent()->getParent() == oldFunction)
-        TI->replaceUsesOfWith(header, newHeader);
+        TI->replaceUsesOfWith(header, CodeReplacer);
 
   return newFunction;
 }
 
 /// FindPhiPredForUseInBlock - Given a value and a basic block, find a PHI
 /// that uses the value within the basic block, and return the predecessor
 /// block associated with that use, or return 0 if none is found.
 static BasicBlock* FindPhiPredForUseInBlock(Value* Used, BasicBlock* BB) {
   for (Use &U : Used->uses()) {
      PHINode *P = dyn_cast<PHINode>(U.getUser());
      if (P && P->getParent() == BB)
        return P->getIncomingBlock(U);
   }
 
   return nullptr;
 }
 
 /// emitCallAndSwitchStatement - This method sets up the caller side by adding
 /// the call instruction, splitting any PHI nodes in the header block as
 /// necessary.
 void CodeExtractor::
-emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
-                           ValueSet &inputs, ValueSet &outputs) {
+emitCallAndSwitchStatement(Function *newFunction, ValueSet &inputs, 
+                           ValueSet &outputs) {
   // Emit a call to the new function, passing in: *pointer to struct (if
   // aggregating parameters), or plan inputs and allocated memory for outputs
   std::vector<Value*> params, StructValues, ReloadOutputs, Reloads;
   
   LLVMContext &Context = newFunction->getContext();
 
   // Add inputs as params, or to be filled into the struct
   for (ValueSet::iterator i = inputs.begin(), e = inputs.end(); i != e; ++i)
     if (AggregateArgs)
       StructValues.push_back(*i);
     else
       params.push_back(*i);
 
   // Create allocas for the outputs
   for (ValueSet::iterator i = outputs.begin(), e = outputs.end(); i != e; ++i) {
     if (AggregateArgs) {
       StructValues.push_back(*i);
     } else {
       AllocaInst *alloca =
         new AllocaInst((*i)->getType(), nullptr, (*i)->getName()+".loc",
-                       codeReplacer->getParent()->begin()->begin());
+                       CodeReplacer->getParent()->begin()->begin());
       ReloadOutputs.push_back(alloca);
       params.push_back(alloca);
     }
   }
 
   AllocaInst *Struct = nullptr;
   if (AggregateArgs && (inputs.size() + outputs.size() > 0)) {
     std::vector<Type*> ArgTypes;
     for (ValueSet::iterator v = StructValues.begin(),
            ve = StructValues.end(); v != ve; ++v)
       ArgTypes.push_back((*v)->getType());
 
     // Allocate a struct at the beginning of this function
     Type *StructArgTy = StructType::get(newFunction->getContext(), ArgTypes);
     Struct =
       new AllocaInst(StructArgTy, nullptr, "structArg",
-                     codeReplacer->getParent()->begin()->begin());
+                     CodeReplacer->getParent()->begin()->begin());
     params.push_back(Struct);
 
     for (unsigned i = 0, e = inputs.size(); i != e; ++i) {
       Value *Idx[2];
       Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
       Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), i);
       GetElementPtrInst *GEP =
         GetElementPtrInst::Create(Struct, Idx,
                                   "gep_" + StructValues[i]->getName());
-      codeReplacer->getInstList().push_back(GEP);
+      CodeReplacer->getInstList().push_back(GEP);
       StoreInst *SI = new StoreInst(StructValues[i], GEP);
-      codeReplacer->getInstList().push_back(SI);
+      CodeReplacer->getInstList().push_back(SI);
     }
   }
 
   // Emit the call to the function
   CallInst *call = CallInst::Create(newFunction, params,
                                     NumExitBlocks > 1 ? "targetBlock" : "");
-  codeReplacer->getInstList().push_back(call);
+  CodeReplacer->getInstList().push_back(call);
 
   Function::arg_iterator OutputArgBegin = newFunction->arg_begin();
   unsigned FirstOut = inputs.size();
   if (!AggregateArgs)
     std::advance(OutputArgBegin, inputs.size());
 
   // Reload the outputs passed in by reference
   for (unsigned i = 0, e = outputs.size(); i != e; ++i) {
     Value *Output = nullptr;
     if (AggregateArgs) {
       Value *Idx[2];
       Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
       Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), FirstOut + i);
       GetElementPtrInst *GEP
         = GetElementPtrInst::Create(Struct, Idx,
                                     "gep_reload_" + outputs[i]->getName());
-      codeReplacer->getInstList().push_back(GEP);
+      CodeReplacer->getInstList().push_back(GEP);
       Output = GEP;
     } else {
       Output = ReloadOutputs[i];
     }
     LoadInst *load = new LoadInst(Output, outputs[i]->getName()+".reload");
     Reloads.push_back(load);
-    codeReplacer->getInstList().push_back(load);
+    CodeReplacer->getInstList().push_back(load);
     std::vector<User*> Users(outputs[i]->user_begin(), outputs[i]->user_end());
     for (unsigned u = 0, e = Users.size(); u != e; ++u) {
       Instruction *inst = cast<Instruction>(Users[u]);
       if (!Blocks.count(inst->getParent()))
         inst->replaceUsesOfWith(outputs[i], load);
     }
   }
 
   // Now we can emit a switch statement using the call as a value.
   SwitchInst *TheSwitch =
       SwitchInst::Create(Constant::getNullValue(Type::getInt16Ty(Context)),
-                         codeReplacer, 0, codeReplacer);
+                         CodeReplacer, 0, CodeReplacer);
 
   // Since there may be multiple exits from the original region, make the new
   // function return an unsigned, switch on that number.  This loop iterates
   // over all of the blocks in the extracted region, updating any terminator
   // instructions in the to-be-extracted region that branch to blocks that are
   // not in the region to be extracted.
   std::map<BasicBlock*, BasicBlock*> ExitBlockMap;
 
for (SetVector<BasicBlock*>::const_iterator i = Blocks.begin(),
     // Delete the basic block from the old function, and the list of blocks
     oldBlocks.remove(*i);
 
     // Insert this basic block into the new function
     newBlocks.push_back(*i);
   }
 }
 
+// Return CodeReplacer block.
+BasicBlock * CodeExtractor::getCodeReplacerIfAvailable() const
+{
+   return CodeReplacer;
+}
+
 Function *CodeExtractor::extractCodeRegion() {
   if (!isEligible())
     return nullptr;
 
   ValueSet inputs, outputs;
 
   // Assumption: this is a single-entry code region, and the header is the first
   // block in the region.
   BasicBlock *header = *Blocks.begin();
 
   // If we have to split PHI nodes or the entry block, do so now.
   severSplitPHINodes(header);
 
   // If we have any return instructions in the region, split those blocks so
   // that the return is not in the region.
   splitReturnBlocks();
 
   Function *oldFunction = header->getParent();
 
   // This takes place of the original loop
-  BasicBlock *codeReplacer = BasicBlock::Create(header->getContext(), 
-                                                "codeRepl", oldFunction,
-                                                header);
+  CodeReplacer = BasicBlock::Create(header->getContext(), 
 
   // The new function needs a root node because other nodes can branch to the
   // head of the region, but the entry node of a function cannot have preds.
   BasicBlock *newFuncRoot = BasicBlock::Create(header->getContext(), 
                                                "newFuncRoot");
   newFuncRoot->getInstList().push_back(BranchInst::Create(header));
 
   // Find inputs to, outputs from the code region.
   findInputsOutputs(inputs, outputs);
 
   SmallPtrSet<BasicBlock *, 1> ExitBlocks;
   for (SetVector<BasicBlock *>::iterator I = Blocks.begin(), E = Blocks.end();
        I != E; ++I)
     for (succ_iterator SI = succ_begin(*I), SE = succ_end(*I); SI != SE; ++SI)
       if (!Blocks.count(*SI))
         ExitBlocks.insert(*SI);
   NumExitBlocks = ExitBlocks.size();
 
   // Construct new function based on inputs/outputs & add allocas for all defs.
   Function *newFunction = constructFunction(inputs, outputs, header,
                                             newFuncRoot,
-                                            codeReplacer, oldFunction,
+                                            oldFunction,
                                             oldFunction->getParent());
 
-  emitCallAndSwitchStatement(newFunction, codeReplacer, inputs, outputs);
+  emitCallAndSwitchStatement(newFunction, inputs, outputs);
 
   moveCodeToFunction(newFunction);
 
   // Loop over all of the PHI nodes in the header block, and change any
   // references to the old incoming edge to be the new incoming edge.
   for (BasicBlock::iterator I = header->begin(); isa<PHINode>(I); ++I) {
     PHINode *PN = cast<PHINode>(I);
     for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
       if (!Blocks.count(PN->getIncomingBlock(i)))
         PN->setIncomingBlock(i, newFuncRoot);
   }
 
-  // Look at all successors of the codeReplacer block.  If any of these blocks
+  // Look at all successors of the CodeReplacer block.  If any of these blocks
   // had PHI nodes in them, we need to update the "from" block to be the code
   // replacer, not the original block in the extracted region.
-  std::vector<BasicBlock*> Succs(succ_begin(codeReplacer),
-                                 succ_end(codeReplacer));
+  std::vector<BasicBlock*> Succs(succ_begin(CodeReplacer),
+                                 succ_end(CodeReplacer));
   for (unsigned i = 0, e = Succs.size(); i != e; ++i)
     for (BasicBlock::iterator I = Succs[i]->begin(); isa<PHINode>(I); ++I) {
       PHINode *PN = cast<PHINode>(I);
       std::set<BasicBlock*> ProcessedPreds;
       for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
         if (Blocks.count(PN->getIncomingBlock(i))) {
           if (ProcessedPreds.insert(PN->getIncomingBlock(i)).second)
-            PN->setIncomingBlock(i, codeReplacer);
+            PN->setIncomingBlock(i, CodeReplacer);
           else {
             // There were multiple entries in the PHI for this block, now there
             // is only one, so remove the duplicated entries.
             PN->removeIncomingValue(i, false);
             --i; --e;
           }
         }
     }