From 01889de8e9b16eeed7ed9f6cdc18636ad20a01ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Valentin=20Clement=20=28=E3=83=90=E3=83=AC=E3=83=B3?=
 =?UTF-8?q?=E3=82=BF=E3=82=A4=E3=83=B3=20=E3=82=AF=E3=83=AC=E3=83=A1?=
 =?UTF-8?q?=E3=83=B3=29?= <clementval@gmail.com>
Date: Tue, 1 Apr 2025 10:06:45 -0700
Subject: [PATCH] [flang][device] Enable Stop functions on device build
 (#133803)

Update `StopStatement` and `StopStatementText` to be build for the
device.
---
 flang-rt/lib/runtime/CMakeLists.txt |  1 +
 flang-rt/lib/runtime/stop.cpp       | 52 +++++++++++++++++++++++++++--
 flang/include/flang/Runtime/stop.h  |  7 ++--
 3 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/flang-rt/lib/runtime/CMakeLists.txt b/flang-rt/lib/runtime/CMakeLists.txt
index 572b4d54552c..c5e7bdce5b2f 100644
--- a/flang-rt/lib/runtime/CMakeLists.txt
+++ b/flang-rt/lib/runtime/CMakeLists.txt
@@ -57,6 +57,7 @@ set(supported_sources
   pseudo-unit.cpp
   ragged.cpp
   stat.cpp
+  stop.cpp
   sum.cpp
   support.cpp
   terminator.cpp
diff --git a/flang-rt/lib/runtime/stop.cpp b/flang-rt/lib/runtime/stop.cpp
index 1d70a137377a..a4ef7104442f 100644
--- a/flang-rt/lib/runtime/stop.cpp
+++ b/flang-rt/lib/runtime/stop.cpp
@@ -65,8 +65,33 @@ static void CloseAllExternalUnits(const char *why) {
   Fortran::runtime::io::ExternalFileUnit::CloseAll(handler);
 }
 
-[[noreturn]] void RTNAME(StopStatement)(
+[[noreturn]] RT_API_ATTRS void RTNAME(StopStatement)(
     int code, bool isErrorStop, bool quiet) {
+#if defined(RT_DEVICE_COMPILATION)
+  if (Fortran::runtime::executionEnvironment.noStopMessage && code == 0) {
+    quiet = true;
+  }
+  if (!quiet) {
+    if (isErrorStop) {
+      std::printf("Fortran ERROR STOP");
+    } else {
+      std::printf("Fortran STOP");
+    }
+    if (code != EXIT_SUCCESS) {
+      std::printf(": code %d\n", code);
+    }
+    std::printf('\n');
+  }
+#if defined(__CUDACC__)
+  // NVCC supports __trap().
+  __trap();
+#elif defined(__clang__)
+  // Clang supports __builtin_trap().
+  __builtin_trap();
+#else
+#error "unsupported compiler"
+#endif
+#else
   CloseAllExternalUnits("STOP statement");
   if (Fortran::runtime::executionEnvironment.noStopMessage && code == 0) {
     quiet = true;
@@ -80,10 +105,32 @@ static void CloseAllExternalUnits(const char *why) {
     DescribeIEEESignaledExceptions();
   }
   std::exit(code);
+#endif
 }
 
-[[noreturn]] void RTNAME(StopStatementText)(
+[[noreturn]] RT_API_ATTRS void RTNAME(StopStatementText)(
     const char *code, std::size_t length, bool isErrorStop, bool quiet) {
+#if defined(RT_DEVICE_COMPILATION)
+  if (!quiet) {
+    if (Fortran::runtime::executionEnvironment.noStopMessage && !isErrorStop) {
+      std::printf("%s\n", code);
+    } else {
+      std::printf(
+          "Fortran %s: %s\n", isErrorStop ? "ERROR STOP" : "STOP", code);
+    }
+  }
+  if (isErrorStop) {
+#if defined(__CUDACC__)
+    // NVCC supports __trap().
+    __trap();
+#elif defined(__clang__)
+    // Clang supports __builtin_trap().
+    __builtin_trap();
+#else
+#error "unsupported compiler"
+#endif
+  }
+#else
   CloseAllExternalUnits("STOP statement");
   if (!quiet) {
     if (Fortran::runtime::executionEnvironment.noStopMessage && !isErrorStop) {
@@ -99,6 +146,7 @@ static void CloseAllExternalUnits(const char *why) {
   } else {
     std::exit(EXIT_SUCCESS);
   }
+#endif
 }
 
 static bool StartPause() {
diff --git a/flang/include/flang/Runtime/stop.h b/flang/include/flang/Runtime/stop.h
index 24ae2cbe01ec..02bce6576590 100644
--- a/flang/include/flang/Runtime/stop.h
+++ b/flang/include/flang/Runtime/stop.h
@@ -17,9 +17,10 @@
 FORTRAN_EXTERN_C_BEGIN
 
 // Program-initiated image stop
-NORETURN void RTNAME(StopStatement)(int code DEFAULT_VALUE(EXIT_SUCCESS),
-    bool isErrorStop DEFAULT_VALUE(false), bool quiet DEFAULT_VALUE(false));
-NORETURN void RTNAME(StopStatementText)(const char *, size_t,
+NORETURN RT_API_ATTRS void RTNAME(StopStatement)(
+    int code DEFAULT_VALUE(EXIT_SUCCESS), bool isErrorStop DEFAULT_VALUE(false),
+    bool quiet DEFAULT_VALUE(false));
+NORETURN RT_API_ATTRS void RTNAME(StopStatementText)(const char *, size_t,
     bool isErrorStop DEFAULT_VALUE(false), bool quiet DEFAULT_VALUE(false));
 void RTNAME(PauseStatement)(NO_ARGUMENTS);
 void RTNAME(PauseStatementInt)(int);