struct work_struct update_wk;
        struct work_struct set_4addr_wk;
        struct rate_info txrate;
+       u32 peer_nss;
        struct rate_info last_txrate;
        u64 rx_duration;
        u64 tx_duration;
 
 }
 
 void ath11k_debugfs_sta_update_txcompl(struct ath11k *ar,
-                                      struct sk_buff *msdu,
                                       struct hal_tx_status *ts)
 {
-       struct ath11k_base *ab = ar->ab;
-       struct ath11k_per_peer_tx_stats *peer_stats = &ar->cached_stats;
-       enum hal_tx_rate_stats_pkt_type pkt_type;
-       enum hal_tx_rate_stats_sgi sgi;
-       enum hal_tx_rate_stats_bw bw;
-       struct ath11k_peer *peer;
-       struct ath11k_sta *arsta;
-       struct ieee80211_sta *sta;
-       u16 rate;
-       u8 rate_idx = 0;
-       int ret;
-       u8 mcs;
-
-       rcu_read_lock();
-       spin_lock_bh(&ab->base_lock);
-       peer = ath11k_peer_find_by_id(ab, ts->peer_id);
-       if (!peer || !peer->sta) {
-               ath11k_warn(ab, "failed to find the peer\n");
-               spin_unlock_bh(&ab->base_lock);
-               rcu_read_unlock();
-               return;
-       }
-
-       sta = peer->sta;
-       arsta = (struct ath11k_sta *)sta->drv_priv;
-
-       memset(&arsta->txrate, 0, sizeof(arsta->txrate));
-       pkt_type = FIELD_GET(HAL_TX_RATE_STATS_INFO0_PKT_TYPE,
-                            ts->rate_stats);
-       mcs = FIELD_GET(HAL_TX_RATE_STATS_INFO0_MCS,
-                       ts->rate_stats);
-       sgi = FIELD_GET(HAL_TX_RATE_STATS_INFO0_SGI,
-                       ts->rate_stats);
-       bw = FIELD_GET(HAL_TX_RATE_STATS_INFO0_BW, ts->rate_stats);
-
-       if (pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11A ||
-           pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11B) {
-               ret = ath11k_mac_hw_ratecode_to_legacy_rate(mcs,
-                                                           pkt_type,
-                                                           &rate_idx,
-                                                           &rate);
-               if (ret < 0)
-                       goto err_out;
-               arsta->txrate.legacy = rate;
-       } else if (pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11N) {
-               if (mcs > 7) {
-                       ath11k_warn(ab, "Invalid HT mcs index %d\n", mcs);
-                       goto err_out;
-               }
-
-               arsta->txrate.mcs = mcs + 8 * (arsta->last_txrate.nss - 1);
-               arsta->txrate.flags = RATE_INFO_FLAGS_MCS;
-               if (sgi)
-                       arsta->txrate.flags |= RATE_INFO_FLAGS_SHORT_GI;
-       } else if (pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11AC) {
-               if (mcs > 9) {
-                       ath11k_warn(ab, "Invalid VHT mcs index %d\n", mcs);
-                       goto err_out;
-               }
-
-               arsta->txrate.mcs = mcs;
-               arsta->txrate.flags = RATE_INFO_FLAGS_VHT_MCS;
-               if (sgi)
-                       arsta->txrate.flags |= RATE_INFO_FLAGS_SHORT_GI;
-       } else if (pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11AX) {
-               /* TODO */
-       }
-
-       arsta->txrate.nss = arsta->last_txrate.nss;
-       arsta->txrate.bw = ath11k_mac_bw_to_mac80211_bw(bw);
-
-       ath11k_debugfs_sta_add_tx_stats(arsta, peer_stats, rate_idx);
-
-err_out:
-       spin_unlock_bh(&ab->base_lock);
-       rcu_read_unlock();
+       ath11k_dp_tx_update_txcompl(ar, ts);
 }
 
 static ssize_t ath11k_dbg_sta_dump_tx_stats(struct file *file,
 
                                     struct ath11k_per_peer_tx_stats *peer_stats,
                                     u8 legacy_rate_idx);
 void ath11k_debugfs_sta_update_txcompl(struct ath11k *ar,
-                                      struct sk_buff *msdu,
                                       struct hal_tx_status *ts);
 
 #else /* CONFIG_ATH11K_DEBUGFS */
 }
 
 static inline void ath11k_debugfs_sta_update_txcompl(struct ath11k *ar,
-                                                    struct sk_buff *msdu,
                                                     struct hal_tx_status *ts)
 {
 }
 
        return 0;
 }
 
