[LoopVectorize] Annotate versioned loop with noalias metadata

[anemet]

Summary

Use the new LoopVersioning facility (D16712) to add noalias metadata in
the vector loop if we versioned with memchecks. This can enable some
optimization opportunities further down the pipeline (see the included
test or the benchmark improvement quoted in D16712).

The test also covers the bug I had in the initial version in D16712.

The vectorizer did not previously use LoopVersioning. The reason is
that the vectorizer performs its transformations in single shot. It
creates an empty single-block vector loop that it then populates with
the widened, if-converted instructions. Thus creating an intermediate
versioned scalar loop seems wasteful.

So this patch (rather than bringing in LoopVersioning fully) adds a
special interface to LoopVersioning to allow the vectorizer to add
no-alias annotation while still performing its own versioning.

As the vectorizer propagates metadata from the instructions in the
original loop to the vector instructions we also check the pointer in
the original instruction and see if LoopVersioning can add no-alias
metadata based on the issued memchecks.

[anemet]

I forgot to clang-format

[mzolotukhin]

Hi Adam,

Please find some comments inline.

Thanks,
Michael

include/llvm/Transforms/Utils/LoopVersioning.h

92–96

Formatting looks weird here.

lib/Transforms/Vectorize/LoopVectorize.cpp

449–455

These functions are generally useful, not only in LoopVectorizer - e.g. we have their duplicates in SLPVectorizer. While you're at it, could you please move them out to a common place (and commit as a separate change)?

659–663

Does it belong here? I.e. should it really be a part of propagateMetadata? Probably that's fine with the current usage of propagateMetadata but it might become surprising if one decides to use this function somewhere else.

[anemet]

lib/Transforms/Vectorize/LoopVectorize.cpp

449–455

They are not really duplicates. The one in LV copies from instruction to one/many. The one in SLP copies many to one and merges them in the process in a metadata-specific way.

You can obviously still refactor the common parts or create a superset but I think that will probably be harder to read at the end.

What do you think?

659–663

Well, that depends on what you respond to the above but semantically, I do think it belongs here.

Here we're propagating the metadata from the "versioned" scalar loop (which is never created) to the vector loop. Let me know if this is not clear and I should add a comment.

[mzolotukhin]

lib/Transforms/Vectorize/LoopVectorize.cpp

449–455

That's true, but to me it looks like SLP version is just more general, where we need to merge several, potentially different sets of attributes. In LV case we can also think of such merging, but all the sets are the same, so the merge is trivial.

The reason I think it might be useful is that it's currently easy to forget to update both versions. E.g. when I added propagation of 'nontemporal' hints, I had to fix that in two very similar places.

659–663

Yeah, if you feel convinced that there is no reason in factoring propagateMetadata out, then it's fine to keep this code here I guess.

test/Transforms/LoopVectorize/noalias-md.ll

1–2

Don't you need -scoped-alias in the first command as well?

[anemet]

lib/Transforms/Vectorize/LoopVectorize.cpp

449–455

I would actually have a slight preference that for a new metadata each pass is carefully considered separately. Look at the if-conversion comment in the LV version as an example.

But I don't feel very strongly about this, so if you want you can factor them out.

test/Transforms/LoopVectorize/noalias-md.ll

1–2

No, -scoped-noalias is necessary to be able to query scoped noalias info. The first case only generates the metadata.

[mzolotukhin]

lib/Transforms/Vectorize/LoopVectorize.cpp

449–455

That sounds good to me.

[anemet]

Thanks, Michael!

[anemet]

Further performance analysis revealed that we were missing (valid)
optimizations cases compared to the original version in D16712.

Uniform loads are "scalarized" i.e. they are not vectorized by the original
instruction but cloned into the vector loop and then the splat value is
constructed (see the noalias-md-licm.ll test). Because of this
propagateMetadata was not invoked on them so we missed annotating these.

This new version splits annotation between cloned and newly created
instructions. As a side effect, the original propagateMetadata function is
unchanged which should help if we wanted to share this the SLP vectorizer.

Michael, still LGTY?

[mzolotukhin]

Yes, the change still looks fine.

include/llvm/Transforms/Utils/LoopVersioning.h

public:
   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.
+  ///
+  /// This is just wrapper that calls prepareNoAliasMetadata and
+  /// annotateInstWithNoAlias on the instructions of the versioned loop.
   void annotateLoopWithNoAlias();
 
+  /// \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 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);
+
 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);
