[llvm] Allow strided acceses to be grouped together for memory checks in loop accesses analysis

[SBallantyne]

Summary

This patch changes the groupChecks method to allow strided loop access over the same object to be grouped together.

This is targeting a particular issue from flang where certain loops would generate unneeded memory checks that lead to the generated vectorised code to never be taken, resulting in heavy performance losses.

Originally written by the co-authors.
Co-authored-by: David Sherwood <[e-mail]>
Co-authored-by: Maciej Gabka <[e-mail]>

[SBallantyne]

Add opt test

llvm/include/llvm/Analysis/LoopAccessAnalysis.h

struct PointerInfo {
     /// Holds the smallest byte address accessed by the pointer throughout all
     /// iterations of the loop.
     const SCEV *Start;
     /// Holds the largest byte address accessed by the pointer throughout all
     /// iterations of the loop, plus 1.
     const SCEV *End;
     /// Holds the information if this pointer is used for writing to memory.
     bool IsWritePtr;
+    /// Contains information regarding the stride of the access, which could be
+    /// 0 (unknown), 1 (a normal contiguous access), or abs(Stride) > 1 (this
+    /// is considered in the vectorizer as a strided access).
+    int Stride;
+    /// Loop Invariant
+    bool LoopInv;
     /// Holds the id of the set of pointers that could be dependent because of a
     /// shared underlying object.
     unsigned DependencySetId;
     /// Holds the id of the disjoint alias set to which this pointer belongs.
     unsigned AliasSetId;
     /// SCEV for the access.
     const SCEV *Expr;
     /// True if the pointer expressions needs to be frozen after expansion.
     bool NeedsFreeze;
 
     PointerInfo(Value *PointerValue, const SCEV *Start, const SCEV *End,
-                bool IsWritePtr, unsigned DependencySetId, unsigned AliasSetId,
-                const SCEV *Expr, bool NeedsFreeze)
+                bool IsWritePtr, int Stride, bool LoopInv,
+                unsigned DependencySetId, unsigned AliasSetId, const SCEV *Expr,
+                bool NeedsFreeze)
         : PointerValue(PointerValue), Start(Start), End(End),
-          IsWritePtr(IsWritePtr), DependencySetId(DependencySetId),
-          AliasSetId(AliasSetId), Expr(Expr), NeedsFreeze(NeedsFreeze) {}
+          IsWritePtr(IsWritePtr), Stride(Stride), LoopInv(LoopInv),
+          DependencySetId(DependencySetId), AliasSetId(AliasSetId), Expr(Expr),
+          NeedsFreeze(NeedsFreeze) {}
   };
 
   RuntimePointerChecking(MemoryDepChecker &DC, ScalarEvolution *SE)
       : DC(DC), SE(SE) {}
 
   /// Reset the state of the pointer runtime information.
   void reset() {
     Need = false;
     Pointers.clear();
     Checks.clear();
   }
 
   /// Insert a pointer and calculate the start and end SCEVs.
   /// We need \p PSE in order to compute the SCEV expression of the pointer
   /// according to the assumptions that we've made during the analysis.
   /// The method might also version the pointer stride according to \p Strides,
   /// and add new predicates to \p PSE.
   void insert(Loop *Lp, Value *Ptr, const SCEV *PtrExpr, Type *AccessTy,
-              bool WritePtr, unsigned DepSetId, unsigned ASId,
+              bool WritePtr, int Stride, unsigned DepSetId, unsigned ASId,
               PredicatedScalarEvolution &PSE, bool NeedsFreeze);
 
   /// No run-time memory checking is necessary.
   bool empty() const { return Pointers.empty(); }
 
