[lfnst] [dual-tree] LFNST should work with dual tree

src/encode_coding_tree.c

@@ -139,9 +139,9 @@ static int get_isp_split_dim(const int width, const int height, const int isp_sp
   return partition_size;
 }
 
-static bool can_use_lfnst_with_isp(const int width, const int height, const int isp_split_type, const int color)
+static bool can_use_lfnst_with_isp(const int width, const int height, const int isp_split_type, const enum uvg_tree_type tree_type)
 {
-  if (color != COLOR_Y) {
+  if (tree_type == UVG_CHROMA_T) {
     return false;
   }
   if (isp_split_type == NOT_INTRA_SUBPARTITIONS) {
@@ -161,11 +161,11 @@ static bool can_use_lfnst_with_isp(const int width, const int height, const int
  bool uvg_is_lfnst_allowed(
   const encoder_state_t* const state,
   const cu_info_t* const pred_cu,
-  const int color,
   const int width,
   const int height,
   const int x,
-  const int y) 
+  const int y,
+  enum uvg_tree_type tree_type) 
 {
   if (state->encoder_control->cfg.lfnst && pred_cu->type == CU_INTRA) {
     const int isp_mode = 0; // ISP_TODO: assign proper ISP mode when ISP is implemented
@@ -173,21 +173,20 @@ static bool can_use_lfnst_with_isp(const int width, const int height, const int
     const int depth = pred_cu->depth;
     const int chroma_width = width >> 1;
     const int chroma_height = height >> 1;
-    const int cu_width = color == COLOR_Y || depth == 4 ? width : chroma_width;
+    const int cu_width = tree_type != UVG_LUMA_T || depth == 4 ? width : chroma_width;
-    const int cu_height = color == COLOR_Y || depth == 4 ? height : chroma_height;
+    const int cu_height = tree_type != UVG_LUMA_T || depth == 4 ? height : chroma_height;
     bool can_use_lfnst_with_mip = (width >= 16 && height >= 16);
-    bool is_sep_tree = depth == 4; // TODO: if/when separate tree structure is implemented, add proper boolean here
+    bool is_sep_tree = depth == 4 || tree_type != UVG_BOTH_T;
     bool mip_flag = pred_cu->type == CU_INTRA ? pred_cu->intra.mip_flag : false;
 
-    if ((isp_mode && !can_use_lfnst_with_isp(width, height, isp_split_type, color)) ||
+    if ((isp_mode && !can_use_lfnst_with_isp(width, height, isp_split_type, tree_type)) ||
       (pred_cu->type == CU_INTRA && mip_flag && !can_use_lfnst_with_mip) || 
-      (is_sep_tree && color != COLOR_Y && MIN(cu_width, cu_height) < 4) || 
+      (is_sep_tree && MIN(cu_width, cu_height) < 4) || 
       (cu_width > TR_MAX_WIDTH || cu_height > TR_MAX_WIDTH)) {
       return false;
     }
-    bool is_separate_tree = depth == 4; // TODO: if/when separate/dual tree structure is implemented, get proper value for this
+    bool luma_flag = (depth == 4 && tree_type == UVG_BOTH_T) || tree_type == UVG_LUMA_T;
-    bool luma_flag = is_separate_tree ? (color == COLOR_Y ? true : false) : true;
+    bool chroma_flag = (depth == 4 && tree_type == UVG_BOTH_T) || tree_type == UVG_CHROMA_T;
-    bool chroma_flag = is_separate_tree ? (color != COLOR_Y ? true : false) : true;
     bool non_zero_coeff_non_ts_corner_8x8 = (luma_flag && pred_cu->violates_lfnst_constrained_luma) || (chroma_flag && pred_cu->violates_lfnst_constrained_chroma);
     bool is_tr_skip = false;
 
@@ -201,7 +200,7 @@ static bool can_use_lfnst_with_isp(const int width, const int height, const int
     const int tu_height = tu_width; // TODO: height for non-square blocks
 
     // TODO: chroma transform skip
-    if (color == COLOR_Y) {
+    if (tree_type != UVG_BOTH_T) {
       for (int i = 0; i < num_transform_units; i++) {
         // TODO: this works only for square blocks
         const int pu_x = x + ((i % tu_row_length) * tu_width);
@@ -233,17 +232,17 @@ static bool encode_lfnst_idx(
   const int x,
   const int y,
   const int depth,
-  const int color,
   const int width,
-  const int height)
+  const int height,
+  enum uvg_tree_type tree_type)
 {
   
-  if (uvg_is_lfnst_allowed(state, pred_cu, color, width, height, x, y)) {
+  if (uvg_is_lfnst_allowed(state, pred_cu, width, height, x, y, tree_type)) {
     // Getting separate tree bool from block size is a temporary fix until a proper dual tree check is possible (there is no dual tree structure at time of writing this).
     // VTM seems to force explicit dual tree structure for small 4x4 blocks
-    bool is_separate_tree = depth == 4; // TODO: if/when separate/dual tree structure is implemented, get proper value for this
+    bool is_separate_tree = depth == 4 || tree_type != UVG_BOTH_T;
 
-    const int lfnst_index = color == COLOR_Y ? pred_cu->lfnst_idx : pred_cu->cr_lfnst_idx;
+    const int lfnst_index = tree_type != UVG_CHROMA_T ? pred_cu->lfnst_idx : pred_cu->cr_lfnst_idx;
     assert((lfnst_index >= 0 && lfnst_index < 3) && "Invalid LFNST index.");
 
     uint16_t ctx_idx = 0;
@@ -1718,8 +1717,7 @@ void uvg_encode_coding_tree(
       encode_transform_coeff(state, x, y, depth, 0, 0, 0, 0, coeff, tree_type);
     }
 
-    bool lfnst_written = encode_lfnst_idx(state, cabac, cur_cu, x, y, depth, COLOR_Y, cu_width, cu_height);
+    bool lfnst_written = encode_lfnst_idx(state, cabac, cur_cu, x, y, depth, cu_width, cu_height, tree_type);
-    bool is_separate_tree = depth == 4; // TODO: proper value for dual tree when dual tree structure is implemented
 
     encode_mts_idx(state, cabac, cur_cu);
 
@@ -1729,14 +1727,14 @@ void uvg_encode_coding_tree(
       tree_type != UVG_LUMA_T)   {
       encode_chroma_intra_cu(cabac, cur_cu, state->encoder_control->cfg.cclm, NULL);
       // LFNST constraints must be reset here. Otherwise the left over values will interfere when calculating new constraints
-      cu_info_t* tmp = uvg_cu_array_at(frame->cu_array, x, y);
+      cu_info_t* tmp = uvg_cu_array_at((cu_array_t*)used_array, x, y);
       tmp->violates_lfnst_constrained_luma = false;
       tmp->violates_lfnst_constrained_chroma = false;
       tmp->lfnst_last_scan_pos = false;
       encode_transform_coeff(state, x, y, depth, 0, 0, 0, 1, coeff, tree_type);
       // Write LFNST only once for single tree structure
-      if (!lfnst_written || is_separate_tree) {
+      if (!lfnst_written || tree_type == UVG_CHROMA_T) {
-        encode_lfnst_idx(state, cabac, tmp, x, y, depth, COLOR_UV, cu_width, cu_height);
+        encode_lfnst_idx(state, cabac, tmp, x, y, depth, cu_width, cu_height, tree_type);
       }
     }
   }

src/encode_coding_tree.h

@@ -44,11 +44,11 @@ bool uvg_is_mts_allowed(const encoder_state_t* const state, cu_info_t* const pre
 bool uvg_is_lfnst_allowed(
   const encoder_state_t* const state,
   const cu_info_t* const pred_cu,
-  const int color,
   const int width,
   const int height,
   const int x,
-  const int y);
+  const int y,
+  enum uvg_tree_type tree_type);
 
 void uvg_encode_coding_tree(
   encoder_state_t * const state,

src/search.c

@@ -566,12 +566,12 @@ static double cu_rd_cost_tr_split_accurate(
     coeff_bits += uvg_get_coeff_cost(state, coeffs, tr_cu, width, 0, luma_scan_mode, tr_cu->tr_skip & 1);
   }
 
-  if(depth == 4) {
+  if(depth == 4 || tree_type == UVG_LUMA_T) {
-    if (uvg_is_lfnst_allowed(state, tr_cu, COLOR_Y, width, width, x_px, y_px)) {
+    if (uvg_is_lfnst_allowed(state, tr_cu, width, width, x_px, y_px, tree_type)) {
       const int lfnst_idx = tr_cu->lfnst_idx;
       CABAC_FBITS_UPDATE(
         cabac,
-        &cabac->ctx.lfnst_idx_model[tr_cu->depth == 4],
+        &cabac->ctx.lfnst_idx_model[1],
         lfnst_idx != 0,
         tr_tree_bits,
         "lfnst_idx");
@@ -634,11 +634,11 @@ static double cu_rd_cost_tr_split_accurate(
     }
   }
 
-  if (uvg_is_lfnst_allowed(state, tr_cu, depth == 4 ? COLOR_UV : COLOR_Y, width, width, x_px, y_px)) {
+  if (uvg_is_lfnst_allowed(state, tr_cu, width, width, x_px, y_px, tree_type)) {
-    const int lfnst_idx = depth != 4 ? tr_cu->lfnst_idx : tr_cu->cr_lfnst_idx;
+    const int lfnst_idx = depth != 4 || tree_type != UVG_CHROMA_T ? tr_cu->lfnst_idx : tr_cu->cr_lfnst_idx;
     CABAC_FBITS_UPDATE(
       cabac,
-      &cabac->ctx.lfnst_idx_model[tr_cu->depth == 4],
+      &cabac->ctx.lfnst_idx_model[tr_cu->depth == 4 || tree_type != UVG_BOTH_T],
       lfnst_idx != 0,
       tr_tree_bits,
       "lfnst_idx");
@@ -913,7 +913,8 @@ static double search_cu(
       if(tree_type != UVG_CHROMA_T) {
         intra_search.pred_cu.joint_cb_cr = 4;
         uvg_search_cu_intra(state, x, y, depth, &intra_search,
-                            lcu);
+                            lcu,
+                            tree_type);
       }
 #ifdef COMPLETE_PRED_MODE_BITS
       // Technically counting these bits would be correct, however counting

src/search_intra.c

@@ -268,7 +268,8 @@ static double search_intra_trdepth(
   int max_depth,
   double cost_treshold,
   intra_search_data_t *const search_data,
-  lcu_t *const lcu)
+  lcu_t *const lcu,
+  enum uvg_tree_type tree_type)
 {
   assert(depth >= 0 && depth <= MAX_PU_DEPTH);
 
@@ -422,7 +423,7 @@ static double search_intra_trdepth(
         double transform_bits = 0;
         if(state->encoder_control->cfg.lfnst && depth == pred_cu->tr_depth) {
           if(!constraints[0] && constraints[1]) {
-            transform_bits += CTX_ENTROPY_FBITS(&state->search_cabac.ctx.lfnst_idx_model[tr_cu->depth == 4], lfnst_idx != 0);
+            transform_bits += CTX_ENTROPY_FBITS(&state->search_cabac.ctx.lfnst_idx_model[tr_cu->depth == 4 || tree_type == UVG_LUMA_T], lfnst_idx != 0);
             if(lfnst_idx > 0) {
               transform_bits += CTX_ENTROPY_FBITS(&state->search_cabac.ctx.lfnst_idx_model[2], lfnst_idx == 2);
             }
@@ -581,15 +582,15 @@ static double search_intra_trdepth(
   if (depth < max_depth && depth < MAX_PU_DEPTH) {
     split_cost = 0;
 
-    split_cost += search_intra_trdepth(state, x_px, y_px, depth + 1, max_depth, nosplit_cost, search_data, lcu);
+    split_cost += search_intra_trdepth(state, x_px, y_px, depth + 1, max_depth, nosplit_cost, search_data, lcu, tree_type);
     if (split_cost < nosplit_cost) {
-      split_cost += search_intra_trdepth(state, x_px + offset, y_px, depth + 1, max_depth, nosplit_cost, search_data, lcu);
+      split_cost += search_intra_trdepth(state, x_px + offset, y_px, depth + 1, max_depth, nosplit_cost, search_data, lcu, tree_type);
     }
     if (split_cost < nosplit_cost) {
-      split_cost += search_intra_trdepth(state, x_px, y_px + offset, depth + 1, max_depth, nosplit_cost, search_data, lcu);
+      split_cost += search_intra_trdepth(state, x_px, y_px + offset, depth + 1, max_depth, nosplit_cost, search_data, lcu, tree_type);
     }
     if (split_cost < nosplit_cost) {
-      split_cost += search_intra_trdepth(state, x_px + offset, y_px + offset, depth + 1, max_depth, nosplit_cost, search_data, lcu);
+      split_cost += search_intra_trdepth(state, x_px + offset, y_px + offset, depth + 1, max_depth, nosplit_cost, search_data, lcu, tree_type);
     }
 
     double cbf_bits = 0.0;
@@ -1329,7 +1330,8 @@ static int8_t search_intra_rdo(
   int depth,
   int modes_to_check,
   intra_search_data_t *search_data,
-  lcu_t *lcu)
+  lcu_t *lcu,
+  enum uvg_tree_type tree_type)
 {
   const int tr_depth = CLIP(1, MAX_PU_DEPTH, depth + state->encoder_control->cfg.tr_depth_intra);
   
@@ -1339,7 +1341,7 @@ static int8_t search_intra_rdo(
     search_data[mode].bits = rdo_bitcost;
     search_data[mode].cost = rdo_bitcost * state->lambda;
 
-    double mode_cost = search_intra_trdepth(state, x_px, y_px, depth, tr_depth, MAX_INT, &search_data[mode], lcu);
+    double mode_cost = search_intra_trdepth(state, x_px, y_px, depth, tr_depth, MAX_INT, &search_data[mode], lcu, tree_type);
     search_data[mode].cost += mode_cost;
     if (state->encoder_control->cfg.intra_rdo_et && !cbf_is_set_any(search_data[mode].pred_cu.cbf, depth)) {
       modes_to_check = mode + 1;
@@ -1761,7 +1763,8 @@ void uvg_search_cu_intra(
   const int y_px,
   const int depth,
   intra_search_data_t* mode_out,
-  lcu_t *lcu)
+  lcu_t *lcu,
+  enum uvg_tree_type tree_type)
 {
   const vector2d_t lcu_px = { SUB_SCU(x_px), SUB_SCU(y_px) };
   const int8_t cu_width = LCU_WIDTH >> depth;
@@ -1973,7 +1976,8 @@ void uvg_search_cu_intra(
       depth,
       number_of_modes_to_search,
       search_data,
-      lcu);
+      lcu,
+      tree_type);
     search_data[0].pred_cu.mts_last_scan_pos = false;
     search_data[0].pred_cu.violates_mts_coeff_constraint = false;
   }

src/search_intra.h

@@ -63,6 +63,7 @@ void uvg_search_cu_intra(
   const int y_px,
   const int depth,
   intra_search_data_t* search_data,
-  lcu_t *lcu);
+  lcu_t *lcu,
+  enum uvg_tree_type tree_type);
 
 #endif // SEARCH_INTRA_H_

src/transform.c

@@ -700,7 +700,7 @@ void uvg_chroma_transform_search(
         transforms[i] == CHROMA_TS);
     }
     if(depth == 4 && state->encoder_control->cfg.lfnst && 0) {
-      if(uvg_is_lfnst_allowed(state, pred_cu, COLOR_UV, width, height, 0 ,0)) {
+      if(uvg_is_lfnst_allowed(state, pred_cu, width, height, 0, 0 , UVG_CHROMA_T)) {
         const int lfnst_idx = pred_cu->cr_lfnst_idx;
         CABAC_FBITS_UPDATE(
           &state->search_cabac,