-static inline u32 ath11k_he_gi_to_nl80211_he_gi(u8 sgi)
-{
-       u32 ret = 0;
-
-       switch (sgi) {
-       case RX_MSDU_START_SGI_0_8_US:
-               ret = NL80211_RATE_INFO_HE_GI_0_8;
-               break;
-       case RX_MSDU_START_SGI_1_6_US:
-               ret = NL80211_RATE_INFO_HE_GI_1_6;
-               break;
-       case RX_MSDU_START_SGI_3_2_US:
-               ret = NL80211_RATE_INFO_HE_GI_3_2;
-               break;
-       }
-
-       return ret;
-}
-
 static void
 ath11k_update_per_peer_tx_stats(struct ath11k *ar,
                                struct htt_ppdu_stats *ppdu_stats, u8 user)
                arsta->txrate.mcs = mcs;
                arsta->txrate.flags = RATE_INFO_FLAGS_HE_MCS;
                arsta->txrate.he_dcm = dcm;
-               arsta->txrate.he_gi = ath11k_he_gi_to_nl80211_he_gi(sgi);
-               arsta->txrate.he_ru_alloc = ath11k_he_ru_tones_to_nl80211_he_ru_alloc(
-                                               (user_rate->ru_end -
+               arsta->txrate.he_gi = ath11k_mac_he_gi_to_nl80211_he_gi(sgi);
+               arsta->txrate.he_ru_alloc = ath11k_mac_phy_he_ru_to_nl80211_he_ru_alloc
+                                               ((user_rate->ru_end -
                                                 user_rate->ru_start) + 1);
                break;
        }
 
        arsta->txrate.nss = nss;
+
        arsta->txrate.bw = ath11k_mac_bw_to_mac80211_bw(bw);
        arsta->tx_duration += tx_duration;
        memcpy(&arsta->last_txrate, &arsta->txrate, sizeof(struct rate_info));
                }
                rx_status->encoding = RX_ENC_HE;
                rx_status->nss = nss;
-               rx_status->he_gi = ath11k_he_gi_to_nl80211_he_gi(sgi);
+               rx_status->he_gi = ath11k_mac_he_gi_to_nl80211_he_gi(sgi);
                rx_status->bw = ath11k_mac_bw_to_mac80211_bw(bw);
                break;
        }
 
        }
 }
 
