diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index 8ba7dcf..440409f 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -350,7 +350,7 @@ struct kvm_gsi_route_kernel_entry {
 	struct hlist_node link;
 };
 
-void kvm_set_irq(struct kvm *kvm, int irq_source_id, int irq, int level);
+void kvm_set_irq(struct kvm *kvm, int irq_source_id, u32 gsi, int level);
 void kvm_notify_acked_irq(struct kvm *kvm, unsigned gsi);
 void kvm_register_irq_ack_notifier(struct kvm *kvm,
 				   struct kvm_irq_ack_notifier *kian);
diff --git a/virt/kvm/irq_comm.c b/virt/kvm/irq_comm.c
index 0c3266a..9e6c622 100644
--- a/virt/kvm/irq_comm.c
+++ b/virt/kvm/irq_comm.c
@@ -24,10 +24,81 @@
 
 #include "ioapic.h"
 
+#ifdef CONFIG_X86
+#include <asm/msidef.h>
+#endif
+
+static void gsi_dispatch(struct kvm *kvm, u32 gsi)
+{
+	int vcpu_id;
+	struct kvm_vcpu *vcpu;
+	struct kvm_ioapic *ioapic = ioapic_irqchip(kvm);
+	struct kvm_gsi_route_kernel_entry *gsi_entry;
+	int dest_id, vector, dest_mode, trig_mode, delivery_mode;
+	u32 deliver_bitmask;
+
+	BUG_ON(!ioapic);
+
+	gsi_entry = kvm_find_gsi_route_entry(kvm, gsi);
+	if (!gsi_entry) {
+		printk(KERN_WARNING "kvm: fail to find correlated gsi entry\n");
+		return;
+	}
+
+#ifdef CONFIG_X86
+	if (gsi_entry->type == KVM_GSI_ROUTE_TYPE_MSI) {
+		dest_id = (gsi_entry->msi.address_lo & MSI_ADDR_DEST_ID_MASK)
+			>> MSI_ADDR_DEST_ID_SHIFT;
+		vector = (gsi_entry->msi.data & MSI_DATA_VECTOR_MASK)
+			>> MSI_DATA_VECTOR_SHIFT;
+		dest_mode = test_bit(MSI_ADDR_DEST_MODE_SHIFT,
+				(unsigned long *)&gsi_entry->msi.address_lo);
+		trig_mode = test_bit(MSI_DATA_TRIGGER_SHIFT,
+				(unsigned long *)&gsi_entry->msi.data);
+		delivery_mode = test_bit(MSI_DATA_DELIVERY_MODE_SHIFT,
+				(unsigned long *)&gsi_entry->msi.data);
+		deliver_bitmask = kvm_ioapic_get_delivery_bitmask(ioapic,
+				dest_id, dest_mode);
+		/* IOAPIC delivery mode value is the same as MSI here */
+		switch (delivery_mode) {
+		case IOAPIC_LOWEST_PRIORITY:
+			vcpu = kvm_get_lowest_prio_vcpu(ioapic->kvm, vector,
+					deliver_bitmask);
+			if (vcpu != NULL)
+				kvm_apic_set_irq(vcpu, vector, trig_mode);
+			else
+				printk(KERN_INFO
+				       "kvm: null lowest priority vcpu!\n");
+			break;
+		case IOAPIC_FIXED:
+			for (vcpu_id = 0; deliver_bitmask != 0; vcpu_id++) {
+				if (!(deliver_bitmask & (1 << vcpu_id)))
+					continue;
+				deliver_bitmask &= ~(1 << vcpu_id);
+				vcpu = ioapic->kvm->vcpus[vcpu_id];
+				if (vcpu)
+					kvm_apic_set_irq(vcpu, vector,
+							trig_mode);
+			}
+			break;
+		default:
+			break;
+		}
+	}
+#endif /* CONFIG_X86 */
+}
+
 /* This should be called with the kvm->lock mutex held */
-void kvm_set_irq(struct kvm *kvm, int irq_source_id, int irq, int level)
+void kvm_set_irq(struct kvm *kvm, int irq_source_id, u32 gsi, int level)
 {
-	unsigned long *irq_state = (unsigned long *)&kvm->arch.irq_states[irq];
+	unsigned long *irq_state;
+
+	if (gsi & KVM_GSI_ROUTE_ENTRY_MASK) {
+		gsi_dispatch(kvm, gsi);
+		return;
+	}
+
+	irq_state = (unsigned long *)&kvm->arch.irq_states[gsi];
 
 	/* Logical OR for level trig interrupt */
 	if (level)
@@ -39,9 +110,9 @@ void kvm_set_irq(struct kvm *kvm, int irq_source_id, int irq, int level)
 	 * IOAPIC.  So set the bit in both. The guest will ignore
 	 * writes to the unused one.
 	 */
-	kvm_ioapic_set_irq(ioapic_irqchip(kvm), irq, !!(*irq_state));
+	kvm_ioapic_set_irq(ioapic_irqchip(kvm), gsi, !!(*irq_state));
 #ifdef CONFIG_X86
-	kvm_pic_set_irq(pic_irqchip(kvm), irq, !!(*irq_state));
+	kvm_pic_set_irq(pic_irqchip(kvm), gsi, !!(*irq_state));
 #endif
 }
 
diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index b0b97ba..8dba0b2 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -47,10 +47,6 @@
 #include <asm/uaccess.h>
 #include <asm/pgtable.h>
 
-#ifdef CONFIG_X86
-#include <asm/msidef.h>
-#endif
-
 #ifdef KVM_COALESCED_MMIO_PAGE_OFFSET
 #include "coalesced_mmio.h"
 #endif
@@ -86,68 +82,6 @@ static bool kvm_rebooting;
 
 #ifdef KVM_CAP_DEVICE_ASSIGNMENT
 
-#ifdef CONFIG_X86
-static void assigned_device_msi_dispatch(struct kvm_assigned_dev_kernel *dev, u32 gsi)
-{
-	int vcpu_id;
-	struct kvm_vcpu *vcpu;
-	struct kvm_ioapic *ioapic = ioapic_irqchip(dev->kvm);
-	struct kvm_gsi_route_kernel_entry *gsi_entry;
-	int dest_id, vector, dest_mode, trig_mode, delivery_mode;
-	u32 deliver_bitmask;
-
-	BUG_ON(!ioapic);
-
-	gsi_entry = kvm_find_gsi_route_entry(dev->kvm, gsi);
-	if (!gsi_entry) {
-		printk(KERN_WARNING "kvm: fail to find correlated gsi entry\n");
-		return;
-	}
-
-	if (gsi_entry->type == KVM_GSI_ROUTE_TYPE_MSI) {
-		dest_id = (gsi_entry->msi.address_lo & MSI_ADDR_DEST_ID_MASK)
-			>> MSI_ADDR_DEST_ID_SHIFT;
-		vector = (gsi_entry->msi.data & MSI_DATA_VECTOR_MASK)
-			>> MSI_DATA_VECTOR_SHIFT;
-		dest_mode = test_bit(MSI_ADDR_DEST_MODE_SHIFT,
-				(unsigned long *)&gsi_entry->msi.address_lo);
-		trig_mode = test_bit(MSI_DATA_TRIGGER_SHIFT,
-				(unsigned long *)&gsi_entry->msi.data);
-		delivery_mode = test_bit(MSI_DATA_DELIVERY_MODE_SHIFT,
-				(unsigned long *)&gsi_entry->msi.data);
-		deliver_bitmask = kvm_ioapic_get_delivery_bitmask(ioapic,
-				dest_id, dest_mode);
-		/* IOAPIC delivery mode value is the same as MSI here */
-		switch (delivery_mode) {
-		case IOAPIC_LOWEST_PRIORITY:
-			vcpu = kvm_get_lowest_prio_vcpu(ioapic->kvm, vector,
-					deliver_bitmask);
-			if (vcpu != NULL)
-				kvm_apic_set_irq(vcpu, vector, trig_mode);
-			else
-				printk(KERN_INFO
-				       "kvm: null lowest priority vcpu!\n");
-			break;
-		case IOAPIC_FIXED:
-			for (vcpu_id = 0; deliver_bitmask != 0; vcpu_id++) {
-				if (!(deliver_bitmask & (1 << vcpu_id)))
-					continue;
-				deliver_bitmask &= ~(1 << vcpu_id);
-				vcpu = ioapic->kvm->vcpus[vcpu_id];
-				if (vcpu)
-					kvm_apic_set_irq(vcpu, vector,
-							trig_mode);
-			}
-			break;
-		default:
-			break;
-		}
-	}
-}
-#else
-static void assigned_device_msi_dispatch(struct kvm_assigned_dev_kernel *dev, u32 gsi) {}
-#endif
-
 static struct kvm_assigned_dev_kernel *kvm_find_assigned_dev(struct list_head *head,
 						      int assigned_dev_id)
 {
@@ -174,13 +108,11 @@ static void kvm_assigned_dev_interrupt_work_handler(struct work_struct *work)
 	 * finer-grained lock, update this
 	 */
 	mutex_lock(&assigned_dev->kvm->lock);
-	if (assigned_dev->irq_requested_type & KVM_ASSIGNED_DEV_GUEST_INTX)
-		kvm_set_irq(assigned_dev->kvm,
-			    assigned_dev->irq_source_id,
-			    assigned_dev->guest_irq, 1);
-	else if (assigned_dev->irq_requested_type &
-				KVM_ASSIGNED_DEV_GUEST_MSI) {
-		assigned_device_msi_dispatch(assigned_dev, assigned_dev->guest_irq);
+
+	kvm_set_irq(assigned_dev->kvm, assigned_dev->irq_source_id,
+			assigned_dev->guest_irq, 1);
+
+	if (assigned_dev->irq_requested_type & KVM_ASSIGNED_DEV_GUEST_MSI) {
 		enable_irq(assigned_dev->host_irq);
 		assigned_dev->host_irq_disabled = false;
 	}
