[IROutliner] Using Canonical Values to find Corresponding Items between Regions

[AndrewLitteken]

Summary

D104143 introduced canonical value numbering between regions, which allows for the easy identification of items across a region, eliminating the need in the outliner to create parallel lists of instructions for each region, and replace output values in a less convoluted way,

[jroelofs]

llvm/include/llvm/Transforms/IPO/IROutliner.h

166

I'd drop the Optional<> and use nullptr to represent "didn't find one". That removes ambiguity around Optional<Value*>(nullptr), and would improve the call sites.

llvm/lib/Transforms/IPO/IROutliner.cpp

1171

1178

... to fix the dereference after dyn_cast without checking it for null.

1184

[paquette]

llvm/include/llvm/Transforms/IPO/IROutliner.h

166

+1

[AndrewLitteken]

Updating function return type, and other comments

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

174

I think you can return FoundValueOpt.getValueOr(nullptr).

[paquette]

This looks okay with nits.

I think the commit message can be more descriptive here though. It should explain why you're changing behaviour surrounding PHIs/stores etc.

llvm/lib/Transforms/IPO/IROutliner.cpp

1173

https://llvm.org/docs/ProgrammersManual.html#the-isa-cast-and-dyn-cast-templates

Seems like it's more appropriate to use cast here than dyn_cast, considering you assert right after.

1174

llvm/include/llvm/Transforms/IPO/IROutliner.h

struct OutlinableRegion {
   void splitCandidate();
 
   /// For the contained region, reattach the BasicBlock at the starting and
   /// ending instructions of the contained IRSimilarityCandidate, or if the
   /// function has been extracted, the start and end of the BasicBlock
   /// containing the called function.
   void reattachCandidate();
 
+  /// Find a corresponding value for \p V in similar OutlinableRegion \p Other.
+  ///
+  /// \param Other [in] - The OutlinableRegion to find the corresponding Value
+  /// in.
+  /// \param V [in] - The Value to look for in the other region.
+  /// \return An optional that contains the corresponding Value to \p V if it
+  /// exists, otherwise None.
+  Optional<Value *> findCorrespondingValueIn(const OutlinableRegion &Other,
+                                             Value *V);
+
   /// Get the size of the code removed from the region.
   ///
   /// \param [in] TTI - The TargetTransformInfo for the parent function.
   /// \returns the code size of the region
   InstructionCost getBenefit(TargetTransformInfo &TTI);
 };
 
 /// This class is a pass that identifies similarity in a Module, extracts

llvm/lib/Transforms/IPO/IROutliner.cpp

stable_sort(SortedKeys, [](const Value *LHS, const Value *RHS) {
       assert(LHSC && "Not a constant integer in return value?");
       const uint64_t LHSVal = LHSC->getLimitedValue();
       const uint64_t RHSVal = RHSC->getLimitedValue();
 
       return LHSVal < RHSVal;
     });
 }
 
+Optional<Value *>
+OutlinableRegion::findCorrespondingValueIn(const OutlinableRegion &Other,
+                                           Value *V) {
+  Optional<unsigned> GVN = Candidate->getGVN(V);
+  assert(GVN.hasValue() && "No GVN for incoming value");
+  Optional<unsigned> CanonNum = Candidate->getCanonicalNum(*GVN);
+  Optional<unsigned> FirstGVN = Other.Candidate->fromCanonicalNum(*CanonNum);
+  return Other.Candidate->fromGVN(*FirstGVN);
+}
+
 void OutlinableRegion::splitCandidate() {
   assert(!CandidateSplit && "Candidate already split!");
 
   Instruction *BackInst = Candidate->backInstruction();
 
   Instruction *EndInst = nullptr;
   // Check whether the last instruction is a terminator, if it is, we do
   // not split on the following instruction. We leave the block as it is.  We
for (BasicBlock *DescendBB : Descendants) {
       if (!isa<ReturnInst>(Term))
         continue;
       ReturnInst *RI = cast<ReturnInst>(Term);
       Value *RetVal = RI->getReturnValue();
       DenseMap<Value *, BasicBlock *>::iterator VBBIt, PHIVBBIt;
       VBBIt = OutputBBs.find(RetVal);
       assert(VBBIt != OutputBBs.end() && "Could not find output value!");
 
+      // If this is storing a PHINode, we must make sure it is included in the
+      // overall function.
+      StoreInst *SI = dyn_cast<StoreInst>(I);
+      assert(SI && "Is not a store instrution");
+
+      Value *ValueOperand = SI->getValueOperand();
+
       Instruction *NewI = I->clone();
       NewI->setDebugLoc(DebugLoc());
       BasicBlock *OutputBB = VBBIt->second;
       OutputBB->getInstList().push_back(NewI);
       LLVM_DEBUG(dbgs() << "Move store for instruction " << *I << " to "
                         << *OutputBB << "\n");
+
+      StoreInst *NewSI = dyn_cast<StoreInst>(NewI);
+      if (FirstFunction)
+        continue;
+      Optional<Value *> CorrVal =
+          Region.findCorrespondingValueIn(*Group.Regions[0], ValueOperand);
+      assert(CorrVal.hasValue() && "Value is none?");
+      NewSI->setOperand(0, *CorrVal);
     }
 
     // If we added an edge for basic blocks without a predecessor, we remove it
     // here.
     if (EdgeAdded)
       DT.deleteEdge(&DominatingFunction->getEntryBlock(), BB);
     I->eraseFromParent();
 
for (std::pair<unsigned, Constant *> &Const : Region.AggArgToConstant) {
     CST->replaceUsesWithIf(Arg, [OutlinedFunction](Use &U) {
       if (Instruction *I = dyn_cast<Instruction>(U.getUser()))
         return I->getFunction() == OutlinedFunction;
       return false;
     });
   }
 }
 
