diff --git a/discover/gpu_info_vulkan.c b/discover/gpu_info_vulkan.c index 0929fdee5..8bb0a9d3a 100644 --- a/discover/gpu_info_vulkan.c +++ b/discover/gpu_info_vulkan.c @@ -2,6 +2,85 @@ #include "gpu_info_vulkan.h" #include +#include +#include + +#define INITIAL_ARRAY_SIZE 10 + +// Function to parse an environment variable into a list of int values. +// Returns a pointer to the allocated array, and stores the count in out_count. +// Returns NULL in case of any error. +int* parse_envvar_to_int_list(const char* envvar_name, size_t *out_count) { + char *env_str = getenv(envvar_name); + if (env_str == NULL) { + *out_count = 0; + return NULL; + } + + // Duplicate the string since strtok modifies it. + char *tmp = strdup(env_str); + if (!tmp) { + *out_count = 0; + return NULL; + } + + size_t capacity = INITIAL_ARRAY_SIZE; + size_t count = 0; + int *list = malloc(capacity * sizeof(uint32_t)); + if (!list) { + free(tmp); + *out_count = 0; + return NULL; + } + + char *token = strtok(tmp, ","); + while (token != NULL) { + char *endptr = NULL; + errno = 0; + unsigned long val = strtoul(token, &endptr, 10); + if (errno != 0 || endptr == token) { + free(list); + free(tmp); + *out_count = 0; + return NULL; + } + // Optional: Check trailing characters. + while (*endptr != '\0') { + if (!isspace((unsigned char)*endptr)) { + free(list); + free(tmp); + *out_count = 0; + return NULL; + } + endptr++; + } + if (val > UINT32_MAX) { + free(list); + free(tmp); + *out_count = 0; + return NULL; + } + + // Save the value, reallocating if necessary. + if (count == capacity) { + capacity *= 2; + int *temp = realloc(list, capacity * sizeof(uint32_t)); + if (!temp) { + free(list); + free(tmp); + *out_count = 0; + return NULL; + } + list = temp; + } + list[count++] = (int)val; + token = strtok(NULL, ","); + } + + free(tmp); + *out_count = count; + return list; +} int is_extension_supported(vk_handle_t* rh, VkPhysicalDevice device, char* extension) { VkPhysicalDeviceProperties properties = {}; @@ -112,10 +191,21 @@ void vk_init(char* vk_lib_path, vk_init_resp_t *resp) { return; } + size_t visDevIdCount; + int* visDevIds = parse_envvar_to_int_list("GGML_VK_VISIBLE_DEVICES", &visDevIdCount); + resp->err = NULL; resp->ch.vk = instance; resp->ch.num_devices = deviceCount; resp->num_devices = deviceCount; + + if (visDevIds && visDevIdCount > 0) { + resp->ch.num_visible_devices = visDevIdCount; + resp->ch.visible_devices = visDevIds; + } else { + resp->ch.num_visible_devices = -1; + resp->ch.visible_devices = NULL; + } } int vk_device_is_supported(vk_handle_t rh, int i) { @@ -192,6 +282,24 @@ void vk_check_vram(vk_handle_t rh, int i, mem_info_t *resp) { device_props2.pNext = &id_props; (*rh.vkGetPhysicalDeviceProperties2)(devices[i], &device_props2); + if (rh.num_visible_devices > 0) { + LOG(rh.verbose, "Checking if device %d is visible\n", i); + int is_visible = 0; + for (uint32_t visDevId = 0; visDevId < rh.num_visible_devices; visDevId++) { + if (i == rh.visible_devices[visDevId]) { + LOG(rh.verbose, "Device %d is visible!\n", i); + is_visible = 1; + break; + } + } + if (!is_visible) { + LOG(rh.verbose, "Device %d is NOT visible!\n", i); + free(devices); + resp->err = strdup("device is hidden with GGML_VK_VISIBLE_DEVICES"); + return; + } + } + VkPhysicalDeviceMemoryBudgetPropertiesEXT physical_device_memory_budget_properties = {}; physical_device_memory_budget_properties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_BUDGET_PROPERTIES_EXT; physical_device_memory_budget_properties.pNext = NULL; @@ -241,6 +349,10 @@ void vk_release(vk_handle_t rh) { (*rh.vkDestroyInstance)(rh.vk, NULL); UNLOAD_LIBRARY(rh.vk_handle); rh.vk_handle = NULL; + + if (rh.visible_devices) { + free(rh.visible_devices); + } } #endif // __APPLE__ diff --git a/discover/gpu_info_vulkan.h b/discover/gpu_info_vulkan.h index 3cd8b0b39..c249d9855 100644 --- a/discover/gpu_info_vulkan.h +++ b/discover/gpu_info_vulkan.h @@ -434,6 +434,9 @@ typedef struct { VkInstance vk; int num_devices; + int num_visible_devices; + int* visible_devices; + void (*vkGetPhysicalDeviceProperties)( VkPhysicalDevice physicalDevice, VkPhysicalDeviceProperties* pProperties);