[Delinearization] Refactoring of fixed-size array delinearization

This is a follow-up patch to D122857 where we added delinearization of
fixed-size arrays to loop cache analysis, which resulted in some duplicate
code, i.e., "tryDelinearizeFixedSize()", in LoopCacheCost.cpp and
DependenceAnalysis.cpp. Refactoring is done in this patch.

This patch refactors out the main logic of "tryDelinearizeFixedSize()" as
"tryDelinearizeFixedSizeImpl()" and moves it to Delinearization.cpp, such that
clients can reuse "llvm::tryDelinearizeFixedSizeImpl()" wherever they would
like to delinearize fixed-size arrays. Currently it has two users, i.e.,
DependenceAnalysis.cpp and LoopCacheCost.cpp.

Reviewed By: Meinersbur, #loopoptwg

Differential Revision: https://reviews.llvm.org/D124745

Author: CongzheUalberta

Diff for: llvm/include/llvm/Analysis/Delinearization.h

@@ -125,6 +125,17 @@ bool getIndexExpressionsFromGEP(ScalarEvolution &SE,
                                 SmallVectorImpl<const SCEV *> &Subscripts,
                                 SmallVectorImpl<int> &Sizes);
 
+/// Implementation of fixed size array delinearization. Try to delinearize
+/// access function for a fixed size multi-dimensional array, by deriving
+/// subscripts from GEP instructions. Returns true upon success and false
+/// otherwise. \p Inst is the load/store instruction whose pointer operand is
+/// the one we want to delinearize. \p AccessFn is its corresponding SCEV
+/// expression w.r.t. the surrounding loop.
+bool tryDelinearizeFixedSizeImpl(ScalarEvolution *SE, Instruction *Inst,
+                                 const SCEV *AccessFn,
+                                 SmallVectorImpl<const SCEV *> &Subscripts,
+                                 SmallVectorImpl<int> &Sizes);
+
 struct DelinearizationPrinterPass
     : public PassInfoMixin<DelinearizationPrinterPass> {
   explicit DelinearizationPrinterPass(raw_ostream &OS);

Diff for: llvm/include/llvm/Analysis/DependenceAnalysis.h

@@ -927,9 +927,9 @@ namespace llvm {
     bool tryDelinearize(Instruction *Src, Instruction *Dst,
                         SmallVectorImpl<Subscript> &Pair);
 
-    /// Tries to delinearize access function for a fixed size multi-dimensional
-    /// array, by deriving subscripts from GEP instructions. Returns true upon
-    /// success and false otherwise.
+    /// Tries to delinearize \p Src and \p Dst access functions for a fixed size
+    /// multi-dimensional array. Calls tryDelinearizeFixedSizeImpl() to
+    /// delinearize \p Src and \p Dst separately,
     bool tryDelinearizeFixedSize(Instruction *Src, Instruction *Dst,
                                  const SCEV *SrcAccessFn,
                                  const SCEV *DstAccessFn,

Diff for: llvm/include/llvm/Analysis/LoopCacheAnalysis.h

@@ -98,9 +98,9 @@ class IndexedReference {
   /// Attempt to delinearize the indexed reference.
   bool delinearize(const LoopInfo &LI);
 
-  bool tryDelinearizeFixedSize(ScalarEvolution *SE, Instruction *Src,
-                               const SCEV *SrcAccessFn,
-                               SmallVectorImpl<const SCEV *> &SrcSubscripts);
+  /// Attempt to delinearize \p AccessFn for fixed-size arrays.
+  bool tryDelinearizeFixedSize(const SCEV *AccessFn,
+                               SmallVectorImpl<const SCEV *> &Subscripts);
 
   /// Return true if the index reference is invariant with respect to loop \p L.
   bool isLoopInvariant(const Loop &L) const;

Diff for: llvm/lib/Analysis/Delinearization.cpp

@@ -521,6 +521,44 @@ bool llvm::getIndexExpressionsFromGEP(ScalarEvolution &SE,
   return !Subscripts.empty();
 }
 
+bool llvm::tryDelinearizeFixedSizeImpl(
+    ScalarEvolution *SE, Instruction *Inst, const SCEV *AccessFn,
+    SmallVectorImpl<const SCEV *> &Subscripts, SmallVectorImpl<int> &Sizes) {
+  Value *SrcPtr = getLoadStorePointerOperand(Inst);
+
+  // Check the simple case where the array dimensions are fixed size.
+  auto *SrcGEP = dyn_cast<GetElementPtrInst>(SrcPtr);
+  if (!SrcGEP)
+    return false;
+
+  getIndexExpressionsFromGEP(*SE, SrcGEP, Subscripts, Sizes);
+
+  // Check that the two size arrays are non-empty and equal in length and
+  // value.
+  // TODO: it would be better to let the caller to clear Subscripts, similar
+  // to how we handle Sizes.
+  if (Sizes.empty() || Subscripts.size() <= 1) {
+    Subscripts.clear();
+    return false;
+  }
+
+  // Check that for identical base pointers we do not miss index offsets
+  // that have been added before this GEP is applied.
+  Value *SrcBasePtr = SrcGEP->getOperand(0)->stripPointerCasts();
+  const SCEVUnknown *SrcBase =
+      dyn_cast<SCEVUnknown>(SE->getPointerBase(AccessFn));
+  if (!SrcBase || SrcBasePtr != SrcBase->getValue()) {
+    Subscripts.clear();
+    return false;
+  }
+
+  assert(Subscripts.size() == Sizes.size() + 1 &&
+         "Expected equal number of entries in the list of size and "
+         "subscript.");
+
+  return true;
+}
+
 namespace {
 
 class Delinearization : public FunctionPass {

Diff for: llvm/lib/Analysis/DependenceAnalysis.cpp

@@ -3334,55 +3334,43 @@ bool DependenceInfo::tryDelinearize(Instruction *Src, Instruction *Dst,
   return true;
 }
 
+/// Try to delinearize \p SrcAccessFn and \p DstAccessFn if the underlying
+/// arrays accessed are fixed-size arrays. Return true if delinearization was
+/// successful.
 bool DependenceInfo::tryDelinearizeFixedSize(
     Instruction *Src, Instruction *Dst, const SCEV *SrcAccessFn,
     const SCEV *DstAccessFn, SmallVectorImpl<const SCEV *> &SrcSubscripts,
     SmallVectorImpl<const SCEV *> &DstSubscripts) {
-
-  Value *SrcPtr = getLoadStorePointerOperand(Src);
-  Value *DstPtr = getLoadStorePointerOperand(Dst);
   const SCEVUnknown *SrcBase =
       dyn_cast<SCEVUnknown>(SE->getPointerBase(SrcAccessFn));
   const SCEVUnknown *DstBase =
       dyn_cast<SCEVUnknown>(SE->getPointerBase(DstAccessFn));
   assert(SrcBase && DstBase && SrcBase == DstBase &&
          "expected src and dst scev unknowns to be equal");
 
-  // Check the simple case where the array dimensions are fixed size.
-  auto *SrcGEP = dyn_cast<GetElementPtrInst>(SrcPtr);
-  auto *DstGEP = dyn_cast<GetElementPtrInst>(DstPtr);
-  if (!SrcGEP || !DstGEP)
+  SmallVector<int, 4> SrcSizes;
+  SmallVector<int, 4> DstSizes;
+  if (!tryDelinearizeFixedSizeImpl(SE, Src, SrcAccessFn, SrcSubscripts,
+                                   SrcSizes) ||
+      !tryDelinearizeFixedSizeImpl(SE, Dst, DstAccessFn, DstSubscripts,
+                                   DstSizes))
     return false;
 
-  SmallVector<int, 4> SrcSizes, DstSizes;
-  getIndexExpressionsFromGEP(*SE, SrcGEP, SrcSubscripts, SrcSizes);
-  getIndexExpressionsFromGEP(*SE, DstGEP, DstSubscripts, DstSizes);
-
   // Check that the two size arrays are non-empty and equal in length and
   // value.
-  if (SrcSizes.empty() || SrcSubscripts.size() <= 1 ||
-      SrcSizes.size() != DstSizes.size() ||
+  if (SrcSizes.size() != DstSizes.size() ||
       !std::equal(SrcSizes.begin(), SrcSizes.end(), DstSizes.begin())) {
     SrcSubscripts.clear();
     DstSubscripts.clear();
     return false;
   }
 
-  Value *SrcBasePtr = SrcGEP->getOperand(0)->stripPointerCasts();
-  Value *DstBasePtr = DstGEP->getOperand(0)->stripPointerCasts();
-
-  // Check that for identical base pointers we do not miss index offsets
-  // that have been added before this GEP is applied.
-  if (SrcBasePtr != SrcBase->getValue() || DstBasePtr != DstBase->getValue()) {
-    SrcSubscripts.clear();
-    DstSubscripts.clear();
-    return false;
-  }
-
   assert(SrcSubscripts.size() == DstSubscripts.size() &&
-         SrcSubscripts.size() == SrcSizes.size() + 1 &&
-         "Expected equal number of entries in the list of sizes and "
-         "subscripts.");
+         "Expected equal number of entries in the list of SrcSubscripts and "
+         "DstSubscripts.");
+
+  Value *SrcPtr = getLoadStorePointerOperand(Src);
+  Value *DstPtr = getLoadStorePointerOperand(Dst);
 
   // In general we cannot safely assume that the subscripts recovered from GEPs
   // are in the range of values defined for their corresponding array
@@ -3418,8 +3406,8 @@ bool DependenceInfo::tryDelinearizeFixedSize(
   }
   LLVM_DEBUG({
     dbgs() << "Delinearized subscripts of fixed-size array\n"
-           << "SrcGEP:" << *SrcGEP << "\n"
-           << "DstGEP:" << *DstGEP << "\n";
+           << "SrcGEP:" << *SrcPtr << "\n"
+           << "DstGEP:" << *DstPtr << "\n";
   });
   return true;
 }

Diff for: llvm/lib/Analysis/LoopCacheAnalysis.cpp

@@ -349,46 +349,21 @@ CacheCostTy IndexedReference::computeRefCost(const Loop &L,
 }
 
 bool IndexedReference::tryDelinearizeFixedSize(
-    ScalarEvolution *SE, Instruction *Src, const SCEV *SrcAccessFn,
-    SmallVectorImpl<const SCEV *> &SrcSubscripts) {
-  Value *SrcPtr = getLoadStorePointerOperand(Src);
-  const SCEVUnknown *SrcBase =
-      dyn_cast<SCEVUnknown>(SE->getPointerBase(SrcAccessFn));
-
-  // Check the simple case where the array dimensions are fixed size.
-  auto *SrcGEP = dyn_cast<GetElementPtrInst>(SrcPtr);
-  if (!SrcGEP)
+    const SCEV *AccessFn, SmallVectorImpl<const SCEV *> &Subscripts) {
+  SmallVector<int, 4> ArraySizes;
+  if (!tryDelinearizeFixedSizeImpl(&SE, &StoreOrLoadInst, AccessFn, Subscripts,
+                                   ArraySizes))
     return false;
 
-  SmallVector<int, 4> SrcSizes;
-  getIndexExpressionsFromGEP(*SE, SrcGEP, SrcSubscripts, SrcSizes);
-
-  // Check that the two size arrays are non-empty and equal in length and
-  // value.
-  if (SrcSizes.empty() || SrcSubscripts.size() <= 1) {
-    SrcSubscripts.clear();
-    return false;
-  }
-
-  Value *SrcBasePtr = SrcGEP->getOperand(0)->stripPointerCasts();
-
-  // Check that for identical base pointers we do not miss index offsets
-  // that have been added before this GEP is applied.
-  if (SrcBasePtr != SrcBase->getValue()) {
-    SrcSubscripts.clear();
-    return false;
-  }
-
-  assert(SrcSubscripts.size() == SrcSizes.size() + 1 &&
-         "Expected equal number of entries in the list of size and "
-         "subscript.");
-
+  // Populate Sizes with scev expressions to be used in calculations later.
   for (auto Idx : seq<unsigned>(1, Subscripts.size()))
-    Sizes.push_back(SE->getConstant(Subscripts[Idx]->getType(), SrcSizes[Idx - 1]));
+    Sizes.push_back(
+        SE.getConstant(Subscripts[Idx]->getType(), ArraySizes[Idx - 1]));
 
   LLVM_DEBUG({
     dbgs() << "Delinearized subscripts of fixed-size array\n"
-           << "SrcGEP:" << *SrcGEP << "\n";
+           << "GEP:" << *getLoadStorePointerOperand(&StoreOrLoadInst)
+           << "\n";
   });
   return true;
 }
@@ -416,9 +391,9 @@ bool IndexedReference::delinearize(const LoopInfo &LI) {
 
     bool IsFixedSize = false;
     // Try to delinearize fixed-size arrays.
-    if (tryDelinearizeFixedSize(&SE, &StoreOrLoadInst, AccessFn, Subscripts)) {
+    if (tryDelinearizeFixedSize(AccessFn, Subscripts)) {
       IsFixedSize = true;
-      /// The last element of \p Sizes is the element size.
+      // The last element of Sizes is the element size.
       Sizes.push_back(ElemSize);
       LLVM_DEBUG(dbgs().indent(2) << "In Loop '" << L->getName()
                                   << "', AccessFn: " << *AccessFn << "\n");