-/// For the given function, find all the nondebug or lifetime instructions,
-/// and return them as a vector. Exclude any blocks in \p ExludeBlocks.
-///
-/// \param [in] F - The function we collect the instructions from.
-/// \param [in] ExcludeBlocks - BasicBlocks to ignore.
-/// \returns the list of instructions extracted.
-static std::vector<Instruction *>
-collectRelevantInstructions(Function &F,
-                            DenseSet<BasicBlock *> &ExcludeBlocks) {
-  std::vector<Instruction *> RelevantInstructions;
-
-  for (BasicBlock &BB : F) {
-    if (ExcludeBlocks.contains(&BB))
-      continue;
-
-    for (Instruction &Inst : BB) {
-      if (Inst.isLifetimeStartOrEnd())
-        continue;
-      if (isa<DbgInfoIntrinsic>(Inst))
-        continue;
-
-      RelevantInstructions.push_back(&Inst);
-    }
-  }
-
-  return RelevantInstructions;
-}
-
 /// It is possible that there is a basic block that already performs the same
 /// stores. This returns a duplicate block, if it exists
 ///
 /// \param OutputBB [in] the block we are looking for a duplicate of.
 /// \param OutputStoreBBs [in] The existing output blocks.
 /// \returns an optional value with the number output block if there is a match.
 Optional<unsigned> findDuplicateOutputBlock(
     DenseMap<Value *, BasicBlock *> &OutputBBs,

 /// been replaced by a new output value.
 /// \param [in,out] OutputStoreBBs - The existing output blocks.
 static void alignOutputBlockWithAggFunc(
     OutlinableGroup &OG, OutlinableRegion &Region,
     DenseMap<Value *, BasicBlock *> &OutputBBs,
     DenseMap<Value *, BasicBlock *> &EndBBs,
     const DenseMap<Value *, Value *> &OutputMappings,
     std::vector<DenseMap<Value *, BasicBlock *>> &OutputStoreBBs) {
-  DenseSet<unsigned> ValuesToFind(Region.GVNStores.begin(),
-                                  Region.GVNStores.end());
-
-  // We iterate over the instructions in the extracted function, and find the
-  // global value number of the instructions.  If we find a value that should
-  // be contained in a store, we replace the uses of the value with the value
-  // from the overall function, so that the store is storing the correct
-  // value from the overall function.
-  DenseSet<BasicBlock *> ExcludeBBs;
-  for (DenseMap<Value *, BasicBlock *> &BBMap : OutputStoreBBs)
-    for (std::pair<Value *, BasicBlock *> &VBPair : BBMap)
-      ExcludeBBs.insert(VBPair.second);
-  for (std::pair<Value *, BasicBlock *> &VBPair : OutputBBs)
-    ExcludeBBs.insert(VBPair.second);
-
-  std::vector<Instruction *> ExtractedFunctionInsts =
-      collectRelevantInstructions(*(Region.ExtractedFunction), ExcludeBBs);
-  std::vector<Instruction *> OverallFunctionInsts =
-      collectRelevantInstructions(*OG.OutlinedFunction, ExcludeBBs);
-
-  assert(ExtractedFunctionInsts.size() == OverallFunctionInsts.size() &&
-         "Number of relevant instructions not equal!");
-
-  unsigned NumInstructions = ExtractedFunctionInsts.size();
-  for (unsigned Idx = 0; Idx < NumInstructions; Idx++) {
-    Value *V = ExtractedFunctionInsts[Idx];
-
-    if (OutputMappings.find(V) != OutputMappings.end())
-      V = OutputMappings.find(V)->second;
-    Optional<unsigned> GVN = Region.Candidate->getGVN(V);
-
-    // If we have found one of the stored values for output, replace the value
-    // with the corresponding one from the overall function.
-    if (GVN.hasValue() && ValuesToFind.erase(GVN.getValue())) {
-      V->replaceAllUsesWith(OverallFunctionInsts[Idx]);
-      if (ValuesToFind.size() == 0)
-        break;
-    }
-
-    if (ValuesToFind.size() == 0)
-      break;
-  }
-
-  assert(ValuesToFind.size() == 0 && "Not all store values were handled!");
-
   // If the size of the block is 0, then there are no stores, and we do not
   // need to save this block.
   bool AllRemoved = true;
   Value *RetValueForBB;
   BasicBlock *NewBB;
   SmallVector<Value *, 4> ToRemove;
   for (std::pair<Value *, BasicBlock *> &VtoBB : OutputBBs) {
     RetValueForBB = VtoBB.first;