// SPDX-License-Identifier: GPL-3.0-only
/*
 * MVGAL Vulkan Layer Implementation
 *
 * Multi-Vendor GPU Aggregation Layer for Linux
 * Vulkan interception layer for multi-GPU aggregation
 */
#include "vk_layer.h"
#include "mvgal_shader_translate.h"

#include <inttypes.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>

#include "mvgal/mvgal.h"
#include "mvgal/mvgal_scheduler.h"
#include "mvgal/mvgal_gpu.h"

static uint32_t mvgal_vk_layer_active_gpu_count(void)
{
    static uint32_t cached_count;
    static bool cached;

    if (cached) {
        return cached_count > 0 ? cached_count : 1U;
    }

    if (mvgal_init(0) == MVGAL_SUCCESS) {
        int32_t count = mvgal_gpu_get_count();
        if (count > 0) {
            cached_count = (uint32_t)count;
        }
    }

    if (cached_count == 0) {
        cached_count = 1U;
    }

    cached = true;
    return cached_count;
}

static bool env_truthy(const char *name)
{
    const char *value = getenv(name);

    if (value == NULL || value[0] == '\0') {
        return false;
    }

    if (strcmp(value, "0") == 0 || strcasecmp(value, "false") == 0 ||
        strcasecmp(value, "off") == 0 || strcasecmp(value, "no") == 0) {
        return false;
    }

    return true;
}

static const char *mvgal_vk_layer_strategy_name(void)
{
    const char *strategy = getenv("MVGAL_STRATEGY");

    if (strategy != NULL && strategy[0] != '\0') {
        return strategy;
    }

    return "single";
}

static bool mvgal_vk_layer_strategy_is(const char *name)
{
    return strcasecmp(mvgal_vk_layer_strategy_name(), name) == 0;
}

/**
 * Multi-GPU virtualization is opt-in. Default is pass-through to the vendor ICD.
 * Enable with MVGAL_AGGREGATE=1 or an explicit non-single MVGAL_STRATEGY.
 */
static bool mvgal_vk_layer_multi_gpu_enabled(void)
{
    if (env_truthy("MVGAL_AGGREGATE")) {
        return true;
    }

    if (mvgal_vk_layer_strategy_is("single") ||
        mvgal_vk_layer_strategy_is("single_gpu")) {
        return false;
    }

    /* Only aggregate when the user sets a multi-GPU strategy explicitly. */
    return getenv("MVGAL_STRATEGY") != NULL && getenv("MVGAL_STRATEGY")[0] != '\0';
}

static bool mvgal_vk_layer_afr_enabled(void)
{
    return mvgal_vk_layer_multi_gpu_enabled() && mvgal_vk_layer_strategy_is("afr");
}

mvgal_layer_state_t g_mvgal_layer_state = {
    .lock = PTHREAD_MUTEX_INITIALIZER,
    .instances = NULL,
    .devices = NULL,
    .queues = NULL,
    .physical_devices = NULL,
    .submit_count = 0,
    .physical_device_count = 0,
};

bool mvgal_vk_layer_debug_enabled(void)
{
    return env_truthy("MVGAL_VULKAN_DEBUG") || env_truthy("UMGAL_VULKAN_DEBUG");
}

void mvgal_vk_layer_log(const char *fmt, ...)
{
    FILE *stream = stderr;
    const char *log_path = getenv("MVGAL_VULKAN_LOG_PATH");
    va_list ap;

    if (!mvgal_vk_layer_debug_enabled() && log_path == NULL) {
        return;
    }

    if (log_path != NULL && log_path[0] != '\0') {
        stream = fopen(log_path, "a");
        if (stream == NULL) {
            stream = stderr;
        }
    }

    fprintf(stream, "[mvgal-vk-layer] ");
    va_start(ap, fmt);
    vfprintf(stream, fmt, ap);
    va_end(ap);
    fputc('\n', stream);
    fflush(stream);

    if (stream != stderr) {
        fclose(stream);
    }
}

void mvgal_vk_layer_note_submit(const char *entrypoint, uint32_t submit_count)
{
    uint64_t total = atomic_fetch_add_explicit(
        &g_mvgal_layer_state.submit_count,
        1U,
        memory_order_relaxed
    ) + 1U;

    mvgal_vk_layer_log(
        "%s submit_count=%u total_submissions=%" PRIu64,
        entrypoint,
        submit_count,
        total
    );
}

static void fill_layer_properties(VkLayerProperties *properties)
{
    memset(properties, 0, sizeof(*properties));
    snprintf(properties->layerName, sizeof(properties->layerName), "%s",
             MVGAL_VK_LAYER_NAME);
    properties->specVersion = VK_API_VERSION_1_3;
    properties->implementationVersion = 1;
    snprintf(properties->description, sizeof(properties->description), "%s",
             MVGAL_VK_LAYER_DESCRIPTION);
}

static bool layer_name_matches(const char *layer_name)
{
    return layer_name != NULL &&
           strcmp(layer_name, MVGAL_VK_LAYER_NAME) == 0;
}

static void remove_queues_for_device_locked(mvgal_device_dispatch_t *device_dispatch)
{
    mvgal_queue_dispatch_t **cursor = &g_mvgal_layer_state.queues;

    while (*cursor != NULL) {
        if ((*cursor)->device_dispatch == device_dispatch) {
            mvgal_queue_dispatch_t *victim = *cursor;

            *cursor = victim->next;
            free(victim);
            continue;
        }

        cursor = &(*cursor)->next;
    }
}

static void remove_physical_devices_for_instance_locked(VkInstance instance)
{
    mvgal_physical_device_dispatch_t **cursor = &g_mvgal_layer_state.physical_devices;

    while (*cursor != NULL) {
        if ((*cursor)->instance == instance) {
            mvgal_physical_device_dispatch_t *victim = *cursor;

            *cursor = victim->next;
            free(victim);
            continue;
        }

        cursor = &(*cursor)->next;
    }
}

mvgal_instance_dispatch_t *mvgal_vk_find_instance_dispatch(VkInstance instance)
{
    mvgal_instance_dispatch_t *dispatch;

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    dispatch = g_mvgal_layer_state.instances;
    while (dispatch != NULL && dispatch->instance != instance) {
        dispatch = dispatch->next;
    }
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);

    return dispatch;
}

mvgal_device_dispatch_t *mvgal_vk_find_device_dispatch(VkDevice device)
{
    mvgal_device_dispatch_t *dispatch;

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    dispatch = g_mvgal_layer_state.devices;
    while (dispatch != NULL && dispatch->device != device) {
        dispatch = dispatch->next;
    }
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);

    return dispatch;
}

mvgal_queue_dispatch_t *mvgal_vk_find_queue_dispatch(VkQueue queue)
{
    mvgal_queue_dispatch_t *dispatch;

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    dispatch = g_mvgal_layer_state.queues;
    while (dispatch != NULL && dispatch->queue != queue) {
        dispatch = dispatch->next;
    }
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);

    return dispatch;
}

