[OpenMP] Fix carefully track SPMDCompatibilityTracker

We did not properly use SPMDCompatibilityTracker in various places.
This patch makes sure we look at the validity properly and also fix
the state if we can.

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

diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index cc385b5..6862a3d 100644 (file)
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -2759,7 +2759,7 @@ struct AAKernelInfoFunction : AAKernelInfo {
       // __kmpc_target_init or
       // __kmpc_target_deinit call. We will answer this one with the internal
       // state.
-      if (!isValidState())
+      if (!SPMDCompatibilityTracker.isValidState())
         return nullptr;
       if (!SPMDCompatibilityTracker.isAtFixpoint()) {
         if (AA)
@@ -3162,20 +3162,22 @@ struct AAKernelInfoFunction : AAKernelInfo {
     };
 
     bool UsedAssumedInformationInCheckRWInst = false;
-    if (!A.checkForAllReadWriteInstructions(
-            CheckRWInst, *this, UsedAssumedInformationInCheckRWInst))
-      SPMDCompatibilityTracker.indicatePessimisticFixpoint();
+    if (!SPMDCompatibilityTracker.isAtFixpoint())
+      if (!A.checkForAllReadWriteInstructions(
+              CheckRWInst, *this, UsedAssumedInformationInCheckRWInst))
+        SPMDCompatibilityTracker.indicatePessimisticFixpoint();
 
     if (!IsKernelEntry)
       updateReachingKernelEntries(A);
 
     // Callback to check a call instruction.
+    bool AllSPMDStatesWereFixed = true;
     auto CheckCallInst = [&](Instruction &I) {
       auto &CB = cast<CallBase>(I);
       auto &CBAA = A.getAAFor<AAKernelInfo>(
           *this, IRPosition::callsite_function(CB), DepClassTy::OPTIONAL);
-      if (CBAA.getState().isValidState())
-        getState() ^= CBAA.getState();
+      getState() ^= CBAA.getState();
+      AllSPMDStatesWereFixed &= CBAA.SPMDCompatibilityTracker.isAtFixpoint();
       return true;
     };
 
@@ -3184,6 +3186,12 @@ struct AAKernelInfoFunction : AAKernelInfo {
             CheckCallInst, *this, UsedAssumedInformationInCheckCallInst))
       return indicatePessimisticFixpoint();
 
+    // If we haven't used any assumed information for the SPMD state we can fix
+    // it.
+    if (!UsedAssumedInformationInCheckRWInst &&
+        !UsedAssumedInformationInCheckCallInst && AllSPMDStatesWereFixed)
+      SPMDCompatibilityTracker.indicateOptimisticFixpoint();
+
     return StateBefore == getState() ? ChangeStatus::UNCHANGED
                                      : ChangeStatus::CHANGED;
   }