raw_spin_unlock_irqrestore(&vpe->vpe_lock, flags);
 }
 
+static struct irq_chip its_vpe_irq_chip;
+
 static int irq_to_cpuid_lock(struct irq_data *d, unsigned long *flags)
 {
-       struct its_vlpi_map *map = get_vlpi_map(d);
+       struct its_vpe *vpe = NULL;
        int cpu;
 
-       if (map) {
-               cpu = vpe_to_cpuid_lock(map->vpe, flags);
+       if (d->chip == &its_vpe_irq_chip) {
+               vpe = irq_data_get_irq_chip_data(d);
+       } else {
+               struct its_vlpi_map *map = get_vlpi_map(d);
+               if (map)
+                       vpe = map->vpe;
+       }
+
+       if (vpe) {
+               cpu = vpe_to_cpuid_lock(vpe, flags);
        } else {
                /* Physical LPIs are already locked via the irq_desc lock */
                struct its_device *its_dev = irq_data_get_irq_chip_data(d);
 
 static void irq_to_cpuid_unlock(struct irq_data *d, unsigned long flags)
 {
-       struct its_vlpi_map *map = get_vlpi_map(d);
+       struct its_vpe *vpe = NULL;
+
+       if (d->chip == &its_vpe_irq_chip) {
+               vpe = irq_data_get_irq_chip_data(d);
+       } else {
+               struct its_vlpi_map *map = get_vlpi_map(d);
+               if (map)
+                       vpe = map->vpe;
+       }
 
-       if (map)
-               vpe_to_cpuid_unlock(map->vpe, flags);
+       if (vpe)
+               vpe_to_cpuid_unlock(vpe, flags);
 }
 
 static struct its_collection *valid_col(struct its_collection *col)
                cpu_relax();
 }
 
-static void direct_lpi_inv(struct irq_data *d)
+static void __direct_lpi_inv(struct irq_data *d, u64 val)
 {
-       struct its_vlpi_map *map = get_vlpi_map(d);
        void __iomem *rdbase;
        unsigned long flags;
-       u64 val;
        int cpu;
 
+       /* Target the redistributor this LPI is currently routed to */
+       cpu = irq_to_cpuid_lock(d, &flags);
+       raw_spin_lock(&gic_data_rdist_cpu(cpu)->rd_lock);
+
+       rdbase = per_cpu_ptr(gic_rdists->rdist, cpu)->rd_base;
+       gic_write_lpir(val, rdbase + GICR_INVLPIR);
+       wait_for_syncr(rdbase);
+
+       raw_spin_unlock(&gic_data_rdist_cpu(cpu)->rd_lock);
+       irq_to_cpuid_unlock(d, flags);
+}
+
+static void direct_lpi_inv(struct irq_data *d)
+{
+       struct its_vlpi_map *map = get_vlpi_map(d);
+       u64 val;
+
        if (map) {
                struct its_device *its_dev = irq_data_get_irq_chip_data(d);
 
                val = d->hwirq;
        }
 
-       /* Target the redistributor this LPI is currently routed to */
-       cpu = irq_to_cpuid_lock(d, &flags);
-       raw_spin_lock(&gic_data_rdist_cpu(cpu)->rd_lock);
-       rdbase = per_cpu_ptr(gic_rdists->rdist, cpu)->rd_base;
-       gic_write_lpir(val, rdbase + GICR_INVLPIR);
-
-       wait_for_syncr(rdbase);
-       raw_spin_unlock(&gic_data_rdist_cpu(cpu)->rd_lock);
-       irq_to_cpuid_unlock(d, flags);
+       __direct_lpi_inv(d, val);
 }
 
 static void lpi_update_config(struct irq_data *d, u8 clr, u8 set)
 {
        struct its_vpe *vpe = irq_data_get_irq_chip_data(d);
 
-       if (gic_rdists->has_direct_lpi) {
-               void __iomem *rdbase;
-
-               /* Target the redistributor this VPE is currently known on */
-               raw_spin_lock(&gic_data_rdist_cpu(vpe->col_idx)->rd_lock);
-               rdbase = per_cpu_ptr(gic_rdists->rdist, vpe->col_idx)->rd_base;
-               gic_write_lpir(d->parent_data->hwirq, rdbase + GICR_INVLPIR);
-               wait_for_syncr(rdbase);
-               raw_spin_unlock(&gic_data_rdist_cpu(vpe->col_idx)->rd_lock);
-       } else {
+       if (gic_rdists->has_direct_lpi)
+               __direct_lpi_inv(d, d->parent_data->hwirq);
+       else
                its_vpe_send_cmd(vpe, its_send_inv);
-       }
 }
 
 static void its_vpe_mask_irq(struct irq_data *d)