[isp] Set cbfs for isp splits after search. Add helper function for isp split number.

src/encode_coding_tree.c

@@ -1623,24 +1623,14 @@ void uvg_encode_coding_tree(
       // 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 split_limit = split_type == ISP_MODE_NO_ISP ? 1 : uvg_get_isp_split_num(cu_width, cu_height, split_type);
 
-      int part_dim = cu_width;
+      for (int i = 0; i < split_limit; ++i) {
-      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;
-        const int part_w = split_type == ISP_MODE_HOR ? cu_width : part_dim;
-        const int part_h = split_type == ISP_MODE_HOR ? part_dim : cu_height;
-        
         cu_loc_t loc;
-        uvg_cu_loc_ctor(&loc, part_x, part_y, part_w, part_h);
+        uvg_get_isp_split_loc(&loc, x, y, cu_width, cu_height, i, split_type);
 
         // Check if last split to write chroma
-        bool last_split = (part + part_dim) == limit;
+        bool last_split = (i + 1) == split_limit;
         encode_transform_coeff(state, &loc, depth, 0, 0, 0, 0, coeff, tree_type, last_split);
       }
     }

src/global.h

@@ -129,8 +129,8 @@ typedef int16_t coeff_t;
 typedef int32_t mv_t;
 
 //#define VERBOSE 1
-#define UVG_DEBUG_PRINT_CABAC 1
+//#define UVG_DEBUG_PRINT_CABAC 1
-#define UVG_DEBUG 1
+//#define UVG_DEBUG 1
 
 //#define UVG_DEBUG_PRINT_YUVIEW_CSV 1
 //#define UVG_DEBUG_PRINT_MV_INFO 1

src/intra.c

@@ -1507,6 +1507,34 @@ int uvg_get_isp_split_dim(const int width, const int height, const int split_typ
 }
 
 
+int uvg_get_isp_split_num(const int width, const int height, const int split_type)
+{
+  assert((split_type != ISP_MODE_NO_ISP) && "This function cannot be called if ISP mode is 0.");
+  int split_dim = uvg_get_isp_split_dim(width, height, split_type);
+  int num = split_type == ISP_MODE_HOR ? height / split_dim : width / split_dim;
+
+  return num;
+}
+
+
+void uvg_get_isp_split_loc(cu_loc_t *loc, const int x, const int y, const int block_w, const int block_h, const int split_idx, const int split_type)
+{
+  assert((split_idx >= 0 && split_idx <= 3) && "ISP split index must be in [0, 3].");
+  int part_dim = block_w;
+  if (split_type != ISP_MODE_NO_ISP) {
+    part_dim = uvg_get_isp_split_dim(block_w, block_h, split_type);
+  }
+  const int offset = part_dim * split_idx;
+
+  const int part_x = split_type == ISP_MODE_HOR ? x : x + offset;
+  const int part_y = split_type == ISP_MODE_HOR ? y + offset : y;
+  const int part_w = split_type == ISP_MODE_HOR ? block_w : part_dim;
+  const int part_h = split_type == ISP_MODE_HOR ? part_dim : block_h;
+
+  uvg_cu_loc_ctor(loc, part_x, part_y, part_w, part_h);
+}
+
+
 static void intra_recon_tb_leaf(
   encoder_state_t* const state,
   const cu_loc_t* cu_loc,
@@ -1667,25 +1695,17 @@ 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 = uvg_get_isp_split_dim(width, height, split_type);
+    int split_limit = split_type == ISP_MODE_NO_ISP ? 1 : uvg_get_isp_split_num(width, height, split_type);
-    int limit = split_type == ISP_MODE_HOR ? height : width;
-    int split_num = 0;
-    for (int part = 0; part < limit; part += part_dim) {
-      cbf_clear(&cur_cu->cbf, depth, COLOR_Y);
-      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 ? width : part_dim;
-      const int part_h = split_type == ISP_MODE_HOR ? part_dim : height;
 
+    for (int i = 0; i < split_limit; ++i) {
       cu_loc_t loc;
-      uvg_cu_loc_ctor(&loc, part_x, part_y, part_w, part_h);
+      uvg_get_isp_split_loc(&loc, x, y, width, height, i, split_type);
 
       intra_recon_tb_leaf(state, &loc, lcu, COLOR_Y, search_data, tree_type);
       uvg_quantize_lcu_residual(state, true, false, false,
-                                &loc, depth, cur_cu, lcu,
+        &loc, depth, cur_cu, lcu,
-                                false, tree_type);
+        false, tree_type);
-      search_data->best_isp_cbfs |= cbf_is_set(cur_cu->cbf, depth, COLOR_Y) << (split_num++);
+      search_data->best_isp_cbfs |= cbf_is_set(cur_cu->cbf, depth, COLOR_Y) << (i++);
-   
     }
   }
   const bool has_luma = recon_luma && search_data->pred_cu.intra.isp_mode == ISP_MODE_NO_ISP;

src/intra.h

