ath11k: Update tx and rx chain count properly on drv_set_antenna

Set the number of tx and rx chains properly on drv_set_antenna().
This will ensure the related ht/vht/he caps are properly recalculated
based on the tx/rx chains set.
Signed-off-by: Sriram R <srirrama@codeaurora.org>
Signed-off-by: Kalle Valo <kvalo@codeaurora.org>

drivers/net/wireless/ath/ath11k/mac.c

@@ -3480,6 +3480,8 @@ static int __ath11k_set_antenna(struct ath11k *ar, u32 tx_ant, u32 rx_ant)
 		return ret;
 	}
 
+	ar->num_tx_chains = get_num_chains(tx_ant);
+
 	ret = ath11k_wmi_pdev_set_param(ar, WMI_PDEV_PARAM_RX_CHAIN_MASK,
 					rx_ant, ar->pdev->pdev_id);
 	if (ret) {
@@ -3488,6 +3490,8 @@ static int __ath11k_set_antenna(struct ath11k *ar, u32 tx_ant, u32 rx_ant)
 		return ret;
 	}
 
+	ar->num_rx_chains = get_num_chains(rx_ant);
+
 	/* Reload HT/VHT/HE capability */
 	ath11k_mac_setup_ht_vht_cap(ar, &ar->pdev->cap, NULL);
 	ath11k_mac_setup_he_cap(ar, &ar->pdev->cap);