[DepQuant] Isp and chroma

src/cu.h

@@ -169,6 +169,8 @@ typedef struct
       int8_t mip_flag;
       int8_t mip_is_transposed;
       int8_t isp_mode;
+      uint8_t isp_cbfs : 4;
+      uint8_t isp_index : 2;
     } intra;
     struct {
       mv_t    mv[2][2];  // \brief Motion vectors for L0 and L1

src/dep_quant.c

@@ -440,7 +440,6 @@ static void xSetLastCoeffOffset(
   const int width,
   const int height,
   rate_estimator* rate_estimator,
-  const bool cb_cbf,
   const color_t compID)
 {
   int32_t cbfDeltaBits = 0;
@@ -451,25 +450,18 @@ static void xSetLastCoeffOffset(
     bool        lastCbfIsInferred     = false;
     bool        useIntraSubPartitions = cur_tu->type == CU_INTRA && cur_tu->intra.isp_mode && compID == COLOR_Y;
     if (useIntraSubPartitions) {
-      bool     rootCbfSoFar       = false;
-      bool     isLastSubPartition = false; //TODO: isp check
       uint32_t nTus = uvg_get_isp_split_num(width, height, cur_tu->intra.isp_mode, true);
+      bool     isLastSubPartition = cur_tu->intra.isp_index +1 == nTus; //TODO: isp check
       if (isLastSubPartition) {
-        //TransformUnit* tuPointer = tu.cu->firstTU;
+        lastCbfIsInferred = cur_tu->intra.isp_cbfs == 0;
-        //for (int tuIdx = 0; tuIdx < nTus - 1; tuIdx++) {
-        //  rootCbfSoFar |= TU::getCbfAtDepth(*tuPointer, COMPONENT_Y, tu.depth);
-        //  tuPointer = tuPointer->next;
-        //}
-        if (!rootCbfSoFar) {
-          lastCbfIsInferred = true;
-        }
       }
       if (!lastCbfIsInferred) {
-        prevLumaCbf = false;
+        prevLumaCbf = cur_tu->intra.isp_index != 0 && (cur_tu->intra.isp_cbfs & (1 << (cur_tu->intra.isp_index - 1)));
       }
       const cabac_ctx_t * const cbf_ctx = &state->search_cabac.ctx.qt_cbf_model_luma[2 + prevLumaCbf];
       cbfDeltaBits = lastCbfIsInferred ? 0 : (int32_t)CTX_ENTROPY_BITS(cbf_ctx, 1) - (int32_t)CTX_ENTROPY_BITS(cbf_ctx, 0);
-    } else {
+    }
+    else {
       const cabac_ctx_t* cbf_ctx;
       switch (compID) {
         case COLOR_Y:
@@ -479,7 +471,7 @@ static void xSetLastCoeffOffset(
           cbf_ctx = &state->search_cabac.ctx.qt_cbf_model_cb[0];
           break;
         case COLOR_V:
-          cbf_ctx = &state->search_cabac.ctx.qt_cbf_model_cr[cb_cbf];
+          cbf_ctx = &state->search_cabac.ctx.qt_cbf_model_cr[cbf_is_set(cur_tu->cbf, COLOR_U)];
           break;
       }
       cbfDeltaBits = (int32_t)CTX_ENTROPY_BITS(cbf_ctx, 1) - (int32_t)CTX_ENTROPY_BITS(cbf_ctx, 0);
@@ -1182,7 +1174,7 @@ int uvg_dep_quant(
   int             effWidth = width, effHeight = height;
   if (
     (is_mts ||
-     (state->encoder_control->cfg.mts && 1 /*sbt not used by block*/ &&
+     (state->encoder_control->cfg.mts && 0 /*sbt used by block*/ &&
       height <= 32 && width <= 32)) &&
     compID == COLOR_Y) {
     effHeight = (height == 32) ? 16 : height;
@@ -1212,7 +1204,7 @@ int uvg_dep_quant(
   //===== real init =====
   rate_estimator rate_estimator;
   init_rate_esimator(&rate_estimator, &state->search_cabac, compID);
-  xSetLastCoeffOffset(state, cur_tu, width, height, &rate_estimator, cbf_is_set(cur_tu->cbf, COLOR_U), compID);
+  xSetLastCoeffOffset(state, cur_tu, width, height, &rate_estimator, compID);
 
   reset_common_context(&dep_quant_context.m_common_context, &rate_estimator, (width * height) >> 4, numCoeff);
   dep_quant_context.m_common_context.m_nbInfo = encoder->m_scanId2NbInfoOutArray[log2_tr_width][log2_tr_height];

src/intra.c

@@ -1941,7 +1941,7 @@ void uvg_intra_recon_cu(
       uvg_get_isp_split_loc(&tu_loc,  cu_loc->x, cu_loc->y, width, height, i, split_type, true);
       cu_loc_t pu_loc;
       uvg_get_isp_split_loc(&pu_loc, cu_loc->x, cu_loc->y, width, height, i, split_type, false);
-
+      cur_cu->intra.isp_index = 0;
       if(tu_loc.x % 4 == 0) {
         intra_recon_tb_leaf(state, &pu_loc, cu_loc, lcu, COLOR_Y, search_data, tree_type);
       }
@@ -1949,6 +1949,7 @@ void uvg_intra_recon_cu(
         &tu_loc, cur_cu, lcu,
         false, tree_type);
       search_data->best_isp_cbfs |= cbf_is_set(cur_cu->cbf, COLOR_Y) << i;
+      cur_cu->intra.isp_cbfs = search_data->best_isp_cbfs;
     }
   }
   const bool has_luma = recon_luma && search_data->pred_cu.intra.isp_mode == ISP_MODE_NO_ISP;

src/transform.c

@@ -417,6 +417,7 @@ static void generate_jccr_transforms(
 
 static void quantize_chroma(
   encoder_state_t* const state,
+  cu_info_t * const cur_tu,
   int8_t width,
   int8_t height,
   coeff_t u_coeff[5120],
@@ -427,8 +428,48 @@ static void quantize_chroma(
   const coeff_scan_order_t scan_order,
   bool* u_has_coeffs,
   bool* v_has_coeffs,
-  uint8_t lfnst_idx)
+  uint8_t lfnst_idx, 
+  enum uvg_tree_type tree_type)
 {
+  if(state->encoder_control->cfg.dep_quant && transform != CHROMA_TS) {
+    int abs_sum = 0;
+    uvg_dep_quant(
+      state,
+      cur_tu,
+      width,
+      height,
+      u_coeff,
+      u_quant_coeff,
+      COLOR_U,
+      tree_type,
+      &abs_sum,
+      state->encoder_control->cfg.scaling_list
+    );
+    if (abs_sum > 0) {
+      *u_has_coeffs = 1;
+      cbf_set(&cur_tu->cbf, COLOR_U);
+    }
+    if (transform == DCT7_CHROMA) {
+      abs_sum = 0;
+      uvg_dep_quant(
+        state,
+        cur_tu,
+        width,
+        height,
+        v_coeff,
+        v_quant_coeff,
+        COLOR_V,
+        tree_type,
+        &abs_sum,
+        state->encoder_control->cfg.scaling_list
+      );
+      if (abs_sum > 0) {
+        *v_has_coeffs = 1;
+      }
+      cbf_clear(&cur_tu->cbf, COLOR_U);
+    }
+    return;
+  }
   if (state->encoder_control->cfg.rdoq_enable &&
     (transform != CHROMA_TS || !state->encoder_control->cfg.rdoq_skip))
   {
@@ -561,6 +602,7 @@ void uvg_chroma_transform_search(
     }
     quantize_chroma(
       state,
+      pred_cu,
       width,
       height,
       &u_coeff[i * trans_offset],
@@ -570,8 +612,7 @@ void uvg_chroma_transform_search(
       v_quant_coeff,
       SCAN_DIAG,
       &u_has_coeffs,
-      &v_has_coeffs,
+      &v_has_coeffs, tree_type == UVG_CHROMA_T ?  pred_cu->cr_lfnst_idx : pred_cu->lfnst_idx, tree_type);
-      tree_type == UVG_CHROMA_T ?  pred_cu->cr_lfnst_idx : pred_cu->lfnst_idx);
     if(pred_cu->cr_lfnst_idx !=0 && !u_has_coeffs && !v_has_coeffs) continue;
     
     if(pred_cu->type == CU_INTRA && transforms[i] != CHROMA_TS && tree_type == UVG_CHROMA_T) {