Allow CSE to cross omp.{teams, distribute, parallel} regions

ROCm-Developer-Tools · Sep 28, 2023 · ced8645 · ced8645
1 parent ccde72e
commit ced8645
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/flang/test/Lower/OpenMP/hlfir-wsloop.f90 b/flang/test/Lower/OpenMP/hlfir-wsloop.f90
@@ -6,10 +6,10 @@
 !CHECK-LABEL: func @_QPsimple_loop()
 subroutine simple_loop
   integer :: i
-  ! CHECK:  omp.parallel
-  !$OMP PARALLEL
   ! CHECK-DAG:     %[[WS_ST:.*]] = arith.constant 1 : i32
   ! CHECK-DAG:     %[[WS_END:.*]] = arith.constant 9 : i32
+  ! CHECK:  omp.parallel
+  !$OMP PARALLEL
   ! CHECK-DAG:     %[[ALLOCA_IV:.*]] = fir.alloca i32 {{{.*}}, pinned}
   ! CHECK:     %[[IV:.*]]    = fir.declare %[[ALLOCA_IV]] {uniq_name = "_QFsimple_loopEi"} : (!fir.ref<i32>) -> !fir.ref<i32>
   ! CHECK:     omp.wsloop for (%[[I:.*]]) : i32 = (%[[WS_ST]]) to (%[[WS_END]]) inclusive step (%[[WS_ST]])

diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -53,7 +53,7 @@ struct OpenMPDialectCSEInterface : public DialectCSEInterface {
 
   bool subexpressionExtractionAllowed(Operation *op) const final {
     // Avoid extracting common subexpressions across op boundaries
-    return !isa<TargetOp, TeamsOp, DistributeOp, ParallelOp>(op);
+    return !isa<TargetOp>(op);
   }
 };
 
@@ -62,8 +62,7 @@ struct OpenMPDialectFoldInterface : public DialectFoldInterface {
 
   bool shouldMaterializeInto(Region *region) const final {
     // Avoid folding constants across target regions
-    return isa<TargetOp, TeamsOp, DistributeOp, ParallelOp>(
-        region->getParentOp());
+    return isa<TargetOp>(region->getParentOp());
   }
 };
 } // namespace