WIP - shared CUDA memory

Signed-off-by: Cocoa <[email protected]>
cocoa-xu · Jun 8, 2024 · 5801ce3 · 5801ce3
1 parent 9b67f7a
commit 5801ce3
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 17 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -66,28 +66,39 @@ endif()
 if(WIN32)
     string(REPLACE "\\" "/" C_SRC "${C_SRC}")
 endif()
-file(GLOB evision_bindings_main "${C_SRC}/*.hpp" "${C_SRC}/*.h" "${C_SRC}/*.cpp" "${C_SRC}/*.cc")
+file(GLOB evision_bindings_main "${C_SRC}/*.hpp" "${C_SRC}/*.h" "${C_SRC}/*.cpp")
 file(GLOB evision_custom_headers "${C_SRC}/evision_custom_headers/*.hpp")
 file(GLOB evision_modules "${C_SRC}/modules/*.hpp" "${C_SRC}/modules/*.h")
 file(GLOB evision_backend "${C_SRC}/modules/evision_backend/*.h")
+set(EVISION_LIBRARIES "${OpenCV_LIBS}")
+
+if(EVISION_ENABLE_CUDA)
+    find_package(CUDA REQUIRED)
+    include_directories("${CUDA_INCLUDE_DIRS}")
+    execute_process(COMMAND bash -c "nvcc -shared -Xcompiler -fPIC -DCUDA_ENABLED -c \"${C_SRC}/modules/evision_cuda.cc\" -o \"${C_SRC}/modules/evision_cuda.o\"" RESULT_VARIABLE STATUS)
+    if(STATUS STREQUAL "0")
+        message(STATUS "Successfully compiled evision_cuda.cc using nvcc")
+    else()
+        message(FATAL_ERROR "Failed to compile evision_cuda.cc using nvcc")
+    endif()
+    list(APPEND evision_bindings_main "${C_SRC}/modules/evision_cuda.o")
+    list(APPEND EVISION_LIBRARIES cudart)
+    add_compile_definitions(CUDA_ENABLED)
+else()
+    list(APPEND evision_bindings_main "${C_SRC}/modules/evision_cuda.cc" "${C_SRC}/modules/evision_cuda.h")
+endif()
 
 add_library(evision SHARED
     ${evision_bindings_main}
     ${evision_custom_headers}
     ${evision_modules}
     ${evision_backend}
 )
-
-if(EVISION_ENABLE_CUDA)
-    find_package(CUDA REQUIRED)
-    include_directories("${CUDA_INCLUDE_DIRS}")
-    target_compile_definitions(evision PRIVATE EVISION_HAVE_CUDA)
-endif()
 install(
     TARGETS evision
     RUNTIME DESTINATION "${PRIV_DIR}"
 )
-target_link_libraries(evision "${OpenCV_LIBS}")
+target_link_libraries(evision "${EVISION_LIBRARIES}")
 set_property(TARGET evision PROPERTY CXX_STANDARD 14)
 set_property(TARGET evision PROPERTY POSITION_INDEPENDENT_CODE ON)
 set_target_properties(evision PROPERTIES PREFIX "")

diff --git a/c_src/modules/evision_cuda.cc b/c_src/modules/evision_cuda.cc
@@ -4,6 +4,8 @@
 #include <cuda_runtime.h>
 
 #include <cstring>
+#include <cstdint>
+#include <vector>
 #include <iostream>
 
 std::pair<std::vector<unsigned char>, int> get_cuda_ipc_handle(std::uintptr_t ptr) {

diff --git a/c_src/modules/evision_mat.h b/c_src/modules/evision_mat.h
@@ -65,15 +65,13 @@ static ERL_NIF_TERM evision_cv_mat_to_pointer(ErlNifEnv *env, int argc, const ER
                     pointer_vec.push_back(bytePtr[i]);
                 }
             }
-            // todo: use nvcc to compile `evision_cuda.{cc,h}`
-            //
-            // else if (pointer_kind == "cuda_ipc") {
-            //     auto result = get_cuda_ipc_handle(ptr);
-            //     if (result.second) {
-            //         return evision::nif::error(env, "Unable to get cuda IPC handle");
-            //     }
-            //     pointer_vec = result.first;
-            // }
+            else if (pointer_kind == "cuda_ipc") {
+                auto result = get_cuda_ipc_handle(ptr);
+                if (result.second) {
+                    return evision::nif::error(env, "Unable to get cuda IPC handle");
+                }
+                pointer_vec = result.first;
+            }
             if (pointer_vec.size() == 0) {
                 return enif_make_badarg(env);
             }