[LoopVersioning] Annotate versioned loop with noalias metadata

[anemet]

Summary

If we decide to version a loop to benefit a transformation, it makes
sense to record the now non-aliasing accesses in the newly versioned
loop. This allows non-aliasing information to be used by subsequent
passes.

One example is 456.hmmer in SPECint2006 where after loop distribution,
we vectorize one of the newly distributed loops. To vectorize we
version this loop to fully disambiguate may-aliasing accesses. If we
add the noalias markers, we can use the same information in a later DSE
pass to eliminate some dead stores which amounts to ~25% of the
instructions of this hot memory-pipeline-bound loop. The overall
performance improves by 18% on our ARM64.

The scoped noalias annotation is added in LoopVersioning. The patch
then enables this for loop distribution. A follow-on patch will enable
it for the vectorizer. Eventually this should be run by default when
versioning the loop but first I'd like to get some feedback whether my
understanding and application of scoped noalias metadata is correct.

Essentially my approach was to have a separate alias domain for each
versioning of the loop. For example, if we first version in loop
distribution and then in vectorization of the distributed loops, we have
a different set of memchecks for each versioning. By keeping the scopes
in different domains they can conveniently be defined independently
since different alias domains don't affect each other.

As written, I also have a separate domain for each loop. This is not
necessary and we could save some metadata here by using the same domain
across the different loops. I don't think it's a big deal either way.

Probably the best is to review the tests first to see if I mapped this
problem correctly to scoped noalias markers. I have plenty of comments
in the tests.

Note that the interface is prepared for the vectorizer which needs the
annotateInstWithNoAlias API. The vectorizer does not use LoopVersioning
so we need a way to pass in the versioned instructions. This is also
why the maps have to become part of the object state.

Also currently, we only have an AA-aware DSE after the vectorizer if we
also run the LTO pipeline. Depending how widely this triggers we may
want to schedule a DSE toward the end of the regular pass pipeline.

[anemet]

More testing revealed that the vectorizer part of this patch is incorrect.

It annotated the loop *before* versioning so the scalar loop also carried
these annotations. This is incorrect because we bail to the scalar loop if
the memchecks fail.

So I will add LoopVectorizer support for this separately in a follow-on patch.
I will also include a testcase for the above bug.

Also updating the title and the description accordingly.

[anemet]

Ping.

[anemet]

I implemented the LoopVectorizer part of this work under D17191. For the
vectorizer, the new annotation interface had to be tweaked. It now consists
of smaller functions which are probably easier to follow so I folded most of
those changes into this patch.

[anemet]

Ping 2.

[hfinkel]

LGTM, thanks!

test/Transforms/LoopVersioning/noalias-version-twice.ll

3 ↗

(On Diff #47777)

generate -> generated

llvm/trunk/include/llvm/Transforms/Utils/LoopVersioning.h

public:
 
   /// \brief Sets the runtime alias checks for versioning the loop.
   void setAliasChecks(
       const 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.
+  void annotateLoopWithNoAlias();
+
 private:
   /// \brief Adds the necessary PHI nodes for the versioned loops based on the
   /// loop-defined values used outside of the loop.
   ///
   /// This needs to be called after versionLoop if there are defs in the loop
   /// that are used outside the loop.
   void addPHINodes(const SmallVectorImpl<Instruction *> &DefsUsedOutside);
 
+  /// \brief Set up the aliasing scopes based on the memchecks.  This needs to
+  /// be called before the first call to annotateInstWithNoAlias.
+  void prepareNoAliasMetadata();
+
+  /// \brief Add the noalias annotations to \p I.  Initialize the aliasing
+  /// scopes with prepareNoAliasMetadata once before this can be called.
+  void annotateInstWithNoAlias(Instruction *I);
+
   /// \brief The original loop.  This becomes the "versioned" one.  I.e.,
   /// control flows here if pointers in the loop don't alias.
   Loop *VersionedLoop;
   /// \brief The fall-back loop.  I.e. control flows here if pointers in the
   /// loop may alias (memchecks failed).
   Loop *NonVersionedLoop;
 
   /// \brief This maps the instructions from VersionedLoop to their counterpart
   /// in NonVersionedLoop.
   ValueToValueMapTy VMap;
 
   /// \brief The set of alias checks that we are versioning for.
   SmallVector<RuntimePointerChecking::PointerCheck, 4> AliasChecks;
 
   /// \brief The set of SCEV checks that we are versioning for.
   SCEVUnionPredicate Preds;
 
+  /// \brief Maps a pointer to the pointer checking group that the pointer
+  /// belongs to.
+  DenseMap<const Value *, const RuntimePointerChecking::CheckingPtrGroup *>
+      PtrToGroup;
+
+  /// \brief The alias scope corresponding to a pointer checking group.
+  DenseMap<const RuntimePointerChecking::CheckingPtrGroup *, MDNode *>
+      GroupToScope;
+
+  /// \brief The list of alias scopes that a pointer checking group can't alias.
+  DenseMap<const RuntimePointerChecking::CheckingPtrGroup *, MDNode *>
+      GroupToNonAliasingScopeList;
+
   /// \brief Analyses used.
   const LoopAccessInfo &LAI;
   LoopInfo *LI;
   DominatorTree *DT;
   ScalarEvolution *SE;
 };
 }
 
 #endif

