[OpenMP][Fix] Fix disable spmdization option

[ggeorgakoudis]

Summary

Besides SPMDization, other analysis and optimization for original, frontend-generated SPMD regions uses information from the AAKernelInfoFunction attribute. This fix makes sure disabling SPMDization through the corresponding option applies only to generic mode regions, which should not be SPMDized, while it leaves unaffected the attribute state of original SPMD regions.

[jdoerfert]

LG, thanks for tracking this down.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

struct AAKernelInfoFunction : AAKernelInfo {
 
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     // This is a high-level transform that might change the constant arguments
     // of the init and dinit calls. We need to tell the Attributor about this
     // to avoid other parts using the current constant value for simpliication.
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
 
-    // If we have disabled SPMD-ization, stop
-    if (DisableOpenMPOptSPMDization)
-      SPMDCompatibilityTracker.indicatePessimisticFixpoint();
-
     Function *Fn = getAnchorScope();
     if (!OMPInfoCache.Kernels.count(Fn))
       return;
 
     // Add itself to the reaching kernel and set IsKernelEntry.
     ReachingKernelEntries.insert(Fn);
     IsKernelEntry = true;
 
A.registerSimplificationCallback(
                                       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();
+    // This is a generic region but SPMDization is disabled so stop tracking.
+    else if (DisableOpenMPOptSPMDization)
+      SPMDCompatibilityTracker.indicatePessimisticFixpoint();
   }
 
   /// Modify the IR based on the KernelInfoState as the fixpoint iteration is
   /// finished now.
   ChangeStatus manifest(Attributor &A) override {
     // If we are not looking at a kernel with __kmpc_target_init and
     // __kmpc_target_deinit call we cannot actually manifest the information.
     if (!KernelInitCB || !KernelDeinitCB)