[flang][OpenMP] Fix handling of nested loop wrappers in LowerWorkshare (#117275)

Author: ivanradanov

flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp

@@ -126,9 +126,9 @@ static bool mustParallelizeOp(Operation *op) {
         //         omp.workshare.loop_wrapper {}
         //
         // Therefore, we skip if we encounter a nested omp.workshare.
-        if (isa<omp::WorkshareOp>(op))
+        if (isa<omp::WorkshareOp>(nested))
           return WalkResult::skip();
-        if (isa<omp::WorkshareLoopWrapperOp>(op))
+        if (isa<omp::WorkshareLoopWrapperOp>(nested))
           return WalkResult::interrupt();
         return WalkResult::advance();
       })
@@ -253,8 +253,7 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
               // Either we have already remapped it
               bool remapped = rootMapping.contains(opr);
               // Or it is available because it dominates `sr`
-              bool dominates =
-                  di.properlyDominates(opr.getDefiningOp(), &*sr.begin);
+              bool dominates = di.properlyDominates(opr, &*sr.begin);
               return remapped || dominates;
             })) {
           // Safe to parallelize operations which have all operands available in
@@ -405,7 +404,7 @@ static void parallelizeRegion(Region &sourceRegion, Region &targetRegion,
 
   if (sourceRegion.hasOneBlock()) {
     handleOneBlock(sourceRegion.front());
-  } else {
+  } else if (!sourceRegion.empty()) {
     auto &domTree = di.getDomTree(&sourceRegion);
     for (auto node : llvm::breadth_first(domTree.getRootNode())) {
       handleOneBlock(*node->getBlock());

flang/test/Transforms/OpenMP/lower-workshare-nested.mlir

@@ -0,0 +1,22 @@
+// RUN: fir-opt --lower-workshare --allow-unregistered-dialect %s | FileCheck %s
+
+// Checks that the nested loop_wrapper gets parallelized
+func.func @wsfunc(%cond : i1) {
+  omp.workshare {
+    %c1 = arith.constant 1 : index
+    %c42 = arith.constant 42 : index
+    fir.if %cond {
+      omp.workshare.loop_wrapper {
+        omp.loop_nest (%arg1) : index = (%c1) to (%c42) inclusive step (%c1) {
+          "test.test1"() : () -> ()
+          omp.yield
+        }
+      }
+    }
+    omp.terminator
+  }
+  return
+}
+
+// CHECK:     fir.if
+// CHECK:       omp.wsloop nowait