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

[via flang-commits]

Author: Ivan R. Ivanov
Date: 2024-11-26T09:30:27+09:00
New Revision: 2ed8c5de585491182486f392a5a570cc5c4474df

URL: https://github.com/llvm/llvm-project/commit/2ed8c5de585491182486f392a5a570cc5c4474df
DIFF: https://github.com/llvm/llvm-project/commit/2ed8c5de585491182486f392a5a570cc5c4474df.diff

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

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

Modified: 
    flang/lib/Optimizer/OpenMP/LowerWorkshare.cpp

Removed: 
    


################################################################################
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