[mip] Fix mpm mode selection. When neighboring CU uses MIP, signaled intra mode must be planar. Fix chroma reconstruction when MIP is enabled. Only allow MIP to be used if chroma scheme is 444. Otherwise use planar mode.

src/intra.c

@@ -100,15 +100,24 @@ int8_t kvz_intra_get_dir_luma_predictor(
   int8_t number_of_candidates = 0;
 
   // The default mode if block is not coded yet is INTRA_PLANAR.
+  // If the neighboring blocks were MIP blocks, intra mode is set to planar.
   int8_t left_intra_dir  = 0;
   if (left_pu && left_pu->type == CU_INTRA) {
+    if (left_pu->intra.mip_flag) {
+      left_intra_dir = PLANAR_IDX;
+    } else {
       left_intra_dir = left_pu->intra.mode;
     }
+  }
 
   int8_t above_intra_dir = 0;
   if (above_pu && above_pu->type == CU_INTRA && y % LCU_WIDTH != 0) {
+    if (above_pu->intra.mip_flag) {
+      above_intra_dir = PLANAR_IDX;
+    } else {
       above_intra_dir = above_pu->intra.mode;
     }
+  }
 
   const int offset = 61;
   const int mod = 64;
@@ -715,6 +724,7 @@ void kvz_mip_pred_upsampling_1D(int* const dst, const int* const src, const int*
 
 /** \brief Matrix weighted intra prediction.
 */
+// MIP_TODO: remove color parameter if it is not used
 void kvz_mip_predict(encoder_state_t const* const state, kvz_intra_references* const refs,
                      const uint16_t pred_block_width, const uint16_t pred_block_height,
                      const color_t color,
@@ -1347,7 +1357,7 @@ static void intra_recon_tb_leaf(
   lcu_t *lcu,
   color_t color,
   uint8_t multi_ref_idx,
-  bool use_mip,
+  bool mip_flag,
   bool mip_transp)
 {
   const kvz_config *cfg = &state->encoder_control->cfg;
@@ -1395,6 +1405,21 @@ static void intra_recon_tb_leaf(
   kvz_pixel pred[32 * 32];
   int stride = state->tile->frame->source->stride;
   const bool filter_boundary = color == COLOR_Y && !(cfg->lossless && cfg->implicit_rdpcm);
+  bool use_mip = false;
+  if (mip_flag) {
+    if (color == COLOR_Y) {
+      use_mip = true;
+    } else {
+      // MIP can be used for chroma if the chroma scheme is 444
+      if (state->encoder_control->chroma_format == KVZ_CSP_444) {
+        use_mip = true;
+      } else {
+        // If MIP cannot be used for chroma, set mode to planar
+        intra_mode = 0;
+      }
+    }
+  }
+
   if(intra_mode < 68) {
     if (use_mip) {
       assert(intra_mode >= 0 && intra_mode < 16 && "MIP mode must be between [0, 15]");
@@ -1477,6 +1502,12 @@ void kvz_intra_recon_cu(
   bool use_mip = mip_flag;
   bool mip_transposed = mip_transp;
   
+  if (mode_luma != -1 && mode_chroma != -1) {
+    if (use_mip) {
+      assert(mode_luma == mode_chroma && "Chroma mode must be derived from luma mode if block uses MIP.");
+    }
+  }
+
   // Reset CBFs because CBFs might have been set
   // for depth earlier
   if (mode_luma >= 0) {
@@ -1521,8 +1552,8 @@ void kvz_intra_recon_cu(
       intra_recon_tb_leaf(state, x, y, depth, mode_luma, cclm_params, lcu, COLOR_Y, multi_ref_index, use_mip, mip_transposed);
     }
     if (has_chroma) {
-      intra_recon_tb_leaf(state, x, y, depth, mode_chroma, cclm_params, lcu, COLOR_U, 0, false, false);
-      intra_recon_tb_leaf(state, x, y, depth, mode_chroma, cclm_params, lcu, COLOR_V, 0, false, false);
+      intra_recon_tb_leaf(state, x, y, depth, mode_chroma, cclm_params, lcu, COLOR_U, 0, use_mip, mip_transposed);
+      intra_recon_tb_leaf(state, x, y, depth, mode_chroma, cclm_params, lcu, COLOR_V, 0, use_mip, mip_transposed);
     }
 
     kvz_quantize_lcu_residual(state, has_luma, has_chroma, x, y, depth, cur_cu, lcu, false);

src/search.c

@@ -742,10 +742,6 @@ static double search_cu(encoder_state_t * const state, int x, int y, int depth,
         cur_cu->intra.multi_ref_idx = multi_ref_index;
         cur_cu->intra.mip_flag = mip_flag;
         cur_cu->intra.mip_is_transposed = mip_transposed;
-        // If a MIP mode is selected, set chroma mode to planar and skip further chroma search
-        if (mip_flag) {
-          cur_cu->intra.mode_chroma = 0;
-        }
 
         //If the CU is not split from 64x64 block, the MTS is disabled for that CU.
         cur_cu->tr_idx = (depth > 0) ? intra_trafo : 0;
@@ -756,10 +752,8 @@ static double search_cu(encoder_state_t * const state, int x, int y, int depth,
     // mode search of adjacent CUs.
     if (cur_cu->type == CU_INTRA) {
       assert(cur_cu->part_size == SIZE_2Nx2N || cur_cu->part_size == SIZE_NxN);
-      // Chroma mode must be planar if mip_flag is set.
-      if (!cur_cu->intra.mip_flag) {
       cur_cu->intra.mode_chroma = cur_cu->intra.mode;
-      }
+      
       lcu_fill_cu_info(lcu, x_local, y_local, cu_width, cu_width, cur_cu);
       kvz_intra_recon_cu(state,
                          x, y,
@@ -789,7 +783,8 @@ static double search_cu(encoder_state_t * const state, int x, int y, int depth,
                            x & ~7, y & ~7, // TODO: as does this
                            depth,
                            -1, cur_cu->intra.mode_chroma, // skip luma
-                           NULL, cclm_params, 0, false, false,
+                           NULL, cclm_params, 0, 
+                           cur_cu->intra.mip_flag, cur_cu->intra.mip_is_transposed,
                            lcu);
       }
     } else if (cur_cu->type == CU_INTER) {
@@ -925,6 +920,7 @@ static double search_cu(encoder_state_t * const state, int x, int y, int depth,
     // of the top left CU from the next depth. This should ensure that 64x64
     // gets used, at least in the most obvious cases, while avoiding any
     // searching.
+    
     if (cur_cu->type == CU_NOTSET && depth < MAX_PU_DEPTH
         && x + cu_width <= frame->width && y + cu_width <= frame->height && 0)
     {

src/search_intra.c

@@ -364,7 +364,7 @@ static double search_intra_trdepth(encoder_state_t * const state,
         depth,
         -1, chroma_mode,
         pred_cu, cclm_params, 0, 
-        false, false,
+        pred_cu->intra.mip_flag, pred_cu->intra.mip_is_transposed,
         lcu);
       best_rd_cost += kvz_cu_rd_cost_chroma(state, lcu_px.x, lcu_px.y, depth, pred_cu, lcu);
     }