[tr-skip] Add cmdline option for chroma transform skip

src/cfg.c

@@ -76,6 +76,7 @@ int uvg_config_init(uvg_config *cfg)
   cfg->mv_rdo          = 0;
   cfg->full_intra_search = 0;
   cfg->trskip_enable   = 0;
+  cfg->chroma_trskip_enable = 0;
   cfg->trskip_max_size = 2; //Default to 4x4
   cfg->mts             = 0;
   cfg->mts_implicit    = 0;
@@ -906,6 +907,8 @@ int uvg_config_parse(uvg_config *cfg, const char *name, const char *value)
     cfg->full_intra_search = atobool(value);
   else if OPT("transform-skip")
     cfg->trskip_enable = atobool(value);
+  else if OPT("chroma-transform-skip")
+    cfg->chroma_trskip_enable = atobool(value);
   else if OPT("tr-skip-max-size") {
     cfg->trskip_max_size = atoi(value);
     if (cfg->trskip_max_size < 2 || cfg->trskip_max_size > 5) {
@@ -1825,6 +1828,11 @@ int uvg_config_validate(const uvg_config *const cfg)
     error = 1;
   }
 
+  if(cfg->chroma_trskip_enable && !cfg->trskip_enable) {
+    fprintf(stderr, "Transform skip has to be enabled when chroma transform skip is enabled.\n");
+    error = 1;
+  }
+
   return !error;
 }
 

src/cli.c

@@ -71,6 +71,8 @@ static const struct option long_options[] = {
   { "no-full-intra-search",     no_argument, NULL, 0 },
   { "transform-skip",           no_argument, NULL, 0 },
   { "no-transform-skip",        no_argument, NULL, 0 },
+  { "chroma-transform-skip",    no_argument, NULL, 0 },
+  { "no-chroma-transform-skip", no_argument, NULL, 0 },
   { "tr-skip-max-size",   required_argument, NULL, 0 },
   { "mts",                required_argument, NULL, 0 },
   { "no-mts",                   no_argument, NULL, 0 },
@@ -577,6 +579,8 @@ void print_help(void)
     "      --(no-)full-intra-search : Try all intra modes during rough search.\n"
     "                               [disabled]\n"
     "      --(no-)transform-skip  : Try transform skip [disabled]\n"
+    "      --(no-)chroma-transform-skip : Try transform skip for chroma \n"
+    "                                     blocks. [disabled]\n"
     "      --tr-skip-max-size     : Max log2 size of transform skip 2..5 [2]\n"
     "      --me <string>          : Integer motion estimation algorithm [hexbs]\n"
     "                                   - hexbs: Hexagon Based Search\n"

src/search.c

@@ -540,6 +540,8 @@ static double cu_rd_cost_tr_split_accurate(const encoder_state_t* const state,
       LCU_WIDTH, LCU_WIDTH,
       width);
   }
+  // Chroma transform skip enable/disable is non-normative, so we need to count the chroma
+  // tr-skip bits even when we are never using it.
   const bool can_use_tr_skip = state->encoder_control->cfg.trskip_enable && width <= (1 << state->encoder_control->cfg.trskip_max_size);
 
   if(cb_flag_y){

src/search_intra.c

@@ -1634,7 +1634,8 @@ int8_t uvg_search_intra_chroma_rdo(
         const int trans_offset = width * height;
         int num_transforms = 1;
         const int can_use_tr_skip = state->encoder_control->cfg.trskip_enable &&
-          (1 << state->encoder_control->cfg.trskip_max_size) >= width;
+          (1 << state->encoder_control->cfg.trskip_max_size) >= width && 
+          state->encoder_control->cfg.chroma_trskip_enable;
         if(can_use_tr_skip) {
           uvg_transformskip(state->encoder_control, u_resi, u_coeff + num_transforms * trans_offset, width);
           uvg_transformskip(state->encoder_control, v_resi, v_coeff + num_transforms * trans_offset, width);
@@ -1777,7 +1778,7 @@ int8_t uvg_search_intra_chroma_rdo(
           }
 
           if (cbf_u || (transforms[i] == JCCR_1 && u_has_coeffs)) {
-            if(can_use_tr_skip && !IS_JCCR_MODE(transforms[i])) {
+            if(can_use_tr_skip) {
               CABAC_FBITS_UPDATE(&state->search_cabac, &state->search_cabac.ctx.transform_skip_model_chroma,
                 transforms[i] == CHROMA_TS, u_bits, "tr_skip_u"
               );

src/uvg266.h

@@ -334,6 +334,7 @@ typedef struct uvg_config
   int32_t rdo;            /*!< \brief RD-calculation level (0..2) */
   int32_t full_intra_search; /*!< \brief If true, don't skip modes in intra search. */
   int32_t trskip_enable;    /*!< \brief Flag to enable transform skip. */
+  int32_t chroma_trskip_enable; /*!< \brief Flag to enable transform skip for chroma blocks. */
   int32_t trskip_max_size;    /*!< \brief Transform skip max block size. */
   enum uvg_mts mts;        /*< \brief flag to enable multiple transform selection*/
   int32_t mts_implicit;        /*< \brief flag to enable implicit multiple transform selection*/