[OpenMP] Use AAHeapToStack/AAHeapToShared analysis in SPMDization

[ggeorgakoudis]

Summary

SPMDization D102307 detects incompatible OpenMP runtime calls to abort converting a target region to SPMD mode. Calls to memory allocation/de-allocation routines kmpc_alloc_shared, kmpc_free_shared are incompatible unless they are removed by AAHeapToStack/AAHeapToShared analysis. This patch extends SPMDization detection to include AAHeapToStack/AAHeapToShared analysis results for enlarging the scope of possible SPMDized regions detected.

[jhuber6]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2589

There should only ever be one free call per variable, it's like a stack. This is here just to make sure we don't accidentally remove something that's behind a PHI or something.

2641–2648

We could remove this if we track the free calls separately, but we'd need to have some mapping from the allocation to its corresponding free.

[jdoerfert]

Test ;)

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2589

Initially we have one free per alloc and the free uses the alloc, transformations can change that, e.g. introduce phi nodes or sink free into conditionals. We need to be conservative if the structure is not as we expect it.

3432–3435

One is sufficient, if the other becomes invalid we don't have to give up. In general, we should not give up on all of kernel info if SPMD-zation can't happen. Also, this is not even needed here ;)

3479–3481

3509

indicateOptimisticFixpoint is not correct above.
You use "assumed" information, thus it can change over time. If the information is known you could indicate a fixpoint. Instead, just don't add this call site to the SPMDCompatibilityTracker if we assume it is removed (by one of the h2s AAs).

[ggeorgakoudis]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2589

Got it

2641–2648

Agreed

[ggeorgakoudis]

Fix for comments

[jhuber6]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3478

For HeapToShared I used lookupAAFor because it returns null if the instance doesn't exist or is invalid. Not sure if that is an issue here.

[jdoerfert]

Test?

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3478

heap2shared doesn't depend on heap2stack, it could but it doesn't. Here we want getorcreate because there is a real dependence

[ggeorgakoudis]

Add tests

[jdoerfert]

Some minor things to address but LGTM.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

5660

Nit: early exit please.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2576

Add brief doxygen description please.

2589

Remove the todo, keep the check though.

2609–2613

2620

same as above.

3474

Nit: Add a message.

llvm/test/Transforms/OpenMP/spmdization.ll

68 ↗