+void ath11k_dp_tx_update_txcompl(struct ath11k *ar, struct hal_tx_status *ts)
+{
+       struct ath11k_base *ab = ar->ab;
+       struct ath11k_per_peer_tx_stats *peer_stats = &ar->cached_stats;
+       enum hal_tx_rate_stats_pkt_type pkt_type;
+       enum hal_tx_rate_stats_sgi sgi;
+       enum hal_tx_rate_stats_bw bw;
+       struct ath11k_peer *peer;
+       struct ath11k_sta *arsta;
+       struct ieee80211_sta *sta;
+       u16 rate, ru_tones;
+       u8 mcs, rate_idx, ofdma;
+       int ret;
+
+       spin_lock_bh(&ab->base_lock);
+       peer = ath11k_peer_find_by_id(ab, ts->peer_id);
+       if (!peer || !peer->sta) {
+               ath11k_dbg(ab, ATH11K_DBG_DP_TX,
+                          "failed to find the peer by id %u\n", ts->peer_id);
+               goto err_out;
+       }
+
+       sta = peer->sta;
+       arsta = (struct ath11k_sta *)sta->drv_priv;
+
+       memset(&arsta->txrate, 0, sizeof(arsta->txrate));
+       pkt_type = FIELD_GET(HAL_TX_RATE_STATS_INFO0_PKT_TYPE,
+                            ts->rate_stats);
+       mcs = FIELD_GET(HAL_TX_RATE_STATS_INFO0_MCS,
+                       ts->rate_stats);
+       sgi = FIELD_GET(HAL_TX_RATE_STATS_INFO0_SGI,
+                       ts->rate_stats);
+       bw = FIELD_GET(HAL_TX_RATE_STATS_INFO0_BW, ts->rate_stats);
+       ru_tones = FIELD_GET(HAL_TX_RATE_STATS_INFO0_TONES_IN_RU, ts->rate_stats);
+       ofdma = FIELD_GET(HAL_TX_RATE_STATS_INFO0_OFDMA_TX, ts->rate_stats);
+
+       /* This is to prefer choose the real NSS value arsta->last_txrate.nss,
+        * if it is invalid, then choose the NSS value while assoc.
+        */
+       if (arsta->last_txrate.nss)
+               arsta->txrate.nss = arsta->last_txrate.nss;
+       else
+               arsta->txrate.nss = arsta->peer_nss;
+
+       if (pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11A ||
+           pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11B) {
+               ret = ath11k_mac_hw_ratecode_to_legacy_rate(mcs,
+                                                           pkt_type,
+                                                           &rate_idx,
+                                                           &rate);
+               if (ret < 0)
+                       goto err_out;
+               arsta->txrate.legacy = rate;
+       } else if (pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11N) {
+               if (mcs > 7) {
+                       ath11k_warn(ab, "Invalid HT mcs index %d\n", mcs);
+                       goto err_out;
+               }
+
+               if (arsta->txrate.nss != 0)
+                       arsta->txrate.mcs = mcs + 8 * (arsta->txrate.nss - 1);
+               arsta->txrate.flags = RATE_INFO_FLAGS_MCS;
+               if (sgi)
+                       arsta->txrate.flags |= RATE_INFO_FLAGS_SHORT_GI;
+       } else if (pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11AC) {
+               if (mcs > 9) {
+                       ath11k_warn(ab, "Invalid VHT mcs index %d\n", mcs);
+                       goto err_out;
+               }
+
+               arsta->txrate.mcs = mcs;
+               arsta->txrate.flags = RATE_INFO_FLAGS_VHT_MCS;
+               if (sgi)
+                       arsta->txrate.flags |= RATE_INFO_FLAGS_SHORT_GI;
+       } else if (pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11AX) {
+               if (mcs > 11) {
+                       ath11k_warn(ab, "Invalid HE mcs index %d\n", mcs);
+                       goto err_out;
+               }
+
+               arsta->txrate.mcs = mcs;
+               arsta->txrate.flags = RATE_INFO_FLAGS_HE_MCS;
+               arsta->txrate.he_gi = ath11k_mac_he_gi_to_nl80211_he_gi(sgi);
+       }
+
+       arsta->txrate.bw = ath11k_mac_bw_to_mac80211_bw(bw);
+       if (ofdma && pkt_type == HAL_TX_RATE_STATS_PKT_TYPE_11AX) {
+               arsta->txrate.bw = RATE_INFO_BW_HE_RU;
+               arsta->txrate.he_ru_alloc =
+                       ath11k_mac_he_ru_tones_to_nl80211_he_ru_alloc(ru_tones);
+       }
+
+       if (ath11k_debugfs_is_extd_tx_stats_enabled(ar))
+               ath11k_debugfs_sta_add_tx_stats(arsta, peer_stats, rate_idx);
+
+err_out:
+       spin_unlock_bh(&ab->base_lock);
+}
+
 static void ath11k_dp_tx_complete_msdu(struct ath11k *ar,
                                       struct sk_buff *msdu,
                                       struct hal_tx_status *ts)
            (info->flags & IEEE80211_TX_CTL_NO_ACK))
                info->flags |= IEEE80211_TX_STAT_NOACK_TRANSMITTED;
 
