From 4b263c3fa567779e7bfccc4a783284fe767f2cc2 Mon Sep 17 00:00:00 2001
From: Roberto Louro Magueta <rmagueta@allbesmart.pt>
Date: Tue, 21 Feb 2023 19:37:07 +0000
Subject: [PATCH] MMSE working for 2-layers UL with a hardcoded noise variance

---
 .../PHY/NR_TRANSPORT/nr_ulsch_demodulation.c  | 70 ++++++++++++-------
 1 file changed, 43 insertions(+), 27 deletions(-)

diff --git a/openair1/PHY/NR_TRANSPORT/nr_ulsch_demodulation.c b/openair1/PHY/NR_TRANSPORT/nr_ulsch_demodulation.c
index f0a52b9c8e..e85ddd682e 100644
--- a/openair1/PHY/NR_TRANSPORT/nr_ulsch_demodulation.c
+++ b/openair1/PHY/NR_TRANSPORT/nr_ulsch_demodulation.c
@@ -1468,21 +1468,21 @@ void nr_ulsch_construct_HhH_elements(int *conjch00_ch00,
   _m_empty();
 }
 
-/* Zero Forcing Rx function: nr_ulsch_zero_forcing_rx_2layers()
- *
- *
- * */
-uint8_t nr_ulsch_zero_forcing_rx_2layers(NR_DL_FRAME_PARMS *frame_parms,
-                                         int **rxdataF_comp,
-                                         int **ul_ch_mag,
-                                         int **ul_ch_magb,
-                                         int **ul_ch_estimates_ext,
-                                         unsigned short nb_rb,
-                                         unsigned char n_rx,
-                                         unsigned char mod_order,
-                                         int shift,
-                                         unsigned char symbol,
-                                         int length)
+/*
+ * MMSE Rx function: nr_ulsch_mmse_2layers()
+ */
+uint8_t nr_ulsch_mmse_2layers(NR_DL_FRAME_PARMS *frame_parms,
+                              int **rxdataF_comp,
+                              int **ul_ch_mag,
+                              int **ul_ch_magb,
+                              int **ul_ch_estimates_ext,
+                              unsigned short nb_rb,
+                              unsigned char n_rx,
+                              unsigned char mod_order,
+                              int shift,
+                              unsigned char symbol,
+                              int length,
+                              int noise_var)
 {
   int *ch00, *ch01, *ch10, *ch11;
   int *ch20, *ch30, *ch21, *ch31;
@@ -1702,6 +1702,20 @@ uint8_t nr_ulsch_zero_forcing_rx_2layers(NR_DL_FRAME_PARMS *frame_parms,
                               nb_rb_0,
                               symbol);
   }
+
+  // Add noise_var such that: H^h * H + noise_var * I
+  if (noise_var != 0) {
+    __m128i nvar_128i = simde_mm_set1_epi32(noise_var);
+    __m128i *af_mf_00_128i = (__m128i *)af_mf_00;
+    __m128i *af_mf_11_128i = (__m128i *)af_mf_11;
+    for (int k = 0; k < 3 * nb_rb_0; k++) {
+      af_mf_00_128i[0] = simde_mm_add_epi32(af_mf_00_128i[0], nvar_128i);
+      af_mf_11_128i[0] = simde_mm_add_epi32(af_mf_11_128i[0], nvar_128i);
+      af_mf_00_128i++;
+      af_mf_11_128i++;
+    }
+  }
+
   //det_HhH = ad -bc
   nr_ulsch_det_HhH(af_mf_00,//a
              af_mf_01,//b
@@ -2039,19 +2053,21 @@ void nr_rx_pusch(PHY_VARS_gNB *gNB,
                              rel15_ul->rb_size,
                              nb_re_pusch);
 
-      //Apply zero forcing for 2 Tx layers
+      // Apply MMSE for 2 Tx layers
       if (rel15_ul->nrOfLayers == 2) {
-        nr_ulsch_zero_forcing_rx_2layers(frame_parms,
-                                         pusch_vars->rxdataF_comp,
-                                         pusch_vars->ul_ch_mag0,
-                                         pusch_vars->ul_ch_magb0,
-                                         pusch_vars->ul_ch_estimates_ext,
-                                         rel15_ul->rb_size,
-                                         frame_parms->nb_antennas_rx,
-                                         rel15_ul->qam_mod_order,
-                                         pusch_vars->log2_maxh,
-                                         symbol,
-                                         nb_re_pusch);
+        int noise_var = 100; // TODO: Get noise variance
+        nr_ulsch_mmse_2layers(frame_parms,
+                              pusch_vars->rxdataF_comp,
+                              pusch_vars->ul_ch_mag0,
+                              pusch_vars->ul_ch_magb0,
+                              pusch_vars->ul_ch_estimates_ext,
+                              rel15_ul->rb_size,
+                              frame_parms->nb_antennas_rx,
+                              rel15_ul->qam_mod_order,
+                              pusch_vars->log2_maxh,
+                              symbol,
+                              nb_re_pusch,
+                              noise_var);
       }
 
       stop_meas(&gNB->ulsch_mrc_stats);