diff --git a/drivers/net/wireless/ath/ath11k/core.c b/drivers/net/wireless/ath/ath11k/core.c
index 2a867237e817..02d407011a8b 100644
--- a/drivers/net/wireless/ath/ath11k/core.c
+++ b/drivers/net/wireless/ath/ath11k/core.c
@@ -96,7 +96,6 @@ static const struct ath11k_hw_params ath11k_hw_params[] = {
 		.hal_params = &ath11k_hw_hal_params_ipq8074,
 		.supports_dynamic_smps_6ghz = false,
 		.alloc_cacheable_memory = true,
-		.wakeup_mhi = false,
 		.supports_rssi_stats = false,
 		.fw_wmi_diag_event = false,
 		.current_cc_support = false,
@@ -163,7 +162,6 @@ static const struct ath11k_hw_params ath11k_hw_params[] = {
 		.hal_params = &ath11k_hw_hal_params_ipq8074,
 		.supports_dynamic_smps_6ghz = false,
 		.alloc_cacheable_memory = true,
-		.wakeup_mhi = false,
 		.supports_rssi_stats = false,
 		.fw_wmi_diag_event = false,
 		.current_cc_support = false,
@@ -229,7 +227,6 @@ static const struct ath11k_hw_params ath11k_hw_params[] = {
 		.hal_params = &ath11k_hw_hal_params_qca6390,
 		.supports_dynamic_smps_6ghz = false,
 		.alloc_cacheable_memory = false,
-		.wakeup_mhi = true,
 		.supports_rssi_stats = true,
 		.fw_wmi_diag_event = true,
 		.current_cc_support = true,
@@ -295,7 +292,6 @@ static const struct ath11k_hw_params ath11k_hw_params[] = {
 		.hal_params = &ath11k_hw_hal_params_ipq8074,
 		.supports_dynamic_smps_6ghz = true,
 		.alloc_cacheable_memory = true,
-		.wakeup_mhi = false,
 		.supports_rssi_stats = false,
 		.fw_wmi_diag_event = false,
 		.current_cc_support = false,
@@ -361,7 +357,6 @@ static const struct ath11k_hw_params ath11k_hw_params[] = {
 		.hal_params = &ath11k_hw_hal_params_qca6390,
 		.supports_dynamic_smps_6ghz = false,
 		.alloc_cacheable_memory = false,
-		.wakeup_mhi = true,
 		.supports_rssi_stats = true,
 		.fw_wmi_diag_event = true,
 		.current_cc_support = true,
@@ -426,7 +421,6 @@ static const struct ath11k_hw_params ath11k_hw_params[] = {
 		.hal_params = &ath11k_hw_hal_params_qca6390,
 		.supports_dynamic_smps_6ghz = false,
 		.alloc_cacheable_memory = false,
-		.wakeup_mhi = true,
 		.supports_rssi_stats = true,
 		.fw_wmi_diag_event = true,
 		.current_cc_support = true,
diff --git a/drivers/net/wireless/ath/ath11k/core.h b/drivers/net/wireless/ath/ath11k/core.h
index 64f119e6882f..406d16b8735d 100644
--- a/drivers/net/wireless/ath/ath11k/core.h
+++ b/drivers/net/wireless/ath/ath11k/core.h
@@ -231,6 +231,8 @@ enum ath11k_dev_flags {
 	ATH11K_FLAG_CE_IRQ_ENABLED,
 	ATH11K_FLAG_EXT_IRQ_ENABLED,
 	ATH11K_FLAG_FIXED_MEM_RGN,
+	ATH11K_FLAG_DEVICE_INIT_DONE,
+	ATH11K_FLAG_MULTI_MSI_VECTORS,
 };
 
 enum ath11k_monitor_flags {
@@ -720,6 +722,14 @@ struct ath11k_bus_params {
 	bool static_window_map;
 };
 
+struct ath11k_pci_ops {
+	int (*wakeup)(struct ath11k_base *ab);
+	void (*release)(struct ath11k_base *ab);
+	int (*get_msi_irq)(struct ath11k_base *ab, unsigned int vector);
+	void (*window_write32)(struct ath11k_base *ab, u32 offset, u32 value);
+	u32 (*window_read32)(struct ath11k_base *ab, u32 offset);
+};
+
 /* IPQ8074 HW channel counters frequency value in hertz */
 #define IPQ8074_CC_FREQ_HERTZ 320000
 
@@ -917,6 +927,8 @@ struct ath11k_base {
 			u32 addr_lo;
 			u32 addr_hi;
 		} msi;
+
+		const struct ath11k_pci_ops *ops;
 	} pci;
 
 	/* must be last */
diff --git a/drivers/net/wireless/ath/ath11k/hw.h b/drivers/net/wireless/ath/ath11k/hw.h
index 5955c6f5da24..e1bcdc538b0e 100644
--- a/drivers/net/wireless/ath/ath11k/hw.h
+++ b/drivers/net/wireless/ath/ath11k/hw.h
@@ -1,6 +1,7 @@
 /* SPDX-License-Identifier: BSD-3-Clause-Clear */
 /*
  * Copyright (c) 2018-2019 The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021-2022, Qualcomm Innovation Center, Inc. All rights reserved.
  */
 
 #ifndef ATH11K_HW_H
@@ -189,7 +190,6 @@ struct ath11k_hw_params {
 	const struct ath11k_hw_hal_params *hal_params;
 	bool supports_dynamic_smps_6ghz;
 	bool alloc_cacheable_memory;
-	bool wakeup_mhi;
 	bool supports_rssi_stats;
 	bool fw_wmi_diag_event;
 	bool current_cc_support;
diff --git a/drivers/net/wireless/ath/ath11k/mhi.c b/drivers/net/wireless/ath/ath11k/mhi.c
index 5de255969364..1aa1e0f01b85 100644
--- a/drivers/net/wireless/ath/ath11k/mhi.c
+++ b/drivers/net/wireless/ath/ath11k/mhi.c
@@ -273,10 +273,10 @@ static int ath11k_mhi_get_msi(struct ath11k_pci *ab_pci)
 	for (i = 0; i < num_vectors; i++) {
 		msi_data = base_vector;
 
-		if (test_bit(ATH11K_PCI_FLAG_MULTI_MSI_VECTORS, &ab_pci->flags))
+		if (test_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags))
 			msi_data += i;
 
-		irq[i] = ath11k_pcic_get_msi_irq(ab->dev, msi_data);
+		irq[i] = ath11k_pci_get_msi_irq(ab, msi_data);
 	}
 
 	ab_pci->mhi_ctrl->irq = irq;
@@ -406,7 +406,7 @@ int ath11k_mhi_register(struct ath11k_pci *ab_pci)
 		return ret;
 	}
 
-	if (!test_bit(ATH11K_PCI_FLAG_MULTI_MSI_VECTORS, &ab_pci->flags))
+	if (!test_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags))
 		mhi_ctrl->irq_flags = IRQF_SHARED | IRQF_NOBALANCING;
 
 	if (test_bit(ATH11K_FLAG_FIXED_MEM_RGN, &ab->dev_flags)) {
diff --git a/drivers/net/wireless/ath/ath11k/pci.c b/drivers/net/wireless/ath/ath11k/pci.c
index fd5487d7045e..3fd5b416a564 100644
--- a/drivers/net/wireless/ath/ath11k/pci.c
+++ b/drivers/net/wireless/ath/ath11k/pci.c
@@ -36,6 +36,85 @@ static const struct pci_device_id ath11k_pci_id_table[] = {
 
 MODULE_DEVICE_TABLE(pci, ath11k_pci_id_table);
 
+static int ath11k_pci_bus_wake_up(struct ath11k_base *ab)
+{
+	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
+
+	return mhi_device_get_sync(ab_pci->mhi_ctrl->mhi_dev);
+}
+
+static void ath11k_pci_bus_release(struct ath11k_base *ab)
+{
+	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
+
+	mhi_device_put(ab_pci->mhi_ctrl->mhi_dev);
+}
+
+static inline void ath11k_pci_select_window(struct ath11k_pci *ab_pci, u32 offset)
+{
+	struct ath11k_base *ab = ab_pci->ab;
+
+	u32 window = FIELD_GET(ATH11K_PCI_WINDOW_VALUE_MASK, offset);
+
+	lockdep_assert_held(&ab_pci->window_lock);
+
+	if (window != ab_pci->register_window) {
+		iowrite32(ATH11K_PCI_WINDOW_ENABLE_BIT | window,
+			  ab->mem + ATH11K_PCI_WINDOW_REG_ADDRESS);
+		ioread32(ab->mem + ATH11K_PCI_WINDOW_REG_ADDRESS);
+		ab_pci->register_window = window;
+	}
+}
+
+static void
+ath11k_pci_window_write32(struct ath11k_base *ab, u32 offset, u32 value)
+{
+	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
+	u32 window_start = ATH11K_PCI_WINDOW_START;
+
+	spin_lock_bh(&ab_pci->window_lock);
+	ath11k_pci_select_window(ab_pci, offset);
+	iowrite32(value, ab->mem + window_start +
+		  (offset & ATH11K_PCI_WINDOW_RANGE_MASK));
+	spin_unlock_bh(&ab_pci->window_lock);
+}
+
+static u32 ath11k_pci_window_read32(struct ath11k_base *ab, u32 offset)
+{
+	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
+	u32 window_start = ATH11K_PCI_WINDOW_START;
+	u32 val;
+
+	spin_lock_bh(&ab_pci->window_lock);
+	ath11k_pci_select_window(ab_pci, offset);
+	val = ioread32(ab->mem + window_start +
+		       (offset & ATH11K_PCI_WINDOW_RANGE_MASK));
+	spin_unlock_bh(&ab_pci->window_lock);
+
+	return val;
+}
+
+int ath11k_pci_get_msi_irq(struct ath11k_base *ab, unsigned int vector)
+{
+	struct pci_dev *pci_dev = to_pci_dev(ab->dev);
+
+	return pci_irq_vector(pci_dev, vector);
+}
+
+static const struct ath11k_pci_ops ath11k_pci_ops_qca6390 = {
+	.wakeup = ath11k_pci_bus_wake_up,
+	.release = ath11k_pci_bus_release,
+	.get_msi_irq = ath11k_pci_get_msi_irq,
+	.window_write32 = ath11k_pci_window_write32,
+	.window_read32 = ath11k_pci_window_read32,
+};
+
+static const struct ath11k_pci_ops ath11k_pci_ops_qcn9074 = {
+	.get_msi_irq = ath11k_pci_get_msi_irq,
+	.window_write32 = ath11k_pci_window_write32,
+	.window_read32 = ath11k_pci_window_read32,
+};
+
 static const struct ath11k_bus_params ath11k_pci_bus_params = {
 	.mhi_support = true,
 	.m3_fw_support = true,
@@ -318,8 +397,7 @@ static int ath11k_pci_alloc_msi(struct ath11k_pci *ab_pci)
 					    msi_config->total_vectors,
 					    PCI_IRQ_MSI);
 	if (num_vectors == msi_config->total_vectors) {
-		set_bit(ATH11K_PCI_FLAG_MULTI_MSI_VECTORS, &ab_pci->flags);
-		ab_pci->irq_flags = IRQF_SHARED;
+		set_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags);
 	} else {
 		num_vectors = pci_alloc_irq_vectors(ab_pci->pdev,
 						    1,
@@ -329,9 +407,8 @@ static int ath11k_pci_alloc_msi(struct ath11k_pci *ab_pci)
 			ret = -EINVAL;
 			goto reset_msi_config;
 		}
-		clear_bit(ATH11K_PCI_FLAG_MULTI_MSI_VECTORS, &ab_pci->flags);
+		clear_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags);
 		ab->pci.msi.config = &msi_config_one_msi;
-		ab_pci->irq_flags = IRQF_SHARED | IRQF_NOBALANCING;
 		ath11k_dbg(ab, ATH11K_DBG_PCI, "request MSI one vector\n");
 	}
 	ath11k_info(ab, "MSI vectors: %d\n", num_vectors);
@@ -487,13 +564,20 @@ static void ath11k_pci_aspm_disable(struct ath11k_pci *ab_pci)
 	set_bit(ATH11K_PCI_ASPM_RESTORE, &ab_pci->flags);
 }
 
+static void ath11k_pci_aspm_restore(struct ath11k_pci *ab_pci)
+{
+	if (test_and_clear_bit(ATH11K_PCI_ASPM_RESTORE, &ab_pci->flags))
+		pcie_capability_write_word(ab_pci->pdev, PCI_EXP_LNKCTL,
+					   ab_pci->link_ctl);
+}
+
 static int ath11k_pci_power_up(struct ath11k_base *ab)
 {
 	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
 	int ret;
 
 	ab_pci->register_window = 0;
-	clear_bit(ATH11K_PCI_FLAG_INIT_DONE, &ab_pci->flags);
+	clear_bit(ATH11K_FLAG_DEVICE_INIT_DONE, &ab->dev_flags);
 	ath11k_pci_sw_reset(ab_pci->ab, true);
 
 	/* Disable ASPM during firmware download due to problems switching
@@ -520,14 +604,14 @@ static void ath11k_pci_power_down(struct ath11k_base *ab)
 	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
 
 	/* restore aspm in case firmware bootup fails */
-	ath11k_pcic_aspm_restore(ab_pci);
+	ath11k_pci_aspm_restore(ab_pci);
 
 	ath11k_pci_force_wake(ab_pci->ab);
 
 	ath11k_pci_msi_disable(ab_pci);
 
 	ath11k_mhi_stop(ab_pci);
-	clear_bit(ATH11K_PCI_FLAG_INIT_DONE, &ab_pci->flags);
+	clear_bit(ATH11K_FLAG_DEVICE_INIT_DONE, &ab->dev_flags);
 	ath11k_pci_sw_reset(ab_pci->ab, false);
 }
 
@@ -559,8 +643,25 @@ static void ath11k_pci_hif_ce_irq_disable(struct ath11k_base *ab)
 	ath11k_pcic_ce_irq_disable_sync(ab);
 }
 
+static int ath11k_pci_start(struct ath11k_base *ab)
+{
+	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
+
+	/* TODO: for now don't restore ASPM in case of single MSI
+	 * vector as MHI register reading in M2 causes system hang.
+	 */
+	if (test_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags))
+		ath11k_pci_aspm_restore(ab_pci);
+	else
+		ath11k_info(ab, "leaving PCI ASPM disabled to avoid MHI M2 problems\n");
+
+	ath11k_pcic_start(ab);
+
+	return 0;
+}
+
 static const struct ath11k_hif_ops ath11k_pci_hif_ops = {
-	.start = ath11k_pcic_start,
+	.start = ath11k_pci_start,
 	.stop = ath11k_pcic_stop,
 	.read32 = ath11k_pcic_read32,
 	.write32 = ath11k_pcic_write32,
@@ -592,6 +693,15 @@ static void ath11k_pci_read_hw_version(struct ath11k_base *ab, u32 *major, u32 *
 		   *major, *minor);
 }
 
+static int ath11k_pci_set_irq_affinity_hint(struct ath11k_pci *ab_pci,
+					    const struct cpumask *m)
+{
+	if (test_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab_pci->ab->dev_flags))
+		return 0;
+
+	return irq_set_affinity_hint(ab_pci->pdev->irq, m);
+}
+
 static int ath11k_pci_probe(struct pci_dev *pdev,
 			    const struct pci_device_id *pci_dev)
 {
@@ -654,9 +764,12 @@ static int ath11k_pci_probe(struct pci_dev *pdev,
 			ret = -EOPNOTSUPP;
 			goto err_pci_free_region;
 		}
+
+		ab->pci.ops = &ath11k_pci_ops_qca6390;
 		break;
 	case QCN9074_DEVICE_ID:
 		ab->bus_params.static_window_map = true;
+		ab->pci.ops = &ath11k_pci_ops_qcn9074;
 		ab->hw_rev = ATH11K_HW_QCN9074_HW10;
 		break;
 	case WCN6855_DEVICE_ID:
@@ -685,6 +798,8 @@ static int ath11k_pci_probe(struct pci_dev *pdev,
 			ret = -EOPNOTSUPP;
 			goto err_pci_free_region;
 		}
+
+		ab->pci.ops = &ath11k_pci_ops_qca6390;
 		break;
 	default:
 		dev_err(&pdev->dev, "Unknown PCI device found: 0x%x\n",
@@ -733,6 +848,12 @@ static int ath11k_pci_probe(struct pci_dev *pdev,
 		goto err_ce_free;
 	}
 
+	ret = ath11k_pci_set_irq_affinity_hint(ab_pci, cpumask_of(0));
+	if (ret) {
+		ath11k_err(ab, "failed to set irq affinity %d\n", ret);
+		goto err_free_irq;
+	}
+
 	/* kernel may allocate a dummy vector before request_irq and
 	 * then allocate a real vector when request_irq is called.
 	 * So get msi_data here again to avoid spurious interrupt
@@ -741,16 +862,19 @@ static int ath11k_pci_probe(struct pci_dev *pdev,
 	ret = ath11k_pci_config_msi_data(ab_pci);
 	if (ret) {
 		ath11k_err(ab, "failed to config msi_data: %d\n", ret);
-		goto err_free_irq;
+		goto err_irq_affinity_cleanup;
 	}
 
 	ret = ath11k_core_init(ab);
 	if (ret) {
 		ath11k_err(ab, "failed to init core: %d\n", ret);
-		goto err_free_irq;
+		goto err_irq_affinity_cleanup;
 	}
 	return 0;
 
+err_irq_affinity_cleanup:
+	ath11k_pci_set_irq_affinity_hint(ab_pci, NULL);
+
 err_free_irq:
 	ath11k_pcic_free_irq(ab);
 
@@ -780,7 +904,7 @@ static void ath11k_pci_remove(struct pci_dev *pdev)
 	struct ath11k_base *ab = pci_get_drvdata(pdev);
 	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
 
-	ath11k_pcic_set_irq_affinity_hint(ab_pci, NULL);
+	ath11k_pci_set_irq_affinity_hint(ab_pci, NULL);
 
 	if (test_bit(ATH11K_FLAG_QMI_FAIL, &ab->dev_flags)) {
 		ath11k_pci_power_down(ab);
diff --git a/drivers/net/wireless/ath/ath11k/pci.h b/drivers/net/wireless/ath/ath11k/pci.h
index c9cae48d784e..16a000b9cc5e 100644
--- a/drivers/net/wireless/ath/ath11k/pci.h
+++ b/drivers/net/wireless/ath/ath11k/pci.h
@@ -54,9 +54,7 @@
 #define QFPROM_PWR_CTRL_VDD4BLOW_MASK		0x4
 
 enum ath11k_pci_flags {
-	ATH11K_PCI_FLAG_INIT_DONE,
 	ATH11K_PCI_ASPM_RESTORE,
-	ATH11K_PCI_FLAG_MULTI_MSI_VECTORS,
 };
 
 struct ath11k_pci {
@@ -74,8 +72,6 @@ struct ath11k_pci {
 	/* enum ath11k_pci_flags */
 	unsigned long flags;
 	u16 link_ctl;
-
-	unsigned long irq_flags;
 };
 
 static inline struct ath11k_pci *ath11k_pci_priv(struct ath11k_base *ab)
@@ -83,4 +79,5 @@ static inline struct ath11k_pci *ath11k_pci_priv(struct ath11k_base *ab)
 	return (struct ath11k_pci *)ab->drv_priv;
 }
 