mvgal_physical_device_dispatch_t *mvgal_vk_find_physical_device_dispatch(
    VkPhysicalDevice physical_device
)
{
    mvgal_physical_device_dispatch_t *dispatch;

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    dispatch = g_mvgal_layer_state.physical_devices;
    while (dispatch != NULL && dispatch->physical_device != physical_device) {
        dispatch = dispatch->next;
    }
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);

    return dispatch;
}

VkResult mvgal_vk_register_instance(
    VkInstance instance,
    PFN_vkGetInstanceProcAddr next_gipa,
    PFN_GetPhysicalDeviceProcAddr next_gpipa
)
{
    mvgal_instance_dispatch_t *dispatch;

    dispatch = calloc(1, sizeof(*dispatch));
    if (dispatch == NULL) {
        return VK_ERROR_OUT_OF_HOST_MEMORY;
    }

    dispatch->instance = instance;
    dispatch->next_gipa = next_gipa;
    dispatch->next_gpipa = next_gpipa;
    dispatch->destroy_instance =
        (PFN_vkDestroyInstance)next_gipa(instance, "vkDestroyInstance");
    dispatch->create_device =
        (PFN_vkCreateDevice)next_gipa(instance, "vkCreateDevice");
    dispatch->enumerate_physical_devices =
        (PFN_vkEnumeratePhysicalDevices)next_gipa(instance, "vkEnumeratePhysicalDevices");
    dispatch->get_physical_device_properties =
        (PFN_vkGetPhysicalDeviceProperties)next_gipa(instance, "vkGetPhysicalDeviceProperties");
    dispatch->get_physical_device_features =
        (PFN_vkGetPhysicalDeviceFeatures)next_gipa(instance, "vkGetPhysicalDeviceFeatures");
    dispatch->get_physical_device_memory_properties =
        (PFN_vkGetPhysicalDeviceMemoryProperties)next_gipa(instance, "vkGetPhysicalDeviceMemoryProperties");
    dispatch->get_physical_device_queue_family_properties =
        (PFN_vkGetPhysicalDeviceQueueFamilyProperties)next_gipa(instance, "vkGetPhysicalDeviceQueueFamilyProperties");
    dispatch->get_physical_device_properties2 =
        (PFN_vkGetPhysicalDeviceProperties2)next_gipa(instance, "vkGetPhysicalDeviceProperties2");
    dispatch->get_physical_device_features2 =
        (PFN_vkGetPhysicalDeviceFeatures2)next_gipa(instance, "vkGetPhysicalDeviceFeatures2");
    dispatch->get_physical_device_memory_properties2 =
        (PFN_vkGetPhysicalDeviceMemoryProperties2)next_gipa(instance, "vkGetPhysicalDeviceMemoryProperties2");
    dispatch->enumerate_physical_device_groups =
        (PFN_vkEnumeratePhysicalDeviceGroups)next_gipa(instance, "vkEnumeratePhysicalDeviceGroups");
    dispatch->enumerate_physical_device_groups_khx =
        (PFN_vkEnumeratePhysicalDeviceGroupsKHX)next_gipa(instance, "vkEnumeratePhysicalDeviceGroupsKHX");

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    dispatch->next = g_mvgal_layer_state.instances;
    g_mvgal_layer_state.instances = dispatch;
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);

    return VK_SUCCESS;
}

VkResult mvgal_vk_register_device(
    VkDevice device,
    PFN_vkGetDeviceProcAddr next_gdpa
)
{
    mvgal_device_dispatch_t *dispatch;

    dispatch = calloc(1, sizeof(*dispatch));
    if (dispatch == NULL) {
        return VK_ERROR_OUT_OF_HOST_MEMORY;
    }

    dispatch->device = device;
    dispatch->next_gdpa = next_gdpa;
    dispatch->destroy_device =
        (PFN_vkDestroyDevice)next_gdpa(device, "vkDestroyDevice");
    dispatch->get_device_queue =
        (PFN_vkGetDeviceQueue)next_gdpa(device, "vkGetDeviceQueue");
    dispatch->get_device_queue2 =
        (PFN_vkGetDeviceQueue2)next_gdpa(device, "vkGetDeviceQueue2");
    dispatch->queue_submit =
        (PFN_vkQueueSubmit)next_gdpa(device, "vkQueueSubmit");
#ifdef VK_VERSION_1_3
    dispatch->queue_submit2 =
        (PFN_vkQueueSubmit2)next_gdpa(device, "vkQueueSubmit2");
    dispatch->queue_submit2_khr =
        (PFN_vkQueueSubmit2KHR)next_gdpa(device, "vkQueueSubmit2KHR");
#endif
    dispatch->create_shader_module =
        (PFN_vkCreateShaderModule)next_gdpa(device, "vkCreateShaderModule");

    mvgal_spirv_cache_init(&dispatch->spirv_cache);

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    dispatch->next = g_mvgal_layer_state.devices;
    g_mvgal_layer_state.devices = dispatch;
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);

    return VK_SUCCESS;
}

VkResult mvgal_vk_register_queue(
    VkQueue queue,
    mvgal_device_dispatch_t *device_dispatch
)
{
    mvgal_queue_dispatch_t *dispatch;

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    dispatch = g_mvgal_layer_state.queues;
    while (dispatch != NULL) {
        if (dispatch->queue == queue) {
            dispatch->device_dispatch = device_dispatch;
            pthread_mutex_unlock(&g_mvgal_layer_state.lock);
            return VK_SUCCESS;
        }
        dispatch = dispatch->next;
    }

    dispatch = calloc(1, sizeof(*dispatch));
    if (dispatch == NULL) {
        pthread_mutex_unlock(&g_mvgal_layer_state.lock);
        return VK_ERROR_OUT_OF_HOST_MEMORY;
    }

    dispatch->queue = queue;
    dispatch->device_dispatch = device_dispatch;
    dispatch->next = g_mvgal_layer_state.queues;
    g_mvgal_layer_state.queues = dispatch;
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);

    return VK_SUCCESS;
}

VkResult mvgal_vk_register_physical_devices(
    VkInstance instance,
    mvgal_instance_dispatch_t *instance_dispatch,
    uint32_t physical_device_count,
    const VkPhysicalDevice *physical_devices
)
{
    uint32_t i;

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    for (i = 0; i < physical_device_count; ++i) {
        mvgal_physical_device_dispatch_t *dispatch =
            g_mvgal_layer_state.physical_devices;

        while (dispatch != NULL) {
            if (dispatch->physical_device == physical_devices[i]) {
                dispatch->instance = instance;
                dispatch->instance_dispatch = instance_dispatch;
                break;
            }
            dispatch = dispatch->next;
        }

        if (dispatch == NULL) {
            dispatch = calloc(1, sizeof(*dispatch));
            if (dispatch == NULL) {
                pthread_mutex_unlock(&g_mvgal_layer_state.lock);
                return VK_ERROR_OUT_OF_HOST_MEMORY;
            }

            dispatch->physical_device = physical_devices[i];
            dispatch->instance = instance;
            dispatch->instance_dispatch = instance_dispatch;
            dispatch->next = g_mvgal_layer_state.physical_devices;
            g_mvgal_layer_state.physical_devices = dispatch;
        }
    }
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);

    return VK_SUCCESS;
}

