chore: check parent alloc

kevmo314 · Dec 24, 2024 · 322f75c · 322f75c
1 parent da5e408
commit 322f75c
Show file tree

Hide file tree

Showing 3 changed files with 358 additions and 346 deletions.
diff --git a/client.cpp b/client.cpp
@@ -96,6 +96,16 @@ static void segfault(int sig, siginfo_t* info, void* unused) {
     raise(SIGSEGV);
 }
 
+int is_unified_pointer(const int index, void* arg)
+{
+    auto& unified_devices = conns[index].unified_devices;
+    auto found = unified_devices.find(arg);
+    if (found != unified_devices.end())
+        return 1;
+
+    return 0;
+}
+
 void maybe_copy_unified_arg(const int index, void* arg, enum cudaMemcpyKind kind)
 {
     auto& unified_devices = conns[index].unified_devices;

diff --git a/codegen/codegen.py b/codegen/codegen.py
@@ -217,18 +217,18 @@ 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("    for (int i = 0; i < {name} && is_unified_pointer(0, (void*){param}); i++)\n".format(param=self.parameter.name, 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("    for (int i = 0; i < static_cast<int>(*{name}) && is_unified_pointer(0, (void*){param}); i++)\n".format(param=self.parameter.name, 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("    for (int i = 0; i < static_cast<int>({name}) && is_unified_pointer(0, (void*){param}); i++)\n".format(param=self.parameter.name, 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("    for (int i = 0; i < static_cast<int>({name}) && is_unified_pointer(0, (void*){param}); i++)\n".format(param=self.parameter.name, name=self.length.name))
                     f.write("       maybe_copy_unified_arg(0, (void*){name}[i], {direction});\n".format(name=self.parameter.name, direction=direction))
 
     @property
@@ -791,6 +791,7 @@ def main():
             "extern int rpc_write(const int index, const void *data, const std::size_t size);\n"
             "extern int rpc_end_request(const int index);\n"
             "extern int rpc_wait_for_response(const int index);\n"
+            "extern int is_unified_pointer(const int index, void* arg);\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"