.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,
                .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,
                .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,
                .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,
                .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,
                .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,
 
        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 {
        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
 
                        u32 addr_lo;
                        u32 addr_hi;
                } msi;
+
+               const struct ath11k_pci_ops *ops;
        } pci;
 
        /* must be last */
 
 /* 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
        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;
 
        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;
                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)) {
 
 
 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,
                                            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,
                        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);
        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
        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);
 }
 
        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,
                   *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)
 {
                        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:
                        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",
                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
        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);
 
        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);
 
 #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 {
        /* 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)
        return (struct ath11k_pci *)ab->drv_priv;
 }
 
+int ath11k_pci_get_msi_irq(struct ath11k_base *ab, unsigned int vector);
 #endif
 
  * Copyright (c) 2021-2022, Qualcomm Innovation Center, Inc. All rights reserved.
  */
 
-#include <linux/pci.h>
 #include "core.h"
 #include "pcic.h"
 #include "debug.h"
 }
 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)
 {
 
 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);
                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);
                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)
 {
 
 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;
 
 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;
 
 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++)
 
 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++)
        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,
        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;
                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;
 
 
                        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);
        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++) {
                        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;
                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;
 
        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)
 
 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);
 
 #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
 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,
                                    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