[libc-commits] [libc] 667c103 - [libc] Fix the implementation of exit on the GPU

[Joseph Huber via libc-commits]

Author: Joseph Huber
Date: 2023-06-29T13:22:23-05:00
New Revision: 667c10353e87d2be5b981cdac0a698c93946cc92

URL: https://github.com/llvm/llvm-project/commit/667c10353e87d2be5b981cdac0a698c93946cc92
DIFF: https://github.com/llvm/llvm-project/commit/667c10353e87d2be5b981cdac0a698c93946cc92.diff

LOG: [libc] Fix the implementation of exit on the GPU

The RPC calls all have delays associated with them. Currently the `exit`
function does an async send and immediately exits the GPU. This can have
the effect that the RPC server never sees the exit call and we continue.
This patch changes that to first sync with the server before continuing
to perform its exit. There is still a hazard here, where the kernel can
complete before the RPC call reads back its response, but this is simply
multi-threaded hazards. This change ensures that the server *will*
always exit some time after the GPU exits.

Reviewed By: JonChesterfield

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

Added: 
    

Modified: 
    libc/docs/gpu/support.rst
    libc/src/__support/OSUtil/gpu/quick_exit.cpp
    libc/utils/gpu/server/rpc_server.cpp

Removed: 
    


################################################################################
diff  --git a/libc/docs/gpu/support.rst b/libc/docs/gpu/support.rst
index 816d9616403e7a..a4ec0e658186a3 100644
--- a/libc/docs/gpu/support.rst
+++ b/libc/docs/gpu/support.rst
@@ -89,6 +89,7 @@ atoi           |check|
 atof           |check|
 atol           |check|
 atoll          |check|
+exit           |check|    |check|
 labs           |check|
 llabs          |check|
 strtod         |check|

diff  --git a/libc/src/__support/OSUtil/gpu/quick_exit.cpp b/libc/src/__support/OSUtil/gpu/quick_exit.cpp
index ec688f932b4f56..470b2bdbdaa70c 100644
--- a/libc/src/__support/OSUtil/gpu/quick_exit.cpp
+++ b/libc/src/__support/OSUtil/gpu/quick_exit.cpp
@@ -17,7 +17,9 @@
 namespace __llvm_libc {
 
 void quick_exit(int status) {
+  // We want to first make sure the server is listening before we exit.
   rpc::Client::Port port = rpc::client.open<RPC_EXIT>();
+  port.send_and_recv([](rpc::Buffer *) {}, [](rpc::Buffer *) {});
   port.send([&](rpc::Buffer *buffer) {
     reinterpret_cast<uint32_t *>(buffer->data)[0] = status;
   });

diff  --git a/libc/utils/gpu/server/rpc_server.cpp b/libc/utils/gpu/server/rpc_server.cpp
index a440efddeb8abb..b5103f754392b4 100644
--- a/libc/utils/gpu/server/rpc_server.cpp
+++ b/libc/utils/gpu/server/rpc_server.cpp
@@ -102,6 +102,8 @@ struct Server {
       break;
     }
     case RPC_EXIT: {
+      // Send a response to the client to signal that we are ready to exit.
+      port->recv_and_send([](rpc::Buffer *) {});
       port->recv([](rpc::Buffer *buffer) {
         int status = 0;
         std::memcpy(&status, buffer->data, sizeof(int));