[libc] Add basic support for calling host functions from the GPU

[jhuber6]

Summary

This patch adds the rpc_host_call function as a GPU extension. This is
exported from the libc project to use the RPC interface to call a
function pointer via RPC any copying the arguments by-value. The
interface can only support a single void pointer argument much like
pthreads. The function call here is the bare-bones version of what's
required for OpenMP reverse offloading. Full support will require
interfacing with the mapping table, nowait support, etc.

I decided to test this interface in libomptarget as that will be the
primary consumer and it would be more difficult to make a test in libc
due to the testing infrastructure not really having a concept of the
"host" as it runs directly on the GPU as if it were a CPU target.

[jhuber6]

Make the test work for the CPU targets.

[JonChesterfield]

Could you find somewhere to document very explicitly that the pointer in question is to some function on the host and not to some function on the gpu, and not some magic thing which can be called on either?

[jhuber6]

In D155003#4490885, @JonChesterfield wrote:

Could you find somewhere to document very explicitly that the pointer in question is to some function on the host and not to some function on the gpu, and not some magic thing which can be called on either?

I probably should have some documentation for these "GPU" extensions. Though I only expect one more for the H2D / D2H memcpy. I'll just put it in the implementation for now.

[jhuber6]

Add comment

[jhuber6]

Ping

[jplehr]

LGTM

openmp/libomptarget/test/libc/host_call.c

55

Nit: Should main return something?

libc/config/gpu/entrypoints.txt

set(TARGET_LIBC_ENTRYPOINTS
     # inttypes.h entrypoints
     libc.src.inttypes.imaxabs
     libc.src.inttypes.imaxdiv
     libc.src.inttypes.strtoimax
     libc.src.inttypes.strtoumax
 
     # gpu/rpc.h entrypoints
     libc.src.gpu.rpc_reset
+    libc.src.gpu.rpc_host_call
 )
 
 set(TARGET_LIBM_ENTRYPOINTS
     # math.h entrypoints
     libc.src.math.fma
     libc.src.math.fmaf
     libc.src.math.sin
     libc.src.math.round
     libc.src.math.roundf
     libc.src.math.roundl
 )
 
 set(TARGET_LLVMLIBC_ENTRYPOINTS
   ${TARGET_LIBC_ENTRYPOINTS}
   ${TARGET_LIBM_ENTRYPOINTS}
 )

libc/include/llvm-libc-types/rpc_opcodes_t.h

typedef enum : unsigned short {
   RPC_EXIT = 1,
   RPC_WRITE_TO_STDOUT = 2,
   RPC_WRITE_TO_STDERR = 3,
   RPC_WRITE_TO_STREAM = 4,
   RPC_OPEN_FILE = 5,
   RPC_CLOSE_FILE = 6,
   RPC_MALLOC = 7,
   RPC_FREE = 8,
+  RPC_HOST_CALL = 9,
   // TODO: Move these out of here and handle then with custom handlers in the
   // loader.
   RPC_TEST_INCREMENT = 1000,
   RPC_TEST_INTERFACE = 1001,
   RPC_TEST_STREAM = 1002,
 } rpc_opcode_t;
 
 #endif // __LLVM_LIBC_TYPES_RPC_OPCODE_H__

libc/spec/gpu_ext.td

 def GPUExtensions : StandardSpec<"GPUExtensions"> {
   HeaderSpec RPC = HeaderSpec<
     "gpu/rpc.h",
     [], // Macros
     [], // Types
     [], // Enumerations
     [
         FunctionSpec<
             "rpc_reset",
             RetValSpec<VoidType>,
             [ArgSpec<UnsignedIntType>, ArgSpec<VoidPtr>]
         >,
+        FunctionSpec<
+            "rpc_host_call",
+            RetValSpec<VoidType>,
+            [ArgSpec<VoidPtr>, ArgSpec<VoidPtr>, ArgSpec<SizeTType>]
+        >,
     ]
   >;
   let Headers = [
     RPC,
   ];
 }

libc/src/gpu/CMakeLists.txt

 add_entrypoint_object(
   rpc_reset
   SRCS
     rpc_reset.cpp
   HDRS
     rpc_reset.h
   DEPENDS
     libc.src.__support.RPC.rpc_client
     libc.src.__support.GPU.utils
 )
+
+add_entrypoint_object(
+  rpc_host_call
+  SRCS
+    rpc_host_call.cpp
+  HDRS
+    rpc_host_call.h
+  DEPENDS
+    libc.src.__support.RPC.rpc_client
+    libc.src.__support.GPU.utils
+)

libc/src/gpu/rpc_host_call.h

+//===-- Implementation header for RPC functions -----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC_GPU_RPC_HOST_CALL_H
+#define LLVM_LIBC_SRC_GPU_RPC_HOST_CALL_H
+
+#include <stddef.h> // size_t
+
+namespace __llvm_libc {
+
+void rpc_host_call(void *fn, void *buffer, size_t size);
+
+} // namespace __llvm_libc
+
+#endif // LLVM_LIBC_SRC_GPU_RPC_H_HOST_CALL

libc/src/gpu/rpc_host_call.cpp

+//===---------- GPU implementation of the external RPC call function ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/gpu/rpc_host_call.h"
+
+#include "llvm-libc-types/rpc_opcodes_t.h"
+#include "src/__support/GPU/utils.h"
+#include "src/__support/RPC/rpc_client.h"
+#include "src/__support/common.h"
+
+namespace __llvm_libc {
+
+// This calls the associated function pointer on the RPC server with the given
+// arguments. We expect that the pointer here is a valid pointer on the server.
+LLVM_LIBC_FUNCTION(void, rpc_host_call, (void *fn, void *data, size_t size)) {
+  rpc::Client::Port port = rpc::client.open<RPC_HOST_CALL>();
+  port.send_n(data, size);
+  port.send([=](rpc::Buffer *buffer) {
+    buffer->data[0] = reinterpret_cast<uintptr_t>(fn);
+  });
+  port.recv([](rpc::Buffer *) {});
+  port.close();
+}
+
+} // namespace __llvm_libc

