[OpenMP] Set RequiresFullRuntime false in SPMDization

[ggeorgakoudis]

Summary

SPMDization in D102307 does not change the RequiresFullRuntime argument of kmpc_target_init/deinit calls. However, the constraints of SPMDization detection for converting a target region to SPMD mode should guarantee that the region does not require full runtime support. Hence, this patch sets RequiresFullRuntime to false for improved execution performance.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2841–2849

You need this callback stuff for the full-runtime flag too

[ggeorgakoudis]

Fix to call registerSimplificationCallback

[ggeorgakoudis]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2841–2849

Fixed

[jdoerfert]

There should be test changes.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2849

The callback you need here should return "the opposite" of IsSPMDModeSimplifyCB.

[ggeorgakoudis]

Fix the simplification callback

[jdoerfert]

All good, just update the spmd-zation tests (with the script)

[ggeorgakoudis]

Rebase

[ggeorgakoudis]

Update tests

[jdoerfert]

LG, one minor thing below.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2814

[jdoerfert]

Two test fail, check with upstream again wrt. remarks. Should not require a change to the remarks, even though an update might "improve" the check lines.

[ggeorgakoudis]

Rebase, fix for comments, update test

[ggeorgakoudis]

Fixing remark tests

[ggeorgakoudis]

Fixing remark tests

[ggeorgakoudis]

Fix remark tests

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

struct KernelInfoState : AbstractState {
   /// can be reached from the associated function.
   BooleanStateWithPtrSetVector<Function, /* InsertInvalidates */ false>
       ReachedKnownParallelRegions;
 
   /// State to track what parallel region we might reach.
   BooleanStateWithPtrSetVector<CallBase> ReachedUnknownParallelRegions;
 
   /// State to track if we are in SPMD-mode, assumed or know, and why we decided
-  /// we cannot be.
+  /// we cannot be. If it is assumed, then RequiresFullRuntime should also be
+  /// false.
   BooleanStateWithPtrSetVector<Instruction> SPMDCompatibilityTracker;
 
   /// The __kmpc_target_init call in this kernel, if any. If we find more than
   /// one we abort as the kernel is malformed.
   CallBase *KernelInitCB = nullptr;
 
   /// The __kmpc_target_deinit call in this kernel, if any. If we find more than
   /// one we abort as the kernel is malformed.
Attributor::SimplifictionCallbackTy IsSPMDModeSimplifyCB =
       } else {
         UsedAssumedInformation = false;
       }
       auto *Val = ConstantInt::getBool(IRP.getAnchorValue().getContext(),
                                        SPMDCompatibilityTracker.isAssumed());
       return Val;
     };
 
+    Attributor::SimplifictionCallbackTy IsGenericModeSimplifyCB =
+        [&](const IRPosition &IRP, const AbstractAttribute *AA,
+            bool &UsedAssumedInformation) -> Optional<Value *> {
+      // IRP represents the "RequiresFullRuntime" argument of an
+      // __kmpc_target_init or __kmpc_target_deinit call. We will answer this
+      // one with the internal state of the SPMDCompatibilityTracker, so if
+      // generic then true, if SPMD then false.
+      if (!isValidState())
+        return nullptr;
+      if (!SPMDCompatibilityTracker.isAtFixpoint()) {
+        if (AA)
+          A.recordDependence(*this, *AA, DepClassTy::OPTIONAL);
+        UsedAssumedInformation = true;
+      } else {
+        UsedAssumedInformation = false;
+      }
+      auto *Val = ConstantInt::getBool(IRP.getAnchorValue().getContext(),
+                                       !SPMDCompatibilityTracker.isAssumed());
+      return Val;
+    };
+
     constexpr const int InitIsSPMDArgNo = 1;
     constexpr const int DeinitIsSPMDArgNo = 1;
     constexpr const int InitUseStateMachineArgNo = 2;
+    constexpr const int InitRequiresFullRuntimeArgNo = 3;
+    constexpr const int DeinitRequiresFullRuntimeArgNo = 2;
     A.registerSimplificationCallback(
         IRPosition::callsite_argument(*KernelInitCB, InitUseStateMachineArgNo),
         StateMachineSimplifyCB);
     A.registerSimplificationCallback(
         IRPosition::callsite_argument(*KernelInitCB, InitIsSPMDArgNo),
         IsSPMDModeSimplifyCB);
     A.registerSimplificationCallback(
         IRPosition::callsite_argument(*KernelDeinitCB, DeinitIsSPMDArgNo),
         IsSPMDModeSimplifyCB);
+    A.registerSimplificationCallback(
+        IRPosition::callsite_argument(*KernelInitCB,
+                                      InitRequiresFullRuntimeArgNo),
+        IsGenericModeSimplifyCB);
+    A.registerSimplificationCallback(
+        IRPosition::callsite_argument(*KernelDeinitCB,
+                                      DeinitRequiresFullRuntimeArgNo),
+        IsGenericModeSimplifyCB);
 
     // Check if we know we are in SPMD-mode already.
     ConstantInt *IsSPMDArg =
         dyn_cast<ConstantInt>(KernelInitCB->getArgOperand(InitIsSPMDArgNo));
     if (IsSPMDArg && !IsSPMDArg->isZero())
       SPMDCompatibilityTracker.indicateOptimisticFixpoint();
   }
 
assert(ExecMode->getInitializer() &&
            "Initially non-SPMD kernel has SPMD exec mode!");
     ExecMode->setInitializer(
         ConstantInt::get(ExecMode->getInitializer()->getType(), 0));
 
     // Next rewrite the init and deinit calls to indicate we use SPMD-mode now.
     const int InitIsSPMDArgNo = 1;
     const int DeinitIsSPMDArgNo = 1;
     const int InitUseStateMachineArgNo = 2;
+    const int InitRequiresFullRuntimeArgNo = 3;
+    const int DeinitRequiresFullRuntimeArgNo = 2;
 
     auto &Ctx = getAnchorValue().getContext();
     A.changeUseAfterManifest(KernelInitCB->getArgOperandUse(InitIsSPMDArgNo),
                              *ConstantInt::getBool(Ctx, 1));
     A.changeUseAfterManifest(
         KernelInitCB->getArgOperandUse(InitUseStateMachineArgNo),
         *ConstantInt::getBool(Ctx, 0));
     A.changeUseAfterManifest(
         KernelDeinitCB->getArgOperandUse(DeinitIsSPMDArgNo),
         *ConstantInt::getBool(Ctx, 1));
+    A.changeUseAfterManifest(
+        KernelInitCB->getArgOperandUse(InitRequiresFullRuntimeArgNo),
+        *ConstantInt::getBool(Ctx, 0));
+    A.changeUseAfterManifest(
+        KernelDeinitCB->getArgOperandUse(DeinitRequiresFullRuntimeArgNo),
+        *ConstantInt::getBool(Ctx, 0));
+
     ++NumOpenMPTargetRegionKernelsSPMD;
 
     auto Remark = [&](OptimizationRemark OR) {
       return OR << "Generic-mode kernel is changed to SPMD-mode.";
     };
     A.emitRemark<OptimizationRemark>(KernelInitCB, "OpenMPKernelSPMDMode",
                                      Remark);
     return true;