+  void annotateInstWithNoAlias(Instruction *I) {
+    annotateInstWithNoAlias(I, 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;
 

lib/Transforms/Utils/LoopVersioning.cpp

void LoopVersioning::annotateLoopWithNoAlias() {
   prepareNoAliasMetadata();
 
   // Add the scope and no-alias metadata to the instructions.
   for (Instruction *I : LAI.getDepChecker().getMemoryInstructions()) {
     annotateInstWithNoAlias(I);
   }
 }
 
-void LoopVersioning::annotateInstWithNoAlias(Instruction *I) {
+void LoopVersioning::annotateInstWithNoAlias(Instruction *VersionedInst,
+                                             const Instruction *OrigInst) {
   LLVMContext &Context = VersionedLoop->getHeader()->getContext();
-  Value *Ptr = isa<LoadInst>(I) ? cast<LoadInst>(I)->getPointerOperand()
-                                : cast<StoreInst>(I)->getPointerOperand();
+  const Value *Ptr = isa<LoadInst>(OrigInst)
+                         ? cast<LoadInst>(OrigInst)->getPointerOperand()
+                         : cast<StoreInst>(OrigInst)->getPointerOperand();
 
   // Find the group for the pointer and then add the scope metadata.
   auto Group = PtrToGroup.find(Ptr);
   if (Group != PtrToGroup.end()) {
-    I->setMetadata(
+    VersionedInst->setMetadata(
         LLVMContext::MD_alias_scope,
-        MDNode::concatenate(I->getMetadata(LLVMContext::MD_alias_scope),
-                            MDNode::get(Context, GroupToScope[Group->second])));
+        MDNode::concatenate(
+            VersionedInst->getMetadata(LLVMContext::MD_alias_scope),
 
     // Add the no-alias metadata.
     auto NonAliasingScopeList = GroupToNonAliasingScopeList.find(Group->second);
     if (NonAliasingScopeList != GroupToNonAliasingScopeList.end())
-      I->setMetadata(
+      VersionedInst->setMetadata(
           LLVMContext::MD_noalias,
-          MDNode::concatenate(I->getMetadata(LLVMContext::MD_noalias),
-                              NonAliasingScopeList->second));
+          MDNode::concatenate(
+              VersionedInst->getMetadata(LLVMContext::MD_noalias),
   }
 }
 
 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 {

lib/Transforms/Vectorize/LoopVectorize.cpp

 #include "llvm/Pass.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/LoopVersioning.h"
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include <algorithm>
 #include <functional>
 #include <map>
 #include <tuple>
 
 using namespace llvm;
protected:
   /// Emit a bypass check to see if the vector trip count is nonzero.
   void emitVectorLoopEnteredCheck(Loop *L, BasicBlock *Bypass);
   /// Emit a bypass check to see if all of the SCEV assumptions we've
   /// had to make are correct.
   void emitSCEVChecks(Loop *L, BasicBlock *Bypass);
   /// Emit bypass checks to check any memory assumptions we may have made.
   void emitMemRuntimeChecks(Loop *L, BasicBlock *Bypass);
 
+  /// \brief Propagate known metadata from one instruction to another.
+  void propagateMetadata(Instruction *To, const Instruction *From);
+
+  /// \brief Propagate known metadata from one instruction to a vector of
+  /// others.
+  void propagateMetadata(SmallVectorImpl<Value *> &To, const Instruction *From);
+
   /// This is a helper class that holds the vectorizer state. It maps scalar
   /// instructions to vector instructions. When the code is 'unrolled' then
   /// then a single scalar value is mapped to multiple vector parts. The parts
   /// are stored in the VectorPart type.
   struct ValueMap {
     /// C'tor.  UnrollFactor controls the number of vectors ('parts') that
     /// are mapped.
     ValueMap(unsigned UnrollFactor) : UF(UnrollFactor) {}
protected:
   DominatorTree *DT;
   /// Alias Analysis.
   AliasAnalysis *AA;
   /// Target Library Info.
   const TargetLibraryInfo *TLI;
   /// Target Transform Info.
   const TargetTransformInfo *TTI;
 
+  /// \brief LoopVersioning.  It's only set up (non-null) if memchecks were
+  /// used.
+  ///
+  /// This is currently only used to add no-alias metadata based on the
+  /// memchecks.  The actually versioning is performed manually.
+  std::unique_ptr<LoopVersioning> LVer;
+
   /// The vectorization SIMD factor to use. Each vector will have this many
   /// vector elements.
   unsigned VF;
 
 protected:
   /// The vectorization unroll factor to use. Each scalar is vectorized to this
   /// many different vector instructions.
   unsigned UF;
else
       OS << L->getHeader()->getParent()->getParent()->getModuleIdentifier();
     OS.flush();
   }
   return Result;
 }
 #endif
 
 /// \brief Propagate known metadata from one instruction to another.
-static void propagateMetadata(Instruction *To, const Instruction *From) {
+void InnerLoopVectorizer::propagateMetadata(Instruction *To,
+                                            const Instruction *From) {
   SmallVector<std::pair<unsigned, MDNode *>, 4> Metadata;
   From->getAllMetadataOtherThanDebugLoc(Metadata);
 
   for (auto M : Metadata) {
     unsigned Kind = M.first;
 
     // These are safe to transfer (this is safe for TBAA, even when we
     // if-convert, because should that metadata have had a control dependency
     // on the condition, and thus actually aliased with some other
     // non-speculated memory access when the condition was false, this would be
     // caught by the runtime overlap checks).
     if (Kind != LLVMContext::MD_tbaa &&
         Kind != LLVMContext::MD_alias_scope &&
         Kind != LLVMContext::MD_noalias &&
         Kind != LLVMContext::MD_fpmath &&
         Kind != LLVMContext::MD_nontemporal)
       continue;
 
     To->setMetadata(Kind, M.second);
   }
+
+  // If the loop was versioned with memchecks, add the corresponding no-alias
+  // metadata.
+  if (LVer && (isa<LoadInst>(From) || isa<StoreInst>(From)))
+    LVer->annotateInstWithNoAlias(To, From);
 }
 
 /// \brief Propagate known metadata from one instruction to a vector of others.
-static void propagateMetadata(SmallVectorImpl<Value *> &To,
+void InnerLoopVectorizer::propagateMetadata(SmallVectorImpl<Value *> &To,
                                             const Instruction *From) {
   for (Value *V : To)
     if (Instruction *I = dyn_cast<Instruction>(V))
       propagateMetadata(I, From);
 }
 
 /// \brief The group of interleaved loads/stores sharing the same stride and
 /// close to each other.
 ///
void InnerLoopVectorizer::emitMemRuntimeChecks(Loop *L,
   // checks may query it before the current function is finished.
   DT->addNewBlock(NewBB, BB);
   if (L->getParentLoop())
     L->getParentLoop()->addBasicBlockToLoop(NewBB, *LI);
   ReplaceInstWithInst(BB->getTerminator(),
                       BranchInst::Create(Bypass, NewBB, MemRuntimeCheck));
   LoopBypassBlocks.push_back(BB);
   AddedSafetyChecks = true;
+
+  // We currently don't use LoopVersioning for the actual loop cloning but we
+  // still use it to add the noalias metadata.
+  LVer = llvm::make_unique<LoopVersioning>(*Legal->getLAI(), OrigLoop, LI, DT,
+                                           PSE.getSE());
+  LVer->prepareNoAliasMetadata();
 }
 
 
 void InnerLoopVectorizer::createEmptyLoop() {
   /*
    In this function we generate a new loop. The new loop will contain
    the vectorized instructions while the old loop will continue to run the
    scalar remainder.

test/Transforms/LoopVectorize/noalias-md.ll

+; RUN: opt -basicaa -loop-vectorize -force-vector-width=2 \
+; RUN:     -S < %s | FileCheck %s -check-prefix=BOTH -check-prefix=LV
+; RUN: opt -basicaa -scoped-noalias -loop-vectorize -dse -force-vector-width=2 \
+; RUN:     -S < %s | FileCheck %s -check-prefix=BOTH -check-prefix=DSE
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+
+; This loop needs to be versioned with memchecks between {A, B} x {C} before
+; it can be vectorized.
+;
+;   for (i = 0; i < n; i++) {
+;     C[i] = A[i] + 1;
+;     C[i] += B[i];
+;   }
+;
+; Check that the corresponding noalias metadata is added to the vector loop
+; but not to the scalar loop.
+;
+; Since in the versioned vector loop C and B can no longer alias, the first
+; store to C[i] can be DSE'd.
+
+
+define void @f(i32* %a, i32* %b, i32* %c) {
+entry:
+  br label %for.body
+
+; BOTH: vector.memcheck:
+; BOTH: vector.body:
+for.body:                                         ; preds = %for.body, %entry
+  %ind = phi i64 [ 0, %entry ], [ %inc, %for.body ]
+
+  %arrayidxA = getelementptr inbounds i32, i32* %a, i64 %ind
+; Scope 1
+; LV: = load {{.*}} !alias.scope !0
+  %loadA = load i32, i32* %arrayidxA, align 4
+
+  %add = add nuw i32 %loadA, 2
+
+  %arrayidxC = getelementptr inbounds i32, i32* %c, i64 %ind
+; Noalias with scope 1 and 6
+; LV: store {{.*}} !alias.scope !3, !noalias !5
+; DSE-NOT: store
+  store i32 %add, i32* %arrayidxC, align 4
+
+  %arrayidxB = getelementptr inbounds i32, i32* %b, i64 %ind
+; Scope 6
+; LV: = load {{.*}} !alias.scope !7
+  %loadB = load i32, i32* %arrayidxB, align 4
+
+  %add2 = add nuw i32 %add, %loadB
+
+; Noalias with scope 1 and 6
+; LV: store {{.*}} !alias.scope !3, !noalias !5
+; DSE: store
+  store i32 %add2, i32* %arrayidxC, align 4
+
+  %inc = add nuw nsw i64 %ind, 1
+  %exitcond = icmp eq i64 %inc, 20
+  br i1 %exitcond, label %for.end, label %for.body
+
+; BOTH: for.body:
+; BOTH-NOT: !alias.scope
+; BOTH-NOT: !noalias
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+
+; LV: !0 = !{!1}
+; LV: !1 = distinct !{!1, !2}
+; LV: !2 = distinct !{!2, !"LVerDomain"}
+; LV: !3 = !{!4}
+; LV: !4 = distinct !{!4, !2}
+; LV: !5 = !{!1, !6}
+; LV: !6 = distinct !{!6, !2}
+; LV: !7 = !{!6}