[isp] Implement coefficient encoding for isp splits. Make get_split_dim non static, it was needed elsewhere after all.

src/encode_coding_tree.c

@@ -525,7 +525,8 @@ static void encode_transform_unit(
   int depth,
   bool only_chroma,
   lcu_coeff_t* coeff,
-  enum uvg_tree_type tree_type)
+  enum uvg_tree_type tree_type,
+  bool last_split)
 {
   assert(depth >= 1 && depth <= MAX_PU_DEPTH);
 
@@ -586,7 +587,7 @@ static void encode_transform_unit(
 
   bool chroma_cbf_set = cbf_is_set(cur_pu->cbf, depth, COLOR_U) ||
                         cbf_is_set(cur_pu->cbf, depth, COLOR_V);
-  if (chroma_cbf_set || joint_chroma) {
+  if ((chroma_cbf_set || joint_chroma) && last_split) {
     //Need to drop const to get lfnst constraints
     encode_chroma_tu(state, x, y, depth, width_c, (cu_info_t*)cur_pu, &scan_idx, coeff, joint_chroma, tree_type);
   }
@@ -611,7 +612,8 @@ static void encode_transform_coeff(
   uint8_t parent_coeff_v,
   bool only_chroma,
   lcu_coeff_t* coeff,
-  enum uvg_tree_type tree_type)
+  enum uvg_tree_type tree_type,
+  bool last_split)                // Always true except when writing sub partition coeffs (ISP)
 {
   cabac_data_t * const cabac = &state->cabac;
   //const encoder_control_t *const ctrl = state->encoder_control;
@@ -642,8 +644,6 @@ static void encode_transform_coeff(
 
   int8_t split = (LCU_WIDTH >> depth > TR_MAX_WIDTH);
 
- 
-
   const int cb_flag_y = tree_type != UVG_CHROMA_T ? cbf_is_set(cur_pu->cbf, depth, COLOR_Y) : 0;
   const int cb_flag_u = tree_type != UVG_LUMA_T ?( cur_pu->joint_cb_cr ? (cur_pu->joint_cb_cr >> 1) & 1 : cbf_is_set(cur_cu->cbf, depth, COLOR_U)) : 0;
   const int cb_flag_v = tree_type != UVG_LUMA_T ? (cur_pu->joint_cb_cr ? cur_pu->joint_cb_cr & 1 : cbf_is_set(cur_cu->cbf, depth, COLOR_V)) : 0;
@@ -671,7 +671,7 @@ static void encode_transform_coeff(
   // - they have already been signaled to 0 previously
   // When they are not present they are inferred to be 0, except for size 4
   // when the flags from previous level are used.
-  if (state->encoder_control->chroma_format != UVG_CSP_400 && (depth != 4 || only_chroma) && tree_type != UVG_LUMA_T) {
+  if (state->encoder_control->chroma_format != UVG_CSP_400 && (depth != 4 || only_chroma) && tree_type != UVG_LUMA_T && last_split) {
     
     if (!split) {
       if (true) {
@@ -690,10 +690,10 @@ static void encode_transform_coeff(
     uint8_t offset = LCU_WIDTH >> (depth + 1);
     int x2 = x + offset;
     int y2 = y + offset;
-    encode_transform_coeff(state, x,  y,  depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type);
+    encode_transform_coeff(state, x,  y,  depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true);
-    encode_transform_coeff(state, x2, y,  depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type);
+    encode_transform_coeff(state, x2, y,  depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true);
-    encode_transform_coeff(state, x,  y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type);
+    encode_transform_coeff(state, x,  y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true);
-    encode_transform_coeff(state, x2, y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type);
+    encode_transform_coeff(state, x2, y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true);
     return;
   }
 
@@ -737,12 +737,13 @@ static void encode_transform_coeff(
         || (cb_flag_u && cb_flag_v)) 
       && (depth != 4 || only_chroma || tree_type == UVG_CHROMA_T) 
       && state->encoder_control->cfg.jccr
+      && last_split
       ) {
       assert(cur_pu->joint_cb_cr < 4 && "JointCbCr is in search state.");
       cabac->cur_ctx = &cabac->ctx.joint_cb_cr[cb_flag_u * 2 + cb_flag_v - 1];
       CABAC_BIN(cabac, cur_pu->joint_cb_cr != 0, "tu_joint_cbcr_residual_flag");
     }
-    encode_transform_unit(state, x, y, depth, only_chroma, coeff, tree_type);
+    encode_transform_unit(state, x, y, depth, only_chroma, coeff, tree_type, last_split);
   }
 }
 
@@ -1671,7 +1672,7 @@ void uvg_encode_coding_tree(
       // Code (possible) coeffs to bitstream
 
       if (cbf) {
-        encode_transform_coeff(state, x, y, depth, 0, 0, 0, 0, coeff, tree_type);
+        encode_transform_coeff(state, x, y, depth, 0, 0, 0, 0, coeff, tree_type, true);
       }
 
       encode_mts_idx(state, cabac, cur_cu);
@@ -1688,7 +1689,25 @@ void uvg_encode_coding_tree(
     }
 
     if (tree_type != UVG_CHROMA_T) {
-      encode_transform_coeff(state, x, y, depth, 0, 0, 0, 0, coeff, tree_type);
+      // Cycle through sub partitions if ISP enabled.
+      // ISP split is done horizontally or vertically depending on ISP mode, 2 or 4 times depending on block dimensions.
+      // Small blocks are split only twice.
+      int split_type = cur_cu->intra.isp_mode;
+
+      int part_dim = cu_width;
+      if (split_type != ISP_MODE_NO_ISP) {
+        part_dim = uvg_get_isp_split_dim(cu_width, cu_height, split_type);
+      }
+      int limit = split_type == ISP_MODE_HOR ? cu_height : cu_width;
+
+      for (int part = 0; part < limit; part += part_dim) {
+        const int part_x = split_type == ISP_MODE_HOR ? x : x + part;
+        const int part_y = split_type == ISP_MODE_HOR ? y + part : y;
+        
+        // Check if last split to write chroma
+        bool last_split = (part + part_dim) == limit;
+        encode_transform_coeff(state, part_x, part_y, depth, 0, 0, 0, 0, coeff, tree_type, last_split);
+      }
     }
 
     if (tree_type != UVG_CHROMA_T) {
@@ -1706,7 +1725,7 @@ void uvg_encode_coding_tree(
       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);
+      encode_transform_coeff(state, x, y, depth, 0, 0, 0, 1, coeff, tree_type, true);
       // Write LFNST only once for single tree structure
       encode_lfnst_idx(state, cabac, tmp, x, y, depth, cu_width, cu_height, tree_type, COLOR_UV);
     }

src/intra.c

@@ -1481,8 +1481,10 @@ const cu_info_t* uvg_get_co_located_luma_cu(
 * \param height       Block height.
 * \param split_type   Horizontal or vertical split.
 */
-static int get_isp_split_dim(const int width, const int height, const int split_type)
+int uvg_get_isp_split_dim(const int width, const int height, const int split_type)
 {
+  assert(split_type != ISP_MODE_NO_ISP && "Cannot calculate split dimension if no split type is set. Make sure this function is not called in this case.");
+
   bool divide_in_rows = split_type == SPLIT_TYPE_HOR;
   int split_dim_size, non_split_dim_size, partition_size, div_shift = 2;
 
@@ -1666,9 +1668,9 @@ void uvg_intra_recon_cu(
     // ISP split is done horizontally or vertically depending on ISP mode, 2 or 4 times depending on block dimensions.
     // Small blocks are split only twice.
     int split_type = search_data->pred_cu.intra.isp_mode;
-    int part_dim = get_isp_split_dim(width, height, split_type);
+    int part_dim = uvg_get_isp_split_dim(width, height, split_type);
     int limit = split_type == ISP_MODE_HOR ? height : width;
-    for (int part = 0; part < limit; part + part_dim) {
+    for (int part = 0; part < limit; part += part_dim) {
       const int part_x = split_type == ISP_MODE_HOR ? x : x + part;
       const int part_y = split_type == ISP_MODE_HOR ? y + part: y;
       const int part_w = split_type == ISP_MODE_HOR ? part_dim : width;
@@ -1749,8 +1751,8 @@ bool uvg_can_use_isp_with_lfnst(const int width, const int height, const int isp
     return false;
   }
 
-  const int tu_width = (isp_split_type == ISP_MODE_HOR) ? width : get_isp_split_dim(width, height, SPLIT_TYPE_VER);
+  const int tu_width = (isp_split_type == ISP_MODE_HOR) ? width : uvg_get_isp_split_dim(width, height, SPLIT_TYPE_VER);
-  const int tu_height = (isp_split_type == ISP_MODE_HOR) ? get_isp_split_dim(width, height, SPLIT_TYPE_HOR) : height;
+  const int tu_height = (isp_split_type == ISP_MODE_HOR) ? uvg_get_isp_split_dim(width, height, SPLIT_TYPE_HOR) : height;
 
   if (!(tu_width >= TR_MIN_WIDTH && tu_height >= TR_MIN_WIDTH))
   {