[DA] batch delinearization

llvm/include/llvm/Analysis/Delinearization.h

@@ -16,16 +16,23 @@
 #ifndef LLVM_ANALYSIS_DELINEARIZATION_H
 #define LLVM_ANALYSIS_DELINEARIZATION_H
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Value.h"
+#include "llvm/Support/Compiler.h"
 
 namespace llvm {
+class Function;
 class raw_ostream;
 template <typename T> class SmallVectorImpl;
 class GetElementPtrInst;
 class Instruction;
+class LoopInfo;
 class ScalarEvolution;
 class SCEV;
+class SCEVUnknown;
 
 /// Compute the array dimensions Sizes from the set of Terms extracted from
 /// the memory access function of this SCEVAddRecExpr (second step of
@@ -164,6 +171,84 @@ bool getIndexExpressionsFromGEP(ScalarEvolution &SE,
                                 SmallVectorImpl<const SCEV *> &Subscripts,
                                 SmallVectorImpl<const SCEV *> &Sizes);
 
+/// BatchDelinearization - A wrapper for batch delinearization that caches
+/// results across multiple queries. Similar to BatchAAResults, this class
+/// should be used when analyzing multiple memory accesses to the same base
+/// pointers, as it computes array dimensions once using terms from all
+/// accesses, leading to better precision.
+///
+/// This class collects all memory accesses in a function, groups them by base
+/// pointer, and computes array dimensions for each base pointer using terms
+/// from all accesses. The results are cached for efficient lookups during
+/// dependence analysis.
+///
+/// Usage:
+///   BatchDelinearization BD(F, SE, LI);
+///   BD.populate();  // Compute and cache delinearization info.
+///   // Then pass BD to DependenceInfo or query it directly.
+class LLVM_ABI BatchDelinearization {
+public:
+  BatchDelinearization(Function &F, ScalarEvolution &SE, LoopInfo &LI)
+      : F(F), SE(SE), LI(LI) {}
+
+  /// Populate the cache with delinearization information for all memory
+  /// accesses in the function.
+  void populate();
+
+  /// Check if the cache has been populated.
+  bool isPopulated() const { return Populated; }
+
+  /// Get the cached array sizes for a base pointer.
+  /// Returns nullptr if not found.
+  const SmallVector<const SCEV *, 4> *
+  getArraySizes(const SCEVUnknown *Base) const {
+    auto It = ArraySizes.find(Base);
+    return It != ArraySizes.end() ? &It->second : nullptr;
+  }
+
+  /// Get the cached subscripts for an instruction.
+  /// Returns nullptr if not found.
+  const SmallVector<const SCEV *, 4> *
+  getSubscripts(const Instruction *I) const {
+    auto It = Subscripts.find(I);
+    return It != Subscripts.end() ? &It->second : nullptr;
+  }
+
+  /// Get the cached element size for a base pointer.
+  /// Returns nullptr if not found.
+  const SCEV *getElementSize(const SCEVUnknown *Base) const {
+    auto It = ElementSizes.find(Base);
+    return It != ElementSizes.end() ? It->second : nullptr;
+  }
+
+  /// Get the ScalarEvolution instance.
+  ScalarEvolution &getSE() { return SE; }
+  const ScalarEvolution &getSE() const { return SE; }
+
+  /// Get the LoopInfo instance.
+  LoopInfo &getLI() { return LI; }
+  const LoopInfo &getLI() const { return LI; }
+
+private:
+  Function &F;
+  ScalarEvolution &SE;
+  LoopInfo &LI;
+
+  /// Map from base pointer to computed array dimension sizes.
+  SmallDenseMap<const SCEVUnknown *, SmallVector<const SCEV *, 4>, 8>
+      ArraySizes;
+
+  /// Map from instruction to pre-computed subscripts.
+  SmallDenseMap<const Instruction *, SmallVector<const SCEV *, 4>, 16>
+      Subscripts;
+
+  /// Element size for the array (used for validation).
+  SmallDenseMap<const SCEVUnknown *, const SCEV *, 8> ElementSizes;
+
+  /// Flag indicating whether the cache has been populated.
+  bool Populated = false;
+};
+
 struct DelinearizationPrinterPass
     : public PassInfoMixin<DelinearizationPrinterPass> {
   explicit DelinearizationPrinterPass(raw_ostream &OS);

llvm/include/llvm/Analysis/DependenceAnalysis.h

@@ -49,6 +49,7 @@
 namespace llvm {
 class AAResults;
 template <typename T> class ArrayRef;
+class BatchDelinearization;
 class Loop;
 class LoopInfo;
 class SCEVConstant;
@@ -335,8 +336,9 @@ class LLVM_ABI FullDependence final : public Dependence {
 /// DependenceInfo - This class is the main dependence-analysis driver.
 class DependenceInfo {
 public:
-  DependenceInfo(Function *F, AAResults *AA, ScalarEvolution *SE, LoopInfo *LI)
-      : AA(AA), SE(SE), LI(LI), F(F) {}
+  DependenceInfo(Function *F, AAResults *AA, ScalarEvolution *SE, LoopInfo *LI,
+                 BatchDelinearization *BD = nullptr)
+      : AA(AA), SE(SE), LI(LI), F(F), BatchDelin(BD) {}
 
   /// Handle transitive invalidation when the cached analysis results go away.
   LLVM_ABI bool invalidate(Function &F, const PreservedAnalyses &PA,
@@ -355,11 +357,19 @@ class DependenceInfo {
 
   Function *getFunction() const { return F; }
 
+  /// setBatchDelinearization - Set the BatchDelinearization instance to use
+  /// for cached delinearization results.
+  void setBatchDelinearization(BatchDelinearization *BD) { BatchDelin = BD; }
+
+  /// getBatchDelinearization - Get the BatchDelinearization instance.
+  BatchDelinearization *getBatchDelinearization() const { return BatchDelin; }
+
 private:
   AAResults *AA;
   ScalarEvolution *SE;
   LoopInfo *LI;
   Function *F;
+  BatchDelinearization *BatchDelin;
 
   /// Subscript - This private struct represents a pair of subscripts from
   /// a pair of potentially multi-dimensional array references. We use a

llvm/lib/Analysis/DDG.cpp

@@ -10,6 +10,7 @@
 //===----------------------------------------------------------------------===//
 #include "llvm/Analysis/DDG.h"
 #include "llvm/ADT/SCCIterator.h"
+#include "llvm/Analysis/Delinearization.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopIterator.h"
 #include "llvm/Support/CommandLine.h"
@@ -308,7 +309,9 @@ bool DDGBuilder::shouldCreatePiBlocks() const { return CreatePiBlocks; }
 DDGAnalysis::Result DDGAnalysis::run(Loop &L, LoopAnalysisManager &AM,
                                      LoopStandardAnalysisResults &AR) {
   Function *F = L.getHeader()->getParent();
-  DependenceInfo DI(F, &AR.AA, &AR.SE, &AR.LI);
+  BatchDelinearization BD(*F, AR.SE, AR.LI);
+  BD.populate();
+  DependenceInfo DI(F, &AR.AA, &AR.SE, &AR.LI, &BD);
   return std::make_unique<DataDependenceGraph>(L, AR.LI, DI);
 }
 AnalysisKey DDGAnalysis::Key;

llvm/lib/Analysis/Delinearization.cpp

@@ -960,3 +960,140 @@ PreservedAnalyses DelinearizationPrinterPass::run(Function &F,
                        &AM.getResult<ScalarEvolutionAnalysis>(F));
   return PreservedAnalyses::all();
 }
+
+//===----------------------------------------------------------------------===//
+// BatchDelinearization Implementation
+//===----------------------------------------------------------------------===//
+
+/// Return true for a Load or Store instruction.
+static bool isLoadOrStore(const Instruction *I) {
+  return isa<LoadInst>(I) || isa<StoreInst>(I);
+}
+
+void BatchDelinearization::populate() {
+  if (Populated)
+    return;
+
+  Populated = true;
+
+  // Step 1: Collect all memory accesses grouped by base pointer.
+  // Map from base pointer to list of (Instruction, AccessFunction) pairs.
+  SmallDenseMap<const SCEVUnknown *,
+                SmallVector<std::pair<Instruction *, const SCEV *>, 4>, 8>
+      AccessesByBase;
+
+  for (Instruction &I : instructions(F)) {
+    if (!isLoadOrStore(&I))
+      continue;
+
+    Value *Ptr = getLoadStorePointerOperand(&I);
+    Loop *L = LI.getLoopFor(I.getParent());
+    const SCEV *AccessFn = SE.getSCEVAtScope(Ptr, L);
+    const SCEVUnknown *Base =
+        dyn_cast<SCEVUnknown>(SE.getPointerBase(AccessFn));
+
+    if (!Base)
+      continue;
+
+    // Only consider accesses where the base is loop invariant.
+    if (L && !SE.isLoopInvariant(Base, L))
+      continue;
+
+    AccessesByBase[Base].push_back({&I, AccessFn});
+  }
+
+  // Step 2: For each base pointer, collect terms from ALL accesses and
+  // compute array dimensions once.
+  for (auto &Entry : AccessesByBase) {
+    const SCEVUnknown *Base = Entry.first;
+    auto &Accesses = Entry.second;
+
+    // Skip if there's only one access - no benefit from batch processing.
+    if (Accesses.size() < 2)
+      continue;
+
+    // Determine element size - use the smallest among all accesses.
+    const SCEV *ElemSize = nullptr;
+    for (auto &Access : Accesses) {
+      const SCEV *EltSize = SE.getElementSize(Access.first);
+      if (!ElemSize)
+        ElemSize = EltSize;
+      else if (SE.isKnownPredicate(ICmpInst::ICMP_ULT, EltSize, ElemSize))
+        ElemSize = EltSize;
+    }
+
+    if (!ElemSize)
+      continue;
+
+    ElementSizes[Base] = ElemSize;
+
+    // Collect parametric terms from all accesses to this base.
+    SmallVector<const SCEV *, 8> Terms;
+    for (auto &Access : Accesses) {
+      const SCEV *AccessFn = Access.second;
+      const SCEV *OffsetSCEV = SE.getMinusSCEV(AccessFn, Base);
+      const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(OffsetSCEV);
+      if (AR && AR->isAffine())
+        collectParametricTerms(SE, AR, Terms);
+    }
+
+    // Find array dimensions using all collected terms.
+    SmallVector<const SCEV *, 4> Sizes;
+    findArrayDimensions(SE, Terms, Sizes, ElemSize);
+
+    // Skip if we couldn't determine dimensions.
+    if (Sizes.size() < 2)
+      continue;
+
+    ArraySizes[Base] = Sizes;
+
+    // Pre-compute subscripts for each access using parametric sizes.
+    for (auto &Access : Accesses) {
+      Instruction *Inst = Access.first;
+      const SCEV *AccessFn = Access.second;
+      const SCEV *OffsetSCEV = SE.getMinusSCEV(AccessFn, Base);
+      const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(OffsetSCEV);
+
+      if (!AR || !AR->isAffine())
+        continue;
+
+      SmallVector<const SCEV *, 4> Subs;
+      computeAccessFunctions(SE, AR, Subs, Sizes);
+
+      if (Subs.size() >= 2)
+        Subscripts[Inst] = std::move(Subs);
+    }
+  }
+
+  // Step 3: Try fixed-size array delinearization for accesses not yet cached.
+  // This handles arrays with known compile-time dimensions.
+  for (auto &Entry : AccessesByBase) {
+    auto &Accesses = Entry.second;
+
+    for (auto &Access : Accesses) {
+      Instruction *Inst = Access.first;
+
+      // Skip if already cached from parametric delinearization.
+      if (Subscripts.count(Inst))
+        continue;
+
+      const SCEV *AccessFn = Access.second;
+      const SCEV *ElemSize = SE.getElementSize(Inst);
+      SmallVector<const SCEV *, 4> Subs, Sizes;
+
+      if (delinearizeFixedSizeArray(SE, SE.removePointerBase(AccessFn), Subs,
+                                    Sizes, ElemSize) &&
+          Subs.size() >= 2) {
+        Subscripts[Inst] = std::move(Subs);
+      }
+    }
+  }
+
+  LLVM_DEBUG({
+    dbgs() << "Batch delinearization cache populated:\n";
+    dbgs() << "  Base pointers with cached dimensions: " << ArraySizes.size()
+           << "\n";
+    dbgs() << "  Instructions with cached subscripts: " << Subscripts.size()
+           << "\n";
+  });
+}