diff --git a/drivers/pci/host/pci-msm-msi.c b/drivers/pci/host/pci-msm-msi.c
index 83f0b112e26fb728ada65ad684a689431a977765..57603e7f0b2dff142b5e9474ef8e73020ae3c05c 100644
--- a/drivers/pci/host/pci-msm-msi.c
+++ b/drivers/pci/host/pci-msm-msi.c
@@ -1,4 +1,4 @@
-/* Copyright (c) 2014-2019, The Linux Foundation. All rights reserved.
+/* Copyright (c) 2014-2020, The Linux Foundation. All rights reserved.
  *
  * This program is free software; you can redistribute it and/or modify
  * it under the terms of the GNU General Public License version 2 and
@@ -34,7 +34,6 @@
 #define PCIE_MSI_CTRL_INT_N_MASK_OFFS(n) (PCIE_MSI_CTRL_BASE + 0xc + 0xc * n)
 #define PCIE_MSI_CTRL_INT_N_STATUS_OFFS(n) (PCIE_MSI_CTRL_BASE + 0x10 + 0xc * n)
 
-#define MSI_IRQ_NR_GRP (2)
 #define MSI_IRQ_PER_GRP (32)
 
 enum msi_type {
@@ -95,35 +94,27 @@ struct msm_msi_client {
 static void msm_msi_snps_handler(struct irq_desc *desc)
 {
 	struct irq_chip *chip = irq_desc_get_chip(desc);
-	struct msm_msi *msi;
+	struct msm_msi_grp *msi_grp;
 	int i;
+	u32 status, mask;
 
 	chained_irq_enter(chip, desc);
 
-	msi = irq_desc_get_handler_data(desc);
+	msi_grp = irq_desc_get_handler_data(desc);
 
-	for (i = 0; i < msi->nr_grps; i++) {
-		struct msm_msi_grp *msi_grp = &msi->grps[i];
-		u32 mask = msi_grp->mask;
-		u32 status;
-		u32 index;
+	status = readl_relaxed(msi_grp->int_status_reg);
 
-		status = readl_relaxed(msi_grp->int_status_reg);
-		if (!status)
-			continue;
+	/* always update the mask set in msm_msi_snps_mask_irq */
+	mask = msi_grp->mask;
+	writel_relaxed(mask, msi_grp->int_mask_reg);
 
-		/* always update the mask set in msm_msi_snps_mask_irq */
-		mask = msi_grp->mask;
-		writel_relaxed(mask, msi_grp->int_mask_reg);
+	/* process only interrupts which are not masked */
+	status ^= (status & mask);
+	writel_relaxed(status, msi_grp->int_status_reg);
 
-		/* process only interrupts which are not masked */
-		status ^= (status & mask);
-		writel_relaxed(status, msi_grp->int_status_reg);
-
-		for (index = 0; status; index++, status >>= 1)
-			if (status & 0x1)
-				generic_handle_irq(msi_grp->irqs[index].virq);
-	}
+	for (i = 0; status; i++, status >>= 1)
+		if (status & 0x1)
+			generic_handle_irq(msi_grp->irqs[i].virq);
 
 	chained_irq_exit(chip, desc);
 }
@@ -483,6 +474,103 @@ static int msm_msi_alloc_domains(struct msm_msi *msi)
 	return 0;
 }
 
