[flang][OpenMP] Fix handling of nested loop wrappers in LowerWorkshare

[ivanradanov]

No description provided.

[llvmbot]

@llvm/pr-subscribers-flang-openmp

Author: Ivan R. Ivanov (ivanradanov)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/117275.diff

2 Files Affected:

• (modified) flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp (+4-5)
• (added) flang/test/Transforms/OpenMP/lower-workshare-nested.mlir (+22)

diff --git a/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp b/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
index 225c585a02d913..6e130a96eb8dd3 100644
--- a/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
+++ b/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());
diff --git a/flang/test/Transforms/OpenMP/lower-workshare-nested.mlir b/flang/test/Transforms/OpenMP/lower-workshare-nested.mlir
new file mode 100644
index 00000000000000..bfd65f04d94b12
--- /dev/null
+++ b/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

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Ivan R. Ivanov (ivanradanov)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/117275.diff

2 Files Affected:

• (modified) flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp (+4-5)
• (added) flang/test/Transforms/OpenMP/lower-workshare-nested.mlir (+22)

diff --git a/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp b/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
index 225c585a02d913..6e130a96eb8dd3 100644
--- a/flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp
+++ b/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());
diff --git a/flang/test/Transforms/OpenMP/lower-workshare-nested.mlir b/flang/test/Transforms/OpenMP/lower-workshare-nested.mlir
new file mode 100644
index 00000000000000..bfd65f04d94b12
--- /dev/null
+++ b/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

[tblah]

LGTM