chore: logs

Author: brodeynewman

codegen/gen_client.cpp

@@ -19,6 +19,8 @@ extern conn_t *rpc_client_get_connection(unsigned int index);
 int is_unified_pointer(conn_t *conn, void *arg);
 int maybe_copy_unified_arg(conn_t *conn, void *arg, enum cudaMemcpyKind kind);
 extern void rpc_close(conn_t *conn);
+extern void increment_host_nodes();
+extern void wait_for_callbacks();
 void invoke_host_funcs(const int index, void *udata);
 
 nvmlReturn_t nvmlInit_v2() {
@@ -23092,6 +23094,7 @@ cudaError_t cudaGraphAddHostNode(cudaGraphNode_t *pGraphNode, cudaGraph_t graph,
                                  size_t numDependencies,
                                  const struct cudaHostNodeParams *pNodeParams) {
   conn_t *conn = rpc_client_get_connection(0);
+  increment_host_nodes();
   printf("hmmmm %p\n", pNodeParams->fn);
   if (maybe_copy_unified_arg(conn, (void *)&numDependencies,
                              cudaMemcpyHostToDevice) < 0)
@@ -24972,6 +24975,8 @@ cudaError_t cudaGraphLaunch(cudaGraphExec_t graphExec, cudaStream_t stream) {
       rpc_read_end(conn) < 0)
     return cudaErrorDevicesUnavailable;
 
+  wait_for_callbacks();
+
   if (maybe_copy_unified_arg(conn, (void *)&graphExec, cudaMemcpyDeviceToHost) <
       0)
     return cudaErrorDevicesUnavailable;

codegen/gen_server.cpp

@@ -21194,8 +21194,8 @@ int handle_cudaGraphAddHostNode(conn_t *conn) {
 
   hostFnData = (callBackData_t *)malloc(sizeof(callBackData_t));
   // assign the previous function pointer so we can map back to it
+  store_conn(conn);
   hostFnData->callback = pNodeParams.fn;
-  hostFnData->conn = conn;
   hostFnData->data = pNodeParams.userData;
 
   pNodeParams.fn = invoke_host_func;
@@ -22714,11 +22714,15 @@ int handle_cudaGraphLaunch(conn_t *conn) {
 
   scuda_intercept_result = cudaGraphLaunch(graphExec, stream);
 
+  std::cout << "RESPONDING TO CUDAGRAPH" << std::endl;
+
   if (rpc_write_start_response(conn, request_id) < 0 ||
       rpc_write(conn, &scuda_intercept_result, sizeof(cudaError_t)) < 0 ||
       rpc_write_end(conn) < 0)
     goto ERROR_0;
 
+  std::cout << "DONE CUDAGRAPH" << std::endl;
+
   return 0;
 ERROR_0:
   return -1;

rpc.cpp

@@ -1,6 +1,7 @@
 #include <sys/socket.h>
 
 #include "rpc.h"
+#include <string.h>
 #include <iostream>
 #include <unistd.h>
 
@@ -10,17 +11,19 @@
 //
 // it is not necessary to call rpc_read_start() if it is the first call in
 // the sequence because by convention, the handler owns the read lock on entry.
-unsigned int rpc_dispatch(conn_t *conn, int parity) {
+int rpc_dispatch(conn_t *conn, int parity) {
   if (pthread_mutex_lock(&conn->read_mutex) < 0) {
     return -1;
   }
 
-  while (true) {
+  while (1) {
     std::cout << "rpc_dispatch: waiting for read id" << std::endl;
 
     while (conn->read_id != 0)
       pthread_cond_wait(&conn->read_cond, &conn->read_mutex);
 
+    std::cout << "reading req id... " << std::endl;
+
     // the read id is zero so it's our turn to read the next int which is the
     // request id of the next request.
     if (rpc_read(conn, &conn->read_id, sizeof(int)) < 0) {
@@ -33,13 +36,12 @@ unsigned int rpc_dispatch(conn_t *conn, int parity) {
     if (conn->read_id % 2 == parity) {
       std::cout << "rpc_dispatch: new read id: " << conn->read_id << std::endl;
       // this request is the one to be dispatched, read the op and return it
-      unsigned int op;
-      if (rpc_read(conn, &op, sizeof(unsigned int)) < 0) {
+      int op;
+      if (rpc_read(conn, &op, sizeof(int)) < 0) {
         pthread_mutex_unlock(&conn->read_mutex);
         return -1;
       }
 
-      pthread_mutex_unlock(&conn->read_mutex);
       return op;
     } else {
       // this is a response to a request that so signal the update and wait for
@@ -62,6 +64,7 @@ int rpc_read_start(conn_t *conn, int write_id) {
   if (pthread_mutex_lock(&conn->read_mutex) < 0)
     return -1;
 
+
   // wait for the active read id to be the request id we are waiting for
   while (conn->read_id != write_id)
     if (pthread_cond_wait(&conn->read_cond, &conn->read_mutex) < 0)
@@ -71,7 +74,11 @@ int rpc_read_start(conn_t *conn, int write_id) {
 }
 
 int rpc_read(conn_t *conn, void *data, size_t size) {
-  return recv(conn->connfd, data, size, MSG_WAITALL);
+    int bytes_read = recv(conn->connfd, data, size, MSG_WAITALL);
+    if (bytes_read == -1) {
+        printf("recv error: %s\n", strerror(errno));
+    }
+    return bytes_read;
 }
 
 // rpc_read_end releases the response lock on the given connection.
@@ -99,7 +106,7 @@ int rpc_wait_for_response(conn_t *conn) {
 //
 // only one request can be active at a time, so this function will take the
 // request lock from the connection.
-int rpc_write_start_request(conn_t *conn, const unsigned int op) {
+int rpc_write_start_request(conn_t *conn, const int op) {
   if (pthread_mutex_lock(&conn->write_mutex) < 0) {
 #ifdef VERBOSE
     std::cout << "rpc_write_start failed due to rpc_open() < 0 || "

rpc.h

@@ -19,14 +19,14 @@ typedef struct {
   int write_iov_count = 0;
 } conn_t;
 
-extern unsigned int rpc_dispatch(conn_t *conn, int parity);
+extern int rpc_dispatch(conn_t *conn, int parity);
 extern int rpc_read_start(conn_t *conn, int write_id);
 extern int rpc_read(conn_t *conn, void *data, size_t size);
 extern int rpc_read_end(conn_t *conn);
 
 extern int rpc_wait_for_response(conn_t *conn);
 
-extern int rpc_write_start_request(conn_t *conn, const unsigned int op);
+extern int rpc_write_start_request(conn_t *conn, const int op);
 extern int rpc_write_start_response(conn_t *conn, const int read_id);
 extern int rpc_write(conn_t *conn, const void *data, const size_t size);
 extern int rpc_write_end(conn_t *conn);

server.cpp

@@ -114,6 +114,12 @@ static void segfault(int sig, siginfo_t *info, void *unused) {
   raise(SIGSEGV);
 }
 
+conn_t* stored_conn;
+
+void store_conn(const void *conn) {
+  stored_conn = (conn_t *)conn;
+}
+
 typedef struct callBackData {
   conn_t *conn;
   void (*callback)(void *);
@@ -122,15 +128,53 @@ typedef struct callBackData {
 
 void invoke_host_func(void *data) {
   callBackData_t *tmp = (callBackData_t *)(data);
-  void *res;
+  void* scuda_intercept_result;
+
+  // Validate connection
+  if (!stored_conn) {
+      std::cerr << "Error: Connection is NULL in invoke_host_func" << std::endl;
+      return;
+  }
+
+  printf("Invoking host function %p\n", tmp->callback);
+
+  if (rpc_write_start_request(stored_conn, 1) < 0) {
+      std::cerr << "Error: rpc_write_start_request failed" << std::endl;
+      return;
+  }
+  if (rpc_write(stored_conn, &tmp->callback, sizeof(void *)) < 0) {
+      std::cerr << "Error: rpc_write failed on callback" << std::endl;
+      return;
+  }
+  if (rpc_write(stored_conn, &tmp->data, sizeof(void *)) < 0) {
+      std::cerr << "Error: rpc_write failed on data" << std::endl;
+      return;
+  }
 
-  printf("invoking host function %p\n", tmp->callback);
+  // Ensure request is fully sent before waiting for response
+  if (rpc_write_end(stored_conn) < 0) {
+      std::cerr << "Error: rpc_write_end failed" << std::endl;
+      return;
+  }
+
+  printf("hereeee %p\n", tmp->callback);
+
+  if (rpc_wait_for_response(stored_conn) < 0) {
+      std::cerr << "Error: rpc_wait_for_response failed" << std::endl;
+      return;
+  }
+
+  if (rpc_read(stored_conn, &scuda_intercept_result, sizeof(void*)) < 0) {
+      std::cerr << "Error: rpc_read failed on scuda_intercept_result" << std::endl;
+      return;
+  }
+
+  if (rpc_read_end(stored_conn) < 0) {
+      std::cerr << "Error: rpc_read_end failed" << std::endl;
+      return;
+  }
 
-  if (rpc_write_start_request(tmp->conn, 1) < 0 ||
-      rpc_write(tmp->conn, &tmp->callback, sizeof(void *)) < 0 ||
-      rpc_write(tmp->conn, &tmp->data, sizeof(void *)) < 0 ||
-      rpc_wait_for_response(tmp->conn) < 0 || rpc_read_end(tmp->conn) < 0)
-    std::cout << "failed to write memory: " << &faulting_address << std::endl;
+  std::cout << "RESULT IS: " << scuda_intercept_result << std::endl;
 }
 
 void append_host_func_ptr(const void *conn, void *ptr) {
@@ -172,7 +216,9 @@ void client_handler(int connfd) {
   printf("Client connected.\n");
 
   while (1) {
-    unsigned int op = rpc_dispatch(&conn, 0);
+    int op = rpc_dispatch(&conn, 0);
+
+    std::cout << "GOT OP: " << op << std::endl;
 
     auto opHandler = get_handler(op);
     if (opHandler(&conn) < 0) {
@@ -183,6 +229,7 @@ void client_handler(int connfd) {
   if (pthread_mutex_destroy(&conn.read_mutex) < 0 ||
       pthread_mutex_destroy(&conn.write_mutex) < 0)
     std::cerr << "Error destroying mutex." << std::endl;
+
   close(connfd);
 }