[LoopVersioning] Get versionLoop to return a basic block with RT checks + allow user's RuntimePointerChecking in annotateLoopWithNoAlias and prepareNoAliasMetadata

[eastig]

Summary

• A user of LoopVersioning might need a basic block where RT checks have been created. This patch provides such functionality.
• LoopVersioning allows a user to provide own checks. Also LoopVersioning can annotate a loop with no-alias metadata but it uses RuntimePointerChecking provided by LoopAccessInfo. This RuntimePointerChecking must correspond to the checks provided by the user which is not always possible. This patch allows a user to provide RuntimePointerChecking which corresponds to the provided checks.

include/llvm/Transforms/Utils/LoopVersioning.h

public:
   ///
   /// This allows the loop transform pass to operate on the same loop regardless
   /// of whether versioning was necessary or not:
   ///
   ///    for each loop L:
   ///        analyze L
   ///        if versioning is necessary version L
   ///        transform L
-  void versionLoop() { versionLoop(findDefsUsedOutsideOfLoop(VersionedLoop)); }
+  ///
+  /// \return BasicBlock which contains runtime checks.
+  BasicBlock *versionLoop() {
+    return versionLoop(findDefsUsedOutsideOfLoop(VersionedLoop));
+  }
 
   /// \brief Same but if the client has already precomputed the set of values
   /// used outside the loop, this API will allows passing that.
-  void versionLoop(const SmallVectorImpl<Instruction *> &DefsUsedOutside);
+  ///
+  /// \return BasicBlock which contains runtime checks.
+  BasicBlock *versionLoop(const SmallVectorImpl<Instruction *> &DefsUsedOutside);
 
   /// \brief Returns the versioned loop.  Control flows here if pointers in the
   /// loop don't alias (i.e. all memchecks passed).  (This loop is actually the
   /// same as the original loop that we got constructed with.)
   Loop *getVersionedLoop() { return VersionedLoop; }
 
   /// \brief Returns the fall-back loop.  Control flows here if pointers in the
   /// loop may alias (i.e. one of the memchecks failed).
   Loop *getNonVersionedLoop() { return NonVersionedLoop; }
 
   /// \brief Sets the runtime alias checks for versioning the loop.
   void setAliasChecks(
       SmallVector<RuntimePointerChecking::PointerCheck, 4> Checks);
 
   /// \brief Sets the runtime SCEV checks for versioning the loop.
   void setSCEVChecks(SCEVUnionPredicate Check);
 
   /// \brief Annotate memory instructions in the versioned loop with no-alias
   /// metadata based on the memchecks issued.
   ///
   /// This is just wrapper that calls prepareNoAliasMetadata and
   /// annotateInstWithNoAlias on the instructions of the versioned loop.
-  void annotateLoopWithNoAlias();
+  void annotateLoopWithNoAlias(
+      const RuntimePointerChecking *RtPtrChecking = nullptr);
 
   /// \brief Set up the aliasing scopes based on the memchecks.  This needs to
   /// be called before the first call to annotateInstWithNoAlias.
-  void prepareNoAliasMetadata();
+  void
+  prepareNoAliasMetadata(const RuntimePointerChecking *RtPtrChecking = nullptr);
 
   /// \brief Add the noalias annotations to \p VersionedInst.
   ///
   /// \p OrigInst is the instruction corresponding to \p VersionedInst in the
   /// original loop.  Initialize the aliasing scopes with
   /// prepareNoAliasMetadata once before this can be called.
   void annotateInstWithNoAlias(Instruction *VersionedInst,
                                const Instruction *OrigInst);

lib/Transforms/Utils/LoopVersioning.cpp

void LoopVersioning::setAliasChecks(
     SmallVector<RuntimePointerChecking::PointerCheck, 4> Checks) {
   AliasChecks = std::move(Checks);
 }
 
 void LoopVersioning::setSCEVChecks(SCEVUnionPredicate Check) {
   Preds = std::move(Check);
 }
 
-void LoopVersioning::versionLoop(
+BasicBlock* LoopVersioning::versionLoop(
     const SmallVectorImpl<Instruction *> &DefsUsedOutside) {
   Instruction *FirstCheckInst;
   Instruction *MemRuntimeCheck;
   Value *SCEVRuntimeCheck;
   Value *RuntimeCheck = nullptr;
 
   // Add the memcheck in the original preheader (this is empty initially).
   BasicBlock *RuntimeCheckBB = VersionedLoop->getLoopPreheader();
BasicBlock* LoopVersioning::versionLoop(
 
   // The loops merge in the original exit block.  This is now dominated by the
   // memchecking block.
   DT->changeImmediateDominator(VersionedLoop->getExitBlock(), RuntimeCheckBB);
 
   // Adds the necessary PHI nodes for the versioned loops based on the
   // loop-defined values used outside of the loop.
   addPHINodes(DefsUsedOutside);
+
+  return RuntimeCheckBB;
 }
 
 void LoopVersioning::addPHINodes(
     const SmallVectorImpl<Instruction *> &DefsUsedOutside) {
   BasicBlock *PHIBlock = VersionedLoop->getExitBlock();
   assert(PHIBlock && "No single successor to loop exit block");
   PHINode *PN;
 
for (auto I = PHIBlock->begin(); (PN = dyn_cast<PHINode>(I)); ++I) {
     auto Mapped = VMap.find(ClonedValue);
     if (Mapped != VMap.end())
       ClonedValue = Mapped->second;
 
     PN->addIncoming(ClonedValue, NonVersionedLoop->getExitingBlock());
   }
 }
 
-void LoopVersioning::prepareNoAliasMetadata() {
+void LoopVersioning::prepareNoAliasMetadata(
+    const RuntimePointerChecking *RtPtrChecking) {
   // We need to turn the no-alias relation between pointer checking groups into
   // no-aliasing annotations between instructions.
   //
   // We accomplish this by mapping each pointer checking group (a set of
   // pointers memchecked together) to an alias scope and then also mapping each
   // group to the list of scopes it can't alias.
 
-  const RuntimePointerChecking *RtPtrChecking = LAI.getRuntimePointerChecking();
+  if (!RtPtrChecking)
+    RtPtrChecking = LAI.getRuntimePointerChecking();
+
   LLVMContext &Context = VersionedLoop->getHeader()->getContext();
 
   // First allocate an aliasing scope for each pointer checking group.
   //
   // While traversing through the checking groups in the loop, also create a
   // reverse map from pointers to the pointer checking group they were assigned
   // to.
   MDBuilder MDB(Context);
void LoopVersioning::prepareNoAliasMetadata(
 
   // Finally, transform the above to actually map to scope list which is what
   // the metadata uses.
 
   for (auto Pair : GroupToNonAliasingScopes)
     GroupToNonAliasingScopeList[Pair.first] = MDNode::get(Context, Pair.second);
 }
 
-void LoopVersioning::annotateLoopWithNoAlias() {
+void LoopVersioning::annotateLoopWithNoAlias(
+    const RuntimePointerChecking *RtPtrChecking) {
   if (!AnnotateNoAlias)
     return;
 
   // First prepare the maps.
-  prepareNoAliasMetadata();
+  prepareNoAliasMetadata(RtPtrChecking);
 
   // Add the scope and no-alias metadata to the instructions.
   for (Instruction *I : LAI.getDepChecker().getMemoryInstructions()) {
     annotateInstWithNoAlias(I);
   }
 }
 
 void LoopVersioning::annotateInstWithNoAlias(Instruction *VersionedInst,