+int ath11k_pci_get_msi_irq(struct ath11k_base *ab, unsigned int vector);
 #endif
diff --git a/drivers/net/wireless/ath/ath11k/pcic.c b/drivers/net/wireless/ath/ath11k/pcic.c
index c743095387f1..63c678aea29e 100644
--- a/drivers/net/wireless/ath/ath11k/pcic.c
+++ b/drivers/net/wireless/ath/ath11k/pcic.c
@@ -4,7 +4,6 @@
  * Copyright (c) 2021-2022, Qualcomm Innovation Center, Inc. All rights reserved.
  */
 
-#include <linux/pci.h>
 #include "core.h"
 #include "pcic.h"
 #include "debug.h"
@@ -132,29 +131,6 @@ int ath11k_pcic_init_msi_config(struct ath11k_base *ab)
 }
 EXPORT_SYMBOL(ath11k_pcic_init_msi_config);
 
-void ath11k_pcic_aspm_restore(struct ath11k_pci *ab_pci)
-{
-	if (test_and_clear_bit(ATH11K_PCI_ASPM_RESTORE, &ab_pci->flags))
-		pcie_capability_write_word(ab_pci->pdev, PCI_EXP_LNKCTL,
-					   ab_pci->link_ctl);
-}
-
-static inline void ath11k_pcic_select_window(struct ath11k_pci *ab_pci, u32 offset)
-{
-	struct ath11k_base *ab = ab_pci->ab;
-
-	u32 window = FIELD_GET(ATH11K_PCI_WINDOW_VALUE_MASK, offset);
-
-	lockdep_assert_held(&ab_pci->window_lock);
-
-	if (window != ab_pci->register_window) {
-		iowrite32(ATH11K_PCI_WINDOW_ENABLE_BIT | window,
-			  ab->mem + ATH11K_PCI_WINDOW_REG_ADDRESS);
-		ioread32(ab->mem + ATH11K_PCI_WINDOW_REG_ADDRESS);
-		ab_pci->register_window = window;
-	}
-}
-
 static inline u32 ath11k_pcic_get_window_start(struct ath11k_base *ab,
 					       u32 offset)
 {
@@ -174,17 +150,15 @@ static inline u32 ath11k_pcic_get_window_start(struct ath11k_base *ab,
 
 void ath11k_pcic_write32(struct ath11k_base *ab, u32 offset, u32 value)
 {
-	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
 	u32 window_start;
 	int ret = 0;
 
 	/* for offset beyond BAR + 4K - 32, may
-	 * need to wakeup MHI to access.
+	 * need to wakeup the device to access.
 	 */
-	if (ab->hw_params.wakeup_mhi &&
-	    test_bit(ATH11K_PCI_FLAG_INIT_DONE, &ab_pci->flags) &&
-	    offset >= ATH11K_PCI_ACCESS_ALWAYS_OFF)
-		ret = mhi_device_get_sync(ab_pci->mhi_ctrl->mhi_dev);
+	if (test_bit(ATH11K_FLAG_DEVICE_INIT_DONE, &ab->dev_flags) &&
+	    offset >= ATH11K_PCI_ACCESS_ALWAYS_OFF && ab->pci.ops->wakeup)
+		ret = ab->pci.ops->wakeup(ab);
 
 	if (offset < ATH11K_PCI_WINDOW_START) {
 		iowrite32(value, ab->mem  + offset);
@@ -194,38 +168,32 @@ void ath11k_pcic_write32(struct ath11k_base *ab, u32 offset, u32 value)
 		else
 			window_start = ATH11K_PCI_WINDOW_START;
 
-		if (window_start == ATH11K_PCI_WINDOW_START) {
-			spin_lock_bh(&ab_pci->window_lock);
-			ath11k_pcic_select_window(ab_pci, offset);
-			iowrite32(value, ab->mem + window_start +
-				  (offset & ATH11K_PCI_WINDOW_RANGE_MASK));
-			spin_unlock_bh(&ab_pci->window_lock);
+		if (window_start == ATH11K_PCI_WINDOW_START &&
+		    ab->pci.ops->window_write32) {
+			ab->pci.ops->window_write32(ab, offset, value);
 		} else {
 			iowrite32(value, ab->mem + window_start +
 				  (offset & ATH11K_PCI_WINDOW_RANGE_MASK));
 		}
 	}
 
-	if (ab->hw_params.wakeup_mhi &&
-	    test_bit(ATH11K_PCI_FLAG_INIT_DONE, &ab_pci->flags) &&
-	    offset >= ATH11K_PCI_ACCESS_ALWAYS_OFF &&
+	if (test_bit(ATH11K_FLAG_DEVICE_INIT_DONE, &ab->dev_flags) &&
+	    offset >= ATH11K_PCI_ACCESS_ALWAYS_OFF && ab->pci.ops->release &&
 	    !ret)
-		mhi_device_put(ab_pci->mhi_ctrl->mhi_dev);
+		ab->pci.ops->release(ab);
 }
 
 u32 ath11k_pcic_read32(struct ath11k_base *ab, u32 offset)
 {
-	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
 	u32 val, window_start;
 	int ret = 0;
 
 	/* for offset beyond BAR + 4K - 32, may
-	 * need to wakeup MHI to access.
+	 * need to wakeup the device to access.
 	 */
-	if (ab->hw_params.wakeup_mhi &&
-	    test_bit(ATH11K_PCI_FLAG_INIT_DONE, &ab_pci->flags) &&
-	    offset >= ATH11K_PCI_ACCESS_ALWAYS_OFF)
-		ret = mhi_device_get_sync(ab_pci->mhi_ctrl->mhi_dev);
+	if (test_bit(ATH11K_FLAG_DEVICE_INIT_DONE, &ab->dev_flags) &&
+	    offset >= ATH11K_PCI_ACCESS_ALWAYS_OFF && ab->pci.ops->wakeup)
+		ret = ab->pci.ops->wakeup(ab);
 
 	if (offset < ATH11K_PCI_WINDOW_START) {
 		val = ioread32(ab->mem + offset);
@@ -235,34 +203,23 @@ u32 ath11k_pcic_read32(struct ath11k_base *ab, u32 offset)
 		else
 			window_start = ATH11K_PCI_WINDOW_START;
 
-		if (window_start == ATH11K_PCI_WINDOW_START) {
-			spin_lock_bh(&ab_pci->window_lock);
-			ath11k_pcic_select_window(ab_pci, offset);
-			val = ioread32(ab->mem + window_start +
-				       (offset & ATH11K_PCI_WINDOW_RANGE_MASK));
-			spin_unlock_bh(&ab_pci->window_lock);
+		if (window_start == ATH11K_PCI_WINDOW_START &&
+		    ab->pci.ops->window_read32) {
+			val = ab->pci.ops->window_read32(ab, offset);
 		} else {
 			val = ioread32(ab->mem + window_start +
 				       (offset & ATH11K_PCI_WINDOW_RANGE_MASK));
 		}
 	}
 
-	if (ab->hw_params.wakeup_mhi &&
-	    test_bit(ATH11K_PCI_FLAG_INIT_DONE, &ab_pci->flags) &&
-	    offset >= ATH11K_PCI_ACCESS_ALWAYS_OFF &&
+	if (test_bit(ATH11K_FLAG_DEVICE_INIT_DONE, &ab->dev_flags) &&
+	    offset >= ATH11K_PCI_ACCESS_ALWAYS_OFF && ab->pci.ops->release &&
 	    !ret)
-		mhi_device_put(ab_pci->mhi_ctrl->mhi_dev);
+		ab->pci.ops->release(ab);
 
 	return val;
 }
 
-int ath11k_pcic_get_msi_irq(struct device *dev, unsigned int vector)
-{
-	struct pci_dev *pci_dev = to_pci_dev(dev);
-
-	return pci_irq_vector(pci_dev, vector);
-}
-
 void ath11k_pcic_get_msi_address(struct ath11k_base *ab, u32 *msi_addr_lo,
 				 u32 *msi_addr_hi)
 {
@@ -343,13 +300,12 @@ void ath11k_pcic_free_irq(struct ath11k_base *ab)
 
 static void ath11k_pcic_ce_irq_enable(struct ath11k_base *ab, u16 ce_id)
 {
-	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
 	u32 irq_idx;
 
 	/* In case of one MSI vector, we handle irq enable/disable in a
 	 * uniform way since we only have one irq
 	 */
-	if (!test_bit(ATH11K_PCI_FLAG_MULTI_MSI_VECTORS, &ab_pci->flags))
+	if (!test_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags))
 		return;
 
 	irq_idx = ATH11K_PCI_IRQ_CE0_OFFSET + ce_id;
@@ -358,13 +314,12 @@ static void ath11k_pcic_ce_irq_enable(struct ath11k_base *ab, u16 ce_id)
 
 static void ath11k_pcic_ce_irq_disable(struct ath11k_base *ab, u16 ce_id)
 {
-	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
 	u32 irq_idx;
 
 	/* In case of one MSI vector, we handle irq enable/disable in a
 	 * uniform way since we only have one irq
 	 */
-	if (!test_bit(ATH11K_PCI_FLAG_MULTI_MSI_VECTORS, &ab_pci->flags))
+	if (!test_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags))
 		return;
 
 	irq_idx = ATH11K_PCI_IRQ_CE0_OFFSET + ce_id;
@@ -429,13 +384,13 @@ static irqreturn_t ath11k_pcic_ce_interrupt_handler(int irq, void *arg)
 
 static void ath11k_pcic_ext_grp_disable(struct ath11k_ext_irq_grp *irq_grp)
 {
-	struct ath11k_pci *ab_pci = ath11k_pci_priv(irq_grp->ab);
+	struct ath11k_base *ab = irq_grp->ab;
 	int i;
 
 	/* In case of one MSI vector, we handle irq enable/disable
 	 * in a uniform way since we only have one irq
 	 */
-	if (!test_bit(ATH11K_PCI_FLAG_MULTI_MSI_VECTORS, &ab_pci->flags))
+	if (!test_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags))
 		return;
 
 	for (i = 0; i < irq_grp->num_irq; i++)
@@ -463,13 +418,13 @@ static void __ath11k_pcic_ext_irq_disable(struct ath11k_base *sc)
 
 static void ath11k_pcic_ext_grp_enable(struct ath11k_ext_irq_grp *irq_grp)
 {
-	struct ath11k_pci *ab_pci = ath11k_pci_priv(irq_grp->ab);
+	struct ath11k_base *ab = irq_grp->ab;
 	int i;
 
 	/* In case of one MSI vector, we handle irq enable/disable in a
 	 * uniform way since we only have one irq
 	 */
-	if (!test_bit(ATH11K_PCI_FLAG_MULTI_MSI_VECTORS, &ab_pci->flags))
+	if (!test_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags))
 		return;
 
 	for (i = 0; i < irq_grp->num_irq; i++)
@@ -557,11 +512,22 @@ static irqreturn_t ath11k_pcic_ext_interrupt_handler(int irq, void *arg)
 	return IRQ_HANDLED;
 }
 
+static int
+ath11k_pcic_get_msi_irq(struct ath11k_base *ab, unsigned int vector)
+{
+	if (!ab->pci.ops->get_msi_irq) {
+		WARN_ONCE(1, "get_msi_irq pci op not defined");
+		return -EOPNOTSUPP;
+	}
+
+	return ab->pci.ops->get_msi_irq(ab, vector);
+}
+
 static int ath11k_pcic_ext_irq_config(struct ath11k_base *ab)
 {
-	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
 	int i, j, ret, num_vectors = 0;
 	u32 user_base_data = 0, base_vector = 0;
+	unsigned long irq_flags;
 
 	ret = ath11k_pcic_get_user_msi_assignment(ab, "DP", &num_vectors,
 						  &user_base_data,
@@ -569,6 +535,10 @@ static int ath11k_pcic_ext_irq_config(struct ath11k_base *ab)
 	if (ret < 0)
 		return ret;
 
+	irq_flags = IRQF_SHARED;
+	if (!test_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags))
+		irq_flags |= IRQF_NOBALANCING;
+
 	for (i = 0; i < ATH11K_EXT_IRQ_GRP_NUM_MAX; i++) {
 		struct ath11k_ext_irq_grp *irq_grp = &ab->ext_irq_grp[i];
 		u32 num_irq = 0;
@@ -596,7 +566,10 @@ static int ath11k_pcic_ext_irq_config(struct ath11k_base *ab)
 		for (j = 0; j < irq_grp->num_irq; j++) {
 			int irq_idx = irq_grp->irqs[j];
 			int vector = (i % num_vectors) + base_vector;
-			int irq = ath11k_pcic_get_msi_irq(ab->dev, vector);
+			int irq = ath11k_pcic_get_msi_irq(ab, vector);
+
+			if (irq < 0)
+				return irq;
 
 			ab->irq_num[irq_idx] = irq;
 
@@ -605,8 +578,7 @@ static int ath11k_pcic_ext_irq_config(struct ath11k_base *ab)
 
 			irq_set_status_flags(irq, IRQ_DISABLE_UNLAZY);
 			ret = request_irq(irq, ath11k_pcic_ext_interrupt_handler,
-					  ab_pci->irq_flags,
-					  "DP_EXT_IRQ", irq_grp);
+					  irq_flags, "DP_EXT_IRQ", irq_grp);
 			if (ret) {
 				ath11k_err(ab, "failed request irq %d: %d\n",
 					   vector, ret);
@@ -619,35 +591,24 @@ static int ath11k_pcic_ext_irq_config(struct ath11k_base *ab)
 	return 0;
 }
 
-int ath11k_pcic_set_irq_affinity_hint(struct ath11k_pci *ab_pci,
-				      const struct cpumask *m)
-{
-	if (test_bit(ATH11K_PCI_FLAG_MULTI_MSI_VECTORS, &ab_pci->flags))
-		return 0;
-
-	return irq_set_affinity_hint(ab_pci->pdev->irq, m);
-}
-
 int ath11k_pcic_config_irq(struct ath11k_base *ab)
 {
-	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
 	struct ath11k_ce_pipe *ce_pipe;
 	u32 msi_data_start;
 	u32 msi_data_count, msi_data_idx;
 	u32 msi_irq_start;
 	unsigned int msi_data;
 	int irq, i, ret, irq_idx;
+	unsigned long irq_flags;
 
 	ret = ath11k_pcic_get_user_msi_assignment(ab, "CE", &msi_data_count,
 						  &msi_data_start, &msi_irq_start);
 	if (ret)
 		return ret;
 
-	ret = ath11k_pcic_set_irq_affinity_hint(ab_pci, cpumask_of(0));
-	if (ret) {
-		ath11k_err(ab, "failed to set irq affinity %d\n", ret);
-		return ret;
-	}
+	irq_flags = IRQF_SHARED;
+	if (!test_bit(ATH11K_FLAG_MULTI_MSI_VECTORS, &ab->dev_flags))
+		irq_flags |= IRQF_NOBALANCING;
 
 	/* Configure CE irqs */
 	for (i = 0, msi_data_idx = 0; i < ab->hw_params.ce_count; i++) {
@@ -655,7 +616,10 @@ int ath11k_pcic_config_irq(struct ath11k_base *ab)
 			continue;
 
 		msi_data = (msi_data_idx % msi_data_count) + msi_irq_start;
-		irq = ath11k_pcic_get_msi_irq(ab->dev, msi_data);
+		irq = ath11k_pcic_get_msi_irq(ab, msi_data);
+		if (irq < 0)
+			return irq;
+
 		ce_pipe = &ab->ce.ce_pipe[i];
 
 		irq_idx = ATH11K_PCI_IRQ_CE0_OFFSET + i;
@@ -663,12 +627,11 @@ int ath11k_pcic_config_irq(struct ath11k_base *ab)
 		tasklet_setup(&ce_pipe->intr_tq, ath11k_pcic_ce_tasklet);
 
 		ret = request_irq(irq, ath11k_pcic_ce_interrupt_handler,
-				  ab_pci->irq_flags, irq_name[irq_idx],
-				  ce_pipe);
+				  irq_flags, irq_name[irq_idx], ce_pipe);
 		if (ret) {
 			ath11k_err(ab, "failed to request irq %d: %d\n",
 				   irq_idx, ret);
-			goto err_irq_affinity_cleanup;
+			return ret;
 		}
 
 		ab->irq_num[irq_idx] = irq;
@@ -679,13 +642,9 @@ int ath11k_pcic_config_irq(struct ath11k_base *ab)
 
 	ret = ath11k_pcic_ext_irq_config(ab);
 	if (ret)
-		goto err_irq_affinity_cleanup;
+		return ret;
 
 	return 0;
-
-err_irq_affinity_cleanup:
-	ath11k_pcic_set_irq_affinity_hint(ab_pci, NULL);
-	return ret;
 }
 
 void ath11k_pcic_ce_irqs_enable(struct ath11k_base *ab)
@@ -730,17 +689,7 @@ void ath11k_pcic_stop(struct ath11k_base *ab)
 
 int ath11k_pcic_start(struct ath11k_base *ab)
 {
-	struct ath11k_pci *ab_pci = ath11k_pci_priv(ab);
-
-	set_bit(ATH11K_PCI_FLAG_INIT_DONE, &ab_pci->flags);
-
-	/* TODO: for now don't restore ASPM in case of single MSI
-	 * vector as MHI register reading in M2 causes system hang.
-	 */
-	if (test_bit(ATH11K_PCI_FLAG_MULTI_MSI_VECTORS, &ab_pci->flags))
-		ath11k_pcic_aspm_restore(ab_pci);
-	else
-		ath11k_info(ab, "leaving PCI ASPM disabled to avoid MHI M2 problems\n");
+	set_bit(ATH11K_FLAG_DEVICE_INIT_DONE, &ab->dev_flags);
 
 	ath11k_pcic_ce_irqs_enable(ab);
 	ath11k_ce_rx_post_buf(ab);
diff --git a/drivers/net/wireless/ath/ath11k/pcic.h b/drivers/net/wireless/ath/ath11k/pcic.h
index 34e8c069d6e4..c53d86289a8e 100644
--- a/drivers/net/wireless/ath/ath11k/pcic.h
+++ b/drivers/net/wireless/ath/ath11k/pcic.h
@@ -8,7 +8,6 @@
 #define _ATH11K_PCI_CMN_H
 
 #include "core.h"
-#include "pci.h"
 
 #define ATH11K_PCI_IRQ_CE0_OFFSET	3
 #define ATH11K_PCI_IRQ_DP_OFFSET	14
@@ -28,7 +27,6 @@
 int ath11k_pcic_get_user_msi_assignment(struct ath11k_base *ab, char *user_name,
 					int *num_vectors, u32 *user_base_data,
 					u32 *base_vector);
-int ath11k_pcic_get_msi_irq(struct device *dev, unsigned int vector);
 void ath11k_pcic_write32(struct ath11k_base *ab, u32 offset, u32 value);
 u32 ath11k_pcic_read32(struct ath11k_base *ab, u32 offset);
 void ath11k_pcic_get_msi_address(struct ath11k_base *ab, u32 *msi_addr_lo,
@@ -44,8 +42,5 @@ int ath11k_pcic_map_service_to_pipe(struct ath11k_base *ab, u16 service_id,
 				    u8 *ul_pipe, u8 *dl_pipe);
 void ath11k_pcic_ce_irqs_enable(struct ath11k_base *ab);
 void ath11k_pcic_ce_irq_disable_sync(struct ath11k_base *ab);
-void ath11k_pcic_aspm_restore(struct ath11k_pci *ab_pci);
-int ath11k_pcic_set_irq_affinity_hint(struct ath11k_pci *ab_pci,
-				      const struct cpumask *m);
 int ath11k_pcic_init_msi_config(struct ath11k_base *ab);
 #endif
