[libc] Check the RPC server once again after the kernel exits

[jhuber6]

Summary

We support asynchronous sends, that means that the kernel can issue a
send, then exit the kernel as we do with the EXIT syscall. Because of
the condition it's therefore possible for the kernel to exit and break
from the loop before we check the server again. This can potentially
cause us to ignore an EXIT call from the GPU.

[JonChesterfield]

If the server gets a value after the kernel has finished, it can send work back to the kernel, which won't do anything useful. Is there a way we can guard against that, perhaps in handle?

[jhuber6]

In D150456#4338163, @JonChesterfield wrote:

If the server gets a value after the kernel has finished, it can send work back to the kernel, which won't do anything useful. Is there a way we can guard against that, perhaps in handle?

I don't think that's possible given the API, the kernel would be blocked on that value to be returned from the server and unable to exit.

libc/utils/gpu/loader/amdgpu/Loader.cpp

hsa_status_t launch_kernel(hsa_agent_t dev_agent, hsa_executable_t executable,
 
   // Wait until the kernel has completed execution on the device. Periodically
   // check the RPC client for work to be performed on the server.
   while (hsa_signal_wait_scacquire(
              packet->completion_signal, HSA_SIGNAL_CONDITION_EQ, 0,
              /*timeout_hint=*/1024, HSA_WAIT_STATE_ACTIVE) != 0)
     handle_server();
 
+  // Handle the server one more time in case the kernel exited with a pending
+  // send still in flight.
+  handle_server();
+
   // Destroy the resources acquired to launch the kernel and return.
   if (hsa_status_t err = hsa_amd_memory_pool_free(args))
     handle_error(err);
   if (hsa_status_t err = hsa_signal_destroy(packet->completion_signal))
     handle_error(err);
 
   return HSA_STATUS_SUCCESS;
 }

libc/utils/gpu/loader/nvptx/Loader.cpp

if (CUresult err = cuLaunchKernel(
           params.num_threads_z, 0, stream, nullptr, args_config))
     handle_error(err);
 
   // Wait until the kernel has completed execution on the device. Periodically
   // check the RPC client for work to be performed on the server.
   while (cuStreamQuery(stream) == CUDA_ERROR_NOT_READY)
     handle_server();
 
+  // Handle the server one more time in case the kernel exited with a pending
+  // send still in flight.
+  handle_server();
+
   return CUDA_SUCCESS;
 }
 
 int load(int argc, char **argv, char **envp, void *image, size_t size,
          const LaunchParameters &params) {
 
   if (CUresult err = cuInit(0))
     handle_error(err);