From e93433a00ee18a27cfb9ef6a9cd0df6c28f96cb7 Mon Sep 17 00:00:00 2001
From: Peter Chiu <chui-hao.chiu@mediatek.com>
Date: Wed, 21 Sep 2022 13:25:26 +0800
Subject: [PATCH] mt76: mt7915: set chainmask to 3 for mt7916

Use chainmask to represent the number of TRx path and set it to 3 for mt7916.
Use antenna_mask to represent the TRx streams.

Signed-off-by: Peter Chiu <chui-hao.chiu@mediatek.com>
---
 mt7915/init.c | 11 ++++++-----
 mt7915/main.c |  6 +++++-
 mt7915/mcu.c  |  2 +-
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/mt7915/init.c b/mt7915/init.c
index 489f6e77..8474c11f 100644
--- a/mt7915/init.c
+++ b/mt7915/init.c
@@ -730,11 +730,12 @@ void mt7915_set_stream_vht_txbf_caps(struct mt7915_phy *phy)
 }
 
 static void
-mt7915_set_stream_he_txbf_caps(struct mt7915_dev *dev,
-			       struct ieee80211_sta_he_cap *he_cap,
-			       int vif, int nss)
+mt7915_set_stream_he_txbf_caps(struct mt7915_phy *phy,
+			       struct ieee80211_sta_he_cap *he_cap, int vif)
 {
+	struct mt7915_dev *dev = phy->dev;
 	struct ieee80211_he_cap_elem *elem = &he_cap->he_cap_elem;
+	int nss = hweight8(phy->mt76->chainmask);
 	u8 c, nss_160;
 
 	/* Can do 1/2 of NSS streams in 160Mhz mode for mt7915 */
@@ -839,7 +840,7 @@ mt7915_init_he_caps(struct mt7915_phy *phy, enum nl80211_band band,
 		    struct ieee80211_sband_iftype_data *data)
 {
 	struct mt7915_dev *dev = phy->dev;
-	int i, idx = 0, nss = hweight8(phy->mt76->chainmask);
+	int i, idx = 0, nss = hweight8(phy->mt76->antenna_mask);
 	u16 mcs_map = 0;
 	u16 mcs_map_160 = 0;
 	u8 nss_160;
@@ -972,7 +973,7 @@ mt7915_init_he_caps(struct mt7915_phy *phy, enum nl80211_band band,
 		he_mcs->rx_mcs_80p80 = cpu_to_le16(mcs_map_160);
 		he_mcs->tx_mcs_80p80 = cpu_to_le16(mcs_map_160);
 
-		mt7915_set_stream_he_txbf_caps(dev, he_cap, i, nss);
+		mt7915_set_stream_he_txbf_caps(phy, he_cap, i);
 
 		memset(he_cap->ppe_thres, 0, sizeof(he_cap->ppe_thres));
 		if (he_cap_elem->phy_cap_info[6] &
diff --git a/mt7915/main.c b/mt7915/main.c
index edd678ce..fc5319da 100644
--- a/mt7915/main.c
+++ b/mt7915/main.c
@@ -952,6 +952,7 @@ mt7915_set_antenna(struct ieee80211_hw *hw, u32 tx_ant, u32 rx_ant)
 	struct mt7915_dev *dev = mt7915_hw_dev(hw);
 	struct mt7915_phy *phy = mt7915_hw_phy(hw);
 	int max_nss = hweight8(hw->wiphy->available_antennas_tx);
+	int tx_nss = hweight8(tx_ant), rx_nss = hweight8(rx_ant);
 	bool ext_phy = phy != &dev->phy;
 
 	if (!tx_ant || tx_ant != rx_ant || ffs(tx_ant) > max_nss)
@@ -964,7 +965,10 @@ mt7915_set_antenna(struct ieee80211_hw *hw, u32 tx_ant, u32 rx_ant)
 	if (ext_phy)
 		tx_ant <<= dev->chainshift;
 
-	phy->mt76->chainmask = tx_ant;
+	if (is_mt7916(&dev->mt76) && tx_nss == max_nss && rx_nss == max_nss)
+		phy->mt76->chainmask = 0x7;
+	else
+		phy->mt76->chainmask = tx_ant;
 
 	mt76_set_stream_caps(phy->mt76, true);
 	mt7915_set_stream_vht_txbf_caps(phy);
diff --git a/mt7915/mcu.c b/mt7915/mcu.c
index 19a33120..f11ca25d 100644
--- a/mt7915/mcu.c
+++ b/mt7915/mcu.c
@@ -485,7 +485,7 @@ static void
 mt7915_mcu_bss_ra_tlv(struct sk_buff *skb, struct ieee80211_vif *vif,
 		      struct mt7915_phy *phy)
 {
-	int max_nss = hweight8(phy->mt76->chainmask);
+	int max_nss = hweight8(phy->mt76->antenna_mask);
 	struct bss_info_ra *ra;
 	struct tlv *tlv;
 
-- 
2.18.0