+  /// All gathered accesses which require checks are strided
+  bool isStrided() const;
+
   /// Generate the checks and store it.  This also performs the grouping
   /// of pointers to reduce the number of memchecks necessary.
   void generateChecks(MemoryDepChecker::DepCandidates &DepCands,
                       bool UseDependencies);
 
   /// Returns the checks that generateChecks created. They can be used to ensure
   /// no read/write accesses overlap across all loop iterations.
   const SmallVectorImpl<RuntimePointerCheck> &getChecks() const {

llvm/lib/Analysis/LoopAccessAnalysis.cpp

 ///     N = (TripCount > 0) ? RoundDown(TripCount -1 , VF) : 0
 /// Start and End points are calculated in the following way:
 /// Start = UMIN(A, B) ; End = UMAX(A, B) + SizeOfElt,
 /// where SizeOfElt is the size of single memory access in bytes.
 ///
 /// There is no conflict when the intervals are disjoint:
 /// NoConflict = (P2.Start >= P1.End) || (P1.Start >= P2.End)
 void RuntimePointerChecking::insert(Loop *Lp, Value *Ptr, const SCEV *PtrExpr,
-                                    Type *AccessTy, bool WritePtr,
+                                    Type *AccessTy, bool WritePtr, int Stride,
                                     unsigned DepSetId, unsigned ASId,
                                     PredicatedScalarEvolution &PSE,
                                     bool NeedsFreeze) {
   ScalarEvolution *SE = PSE.getSE();
 
   const SCEV *ScStart;
   const SCEV *ScEnd;
 
+  bool LoopInv = false;
   if (SE->isLoopInvariant(PtrExpr, Lp)) {
+    LoopInv = true;
     ScStart = ScEnd = PtrExpr;
   } else {
     const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(PtrExpr);
     assert(AR && "Invalid addrec expression");
     const SCEV *Ex = PSE.getBackedgeTakenCount();
 
     ScStart = AR->getStart();
     ScEnd = AR->evaluateAtIteration(Ex, *SE);
if (SE->isLoopInvariant(PtrExpr, Lp)) {
     }
   }
   // Add the size of the pointed element to ScEnd.
   auto &DL = Lp->getHeader()->getModule()->getDataLayout();
   Type *IdxTy = DL.getIndexType(Ptr->getType());
   const SCEV *EltSizeSCEV = SE->getStoreSizeOfExpr(IdxTy, AccessTy);
   ScEnd = SE->getAddExpr(ScEnd, EltSizeSCEV);
 
-  Pointers.emplace_back(Ptr, ScStart, ScEnd, WritePtr, DepSetId, ASId, PtrExpr,
-                        NeedsFreeze);
+  Pointers.emplace_back(Ptr, ScStart, ScEnd, WritePtr, Stride, LoopInv,
+                        DepSetId, ASId, PtrExpr, NeedsFreeze);
 }
 
 void RuntimePointerChecking::tryToCreateDiffCheck(
     const RuntimeCheckingPtrGroup &CGI, const RuntimeCheckingPtrGroup &CGJ) {
   if (!CanUseDiffCheck)
     return;
 
   // If either group contains multiple different pointers, bail out.
void RuntimePointerChecking::groupChecks(
   // FoundNonConstantDistanceDependence is set, and therefore UseDependencies
   // is also false. In this case we will use the fallback path and create
   // separate checking groups for all pointers.
 
   // If we don't have the dependency partitions, construct a new
   // checking pointer group for each pointer. This is also required
   // for correctness, because in this case we can have checking between
   // pointers to the same underlying object.
-  if (!UseDependencies) {
+  // If we have strided accesses we want to always group them to emit checks.
+  // Otherwise we can have two strided pointers with constant offset
+  // being compared against each other for the same underlying object.
+  //
+  // However we can not merge loop invariants into any of the groups.
+  // Loop invariants need to be in separate groups to ensure that checks
+  // against all write accesses are emitted.
+  //
+  // In the example below:
+  // void foo(double _Complex a *a, unsigned int N, unsigned int n) {
+  //   unsigned int i = 0;
+  //   for (i = 0; i < N; i++) {
+  //     a[i] = a[n] + 1;
+  //   }
+  // }
+  //
+  // We have strided writes to: creal(a[i]) and cimag(a[i]) and
+  // reads from loop invariants: creal(a[n]) and creal(a[n]).
+  // If grouping is enabled we end up with just one group,
+  // and there is no checks emitted. It is leading to wrong analysis,
+  // as there is a dependency between a[n] and a[i], if n is within the range
+  // of iteration space.
+  // A memory access is defined as strided if either:
+  //  1. The stride is unknown (0), or
+  //  2. The stride is known and has an absolute value > 1. A stride of 1
+  //     is considered to be contiguous.
+  bool StridedOrInv = all_of(Pointers, [](const PointerInfo &P) {
+    return (std::abs(P.Stride) != 1 || P.LoopInv);
+  });
+
+  if (!StridedOrInv && !UseDependencies) {
     for (unsigned I = 0; I < Pointers.size(); ++I)
       CheckingGroups.push_back(RuntimeCheckingPtrGroup(I, *this));
     return;
   }
 
   unsigned TotalComparisons = 0;
 
   DenseMap<Value *, SmallVector<unsigned>> PositionMap;
if (isDependencyCheckNeeded()) {
       if (!LeaderId)
         LeaderId = RunningDepId++;
       DepId = LeaderId;
     } else
       // Each access has its own dependence set.
       DepId = RunningDepId++;
 
     bool IsWrite = Access.getInt();
-    RtCheck.insert(TheLoop, Ptr, PtrExpr, AccessTy, IsWrite, DepId, ASId, PSE,
-                   NeedsFreeze);
+    int Stride =
+        getPtrStride(PSE, AccessTy, Ptr, TheLoop, StridesMap, true).value_or(0);
+    RtCheck.insert(TheLoop, Ptr, PtrExpr, AccessTy, IsWrite, Stride, DepId,
+                   ASId, PSE, NeedsFreeze);
     LLVM_DEBUG(dbgs() << "LAA: Found a runtime check ptr:" << *Ptr << '\n');
   }
 
   return true;
 }
 
 bool AccessAnalysis::canCheckPtrAtRT(RuntimePointerChecking &RtCheck,
                                      ScalarEvolution *SE, Loop *TheLoop,

llvm/test/Analysis/LoopAccessAnalysis/number-of-memchecks.ll

for.body:                                         ; preds = %for.body, %entry
   store i16 %addres, ptr %storeidx, align 2
 
   %exitcond = icmp eq i64 %add, 5000
   br i1 %exitcond, label %for.end, label %for.body
 
 for.end:                                          ; preds = %for.body
   ret void
 }
+
+; Explanation taken from llvm/lib/Analysis/LoopAccessAnalysis:
+; If we have strided accesses we want to always group them to emit checks.
+; Otherwise we can have two strided pointers with constant offset
+; being compared against each other for the same underlying object.
+;
+; However we can not merge loop invariants into any of the groups.
+; Loop invariants need to be in separate groups to ensure that checks
+; against all write accesses are emitted.
+;
+; In the example below:
+; void foo(double _Complex a *a, unsigned int N, unsigned int n) {
+;   unsigned int i = 0;
+;   for (i = 0; i < N; i++) {
+;     a[i] = a[n] + 1;
+;   }
+; }
+;
+; We have strided writes to: creal(a[i]) and cimag(a[i]) and
+; reads from loop invariants: creal(a[n]) and creal(a[n]).
+; If grouping is enabled we end up with just one group,
+; and there is no checks emitted. It is leading to wrong analysis,
+; as there is a dependency between a[n] and a[i], if n is within the range
+; of iteration space.
+
+; CHECK: function 'testj':
+; CHECK: Run-time memory checks:
+; CHECK-NEXT:  Check 0:
+; CHECK-NEXT:      Comparing group ([[ZERO:.+]]):
+; CHECK-NEXT:        %gep.B = getelementptr { float, float }, ptr %a, i64 %stride.B
+; CHECK-NEXT:      Against group ([[ONE:.+]]):
+; CHECK-NEXT:        %gep.A = getelementptr { float, float }, ptr %a, i64 %stride.A
+; CHECK-NEXT:        %.elt = getelementptr { float, float }, ptr %a, i64 %stride.A, i32 1
+; CHECK-NEXT:        %gep.A = getelementptr { float, float }, ptr %a, i64 %stride.A
+; CHECK-NEXT:    Grouped accesses:
+; CHECK-NEXT:      Group {{.*}}[[ZERO]]:
+; CHECK-NEXT:        (Low: ((8 * %c * %ind.outer) + %a) High: (4 + ((8 + (8 * %ind.outer)) * %c) + %a))
+; CHECK-NEXT:          Member: {((8 * %c * %ind.outer) + %a),+,8}<%loop>
+; CHECK-NEXT:      Group {{.*}}[[ONE]]:
+; CHECK-NEXT:        (Low: ((8 * %b * %c) + %a) High: (8 + ((8 + (8 * %b)) * %c) + %a))
+; CHECK-NEXT:          Member: {((8 * %b * %c) + %a),+,8}<%loop>
+; CHECK-NEXT:          Member: {(4 + (8 * %b * %c) + %a),+,8}<%loop>
+; CHECK-NEXT:          Member: {((8 * %b * %c) + %a),+,8}<%loop>
+
+define void @testj(ptr %a,
+                   i64 %b,
+                   i64 %c) {
+entry:
+  br label %outer
+
+outer:                                            ; preds = %inner_crit_edge, %entry
+  %ind.outer = phi i64 [ 0, %entry ], [ %ind-pre, %inner_crit_edge ]
+  %cmp = icmp sgt i64 %c, 0
+  br i1 %cmp, label %inner, label %outer_crit_edge
+
+outer_crit_edge:                                  ; preds = %outer
+  %.pre = add nuw nsw i64 %ind.outer, 1
+  br label %inner_crit_edge
+
+inner:                                            ; preds = %outer
+  %p = mul nsw i64 %b, %c
+  %q = mul nuw nsw i64 %ind.outer, %c
+  br label %loop
+
+loop:                                             ; preds = %loop, %inner
+  %ind = phi i64 [ 0, %inner ], [ %ind.next, %loop ]
+
+  %stride.A = add nsw i64 %ind, %p
+  %gep.A = getelementptr { float, float }, ptr %a, i64 %stride.A
+  %.unpack = load float, ptr %gep.A, align 4
+  %.elt = getelementptr { float, float }, ptr %a, i64 %stride.A, i32 1
+  %fmul.A = fmul float %.unpack, %.unpack
+
+  %stride.B = add nuw nsw i64 %ind, %q
+  %gep.B = getelementptr { float, float }, ptr %a, i64 %stride.B
+  %.unpack1 = load float, ptr %gep.B, align 4
+  %res = fmul float %fmul.A, %.unpack1
+
+  store float %res, ptr %gep.A, align 4
+  store float %res, ptr %.elt, align 4
+
+  %ind.next = add nuw nsw i64 %ind, 1
+  %exitcond.not = icmp eq i64 %ind, %c
+  br i1 %exitcond.not, label %inner_crit_edge, label %loop
+
+inner_crit_edge:                                  ; preds = %loop, %outer_crit_edge
+  %ind-pre = phi i64 [ %.pre, %outer_crit_edge ], [ %b, %loop ]
+  %exitcond = icmp eq i64 %ind-pre, %b
+  br i1 %exitcond, label %end, label %outer
+
+end:                                              ; preds = %inner_crit_edge
+  ret void
+}