allow users to explicitly specify device id

Signed-off-by: Cocoa <[email protected]>
cocoa-xu · Jun 17, 2024 · e6d8d31 · e6d8d31
1 parent 9a4ebcd
commit e6d8d31
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,7 +4,7 @@
 [Browse the Repository](https://github.com/cocoa-xu/evision/tree/v0.2.3) | [Released Assets](https://github.com/cocoa-xu/evision/releases/tag/v0.2.3)
 
 ### Added
-- Experimental support for use shared CUDA memory via `Evision.CUDA.GpuMat.from_pointer/3`.
+- Experimental support for use shared CUDA memory via `Evision.CUDA.GpuMat.from_pointer/{3,4}`.
 
 ## v0.2.2 (2024-06-15)
 [Browse the Repository](https://github.com/cocoa-xu/evision/tree/v0.2.2) | [Released Assets](https://github.com/cocoa-xu/evision/releases/tag/v0.2.2)

diff --git a/c_src/modules/evision_gpumat.h b/c_src/modules/evision_gpumat.h
@@ -72,11 +72,12 @@ static ERL_NIF_TERM evision_cv_cuda_cuda_GpuMat_from_pointer(ErlNifEnv *env, int
         std::vector<int64_t> device_pointer;
         std::vector<int64_t> shape;
         std::string dtype;
+        int device_id = 0;
 
-        if (!evision_to_safe(env, evision_get_kw(env, erl_terms, "device_pointer"), device_pointer, ArgInfo("device_pointer", 1))) {
+        if (!evision_to_safe(env, evision_get_kw(env, erl_terms, "device_pointer"), device_pointer, ArgInfo("device_pointer", 0))) {
             return enif_make_badarg(env);
         }
-        if (!evision_to_safe(env, evision_get_kw(env, erl_terms, "dtype"), dtype, ArgInfo("dtype", 1))) {
+        if (!evision_to_safe(env, evision_get_kw(env, erl_terms, "dtype"), dtype, ArgInfo("dtype", 0))) {
             return enif_make_badarg(env);
         }
         if (!evision::nif::get_tuple(env, evision_get_kw(env, erl_terms, "shape"), shape)) {
@@ -85,6 +86,7 @@ static ERL_NIF_TERM evision_cv_cuda_cuda_GpuMat_from_pointer(ErlNifEnv *env, int
         if (shape.size() > 3 || shape.size() == 0) {
             return evision::nif::error(env, "GpuMat expects shape to be 1 <= tuple_size(shape) <= 3");
         }
+        evision_to_safe(env, evision_get_kw(env, erl_terms, "device_id"), device_id, ArgInfo("device_id", 0x8));
 
         int height;
         int width = 1;
@@ -111,7 +113,7 @@ static ERL_NIF_TERM evision_cv_cuda_cuda_GpuMat_from_pointer(ErlNifEnv *env, int
                 bytePtr[i] = device_pointer[i];
             }
         } else {
-            auto result = get_pointer_for_ipc_handle(device_pointer);
+            auto result = get_pointer_for_ipc_handle(device_pointer, device_id);
             if (result.second) {
                 return evision::nif::error(env, "Unable to get pointer for IPC handle.");
             }

diff --git a/py_src/evision_extra_functions.py b/py_src/evision_extra_functions.py
@@ -64,7 +64,54 @@ def to_pointer(%{ref: ref}, opts) when is_list(opts) do
   """
   @spec from_pointer(list(integer()), atom() | {atom(), integer()}, tuple()) :: Evision.CUDA.GpuMat.t() | {:error, String.t()}
   def from_pointer(device_pointer, dtype, shape) when is_list(device_pointer) and is_tuple(shape) do
-    :evision_nif.cuda_cuda_GpuMat_from_pointer([device_pointer: device_pointer, dtype: compact_type(dtype), shape: shape])
+    positional = [
+      device_pointer: device_pointer,
+      dtype: compact_type(dtype),
+      shape: shape
+    ]
+    :evision_nif.cuda_cuda_GpuMat_from_pointer(positional)
+    |> to_struct()
+  end
+
+  @doc """
+  Create CUDA GpuMat from a shared CUDA device pointer
+  
+  ##### Positional Arguments
+
+  - **device_pointer**, `list(integer())`.
+  
+    This can be either a local pointer or an IPC pointer.
+    
+    However, please note that IPC pointers have to be generated from 
+    another OS process (Erlang process doesn't count).
+  
+  - **dtype**, `tuple() | atom()`
+  
+    Data type.
+    
+  - **shape**, `tuple()`
+  
+    The shape of the shared image. It's expected to be either
+    
+    - `{height, width, channels}`, for any 2D image that has 1 or multiple channels
+    - `{height, width}`, for any 1-channel 2D image
+    - `{rows}`
+    
+  ##### Keyword Arguments
+
+  - **device_id**, `non_neg_integer`. 
+  
+    GPU Device ID, default to `0`.
+  """
+  @spec from_pointer(list(integer()), atom() | {atom(), integer()}, tuple(), [device_id: non_neg_integer()]) :: Evision.CUDA.GpuMat.t() | {:error, String.t()}
+  def from_pointer(device_pointer, dtype, shape, opts) when is_list(device_pointer) and is_tuple(shape) and is_list(opts) do
+    opts = Keyword.validate!(opts || [], [device_id: 0])
+    positional = [
+      device_pointer: device_pointer,
+      dtype: compact_type(dtype),
+      shape: shape
+    ]
+    :evision_nif.cuda_cuda_GpuMat_from_pointer(positional ++ Evision.Internal.Structurise.from_struct(opts))
     |> to_struct()
   end
 '''