From ca2890d7e903aa85fa7bc6c2364ec292d7c5decd Mon Sep 17 00:00:00 2001
From: siivonek <kari.siivonen@tuni.fi>
Date: Mon, 11 Apr 2022 13:58:14 +0300
Subject: [PATCH] [lfnst] Fix error in inv lfnst call in quant generic. Add
 lfnst calls to quant avx2. Improve lfnst flag sps write.

---
 src/encoder_state-bitstream.c          |  8 +-------
 src/strategies/avx2/quant-avx2.c       | 14 ++++++++++++++
 src/strategies/generic/quant-generic.c | 10 +++++-----
 3 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/src/encoder_state-bitstream.c b/src/encoder_state-bitstream.c
index d425af92..2bc6d9af 100644
--- a/src/encoder_state-bitstream.c
+++ b/src/encoder_state-bitstream.c
@@ -582,13 +582,7 @@ static void encoder_state_write_bitstream_seq_parameter_set(bitstream_t* stream,
     WRITE_U(stream, mts_selection == UVG_MTS_INTER || mts_selection == UVG_MTS_BOTH ? 1 : 0, 1, "sps_explicit_mts_inter_enabled_flag");
   }
 
-  if (state->encoder_control->cfg.lfnst) {
-    WRITE_U(stream, 1, 1, "sps_lfnst_enabled_flag");
-  }
-  else {
-    WRITE_U(stream, 0, 1, "sps_lfnst_enabled_flag");
-  }
-
+  WRITE_U(stream, encoder->cfg.lfnst, 1, "sps_lfnst_enabled_flag");
 
   if (encoder->chroma_format != UVG_CSP_400) {
     WRITE_U(stream, encoder->cfg.jccr, 1, "sps_joint_cbcr_enabled_flag");
diff --git a/src/strategies/avx2/quant-avx2.c b/src/strategies/avx2/quant-avx2.c
index 95ce550f..921e4e5f 100644
--- a/src/strategies/avx2/quant-avx2.c
+++ b/src/strategies/avx2/quant-avx2.c
@@ -673,6 +673,7 @@ int uvg_quantize_residual_avx2(encoder_state_t *const state,
   ALIGNED(64) int16_t residual[TR_MAX_WIDTH * TR_MAX_WIDTH];
   ALIGNED(64) coeff_t coeff[TR_MAX_WIDTH * TR_MAX_WIDTH];
 
+  const int height = width; // TODO: height for non-square blocks
   int has_coeffs = 0;
 
   assert(width <= TR_MAX_WIDTH);
@@ -702,6 +703,14 @@ int uvg_quantize_residual_avx2(encoder_state_t *const state,
     uvg_transform2d(state->encoder_control, residual, coeff, width, color, cur_cu);
   }
 
+  // LFNST_TODO: lfnst index must be derived by performing search, like MTS index. Search is performed during intra search, inside search_tr_depth.
+  const uint16_t lfnst_index = 0;
+
+  if (state->encoder_control->cfg.lfnst) {
+    // Forward low frequency non-separable transform
+    kvz_fwd_lfnst(cur_cu, width, height, color, lfnst_index, coeff);
+  }
+
   // Quantize coeffs. (coeff -> coeff_out)
   
   if (state->encoder_control->cfg.rdoq_enable &&
@@ -732,6 +741,11 @@ int uvg_quantize_residual_avx2(encoder_state_t *const state,
   // rec_out.
   if (has_coeffs && !early_skip) {
 
+    if (state->encoder_control->cfg.lfnst) {
+      // Inverse low frequency non-separable transform
+      kvz_inv_lfnst(cur_cu, width, height, color, lfnst_index, coeff);
+    }
+
     // Get quantized residual. (coeff_out -> coeff -> residual)
     uvg_dequant(state, coeff_out, coeff, width, width, color,
       cur_cu->type, cur_cu->tr_idx == MTS_SKIP && color == COLOR_Y);
diff --git a/src/strategies/generic/quant-generic.c b/src/strategies/generic/quant-generic.c
index c6bacfb2..76fcd7c7 100644
--- a/src/strategies/generic/quant-generic.c
+++ b/src/strategies/generic/quant-generic.c
@@ -502,6 +502,11 @@ int uvg_quantize_residual_generic(encoder_state_t *const state,
   if (has_coeffs && !early_skip) {
     int y, x;
 
+    if (state->encoder_control->cfg.lfnst) {
+      // Inverse low frequency non-separable transform
+      kvz_inv_lfnst(cur_cu, width, height, color, lfnst_index, coeff);
+    }
+
     // Get quantized residual. (coeff_out -> coeff -> residual)
     uvg_dequant(state, coeff_out, coeff, width, width, color,
       cur_cu->type, cur_cu->tr_idx == MTS_SKIP && color == COLOR_Y);
@@ -511,11 +516,6 @@ int uvg_quantize_residual_generic(encoder_state_t *const state,
     else {
       uvg_itransform2d(state->encoder_control, residual, coeff, width, color, cur_cu);
     }
-
-    if (state->encoder_control->cfg.lfnst) {
-      // Inverse low frequency non-separable transform
-      kvz_inv_lfnst(cur_cu, width, height, color, lfnst_index, coeff);
-    }
     
     if (state->tile->frame->lmcs_aps->m_sliceReshapeInfo.enableChromaAdj && color != COLOR_Y) {
       int y, x;