void mvgal_vk_unregister_instance(VkInstance instance)
{
    mvgal_instance_dispatch_t **cursor;

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    cursor = &g_mvgal_layer_state.instances;
    while (*cursor != NULL) {
        if ((*cursor)->instance == instance) {
            mvgal_instance_dispatch_t *victim = *cursor;

            *cursor = victim->next;
            remove_physical_devices_for_instance_locked(instance);
            free(victim);
            break;
        }
        cursor = &(*cursor)->next;
    }
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);
}

void mvgal_vk_unregister_device(VkDevice device)
{
    mvgal_device_dispatch_t **cursor;

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    cursor = &g_mvgal_layer_state.devices;
    while (*cursor != NULL) {
        if ((*cursor)->device == device) {
            mvgal_device_dispatch_t *victim = *cursor;

            *cursor = victim->next;
            remove_queues_for_device_locked(victim);
            free(victim);
            break;
        }
        cursor = &(*cursor)->next;
    }
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);
}

VkLayerInstanceCreateInfo *mvgal_vk_find_instance_layer_info(
    const VkInstanceCreateInfo *create_info,
    VkLayerFunction function
)
{
    VkLayerInstanceCreateInfo *info;

    info = (VkLayerInstanceCreateInfo *)create_info->pNext;
    while (info != NULL) {
        if (info->sType == VK_STRUCTURE_TYPE_LOADER_INSTANCE_CREATE_INFO &&
            info->function == function) {
            return info;
        }

        info = (VkLayerInstanceCreateInfo *)info->pNext;
    }

    return NULL;
}

VkLayerDeviceCreateInfo *mvgal_vk_find_device_layer_info(
    const VkDeviceCreateInfo *create_info,
    VkLayerFunction function
)
{
    VkLayerDeviceCreateInfo *info;

    info = (VkLayerDeviceCreateInfo *)create_info->pNext;
    while (info != NULL) {
        if (info->sType == VK_STRUCTURE_TYPE_LOADER_DEVICE_CREATE_INFO &&
            info->function == function) {
            return info;
        }

        info = (VkLayerDeviceCreateInfo *)info->pNext;
    }

    return NULL;
}

