[flang][OpenMP][OpenACC] Support stop statement in OpenMP/OpenACC region

[flang][OpenMP][OpenACC] Support stop statement in OpenMP/OpenACC region

This supports lowering of stop statement in OpenMP/OpenACC region.
* OpenMP/OpenACC: Emit `fir.unreachable` only if the block is not
  terminated by any terminator. This avoids knocking off an existing
  OpenMP/OpenACC terminator.
* OpenMP: Emit the OpenMP terminator instead of `fir.unreachable` since
  OpenMP regions can only be terminated by OpenMP terminators. This is
  currently skipped for OpenACC since unstructured code is not yet
  handled specially in OpenACC lowering.

Fixes #60737
Fixes #61877

Co-authored-by: Kiran Chandramohan <kiranchandramohan@gmail.com>
Co-authored-by: Val Donaldson <vdonaldson@nvidia.com>

Reviewed By: vdonaldson, peixin

Differential Revision: https://reviews.llvm.org/D129969

5 files changed, 223 insertions, 4 deletions

diff --git a/flang/include/flang/Lower/OpenMP.h b/flang/include/flang/Lower/OpenMP.h
index 342012f..759afc6 100644
--- a/flang/include/flang/Lower/OpenMP.h
+++ b/flang/include/flang/Lower/OpenMP.h
@@ -18,6 +18,7 @@
 namespace mlir {
 class Value;
 class Operation;
+class Location;
 } // namespace mlir
 
 namespace fir {
@@ -42,6 +43,10 @@ struct Evaluation;
 struct Variable;
 } // namespace pft
 
