[LV] Mask off possibly aliasing vector lanes

    When vectorising a loop that uses loads and stores, those pointers could
    overlap if their difference is less than the vector factor. For example,
    if address 20 is being stored to and address 23 is being loaded from, they
    overlap when the vector factor is 4 or higher. Currently LoopVectorize
    branches to a scalar loop in these cases with a runtime check. Howver if
    we construct a mask that disables the overlapping (aliasing) lanes then
    the vectorised loop can be safely entered, as long as the loads and
    stores are masked off.

Author: SamTebbs33

llvm/include/llvm/Analysis/LoopAccessAnalysis.h

@@ -491,11 +491,12 @@ struct PointerDiffInfo {
   const SCEV *SinkStart;
   unsigned AccessSize;
   bool NeedsFreeze;
+  bool WriteAfterRead;
 
   PointerDiffInfo(const SCEV *SrcStart, const SCEV *SinkStart,
-                  unsigned AccessSize, bool NeedsFreeze)
+                  unsigned AccessSize, bool NeedsFreeze, bool WriteAfterRead)
       : SrcStart(SrcStart), SinkStart(SinkStart), AccessSize(AccessSize),
-        NeedsFreeze(NeedsFreeze) {}
+        NeedsFreeze(NeedsFreeze), WriteAfterRead(WriteAfterRead) {}
 };
 
 /// Holds information about the memory runtime legality checks to verify

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -205,6 +205,13 @@ enum class TailFoldingStyle {
   DataWithEVL,
 };
 
+enum class RTCheckStyle {
+  /// Create runtime checks based on the difference between two pointers
+  ScalarDifference,
+  /// Form a mask based on elements which won't be a WAR or RAW hazard.
+  UseSafeEltsMask,
+};
+
 struct TailFoldingInfo {
   TargetLibraryInfo *TLI;
   LoopVectorizationLegality *LVL;
@@ -1357,6 +1364,11 @@ class TargetTransformInfo {
       PartialReductionExtendKind OpBExtend, std::optional<unsigned> BinOp,
       TTI::TargetCostKind CostKind) const;
 
+  /// \return true if a mask should be formed that disables lanes that could
+  /// alias between two pointers. The mask is created by the
+  /// loop_dependence_{war,raw}_mask intrinsics.
+  LLVM_ABI bool useSafeEltsMask() const;
+
   /// \return The maximum interleave factor that any transform should try to
   /// perform for this target. This number depends on the level of parallelism
   /// and the number of execution units in the CPU.

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -659,6 +659,8 @@ class TargetTransformInfoImplBase {
     return InstructionCost::getInvalid();
   }
 
+  virtual bool useSafeEltsMask() const { return false; }
+
   virtual unsigned getMaxInterleaveFactor(ElementCount VF) const { return 1; }
 
   virtual InstructionCost getArithmeticInstrCost(

llvm/include/llvm/Transforms/Utils/LoopUtils.h

@@ -614,6 +614,11 @@ addRuntimeChecks(Instruction *Loc, Loop *TheLoop,
                  const SmallVectorImpl<RuntimePointerCheck> &PointerChecks,
                  SCEVExpander &Expander, bool HoistRuntimeChecks = false);
 
+LLVM_ABI Value *addSafeEltsRuntimeChecks(Instruction *Loc,
+                                         ArrayRef<PointerDiffInfo> Checks,
+                                         SCEVExpander &Expander,
+                                         ElementCount VF);
+
 LLVM_ABI Value *addDiffRuntimeChecks(
     Instruction *Loc, ArrayRef<PointerDiffInfo> Checks, SCEVExpander &Expander,
     function_ref<Value *(IRBuilderBase &, unsigned)> GetVF, unsigned IC);

llvm/lib/Analysis/LoopAccessAnalysis.cpp

@@ -511,11 +511,14 @@ bool RuntimePointerChecking::tryToCreateDiffCheck(
     }
   }
 
+  bool WriteAfterRead = !Src->IsWritePtr && Sink->IsWritePtr;
+
   LLVM_DEBUG(dbgs() << "LAA: Creating diff runtime check for:\n"
                     << "SrcStart: " << *SrcStartInt << '\n'
                     << "SinkStartInt: " << *SinkStartInt << '\n');
   DiffChecks.emplace_back(SrcStartInt, SinkStartInt, AllocSize,
-                          Src->NeedsFreeze || Sink->NeedsFreeze);
+                          Src->NeedsFreeze || Sink->NeedsFreeze,
+                          WriteAfterRead);
   return true;
 }
 