(On Diff #360678)

Rebase and this is probably going to be a 2 not a 0 anymore.

[ggeorgakoudis]

Fix for comments

[GMNGeoffrey]

Looks like you've got some missing override markers:

llvm/lib/Transforms/IPO/OpenMPOpt.cpp:2588:8: error: 'isAssumedHeapToShared' overrides a member function but is not marked 'override' [-Werror,-Winconsistent-missing-override]
  bool isAssumedHeapToShared(CallBase &CB) const {
       ^
llvm/lib/Transforms/IPO/OpenMPOpt.cpp:2520:16: note: overridden virtual function is here
  virtual bool isAssumedHeapToShared(CallBase &CB) const = 0;
               ^
llvm/lib/Transforms/IPO/OpenMPOpt.cpp:2592:8: error: 'isAssumedHeapToSharedRemovedFree' overrides a member function but is not marked 'override' [-Werror,-Winconsistent-missing-override]
  bool isAssumedHeapToSharedRemovedFree(CallBase &CB) const {
       ^
llvm/lib/Transforms/IPO/OpenMPOpt.cpp:2524:16: note: overridden virtual function is here
  virtual bool isAssumedHeapToSharedRemovedFree(CallBase &CB) const = 0;
               ^
2 errors generated.

[jdoerfert]

In D105634#2898821, @GMNGeoffrey wrote:

Looks like you've got some missing override markers:

llvm/lib/Transforms/IPO/OpenMPOpt.cpp:2588:8: error: 'isAssumedHeapToShared' overrides a member function but is not marked 'override' [-Werror,-Winconsistent-missing-override]
  bool isAssumedHeapToShared(CallBase &CB) const {
       ^
llvm/lib/Transforms/IPO/OpenMPOpt.cpp:2520:16: note: overridden virtual function is here
  virtual bool isAssumedHeapToShared(CallBase &CB) const = 0;
               ^
llvm/lib/Transforms/IPO/OpenMPOpt.cpp:2592:8: error: 'isAssumedHeapToSharedRemovedFree' overrides a member function but is not marked 'override' [-Werror,-Winconsistent-missing-override]
  bool isAssumedHeapToSharedRemovedFree(CallBase &CB) const {
       ^
llvm/lib/Transforms/IPO/OpenMPOpt.cpp:2524:16: note: overridden virtual function is here
  virtual bool isAssumedHeapToSharedRemovedFree(CallBase &CB) const = 0;
               ^
2 errors generated.

I'll fix it real quick

[GMNGeoffrey]

In D105634#2898822, @jdoerfert wrote:

I'll fix it real quick

Thanks!

[ggeorgakoudis]

Thanks @GMNGeoffrey, fixed in https://github.com/llvm/llvm-project/commit/eaab880e45860192e5e52dba7f694f639f4dbcb7

llvm/include/llvm/Transforms/IPO/Attributor.h

 
 struct AAHeapToStack : public StateWrapper<BooleanState, AbstractAttribute> {
   using Base = StateWrapper<BooleanState, AbstractAttribute>;
   AAHeapToStack(const IRPosition &IRP, Attributor &A) : Base(IRP) {}
 
   /// Returns true if HeapToStack conversion is assumed to be possible.
   virtual bool isAssumedHeapToStack(CallBase &CB) const = 0;
 
+  /// Returns true if HeapToStack conversion is assumed and the CB is a
+  /// callsite to a free operation to be removed.
+  virtual bool isAssumedHeapToStackRemovedFree(CallBase &CB) const = 0;
+
   /// Create an abstract attribute view for the position \p IRP.
   static AAHeapToStack &createForPosition(const IRPosition &IRP, Attributor &A);
 
   /// See AbstractAttribute::getName()
   const std::string getName() const override { return "AAHeapToStack"; }
 
   /// See AbstractAttribute::getIdAddr()
   const char *getIdAddr() const override { return &ID; }

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

struct AAHeapToStackFunction final : public AAHeapToStack {
 
   bool isAssumedHeapToStack(CallBase &CB) const override {
     if (isValidState())
       if (AllocationInfo *AI = AllocationInfos.lookup(&CB))
         return AI->Status != AllocationInfo::INVALID;
     return false;
   }
 
+  bool isAssumedHeapToStackRemovedFree(CallBase &CB) const override {
+    if (isValidState())
+      for (auto &It : AllocationInfos) {
+        AllocationInfo &AI = *It.second;
+        if (AI.Status == AllocationInfo::INVALID)
+          continue;
+
+        if (AI.PotentialFreeCalls.count(&CB))
+          return true;
+      }
+
+    return false;
+  }
+
   ChangeStatus manifest(Attributor &A) override {
     assert(getState().isValidState() &&
            "Attempted to manifest an invalid state!");
 
     ChangeStatus HasChanged = ChangeStatus::UNCHANGED;
     Function *F = getAnchorScope();
     const auto *TLI = A.getInfoCache().getTargetLibraryInfoForFunction(*F);
 

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 struct AAHeapToShared : public StateWrapper<BooleanState, AbstractAttribute> {
   using Base = StateWrapper<BooleanState, AbstractAttribute>;
   AAHeapToShared(const IRPosition &IRP, Attributor &A) : Base(IRP) {}
 
   /// Create an abstract attribute view for the position \p IRP.
   static AAHeapToShared &createForPosition(const IRPosition &IRP,
                                            Attributor &A);
 
+  /// Returns true if HeapToShared conversion is assumed to be possible.
+  virtual bool isAssumedHeapToShared(CallBase &CB) const = 0;
+
+  /// Returns true if HeapToShared conversion is assumed and the CB is a
+  /// callsite to a free operation to be removed.
+  virtual bool isAssumedHeapToSharedRemovedFree(CallBase &CB) const = 0;
+
   /// See AbstractAttribute::getName().
   const std::string getName() const override { return "AAHeapToShared"; }
 
   /// See AbstractAttribute::getIdAddr().
   const char *getIdAddr() const override { return &ID; }
 
   /// This function should return true if the type of the \p AA is
   /// AAHeapToShared.
struct AAHeapToSharedFunction : public AAHeapToShared {
   const std::string getAsStr() const override {
     return "[AAHeapToShared] " + std::to_string(MallocCalls.size()) +
            " malloc calls eligible.";
   }
 
   /// See AbstractAttribute::trackStatistics().
   void trackStatistics() const override {}
 
+  void findPotentialRemovedFreeCalls(Attributor &A) {
+    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+    auto &FreeRFI = OMPInfoCache.RFIs[OMPRTL___kmpc_free_shared];
+
+    PotentialRemovedFreeCalls.clear();
+    // Update free call users of found malloc calls.
+    for (CallBase *CB : MallocCalls) {
+      SmallVector<CallBase *, 4> FreeCalls;
+      for (auto *U : CB->users()) {
+        CallBase *C = dyn_cast<CallBase>(U);
+        if (C && C->getCalledFunction() == FreeRFI.Declaration)
+          FreeCalls.push_back(C);
+      }
+      // TODO: Do we need to assume there is one, unique free user?
+      if (FreeCalls.size() != 1)
+        continue;
+
+      PotentialRemovedFreeCalls.insert(FreeCalls.front());
+    }
+  }
+
   void initialize(Attributor &A) override {
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
     auto &RFI = OMPInfoCache.RFIs[OMPRTL___kmpc_alloc_shared];
 
     for (User *U : RFI.Declaration->users())
       if (CallBase *CB = dyn_cast<CallBase>(U))
         MallocCalls.insert(CB);
+
+    findPotentialRemovedFreeCalls(A);
+  }
+
+  bool isAssumedHeapToShared(CallBase &CB) const {
+    if (isValidState())
+      if (MallocCalls.count(&CB))
+        return true;
+
+    return false;
+  }
+
+  bool isAssumedHeapToSharedRemovedFree(CallBase &CB) const {
+    if (isValidState())
+      if (PotentialRemovedFreeCalls.count(&CB))
+        return true;
+    return false;
   }
 
   ChangeStatus manifest(Attributor &A) override {
     if (MallocCalls.empty())
       return ChangeStatus::UNCHANGED;
 
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
     auto &FreeCall = OMPInfoCache.RFIs[OMPRTL___kmpc_free_shared];
 
     Function *F = getAnchorScope();
     auto *HS = A.lookupAAFor<AAHeapToStack>(IRPosition::function(*F), this,
                                             DepClassTy::OPTIONAL);
 
     ChangeStatus Changed = ChangeStatus::UNCHANGED;
     for (CallBase *CB : MallocCalls) {
       // Skip replacing this if HeapToStack has already claimed it.
       if (HS && HS->isAssumedHeapToStack(*CB))
         continue;
 
       // Find the unique free call to remove it.
       SmallVector<CallBase *, 4> FreeCalls;
       for (auto *U : CB->users()) {
         CallBase *C = dyn_cast<CallBase>(U);
         if (C && C->getCalledFunction() == FreeCall.Declaration)
           FreeCalls.push_back(C);
       }
       if (FreeCalls.size() != 1)
         continue;
 
       ConstantInt *AllocSize = dyn_cast<ConstantInt>(CB->getArgOperand(0));
 
       LLVM_DEBUG(dbgs() << TAG << "Replace globalization call in "
                         << CB->getCaller()->getName() << " with "
                         << AllocSize->getZExtValue()
                         << " bytes of shared memory\n");
 
for (User *U : RFI.Declaration->users()) {
       const auto &ED = A.getAAFor<AAExecutionDomain>(
           *this, IRPosition::function(*F), DepClassTy::REQUIRED);
       if (CallBase *CB = dyn_cast<CallBase>(U))
         if (!dyn_cast<ConstantInt>(CB->getArgOperand(0)) ||
             !ED.isExecutedByInitialThreadOnly(*CB))
           MallocCalls.erase(CB);
     }
 
+    findPotentialRemovedFreeCalls(A);
+
     if (NumMallocCalls != MallocCalls.size())
       return ChangeStatus::CHANGED;
 
     return ChangeStatus::UNCHANGED;
   }
 
   /// Collection of all malloc calls in a function.
   SmallPtrSet<CallBase *, 4> MallocCalls;
+  /// Collection of potentially removed free calls in a function.
+  SmallPtrSet<CallBase *, 4> PotentialRemovedFreeCalls;
 };
 
 struct AAKernelInfo : public StateWrapper<KernelInfoState, AbstractAttribute> {
   using Base = StateWrapper<KernelInfoState, AbstractAttribute>;
   AAKernelInfo(const IRPosition &IRP, Attributor &A) : Base(IRP) {}
 
   /// Statistics are tracked as part of manifest for now.
   void trackStatistics() const override {}
case OMPRTL___kmpc_parallel_51:
       // 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);
       break;
+    case OMPRTL___kmpc_alloc_shared:
+    case OMPRTL___kmpc_free_shared:
+      A.getAAFor<AAHeapToStack>(*this, IRPosition::function(*CB.getCaller()),
+                                DepClassTy::REQUIRED);
+      A.getAAFor<AAHeapToShared>(*this, IRPosition::function(*CB.getCaller()),
+                                 DepClassTy::REQUIRED);
+      // Return without setting a fixpoint, to be resolved in updateImpl.
+      return;
     default:
       // Unknown OpenMP runtime calls cannot be executed in SPMD-mode,
       // generally.
       SPMDCompatibilityTracker.insert(&CB);
       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 {
     // TODO: Once we have call site specific value information we can provide
     //       call site specific liveness information and then it makes
     //       sense to specialize attributes for call sites arguments instead of
     //       redirecting requests to the callee argument.
     Function *F = getAssociatedFunction();
+
+    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+    const auto &It = OMPInfoCache.RuntimeFunctionIDMap.find(F);
+
+    // If F is not a runtime function, propagate the AAKernelInfo of the callee.
+    if (It == OMPInfoCache.RuntimeFunctionIDMap.end()) {
       const IRPosition &FnPos = IRPosition::function(*F);
       auto &FnAA = A.getAAFor<AAKernelInfo>(*this, FnPos, DepClassTy::REQUIRED);
       if (getState() == FnAA.getState())
         return ChangeStatus::UNCHANGED;
       getState() = FnAA.getState();
       return ChangeStatus::CHANGED;
     }
+
+    // F is a runtime function that allocates or frees memory, check
+    // AAHeapToStack and AAHeapToShared.
+    KernelInfoState StateBefore = getState();
+    assert(It->getSecond() == OMPRTL___kmpc_alloc_shared ||
+           It->getSecond() == OMPRTL___kmpc_free_shared);
+
+    CallBase &CB = cast<CallBase>(getAssociatedValue());
+
+    auto &HeapToStackAA = A.getAAFor<AAHeapToStack>(
+        *this, IRPosition::function(*CB.getCaller()), DepClassTy::REQUIRED);
+    auto &HeapToSharedAA = A.getAAFor<AAHeapToShared>(
+        *this, IRPosition::function(*CB.getCaller()), DepClassTy::REQUIRED);
+
+    RuntimeFunction RF = It->getSecond();
+
+    switch (RF) {
+    case OMPRTL___kmpc_alloc_shared:
+      if (HeapToStackAA.isAssumedHeapToStack(CB))
+        indicateOptimisticFixpoint();
+      else if (HeapToSharedAA.isAssumedHeapToShared(CB))
+        indicateOptimisticFixpoint();
+      else {
+        SPMDCompatibilityTracker.insert(&CB);
+        indicatePessimisticFixpoint();
+      }
+      break;
+    case OMPRTL___kmpc_free_shared:
+      if (HeapToStackAA.isAssumedHeapToStackRemovedFree(CB))
+        indicateOptimisticFixpoint();
+      else if (HeapToSharedAA.isAssumedHeapToSharedRemovedFree(CB))
+        indicateOptimisticFixpoint();
+      else {
+        SPMDCompatibilityTracker.insert(&CB);
+        indicatePessimisticFixpoint();
+      }
+      break;
+    default:
+      SPMDCompatibilityTracker.insert(&CB);
+      indicatePessimisticFixpoint();
+    }
+
+    return StateBefore == getState() ? ChangeStatus::UNCHANGED
+                                     : ChangeStatus::CHANGED;
+  }
 };
 
 } // namespace
 
 void OpenMPOpt::registerAAs(bool IsModulePass) {
   if (SCC.empty())
 
     return;