[lfnst] Fix hash mismatches for depth 4 chroma

src/cu.h

@@ -168,6 +168,7 @@ typedef struct
   uint8_t violates_lfnst_constrained_chroma : 1; // Two types, luma and chroma. Luma index is 0.
   uint8_t lfnst_last_scan_pos : 1;
   uint8_t lfnst_idx : 2;
+  uint8_t cr_lfnst_idx : 2;
 
   union {
     struct {

src/encode_coding_tree.c

@@ -243,7 +243,7 @@ static bool encode_lfnst_idx(
     // 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
 
-    const int lfnst_index = pred_cu->lfnst_idx;
+    const int lfnst_index = color == COLOR_Y ? pred_cu->lfnst_idx : pred_cu->cr_lfnst_idx;
     assert((lfnst_index >= 0 && lfnst_index < 3) && "Invalid LFNST index.");
 
     uint16_t ctx_idx = 0;

src/search_intra.c

@@ -1434,6 +1434,11 @@ int8_t uvg_search_intra_chroma_rdo(
     const int offset = ((lcu_px.x & ~7) >> 1) + ((lcu_px.y & ~7) >> 1)* LCU_WIDTH_C;
 
     int lfnst_modes_to_check[3];
+    if(depth == 4) {
+      for (int i = 0; i < 3; ++i) {
+        lfnst_modes_to_check[i] = i;
+      }
+    }
     if(chroma_data->pred_cu.lfnst_idx) {
       lfnst_modes_to_check[0] = chroma_data->pred_cu.lfnst_idx;
       lfnst_modes_to_check[1] = -1;
@@ -1461,6 +1466,7 @@ int8_t uvg_search_intra_chroma_rdo(
         if (lfnst == -1) {
           continue;
         }
+        pred_cu->cr_lfnst_idx = lfnst;
         chroma_data[mode_i].lfnst_costs[lfnst] += mode_bits * state->lambda;
         if (pred_cu->tr_depth == pred_cu->depth) {
           uvg_intra_predict(
@@ -1555,7 +1561,7 @@ int8_t uvg_search_intra_chroma_rdo(
           memcpy(&state->search_cabac, &temp_cabac, sizeof(cabac_data_t));
         }      
       }
-      pred_cu->lfnst_idx = best_lfnst_index;
+      pred_cu->cr_lfnst_idx = best_lfnst_index;
     }
     sort_modes(chroma_data, num_modes);
     
@@ -1603,7 +1609,9 @@ int8_t uvg_search_cu_intra_chroma(encoder_state_t * const state,
     chroma_data[i].pred_cu.intra.mode_chroma = num_modes == 1 ? intra_mode : modes[i];
     chroma_data[i].pred_cu.intra.mode = -1;
     chroma_data[i].cost = 0;
-    memcpy(chroma_data[i].lfnst_costs, search_data->lfnst_costs, sizeof(double) * 3);
+    if(depth != 4) {
+      memcpy(chroma_data[i].lfnst_costs, search_data->lfnst_costs, sizeof(double) * 3);
+    }
   }
   // Don't do rough mode search if all modes are selected.
   // FIXME: It might make more sense to only disable rough search if

src/strategies/generic/quant-generic.c

@@ -476,7 +476,7 @@ int uvg_quantize_residual_generic(encoder_state_t *const state,
     uvg_transform2d(state->encoder_control, residual, coeff, width, color, cur_cu);
   }
 
-  const uint16_t lfnst_index = cur_cu->lfnst_idx;
+  const uint16_t lfnst_index = color == COLOR_Y ? cur_cu->lfnst_idx : cur_cu->cr_lfnst_idx;
 
   if (state->encoder_control->cfg.lfnst && cur_cu->type == CU_INTRA) {
     // Forward low frequency non-separable transform

src/transform.c

@@ -501,7 +501,7 @@ void uvg_chroma_transform_search(
   const int can_use_tr_skip = state->encoder_control->cfg.trskip_enable &&
     (1 << state->encoder_control->cfg.trskip_max_size) >= width &&
     state->encoder_control->cfg.chroma_trskip_enable && 
-    pred_cu->lfnst_idx == 0;
+    pred_cu->cr_lfnst_idx == 0 ;
 
   if (can_use_tr_skip) {
     uvg_transformskip(state->encoder_control, u_resi, u_coeff + num_transforms * trans_offset, width);
@@ -552,16 +552,16 @@ void uvg_chroma_transform_search(
       scan_order,
       &u_has_coeffs,
       &v_has_coeffs,
-      pred_cu->lfnst_idx);
+      pred_cu->cr_lfnst_idx);
-    /*
+    
-    if(pred_cu->type == CU_INTRA && transforms[i] != CHROMA_TS) {
+    if(pred_cu->type == CU_INTRA && transforms[i] != CHROMA_TS && depth == 4) {
       bool constraints[2] = { false, false };
       uvg_derive_lfnst_constraints(pred_cu, depth, constraints, &u_coeff[i * trans_offset], width, height);
       if(!IS_JCCR_MODE(transforms[i])) {
         uvg_derive_lfnst_constraints(pred_cu, depth, constraints, &v_coeff[i * trans_offset], width, height);        
       }
-      if ((constraints[0] || !constraints[1]) && pred_cu->lfnst_idx != 0) continue;
+      if ((constraints[0] || !constraints[1]) && pred_cu->cr_lfnst_idx != 0) continue;
-    }*/
+    }
 
     if (IS_JCCR_MODE(transforms[i]) && !u_has_coeffs) continue;