-       if (unlikely(ath11k_debugfs_is_extd_tx_stats_enabled(ar))) {
+       if (unlikely(ath11k_debugfs_is_extd_tx_stats_enabled(ar)) ||
+           ab->hw_params.single_pdev_only) {
                if (ts->flags & HAL_TX_STATUS_FLAGS_FIRST_MSDU) {
                        if (ar->last_ppdu_id == 0) {
                                ar->last_ppdu_id = ts->ppdu_id;
                                   ar->cached_ppdu_id == ar->last_ppdu_id) {
                                ar->cached_ppdu_id = ar->last_ppdu_id;
                                ar->cached_stats.is_ampdu = true;
-                               ath11k_debugfs_sta_update_txcompl(ar, msdu, ts);
+                               ath11k_dp_tx_update_txcompl(ar, ts);
                                memset(&ar->cached_stats, 0,
                                       sizeof(struct ath11k_per_peer_tx_stats));
                        } else {
                                ar->cached_stats.is_ampdu = false;
-                               ath11k_debugfs_sta_update_txcompl(ar, msdu, ts);
+                               ath11k_dp_tx_update_txcompl(ar, ts);
                                memset(&ar->cached_stats, 0,
                                       sizeof(struct ath11k_per_peer_tx_stats));
                        }
 
        int ack_rssi;
 };
 
+void ath11k_dp_tx_update_txcompl(struct ath11k *ar, struct hal_tx_status *ts);
 int ath11k_dp_tx_htt_h2t_ver_req_msg(struct ath11k_base *ab);
 int ath11k_dp_tx(struct ath11k *ar, struct ath11k_vif *arvif,
                 struct ath11k_sta *arsta, struct sk_buff *skb);
 
 
                ru_tones = FIELD_GET(HAL_RX_HE_SIG_B1_MU_INFO_INFO0_RU_ALLOCATION,
                                     info0);
-               ppdu_info->ru_alloc = ath11k_he_ru_tones_to_nl80211_he_ru_alloc(ru_tones);
+               ppdu_info->ru_alloc =
+                       ath11k_mac_phy_he_ru_to_nl80211_he_ru_alloc(ru_tones);
+
                ppdu_info->reception_type = HAL_RX_RECEPTION_TYPE_MU_MIMO;
                break;
        }
 
                               struct hal_rx_mon_ppdu_info *ppdu_info,
                               struct sk_buff *skb);
 
-static inline u32 ath11k_he_ru_tones_to_nl80211_he_ru_alloc(u16 ru_tones)
-{
-       u32 ret = 0;
-
-       switch (ru_tones) {
-       case RU_26:
-               ret = NL80211_RATE_INFO_HE_RU_ALLOC_26;
-               break;
-       case RU_52:
-               ret = NL80211_RATE_INFO_HE_RU_ALLOC_52;
-               break;
-       case RU_106:
-               ret = NL80211_RATE_INFO_HE_RU_ALLOC_106;
-               break;
-       case RU_242:
-               ret = NL80211_RATE_INFO_HE_RU_ALLOC_242;
-               break;
-       case RU_484:
-               ret = NL80211_RATE_INFO_HE_RU_ALLOC_484;
-               break;
-       case RU_996:
-               ret = NL80211_RATE_INFO_HE_RU_ALLOC_996;
-               break;
-       }
-       return ret;
-}
-
 #define REO_QUEUE_DESC_MAGIC_DEBUG_PATTERN_0 0xDDBEEF
 #define REO_QUEUE_DESC_MAGIC_DEBUG_PATTERN_1 0xADBEEF
 #define REO_QUEUE_DESC_MAGIC_DEBUG_PATTERN_2 0xBDBEEF
 
 static int ath11k_start_vdev_delay(struct ieee80211_hw *hw,
                                   struct ieee80211_vif *vif);
 
