[LoopDist] Add some runtime checks for cross partition loads - p2

Add back in getInstructionsForAccess as query for Instructions modifying
memory to mark cross partition accesses.

Author: mcberg2021

llvm/lib/Analysis/LoopAccessAnalysis.cpp

@@ -2404,12 +2404,14 @@ bool MemoryDepChecker::areDepsSafe(const DepCandidates &DepCands,
 SmallVector<Instruction *, 4>
 MemoryDepChecker::getInstructionsForAccess(Value *Ptr, bool IsWrite) const {
   MemAccessInfo Access(Ptr, IsWrite);
-  auto &IndexVector = Accesses.find(Access)->second;
-
+  auto I = Accesses.find(Access);
   SmallVector<Instruction *, 4> Insts;
-  transform(IndexVector,
-                 std::back_inserter(Insts),
-                 [&](unsigned Idx) { return this->InstMap[Idx]; });
+  if (I != Accesses.end()) {
+    transform(I->second,
+                   std::back_inserter(Insts),
+                   [&](unsigned Idx) { return this->InstMap[Idx]; });
+  }
+
   return Insts;
 }
 

llvm/lib/Transforms/Scalar/LoopDistribute.cpp

@@ -502,8 +502,10 @@ class InstPartitionContainer {
     SmallVector<int, 8> PtrToPartitions(N);
     for (unsigned I = 0; I < N; ++I) {
       Value *Ptr = RtPtrCheck->Pointers[I].PointerValue;
-      auto Instructions =
-          LAI.getInstructionsForAccess(Ptr, RtPtrCheck->Pointers[I].IsWritePtr);
+      auto Instructions = LAI.getInstructionsForAccess(Ptr, /* IsWrite */ true);
+      auto ReadInstructions =
+          LAI.getInstructionsForAccess(Ptr, /* IsWrite */ false);
+      Instructions.append(ReadInstructions.begin(), ReadInstructions.end());
 
       int &Partition = PtrToPartitions[I];
       // First set it to uninitialized.
@@ -521,21 +523,6 @@ class InstPartitionContainer {
           Partition = -1;
       }
       assert(Partition != -2 && "Pointer not belonging to any partition");
-      // All the store context uses of our address were processed,
-      // Now make sure we don't have cross partition loads.
-      if (RtPtrCheck->Pointers[I].IsWritePtr) {
-        if (Ptr->hasOneUse() || Partition == -1)
-          continue;
-
-        for (User *U : Ptr->users())
-          if (auto *CurLoad = dyn_cast<LoadInst>(U))
-            if (L->contains(CurLoad->getParent()))
-              if (Partition != (int)this->InstToPartitionId[CurLoad]) {
-                // -1 means belonging to multiple partitions.
-                Partition = -1;
-                break;
-              }
-      }
     }
 
     return PtrToPartitions;