[mtt] Fix implicit splits for non ctu divisible frames.

src/cu.c

@@ -370,10 +370,17 @@ int uvg_get_split_locs(
 }
 
 
-int uvg_get_implicit_split(const encoder_state_t* const state, const cu_loc_t* const cu_loc)
+int uvg_get_implicit_split(const encoder_state_t* const state, const cu_loc_t* const cu_loc, enum
+                           uvg_tree_type tree_type)
 {
-  bool right_ok = state->tile->frame->width >= cu_loc->x + cu_loc->width;
+  // This checking if cabac is in update state is a very dirty way of checking
-  bool bottom_ok = state->tile->frame->height >= cu_loc->y + cu_loc->height;
+  // whether we are in the search or writing the bitstream, and unfortunately the
+  // coordinates are different for chroma tree in those two conditions. It might be
+  // possible to pass the chroma loc for uvg_get_possible_splits in the search but
+  // then all of the conditions need to be checked in that function.
+  // This current solutions *might* not work with alf enabled but I think it should work
+  bool right_ok = (state->tile->frame->width >> (tree_type == UVG_CHROMA_T && state->cabac.update)) >= cu_loc->x + cu_loc->width;
+  bool bottom_ok = (state->tile->frame->height >> (tree_type == UVG_CHROMA_T && state->cabac.update)) >= cu_loc->y + cu_loc->height;
 
   if (right_ok && bottom_ok) return NO_SPLIT;
   if (right_ok) return BT_HOR_SPLIT;
@@ -387,10 +394,11 @@ int uvg_get_possible_splits(const encoder_state_t * const state,
 {
   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 enum split_type implicitSplit = uvg_get_implicit_split(state, cu_loc, tree_type);
   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_btd =
+    state->encoder_control->cfg.max_btt_depth[slice_type] + split_tree.implicit_mtt_depth;
   const unsigned max_bt_size = state->encoder_control->cfg.max_bt_size[slice_type] >> (tree_type == UVG_CHROMA_T);
   const unsigned min_bt_size = 1 << MIN_SIZE >> (tree_type == UVG_CHROMA_T);
   const unsigned max_tt_size = state->encoder_control->cfg.max_tt_size[slice_type] >> (tree_type == UVG_CHROMA_T);

src/cu.h

@@ -105,6 +105,7 @@ typedef struct  {
   uint32_t split_tree;
   uint8_t current_depth;
   uint8_t mtt_depth;
+  uint8_t implicit_mtt_depth;
   uint8_t part_index;
 } split_tree_t;
 

src/encode_coding_tree.c

@@ -1239,6 +1239,7 @@ uint8_t uvg_write_split_flag(
   const cu_loc_t* const cu_loc,
   split_tree_t split_tree,
   enum uvg_tree_type tree_type,
+  bool* is_implicit_out,
   double* bits_out)
 {
   double bits = 0;
@@ -1257,7 +1258,10 @@ uint8_t uvg_write_split_flag(
 
   enum split_type split_flag = (split_tree.split_tree >> (split_tree.current_depth * 3)) & 7;
 
-  split_flag = is_implicit ? (can_split[QT_SPLIT] ? QT_SPLIT : (can_split[BT_HOR_SPLIT] ? BT_HOR_SPLIT : BT_VER_SPLIT)) : split_flag;
+  assert(can_split[split_flag] && "Trying to write an illegal split");
+
+  // split_flag = is_implicit ? (can_split[QT_SPLIT] ? QT_SPLIT : (can_split[BT_HOR_SPLIT] ? BT_HOR_SPLIT : BT_VER_SPLIT)) : split_flag;
+  *is_implicit_out = is_implicit;
 
   int split_model = 0;
   if (can_split[NO_SPLIT] && allow_split) {
@@ -1287,7 +1291,9 @@ uint8_t uvg_write_split_flag(
   }
 
 
-  if (!is_implicit && (can_split[BT_HOR_SPLIT] || can_split[BT_VER_SPLIT] || can_split[TT_HOR_SPLIT] || can_split[TT_VER_SPLIT]) && split_flag != NO_SPLIT) {
+  if ((!is_implicit || (can_split[QT_SPLIT] && (can_split[BT_HOR_SPLIT] || can_split[BT_VER_SPLIT]))) 
+    && (can_split[BT_HOR_SPLIT] || can_split[BT_VER_SPLIT] || can_split[TT_HOR_SPLIT] || can_split[TT_VER_SPLIT]) 
+    && split_flag != NO_SPLIT) {
     bool qt_split = split_flag == QT_SPLIT;
     if((can_split[BT_VER_SPLIT] || can_split[BT_HOR_SPLIT] || can_split[TT_VER_SPLIT] || can_split[TT_HOR_SPLIT]) && can_split[QT_SPLIT]) {
       unsigned left_qt_depth = 0;
@@ -1374,12 +1380,9 @@ void uvg_encode_coding_tree(
 
   int32_t frame_width = tree_type !=  UVG_CHROMA_T ? ctrl->in.width : ctrl->in.width / 2;
   int32_t frame_height = tree_type != UVG_CHROMA_T ? ctrl->in.height : ctrl->in.height / 2;
-  // Check for slice border
+
-  bool border_x = frame_width  < abs_x + cu_width;
+  // Stop if we are outside of the frame
-  bool border_y = frame_height < abs_y + cu_height;
+  if (abs_x >= frame_width || abs_y >= frame_height) return;
-  bool border_split_x = frame_width  >= abs_x + (LCU_WIDTH >> MAX_DEPTH) + cu_width / 2;
-  bool border_split_y = frame_height >= abs_y + (LCU_WIDTH >> MAX_DEPTH) + cu_height / 2;
-  bool border = border_x || border_y; /*!< are we in any border CU */
 
   if (depth <= state->frame->max_qp_delta_depth) {
     state->must_code_qp_delta = true;
@@ -1388,6 +1391,7 @@ void uvg_encode_coding_tree(
   // When not in MAX_DEPTH, insert split flag and split the blocks if needed
   if (cu_width + cu_height > 8) {
     split_tree.split_tree = cur_cu->split_tree;
+    bool is_implicit;
     const int split_flag = uvg_write_split_flag(
       state,
       cabac,
@@ -1396,10 +1400,16 @@ void uvg_encode_coding_tree(
       tree_type != UVG_CHROMA_T ? cu_loc : chroma_loc,
       split_tree,
       tree_type,
-      NULL);
+      &is_implicit,
+      NULL
+      );
     
-    if (split_flag || border) {
+    if (split_flag != NO_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};
+      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.implicit_mtt_depth + (split_flag != QT_SPLIT && is_implicit),
+      0};
 
       cu_loc_t new_cu_loc[4];
       cu_loc_t chroma_tree_loc;
@@ -1702,6 +1712,8 @@ double uvg_mock_encode_coding_unit(
 
   // When not in MAX_DEPTH, insert split flag and split the blocks if needed
   if (cur_cu->log2_height + cur_cu->log2_width > 4) {
+    // We do not care about whether the split is implicit or not since there is never split here
+    bool is_implicit;
     uvg_write_split_flag(
       state,
       cabac,
@@ -1709,8 +1721,9 @@ double uvg_mock_encode_coding_unit(
       above_cu,
       cu_loc,
       split_tree,
-      tree_type,
+      tree_type, &is_implicit,
-      &bits);
+      &bits
+      );
   }
 
   // Encode skip flag

src/encode_coding_tree.h

@@ -108,6 +108,7 @@ uint8_t uvg_write_split_flag(
   const cu_loc_t* const cu_loc,
   split_tree_t,
   enum uvg_tree_type tree_type,
+  bool* is_implicit_out,
   double* bits_out);
 
 void uvg_encode_last_significant_xy(cabac_data_t * const cabac,

src/encoderstate.c

@@ -780,7 +780,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, 0 };
 
   uvg_encode_coding_tree(state, lcu->coeff, tree_type, &start, &start, split_tree, true);
 

src/search.c

@@ -1350,19 +1350,13 @@ static double search_cu(
         || (tree_type == UVG_CHROMA_T && split_type == TT_HOR_SPLIT && cu_loc->chroma_height == 8)
         || (tree_type == UVG_CHROMA_T && split_type == BT_HOR_SPLIT && cu_loc->chroma_height == 4))
         continue;
-      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;
       int cbf = cbf_is_set_any(cur_cu->cbf);
       memcpy(&state->search_cabac, &pre_search_cabac, sizeof(post_seach_cabac));
 
 
       double split_bits = 0;
+      bool is_implicit = false;
 
       if (cur_cu->log2_height + cur_cu->log2_width > 4) {
 
@@ -1395,9 +1389,20 @@ static double search_cu(
           tree_type != UVG_CHROMA_T ? cu_loc : &separate_tree_chroma_loc,
           count_tree,
           tree_type,
-          &split_bits);
+          &is_implicit,
+          &split_bits
+          );
       }
 
+
+      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),
+        split_tree.implicit_mtt_depth + (split_type != QT_SPLIT && is_implicit),
+        0
+      };
+
       state->search_cabac.update = 0;
       split_cost += split_bits * state->lambda;
 
@@ -1457,7 +1462,8 @@ static double search_cu(
         double bits = 0;
         uvg_write_split_flag(state, &state->search_cabac,
                              x > 0 ? LCU_GET_CU_AT_PX(lcu, SUB_SCU(x) - 1, SUB_SCU(y)) : NULL,
-                             y > 0 ? LCU_GET_CU_AT_PX(lcu, SUB_SCU(x), SUB_SCU(y) - 1) : NULL, cu_loc, split_tree, tree_type, &bits);
+                             y > 0 ? LCU_GET_CU_AT_PX(lcu, SUB_SCU(x), SUB_SCU(y) - 1) : NULL, cu_loc, split_tree, tree_type, NULL,
+                             &bits);
 
         cur_cu->intra = cu_d1->intra;
         cur_cu->type = CU_INTRA;
@@ -1737,7 +1743,7 @@ void uvg_search_lcu(encoder_state_t * const state, const int x, const int y, con
 
   cu_loc_t start;
   uvg_cu_loc_ctor(&start, x, y, LCU_WIDTH, LCU_WIDTH);
-  split_tree_t split_tree = { 0, 0, 0 };
+  split_tree_t split_tree = { 0, 0, 0, 0, 0 };
   // Start search from depth 0.
   double cost = search_cu(
     state,