+static int msm_msi_snps_irq_setup(struct msm_msi *msi)
+{
+	int i, index, ret;
+	struct msm_msi_grp *msi_grp;
+	struct msm_msi_irq *msi_irq;
+	unsigned int irq = 0;
+
+	/* setup each MSI group. nr_hwirqs == nr_grps */
+	for (i = 0; i < msi->nr_hwirqs; i++) {
+		irq = irq_of_parse_and_map(msi->of_node, i);
+		if (!irq) {
+			dev_err(msi->dev,
+				"MSI: failed to parse/map interrupt\n");
+			ret = -ENODEV;
+			goto free_irqs;
+		}
+
+		msi_grp = &msi->grps[i];
+		msi_grp->int_en_reg = msi->pcie_cfg +
+				PCIE_MSI_CTRL_INT_N_EN_OFFS(i);
+		msi_grp->int_mask_reg = msi->pcie_cfg +
+				PCIE_MSI_CTRL_INT_N_MASK_OFFS(i);
+		msi_grp->int_status_reg = msi->pcie_cfg +
+				PCIE_MSI_CTRL_INT_N_STATUS_OFFS(i);
+
+		for (index = 0; index < MSI_IRQ_PER_GRP; index++) {
+			msi_irq = &msi_grp->irqs[index];
+
+			msi_irq->grp = msi_grp;
+			msi_irq->grp_index = index;
+			msi_irq->pos = (i * MSI_IRQ_PER_GRP) + index;
+			msi_irq->hwirq = irq;
+		}
+
+		irq_set_chained_handler_and_data(irq, msm_msi_snps_handler,
+						msi_grp);
+	}
+
+	return 0;
+
+free_irqs:
+	for (--i; i >= 0; i--) {
+		irq = msi->grps[i].irqs[0].hwirq;
+
+		irq_set_chained_handler_and_data(irq, NULL, NULL);
+		irq_dispose_mapping(irq);
+	}
+
+	return ret;
+}
+
+static int msm_msi_qgic_irq_setup(struct msm_msi *msi)
+{
+	int i, ret;
+	u32 index, grp;
+	struct msm_msi_grp *msi_grp;
+	struct msm_msi_irq *msi_irq;
+	unsigned int irq = 0;
+
+	for (i = 0; i < msi->nr_hwirqs; i++) {
+		irq = irq_of_parse_and_map(msi->of_node, i);
+		if (!irq) {
+			dev_err(msi->dev,
+				"MSI: failed to parse/map interrupt\n");
+			ret = -ENODEV;
+			goto free_irqs;
+		}
+
+		grp = i / MSI_IRQ_PER_GRP;
+		index = i % MSI_IRQ_PER_GRP;
+		msi_grp = &msi->grps[grp];
+		msi_irq = &msi_grp->irqs[index];
+
+		msi_irq->grp = msi_grp;
+		msi_irq->grp_index = index;
+		msi_irq->pos = i;
+		msi_irq->hwirq = irq;
+
+		irq_set_chained_handler_and_data(irq, msm_msi_qgic_handler,
+						msi);
+	}
+
+	return 0;
+
+free_irqs:
+	for (--i; i >= 0; i--) {
+		grp = i / MSI_IRQ_PER_GRP;
+		index = i % MSI_IRQ_PER_GRP;
+		irq = msi->grps[grp].irqs[index].hwirq;
+
+		irq_set_chained_handler_and_data(irq, NULL, NULL);
+		irq_dispose_mapping(irq);
+	}
+
+	return ret;
+}
+
 /* control access to PCIe MSI registers */
 void msm_msi_config_access(struct irq_domain *domain, bool allow)
 {
@@ -524,14 +612,12 @@ EXPORT_SYMBOL(msm_msi_config);
 
 int msm_msi_init(struct device *dev)
 {
-	int i, ret;
+	int ret;
 	struct msm_msi *msi;
 	struct device_node *of_node;
 	const __be32 *prop_val;
 	struct resource *res;
-	void (*msi_handler)(struct irq_desc *);
-	u32 grp;
-	u32 index;
+	int (*msi_irq_setup)(struct msm_msi *);
 
 	if (!dev->of_node) {
 		dev_err(dev, "MSI: missing DT node\n");
@@ -599,17 +685,17 @@ int msm_msi_init(struct device *dev)
 		if (!msi->pcie_cfg)
 			return -ENOMEM;
 
-		msi->nr_virqs = MSI_IRQ_NR_GRP * MSI_IRQ_PER_GRP;
-		msi->nr_grps = MSI_IRQ_NR_GRP;
+		msi->nr_virqs = msi->nr_hwirqs * MSI_IRQ_PER_GRP;
+		msi->nr_grps = msi->nr_hwirqs;
 		msi->mask_irq = msm_msi_snps_mask_irq;
 		msi->unmask_irq = msm_msi_snps_unmask_irq;
-		msi_handler = msm_msi_snps_handler;
+		msi_irq_setup = msm_msi_snps_irq_setup;
 	} else {
 		msi->nr_virqs = msi->nr_hwirqs;
 		msi->nr_grps = 1;
 		msi->mask_irq = msm_msi_qgic_mask_irq;
 		msi->unmask_irq = msm_msi_qgic_unmask_irq;
-		msi_handler = msm_msi_qgic_handler;
+		msi_irq_setup = msm_msi_qgic_irq_setup;
 	}
 
 	msi->grps = devm_kcalloc(msi->dev, msi->nr_grps,
@@ -628,78 +714,15 @@ int msm_msi_init(struct device *dev)
 		return ret;
 	}
 
-	for (i = 0; i < msi->nr_hwirqs; i++) {
-		unsigned int irq = irq_of_parse_and_map(msi->of_node, i);
-		struct msm_msi_grp *msi_grp;
-		struct msm_msi_irq *msi_irq;
-
-		if (!irq) {
-			dev_err(msi->dev,
-				"MSI: failed to parse/map interrupt\n");
-			ret = -ENODEV;
-			goto free_irqs;
-		}
-
-		grp = i / MSI_IRQ_PER_GRP;
-		index = i % MSI_IRQ_PER_GRP;
-		msi_grp = &msi->grps[grp];
-		msi_irq = &msi_grp->irqs[index];
-
-		msi_irq->grp = msi_grp;
-		msi_irq->grp_index = index;
-		msi_irq->pos = i;
-		msi_irq->hwirq = irq;
-
-		irq_set_chained_handler_and_data(irq, msi_handler, msi);
-	}
-
-	if (msi->type == MSM_MSI_TYPE_SNPS) {
-		struct msm_msi_grp *msi_grp;
-		/* all SNPS VIRQs are mapped to one PCIe MSI HWIRQ */
-		u32 snps_hwirq = msi->grps[0].irqs[0].hwirq;
-
-		for (i = 0; i < msi->nr_grps; i++) {
-			msi_grp = &msi->grps[i];
-
-			msi_grp->int_en_reg = msi->pcie_cfg +
-					PCIE_MSI_CTRL_INT_N_EN_OFFS(i);
-			msi_grp->int_mask_reg = msi->pcie_cfg +
-					PCIE_MSI_CTRL_INT_N_MASK_OFFS(i);
-			msi_grp->int_status_reg = msi->pcie_cfg +
-					PCIE_MSI_CTRL_INT_N_STATUS_OFFS(i);
-		}
-
-		for (i = 0; i < msi->nr_virqs; i++) {
-			struct msm_msi_irq *msi_irq;
-
-			grp = i / MSI_IRQ_PER_GRP;
-			index = i % MSI_IRQ_PER_GRP;
-			msi_grp = &msi->grps[grp];
-			msi_irq = &msi_grp->irqs[index];
-
-			msi_irq->grp = msi_grp;
-			msi_irq->grp_index = index;
-			msi_irq->pos = i;
-			msi_irq->hwirq = snps_hwirq;
-		}
-	}
+	ret = msi_irq_setup(msi);
+	if (ret)
+		goto remove_domains;
 
 	msm_msi_config(msi->msi_domain);
 
 	return 0;
 
-free_irqs:
-	for (--i; i >= 0; i--) {
-		u32 hwirq;
-
-		grp = i / MSI_IRQ_PER_GRP;
-		index = i % MSI_IRQ_PER_GRP;
-		hwirq = msi->grps[grp].irqs[index].hwirq;
-
-		irq_set_chained_handler_and_data(hwirq, NULL, NULL);
-		irq_dispose_mapping(hwirq);
-	}
-
+remove_domains:
 	irq_domain_remove(msi->msi_domain);
 	irq_domain_remove(msi->inner_domain);