rpc : add support for multiple devices

Allow rpc-server to expose multiple devices from a single endpoint.
Change RPC protocol to include device identifier where needed.
Add new API to get the device count from an RPC endpoint.

closes: ggml-org#15210

Author: rgerganov

common/arg.cpp

@@ -1721,17 +1721,26 @@ static void add_rpc_devices(const std::string & servers) {
     if (!rpc_reg) {
         throw std::invalid_argument("failed to find RPC backend");
     }
-    typedef ggml_backend_dev_t (*ggml_backend_rpc_add_device_t)(const char * endpoint);
+    typedef int (*ggml_backend_rpc_get_device_count_t)(const char * endpoint);
+    ggml_backend_rpc_get_device_count_t ggml_backend_rpc_get_device_count_fn =
+        (ggml_backend_rpc_get_device_count_t) ggml_backend_reg_get_proc_address(rpc_reg, "ggml_backend_rpc_get_device_count");
+    if (!ggml_backend_rpc_get_device_count_fn) {
+        throw std::invalid_argument("failed to find RPC device count function");
+    }
+    typedef ggml_backend_dev_t (*ggml_backend_rpc_add_device_t)(const char * endpoint, uint32_t device);
     ggml_backend_rpc_add_device_t ggml_backend_rpc_add_device_fn = (ggml_backend_rpc_add_device_t) ggml_backend_reg_get_proc_address(rpc_reg, "ggml_backend_rpc_add_device");
     if (!ggml_backend_rpc_add_device_fn) {
         throw std::invalid_argument("failed to find RPC device add function");
     }
     for (const auto & server : rpc_servers) {
-        ggml_backend_dev_t dev = ggml_backend_rpc_add_device_fn(server.c_str());
-        if (dev) {
-            ggml_backend_device_register(dev);
-        } else {
-            throw std::invalid_argument("failed to register RPC device");
+        int dev_count = ggml_backend_rpc_get_device_count_fn(server.c_str());
+        for (int i = 0; i < dev_count; i++) {
+            ggml_backend_dev_t dev = ggml_backend_rpc_add_device_fn(server.c_str(), i);
+            if (dev) {
+                ggml_backend_device_register(dev);
+            } else {
+                throw std::invalid_argument("failed to register RPC device");
+            }
         }
     }
 }

ggml/include/ggml-rpc.h

@@ -7,26 +7,26 @@
 extern "C" {
 #endif
 
-#define RPC_PROTO_MAJOR_VERSION    2
+#define RPC_PROTO_MAJOR_VERSION    3
 #define RPC_PROTO_MINOR_VERSION    0
 #define RPC_PROTO_PATCH_VERSION    0
 #define GGML_RPC_MAX_SERVERS       16
 
 // backend API
-GGML_BACKEND_API ggml_backend_t ggml_backend_rpc_init(const char * endpoint);
+GGML_BACKEND_API ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device);
 GGML_BACKEND_API bool ggml_backend_is_rpc(ggml_backend_t backend);
 
-GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint);
+GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, uint32_t device);
 
-GGML_BACKEND_API void ggml_backend_rpc_get_device_memory(const char * endpoint, size_t * free, size_t * total);
+GGML_BACKEND_API void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total);
 
-GGML_BACKEND_API void ggml_backend_rpc_start_server(ggml_backend_t backend, const char * endpoint,
-                                                    const char * cache_dir,
-                                                    size_t free_mem, size_t total_mem);
+GGML_BACKEND_API void ggml_backend_rpc_start_server(const char * endpoint, const char * cache_dir,
+                                                    size_t dev_count, ggml_backend_t * backends,
+                                                    size_t * free_mem, size_t * total_mem);
 
 GGML_BACKEND_API ggml_backend_reg_t ggml_backend_rpc_reg(void);
-
-GGML_BACKEND_API ggml_backend_dev_t ggml_backend_rpc_add_device(const char * endpoint);
+GGML_BACKEND_API int ggml_backend_rpc_get_device_count(const char * endpoint);
+GGML_BACKEND_API ggml_backend_dev_t ggml_backend_rpc_add_device(const char * endpoint, uint32_t device);
 
 #ifdef  __cplusplus
 }