[SandboxVec][DAG] Build actual dependencies (#111094)

This patch implements actual dependencies checking using BatchAA. This
adds memory dep edges between MemDGNodes.

llvm/include/llvm/SandboxIR/Utils.h

@@ -12,6 +12,7 @@
 #ifndef LLVM_SANDBOXIR_UTILS_H
 #define LLVM_SANDBOXIR_UTILS_H
 
+#include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/LoopAccessAnalysis.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/ScalarEvolution.h"
@@ -99,6 +100,13 @@ class Utils {
       return false;
     return *Diff > 0;
   }
+
+  /// Equivalent to BatchAA::getModRefInfo().
+  static ModRefInfo
+  aliasAnalysisGetModRefInfo(BatchAAResults &BatchAA, const Instruction *I,
+                             const std::optional<MemoryLocation> &OptLoc) {
+    return BatchAA.getModRefInfo(cast<llvm::Instruction>(I->Val), OptLoc);
+  }
 };
 
 } // namespace llvm::sandboxir

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h

@@ -24,6 +24,7 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/iterator_range.h"
+#include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/SandboxIR/Instruction.h"
 #include "llvm/SandboxIR/IntrinsicInst.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h"
@@ -47,6 +48,7 @@ class DGNode {
   // TODO: Use a PointerIntPair for SubclassID and I.
   /// For isa/dyn_cast etc.
   DGNodeID SubclassID;
+  // TODO: Move MemPreds to MemDGNode.
   /// Memory predecessors.
   DenseSet<MemDGNode *> MemPreds;
 
@@ -86,13 +88,20 @@ class DGNode {
            (!(II = dyn_cast<IntrinsicInst>(I)) || isMemIntrinsic(II));
   }
 
+  /// \Returns true if \p I is fence like. It excludes non-mem intrinsics.
+  static bool isFenceLike(Instruction *I) {
+    IntrinsicInst *II;
+    return I->isFenceLike() &&
+           (!(II = dyn_cast<IntrinsicInst>(I)) || isMemIntrinsic(II));
+  }
+
   /// \Returns true if \p I is a memory dependency candidate instruction.
   static bool isMemDepNodeCandidate(Instruction *I) {
     AllocaInst *Alloca;
     return isMemDepCandidate(I) ||
            ((Alloca = dyn_cast<AllocaInst>(I)) &&
             Alloca->isUsedWithInAlloca()) ||
-           isStackSaveOrRestoreIntrinsic(I);
+           isStackSaveOrRestoreIntrinsic(I) || isFenceLike(I);
   }
 
   Instruction *getInstruction() const { return I; }
@@ -159,8 +168,37 @@ class DependencyGraph {
   /// The DAG spans across all instructions in this interval.
   Interval<Instruction> DAGInterval;
 
+  std::unique_ptr<BatchAAResults> BatchAA;
+
+  enum class DependencyType {
+    RAW,   ///> Read After Write
+    WAW,   ///> Write After Write
+    RAR,   ///> Read After Read
+    WAR,   ///> Write After Read
+    CTRL,  ///> Control-related dependencies, like with PHIs/Terminators
+    OTHER, ///> Currently used for stack related instrs
+    NONE,  ///> No memory/other dependency
+  };
+  /// \Returns the dependency type depending on whether instructions may
+  /// read/write memory or whether they are some specific opcode-related
+  /// restrictions.
+  /// Note: It does not check whether a memory dependency is actually correct,
+  /// as it won't call AA. Therefore it returns the worst-case dep type.
+  static DependencyType getRoughDepType(Instruction *FromI, Instruction *ToI);
+
+  // TODO: Implement AABudget.
+  /// \Returns true if there is a memory/other dependency \p SrcI->DstI.
+  bool alias(Instruction *SrcI, Instruction *DstI, DependencyType DepType);
+
+  bool hasDep(sandboxir::Instruction *SrcI, sandboxir::Instruction *DstI);
+
+  /// Go through all mem nodes in \p SrcScanRange and try to add dependencies to
+  /// \p DstN.
+  void scanAndAddDeps(DGNode &DstN, const Interval<MemDGNode> &SrcScanRange);
+
 public:
-  DependencyGraph() {}
+  DependencyGraph(AAResults &AA)
+      : BatchAA(std::make_unique<BatchAAResults>(AA)) {}
 
   DGNode *getNode(Instruction *I) const {
     auto It = InstrToNodeMap.find(I);

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Interval.h

@@ -58,7 +58,7 @@ template <typename T, typename IntervalType> class IntervalIterator {
   }
   IntervalIterator &operator--() {
     // `I` is nullptr for end() when To is the BB terminator.
-    I = I != nullptr ? I->getPrevNode() : R.To;
+    I = I != nullptr ? I->getPrevNode() : R.bottom();
     return *this;
   }
   IntervalIterator operator--(int) {
@@ -110,14 +110,16 @@ template <typename T> class Interval {
   T *bottom() const { return To; }
 
   using iterator = IntervalIterator<T, Interval>;
-  using const_iterator = IntervalIterator<const T, const Interval>;
   iterator begin() { return iterator(From, *this); }
   iterator end() {
     return iterator(To != nullptr ? To->getNextNode() : nullptr, *this);
   }
-  const_iterator begin() const { return const_iterator(From, *this); }
-  const_iterator end() const {
-    return const_iterator(To != nullptr ? To->getNextNode() : nullptr, *this);
+  iterator begin() const {
+    return iterator(From, const_cast<Interval &>(*this));
+  }
+  iterator end() const {
+    return iterator(To != nullptr ? To->getNextNode() : nullptr,
+                    const_cast<Interval &>(*this));
   }
   /// Equality.
   bool operator==(const Interval &Other) const {

llvm/lib/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.cpp

@@ -8,8 +8,9 @@
 
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/SandboxIR/Utils.h"
 
-using namespace llvm::sandboxir;
+namespace llvm::sandboxir {
 
 #ifndef NDEBUG
 void DGNode::print(raw_ostream &OS, bool PrintDeps) const {
@@ -50,29 +51,135 @@ MemDGNodeIntervalBuilder::make(const Interval<Instruction> &Instrs,
                              cast<MemDGNode>(DAG.getNode(MemBotI)));
 }
 
+DependencyGraph::DependencyType
+DependencyGraph::getRoughDepType(Instruction *FromI, Instruction *ToI) {
+  // TODO: Perhaps compile-time improvement by skipping if neither is mem?
+  if (FromI->mayWriteToMemory()) {
+    if (ToI->mayReadFromMemory())
+      return DependencyType::RAW;
+    if (ToI->mayWriteToMemory())
+      return DependencyType::WAW;
+  } else if (FromI->mayReadFromMemory()) {
+    if (ToI->mayWriteToMemory())
+      return DependencyType::WAR;
+    if (ToI->mayReadFromMemory())
+      return DependencyType::RAR;
+  }
+  if (isa<sandboxir::PHINode>(FromI) || isa<sandboxir::PHINode>(ToI))
+    return DependencyType::CTRL;
+  if (ToI->isTerminator())
+    return DependencyType::CTRL;
+  if (DGNode::isStackSaveOrRestoreIntrinsic(FromI) ||
+      DGNode::isStackSaveOrRestoreIntrinsic(ToI))
+    return DependencyType::OTHER;
+  return DependencyType::NONE;
+}
+
+static bool isOrdered(Instruction *I) {
+  auto IsOrdered = [](Instruction *I) {
+    if (auto *LI = dyn_cast<LoadInst>(I))
+      return !LI->isUnordered();
+    if (auto *SI = dyn_cast<StoreInst>(I))
+      return !SI->isUnordered();
+    if (DGNode::isFenceLike(I))
+      return true;
+    return false;
+  };
+  bool Is = IsOrdered(I);
+  assert((!Is || DGNode::isMemDepCandidate(I)) &&
+         "An ordered instruction must be a MemDepCandidate!");
+  return Is;
+}
+
+bool DependencyGraph::alias(Instruction *SrcI, Instruction *DstI,
+                            DependencyType DepType) {
+  std::optional<MemoryLocation> DstLocOpt =
+      Utils::memoryLocationGetOrNone(DstI);
+  if (!DstLocOpt)
+    return true;
+  // Check aliasing.
+  assert((SrcI->mayReadFromMemory() || SrcI->mayWriteToMemory()) &&
+         "Expected a mem instr");
+  // TODO: Check AABudget
+  ModRefInfo SrcModRef =
+      isOrdered(SrcI)
+          ? ModRefInfo::Mod
+          : Utils::aliasAnalysisGetModRefInfo(*BatchAA, SrcI, *DstLocOpt);
+  switch (DepType) {
+  case DependencyType::RAW:
+  case DependencyType::WAW:
+    return isModSet(SrcModRef);
+  case DependencyType::WAR:
+    return isRefSet(SrcModRef);
+  default:
+    llvm_unreachable("Expected only RAW, WAW and WAR!");
+  }
+}
+
+bool DependencyGraph::hasDep(Instruction *SrcI, Instruction *DstI) {
+  DependencyType RoughDepType = getRoughDepType(SrcI, DstI);
+  switch (RoughDepType) {
+  case DependencyType::RAR:
+    return false;
+  case DependencyType::RAW:
+  case DependencyType::WAW:
+  case DependencyType::WAR:
+    return alias(SrcI, DstI, RoughDepType);
+  case DependencyType::CTRL:
+    // Adding actual dep edges from PHIs/to terminator would just create too
+    // many edges, which would be bad for compile-time.
+    // So we ignore them in the DAG formation but handle them in the
+    // scheduler, while sorting the ready list.
+    return false;
+  case DependencyType::OTHER:
+    return true;
+  case DependencyType::NONE:
+    return false;
+  }
+}
+
+void DependencyGraph::scanAndAddDeps(DGNode &DstN,
+                                     const Interval<MemDGNode> &SrcScanRange) {
+  assert(isa<MemDGNode>(DstN) &&
+         "DstN is the mem dep destination, so it must be mem");
+  Instruction *DstI = DstN.getInstruction();
+  // Walk up the instruction chain from ScanRange bottom to top, looking for
+  // memory instrs that may alias.
+  for (MemDGNode &SrcN : reverse(SrcScanRange)) {
+    Instruction *SrcI = SrcN.getInstruction();
+    if (hasDep(SrcI, DstI))
+      DstN.addMemPred(&SrcN);
+  }
+}
+
 Interval<Instruction> DependencyGraph::extend(ArrayRef<Instruction *> Instrs) {
   if (Instrs.empty())
     return {};
-  // TODO: For now create a chain of dependencies.
-  Interval<Instruction> Interval(Instrs);
-  auto *TopI = Interval.top();
-  auto *BotI = Interval.bottom();
-  DGNode *LastN = getOrCreateNode(TopI);
+
+  Interval<Instruction> InstrInterval(Instrs);
+
+  DGNode *LastN = getOrCreateNode(InstrInterval.top());
+  // Create DGNodes for all instrs in Interval to avoid future Instruction to
+  // DGNode lookups.
   MemDGNode *LastMemN = dyn_cast<MemDGNode>(LastN);
-  for (Instruction *I = TopI->getNextNode(), *E = BotI->getNextNode(); I != E;
-       I = I->getNextNode()) {
-    auto *N = getOrCreateNode(I);
-    N->addMemPred(LastMemN);
+  for (Instruction &I : drop_begin(InstrInterval)) {
+    auto *N = getOrCreateNode(&I);
     // Build the Mem node chain.
     if (auto *MemN = dyn_cast<MemDGNode>(N)) {
       MemN->setPrevNode(LastMemN);
       if (LastMemN != nullptr)
         LastMemN->setNextNode(MemN);
       LastMemN = MemN;
     }
-    LastN = N;
   }
-  return Interval;
+  // Create the dependencies.
+  auto DstRange = MemDGNodeIntervalBuilder::make(InstrInterval, *this);
+  for (MemDGNode &DstN : drop_begin(DstRange)) {
+    auto SrcRange = Interval<MemDGNode>(DstRange.top(), DstN.getPrevNode());
+    scanAndAddDeps(DstN, SrcRange);
+  }
+
+  return InstrInterval;
 }
 
 #ifndef NDEBUG
@@ -95,3 +202,5 @@ void DependencyGraph::dump() const {
   dbgs() << "\n";
 }
 #endif // NDEBUG
+
+} // namespace llvm::sandboxir