VKAPI_ATTR VkResult VKAPI_CALL vkEnumerateInstanceLayerProperties(
    uint32_t *pPropertyCount,
    VkLayerProperties *pProperties
)
{
    if (pPropertyCount == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    if (pProperties == NULL) {
        *pPropertyCount = 1;
        return VK_SUCCESS;
    }

    if (*pPropertyCount == 0) {
        return VK_INCOMPLETE;
    }

    fill_layer_properties(&pProperties[0]);
    *pPropertyCount = 1;
    return VK_SUCCESS;
}

VKAPI_ATTR VkResult VKAPI_CALL vkEnumerateDeviceLayerProperties(
    VkPhysicalDevice physical_device,
    uint32_t *pPropertyCount,
    VkLayerProperties *pProperties
)
{
    (void)physical_device;
    return vkEnumerateInstanceLayerProperties(pPropertyCount, pProperties);
}

VKAPI_ATTR VkResult VKAPI_CALL vkEnumerateInstanceExtensionProperties(
    const char *pLayerName,
    uint32_t *pPropertyCount,
    VkExtensionProperties *pProperties
)
{
    if (!layer_name_matches(pLayerName)) {
        return VK_ERROR_LAYER_NOT_PRESENT;
    }

    if (pPropertyCount != NULL) {
        *pPropertyCount = 0;
    }
    (void)pProperties;
    return VK_SUCCESS;
}

VKAPI_ATTR VkResult VKAPI_CALL vkEnumerateDeviceExtensionProperties(
    VkPhysicalDevice physical_device,
    const char *pLayerName,
    uint32_t *pPropertyCount,
    VkExtensionProperties *pProperties
)
{
    mvgal_physical_device_dispatch_t *dispatch;
    PFN_vkEnumerateDeviceExtensionProperties next_enumerate;

    if (layer_name_matches(pLayerName)) {
        if (pPropertyCount != NULL) {
            *pPropertyCount = 0;
        }
        (void)pProperties;
        return VK_SUCCESS;
    }

    dispatch = mvgal_vk_find_physical_device_dispatch(physical_device);
    if (dispatch == NULL || dispatch->instance_dispatch == NULL ||
        dispatch->instance_dispatch->next_gipa == NULL) {
        return VK_ERROR_LAYER_NOT_PRESENT;
    }

    next_enumerate = (PFN_vkEnumerateDeviceExtensionProperties)
        dispatch->instance_dispatch->next_gipa(
            dispatch->instance,
            "vkEnumerateDeviceExtensionProperties"
        );
    if (next_enumerate == NULL) {
        return VK_ERROR_LAYER_NOT_PRESENT;
    }

    return next_enumerate(
        physical_device,
        pLayerName,
        pPropertyCount,
        pProperties
    );
}

VKAPI_ATTR VkResult VKAPI_CALL vkEnumerateInstanceVersion(uint32_t *pApiVersion)
{
    if (pApiVersion == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    *pApiVersion = VK_API_VERSION_1_3;
    return VK_SUCCESS;
}

VKAPI_ATTR VkResult VKAPI_CALL vkCreateInstance(
    const VkInstanceCreateInfo *pCreateInfo,
    const VkAllocationCallbacks *pAllocator,
    VkInstance *pInstance
)
{
    VkLayerInstanceCreateInfo *layer_info;
    PFN_vkGetInstanceProcAddr next_gipa;
    PFN_GetPhysicalDeviceProcAddr next_gpipa;
    PFN_vkCreateInstance next_create_instance;
    VkResult result;

    layer_info = mvgal_vk_find_instance_layer_info(
        pCreateInfo,
        VK_LAYER_LINK_INFO
    );
    if (layer_info == NULL || layer_info->u.pLayerInfo == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    next_gipa = layer_info->u.pLayerInfo->pfnNextGetInstanceProcAddr;
    next_gpipa = layer_info->u.pLayerInfo->pfnNextGetPhysicalDeviceProcAddr;
    layer_info->u.pLayerInfo = layer_info->u.pLayerInfo->pNext;

    next_create_instance =
        (PFN_vkCreateInstance)next_gipa(VK_NULL_HANDLE, "vkCreateInstance");
    if (next_create_instance == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    result = next_create_instance(pCreateInfo, pAllocator, pInstance);
    if (result != VK_SUCCESS) {
        return result;
    }

    result = mvgal_vk_register_instance(*pInstance, next_gipa, next_gpipa);
    if (result != VK_SUCCESS) {
        mvgal_instance_dispatch_t *dispatch = mvgal_vk_find_instance_dispatch(*pInstance);
        if (dispatch != NULL && dispatch->destroy_instance != NULL) {
            dispatch->destroy_instance(*pInstance, pAllocator);
        }
        return result;
    }

    mvgal_vk_layer_log("created instance=%p", (void *)*pInstance);
    return VK_SUCCESS;
}

VKAPI_ATTR void VKAPI_CALL vkDestroyInstance(
    VkInstance instance,
    const VkAllocationCallbacks *pAllocator
)
{
    mvgal_instance_dispatch_t *dispatch = mvgal_vk_find_instance_dispatch(instance);

    if (dispatch != NULL && dispatch->destroy_instance != NULL) {
        dispatch->destroy_instance(instance, pAllocator);
    }

    mvgal_vk_unregister_instance(instance);
    mvgal_vk_layer_log("destroyed instance=%p", (void *)instance);
}

VKAPI_ATTR VkResult VKAPI_CALL vkCreateDevice(
    VkPhysicalDevice physicalDevice,
    const VkDeviceCreateInfo *pCreateInfo,
    const VkAllocationCallbacks *pAllocator,
    VkDevice *pDevice
)
{
    VkLayerDeviceCreateInfo *layer_info;
    PFN_vkGetInstanceProcAddr next_gipa;
    PFN_vkGetDeviceProcAddr next_gdpa;
    PFN_vkCreateDevice next_create_device;
    VkResult result;

    layer_info = mvgal_vk_find_device_layer_info(
        pCreateInfo,
        VK_LAYER_LINK_INFO
    );
    if (layer_info == NULL || layer_info->u.pLayerInfo == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    next_gipa = layer_info->u.pLayerInfo->pfnNextGetInstanceProcAddr;
    next_gdpa = layer_info->u.pLayerInfo->pfnNextGetDeviceProcAddr;
    layer_info->u.pLayerInfo = layer_info->u.pLayerInfo->pNext;

    next_create_device =
        (PFN_vkCreateDevice)next_gipa(VK_NULL_HANDLE, "vkCreateDevice");
    if (next_create_device == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    result = next_create_device(physicalDevice, pCreateInfo, pAllocator, pDevice);
    if (result != VK_SUCCESS) {
        return result;
    }

    result = mvgal_vk_register_device(*pDevice, next_gdpa);
    if (result != VK_SUCCESS) {
        return result;
    }

    /* Check if this device belongs to an aggregate physical device */
    mvgal_physical_device_dispatch_t *pd_dispatch = 
        mvgal_vk_find_physical_device_dispatch(physicalDevice);
    if (pd_dispatch != NULL && pd_dispatch->is_aggregate) {
        mvgal_device_dispatch_t *dev_dispatch = mvgal_vk_find_device_dispatch(*pDevice);
        if (dev_dispatch != NULL) {
            dev_dispatch->is_aggregate = true;
            mvgal_vk_layer_log("device=%p marked as aggregate", (void *)*pDevice);
        }
    }

    mvgal_vk_layer_log("created device=%p", (void *)*pDevice);
    return VK_SUCCESS;
}

VKAPI_ATTR VkResult VKAPI_CALL vkEnumeratePhysicalDevices(
    VkInstance instance,
    uint32_t *pPhysicalDeviceCount,
    VkPhysicalDevice *pPhysicalDevices
)
{
    mvgal_instance_dispatch_t *dispatch = mvgal_vk_find_instance_dispatch(instance);
    PFN_vkEnumeratePhysicalDevices next_enumerate;
    VkResult result;

    if (dispatch == NULL || dispatch->next_gipa == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    next_enumerate = (PFN_vkEnumeratePhysicalDevices)
        dispatch->next_gipa(instance, "vkEnumeratePhysicalDevices");
    if (next_enumerate == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    result = next_enumerate(instance, pPhysicalDeviceCount, pPhysicalDevices);
    if ((result == VK_SUCCESS || result == VK_INCOMPLETE) &&
        pPhysicalDeviceCount != NULL &&
        pPhysicalDevices != NULL) {
        VkResult map_result = mvgal_vk_register_physical_devices(
            instance,
            dispatch,
            *pPhysicalDeviceCount,
            pPhysicalDevices
        );
        if (map_result != VK_SUCCESS) {
            return map_result;
        }

        /* Log topology for telemetry */
        mvgal_vk_layer_log(
            "vkEnumeratePhysicalDevices instance=%p count=%u",
            (void *)instance,
            *pPhysicalDeviceCount
        );

        /* Cache properties for each newly registered device */
        for (uint32_t i = 0; i < *pPhysicalDeviceCount; ++i) {
            mvgal_physical_device_dispatch_t *pd =
                mvgal_vk_find_physical_device_dispatch(pPhysicalDevices[i]);
            if (pd != NULL && !pd->properties_cached) {
                mvgal_vk_cache_physical_device_properties(pd);
            }
        }
    }

    return result;
}

VKAPI_ATTR void VKAPI_CALL vkDestroyDevice(
    VkDevice device,
    const VkAllocationCallbacks *pAllocator
)
{
    mvgal_device_dispatch_t *dispatch = mvgal_vk_find_device_dispatch(device);

    if (dispatch != NULL) {
        mvgal_spirv_cache_destroy(&dispatch->spirv_cache);
        if (dispatch->destroy_device != NULL) {
            dispatch->destroy_device(device, pAllocator);
        }
    }

    mvgal_vk_unregister_device(device);
    mvgal_vk_layer_log("destroyed device=%p", (void *)device);
}

VKAPI_ATTR void VKAPI_CALL vkGetDeviceQueue(
    VkDevice device,
    uint32_t queueFamilyIndex,
    uint32_t queueIndex,
    VkQueue *pQueue
)
{
    mvgal_device_dispatch_t *dispatch = mvgal_vk_find_device_dispatch(device);

    if (dispatch == NULL || dispatch->get_device_queue == NULL) {
        if (pQueue != NULL) {
            *pQueue = VK_NULL_HANDLE;
        }
        return;
    }

    dispatch->get_device_queue(device, queueFamilyIndex, queueIndex, pQueue);
    if (pQueue != NULL && *pQueue != VK_NULL_HANDLE) {
        (void)mvgal_vk_register_queue(*pQueue, dispatch);
    }
}

VKAPI_ATTR void VKAPI_CALL vkGetDeviceQueue2(
    VkDevice device,
    const VkDeviceQueueInfo2 *pQueueInfo,
    VkQueue *pQueue
)
{
    mvgal_device_dispatch_t *dispatch = mvgal_vk_find_device_dispatch(device);

    if (dispatch == NULL || dispatch->get_device_queue2 == NULL) {
        if (pQueue != NULL) {
            *pQueue = VK_NULL_HANDLE;
        }
        return;
    }

    dispatch->get_device_queue2(device, pQueueInfo, pQueue);
    if (pQueue != NULL && *pQueue != VK_NULL_HANDLE) {
        (void)mvgal_vk_register_queue(*pQueue, dispatch);
    }
}

VKAPI_ATTR VkResult VKAPI_CALL vkQueueSubmit(
    VkQueue queue,
    uint32_t submitCount,
    const VkSubmitInfo *pSubmits,
    VkFence fence
)
{
    mvgal_queue_dispatch_t *queue_dispatch = mvgal_vk_find_queue_dispatch(queue);

    if (queue_dispatch == NULL ||
        queue_dispatch->device_dispatch == NULL ||
        queue_dispatch->device_dispatch->queue_submit == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    mvgal_vk_layer_note_submit("vkQueueSubmit", submitCount);

    /* AFR only when explicitly requested (never the implicit default). */
    if (queue_dispatch->device_dispatch->is_aggregate && mvgal_vk_layer_afr_enabled()) {
            static _Atomic uint32_t s_frame_count = 0;
            uint32_t frame_idx = atomic_fetch_add(&s_frame_count, 1);
            uint32_t gpu_count = mvgal_vk_layer_active_gpu_count();
            uint32_t gpu_idx = frame_idx % gpu_count;
            uint32_t device_mask = (1U << gpu_idx);

            mvgal_vk_layer_log("AFR: Dispatching submission to GPU %u (mask 0x%x)", gpu_idx, device_mask);

            /* Inject VkDeviceGroupSubmitInfo to route to specific GPU */
            VkSubmitInfo *new_submits = (VkSubmitInfo *)malloc(submitCount * sizeof(VkSubmitInfo));
            VkDeviceGroupSubmitInfo *group_infos = (VkDeviceGroupSubmitInfo *)malloc(submitCount * sizeof(VkDeviceGroupSubmitInfo));
            uint32_t **masks_pool = (uint32_t **)malloc(submitCount * sizeof(uint32_t *));

            for (uint32_t i = 0; i < submitCount; ++i) {
                new_submits[i] = pSubmits[i];
                group_infos[i].sType = VK_STRUCTURE_TYPE_DEVICE_GROUP_SUBMIT_INFO;
                group_infos[i].pNext = pSubmits[i].pNext;
                group_infos[i].waitSemaphoreCount = 0;
                group_infos[i].pWaitSemaphoreDeviceIndices = NULL;
                group_infos[i].commandBufferCount = pSubmits[i].commandBufferCount;
                
                uint32_t *masks = (uint32_t *)malloc(pSubmits[i].commandBufferCount * sizeof(uint32_t));
                for (uint32_t j = 0; j < pSubmits[i].commandBufferCount; ++j) {
                    masks[j] = device_mask;
                }
                group_infos[i].pCommandBufferDeviceMasks = masks;
                masks_pool[i] = masks;
                
                group_infos[i].signalSemaphoreCount = 0;
                group_infos[i].pSignalSemaphoreDeviceIndices = NULL;
                
                new_submits[i].pNext = &group_infos[i];
            }

            VkResult result = queue_dispatch->device_dispatch->queue_submit(
                queue,
                submitCount,
                new_submits,
                fence
            );

            /* Cleanup */
            for (uint32_t i = 0; i < submitCount; ++i) {
                free(masks_pool[i]);
            }
            free(masks_pool);
            free(group_infos);
            free(new_submits);

            return result;
    }

    return queue_dispatch->device_dispatch->queue_submit(
        queue,
        submitCount,
        pSubmits,
        fence
    );
}

#ifdef VK_VERSION_1_3
VKAPI_ATTR VkResult VKAPI_CALL vkQueueSubmit2(
    VkQueue queue,
    uint32_t submitCount,
    const VkSubmitInfo2 *pSubmits,
    VkFence fence
)
{
    mvgal_queue_dispatch_t *queue_dispatch = mvgal_vk_find_queue_dispatch(queue);

    if (queue_dispatch == NULL ||
        queue_dispatch->device_dispatch == NULL ||
        queue_dispatch->device_dispatch->queue_submit2 == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    mvgal_vk_layer_note_submit("vkQueueSubmit2", submitCount);
    return queue_dispatch->device_dispatch->queue_submit2(
        queue,
        submitCount,
        pSubmits,
        fence
    );
}
#endif /* VK_VERSION_1_3 */

#ifdef VK_VERSION_1_3
VKAPI_ATTR VkResult VKAPI_CALL vkQueueSubmit2KHR(
    VkQueue queue,
    uint32_t submitCount,
    const VkSubmitInfo2 *pSubmits,
    VkFence fence
)
{
    mvgal_queue_dispatch_t *queue_dispatch = mvgal_vk_find_queue_dispatch(queue);

    if (queue_dispatch == NULL ||
        queue_dispatch->device_dispatch == NULL ||
        queue_dispatch->device_dispatch->queue_submit2_khr == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    mvgal_vk_layer_note_submit("vkQueueSubmit2KHR", submitCount);
    return queue_dispatch->device_dispatch->queue_submit2_khr(
        queue,
        submitCount,
        pSubmits,
        fence
    );
}
#endif /* VK_VERSION_1_3 */

/* ------------------------------------------------------------------ */
/*  Shader module interception — capture SPIR-V into per-device cache */
/* ------------------------------------------------------------------ */

VKAPI_ATTR VkResult VKAPI_CALL vkCreateShaderModule(
    VkDevice device,
    const VkShaderModuleCreateInfo *pCreateInfo,
    const VkAllocationCallbacks *pAllocator,
    VkShaderModule *pShaderModule)
{
    mvgal_device_dispatch_t *dispatch = mvgal_vk_find_device_dispatch(device);
    if (dispatch == NULL || dispatch->create_shader_module == NULL)
        return VK_ERROR_INITIALIZATION_FAILED;

    VkResult result = dispatch->create_shader_module(
        device, pCreateInfo, pAllocator, pShaderModule);

    if (result == VK_SUCCESS && pCreateInfo != NULL &&
        pCreateInfo->pCode != NULL && pShaderModule != NULL) {
        mvgal_shader_capture(&dispatch->spirv_cache,
                             pCreateInfo->pCode,
                             pCreateInfo->codeSize / sizeof(uint32_t),
                             *pShaderModule);

        /* Phase 3: attempt vendor-specific SPIR-V → GLSL route
         * The decompiled source is not yet consumed upstream, but calling
         * shader_route here validates the entire opt+decompile pipeline at
         * runtime and exercises the code path for test coverage. */
        uint32_t *vendor_code = NULL;
        size_t    vendor_size = 0;
        if (mvgal_shader_route(MVGAL_VENDOR_AMD,
                               pCreateInfo->pCode,
                               pCreateInfo->codeSize / sizeof(uint32_t),
                               &vendor_code,
                               &vendor_size) == 0) {
            mvgal_vk_layer_log(
                "vkCreateShaderModule device=%p module=%p "
                "vendor_code=%p vendor_size=%zu",
                (void *)device, (void *)*pShaderModule,
                (void *)vendor_code, vendor_size);
            free(vendor_code);
        }

        mvgal_vk_layer_log(
            "vkCreateShaderModule device=%p module=%p code_size=%zu",
            (void *)device, (void *)*pShaderModule,
            (size_t)pCreateInfo->codeSize);
    }
    return result;
}

VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL vk_layerGetPhysicalDeviceProcAddr(
    VkInstance instance,
    const char *pName
)
{
    mvgal_instance_dispatch_t *dispatch = mvgal_vk_find_instance_dispatch(instance);

    (void)pName;
    if (dispatch == NULL || dispatch->next_gpipa == NULL) {
        return NULL;
    }

    return dispatch->next_gpipa(instance, pName);
}

VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL vkGetDeviceProcAddr(
    VkDevice device,
    const char *pName
)
{
    mvgal_device_dispatch_t *dispatch;

    if (pName == NULL) {
        return NULL;
    }

    if (strcmp(pName, "vkGetDeviceProcAddr") == 0) {
        return (PFN_vkVoidFunction)vkGetDeviceProcAddr;
    }
    if (strcmp(pName, "vkDestroyDevice") == 0) {
        return (PFN_vkVoidFunction)vkDestroyDevice;
    }
    if (strcmp(pName, "vkGetDeviceQueue") == 0) {
        return (PFN_vkVoidFunction)vkGetDeviceQueue;
    }
    if (strcmp(pName, "vkGetDeviceQueue2") == 0) {
        return (PFN_vkVoidFunction)vkGetDeviceQueue2;
    }
    if (strcmp(pName, "vkQueueSubmit") == 0) {
        return (PFN_vkVoidFunction)vkQueueSubmit;
    }
    if (strcmp(pName, "vkCreateShaderModule") == 0) {
        return (PFN_vkVoidFunction)vkCreateShaderModule;
    }
#ifdef VK_VERSION_1_3
    if (strcmp(pName, "vkQueueSubmit2") == 0) {
        return (PFN_vkVoidFunction)vkQueueSubmit2;
    }
    if (strcmp(pName, "vkQueueSubmit2KHR") == 0) {
        return (PFN_vkVoidFunction)vkQueueSubmit2KHR;
    }
#endif

    dispatch = mvgal_vk_find_device_dispatch(device);
    if (dispatch == NULL || dispatch->next_gdpa == NULL) {
        return NULL;
    }

    return dispatch->next_gdpa(device, pName);
}

VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL vkGetInstanceProcAddr(
    VkInstance instance,
    const char *pName
)
{
    mvgal_instance_dispatch_t *dispatch;

    if (pName == NULL) {
        return NULL;
    }

    if (strcmp(pName, "vkGetInstanceProcAddr") == 0) {
        return (PFN_vkVoidFunction)vkGetInstanceProcAddr;
    }
    if (strcmp(pName, "vkGetDeviceProcAddr") == 0) {
        return (PFN_vkVoidFunction)vkGetDeviceProcAddr;
    }
    if (strcmp(pName, "vkNegotiateLoaderLayerInterfaceVersion") == 0) {
        return (PFN_vkVoidFunction)vkNegotiateLoaderLayerInterfaceVersion;
    }
    if (strcmp(pName, "vk_layerGetPhysicalDeviceProcAddr") == 0) {
        return (PFN_vkVoidFunction)vk_layerGetPhysicalDeviceProcAddr;
    }
    if (strcmp(pName, "vkEnumerateInstanceLayerProperties") == 0) {
        return (PFN_vkVoidFunction)vkEnumerateInstanceLayerProperties;
    }
    if (strcmp(pName, "vkEnumerateDeviceLayerProperties") == 0) {
        return (PFN_vkVoidFunction)vkEnumerateDeviceLayerProperties;
    }
    if (strcmp(pName, "vkEnumerateInstanceExtensionProperties") == 0) {
        return (PFN_vkVoidFunction)vkEnumerateInstanceExtensionProperties;
    }
    if (strcmp(pName, "vkEnumerateDeviceExtensionProperties") == 0) {
        return (PFN_vkVoidFunction)vkEnumerateDeviceExtensionProperties;
    }
    if (strcmp(pName, "vkEnumerateInstanceVersion") == 0) {
        return (PFN_vkVoidFunction)vkEnumerateInstanceVersion;
    }
    if (strcmp(pName, "vkCreateInstance") == 0) {
        return (PFN_vkVoidFunction)vkCreateInstance;
    }
    if (strcmp(pName, "vkDestroyInstance") == 0) {
        return (PFN_vkVoidFunction)vkDestroyInstance;
    }
    if (strcmp(pName, "vkCreateDevice") == 0) {
        return (PFN_vkVoidFunction)vkCreateDevice;
    }
    if (strcmp(pName, "vkEnumeratePhysicalDevices") == 0) {
        return (PFN_vkVoidFunction)vkEnumeratePhysicalDevices;
    }
    if (strcmp(pName, "vkDestroyDevice") == 0) {
        return (PFN_vkVoidFunction)vkDestroyDevice;
    }
    if (strcmp(pName, "vkGetDeviceQueue") == 0) {
        return (PFN_vkVoidFunction)vkGetDeviceQueue;
    }
    if (strcmp(pName, "vkGetDeviceQueue2") == 0) {
        return (PFN_vkVoidFunction)vkGetDeviceQueue2;
    }
    if (strcmp(pName, "vkQueueSubmit") == 0) {
        return (PFN_vkVoidFunction)vkQueueSubmit;
    }
#ifdef VK_VERSION_1_3
    if (strcmp(pName, "vkQueueSubmit2") == 0) {
        return (PFN_vkVoidFunction)vkQueueSubmit2;
    }
    if (strcmp(pName, "vkQueueSubmit2KHR") == 0) {
        return (PFN_vkVoidFunction)vkQueueSubmit2KHR;
    }
#endif
    if (strcmp(pName, "vkEnumeratePhysicalDeviceGroups") == 0) {
        return (PFN_vkVoidFunction)vkEnumeratePhysicalDeviceGroups;
    }
    if (strcmp(pName, "vkEnumeratePhysicalDeviceGroupsKHX") == 0) {
        return (PFN_vkVoidFunction)vkEnumeratePhysicalDeviceGroupsKHX;
    }

    if (strcmp(pName, "vkGetPhysicalDeviceProperties") == 0) {
        return (PFN_vkVoidFunction)vkGetPhysicalDeviceProperties;
    }
    if (strcmp(pName, "vkGetPhysicalDeviceFeatures") == 0) {
        return (PFN_vkVoidFunction)vkGetPhysicalDeviceFeatures;
    }
    if (strcmp(pName, "vkGetPhysicalDeviceMemoryProperties") == 0) {
        return (PFN_vkVoidFunction)vkGetPhysicalDeviceMemoryProperties;
    }
    if (strcmp(pName, "vkGetPhysicalDeviceQueueFamilyProperties") == 0) {
        return (PFN_vkVoidFunction)vkGetPhysicalDeviceQueueFamilyProperties;
    }
    if (strcmp(pName, "vkGetPhysicalDeviceProperties2") == 0) {
        return (PFN_vkVoidFunction)vkGetPhysicalDeviceProperties2;
    }
    if (strcmp(pName, "vkGetPhysicalDeviceFeatures2") == 0) {
        return (PFN_vkVoidFunction)vkGetPhysicalDeviceFeatures2;
    }
    if (strcmp(pName, "vkGetPhysicalDeviceMemoryProperties2") == 0) {
        return (PFN_vkVoidFunction)vkGetPhysicalDeviceMemoryProperties2;
    }

    dispatch = mvgal_vk_find_instance_dispatch(instance);
    if (dispatch == NULL || dispatch->next_gipa == NULL) {
        return NULL;
    }

    return dispatch->next_gipa(instance, pName);
}

/* ---------------------------------------------------------------------------
 * Physical device property caching
 * ---------------------------------------------------------------------------
 * Cache VkPhysicalDeviceProperties for each tracked physical device so the
 * layer can log device names and vendor IDs in telemetry without calling back
 * into the driver on every submit.
 */
void mvgal_vk_cache_physical_device_properties(
    mvgal_physical_device_dispatch_t *dispatch
)
{
    if (dispatch == NULL || dispatch->properties_cached) {
        return;
    }

    if (dispatch->instance_dispatch == NULL ||
        dispatch->instance_dispatch->get_physical_device_properties == NULL) {
        return;
    }

    dispatch->instance_dispatch->get_physical_device_properties(
        dispatch->physical_device,
        &dispatch->properties
    );
    dispatch->properties_cached = true;

    mvgal_vk_layer_log(
        "physical_device=%p vendor=0x%04x device=0x%04x name=\"%s\"",
        (void *)dispatch->physical_device,
        (unsigned)dispatch->properties.vendorID,
        (unsigned)dispatch->properties.deviceID,
        dispatch->properties.deviceName
    );

    atomic_fetch_add_explicit(
        &g_mvgal_layer_state.physical_device_count,
        1U,
        memory_order_relaxed
    );
}

/* ---------------------------------------------------------------------------
 * Physical device forwarding interceptors
 *
 * These pass through to the next layer/driver while giving the MVGAL layer
 * an observation point for future aggregation logic.
 * ---------------------------------------------------------------------------
 */

VKAPI_ATTR void VKAPI_CALL vkGetPhysicalDeviceProperties(
    VkPhysicalDevice physicalDevice,
    VkPhysicalDeviceProperties *pProperties
)
{
    mvgal_physical_device_dispatch_t *dispatch =
        mvgal_vk_find_physical_device_dispatch(physicalDevice);

    if (dispatch == NULL || dispatch->instance_dispatch == NULL ||
        dispatch->instance_dispatch->get_physical_device_properties == NULL) {
        return;
    }

    dispatch->instance_dispatch->get_physical_device_properties(
        physicalDevice, pProperties
    );

    /* Opportunistically cache if not yet done */
    if (!dispatch->properties_cached && pProperties != NULL) {
        dispatch->properties = *pProperties;
        dispatch->properties_cached = true;
        mvgal_vk_layer_log(
            "vkGetPhysicalDeviceProperties physical_device=%p vendor=0x%04x name=\"%s\"",
            (void *)physicalDevice,
            (unsigned)pProperties->vendorID,
            pProperties->deviceName
        );
    }
}

VKAPI_ATTR void VKAPI_CALL vkGetPhysicalDeviceFeatures(
    VkPhysicalDevice physicalDevice,
    VkPhysicalDeviceFeatures *pFeatures
)
{
    mvgal_physical_device_dispatch_t *dispatch =
        mvgal_vk_find_physical_device_dispatch(physicalDevice);

    if (dispatch == NULL || dispatch->instance_dispatch == NULL ||
        dispatch->instance_dispatch->get_physical_device_features == NULL) {
        return;
    }

    dispatch->instance_dispatch->get_physical_device_features(
        physicalDevice, pFeatures
    );
}

VKAPI_ATTR void VKAPI_CALL vkGetPhysicalDeviceMemoryProperties(
    VkPhysicalDevice physicalDevice,
    VkPhysicalDeviceMemoryProperties *pMemoryProperties
)
{
    mvgal_physical_device_dispatch_t *dispatch =
        mvgal_vk_find_physical_device_dispatch(physicalDevice);

    if (dispatch == NULL || dispatch->instance_dispatch == NULL ||
        dispatch->instance_dispatch->get_physical_device_memory_properties == NULL) {
        return;
    }

    dispatch->instance_dispatch->get_physical_device_memory_properties(
        physicalDevice, pMemoryProperties
    );
}

VKAPI_ATTR void VKAPI_CALL vkGetPhysicalDeviceQueueFamilyProperties(
    VkPhysicalDevice physicalDevice,
    uint32_t *pQueueFamilyPropertyCount,
    VkQueueFamilyProperties *pQueueFamilyProperties
)
{
    mvgal_physical_device_dispatch_t *dispatch =
        mvgal_vk_find_physical_device_dispatch(physicalDevice);

    if (dispatch == NULL || dispatch->instance_dispatch == NULL ||
        dispatch->instance_dispatch->get_physical_device_queue_family_properties == NULL) {
        return;
    }

    dispatch->instance_dispatch->get_physical_device_queue_family_properties(
        physicalDevice, pQueueFamilyPropertyCount, pQueueFamilyProperties
    );
}

VKAPI_ATTR void VKAPI_CALL vkGetPhysicalDeviceProperties2(
    VkPhysicalDevice physicalDevice,
    VkPhysicalDeviceProperties2 *pProperties
)
{
    mvgal_physical_device_dispatch_t *dispatch =
        mvgal_vk_find_physical_device_dispatch(physicalDevice);

    if (dispatch == NULL || dispatch->instance_dispatch == NULL ||
        dispatch->instance_dispatch->get_physical_device_properties2 == NULL) {
        return;
    }

    dispatch->instance_dispatch->get_physical_device_properties2(
        physicalDevice, pProperties
    );

    /* Opportunistically cache core properties */
    if (!dispatch->properties_cached && pProperties != NULL) {
        dispatch->properties = pProperties->properties;
        dispatch->properties_cached = true;
        mvgal_vk_layer_log(
            "vkGetPhysicalDeviceProperties2 physical_device=%p vendor=0x%04x name=\"%s\"",
            (void *)physicalDevice,
            (unsigned)pProperties->properties.vendorID,
            pProperties->properties.deviceName
        );
    }
}

VKAPI_ATTR void VKAPI_CALL vkGetPhysicalDeviceFeatures2(
    VkPhysicalDevice physicalDevice,
    VkPhysicalDeviceFeatures2 *pFeatures
)
{
    mvgal_physical_device_dispatch_t *dispatch =
        mvgal_vk_find_physical_device_dispatch(physicalDevice);

    if (dispatch == NULL || dispatch->instance_dispatch == NULL ||
        dispatch->instance_dispatch->get_physical_device_features2 == NULL) {
        return;
    }

    dispatch->instance_dispatch->get_physical_device_features2(
        physicalDevice, pFeatures
    );
}

VKAPI_ATTR void VKAPI_CALL vkGetPhysicalDeviceMemoryProperties2(
    VkPhysicalDevice physicalDevice,
    VkPhysicalDeviceMemoryProperties2 *pMemoryProperties
)
{
    mvgal_physical_device_dispatch_t *dispatch =
        mvgal_vk_find_physical_device_dispatch(physicalDevice);

    if (dispatch == NULL || dispatch->instance_dispatch == NULL ||
        dispatch->instance_dispatch->get_physical_device_memory_properties2 == NULL) {
        return;
    }

    dispatch->instance_dispatch->get_physical_device_memory_properties2(
        physicalDevice, pMemoryProperties
    );
}

VKAPI_ATTR VkResult VKAPI_CALL vkEnumeratePhysicalDeviceGroups(
    VkInstance instance,
    uint32_t *pPhysicalDeviceGroupCount,
    VkPhysicalDeviceGroupProperties *pPhysicalDeviceGroupProperties
)
{
    mvgal_instance_dispatch_t *dispatch = mvgal_vk_find_instance_dispatch(instance);
    uint32_t total_phys_devs = 0;
    mvgal_physical_device_dispatch_t *pd_iter;

    if (dispatch == NULL) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    pd_iter = g_mvgal_layer_state.physical_devices;
    while (pd_iter != NULL) {
        if (pd_iter->instance == instance) {
            total_phys_devs++;
        }
        pd_iter = pd_iter->next;
    }
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);

    if (total_phys_devs == 0) {
        /* Fallback to next layer if we haven't tracked anything yet */
        if (dispatch->enumerate_physical_device_groups) {
            return dispatch->enumerate_physical_device_groups(
                instance, pPhysicalDeviceGroupCount, pPhysicalDeviceGroupProperties);
        }
        if (pPhysicalDeviceGroupCount) *pPhysicalDeviceGroupCount = 0;
        return VK_SUCCESS;
    }

    if (!mvgal_vk_layer_multi_gpu_enabled()) {
        if (dispatch->enumerate_physical_device_groups) {
            return dispatch->enumerate_physical_device_groups(
                instance, pPhysicalDeviceGroupCount, pPhysicalDeviceGroupProperties);
        }
        if (pPhysicalDeviceGroupCount) {
            *pPhysicalDeviceGroupCount = total_phys_devs;
        }
        return VK_SUCCESS;
    }

    if (pPhysicalDeviceGroupProperties == NULL) {
        *pPhysicalDeviceGroupCount = 1;
        return VK_SUCCESS;
    }

    if (*pPhysicalDeviceGroupCount == 0) {
        return VK_INCOMPLETE;
    }

    /* Virtualize a single group containing all GPUs */
    pPhysicalDeviceGroupProperties[0].sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_GROUP_PROPERTIES;
    pPhysicalDeviceGroupProperties[0].pNext = NULL;
    pPhysicalDeviceGroupProperties[0].physicalDeviceCount = 
        (total_phys_devs > VK_MAX_DEVICE_GROUP_SIZE) ? VK_MAX_DEVICE_GROUP_SIZE : total_phys_devs;
    
    pthread_mutex_lock(&g_mvgal_layer_state.lock);
    pd_iter = g_mvgal_layer_state.physical_devices;
    uint32_t i = 0;
    while (pd_iter != NULL && i < pPhysicalDeviceGroupProperties[0].physicalDeviceCount) {
        if (pd_iter->instance == instance) {
            pPhysicalDeviceGroupProperties[0].physicalDevices[i] = pd_iter->physical_device;
            pd_iter->is_aggregate = true;
            i++;
        }
        pd_iter = pd_iter->next;
    }
    pthread_mutex_unlock(&g_mvgal_layer_state.lock);
    
    pPhysicalDeviceGroupProperties[0].subsetAllocation = VK_FALSE;
    *pPhysicalDeviceGroupCount = 1;

    mvgal_vk_layer_log("Virtualizing device group with %u GPUs", pPhysicalDeviceGroupProperties[0].physicalDeviceCount);

    return VK_SUCCESS;
}

VKAPI_ATTR VkResult VKAPI_CALL vkEnumeratePhysicalDeviceGroupsKHX(
    VkInstance instance,
    uint32_t *pPhysicalDeviceGroupCount,
    VkPhysicalDeviceGroupProperties *pPhysicalDeviceGroupProperties
)
{
    return vkEnumeratePhysicalDeviceGroups(instance, pPhysicalDeviceGroupCount, pPhysicalDeviceGroupProperties);
}

VKAPI_ATTR VkResult VKAPI_CALL vkNegotiateLoaderLayerInterfaceVersion(
    VkNegotiateLayerInterface *pVersionStruct
)
{
    if (pVersionStruct == NULL ||
        pVersionStruct->sType != LAYER_NEGOTIATE_INTERFACE_STRUCT) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    if (pVersionStruct->loaderLayerInterfaceVersion >
        CURRENT_LOADER_LAYER_INTERFACE_VERSION) {
        pVersionStruct->loaderLayerInterfaceVersion =
            CURRENT_LOADER_LAYER_INTERFACE_VERSION;
    }

    if (pVersionStruct->loaderLayerInterfaceVersion <
        MIN_SUPPORTED_LOADER_LAYER_INTERFACE_VERSION) {
        return VK_ERROR_INITIALIZATION_FAILED;
    }

    pVersionStruct->pfnGetInstanceProcAddr = vkGetInstanceProcAddr;
    pVersionStruct->pfnGetDeviceProcAddr = vkGetDeviceProcAddr;
    pVersionStruct->pfnGetPhysicalDeviceProcAddr =
        vk_layerGetPhysicalDeviceProcAddr;

    return VK_SUCCESS;
}
