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

[Ivan R. Ivanov via flang-commits]

https://github.com/ivanradanov created https://github.com/llvm/llvm-project/pull/117275

None

>From c0f60c5875aa122fdf54959176f6fe2f9daf67d2 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <ivanov.i.aa at m.titech.ac.jp>
Date: Fri, 22 Nov 2024 10:16:23 +0900
Subject: [PATCH] [flang][OpenMP] Fix handling of nested loops in
 LowerWorkshare

---
 flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp |  9 ++++----
 .../OpenMP/lower-workshare-nested.mlir        | 22 +++++++++++++++++++
 2 files changed, 26 insertions(+), 5 deletions(-)
 create mode 100644 flang/test/Transforms/OpenMP/lower-workshare-nested.mlir

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