[mtt] All individual mtt splits should be working + uvg_get_possible_splits

src/cu.c

@@ -34,6 +34,9 @@
 #include <stdlib.h>
 
 #include "cu.h"
+
+#include "alf.h"
+#include "encoderstate.h"
 #include "threads.h"
 
 
@@ -365,6 +368,116 @@ int uvg_get_split_locs(
   return 0;
 }
 
+
+int uvg_get_implicit_split(const encoder_state_t* const state, const cu_loc_t* const cu_loc)
+{
+  bool right_ok = state->tile->frame->width >= cu_loc->x + cu_loc->width;
+  bool bottom_ok = state->tile->frame->height >= cu_loc->y + cu_loc->height;
+
+  if (right_ok && bottom_ok) return NO_SPLIT;
+  if (right_ok) return BT_HOR_SPLIT;
+  if (bottom_ok) return BT_VER_SPLIT;
+  return QT_SPLIT;
+}
+
+
+int uvg_get_possible_splits(const encoder_state_t * const state,
+  const cu_loc_t * const cu_loc, split_tree_t split_tree, enum uvg_tree_type tree_type, bool splits[6])
+{
+  const int width = tree_type != UVG_CHROMA_T ? cu_loc->width : cu_loc->chroma_width;
+  const int height = tree_type != UVG_CHROMA_T ? cu_loc->height : cu_loc->chroma_height;
+  const enum split_type implicitSplit = uvg_get_implicit_split(state, cu_loc);
+  const int slice_type = state->frame->is_irap ? (tree_type == UVG_CHROMA_T ? 2 : 0) : 1;
+
+  const unsigned max_btd = state->encoder_control->cfg.max_btt_depth[slice_type]; // +currImplicitBtDepth;
+  const unsigned max_bt_size = state->encoder_control->cfg.max_bt_size[slice_type];
+  const unsigned min_bt_size = 1 << MIN_SIZE;
+  const unsigned max_tt_size = state->encoder_control->cfg.max_tt_size[slice_type];
+  const unsigned min_tt_size = 1 << MIN_SIZE;
+  const unsigned min_qt_size = state->encoder_control->cfg.min_qt_size[slice_type];
+  
+  splits[NO_SPLIT] = splits[QT_SPLIT] = splits[BT_HOR_SPLIT] = splits[TT_HOR_SPLIT] = splits[BT_VER_SPLIT] = splits[TT_VER_SPLIT] = true;
+  bool can_btt = split_tree.mtt_depth < max_btd;
+  
+  const enum split_type last_split = (split_tree.split_tree >> (split_tree.current_depth * 3)) & 7;
+  const enum split_type parl_split = last_split == BT_HOR_SPLIT ? BT_HOR_SPLIT : BT_VER_SPLIT;
+
+  // don't allow QT-splitting below a BT split
+  if (split_tree.current_depth != 0 && last_split != QT_SPLIT && (width > 64 || height > 64)) splits[QT_SPLIT] = false;
+  if (width <= min_qt_size)                              splits[QT_SPLIT] = false;
+
+  if (tree_type == UVG_CHROMA_T && width <= 4) splits[QT_SPLIT] = false;
+  if (tree_type == UVG_CHROMA_T)
+  {
+    splits[QT_SPLIT] = splits[BT_VER_SPLIT] = splits[TT_HOR_SPLIT] = splits[BT_VER_SPLIT] = splits[TT_VER_SPLIT] = false;
+    return;
+  }
+  if (implicitSplit != NO_SPLIT)
+  {
+    splits[NO_SPLIT] = splits[TT_HOR_SPLIT] = splits[TT_VER_SPLIT] = false;
+
+    splits[BT_HOR_SPLIT] = implicitSplit == BT_HOR_SPLIT;
+    splits[BT_VER_SPLIT] = implicitSplit == BT_VER_SPLIT;
+    if (tree_type == UVG_CHROMA_T && width == 4) splits[BT_VER_SPLIT] = false;
+    if (!splits[BT_HOR_SPLIT] && !splits[BT_VER_SPLIT] && !splits[QT_SPLIT]) splits[QT_SPLIT] = true;
+    return 1;
+  }
+
+  if ((last_split == TT_HOR_SPLIT || last_split == TT_VER_SPLIT) && split_tree.part_index == 1)
+  {
+    splits[BT_HOR_SPLIT] = parl_split != BT_HOR_SPLIT;
+    splits[BT_VER_SPLIT] = parl_split != BT_VER_SPLIT;
+  }
+
+  if (can_btt && (width <= min_bt_size && height <= min_bt_size)
+    && ((width <= min_tt_size && height <= min_tt_size)))
+  {
+    can_btt = false;
+  }
+  if (can_btt && (width > max_bt_size || height > max_bt_size)
+    && ((width > max_tt_size || height > max_tt_size)))
+  {
+    can_btt = false;
+  }
+
+  if (!can_btt)
+  {
+    splits[BT_HOR_SPLIT] = splits[TT_HOR_SPLIT] = splits[BT_VER_SPLIT] = splits[TT_VER_SPLIT] = false;
+
+    return;
+  }
+
+  if (width > max_bt_size || height > max_bt_size)
+  {
+    splits[BT_HOR_SPLIT] = splits[BT_VER_SPLIT] = false;
+  }
+
+  // specific check for BT splits
+  if (height <= min_bt_size)                            splits[BT_HOR_SPLIT] = false;
+  if (width > 64 && height <= 64) splits[BT_HOR_SPLIT] = false;
+  if (tree_type == UVG_CHROMA_T && width * height <= 16)     splits[BT_HOR_SPLIT] = false;
+
+  if (width <= min_bt_size)                              splits[BT_VER_SPLIT] = false;
+  if (width <= 64 && height > 64) splits[BT_VER_SPLIT] = false;
+  if (tree_type == UVG_CHROMA_T && (width * height <= 16 || width == 4))     splits[BT_VER_SPLIT] = false;
+
+  //if (modeType == MODE_TYPE_INTER && width * height == 32)  splits[BT_VER_SPLIT] = splits[BT_HOR_SPLIT] = false;
+
+  if (height <= 2 * min_tt_size || height > max_tt_size || width > max_tt_size)
+    splits[TT_HOR_SPLIT] = false;
+  if (width > 64 || height > 64)  splits[TT_HOR_SPLIT] = false;
+  if (tree_type == UVG_CHROMA_T && width * height <= 16 * 2)     splits[TT_HOR_SPLIT] = false;
+
+  if (width <= 2 * min_tt_size || width > max_tt_size || height > max_tt_size)
+    splits[TT_VER_SPLIT] = false;
+  if (width > 64 || height > 64)  splits[TT_VER_SPLIT] = false;
+  if (tree_type == UVG_CHROMA_T && (width * height <= 16 * 2 || width == 8))     splits[TT_VER_SPLIT] = false;
+
+  //if (modeType == MODE_TYPE_INTER && width * height == 64)  splits[TT_VER_SPLIT] = splits[TT_HOR_SPLIT] = false;
+  return 0;
+}
+
+
 int uvg_count_available_edge_cus(const cu_loc_t* const cu_loc, const lcu_t* const lcu, bool left)
 {
   if ((left && cu_loc->x == 0) || (!left && cu_loc->y == 0)) {

src/cu.h

@@ -106,6 +106,7 @@ typedef struct  {
   uint32_t split_tree;
   uint8_t current_depth;
   uint8_t mtt_depth;
+  uint8_t part_index;
 } split_tree_t;
 
 
@@ -185,12 +186,15 @@ typedef struct {
 } cu_loc_t;
 
 void uvg_cu_loc_ctor(cu_loc_t *loc, int x, int y, int width, int height);
+typedef struct encoder_state_t encoder_state_t;
 
 int uvg_get_split_locs(
   const cu_loc_t* const origin,
   enum split_type split,
   cu_loc_t out[4],
   uint8_t* separate_chroma);
+int uvg_get_possible_splits(const encoder_state_t* const state,
+  const cu_loc_t* const cu_loc, split_tree_t split_tree, enum uvg_tree_type tree_type, bool splits[6]);
 
 
 #define CU_GET_MV_CAND(cu_info_ptr, reflist) \

src/encode_coding_tree.c

@@ -508,7 +508,8 @@ static void encode_transform_unit(
   bool only_chroma,
   enum uvg_tree_type tree_type,
   bool last_split,
-  const cu_loc_t *original_loc)               // Original cu dimensions, before CU split
+  const cu_loc_t *original_loc,
+  const cu_loc_t* const chroma_loc)               // Original cu dimensions, before CU split
 {
   const videoframe_t * const frame = state->tile->frame;
   cabac_data_t* const cabac = &state->cabac;
@@ -581,10 +582,10 @@ static void encode_transform_unit(
 
   bool chroma_cbf_set = cbf_is_set(cur_pu->cbf, COLOR_U) ||
                         cbf_is_set(cur_pu->cbf, COLOR_V);
-  if ((chroma_cbf_set || joint_chroma) && last_split) {
+  if ((chroma_cbf_set || joint_chroma) && last_split && chroma_loc) {
     //Need to drop const to get lfnst constraints
     // Use original dimensions instead of ISP split dimensions
-    encode_chroma_tu(state, original_loc, (cu_info_t*)cur_pu, &scan_idx, coeff, joint_chroma, tree_type);
+    encode_chroma_tu(state, chroma_loc, (cu_info_t*)cur_pu, &scan_idx, coeff, joint_chroma, tree_type);
   }
 }
 
@@ -605,8 +606,10 @@ static void encode_transform_coeff(
   enum uvg_tree_type tree_type,
   bool last_split,
   bool can_skip_last_cbf,
-  int *luma_cbf_ctx,            // Always true except when writing sub partition coeffs (ISP)
+  int *luma_cbf_ctx,
-  const cu_loc_t * const original_loc)       // Original dimensions before ISP split
+  // Always true except when writing sub partition coeffs (ISP)
+  const cu_loc_t * const original_loc,
+  const cu_loc_t* const chroma_loc)       // Original dimensions before ISP split
 {
   cabac_data_t * const cabac = &state->cabac;
 
@@ -647,7 +650,8 @@ static void encode_transform_coeff(
     cu_loc_t split_cu_loc[4];
     const int split_count = uvg_get_split_locs(cu_loc, split, split_cu_loc,NULL);
     for (int i = 0; i < split_count; ++i) {
-      encode_transform_coeff(state, &split_cu_loc[i], only_chroma, coeff, NULL, tree_type, true, false, luma_cbf_ctx, &split_cu_loc[i]);
+      encode_transform_coeff(state, &split_cu_loc[i], only_chroma,
+        coeff, NULL, tree_type, true, false, luma_cbf_ctx, &split_cu_loc[i], &split_cu_loc[i]);
     }
     return;
   }
@@ -658,7 +662,7 @@ static void encode_transform_coeff(
   // Not the last CU for area of 64 pixels cowered by more than one luma CU.
   // Not the last ISP Split
   if (state->encoder_control->chroma_format != UVG_CSP_400
-    && (cur_tu->log2_height + cur_tu->log2_width >= 6 || only_chroma)
+    && (chroma_loc || only_chroma)
     && tree_type != UVG_LUMA_T
     && last_split) {
     cabac->cur_ctx = &(cabac->ctx.qt_cbf_model_cb[0]);
@@ -684,7 +688,7 @@ static void encode_transform_coeff(
   }
 
   if (cb_flag_y | cb_flag_u | cb_flag_v) {
-    if (state->must_code_qp_delta && (only_chroma || cb_flag_y || cur_tu->log2_height + cur_tu->log2_width >= 6) ) {
+    if (state->must_code_qp_delta && (only_chroma || cb_flag_y || chroma_loc) ) {
       const int qp_pred      = uvg_get_cu_ref_qp(state, cu_loc->x, cu_loc->y, state->last_qp);
       const int qp_delta     = cur_tu->qp - qp_pred;
       // Possible deltaQP range depends on bit depth as stated in HEVC specification.
@@ -711,7 +715,7 @@ static void encode_transform_coeff(
         ((cb_flag_u || cb_flag_v ) 
           && cur_tu->type == CU_INTRA)
         || (cb_flag_u && cb_flag_v)) 
-      && (cur_tu->log2_height + cur_tu->log2_width >= 6 || only_chroma || tree_type == UVG_CHROMA_T)
+      && (chroma_loc || only_chroma || tree_type == UVG_CHROMA_T)
       && state->encoder_control->cfg.jccr
       && last_split
       ) {
@@ -720,7 +724,7 @@ static void encode_transform_coeff(
       CABAC_BIN(cabac, cur_tu->joint_cb_cr != 0, "tu_joint_cbcr_residual_flag");
     }
 
-    encode_transform_unit(state, cu_loc, only_chroma ? cur_tu : NULL, coeff, only_chroma, tree_type, last_split, original_loc);
+    encode_transform_unit(state, cu_loc, only_chroma ? cur_tu : NULL, coeff, only_chroma, tree_type, last_split, original_loc, chroma_loc);
   }
 }
 
@@ -855,14 +859,14 @@ int uvg_encode_inter_prediction_unit(
 }
 
 static void encode_chroma_intra_cu(
-  cabac_data_t* const cabac, 
+  cabac_data_t* const cabac,
-  const cu_info_t* const cur_cu, 
+  const cu_info_t* const cur_cu,
-  const int cclm_enabled, 
+  const int cclm_enabled,
+  int8_t luma_intra_dir,
   double* bits_out) {
   unsigned pred_mode = 0;
   unsigned chroma_pred_modes[8] = {0, 50, 18, 1, 67, 81, 82, 83};
   int8_t chroma_intra_dir = cur_cu->intra.mode_chroma;
-  int8_t luma_intra_dir = !cur_cu->intra.mip_flag ? cur_cu->intra.mode : 0;
   for(int i = 0; i < 4; i++) {
     if(chroma_pred_modes[i] == luma_intra_dir) {
       chroma_pred_modes[i] = 66;
@@ -1399,12 +1403,13 @@ void uvg_encode_coding_tree(
       NULL);
     
     if (split_flag || border) {
-      const split_tree_t new_split_tree = { cur_cu->split_tree, split_tree.current_depth + 1, split_tree.mtt_depth + (split_flag != QT_SPLIT)};
+      split_tree_t new_split_tree = { cur_cu->split_tree, split_tree.current_depth + 1, split_tree.mtt_depth + (split_flag != QT_SPLIT), 0};
 
       cu_loc_t new_cu_loc[4];
       uint8_t separate_chroma = 0;
       const int splits = uvg_get_split_locs(cu_loc, split_flag, new_cu_loc, &separate_chroma);
       for (int split = 0; split <splits; ++split) {
+        new_split_tree.part_index = split;
         uvg_encode_coding_tree(state, coeff, tree_type,
           &new_cu_loc[split], separate_chroma ? chroma_loc : &new_cu_loc[split],
           new_split_tree, !separate_chroma || split == splits - 1);
@@ -1586,7 +1591,7 @@ void uvg_encode_coding_tree(
       // Code (possible) coeffs to bitstream
       if (has_coeffs) {
         int luma_cbf_ctx = 0;
-        encode_transform_coeff(state, cu_loc, 0, coeff, cur_cu, tree_type, true, false, &luma_cbf_ctx, cu_loc);
+        encode_transform_coeff(state, cu_loc, 0, coeff, cur_cu, tree_type, true, false, &luma_cbf_ctx, cu_loc, cu_loc);
       }
 
       encode_mts_idx(state, cabac, cur_cu, cu_loc);
@@ -1596,12 +1601,14 @@ void uvg_encode_coding_tree(
     if(tree_type != UVG_CHROMA_T) {
       uvg_encode_intra_luma_coding_unit(state, cabac, cur_cu, cu_loc, NULL, NULL);
     }
-
+    
-    const bool is_local_dual_tree = cu_height * cu_width < 64 && tree_type == UVG_BOTH_T;
+    const bool is_local_dual_tree = (chroma_loc->width != cu_loc->width || chroma_loc->height != cu_loc->height);
 
     // Code chroma prediction mode.
-    if (state->encoder_control->chroma_format != UVG_CSP_400 && cur_cu->log2_height + cur_cu->log2_width >= 6 && tree_type == UVG_BOTH_T) {
+    if (state->encoder_control->chroma_format != UVG_CSP_400 
-      encode_chroma_intra_cu(cabac, cur_cu, state->encoder_control->cfg.cclm, NULL);
+      && (chroma_loc->width == cu_loc->width && chroma_loc->height == cu_loc->height) 
+      && tree_type == UVG_BOTH_T) {
+      encode_chroma_intra_cu(cabac, cur_cu, state->encoder_control->cfg.cclm, !cur_cu->intra.mip_flag ? cur_cu->intra.mode : 0, NULL);
     }
     int luma_cbf_ctx = 0;
 
@@ -1620,7 +1627,9 @@ void uvg_encode_coding_tree(
 
         // Check if last split to write chroma
         bool last_split = (i + 1) == split_limit;
-        encode_transform_coeff(state, &split_loc, 0, coeff, NULL, tree_type, last_split, can_skip_last_cbf, &luma_cbf_ctx, cu_loc);
+        encode_transform_coeff(state, &split_loc,
+          0, coeff, NULL, tree_type, last_split, can_skip_last_cbf, &luma_cbf_ctx, 
+          cu_loc, is_local_dual_tree ? NULL : chroma_loc);
       }
     }
 
@@ -1631,16 +1640,17 @@ void uvg_encode_coding_tree(
 
     // For 4x4 the chroma PU/TU is coded after the last 
     if (state->encoder_control->chroma_format != UVG_CSP_400 &&
-      (((chroma_loc->width != cu_loc->width || chroma_loc->height != cu_loc->height)&&
+      ((is_local_dual_tree &&
       has_chroma) || tree_type == UVG_CHROMA_T) &&
       tree_type != UVG_LUMA_T)   {
-      encode_chroma_intra_cu(cabac, cur_cu, state->encoder_control->cfg.cclm, NULL);
+      int8_t luma_dir = uvg_get_co_located_luma_mode(chroma_loc->x, chroma_loc->y, chroma_loc->width, chroma_loc->height, NULL, frame->cu_array, UVG_CHROMA_T);
+      encode_chroma_intra_cu(cabac, cur_cu, state->encoder_control->cfg.cclm, luma_dir,NULL);
       // LFNST constraints must be reset here. Otherwise the left over values will interfere when calculating new constraints
       cu_info_t* tmp = (cu_info_t*)cur_cu;
       tmp->violates_lfnst_constrained_luma = false;
       tmp->violates_lfnst_constrained_chroma = false;
       tmp->lfnst_last_scan_pos = false;
-      encode_transform_coeff(state, chroma_loc, 1, coeff, cur_cu, tree_type, true, false, &luma_cbf_ctx, chroma_loc);
+      encode_transform_coeff(state, chroma_loc, 1, coeff, cur_cu, tree_type, true, false, &luma_cbf_ctx, chroma_loc, chroma_loc);
       // Write LFNST only once for single tree structure
       encode_lfnst_idx(state, cabac, tmp, is_local_dual_tree ? UVG_CHROMA_T : tree_type, COLOR_UV, chroma_loc);
     }
@@ -1683,6 +1693,7 @@ double uvg_mock_encode_coding_unit(
 
   int x_local = cu_loc->local_x >> (tree_type == UVG_CHROMA_T);
   int y_local = cu_loc->local_y >> (tree_type == UVG_CHROMA_T);
+  const bool is_separate_tree = chroma_loc == NULL || cu_loc->height != chroma_loc->height || cu_loc->width != chroma_loc->width;
     
   const cu_info_t* left_cu = NULL, *above_cu = NULL;
   if (x) {
@@ -1782,7 +1793,10 @@ double uvg_mock_encode_coding_unit(
       uvg_encode_intra_luma_coding_unit(state, cabac, cur_cu, cu_loc, lcu, &bits);
     }
     if((chroma_loc || tree_type == UVG_CHROMA_T) && state->encoder_control->chroma_format != UVG_CSP_400 && tree_type != UVG_LUMA_T) {
-      encode_chroma_intra_cu(cabac, cur_cu, state->encoder_control->cfg.cclm, &bits);
+      int8_t luma_dir = uvg_get_co_located_luma_mode(chroma_loc->x, chroma_loc->y, chroma_loc->width, chroma_loc->height,
+        tree_type != UVG_CHROMA_T ? lcu : NULL,
+        tree_type == UVG_CHROMA_T ? state->tile->frame->cu_array : NULL,   is_separate_tree ? UVG_CHROMA_T : tree_type);
+      encode_chroma_intra_cu(cabac, cur_cu, state->encoder_control->cfg.cclm, luma_dir, &bits);
     }
   }
   else {

src/encoderstate.c

@@ -883,7 +883,7 @@ static void encoder_state_worker_encode_lcu_bitstream(void * opaque)
   //Encode coding tree
   cu_loc_t start;
   uvg_cu_loc_ctor(&start, lcu->position.x * LCU_WIDTH, lcu->position.y * LCU_WIDTH, LCU_WIDTH, LCU_WIDTH);
-  split_tree_t split_tree = { 0, 0, 0 };
+  split_tree_t split_tree = { 0, 0, 0, 0 };
 
   uvg_encode_coding_tree(state, lcu->coeff, tree_type, &start, &start, split_tree, true);
 

src/intra.c

@@ -1602,7 +1602,7 @@ void uvg_intra_predict(
 }
 
 // This function works on luma coordinates 
-const cu_info_t* uvg_get_co_located_luma_cu(
+int8_t uvg_get_co_located_luma_mode(
   int x,
   int y,
   int width,
@@ -1617,12 +1617,17 @@ const cu_info_t* uvg_get_co_located_luma_cu(
     x += width >> 1;
     y += height >> 1;
   }
+  const cu_info_t* cu;
   if(cu_array) {
-    return uvg_cu_array_at_const(cu_array, x, y);
+    cu = uvg_cu_array_at_const(cu_array, x, y);
   }
   else {
-    return LCU_GET_CU_AT_PX(lcu, SUB_SCU(x), SUB_SCU(y));
+    cu = LCU_GET_CU_AT_PX(lcu, SUB_SCU(x), SUB_SCU(y));
   }
+  if (cu->intra.mip_flag) {
+    return 0;
+  }
+  return cu->intra.mode;
 }
 
 

src/intra.h

@@ -151,7 +151,7 @@ void uvg_intra_recon_cu(
   bool recon_luma,
   bool recon_chroma);
 
-const cu_info_t* uvg_get_co_located_luma_cu(
+int8_t uvg_get_co_located_luma_mode(
   int x,
   int y,
   int width,

src/search.c

@@ -1103,17 +1103,20 @@ static double search_cu(
         int8_t intra_mode = intra_search.pred_cu.intra.mode;
         
         if ((has_chroma || tree_type == UVG_CHROMA_T)
-          && state->encoder_control->chroma_format != UVG_CSP_400 && tree_type != UVG_LUMA_T) {
+          && state->encoder_control->chroma_format != UVG_CSP_400) {
 
           intra_search.pred_cu.joint_cb_cr = 0;
-          if(tree_type == UVG_CHROMA_T) {
+          if(tree_type == UVG_CHROMA_T || is_separate_tree) {
-            intra_search.pred_cu.intra = uvg_get_co_located_luma_cu(x, y, luma_width, luma_width, NULL, state->tile->frame->cu_array, UVG_CHROMA_T)->intra;
+            intra_mode = uvg_get_co_located_luma_mode(chroma_loc->x, chroma_loc->y, chroma_loc->width, chroma_loc->height,
-            intra_mode = intra_search.pred_cu.intra.mode;
+                                                    is_separate_tree ? lcu : NULL,
+                                                    tree_type == UVG_CHROMA_T ? state->tile->frame->cu_array : NULL, UVG_CHROMA_T);
             intra_search.pred_cu.type = CU_INTRA;
+          } else  if (intra_search.pred_cu.intra.mip_flag) {
+            intra_mode = 0;
           }
-          intra_search.pred_cu.intra.mode_chroma = intra_search.pred_cu.intra.mode;
+          intra_search.pred_cu.intra.mode_chroma = intra_mode;
           if (ctrl->cfg.rdo >= 2 || ctrl->cfg.jccr || ctrl->cfg.lfnst) {
-            uvg_search_cu_intra_chroma(state, chroma_loc, lcu, &intra_search, tree_type, cu_loc->x != chroma_loc->x || cu_loc->y != chroma_loc->y);
+            uvg_search_cu_intra_chroma(state, chroma_loc, lcu, &intra_search, intra_mode, tree_type, is_separate_tree);
 
             if (intra_search.pred_cu.joint_cb_cr == 0) {
               intra_search.pred_cu.joint_cb_cr = 4;
@@ -1121,7 +1124,7 @@ static double search_cu(
 
           }
           else if (!intra_search.pred_cu.intra.mip_flag) {
-            intra_search.pred_cu.intra.mode_chroma = intra_search.pred_cu.intra.mode;
+            intra_search.pred_cu.intra.mode_chroma = intra_mode;
           }
           else {
             intra_search.pred_cu.intra.mode_chroma = 0;
@@ -1138,14 +1141,12 @@ static double search_cu(
           else {
             intra_cost = intra_search.cost;
           }
-          intra_search.pred_cu.intra.mode = intra_mode;
           intra_search.pred_cu.violates_lfnst_constrained_chroma = false;
           intra_search.pred_cu.lfnst_last_scan_pos = false;
         }
         else {
           intra_search.pred_cu.intra.mode_chroma = intra_mode;
         }
-        intra_search.pred_cu.intra.mode = intra_mode;
       }
       if (intra_cost < cost) {
         cost = intra_cost;
@@ -1207,6 +1208,8 @@ static double search_cu(
                            tree_type,
                            false,
                            true);
+      } else {
+        assert(cur_cu->cr_lfnst_idx == 0 && "If we don't have separate tree chroma lfnst index must be 0");
       }
       if (cur_cu->joint_cb_cr == 4) cur_cu->joint_cb_cr = 0;
 
@@ -1346,10 +1349,11 @@ static double search_cu(
   // Recursively split all the way to max search depth.
   if (can_split_cu) {
     const int split_type = depth == 2 ? TT_VER_SPLIT : QT_SPLIT;
-    const split_tree_t new_split = {
+    split_tree_t new_split = {
       split_tree.split_tree | split_type << (split_tree.current_depth * 3),
       split_tree.current_depth + 1,
       split_tree.mtt_depth + (split_type != QT_SPLIT),
+      0
     };
     
     double split_cost = 0.0;
@@ -1408,6 +1412,7 @@ static double search_cu(
       const int splits = uvg_get_split_locs(cu_loc, split_type, new_cu_loc, &separate_chroma);
       initialize_partial_work_tree(lcu, &split_lcu, cu_loc, tree_type);
       for (int split = 0; split < splits; ++split) {
+        new_split.part_index = split;
         split_cost += search_cu(state, 
           &new_cu_loc[split], separate_chroma ? chroma_loc : &new_cu_loc[split],
           &split_lcu, 

src/search_inter.c

@@ -2169,7 +2169,7 @@ void uvg_cu_cost_inter_rd2(
     depth++;
     splits >>= 3;
   }
-  const split_tree_t splitt_tree = { cur_cu->split_tree, depth, mtt_depth };
+  const split_tree_t splitt_tree = { cur_cu->split_tree, depth, mtt_depth, 0};
   if (cur_cu->merged) {
     no_cbf_bits = CTX_ENTROPY_FBITS(&state->cabac.ctx.cu_skip_flag_model[skip_context], 1) + *inter_bitcost;
     bits += uvg_mock_encode_coding_unit(state, cabac, cu_loc, cu_loc, lcu, cur_cu, UVG_BOTH_T, splitt_tree);

src/search_intra.c

@@ -1452,11 +1452,6 @@ int8_t uvg_search_intra_chroma_rdo(
         lfnst_modes_to_check[i] = i;
       }
     }
-    else if(chroma_data->pred_cu.lfnst_idx) {
-      lfnst_modes_to_check[0] = chroma_data->pred_cu.lfnst_idx;
-      lfnst_modes_to_check[1] = -1;
-      lfnst_modes_to_check[2] = -1;
-    }
     else {
       lfnst_modes_to_check[0] = 0;
       lfnst_modes_to_check[1] = -1;
@@ -1591,17 +1586,17 @@ int8_t uvg_search_cu_intra_chroma(
   const cu_loc_t* const cu_loc,
   lcu_t *lcu,
   intra_search_data_t *search_data,
+  int8_t luma_mode,
   enum uvg_tree_type tree_type,
   bool is_separate)
 {
 
   const cu_info_t *cur_pu = &search_data->pred_cu;
-  int8_t intra_mode = !cur_pu->intra.mip_flag ? cur_pu->intra.mode : 0;
   
-  int8_t modes[8] = { 0, 50, 18, 1, intra_mode, 81, 82, 83 };
+  int8_t modes[8] = { 0, 50, 18, 1, luma_mode, 81, 82, 83 };
   uint8_t total_modes = (state->encoder_control->cfg.cclm ? 8 : 5);
   for(int i = 0; i < 4; i++) {
-    if (modes[i] == intra_mode) {
+    if (modes[i] == luma_mode) {
       modes[i] = 66;
       break;
     }
@@ -1623,7 +1618,7 @@ int8_t uvg_search_cu_intra_chroma(
   FILL(chroma_data, 0);
   for (int i = 0; i < num_modes; i++) {
     chroma_data[i].pred_cu = *cur_pu;
-    chroma_data[i].pred_cu.intra.mode_chroma = num_modes == 1 ? intra_mode : modes[i];
+    chroma_data[i].pred_cu.intra.mode_chroma = num_modes == 1 ? luma_mode : modes[i];
     chroma_data[i].cost = 0;
     if(cu_loc->width != 4 && tree_type == UVG_BOTH_T) {
       memcpy(chroma_data[i].lfnst_costs, search_data->lfnst_costs, sizeof(double) * 3);
@@ -1636,13 +1631,13 @@ int8_t uvg_search_cu_intra_chroma(
   if(state->encoder_control->cfg.cclm && 0){
     
 
-    num_modes = search_intra_chroma_rough(state, chroma_data, lcu, intra_mode,
+    num_modes = search_intra_chroma_rough(state, chroma_data, lcu, luma_mode,
                                           tree_type,
                                           cu_loc);
   }
   
   if (num_modes > 1 || state->encoder_control->cfg.jccr) {
-    uvg_search_intra_chroma_rdo(state, num_modes, lcu, cu_loc, chroma_data, intra_mode, tree_type, is_separate);
+    uvg_search_intra_chroma_rdo(state, num_modes, lcu, cu_loc, chroma_data, luma_mode, tree_type, is_separate);
   }
   else if(cur_pu->lfnst_idx) {
     chroma_data[0].pred_cu.cr_lfnst_idx = cur_pu->lfnst_idx;

src/search_intra.h

@@ -55,6 +55,7 @@ int8_t uvg_search_cu_intra_chroma(
   const cu_loc_t* const cu_loc,
   lcu_t *lcu,
   intra_search_data_t* best_cclm,
+  int8_t luma_mode,
   enum uvg_tree_type tree_type,
   bool is_separate);
 

src/transform.c

@@ -872,7 +872,7 @@ void uvg_fwd_lfnst(
   
   const int scan_order = SCAN_DIAG;
 
-  if (lfnst_index && !mts_skip && (is_separate_tree || color == COLOR_Y))
+  if (lfnst_index && !mts_skip)
   {
     assert(log2_width != -1 && "LFNST: invalid block width.");
     const bool whge3 = width >= 8 && height >= 8;
@@ -1005,7 +1005,7 @@ void uvg_inv_lfnst(
   bool is_mip = block_is_mip(cur_cu, color, is_separate_tree);
   const int scan_order = SCAN_DIAG;
   
-  if (lfnst_index && !mts_skip && (is_separate_tree || color == COLOR_Y)) {
+  if (lfnst_index && !mts_skip) {
     const bool whge3 = width >= 8 && height >= 8;
     const uint32_t* scan = whge3 ? uvg_coef_top_left_diag_scan_8x8[log2_width] : uvg_g_sig_last_scan[scan_order][log2_width - 1];