chore: dynamic unified pointers

kevmo314 · Dec 23, 2024 · da5e408 · da5e408
1 parent 644f22a
commit da5e408
Show file tree

Hide file tree

Showing 5 changed files with 14,497 additions and 5 deletions.
diff --git a/client.cpp b/client.cpp
@@ -96,6 +96,28 @@ static void segfault(int sig, siginfo_t* info, void* unused) {
     raise(SIGSEGV);
 }
 
+void maybe_copy_unified_arg(const int index, void* arg, enum cudaMemcpyKind kind)
+{
+    auto& unified_devices = conns[index].unified_devices;
+    auto found = unified_devices.find(arg);
+    if (found != unified_devices.end())
+    {
+        std::cout << "found unified arg pointer; copying..." << std::endl;
+
+        void* ptr = found->first;
+        size_t size = found->second;
+
+        cudaError_t res = cudaMemcpy(ptr, ptr, size, kind);
+
+        if (res != cudaSuccess) {
+            std::cerr << "cudaMemcpy failed: " << cudaGetErrorString(res) << std::endl;
+        } else {
+            std::cout << "Successfully copied " << size << " bytes" << std::endl;
+        }
+    }
+}
+
+
 static void set_segfault_handlers() {
     if (init > 0) {
         return;

diff --git a/codegen/codegen.py b/codegen/codegen.py
@@ -104,6 +104,9 @@ def client_rpc_write(self, f):
             )
         )
 
+    def client_unified_copy(self, f, direction):
+        f.write("    maybe_copy_unified_arg(0, (void*){name}, cudaMemcpyDeviceToHost);\n".format(name=self.parameter.name, direction=direction))
+
     @property
     def server_declaration(self) -> str:
         c = self.ptr.ptr_to.const
@@ -209,7 +212,24 @@ def client_rpc_write(self, f):
                     length=length,
                 )
             )
-
+
+    def client_unified_copy(self, f, direction):
+        f.write("    maybe_copy_unified_arg(0, (void*){name}, {direction});\n".format(name=self.parameter.name, direction=direction))
+
+        if isinstance(self.length, int):
+            f.write("    for (int i = 0; i < {name}; i++)\n".format(name=self.length))
+            f.write("       maybe_copy_unified_arg(0, (void*)&{name}[i], {direction});\n".format(name=self.parameter.name, direction=direction))
+        else:
+            if hasattr(self.length.type, "ptr_to"):
+                f.write("    for (int i = 0; i < static_cast<int>(*{name}); i++)\n".format(name=self.length.name))
+                f.write("       maybe_copy_unified_arg(0, (void*)&{name}[i], {direction});\n".format(name=self.parameter.name, direction=direction))
+            else:
+                if hasattr(self.parameter.type, "ptr_to"):
+                    f.write("    for (int i = 0; i < static_cast<int>({name}); i++)\n".format(name=self.length.name))
+                    f.write("       maybe_copy_unified_arg(0, (void*)&{name}[i], {direction});\n".format(name=self.parameter.name, direction=direction))
+                else:
+                    f.write("    for (int i = 0; i < static_cast<int>({name}); i++)\n".format(name=self.length.name))
+                    f.write("       maybe_copy_unified_arg(0, (void*){name}[i], {direction});\n".format(name=self.parameter.name, direction=direction))
 
     @property
     def server_declaration(self) -> str:
@@ -330,6 +350,9 @@ def client_rpc_write(self, f):
     def server_declaration(self) -> str:
         return f"    {self.ptr.format()} {self.parameter.name};\n" + \
                 f"    std::size_t {self.parameter.name}_len;\n"
+
+    def client_unified_copy(self, f, direction):
+        f.write("    maybe_copy_unified_arg(0, (void*){name}, {direction});\n".format(name=self.parameter.name, direction=direction))
 
     def server_rpc_read(self, f, index) -> Optional[str]:
         if not self.send:
@@ -415,6 +438,12 @@ def server_declaration(self) -> str:
             return f"   {self.type_.format().replace("const", "")} {self.parameter.name};\n"
         else:
             return f"    {self.type_.format()} {self.parameter.name};\n"
+
+    def client_unified_copy(self, f, direction):
+        if isinstance(self.type_, Pointer):
+            f.write("    maybe_copy_unified_arg(0, (void*){name}, {direction});\n".format(name=self.parameter.name, direction=direction))
+        else:
+            f.write("    maybe_copy_unified_arg(0, (void*)&{name}, {direction});\n".format(name=self.parameter.name, direction=direction))
 
     def server_rpc_read(self, f):
         if not self.send:
@@ -486,6 +515,9 @@ def server_rpc_read(self, f):
                 param_type=self.type_.ptr_to.format(),
             )
         )
+
+    def client_unified_copy(self, f, direction):
+        f.write("    maybe_copy_unified_arg(0, (void*){name}, {direction});\n".format(name=self.parameter.name, direction=direction))
 
     @property
     def server_reference(self) -> str:
@@ -761,6 +793,7 @@ def main():
             "extern int rpc_wait_for_response(const int index);\n"
             "extern int rpc_read(const int index, void *data, const std::size_t size);\n"
             "extern int rpc_end_response(const int index, void *return_value);\n"
+            "void maybe_copy_unified_arg(const int index, void *arg, enum cudaMemcpyKind kind);\n"
             "extern int rpc_close();\n\n"
         )
         for function, annotation, operations, disabled in functions_with_annotations:
@@ -798,6 +831,9 @@ def main():
             )
             f.write("{\n")
 
+            for operation in operations:
+                operation.client_unified_copy(f, "cudaMemcpyHostToDevice")
+
             f.write(
                 "    {return_type} return_value;\n".format(
                     return_type=function.return_type.format()
@@ -841,12 +877,14 @@ def main():
                 )
             )
 
+            for operation in operations:
+                operation.client_unified_copy(f, "cudaMemcpyDeviceToHost")
+
             if function.name.format() == "nvmlShutdown":
                 f.write("    if (rpc_close() < 0)\n")
                 f.write("        return {error_return};\n".format(error_return=error_const(function.return_type.format())))
 
             f.write("    return return_value;\n")
-
             f.write("}\n\n")
 
         f.write("std::unordered_map<std::string, void *> functionMap = {\n")