@@ -170,5 +170,7 @@ int uvg_get_mip_flag_context(int x, int y, int width, int height, const lcu_t* l
 #define SPLIT_TYPE_VER 2
 
 int uvg_get_isp_split_dim(const int width, const int height, const int split_type);
+int uvg_get_isp_split_num(const int width, const int height, const int split_type);
+void uvg_get_isp_split_loc(cu_loc_t *loc, const int x, const int y, const int block_w, const int block_h, const int split_idx, const int split_type);
 bool uvg_can_use_isp(const int width, const int height, const int max_tr_size);
 bool uvg_can_use_isp_with_lfnst(const int width, const int height, const int isp_mode, const enum uvg_tree_type tree_type);

src/search.c

@@ -36,6 +36,7 @@
 #include <string.h>
 
 #include "cabac.h"
+#include "cu.h"
 #include "encoder.h"
 #include "encode_coding_tree.h"
 #include "imagelist.h"
@@ -386,15 +387,17 @@ double uvg_cu_rd_cost_luma(const encoder_state_t *const state,
     }
     else {
       int split_type = pred_cu->intra.isp_mode;
-      int part_dim = uvg_get_isp_split_dim(width, height, split_type);
+      int split_limit = split_type == ISP_MODE_NO_ISP ? 1 : uvg_get_isp_split_num(width, height, split_type);
-      int limit = split_type == ISP_MODE_HOR ? height : width;
-      int split_num = 0;
-      for (int part = 0; part < limit; part += part_dim) {
-        const int part_x = split_type == ISP_MODE_HOR ? x_px : x_px + part;
-        const int part_y = split_type == ISP_MODE_HOR ? y_px + part : y_px;
-        const int part_w = split_type == ISP_MODE_HOR ? width : part_dim;
-        const int part_h = split_type == ISP_MODE_HOR ? part_dim : height;
 
+      for (int i = 0; i < split_limit; ++i) {
+        cu_loc_t loc;
+        uvg_get_isp_split_loc(&loc, x_px, y_px, width, height, i, split_type);
+        const int part_x = loc.x;
+        const int part_y = loc.y;
+        const int part_w = loc.width;
+        const int part_h = loc.height;
+
+        // TODO: maybe just pass the cu_loc_t to these functions
         const coeff_t* coeffs = &lcu->coeff.y[xy_to_zorder(LCU_WIDTH, part_x, part_y)];
 
         coeff_bits += uvg_get_coeff_cost(state, coeffs, NULL, part_w, part_h, 0, luma_scan_mode, pred_cu->tr_idx == MTS_SKIP);
@@ -623,15 +626,17 @@ static double cu_rd_cost_tr_split_accurate(
     }
     else {
       int split_type = pred_cu->intra.isp_mode;
-      int part_dim = uvg_get_isp_split_dim(width, height, split_type);
+      int split_limit = split_type == ISP_MODE_NO_ISP ? 1 : uvg_get_isp_split_num(width, height, split_type);
-      int limit = split_type == ISP_MODE_HOR ? height : width;
-      int split_num = 0;
-      for (int part = 0; part < limit; part += part_dim) {
-        const int part_x = split_type == ISP_MODE_HOR ? x_px : x_px + part;
-        const int part_y = split_type == ISP_MODE_HOR ? y_px + part : y_px;
-        const int part_w = split_type == ISP_MODE_HOR ? width : part_dim;
-        const int part_h = split_type == ISP_MODE_HOR ? part_dim : height;
 
+      for (int i = 0; i < split_limit; ++i) {
+        cu_loc_t loc;
+        uvg_get_isp_split_loc(&loc, x_px, y_px, width, height, i, split_type);
+        const int part_x = loc.x;
+        const int part_y = loc.y;
+        const int part_w = loc.width;
+        const int part_h = loc.height;
+
+        // TODO: maybe just pass the cu_loc_t to these functions
         const coeff_t* coeffs = &lcu->coeff.y[xy_to_zorder(LCU_WIDTH, part_x, part_y)];
 
         coeff_bits += uvg_get_coeff_cost(state, coeffs, NULL, part_w, part_h, 0, luma_scan_mode, pred_cu->tr_idx == MTS_SKIP);
@@ -858,6 +863,7 @@ static double search_cu(
   const encoder_control_t* ctrl = state->encoder_control;
   const videoframe_t * const frame = state->tile->frame;
   const int cu_width = tree_type != UVG_CHROMA_T ? LCU_WIDTH >> depth : LCU_WIDTH_C >> depth;
+  const int cu_height = cu_width; // TODO: height
   const int luma_width = LCU_WIDTH >> depth;
   assert(cu_width >= 4);
   double cost = MAX_DOUBLE;
@@ -1122,6 +1128,21 @@ static double search_cu(
                          depth, &intra_search,
                          NULL, 
                          lcu, tree_type,recon_luma,recon_chroma);
+      // Set isp split cbfs here
+      const int split_type = intra_search.pred_cu.intra.isp_mode;
+      const int split_num = split_type == ISP_MODE_NO_ISP ? 1 : uvg_get_isp_split_num(cu_width, cu_height, split_type);
+      for (int i = 0; i < split_num; ++i) {
+        cu_loc_t isp_loc;
+        uvg_get_isp_split_loc(&isp_loc, x, y, cu_width, cu_height, i, split_type);
+        //search_data->best_isp_cbfs |= cbf_is_set(cur_cu->cbf, depth, COLOR_Y) << (i++);
+        cu_info_t* split_cu = LCU_GET_CU_AT_PX(lcu, isp_loc.x, isp_loc.y);
+        bool cur_cbf = (intra_search.best_isp_cbfs >> i) & 1;
+        cbf_clear(&split_cu->cbf, depth, COLOR_Y);
+        if (cur_cbf) {
+          cbf_set(&split_cu->cbf, depth, COLOR_Y);
+        }
+      }
+
       if(depth == 4 && x % 8 && y % 8 && tree_type != UVG_LUMA_T && state->encoder_control->chroma_format != UVG_CSP_400) {
         intra_search.pred_cu.intra.mode_chroma = cur_cu->intra.mode_chroma;
         uvg_intra_recon_cu(state,