[intra] Fix various issues with cclm, mip, dual-tree, and lfnst

src/encode_coding_tree.c

@@ -178,7 +178,7 @@ static bool can_use_lfnst_with_isp(const int width, const int height, const int
     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 || tree_type != UVG_BOTH_T;
-    bool mip_flag = pred_cu->type == CU_INTRA ? pred_cu->intra.mip_flag : false;
+    bool mip_flag = pred_cu->type == CU_INTRA && color == COLOR_Y ? pred_cu->intra.mip_flag : false;
 
     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) || 
@@ -1720,8 +1720,9 @@ 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, cu_width, cu_height, tree_type, COLOR_Y);
+    if (tree_type != UVG_CHROMA_T) {
-
+      bool lfnst_written = encode_lfnst_idx(state, cabac, cur_cu, x, y, depth, cu_width, cu_height, tree_type, COLOR_Y);
+    }
     encode_mts_idx(state, cabac, cur_cu);
 
     // For 4x4 the chroma PU/TU is coded after the last 
@@ -1736,9 +1737,8 @@ void uvg_encode_coding_tree(
       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 || tree_type == UVG_CHROMA_T || depth == 4) {
+      encode_lfnst_idx(state, cabac, tmp, x, y, depth, cu_width, cu_height, tree_type, COLOR_UV);
-        encode_lfnst_idx(state, cabac, tmp, x, y, depth, cu_width, cu_height, tree_type, COLOR_UV);
+      
-      }
     }
   }
 

src/encoderstate.c

