return ath12k_mac_get_ar_by_chan(hw, ctx->def.chan);
 }
 
+static struct ath12k *ath12k_get_ar_by_vif(struct ieee80211_hw *hw,
+                                          struct ieee80211_vif *vif)
+{
+       struct ath12k_vif *arvif = ath12k_vif_to_arvif(vif);
+       struct ath12k_hw *ah = ath12k_hw_to_ah(hw);
+
+       /* If there is one pdev within ah, then we return
+        * ar directly.
+        */
+       if (ah->num_radio == 1)
+               return ah->radio;
+
+       if (arvif->is_created)
+               return arvif->ar;
+
+       return NULL;
+}
+
 static void ath12k_pdev_caps_update(struct ath12k *ar)
 {
        struct ath12k_base *ab = ar->ab;
                                           struct ieee80211_bss_conf *info,
                                           u64 changed)
 {
-       struct ath12k_hw *ah = ath12k_hw_to_ah(hw);
        struct ath12k *ar;
        struct ath12k_vif *arvif = ath12k_vif_to_arvif(vif);
 
-       ar = ath12k_ah_to_ar(ah, 0);
+       ar = ath12k_get_ar_by_vif(hw, vif);
+
+       /* TODO if the vdev is not created on a certain radio,
+        * cache the info to be updated later on vdev creation
+        */
+
+       if (!ar)
+               return;
 
        mutex_lock(&ar->conf_mutex);
 
                                 struct ieee80211_vif *vif, struct ieee80211_sta *sta,
                                 struct ieee80211_key_conf *key)
 {
-       struct ath12k_hw *ah = ath12k_hw_to_ah(hw);
        struct ath12k *ar;
        struct ath12k_base *ab;
        struct ath12k_vif *arvif = ath12k_vif_to_arvif(vif);
            key->cipher == WLAN_CIPHER_SUITE_BIP_CMAC_256)
                return 1;
 
-       ar = ath12k_ah_to_ar(ah, 0);
+       ar = ath12k_get_ar_by_vif(hw, vif);
+       if (!ar) {
+               WARN_ON_ONCE(1);
+               return -EINVAL;
+       }
        ab = ar->ab;
 
        if (test_bit(ATH12K_FLAG_HW_CRYPTO_DISABLED, &ar->ab->dev_flags))
                                   enum ieee80211_sta_state old_state,
                                   enum ieee80211_sta_state new_state)
 {
-       struct ath12k_hw *ah = ath12k_hw_to_ah(hw);
        struct ath12k *ar;
        struct ath12k_vif *arvif = ath12k_vif_to_arvif(vif);
        struct ath12k_sta *arsta = ath12k_sta_to_arsta(sta);
             new_state == IEEE80211_STA_NOTEXIST))
                cancel_work_sync(&arsta->update_wk);
 
-       ar = ath12k_ah_to_ar(ah, 0);
+       ar = ath12k_get_ar_by_vif(hw, vif);
+       if (!ar) {
+               WARN_ON_ONCE(1);
+               return -EINVAL;
+       }
 
        mutex_lock(&ar->conf_mutex);
 
                                        struct ieee80211_sta *sta,
                                        u32 changed)
 {
-       struct ath12k_hw *ah = ath12k_hw_to_ah(hw);
        struct ath12k *ar;
        struct ath12k_sta *arsta = ath12k_sta_to_arsta(sta);
        struct ath12k_vif *arvif = ath12k_vif_to_arvif(vif);
        struct ath12k_peer *peer;
        u32 bw, smps;
 
-       ar = ath12k_ah_to_ar(ah, 0);
+       ar = ath12k_get_ar_by_vif(hw, vif);
+       if (!ar) {
+               WARN_ON_ONCE(1);
+               return;
+       }
 
        spin_lock_bh(&ar->ab->base_lock);
 
                                 unsigned int link_id, u16 ac,
                                 const struct ieee80211_tx_queue_params *params)
 {
-       struct ath12k_hw *ah = ath12k_hw_to_ah(hw);
        struct ath12k *ar;
        struct ath12k_vif *arvif = ath12k_vif_to_arvif(vif);
        int ret;
 
-       ar = ath12k_ah_to_ar(ah, 0);
+       ar = ath12k_get_ar_by_vif(hw, vif);
+       if (!ar) {
+               /* TODO cache the info and apply after vdev is created */
+               return -EINVAL;
+       }
 
        mutex_lock(&ar->conf_mutex);
        ret = ath12k_mac_conf_tx(arvif, link_id, ac, params);
        struct ath12k_vif *arvif = ath12k_vif_to_arvif(vif);
        int ret = -EINVAL;
 
+       ar = ath12k_get_ar_by_vif(hw, vif);
+       if (!ar)
+               return -EINVAL;
+
        ar = ath12k_ah_to_ar(ah, 0);
 
        mutex_lock(&ar->conf_mutex);
                                u32 queues, bool drop)
 {
        struct ath12k_hw *ah = ath12k_hw_to_ah(hw);
-       struct ath12k *ar = ath12k_ah_to_ar(ah, 0);
+       struct ath12k *ar;
+       int i;
 
        if (drop)
                return;
 
+       /* vif can be NULL when flush() is considered for hw */
+       if (!vif) {
+               for_each_ar(ah, ar, i)
+                       ath12k_mac_flush(ar);
+               return;
+       }
+
+       ar = ath12k_get_ar_by_vif(hw, vif);
+
+       if (!ar)
+               return;
+
        ath12k_mac_flush(ar);
 }