[OpenMP][FIX] Be more deliberate about invalidating the AAKernelInfo state

[jdoerfert]

Summary

This patch fixes a problem when the AAKernelInfo state was invalidated,
e.g., due to optnone for a kernel, but not all parts indicated the
invalidation properly. We further eliminate most full state
invalidations as they should never me necessary.

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

631–632

format

3045

Don't you expect anything else between in the future?

3349

From Attributor's perspective, they behave same, don't they?

3902–3903

IIRC, ReachedUnknownParallelRegions will be invalidated on insertion. With this change, this AA indicates optimistic fixed point?

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3045

I do not understand the question. Generally, if this is not proper in the future we can change it.
For now, we should return the right thing, that is unchanged if we did not do anything.

3349

Not if we look at AAKernelInfo as part of the manifest of another AA, which we might already do.
So folding might not be performed because AAKernelInfo is invalid while we only wanted to disable the state machine customization.

3902–3903

Yes. But optimistic fixpoint != everything is optimistic. It will move the known state to the assumed state. Given that we already invalidated ReachedUnknownParallelRegions before, it won't "un-invalidate" it, ever.

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3045

If we don't expect anything (any other manipulation) in the future between line 3006 and 3008, the change here is fine.

[tianshilei1992]

Sorry forget to accept the patch. Please fix the format before landing it.

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3902–3903

I'm fine as long as this is consistent. But in other hand, I feel it might be better to invalid the AA as well if it is using another state that has already been invalidated.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3902–3903

I don't follow. If an AA (X) uses another one (Y) and Y has an invalid state, X should not use the state of Y, ever. That is not a question nor is it changed here in any way. What is changed here is the fact that we do not need to give up on all information provided by AAKernelInfoCallSite if we see an __kmpc_omp_task call. We cannot argue SPMD and we cannot know all unknown parallel regions but we can certainly still know ReachingKernels (among other things). Indicating a pessimistic fixpoint here will give up on all information including stuff that we can use without problem.

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3902–3903

indicatePessimisticFixpoint invalidates the AA (X), which is AAKernelInfoCallSite here, as well. When another AA (Y) uses its member, such as ReachingKernels, if we only check the state of Y, it's gonna break. If we directly check the member instead, then we are good. That's what I mean by "consistent".

But yeah, I agree now that we could use partial information. However, if it is fixed point, no matter optimistic or pessimistic, it will not be updated, right? What if the partial information we want to use gets updated? Is it possible?

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3902–3903

A pessimistic fixpoint on AAKernelInfo should (almost) never be necessary. This patch removes (almost) all occurrences. A pessimistic fixpoint on AAKernelInfo will invalidate all information in the state, basically reverting everything to the previously known state. Since we track various things deduced in different ways, there is basically no good reason to invalidate everything at once. One could argue we should not track everything in a single state but that is a different discussion.

In your example, if X is invalidated it also invalidates reaching kernels. So Y will realize reaching kernels of X are unusable and deal with it accordingly (probably give up). However, the fact that this is a task runtime call did not impact the reaching kernels so there was no reason to invalidate it in the first place.

[JonChesterfield]

This doesn't seem to apply against main for me, please could someone more familiar with the patch rebase it?

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

struct KernelInfoState : AbstractState {
   bool isValidState() const override { return true; }
 
   /// See AbstractState::isAtFixpoint(...)
   bool isAtFixpoint() const override { return IsAtFixpoint; }
 
   /// See AbstractState::indicatePessimisticFixpoint(...)
   ChangeStatus indicatePessimisticFixpoint() override {
     IsAtFixpoint = true;
+    ReachingKernelEntries.indicatePessimisticFixpoint();
     SPMDCompatibilityTracker.indicatePessimisticFixpoint();
+    ReachedKnownParallelRegions.indicatePessimisticFixpoint();
     ReachedUnknownParallelRegions.indicatePessimisticFixpoint();
     return ChangeStatus::CHANGED;
   }
 