libc/utils/gpu/server/rpc_server.cpp

case RPC_EXIT: {
       port->recv_and_send([](rpc::Buffer *) {});
       port->recv([](rpc::Buffer *buffer) {
         int status = 0;
         std::memcpy(&status, buffer->data, sizeof(int));
         exit(status);
       });
       break;
     }
+    case RPC_HOST_CALL: {
+      uint64_t sizes[rpc::MAX_LANE_SIZE] = {0};
+      void *args[rpc::MAX_LANE_SIZE] = {nullptr};
+      port->recv_n(args, sizes, [&](uint64_t size) { return new char[size]; });
+      port->recv([&](rpc::Buffer *buffer, uint32_t id) {
+        reinterpret_cast<void (*)(void *)>(buffer->data[0])(args[id]);
+      });
+      port->send([&](rpc::Buffer *, uint32_t id) {
+        delete[] reinterpret_cast<uint8_t *>(args[id]);
+      });
+      break;
+    }
     // TODO: Move handling of these  test cases to the loader implementation.
     case RPC_TEST_INCREMENT: {
       port->recv_and_send([](rpc::Buffer *buffer) {
         reinterpret_cast<uint64_t *>(buffer->data)[0] += 1;
       });
       break;
     }
     case RPC_TEST_INTERFACE: {
const void *rpc_get_client_buffer(uint32_t device_id) {
   return &state->devices[device_id]->client;
 }
 
 uint64_t rpc_get_client_size() { return sizeof(rpc::Client); }
 
 using ServerPort = std::variant<rpc::Server<1>::Port *, rpc::Server<32>::Port *,
                                 rpc::Server<64>::Port *>;
 
-ServerPort getPort(rpc_port_t ref) {
+ServerPort get_port(rpc_port_t ref) {
   if (ref.lane_size == 1)
     return reinterpret_cast<rpc::Server<1>::Port *>(ref.handle);
   else if (ref.lane_size == 32)
     return reinterpret_cast<rpc::Server<32>::Port *>(ref.handle);
   else if (ref.lane_size == 64)
     return reinterpret_cast<rpc::Server<64>::Port *>(ref.handle);
   else
     __builtin_unreachable();
 }
 
 void rpc_send(rpc_port_t ref, rpc_port_callback_ty callback, void *data) {
-  auto port = getPort(ref);
+  auto port = get_port(ref);
   std::visit(
       [=](auto &port) {
         port->send([=](rpc::Buffer *buffer) {
           callback(reinterpret_cast<rpc_buffer_t *>(buffer), data);
         });
       },
       port);
 }
 
 void rpc_recv(rpc_port_t ref, rpc_port_callback_ty callback, void *data) {
-  auto port = getPort(ref);
+  auto port = get_port(ref);
   std::visit(
       [=](auto &port) {
         port->recv([=](rpc::Buffer *buffer) {
           callback(reinterpret_cast<rpc_buffer_t *>(buffer), data);
         });
       },
       port);
 }
 
 void rpc_recv_and_send(rpc_port_t ref, rpc_port_callback_ty callback,
                        void *data) {
-  auto port = getPort(ref);
+  auto port = get_port(ref);
   std::visit(
       [=](auto &port) {
         port->recv_and_send([=](rpc::Buffer *buffer) {
           callback(reinterpret_cast<rpc_buffer_t *>(buffer), data);
         });
       },
       port);
 }

openmp/libomptarget/test/libc/host_call.c

+// RUN: %libomptarget-compile-run-and-check-generic
+
+// REQUIRES: libc
+
+#include <assert.h>
+#include <omp.h>
+#include <stdio.h>
+
+#pragma omp begin declare variant match(device = {kind(gpu)})
+// Extension provided by the 'libc' project.
+void rpc_host_call(void *fn, void *args, size_t size);
+#pragma omp declare target to(rpc_host_call) device_type(nohost)
+#pragma omp end declare variant
+
+#pragma omp begin declare variant match(device = {kind(cpu)})
+// Dummy host implementation to make this work for all targets.
+void rpc_host_call(void *fn, void *args, size_t size) {
+  ((void (*)(void *))fn)(args);
+}
+#pragma omp end declare variant
+
+typedef struct args_s {
+  int thread_id;
+  int block_id;
+} args_t;
+
+// CHECK-DAG: Thread: 0, Block: 0
+// CHECK-DAG: Thread: 1, Block: 0
+// CHECK-DAG: Thread: 0, Block: 1
+// CHECK-DAG: Thread: 1, Block: 1
+// CHECK-DAG: Thread: 0, Block: 2
+// CHECK-DAG: Thread: 1, Block: 2
+// CHECK-DAG: Thread: 0, Block: 3
+// CHECK-DAG: Thread: 1, Block: 3
+void foo(void *data) {
+  assert(omp_is_initial_device() && "Not executing on host?");
+  args_t *args = (args_t *)data;
+  printf("Thread: %d, Block: %d\n", args->thread_id, args->block_id);
+}
+
+void *fn_ptr = NULL;
+#pragma omp declare target to(fn_ptr)
+
+int main() {
+  fn_ptr = (void *)&foo;
+#pragma omp target update to(fn_ptr)
+
+#pragma omp target teams num_teams(4)
+#pragma omp parallel num_threads(2)
+  {
+    args_t args = {omp_get_thread_num(), omp_get_team_num()};
+    rpc_host_call(fn_ptr, &args, sizeof(args_t));
+  }
+}