+// Generate the OpenMP terminator for Operation at Location.
+void genOpenMPTerminator(fir::FirOpBuilder &, mlir::Operation *,
+                         mlir::Location);
+
 void genOpenMPConstruct(AbstractConverter &, pft::Evaluation &,
                         const parser::OpenMPConstruct &);
 void genOpenMPDeclarativeConstruct(AbstractConverter &, pft::Evaluation &,
diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index 2a3fc19..141912a 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -27,6 +27,15 @@
 
 using namespace mlir;
 
+void Fortran::lower::genOpenMPTerminator(fir::FirOpBuilder &builder,
+                                         Operation *op, mlir::Location loc) {
+  if (mlir::isa<omp::WsLoopOp, omp::ReductionDeclareOp, omp::AtomicUpdateOp,
+                omp::SimdLoopOp>(op))
+    builder.create<omp::YieldOp>(loc);
+  else
+    builder.create<omp::TerminatorOp>(loc);
+}
+
 int64_t Fortran::lower::getCollapseValue(
     const Fortran::parser::OmpClauseList &clauseList) {
   for (const auto &clause : clauseList.v) {
diff --git a/flang/lib/Lower/Runtime.cpp b/flang/lib/Lower/Runtime.cpp
index 54c3a4e..11cb56c 100644
--- a/flang/lib/Lower/Runtime.cpp
+++ b/flang/lib/Lower/Runtime.cpp
@@ -8,6 +8,7 @@
 
 #include "flang/Lower/Runtime.h"
 #include "flang/Lower/Bridge.h"
+#include "flang/Lower/OpenMP.h"
 #include "flang/Lower/StatementContext.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Runtime/RTBuilder.h"
@@ -20,6 +21,7 @@
 #include "flang/Runtime/stop.h"
 #include "flang/Runtime/time-intrinsic.h"
 #include "flang/Semantics/tools.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "llvm/Support/Debug.h"
 #include <optional>
 
@@ -30,9 +32,14 @@ using namespace Fortran::runtime;
 /// Runtime calls that do not return to the caller indicate this condition by
 /// terminating the current basic block with an unreachable op.
 static void genUnreachable(fir::FirOpBuilder &builder, mlir::Location loc) {
-  builder.create<fir::UnreachableOp>(loc);
-  mlir::Block *newBlock =
-      builder.getBlock()->splitBlock(builder.getInsertionPoint());
+  mlir::Block *curBlock = builder.getBlock();
+  mlir::Operation *parentOp = curBlock->getParentOp();
+  if (parentOp->getDialect()->getNamespace() ==
+      mlir::omp::OpenMPDialect::getDialectNamespace())
+    Fortran::lower::genOpenMPTerminator(builder, parentOp, loc);
+  else
+    builder.create<fir::UnreachableOp>(loc);
+  mlir::Block *newBlock = curBlock->splitBlock(builder.getInsertionPoint());
   builder.setInsertionPointToStart(newBlock);
 }
 
@@ -106,7 +113,13 @@ void Fortran::lower::genStopStatement(
   }
 
   builder.create<fir::CallOp>(loc, callee, operands);
-  genUnreachable(builder, loc);
+  auto blockIsUnterminated = [&builder]() {
+    mlir::Block *currentBlock = builder.getBlock();
+    return currentBlock->empty() ||
+           !currentBlock->back().hasTrait<mlir::OpTrait::IsTerminator>();
+  };
+  if (blockIsUnterminated())
+    genUnreachable(builder, loc);
 }
 
 void Fortran::lower::genFailImageStatement(
diff --git a/flang/test/Lower/OpenACC/stop-stmt-in-region.f90 b/flang/test/Lower/OpenACC/stop-stmt-in-region.f90
new file mode 100644
index 0000000..b36e3b4
--- /dev/null
+++ b/flang/test/Lower/OpenACC/stop-stmt-in-region.f90
@@ -0,0 +1,41 @@
+! This test checks lowering of stop statement in OpenACC region.
+
+! RUN: bbc -fopenacc -emit-fir %s -o - | FileCheck %s
+! RUN: %flang_fc1 -emit-fir -fopenacc %s -o - | FileCheck %s
+
+! CHECK-LABEL: func.func @_QPtest_stop_in_region1() {
+! CHECK:         acc.parallel {
+! CHECK:           %[[VAL_0:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_1:.*]] = arith.constant false
+! CHECK:           %[[VAL_2:.*]] = arith.constant false
+! CHECK:           %[[VAL_3:.*]] = fir.call @_FortranAStopStatement(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) {{.*}} : (i32, i1, i1) -> none
+! CHECK:           acc.yield
+! CHECK:         }
+! CHECK:         return
+! CHECK:       }
+
+subroutine test_stop_in_region1()
+  !$acc parallel
+    stop 1
+  !$acc end parallel
+end
+
+! CHECK-LABEL: func.func @_QPtest_stop_in_region2() {
+! CHECK:         %[[VAL_0:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFtest_stop_in_region2Ex"}
+! CHECK:         acc.parallel {
+! CHECK:           %[[VAL_1:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_2:.*]] = arith.constant false
+! CHECK:           %[[VAL_3:.*]] = arith.constant false
+! CHECK:           %[[VAL_4:.*]] = fir.call @_FortranAStopStatement(%[[VAL_1]], %[[VAL_2]], %[[VAL_3]]) {{.*}} : (i32, i1, i1) -> none
+! CHECK:           acc.yield
+! CHECK:         }
+! CHECK:         return
+! CHECK:       }
+
+subroutine test_stop_in_region2()
+  integer :: x
+  !$acc parallel
+    stop 1
+    x = 2
+  !$acc end parallel
+end
diff --git a/flang/test/Lower/OpenMP/stop-stmt-in-region.f90 b/flang/test/Lower/OpenMP/stop-stmt-in-region.f90
new file mode 100644
index 0000000..75a5c73
--- /dev/null
+++ b/flang/test/Lower/OpenMP/stop-stmt-in-region.f90
@@ -0,0 +1,151 @@
+! This test checks lowering of stop statement in OpenMP region.
+
+! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+! RUN: %flang_fc1 -emit-fir -fopenmp %s -o - | FileCheck %s
+
+! CHECK-LABEL: func.func @_QPtest_stop_in_region1() {
+! CHECK:         omp.parallel   {
+! CHECK:           %[[VAL_0:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_1:.*]] = arith.constant false
+! CHECK:           %[[VAL_2:.*]] = arith.constant false
+! CHECK:           %[[VAL_3:.*]] = fir.call @_FortranAStopStatement(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) {{.*}} : (i32, i1, i1) -> none
+! CHECK-NOT:       fir.unreachable
+! CHECK:           omp.terminator
+! CHECK:         }
+! CHECK:         return
+! CHECK:       }
+
+subroutine test_stop_in_region1()
+  !$omp parallel
+    stop 1
+  !$omp end parallel
+end
+
+! CHECK-LABEL: func.func @_QPtest_stop_in_region2() {
+! CHECK:         %[[VAL_0:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFtest_stop_in_region2Ex"}
+! CHECK:         omp.parallel   {
+! CHECK:           %[[VAL_1:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_2:.*]] = arith.constant false
+! CHECK:           %[[VAL_3:.*]] = arith.constant false
+! CHECK:           %[[VAL_4:.*]] = fir.call @_FortranAStopStatement(%[[VAL_1]], %[[VAL_2]], %[[VAL_3]]) {{.*}} : (i32, i1, i1) -> none
+! CHECK:           omp.terminator
+! CHECK:         }
+! CHECK:         return
+! CHECK:       }
+
+subroutine test_stop_in_region2()
+  integer :: x
+  !$omp parallel
+    stop 1
+    x = 2
+  !$omp end parallel
+end
+
+! CHECK-LABEL: func.func @_QPtest_stop_in_region3() {
+! CHECK:         %[[VAL_0:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFtest_stop_in_region3Ex"}
+! CHECK:         omp.parallel   {
+! CHECK:           %[[VAL_1:.*]] = arith.constant 3 : i32
+! CHECK:           fir.store %[[VAL_1]] to %[[VAL_0]] : !fir.ref<i32>
+! CHECK:           %[[VAL_2:.*]] = fir.load %[[VAL_0]] : !fir.ref<i32>
+! CHECK:           %[[VAL_3:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_4:.*]] = arith.cmpi sgt, %[[VAL_2]], %[[VAL_3]] : i32
+! CHECK:           cf.cond_br %[[VAL_4]], ^bb1, ^bb2
+! CHECK:         ^bb1:
+! CHECK:           %[[VAL_5:.*]] = fir.load %[[VAL_0]] : !fir.ref<i32>
+! CHECK:           %[[VAL_6:.*]] = arith.constant false
+! CHECK:           %[[VAL_7:.*]] = arith.constant false
+! CHECK:           %[[VAL_8:.*]] = fir.call @_FortranAStopStatement(%[[VAL_5]], %[[VAL_6]], %[[VAL_7]]) {{.*}} : (i32, i1, i1) -> none
+! CHECK:           omp.terminator
+! CHECK:         ^bb2:
+! CHECK:           omp.terminator
+! CHECK:         }
+! CHECK:         return
+! CHECK:       }
+
+subroutine test_stop_in_region3()
+  integer :: x
+  !$omp parallel
+    x = 3
+    if (x > 1) stop x
+  !$omp end parallel
+end
+
+! CHECK-LABEL: func.func @_QPtest_stop_in_region4() {
+! CHECK:         %[[VAL_0:.*]] = fir.alloca i32 {adapt.valuebyref}
+! CHECK:         %[[VAL_1:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFtest_stop_in_region4Ei"}
+! CHECK:         %[[VAL_2:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFtest_stop_in_region4Ex"}
+! CHECK:         %[[VAL_3:.*]] = arith.constant 1 : i32
+! CHECK:         %[[VAL_4:.*]] = arith.constant 10 : i32
+! CHECK:         %[[VAL_5:.*]] = arith.constant 1 : i32
+! CHECK:         omp.wsloop   for  (%[[VAL_6:.*]]) : i32 = (%[[VAL_3]]) to (%[[VAL_4]]) inclusive step (%[[VAL_5]]) {
+! CHECK:           fir.store %[[VAL_6]] to %[[VAL_0]] : !fir.ref<i32>
+! CHECK:           cf.br ^bb1
+! CHECK:         ^bb1:
+! CHECK:           %[[VAL_7:.*]] = arith.constant 3 : i32
+! CHECK:           fir.store %[[VAL_7]] to %[[VAL_2]] : !fir.ref<i32>
+! CHECK:           %[[VAL_8:.*]] = fir.load %[[VAL_2]] : !fir.ref<i32>
+! CHECK:           %[[VAL_9:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_10:.*]] = arith.cmpi sgt, %[[VAL_8]], %[[VAL_9]] : i32
+! CHECK:           cf.cond_br %[[VAL_10]], ^bb2, ^bb3
+! CHECK:         ^bb2:
+! CHECK:           %[[VAL_11:.*]] = fir.load %[[VAL_2]] : !fir.ref<i32>
+! CHECK:           %[[VAL_12:.*]] = arith.constant false
+! CHECK:           %[[VAL_13:.*]] = arith.constant false
+! CHECK:           %[[VAL_14:.*]] = fir.call @_FortranAStopStatement(%[[VAL_11]], %[[VAL_12]], %[[VAL_13]]) {{.*}} : (i32, i1, i1) -> none
+! CHECK:           omp.yield
+! CHECK:         ^bb3:
+! CHECK:           omp.yield
+! CHECK:         }
+! CHECK:         cf.br ^bb1
+! CHECK:       ^bb1:
+! CHECK:         return
+! CHECK:       }
+
+subroutine test_stop_in_region4()
+  integer :: x
+  !$omp do
+  do i = 1, 10
+    x = 3
+    if (x > 1) stop x
+  enddo
+  !$omp end do
+end
+
+
+!CHECK-LABEL: func.func @_QPtest_stop_in_region5
+!CHECK:   omp.parallel   {
+!CHECK:     {{.*}} fir.call @_FortranAStopStatement({{.*}}, {{.*}}, {{.*}}) fastmath<contract> : (i32, i1, i1) -> none
+!CHECK:     omp.terminator
+!CHECK:   }
+!CHECK:   return
+
+subroutine test_stop_in_region5()
+  !$omp parallel
+  block
+    stop 1
+  end block
+  !$omp end parallel
+end
+
+!CHECK-LABEL: func.func @_QPtest_stop_in_region6
+!CHECK:  omp.parallel   {
+!CHECK:    cf.cond_br %{{.*}}, ^[[BB1:.*]], ^[[BB2:.*]]
+!CHECK:  ^[[BB1]]:
+!CHECK:    {{.*}}fir.call @_FortranAStopStatement({{.*}}, {{.*}}, {{.*}}) fastmath<contract> : (i32, i1, i1) -> none
+!CHECK:    omp.terminator
+!CHECK:  ^[[BB2]]:
+!CHECK:    {{.*}}fir.call @_FortranAStopStatement({{.*}}, {{.*}}, {{.*}}) fastmath<contract> : (i32, i1, i1) -> none
+!CHECK:    omp.terminator
+!CHECK:  }
+!CHECK:  return
+
+subroutine test_stop_in_region6(x)
+  integer :: x
+  !$omp parallel
+  if (x .gt. 1) then
+    stop 1
+  else
+    stop 2
+  end if
+  !$omp end parallel
+end