llvm/lib/Analysis/TargetTransformInfo.cpp

@@ -878,6 +878,10 @@ InstructionCost TargetTransformInfo::getPartialReductionCost(
                                           BinOp, CostKind);
 }
 
+bool TargetTransformInfo::useSafeEltsMask() const {
+  return TTIImpl->useSafeEltsMask();
+}
+
 unsigned TargetTransformInfo::getMaxInterleaveFactor(ElementCount VF) const {
   return TTIImpl->getMaxInterleaveFactor(VF);
 }

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

@@ -1073,8 +1073,6 @@ AArch64TTIImpl::getIntrinsicInstrCost(const IntrinsicCostAttributes &ICA,
       return LegalCost;
     }
     break;
-  default:
-    break;
   }
   return BaseT::getIntrinsicInstrCost(ICA, CostKind);
 }
@@ -5880,6 +5878,11 @@ InstructionCost AArch64TTIImpl::getPartialReductionCost(
   return Cost + 2;
 }
 
+bool AArch64TTIImpl::useSafeEltsMask() const {
+  // The whilewr/rw instructions require SVE2
+  return ST->hasSVE2() || ST->hasSME();
+}
+
 InstructionCost
 AArch64TTIImpl::getShuffleCost(TTI::ShuffleKind Kind, VectorType *DstTy,
                                VectorType *SrcTy, ArrayRef<int> Mask,

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

@@ -406,6 +406,8 @@ class AArch64TTIImpl final : public BasicTTIImplBase<AArch64TTIImpl> {
       TTI::PartialReductionExtendKind OpBExtend, std::optional<unsigned> BinOp,
       TTI::TargetCostKind CostKind) const override;
 
+  bool useSafeEltsMask() const override;
+
   bool enableOrderedReductions() const override { return true; }
 
   InstructionCost getInterleavedMemoryOpCost(

llvm/lib/Transforms/Utils/LoopUtils.cpp

@@ -2135,6 +2135,58 @@ Value *llvm::addRuntimeChecks(
   return MemoryRuntimeCheck;
 }
 
+Value *llvm::addSafeEltsRuntimeChecks(Instruction *Loc,
+                                      ArrayRef<PointerDiffInfo> Checks,
+                                      SCEVExpander &Expander, ElementCount VF) {
+  IRBuilder ChkBuilder(Loc->getContext(),
+                       InstSimplifyFolder(Loc->getDataLayout()));
+  ChkBuilder.SetInsertPoint(Loc);
+  Value *MemoryRuntimeCheck = nullptr;
+
+  // Map to keep track of created compares, The key is the pair of operands for
+  // the compare, to allow detecting and re-using redundant compares.
+  DenseMap<std::pair<Value *, Value *>, Value *> SeenCompares;
+  Value *AliasLaneMask = nullptr;
+  for (const auto &[SrcStart, SinkStart, AccessSize, NeedsFreeze,
+                    WriteAfterRead] : Checks) {
+    Type *Ty = SinkStart->getType();
+    Value *Sink = Expander.expandCodeFor(SinkStart, Ty, Loc);
+    Value *Src = Expander.expandCodeFor(SrcStart, Ty, Loc);
+    if (SeenCompares.lookup({Sink, Src}))
+      continue;
+
+    unsigned IntOpc = WriteAfterRead ? Intrinsic::loop_dependence_war_mask
+                                     : Intrinsic::loop_dependence_raw_mask;
+    Value *SourceAsPtr = ChkBuilder.CreateCast(Instruction::IntToPtr, Src,
+                                               ChkBuilder.getPtrTy());
+    Value *SinkAsPtr = ChkBuilder.CreateCast(Instruction::IntToPtr, Sink,
+                                             ChkBuilder.getPtrTy());
+    Value *M = ChkBuilder.CreateIntrinsic(
+        IntOpc, {VectorType::get(ChkBuilder.getInt1Ty(), VF)},
+        {SourceAsPtr, SinkAsPtr, ChkBuilder.getInt64(AccessSize)}, nullptr,
+        "alias.lane.mask");
+    SeenCompares.insert({{Sink, Src}, M});
+    if (AliasLaneMask)
+      M = ChkBuilder.CreateAnd(AliasLaneMask, M);
+    else
+      AliasLaneMask = M;
+  }
+  assert(AliasLaneMask && "Expected an alias lane mask to have been created.");
+  auto *VecVT = VectorType::get(ChkBuilder.getInt1Ty(), VF);
+  // Extend to an i8 since i1 is too small to add with
+  Value *PopCount = ChkBuilder.CreateCast(
+      Instruction::ZExt, AliasLaneMask,
+      VectorType::get(ChkBuilder.getInt8Ty(), VecVT->getElementCount()));
+
+  PopCount =
+      ChkBuilder.CreateUnaryIntrinsic(Intrinsic::vector_reduce_add, PopCount);
+  PopCount = ChkBuilder.CreateCast(Instruction::ZExt, PopCount,
+                                   ChkBuilder.getInt64Ty());
+  MemoryRuntimeCheck = ChkBuilder.CreateICmpUGT(
+      PopCount, ConstantInt::get(ChkBuilder.getInt64Ty(), 0));
+  return MemoryRuntimeCheck;
+}
+
 Value *llvm::addDiffRuntimeChecks(
     Instruction *Loc, ArrayRef<PointerDiffInfo> Checks, SCEVExpander &Expander,
     function_ref<Value *(IRBuilderBase &, unsigned)> GetVF, unsigned IC) {
@@ -2149,7 +2201,8 @@ Value *llvm::addDiffRuntimeChecks(
   // Map to keep track of created compares, The key is the pair of operands for
   // the compare, to allow detecting and re-using redundant compares.
   DenseMap<std::pair<Value *, Value *>, Value *> SeenCompares;
-  for (const auto &[SrcStart, SinkStart, AccessSize, NeedsFreeze] : Checks) {
+  for (const auto &[SrcStart, SinkStart, AccessSize, NeedsFreeze,
+                    WriteAfterRead] : Checks) {
     Type *Ty = SinkStart->getType();
     // Compute VF * IC * AccessSize.
     auto *VFTimesICTimesSize =
@@ -2158,8 +2211,8 @@ Value *llvm::addDiffRuntimeChecks(
     Value *Diff =
         Expander.expandCodeFor(SE.getMinusSCEV(SinkStart, SrcStart), Ty, Loc);
 
-    // Check if the same compare has already been created earlier. In that case,
-    // there is no need to check it again.
+    // Check if the same compare has already been created earlier. In that
+    // case, there is no need to check it again.
     Value *IsConflict = SeenCompares.lookup({Diff, VFTimesICTimesSize});
     if (IsConflict)
       continue;

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

@@ -506,7 +506,11 @@ class LoopVectorizationPlanner {
   /// Build VPlans for the specified \p UserVF and \p UserIC if they are
   /// non-zero or all applicable candidate VFs otherwise. If vectorization and
   /// interleaving should be avoided up-front, no plans are generated.
-  void plan(ElementCount UserVF, unsigned UserIC);
+  /// DiffChecks is a list of pointer pairs that should be checked for aliasing,
+  /// combining the resulting predicate with an active lane mask if one is in
+  /// use.
+  void plan(ElementCount UserVF, unsigned UserIC,
+            std::optional<ArrayRef<PointerDiffInfo>> DiffChecks);
 
   /// Use the VPlan-native path to plan how to best vectorize, return the best
   /// VF and its cost.