MODULE_LICENSE("GPL v2");
 MODULE_AUTHOR("Joerg Roedel <jroedel@suse.de>");
 
-#define MAX_DEVICES            0x10000
 #define PRI_QUEUE_SIZE         512
 
 struct pri_queue {
 {
        struct iommu_group *group;
 
+       /* Get rid of any remaining pasid states */
+       free_pasid_states(dev_state);
+
+       /*
+        * Wait until the last reference is dropped before freeing
+        * the device state.
+        */
+       wait_event(dev_state->wq, !atomic_read(&dev_state->count));
+
        /*
         * First detach device from domain - No more PRI requests will arrive
         * from that device after it is unbound from the IOMMUv2 domain.
 
        spin_unlock_irqrestore(&state_lock, flags);
 
-       /* Get rid of any remaining pasid states */
-       free_pasid_states(dev_state);
-
        put_device_state(dev_state);
-       /*
-        * Wait until the last reference is dropped before freeing
-        * the device state.
-        */
-       wait_event(dev_state->wq, !atomic_read(&dev_state->count));
        free_device_state(dev_state);
 }
 EXPORT_SYMBOL(amd_iommu_free_device);
 
 static void __exit amd_iommu_v2_exit(void)
 {
-       struct device_state *dev_state;
-       int i;
+       struct device_state *dev_state, *next;
+       unsigned long flags;
 
        if (!amd_iommu_v2_supported())
                return;
         * The loop below might call flush_workqueue(), so call
         * destroy_workqueue() after it
         */
-       for (i = 0; i < MAX_DEVICES; ++i) {
-               dev_state = get_device_state(i);
-
-               if (dev_state == NULL)
-                       continue;
+       spin_lock_irqsave(&state_lock, flags);
 
+       list_for_each_entry_safe(dev_state, next, &state_list, list) {
                WARN_ON_ONCE(1);
 
                put_device_state(dev_state);
-               amd_iommu_free_device(dev_state->pdev);
+               list_del(&dev_state->list);
+               free_device_state(dev_state);
        }
 
+       spin_unlock_irqrestore(&state_lock, flags);
+
        destroy_workqueue(iommu_wq);
 }