loader: EnumPhysDev fixes

Found a few issues, and I had some concerns about the physical
device values enduring over multiple queries.

Change-Id: Ifaa94a4ecf9edfc79bdd3b3d473db068952e3264

2 files changed, 208 insertions, 148 deletions

diff --git a/loader/loader.c b/loader/loader.c
index 80c1ce36..720e2762 100644
--- a/loader/loader.c
+++ b/loader/loader.c
@@ -4435,93 +4435,112 @@ VKAPI_ATTR VkResult VKAPI_CALL terminator_EnumeratePhysicalDevices(
     VkPhysicalDevice *pPhysicalDevices) {
     struct loader_instance *inst = (struct loader_instance *)instance;
     VkResult res = VK_SUCCESS;
-
-    struct loader_icd_term *icd_term;
-    struct loader_phys_dev_per_icd *phys_devs;
+    struct loader_icd_term *icd_term = NULL;
+    struct loader_phys_dev_per_icd *icd_phys_devs = NULL;
+    uint32_t copy_count = 0;
+    uint32_t new_phys_dev_count = 0;
+    uint32_t i = 0;
+    struct loader_physical_device_term **new_phys_devs = NULL;
 
     inst->total_gpu_count = 0;
-    phys_devs = (struct loader_phys_dev_per_icd *)loader_stack_alloc(
+    icd_phys_devs = (struct loader_phys_dev_per_icd *)loader_stack_alloc(
         sizeof(struct loader_phys_dev_per_icd) * inst->total_icd_count);
-    if (!phys_devs)
-        return VK_ERROR_OUT_OF_HOST_MEMORY;
-
-    icd_term = inst->icd_terms;
-    for (uint32_t i = 0; i < inst->total_icd_count; i++) {
-        assert(icd_term);
-        res = icd_term->EnumeratePhysicalDevices(icd_term->instance,
-                                                 &phys_devs[i].count, NULL);
-        if (res != VK_SUCCESS)
-            return res;
-        icd_term = icd_term->next;
+    if (NULL == icd_phys_devs) {
+        res = VK_ERROR_OUT_OF_HOST_MEMORY;
+        goto out;
     }
 
     icd_term = inst->icd_terms;
-    for (uint32_t i = 0; i < inst->total_icd_count; i++) {
-        assert(icd_term);
-        phys_devs[i].phys_devs = (VkPhysicalDevice *)loader_stack_alloc(
-            phys_devs[i].count * sizeof(VkPhysicalDevice));
-        if (!phys_devs[i].phys_devs) {
-            return VK_ERROR_OUT_OF_HOST_MEMORY;
+    for (i = 0; i < inst->total_icd_count; i++) {
+        if (NULL == icd_term) {
+            loader_log(inst, VK_DEBUG_REPORT_ERROR_BIT_EXT, 0,
+                       "Invalid ICD encountered during"
+                       "vkEnumeratePhysicalDevices");
+            assert(false);
         }
-        res = icd_term->EnumeratePhysicalDevices(
-            icd_term->instance, &(phys_devs[i].count), phys_devs[i].phys_devs);
-        if (res == VK_SUCCESS) {
-            inst->total_gpu_count += phys_devs[i].count;
-        } else {
-            return res;
+
+        // Determine how many physical devices are associated with this ICD.
+        res = icd_term->EnumeratePhysicalDevices(icd_term->instance,
+                                                 &icd_phys_devs[i].count, NULL);
+        if (res != VK_SUCCESS) {
+            goto out;
         }
-        phys_devs[i].this_icd_term = icd_term;
+
+        if (NULL != pPhysicalDevices) {
+            // Create an array to store each physical device for this ICD.
+            icd_phys_devs[i].phys_devs = (VkPhysicalDevice *)loader_stack_alloc(
+                icd_phys_devs[i].count * sizeof(VkPhysicalDevice));
+            if (NULL == icd_phys_devs[i].phys_devs) {
+                res = VK_ERROR_OUT_OF_HOST_MEMORY;
+                goto out;
+            }
+
+            // Query the VkPhysicalDevice values for each of the physical devices
+            // associated with this ICD.
+            res = icd_term->EnumeratePhysicalDevices(
+                icd_term->instance, &(icd_phys_devs[i].count),
+                icd_phys_devs[i].phys_devs);
+            if (res != VK_SUCCESS) {
+                goto out;
+            }
+
+            icd_phys_devs[i].this_icd_term = icd_term;
+        }
+
+        inst->total_gpu_count += icd_phys_devs[i].count;
+
+        // Go to the next ICD
         icd_term = icd_term->next;
     }
+
     if (inst->total_gpu_count == 0) {
-        return VK_ERROR_INITIALIZATION_FAILED;
+        res = VK_ERROR_INITIALIZATION_FAILED;
+        goto out;
     }
 
-    uint32_t copy_count = inst->total_gpu_count;
-    uint32_t new_phys_dev_count = inst->total_gpu_count;
-    struct loader_physical_device_term **new_phys_devs = NULL;
+    copy_count = inst->total_gpu_count;
 
     if (NULL != pPhysicalDevices) {
+        new_phys_dev_count = inst->total_gpu_count;
 
-        // cap the number of devices at pPhysicalDeviceCount
+        // Cap the number of devices at pPhysicalDeviceCount
         if (copy_count > *pPhysicalDeviceCount) {
             copy_count = *pPhysicalDeviceCount;
         }
 
-        // allocate the new devices list
+        // Allocate the new devices list
         new_phys_devs = loader_instance_heap_alloc(
             inst,
             sizeof(struct loader_physical_device_term *) * new_phys_dev_count,
             VK_SYSTEM_ALLOCATION_SCOPE_INSTANCE);
         if (NULL == new_phys_devs) {
-            return VK_ERROR_OUT_OF_HOST_MEMORY;
+            res = VK_ERROR_OUT_OF_HOST_MEMORY;
+            goto out;
         }
         memset(new_phys_devs, 0, sizeof(struct loader_physical_device_term *) *
                                      new_phys_dev_count);
 
-        // copy or create everything to fill the new array of physical devices
+        // Copy or create everything to fill the new array of physical devices
         uint32_t idx = 0;
-        for (uint32_t i = 0;
-             idx < new_phys_dev_count && i < inst->total_icd_count;
-             i++) {
-            for (uint32_t j = 0;
-                 j < phys_devs[i].count && idx < new_phys_dev_count;
-                 j++) {
-
-                // check if this physical device is already in the old buffer
-                new_phys_devs[i] = NULL;
-                for (uint32_t k = 0; k < inst->phys_dev_count_term &&
-                                     NULL != inst->phys_devs_term;
-                     k++) {
-                    if (phys_devs[i].phys_devs[k] ==
-                        inst->phys_devs_term[k]->phys_dev) {
-                        new_phys_devs[i] = inst->phys_devs_term[k];
-                        break;
+        for (uint32_t icd_idx = 0; icd_idx < inst->total_icd_count; icd_idx++) {
+            for (uint32_t pd_idx = 0; pd_idx < icd_phys_devs[icd_idx].count;
+                 pd_idx++) {
+
+                // Check if this physical device is already in the old buffer
+                if (NULL != inst->phys_devs_term) {
+                    for (uint32_t old_idx = 0;
+                         old_idx < inst->phys_dev_count_term;
+                         old_idx++) {
+                        if (icd_phys_devs[icd_idx].phys_devs[pd_idx] ==
+                            inst->phys_devs_term[old_idx]->phys_dev) {
+                            new_phys_devs[idx] = inst->phys_devs_term[old_idx];
+                            break;
+                        }
                     }
                 }
-
-                // if this physical  device isn't in the old buffer, create it
-                if (NULL == new_phys_devs[i]) {
+                // If this physical device isn't in the old buffer, then we
+                // need to create it.
+                if (NULL == new_phys_devs[idx]) {
                     new_phys_devs[idx] = loader_instance_heap_alloc(
                         inst, sizeof(struct loader_physical_device_term),
                         VK_SYSTEM_ALLOCATION_SCOPE_INSTANCE);
@@ -4531,64 +4550,79 @@ VKAPI_ATTR VkResult VKAPI_CALL terminator_EnumeratePhysicalDevices(
                         goto out;
                     }
 
-                    loader_set_dispatch((void *)new_phys_devs[idx],
-                                        inst->disp);
+                    loader_set_dispatch((void *)new_phys_devs[idx], inst->disp);
                     new_phys_devs[idx]->this_icd_term =
-                        phys_devs[i].this_icd_term;
-                    new_phys_devs[idx]->icd_index = (uint8_t)(i);
+                        icd_phys_devs[icd_idx].this_icd_term;
+                    new_phys_devs[idx]->icd_index = (uint8_t)(icd_idx);
                     new_phys_devs[idx]->phys_dev =
-                        phys_devs[i].phys_devs[j];
+                        icd_phys_devs[icd_idx].phys_devs[pd_idx];
                 }
 
-                // copy wrapped object into application provided array
+                // Copy wrapped object into application provided array
                 if (idx < copy_count) {
                     pPhysicalDevices[idx] =
                         (VkPhysicalDevice)new_phys_devs[idx];
                 }
                 idx++;
+                if (idx >= new_phys_dev_count) {
+                    break;
+                }
+            }
+            if (idx >= new_phys_dev_count) {
+                break;
             }
         }
     }
 
 out:
+
     if (NULL != pPhysicalDevices) {
-        // if there was no error, free the old buffer and assign the new one
+        // If there was no error, we still need to free the old buffer and
+        // assign the new one
         if (res == VK_SUCCESS || res == VK_INCOMPLETE) {
-            // free everything that didn't carry over to the new array of
-            // physical
-            // devices
+            // Free everything that didn't carry over to the new array of
+            // physical devices.  Everything else will have been copied over
+            // to the new array.
             if (NULL != inst->phys_devs_term) {
-                for (uint32_t i = 0; i < inst->phys_dev_count_term; i++) {
+                for (uint32_t cur_pd = 0; cur_pd < inst->phys_dev_count_term;
+                     cur_pd++) {
                     bool found = false;
-                    for (uint32_t j = 0; j < new_phys_dev_count; j++) {
-                        if (inst->phys_devs_term[i] == new_phys_devs[j]) {
+                    for (uint32_t new_pd_idx = 0;
+                         new_pd_idx < new_phys_dev_count;
+                         new_pd_idx++) {
+                        if (inst->phys_devs_term[cur_pd] ==
+                                new_phys_devs[new_pd_idx]) {
                             found = true;
                             break;
                         }
                     }
                     if (!found) {
                         loader_instance_heap_free(inst,
-                                                  inst->phys_devs_term[i]);
+                                                  inst->phys_devs_term[cur_pd]);
                     }
                 }
                 loader_instance_heap_free(inst, inst->phys_devs_term);
             }
 
-            // if we didn't load every device, the result is inclomplete
+            // If we didn't load every device, the result is incomplete
             if (copy_count < new_phys_dev_count) {
                 res = VK_INCOMPLETE;
             }
 
-            // swap out old and new devices list
+            // Swap out old and new devices list
             inst->phys_dev_count_term = new_phys_dev_count;
             inst->phys_devs_term = new_phys_devs;
 
-            // if there was an error, free the new buffer
         } else {
+            // Otherwise, we've encountered an error, so we should free the
+            // new buffers.
             for (uint32_t i = 0; i < copy_count; i++) {
                 loader_instance_heap_free(inst, new_phys_devs[i]);
             }
             loader_instance_heap_free(inst, new_phys_devs);
+
+            // Set the copy count to 0 since something bad happened.
+            copy_count = 0;
         }
     }
 
diff --git a/loader/trampoline.c b/loader/trampoline.c
index 20b133ea..e56cbbe3 100644
--- a/loader/trampoline.c
+++ b/loader/trampoline.c
@@ -528,116 +528,142 @@ LOADER_EXPORT VKAPI_ATTR VkResult VKAPI_CALL
 vkEnumeratePhysicalDevices(VkInstance instance, uint32_t *pPhysicalDeviceCount,
                            VkPhysicalDevice *pPhysicalDevices) {
     const VkLayerInstanceDispatchTable *disp;
-    VkResult res;
-    struct loader_instance *inst;
-    disp = loader_get_instance_dispatch(instance);
-
+    VkResult res = VK_SUCCESS;
     loader_platform_thread_lock_mutex(&loader_lock);
-    res = disp->EnumeratePhysicalDevices(instance, pPhysicalDeviceCount,
-                                         pPhysicalDevices);
-
-    if (res != VK_SUCCESS && res != VK_INCOMPLETE) {
-        loader_platform_thread_unlock_mutex(&loader_lock);
-        return res;
+    uint32_t copy_count = 0;
+    struct loader_physical_device_tramp **new_phys_devs = NULL;
+    struct loader_instance *inst = loader_get_instance(instance);
+    VkPhysicalDevice *temp_pd_array = NULL;
+    uint32_t new_pd_count = 0;
+
+    if (NULL == inst) {
+        res = VK_ERROR_INITIALIZATION_FAILED;
+        goto out;
     }
 
-    if (!pPhysicalDevices) {
-        loader_platform_thread_unlock_mutex(&loader_lock);
-        return res;
+    disp = loader_get_instance_dispatch(instance);
+
+    // If they only want the count, just pass it in and return
+    res = disp->EnumeratePhysicalDevices(instance, &copy_count,
+                                            NULL);
+    if (NULL == pPhysicalDevices) {
+        goto out;
     }
 
-    // wrap the PhysDev object for loader usage, return wrapped objects
-    inst = loader_get_instance(instance);
-    if (!inst) {
-        loader_platform_thread_unlock_mutex(&loader_lock);
-        return VK_ERROR_INITIALIZATION_FAILED;
+    // Reset the copy count until we complete.
+    copy_count = 0;
+
+    // If we're querying actual information, then we want to internally
+    // keep track of all GPUs.  So, query them all, and we can trim the
+    // list down later.
+    new_pd_count = inst->total_gpu_count;
+    temp_pd_array = (VkPhysicalDevice *)loader_stack_alloc(
+        sizeof(VkPhysicalDevice) * new_pd_count);
+    if (NULL == temp_pd_array) {
+        res = VK_ERROR_OUT_OF_HOST_MEMORY;
+        goto out;
     }
+    res = disp->EnumeratePhysicalDevices(instance, &new_pd_count,
+                                         temp_pd_array);
+    if (res != VK_SUCCESS && res != VK_INCOMPLETE) {
+        goto out;
+    }
+
+    // Determine how many items we need to create and return (via copying into
+    // the provided array).
+    copy_count = (new_pd_count < *pPhysicalDeviceCount) ? new_pd_count :
+        *pPhysicalDeviceCount;
 
-    // create a new array for the physical devices
-    uint32_t new_phys_dev_count = (inst->total_gpu_count < *pPhysicalDeviceCount)
-                                     ? inst->total_gpu_count
-                                     : *pPhysicalDeviceCount;
-    *pPhysicalDeviceCount = new_phys_dev_count;
-    struct loader_physical_device_tramp **new_phys_devs =
-        (struct loader_physical_device_tramp **)loader_instance_heap_alloc(
-            inst, new_phys_dev_count *
-                      sizeof(struct loader_physical_device_tramp *),
+    // Create a new array for the physical devices
+    new_phys_devs = (struct loader_physical_device_tramp **)
+        loader_instance_heap_alloc(
+            inst, new_pd_count * sizeof(struct loader_physical_device_tramp *),
             VK_SYSTEM_ALLOCATION_SCOPE_INSTANCE);
     if (NULL == new_phys_devs) {
-        loader_platform_thread_unlock_mutex(&loader_lock);
-        return VK_ERROR_OUT_OF_HOST_MEMORY;
+        res = VK_ERROR_OUT_OF_HOST_MEMORY;
+        goto out;
     }
     memset(new_phys_devs, 0,
-           new_phys_dev_count * sizeof(struct loader_physical_device_tramp *));
-
-    // copy or create everything to fill the new array of physical devices
-    for (uint32_t i = 0; i < new_phys_dev_count; i++) {
-
-        // check if this physical device is already in the old buffer
-        new_phys_devs[i] = NULL;
-        for (uint32_t j = 0; j < inst->phys_dev_count_tramp; j++) {
-            if (pPhysicalDevices[i] == inst->phys_devs_tramp[j]->phys_dev) {
-                new_phys_devs[i] = inst->phys_devs_tramp[j];
+           new_pd_count * sizeof(struct loader_physical_device_tramp *));
+
+    // Copy or create everything to fill the new array of physical devices
+    for (uint32_t new_idx = 0; new_idx < new_pd_count; new_idx++) {
+
+        // Check if this physical device is already in the old buffer
+        for (uint32_t old_idx = 0;
+             old_idx < inst->phys_dev_count_tramp;
+             old_idx++) {
+            if (temp_pd_array[new_idx] ==
+                inst->phys_devs_tramp[old_idx]->phys_dev) {
+                new_phys_devs[new_idx] = inst->phys_devs_tramp[old_idx];
                 break;
             }
         }
-
-        // if this physical  device isn't in the old buffer, create it
-        if (NULL == new_phys_devs[i]) {
-            new_phys_devs[i] = (struct loader_physical_device_tramp *)
+        // If this physical device isn't in the old buffer, create it
+        if (NULL == new_phys_devs[new_idx]) {
+            new_phys_devs[new_idx] = (struct loader_physical_device_tramp *)
                 loader_instance_heap_alloc(
                     inst, sizeof(struct loader_physical_device_tramp),
                     VK_SYSTEM_ALLOCATION_SCOPE_INSTANCE);
-            if (NULL == new_phys_devs[i]) {
-                new_phys_dev_count = i;
+            if (NULL == new_phys_devs[new_idx]) {
+                new_pd_count = new_idx;
                 res = VK_ERROR_OUT_OF_HOST_MEMORY;
                 goto out;
             }
 
-            // initialize the loader's physicalDevice object
-            loader_set_dispatch((void *)new_phys_devs[i], inst->disp);
-            new_phys_devs[i]->this_instance = inst;
-            new_phys_devs[i]->phys_dev = pPhysicalDevices[i];
+            // Initialize the new physicalDevice object
+            loader_set_dispatch((void *)new_phys_devs[new_idx], inst->disp);
+            new_phys_devs[new_idx]->this_instance = inst;
+            new_phys_devs[new_idx]->phys_dev = temp_pd_array[new_idx];
         }
 
-        // copy wrapped object into Application provided array
-        pPhysicalDevices[i] = (VkPhysicalDevice)new_phys_devs[i];
+        if (new_idx < copy_count) {
+            // Copy wrapped object into Application provided array
+            pPhysicalDevices[new_idx] =
+                (VkPhysicalDevice)new_phys_devs[new_idx];
+        }
     }
 
 out:
-    // if there was no error, free the old buffer and assign the new one
-    if (res == VK_SUCCESS || res == VK_INCOMPLETE) {
-        // free everything that didn't carry over to the new array of physical
-        // devices
-        if (NULL != inst->phys_devs_tramp) {
-            for (uint32_t i = 0; i < inst->phys_dev_count_tramp; i++) {
-                bool found = false;
-                for (uint32_t j = 0; j < new_phys_dev_count; j++) {
-                    if (inst->phys_devs_tramp[i] == new_phys_devs[j]) {
-                        found = true;
-                        break;
+
+    if (NULL != pPhysicalDevices) {
+        // If there was no error, free the old buffer and assign the new one
+        if (res == VK_SUCCESS || res == VK_INCOMPLETE) {
+            // Free everything that didn't carry over to the new array of
+            // physical devices
+            if (NULL != inst->phys_devs_tramp) {
+                for (uint32_t i = 0; i < inst->phys_dev_count_tramp; i++) {
+                    bool found = false;
+                    for (uint32_t j = 0; j < new_pd_count; j++) {
+                        if (inst->phys_devs_tramp[i] == new_phys_devs[j]) {
+                            found = true;
+                            break;
+                        }
+                    }
+                    if (!found) {
+                        loader_instance_heap_free(inst,
+                                                  inst->phys_devs_tramp[i]);
                     }
                 }
-                if (!found) {
-                    loader_instance_heap_free(inst, inst->phys_devs_tramp[i]);
-                }
+                loader_instance_heap_free(inst, inst->phys_devs_tramp);
             }
-            loader_instance_heap_free(inst, inst->phys_devs_tramp);
-        }
 
-        // swap in the new physical device list
-        inst->phys_dev_count_tramp = new_phys_dev_count;
-        inst->phys_devs_tramp = new_phys_devs;
+            // Swap in the new physical device list
+            inst->phys_dev_count_tramp = new_pd_count;
+            inst->phys_devs_tramp = new_phys_devs;
+        } else {
+            for (uint32_t i = 0; i < new_pd_count; i++) {
+                loader_instance_heap_free(inst, new_phys_devs[i]);
+            }
+            loader_instance_heap_free(inst, new_phys_devs);
 
-        // if there was an error, free the new buffer
-    } else {
-        for (uint32_t i = 0; i < new_phys_dev_count; i++) {
-            loader_instance_heap_free(inst, new_phys_devs[i]);
+            // Set the copy count to 0 since something bad happened.
+            copy_count = 0;
         }
-        loader_instance_heap_free(inst, new_phys_devs);
-        *pPhysicalDeviceCount = new_phys_dev_count;
     }
 
+    *pPhysicalDeviceCount = copy_count;
+
     loader_platform_thread_unlock_mutex(&loader_lock);
     return res;
 }