wifi: Handle nss calculations and checks for MU-MIMO

src/wifi/model/wifi-phy.cc

@@ -1730,7 +1730,25 @@ WifiPhy::Send(WifiConstPsduMap psdus, const WifiTxVector& txVector)
         NS_FATAL_ERROR("TX-VECTOR is invalid!");
     }
 
-    if (txVector.GetNssMax() > GetMaxSupportedTxSpatialStreams())
+    uint8_t nss = 0;
+    if (txVector.IsMu())
+    {
+        // We do not support mixed OFDMA and MU-MIMO
+        if (txVector.IsDlMuMimo())
+        {
+            nss = txVector.GetNssTotal();
+        }
+        else
+        {
+            nss = txVector.GetNssMax();
+        }
+    }
+    else
+    {
+        nss = txVector.GetNss();
+    }
+
+    if (nss > GetMaxSupportedTxSpatialStreams())
     {
         NS_FATAL_ERROR("Unsupported number of spatial streams!");
     }
@@ -2151,7 +2169,7 @@ WifiPhy::GetTxPowerForTransmission(Ptr<const WifiPpdu> ppdu) const
     }
     else
     {
-        if (txVector.GetNssMax() > 1)
+        if (txVector.GetNssMax() > 1 || txVector.GetNssTotal() > 1)
         {
             txPowerDbm = std::min(m_txPowerMaxMimo, GetPowerDbm(txVector.GetTxPowerLevel()));
         }

src/wifi/model/wifi-tx-vector.cc

@@ -27,6 +27,7 @@
 
 #include <algorithm>
 #include <iterator>
+#include <numeric>
 #include <set>
 
 namespace ns3
@@ -218,6 +219,7 @@ WifiTxVector::GetNss(uint16_t staId) const
 uint8_t
 WifiTxVector::GetNssMax() const
 {
+    // We do not support mixed OFDMA and MU-MIMO
     uint8_t nss = 0;
     if (IsMu())
     {
@@ -233,6 +235,26 @@ WifiTxVector::GetNssMax() const
     return nss;
 }
 
+uint8_t
+WifiTxVector::GetNssTotal() const
+{
+    // We do not support mixed OFDMA and MU-MIMO
+    uint8_t nss = 0;
+    if (IsMu())
+    {
+        nss = std::accumulate(
+            m_muUserInfos.cbegin(),
+            m_muUserInfos.cend(),
+            0,
+            [](uint8_t prevNss, const auto& info) { return prevNss + info.second.nss; });
+    }
+    else
+    {
+        nss = m_nss;
+    }
+    return nss;
+}
+
 uint8_t
 WifiTxVector::GetNess() const
 {

src/wifi/model/wifi-tx-vector.h

@@ -249,9 +249,13 @@ class WifiTxVector
      */
     uint8_t GetNss(uint16_t staId = SU_STA_ID) const;
     /**
-     * \returns the maximum number of Nss (namely if MU)
+     * \returns the maximum number of Nss over all RUs of an HE MU (used for OFDMA)
      */
     uint8_t GetNssMax() const;
+    /**
+     * \returns the total number of Nss for a given RU of an HE MU (used for full bandwidth MU-MIMO)
+     */
+    uint8_t GetNssTotal() const;
     /**
      * Sets the number of Nss
      *