[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/Transforms/Scalar/LoopDistribute.cpp

@@ -504,6 +504,12 @@ class InstPartitionContainer {
       Value *Ptr = RtPtrCheck->Pointers[I].PointerValue;
       auto Instructions =
           LAI.getInstructionsForAccess(Ptr, RtPtrCheck->Pointers[I].IsWritePtr);
+      const MemoryDepChecker &DC = LAI.getDepChecker();
+      bool IsWritePtr = !RtPtrCheck->Pointers[I].IsWritePtr;
+      if (!DC.getOrderForAccess(Ptr, IsWritePtr).empty()) {
+        auto AltInstructions = LAI.getInstructionsForAccess(Ptr, IsWritePtr);
+        Instructions.append(AltInstructions.begin(), AltInstructions.end());
+      }
 
       int &Partition = PtrToPartitions[I];
       // First set it to uninitialized.
@@ -521,21 +527,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;