   /// See AbstractState::indicateOptimisticFixpoint(...)
   ChangeStatus indicateOptimisticFixpoint() override {
     IsAtFixpoint = true;
     return ChangeStatus::UNCHANGED;
struct KernelInfoState : AbstractState {
 
   /// Return full set as the worst state of potential values.
   static KernelInfoState getWorstState() { return KernelInfoState(false); }
 
   /// "Clamp" this state with \p KIS.
   KernelInfoState operator^=(const KernelInfoState &KIS) {
     // Do not merge two different _init and _deinit call sites.
     if (KIS.KernelInitCB) {
-      if (KernelInitCB && KernelInitCB != KIS.KernelInitCB)
-        indicatePessimisticFixpoint();
+      if(KernelInitCB && KernelInitCB != KIS.KernelInitCB)
+        llvm_unreachable("Kernel that calls another kernel violates OpenMP-Opt assumptions.");
       KernelInitCB = KIS.KernelInitCB;
     }
     if (KIS.KernelDeinitCB) {
-      if (KernelDeinitCB && KernelDeinitCB != KIS.KernelDeinitCB)
-        indicatePessimisticFixpoint();
+      if(KernelDeinitCB && KernelDeinitCB != KIS.KernelDeinitCB)
+        llvm_unreachable("Kernel that calls another kernel violates OpenMP-Opt assumptions.");
       KernelDeinitCB = KIS.KernelDeinitCB;
     }
     SPMDCompatibilityTracker ^= KIS.SPMDCompatibilityTracker;
     ReachedKnownParallelRegions ^= KIS.ReachedKnownParallelRegions;
     ReachedUnknownParallelRegions ^= KIS.ReachedUnknownParallelRegions;
     return *this;
   }
 
struct AAKernelInfo : public StateWrapper<KernelInfoState, AbstractAttribute> {
   const std::string getAsStr() const override {
     if (!isValidState())
       return "<invalid>";
     return std::string(SPMDCompatibilityTracker.isAssumed() ? "SPMD"
                                                             : "generic") +
            std::string(SPMDCompatibilityTracker.isAtFixpoint() ? " [FIX]"
                                                                : "") +
            std::string(" #PRs: ") +
-           std::to_string(ReachedKnownParallelRegions.size()) +
+           (ReachedKnownParallelRegions.isValidState()
+                ? std::to_string(ReachedKnownParallelRegions.size())
+                : "<invalid>") +
            ", #Unknown PRs: " +
-           std::to_string(ReachedUnknownParallelRegions.size()) +
+           (ReachedUnknownParallelRegions.isValidState()
+                ? std::to_string(ReachedUnknownParallelRegions.size())
+                : "<invalid>") +
            ", #Reaching Kernels: " +
            (ReachingKernelEntries.isValidState()
                 ? std::to_string(ReachingKernelEntries.size())
                 : "<invalid>");
   }
 
   /// Create an abstract attribute biew for the position \p IRP.
   static AAKernelInfo &createForPosition(const IRPosition &IRP, Attributor &A);
ChangeStatus manifest(Attributor &A) override {
 
     // Known SPMD-mode kernels need no manifest changes.
     if (SPMDCompatibilityTracker.isKnown())
       return ChangeStatus::UNCHANGED;
 
     // If we can we change the execution mode to SPMD-mode otherwise we build a
     // custom state machine.
     if (!mayContainParallelRegion() || !changeToSPMDMode(A))
-      buildCustomStateMachine(A);
+      return buildCustomStateMachine(A);
 
     return ChangeStatus::CHANGED;
   }
 
   bool changeToSPMDMode(Attributor &A) {
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
 
     if (!SPMDCompatibilityTracker.isAssumed()) {
bool changeToSPMDMode(Attributor &A) {
     };
     A.emitRemark<OptimizationRemark>(KernelInitCB, "OMP120", Remark);
     return true;
   };
 
   ChangeStatus buildCustomStateMachine(Attributor &A) {
     // If we have disabled state machine rewrites, don't make a custom one
     if (DisableOpenMPOptStateMachineRewrite)
-      return indicatePessimisticFixpoint();
+      return ChangeStatus::UNCHANGED;
 
     assert(ReachedKnownParallelRegions.isValidState() &&
            "Custom state machine with invalid parallel region states?");
 
     const int InitModeArgNo = 1;
     const int InitUseStateMachineArgNo = 2;
 
     // Check if the current configuration is non-SPMD and generic state machine.
auto CheckCallInst = [&](Instruction &I) {
           *this, IRPosition::callsite_function(CB), DepClassTy::OPTIONAL);
       getState() ^= CBAA.getState();
       AllSPMDStatesWereFixed &= CBAA.SPMDCompatibilityTracker.isAtFixpoint();
       return true;
     };
 
     bool UsedAssumedInformationInCheckCallInst = false;
     if (!A.checkForAllCallLikeInstructions(
-            CheckCallInst, *this, UsedAssumedInformationInCheckCallInst))
+            CheckCallInst, *this, UsedAssumedInformationInCheckCallInst)) {
+      LLVM_DEBUG(dbgs() << TAG << "Failed to visit all call-like instructions!\n";);
       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
case OMPRTL___kmpc_parallel_51:
       // The condition above should usually get the parallel region function
       // pointer and record it. In the off chance it doesn't we assume the
       // worst.
       ReachedUnknownParallelRegions.insert(&CB);
       break;
     case OMPRTL___kmpc_omp_task:
       // We do not look into tasks right now, just give up.
       SPMDCompatibilityTracker.insert(&CB);
       ReachedUnknownParallelRegions.insert(&CB);
-      indicatePessimisticFixpoint();
-      return;
+      break;
     case OMPRTL___kmpc_alloc_shared:
     case OMPRTL___kmpc_free_shared:
       // Return without setting a fixpoint, to be resolved in updateImpl.
       return;
     default:
       // Unknown OpenMP runtime calls cannot be executed in SPMD-mode,
-      // generally.
+      // generally. However, they do not hide parallel regions.
       SPMDCompatibilityTracker.insert(&CB);
-      indicatePessimisticFixpoint();
-      return;
+      break;
     }
     // All other OpenMP runtime calls will not reach parallel regions so they
     // can be safely ignored for now. Since it is a known OpenMP runtime call we
     // have now modeled all effects and there is no need for any update.
     indicateOptimisticFixpoint();
   }
 
   ChangeStatus updateImpl(Attributor &A) override {

llvm/test/Transforms/OpenMP/get_hardware_num_threads_in_block_fold_optnone.ll

;
   call void @helper0()
   call void @helper1()
   call void @helper2()
   call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
   ret void
 }
 
 define internal void @helper0() {
-; CHECK-LABEL: define {{[^@]+}}@helper0() {
-; CHECK-NEXT:    [[THREADLIMIT:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block() #[[ATTR1:[0-9]+]]
-; CHECK-NEXT:    store i32 undef, i32* @G, align 4
+; CHECK-LABEL: define {{[^@]+}}@helper0
+; CHECK-SAME: () #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT:    [[THREADLIMIT:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block() #[[ATTR1]]
+; CHECK-NEXT:    store i32 [[THREADLIMIT]], i32* @G, align 4
 ; CHECK-NEXT:    ret void
 ;
   %threadLimit = call i32 @__kmpc_get_hardware_num_threads_in_block()
   store i32 %threadLimit, i32* @G
   ret void
 }
 
 define internal void @helper1() {
 ; CHECK-LABEL: define {{[^@]+}}@helper1() {
 ; CHECK-NEXT:    [[THREADLIMIT:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block() #[[ATTR1]]
-; CHECK-NEXT:    unreachable
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[THREADLIMIT]], 666
+; CHECK-NEXT:    br i1 [[C]], label [[F:%.*]], label [[T:%.*]]
 ; CHECK:       t:
-; CHECK-NEXT:    unreachable
+; CHECK-NEXT:    call void @helper0() #[[ATTR1]]
+; CHECK-NEXT:    ret void
 ; CHECK:       f:
-; CHECK-NEXT:    unreachable
+; CHECK-NEXT:    ret void
 ;
   %threadLimit = call i32 @__kmpc_get_hardware_num_threads_in_block()
   %c = icmp eq i32 %threadLimit, 666
   br i1 %c, label %f, label %t
 t:
   call void @helper0()
   ret void
 f:
   ret void
 }
 
 define internal void @helper2() {
 ; CHECK-LABEL: define {{[^@]+}}@helper2() {
 ; CHECK-NEXT:    [[THREADLIMIT:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block() #[[ATTR1]]
-; CHECK-NEXT:    store i32 undef, i32* @G, align 4
+; CHECK-NEXT:    store i32 [[THREADLIMIT]], i32* @G, align 4
 ; CHECK-NEXT:    ret void
 ;
   %threadLimit = call i32 @__kmpc_get_hardware_num_threads_in_block()
   store i32 %threadLimit, i32* @G
   ret void
 }
 
 declare i32 @__kmpc_get_hardware_num_threads_in_block()