+enum nl80211_he_ru_alloc ath11k_mac_phy_he_ru_to_nl80211_he_ru_alloc(u16 ru_phy)
+{
+       enum nl80211_he_ru_alloc ret;
+
+       switch (ru_phy) {
+       case RU_26:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_26;
+               break;
+       case RU_52:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_52;
+               break;
+       case RU_106:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_106;
+               break;
+       case RU_242:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_242;
+               break;
+       case RU_484:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_484;
+               break;
+       case RU_996:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_996;
+               break;
+       default:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_26;
+               break;
+       }
+
+       return ret;
+}
+
+enum nl80211_he_ru_alloc ath11k_mac_he_ru_tones_to_nl80211_he_ru_alloc(u16 ru_tones)
+{
+       enum nl80211_he_ru_alloc ret;
+
+       switch (ru_tones) {
+       case 26:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_26;
+               break;
+       case 52:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_52;
+               break;
+       case 106:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_106;
+               break;
+       case 242:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_242;
+               break;
+       case 484:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_484;
+               break;
+       case 996:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_996;
+               break;
+       case (996 * 2):
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_2x996;
+               break;
+       default:
+               ret = NL80211_RATE_INFO_HE_RU_ALLOC_26;
+               break;
+       }
+
+       return ret;
+}
+
+enum nl80211_he_gi ath11k_mac_he_gi_to_nl80211_he_gi(u8 sgi)
+{
+       enum nl80211_he_gi ret;
+
+       switch (sgi) {
+       case RX_MSDU_START_SGI_0_8_US:
+               ret = NL80211_RATE_INFO_HE_GI_0_8;
+               break;
+       case RX_MSDU_START_SGI_1_6_US:
+               ret = NL80211_RATE_INFO_HE_GI_1_6;
+               break;
+       case RX_MSDU_START_SGI_3_2_US:
+               ret = NL80211_RATE_INFO_HE_GI_3_2;
+               break;
+       default:
+               ret = NL80211_RATE_INFO_HE_GI_0_8;
+               break;
+       }
+
+       return ret;
+}
+
 u8 ath11k_mac_bw_to_mac80211_bw(u8 bw)
 {
        u8 ret = 0;
                                      struct peer_assoc_params *arg,
                                      bool reassoc)
 {
+       struct ath11k_sta *arsta;
+
        lockdep_assert_held(&ar->conf_mutex);
 
+       arsta = (struct ath11k_sta *)sta->drv_priv;
+
        memset(arg, 0, sizeof(*arg));
 
        reinit_completion(&ar->peer_assoc_done);
        ath11k_peer_assoc_h_qos(ar, vif, sta, arg);
        ath11k_peer_assoc_h_smps(sta, arg);
 
+       arsta->peer_nss = arg->peer_nss;
+
        /* TODO: amsdu_disable req? */
 }
 
 
 void ath11k_mac_peer_cleanup_all(struct ath11k *ar);
 int ath11k_mac_tx_mgmt_pending_free(int buf_id, void *skb, void *ctx);
 u8 ath11k_mac_bw_to_mac80211_bw(u8 bw);
+u32 ath11k_mac_he_gi_to_nl80211_he_gi(u8 sgi);
+enum nl80211_he_ru_alloc ath11k_mac_phy_he_ru_to_nl80211_he_ru_alloc(u16 ru_phy);
+enum nl80211_he_ru_alloc ath11k_mac_he_ru_tones_to_nl80211_he_ru_alloc(u16 ru_tones);
 enum ath11k_supported_bw ath11k_mac_mac80211_bw_to_ath11k_bw(enum rate_info_bw bw);
 enum hal_encrypt_type ath11k_dp_tx_get_encrypt_type(u32 cipher);
 void ath11k_mac_handle_beacon(struct ath11k *ar, struct sk_buff *skb);