@@ -1118,6 +1118,9 @@ static void encoder_state_encode(encoder_state_t * const main_state) {
         sub_state->tile->frame->rec = NULL;
 
         uvg_cu_array_free(&sub_state->tile->frame->cu_array);
+        if(sub_state->tile->frame->chroma_cu_array) {
+          uvg_cu_array_free(&sub_state->tile->frame->chroma_cu_array);
+        }
 
         sub_state->tile->frame->source = uvg_image_make_subimage(
             main_state->tile->frame->source,
@@ -1989,6 +1992,9 @@ void uvg_encoder_prepare(encoder_state_t *state)
 
   if (state->previous_encoder_state != state) {
     uvg_cu_array_free(&state->tile->frame->cu_array);
+    if (state->tile->frame->chroma_cu_array) {
+      uvg_cu_array_free(&state->tile->frame->chroma_cu_array);
+    }
     unsigned width  = state->tile->frame->width_in_lcu  * LCU_WIDTH;
     unsigned height = state->tile->frame->height_in_lcu * LCU_WIDTH;
     state->tile->frame->cu_array = uvg_cu_array_alloc(width, height);
@@ -2011,6 +2017,9 @@ void uvg_encoder_prepare(encoder_state_t *state)
                    prev_state->frame->poc,
                    prev_state->frame->ref_LX);
     uvg_cu_array_free(&state->tile->frame->cu_array);
+    if (state->tile->frame->chroma_cu_array) {
+      uvg_cu_array_free(&state->tile->frame->chroma_cu_array);
+    }
     unsigned height = state->tile->frame->height_in_lcu * LCU_WIDTH;
     unsigned width  = state->tile->frame->width_in_lcu  * LCU_WIDTH;
     state->tile->frame->cu_array = uvg_cu_array_alloc(width, height);
@@ -2032,6 +2041,9 @@ void uvg_encoder_prepare(encoder_state_t *state)
   state->tile->frame->rec = NULL;
 
   uvg_cu_array_free(&state->tile->frame->cu_array);
+  if (state->tile->frame->chroma_cu_array) {
+    uvg_cu_array_free(&state->tile->frame->chroma_cu_array);
+  }
 
   // Update POC and frame count.
   state->frame->num = prev_state->frame->num + 1;

src/search.c

@@ -833,8 +833,8 @@ static double search_cu(
     pu_depth_intra.max = ctrl->cfg.pu_depth_intra.max[gop_layer] >= 0 ? ctrl->cfg.pu_depth_intra.max[gop_layer] : ctrl->cfg.pu_depth_intra.max[0];
   }
   if(tree_type == UVG_CHROMA_T) {
-    pu_depth_intra.max = MIN(3, pu_depth_intra.max);
+    pu_depth_intra.max = CLIP(1, 3, pu_depth_intra.max);
-    pu_depth_intra.min = MIN(3, pu_depth_intra.min);
+    pu_depth_intra.min = CLIP(1, 3, pu_depth_intra.min);
   }
   pu_depth_inter.min = ctrl->cfg.pu_depth_inter.min[gop_layer] >= 0 ? ctrl->cfg.pu_depth_inter.min[gop_layer] : ctrl->cfg.pu_depth_inter.min[0];
   pu_depth_inter.max = ctrl->cfg.pu_depth_inter.max[gop_layer] >= 0 ? ctrl->cfg.pu_depth_inter.max[gop_layer] : ctrl->cfg.pu_depth_inter.max[0];
@@ -926,20 +926,21 @@ static double search_cu(
         intra_cost += pred_mode_type_bits * state->lambda;
       }
 #endif
+      if (state->encoder_control->cfg.cclm && tree_type != UVG_CHROMA_T && state->encoder_control->chroma_format != UVG_CSP_400) {
+        uvg_intra_recon_cu(state,
+          x, y,
+          depth, &intra_search,
+          &intra_search.pred_cu,
+          lcu, tree_type, true, false);
+
+        downsample_cclm_rec(
+          state, x, y, cu_width / 2, cu_width / 2, lcu->rec.y, lcu->left_ref.y[64]
+        );
+      }
       double intra_cost = intra_search.cost;
       if (intra_cost < cost && tree_type != UVG_LUMA_T) {
         int8_t intra_mode = intra_search.pred_cu.intra.mode;
-        if (state->encoder_control->cfg.cclm && tree_type == UVG_BOTH_T) {
-          uvg_intra_recon_cu(state,
-                             x, y,
-                             depth, &intra_search,
-                             &intra_search.pred_cu,
-                             lcu, tree_type, true, false);
 
-          downsample_cclm_rec(
-            state, x, y, cu_width / 2, cu_width / 2, lcu->rec.y, lcu->left_ref.y[64]
-          );
-        }
         // TODO: This heavily relies to square CUs
         if ((depth != 4 || (x % 8 && y % 8)) && state->encoder_control->chroma_format != UVG_CSP_400 && tree_type != UVG_LUMA_T) {
 

src/search_intra.c

@@ -336,9 +336,12 @@ static double search_intra_trdepth(
     
     const int max_tb_size = TR_MAX_WIDTH;
     // LFNST search params
-    const int max_lfnst_idx = width > max_tb_size || height > max_tb_size ?
+    int max_lfnst_idx = width > max_tb_size || height > max_tb_size ?
                                 0 :
                                 2;
+    if(pred_cu->intra.mip_flag && (width < 16 || height < 16)) {
+      max_lfnst_idx = 0;
+    }
 
     int start_idx = 0;
     int end_idx = state->encoder_control->cfg.lfnst && depth == pred_cu->
@@ -989,7 +992,7 @@ static double count_bits(
   return bits;
 }
 
-static int16_t search_intra_rough(
+static uint8_t search_intra_rough(
   encoder_state_t * const state,
   const cu_loc_t* const cu_loc,
   uvg_pixel *orig,
@@ -1041,7 +1044,7 @@ static int16_t search_intra_rough(
   const double planar_mode_flag = CTX_ENTROPY_FBITS(&(state->search_cabac.ctx.luma_planar_model[1]), 0);
   const double not_planar_mode_flag = CTX_ENTROPY_FBITS(&(state->search_cabac.ctx.luma_planar_model[1]), 1);
 
-  const int mode_list_size = state->encoder_control->cfg.mip ? 6 : 3;
+  const uint8_t mode_list_size = state->encoder_control->cfg.mip ? 6 : 3;
   struct mode_cost best_six_modes[6];
   // Initial offset decides how many modes are tried before moving on to the
   // recursive search.

src/transform.c

@@ -943,7 +943,7 @@ void uvg_inv_lfnst(
   bool is_separate_tree = depth == 4 || tree_type != UVG_BOTH_T;
   bool is_cclm_mode = (intra_mode >= 81 && intra_mode <= 83); // CCLM modes are in [81, 83]
 
-  bool is_mip = cur_cu->type == CU_INTRA ? cur_cu->intra.mip_flag : false;
+  bool is_mip = cur_cu->type == CU_INTRA && tree_type != UVG_CHROMA_T ? cur_cu->intra.mip_flag : false;
   bool is_wide_angle = false; // TODO: get wide angle mode when implemented
 
   const int cu_type = cur_cu->type;
@@ -956,7 +956,7 @@ void uvg_inv_lfnst(
     const uint32_t* scan = whge3 ? uvg_coef_top_left_diag_scan_8x8[log2_block_size] : uvg_g_sig_last_scan[scan_order][log2_block_size - 1];
     
     if (is_cclm_mode) {
-      intra_mode = cur_cu->intra.mode;
+      intra_mode = cur_cu->intra.mip_flag ? 0 : cur_cu->intra.mode;
     }
     if (is_mip) {
       intra_mode = 0; // Set to planar mode