llvm/trunk/lib/Transforms/Scalar/LoopDistribute.cpp

bool processLoop(Loop *L) {
 
     if (!Pred.isAlwaysTrue() || !Checks.empty()) {
       DEBUG(dbgs() << "\nPointers:\n");
       DEBUG(LAI.getRuntimePointerChecking()->printChecks(dbgs(), Checks));
       LoopVersioning LVer(LAI, L, LI, DT, SE, false);
       LVer.setAliasChecks(std::move(Checks));
       LVer.setSCEVChecks(LAI.PSE.getUnionPredicate());
       LVer.versionLoop(DefsUsedOutside);
+      LVer.annotateLoopWithNoAlias();
     }
 
     // Create identical copies of the original loop for each partition and hook
     // them up sequentially.
     Partitions.cloneLoops();
 
     // Now, we remove the instruction from each loop that don't belong to that
     // partition.

llvm/trunk/lib/Transforms/Utils/LoopVersioning.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/LoopVersioning.h"
 #include "llvm/Analysis/LoopAccessAnalysis.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/ScalarEvolutionExpander.h"
 #include "llvm/IR/Dominators.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 
 using namespace llvm;
 
+static cl::opt<bool>
+    AnnotateNoAlias("loop-version-annotate-no-alias", cl::init(true),
+                    cl::Hidden,
+                    cl::desc("Add no-alias annotation for instructions that "
+                             "are disambiguated by memchecks"));
+
 LoopVersioning::LoopVersioning(const LoopAccessInfo &LAI, Loop *L, LoopInfo *LI,
                                DominatorTree *DT, ScalarEvolution *SE,
                                bool UseLAIChecks)
     : VersionedLoop(L), NonVersionedLoop(nullptr), LAI(LAI), LI(LI), DT(DT),
       SE(SE) {
   assert(L->getExitBlock() && "No single exit block");
   assert(L->getLoopPreheader() && "No preheader");
   if (UseLAIChecks) {
if (!PN) {
           User->replaceUsesOfWith(Inst, PN);
       PN->addIncoming(Inst, VersionedLoop->getExitingBlock());
     }
     // Add the new incoming value from the non-versioned loop.
     PN->addIncoming(NonVersionedLoopInst, NonVersionedLoop->getExitingBlock());
   }
 }
 
+void LoopVersioning::prepareNoAliasMetadata() {
+  // 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();
+  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);
+  MDNode *Domain = MDB.createAnonymousAliasScopeDomain("LVerDomain");
+
+  for (const auto &Group : RtPtrChecking->CheckingGroups) {
+    GroupToScope[&Group] = MDB.createAnonymousAliasScope(Domain);
+
+    for (unsigned PtrIdx : Group.Members)
+      PtrToGroup[RtPtrChecking->getPointerInfo(PtrIdx).PointerValue] = &Group;
+  }
+
+  // Go through the checks and for each pointer group, collect the scopes for
+  // each non-aliasing pointer group.
+  DenseMap<const RuntimePointerChecking::CheckingPtrGroup *,
+           SmallVector<Metadata *, 4>>
+      GroupToNonAliasingScopes;
+
+  for (const auto &Check : AliasChecks)
+    GroupToNonAliasingScopes[Check.first].push_back(GroupToScope[Check.second]);
+
+  // 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() {
+  if (!AnnotateNoAlias)
+    return;
+
+  // First prepare the maps.
+  prepareNoAliasMetadata();
+
+  // Add the scope and no-alias metadata to the instructions.
+  for (Instruction *I : LAI.getDepChecker().getMemoryInstructions()) {
+    annotateInstWithNoAlias(I);
+  }
+}
+
+void LoopVersioning::annotateInstWithNoAlias(Instruction *I) {
+  if (!AnnotateNoAlias)
+    return;
+
+  LLVMContext &Context = VersionedLoop->getHeader()->getContext();
+  Value *Ptr = isa<LoadInst>(I) ? cast<LoadInst>(I)->getPointerOperand()
+                                : cast<StoreInst>(I)->getPointerOperand();
+
+  // Find the group for the pointer and then add the scope metadata.
+  auto Group = PtrToGroup.find(Ptr);
+  if (Group != PtrToGroup.end()) {
+    I->setMetadata(
+        LLVMContext::MD_alias_scope,
+        MDNode::concatenate(I->getMetadata(LLVMContext::MD_alias_scope),
+                            MDNode::get(Context, GroupToScope[Group->second])));
+
+    // Add the no-alias metadata.
+    auto NonAliasingScopeList = GroupToNonAliasingScopeList.find(Group->second);
+    if (NonAliasingScopeList != GroupToNonAliasingScopeList.end())
+      I->setMetadata(
+          LLVMContext::MD_noalias,
+          MDNode::concatenate(I->getMetadata(LLVMContext::MD_noalias),
+                              NonAliasingScopeList->second));
+  }
+}
+
 namespace {
 /// \brief Also expose this is a pass.  Currently this is only used for
 /// unit-testing.  It adds all memchecks necessary to remove all may-aliasing
 /// array accesses from the loop.
 class LoopVersioningPass : public FunctionPass {
 public:
   LoopVersioningPass() : FunctionPass(ID) {
     initializeLoopVersioningPassPass(*PassRegistry::getPassRegistry());
bool runOnFunction(Function &F) override {
     // Now walk the identified inner loops.
     bool Changed = false;
     for (Loop *L : Worklist) {
       const LoopAccessInfo &LAI = LAA->getInfo(L, ValueToValueMap());
       if (LAI.getNumRuntimePointerChecks() ||
           !LAI.PSE.getUnionPredicate().isAlwaysTrue()) {
         LoopVersioning LVer(LAI, L, LI, DT, SE);
         LVer.versionLoop();
+        LVer.annotateLoopWithNoAlias();
         Changed = true;
       }
     }
 
     return Changed;
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {

llvm/trunk/test/Transforms/LoopVersioning/noalias-version-twice.ll

+; RUN: opt -basicaa -loop-distribute -scoped-noalias -loop-versioning -S < %s | FileCheck %s
+
+; Test the metadata generated when versioning an already versioned loop.  Here
+; we invoke loop distribution to perform the first round of versioning.  It
+; adds memchecks for accesses that can alias across the distribution boundary.
+; Then we further version the distributed loops to fully disambiguate accesses
+; within each.
+;
+; So as an example, we add noalias between C and A during the versioning
+; within loop distribution and then add noalias between C and D during the
+; second explicit versioning step:
+;
+;   for (i = 0; i < n; i++) {
+;     A[i + 1] = A[i] * B[i];
+; -------------------------------
+;     C[i] = D[i] * E[i];
+;   }
+
+; To see it easier what's going on, I expanded every noalias/scope metadata
+; reference below in a comment.  For a scope I use the format scope(domain),
+; e.g. scope 17 in domain 15 is written as 17(15).
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+
+@B = common global i32* null, align 8
+@A = common global i32* null, align 8
+@C = common global i32* null, align 8
+@D = common global i32* null, align 8
+@E = common global i32* null, align 8
+
+define void @f() {
+entry:
+  %a = load i32*, i32** @A, align 8
+  %b = load i32*, i32** @B, align 8
+  %c = load i32*, i32** @C, align 8
+  %d = load i32*, i32** @D, align 8
+  %e = load i32*, i32** @E, align 8
+  br label %for.body
+
+for.body:                                         ; preds = %for.body, %entry
+  %ind = phi i64 [ 0, %entry ], [ %add, %for.body ]
+
+  %arrayidxA = getelementptr inbounds i32, i32* %a, i64 %ind
+
+; CHECK: %loadA.ldist1 = {{.*}} !noalias !25
+; A noalias C: !25 -> { 17(15), 18(15), 19(15), 26(24) }
+;                       ^^^^^^
+  %loadA = load i32, i32* %arrayidxA, align 4
+
+  %arrayidxB = getelementptr inbounds i32, i32* %b, i64 %ind
+  %loadB = load i32, i32* %arrayidxB, align 4
+
+  %mulA = mul i32 %loadB, %loadA
+
+  %add = add nuw nsw i64 %ind, 1
+  %arrayidxA_plus_4 = getelementptr inbounds i32, i32* %a, i64 %add
+  store i32 %mulA, i32* %arrayidxA_plus_4, align 4
+
+; CHECK: for.body:
+
+  %arrayidxD = getelementptr inbounds i32, i32* %d, i64 %ind
+
+; CHECK: %loadD = {{.*}} !alias.scope !31
+; D's scope: !31 -> { 18(15), 32(33) }
+;                             ^^^^^^
+  %loadD = load i32, i32* %arrayidxD, align 4
+
+  %arrayidxE = getelementptr inbounds i32, i32* %e, i64 %ind
+
+; CHECK: %loadE = {{.*}} !alias.scope !34
+; E's scope: !34 -> { 19(15), 35(33) }
+;                             ^^^^^^
+  %loadE = load i32, i32* %arrayidxE, align 4
+
+  %mulC = mul i32 %loadD, %loadE
+
+  %arrayidxC = getelementptr inbounds i32, i32* %c, i64 %ind
+
+; CHECK: store i32 %mulC, {{.*}} !alias.scope !36, !noalias !38
+; C's scope: !36 -> { 17(15), 37(33) }
+;                     ^^^^^^
+; C noalias D and E: !38 -> { 21(15), 32(33), 35(33) }
+;                                     ^^^^^^  ^^^^^^
+  store i32 %mulC, i32* %arrayidxC, align 4
+
+  %exitcond = icmp eq i64 %add, 20
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+
+; Domain for the second loop versioning for the top loop after
+; distribution.
+; CHECK: !15 = distinct !{!15, !"LVerDomain"}
+; CHECK: !17 = distinct !{!17, !15}
+; CHECK: !25 = !{!17, !18, !19, !26}
+; CHECK: !31 = !{!18, !32}
+; CHECK: !32 = distinct !{!32, !33}
+; Domain for the second loop versioning for the bottom loop after
+; distribution.
+; CHECK: !33 = distinct !{!33, !"LVerDomain"}
+; CHECK: !34 = !{!19, !35}
+; CHECK: !35 = distinct !{!35, !33}
+; CHECK: !36 = !{!17, !37}
+; CHECK: !38 = !{!21, !32, !35}

llvm/trunk/test/Transforms/LoopVersioning/noalias.ll

+; RUN: opt -basicaa -loop-versioning -S < %s | FileCheck %s
+
+; A very simple case.  After versioning the %loadA and %loadB can't alias with
+; the store.
+;
+; To see it easier what's going on, I expanded every noalias/scope metadata
+; reference below in a comment.  For a scope I use the format scope(domain),
+; e.g. scope 17 in domain 15 is written as 17(15).
+
+; CHECK_LABEL: @f(
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+
+define void @f(i32* %a, i32* %b, i32* %c) {
+entry:
+  br label %for.body
+
+; CHECK: for.body.lver.orig:
+; CHECK: for.body:
+for.body:                                         ; preds = %for.body, %entry
+  %ind = phi i64 [ 0, %entry ], [ %add, %for.body ]
+
+  %arrayidxA = getelementptr inbounds i32, i32* %a, i64 %ind
+; CHECK: %loadA = {{.*}} !alias.scope !0
+; A's scope: !0 -> { 1(2) }
+  %loadA = load i32, i32* %arrayidxA, align 4
+
+  %arrayidxB = getelementptr inbounds i32, i32* %b, i64 %ind
+; CHECK: %loadB = {{.*}} !alias.scope !3
+; B's scope: !3 -> { 4(2) }
+  %loadB = load i32, i32* %arrayidxB, align 4
+
+  %mulC = mul i32 %loadA, %loadB
+
+  %arrayidxC = getelementptr inbounds i32, i32* %c, i64 %ind
+; CHECK: store {{.*}} !alias.scope !5, !noalias !7
+; C noalias A and B: !7 -> { 1(2), 4(2) }
+  store i32 %mulC, i32* %arrayidxC, align 4
+
+  %add = add nuw nsw i64 %ind, 1
+  %exitcond = icmp eq i64 %add, 20
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+; CHECK: !0 = !{!1}
+; CHECK: !1 = distinct !{!1, !2}
+; CHECK: !2 = distinct !{!2, !"LVerDomain"}
+; CHECK: !3 = !{!4}
+; CHECK: !4 = distinct !{!4, !2}
+; CHECK: !5 = !{!6}
+; CHECK: !6 = distinct !{!6, !2}
+; CHECK: !7 = !{!1, !4}