[dual-tree] Still not working but bitstream valid

This commit is contained in:
Joose Sainio 2022-06-08 13:22:29 +03:00
parent 0adb0846d2
commit abd00d04a1
6 changed files with 107 additions and 59 deletions

View file

@ -511,9 +511,21 @@ void uvg_encode_last_significant_xy(cabac_data_t * const cabac,
if (cabac->only_count && bits_out) *bits_out += bits;
}
static void encode_chroma_tu(encoder_state_t* const state, int x, int y, int depth, const uint8_t width_c, cu_info_t* cur_pu, int8_t* scan_idx, lcu_coeff_t* coeff, uint8_t joint_chroma) {
int x_local = ((x & ~7) >> 1) % LCU_WIDTH_C;
int y_local = ((y & ~7) >> 1) % LCU_WIDTH_C;
static void encode_chroma_tu(
encoder_state_t* const state,
int x,
int y,
int depth,
const uint8_t width_c,
const cu_info_t* cur_pu,
int8_t* scan_idx,
lcu_coeff_t* coeff,
uint8_t joint_chroma,
enum
kvz_tree_type tree_type)
{
int x_local = ((x >> (tree_type != KVZ_CHROMA_T)) & ~3) % LCU_WIDTH_C;
int y_local = ((y >> (tree_type != KVZ_CHROMA_T)) & ~3) % LCU_WIDTH_C;
cabac_data_t* const cabac = &state->cabac;
*scan_idx = uvg_get_scan_order(cur_pu->type, cur_pu->intra.mode_chroma, depth);
if(!joint_chroma){
@ -618,7 +630,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) {
encode_chroma_tu(state, x, y, depth, width_c, cur_pu, &scan_idx, coeff, joint_chroma);
encode_chroma_tu(state, x, y, depth, width_c, cur_pu, &scan_idx, coeff, joint_chroma, tree_type);
}
}
@ -1488,13 +1500,13 @@ void uvg_encode_coding_tree(
uint16_t abs_x = x + (state->tile->offset_x >> (tree_type == KVZ_CHROMA_T));
uint16_t abs_y = y + (state->tile->offset_y >> (tree_type == KVZ_CHROMA_T));
int32_t used_width = tree_type != KVZ_CHROMA_T ? ctrl->in.width : ctrl->in.width / 2;
int32_t used_height = tree_type != KVZ_CHROMA_T ? ctrl->in.height : ctrl->in.height / 2;
int32_t frame_width = tree_type != KVZ_CHROMA_T ? ctrl->in.width : ctrl->in.width / 2;
int32_t frame_height = tree_type != KVZ_CHROMA_T ? ctrl->in.height : ctrl->in.height / 2;
// Check for slice border
bool border_x = used_width < abs_x + cu_width;
bool border_y = used_height < abs_y + cu_width;
bool border_split_x = used_width >= abs_x + (LCU_WIDTH >> MAX_DEPTH) + half_cu;
bool border_split_y = used_height >= abs_y + (LCU_WIDTH >> MAX_DEPTH) + half_cu;
bool border_x = frame_width < abs_x + cu_width;
bool border_y = frame_height < abs_y + cu_width;
bool border_split_x = frame_width >= abs_x + (LCU_WIDTH >> MAX_DEPTH) + half_cu;
bool border_split_y = frame_height >= abs_y + (LCU_WIDTH >> MAX_DEPTH) + half_cu;
bool border = border_x || border_y; /*!< are we in any border CU */
if (depth <= state->frame->max_qp_delta_depth) {

View file

@ -520,6 +520,7 @@ static void predict_cclm(
if (y0) {
for (; available_above_right < width / 2; available_above_right++) {
int x_extension = x_scu + width * 2 + 4 * available_above_right;
// TODO: Dual tree
const cu_info_t* pu = LCU_GET_CU_AT_PX(lcu, x_extension, y_scu - 4);
if (x_extension >= LCU_WIDTH || pu->type == CU_NOTSET) break;
}
@ -1524,9 +1525,10 @@ void uvg_intra_recon_cu(
int depth,
intra_search_data_t* search_data,
cu_info_t *cur_cu,
lcu_t *lcu)
lcu_t *lcu,
enum kvz_tree_type tree_type)
{
const vector2d_t lcu_px = { SUB_SCU(x), SUB_SCU(y) };
const vector2d_t lcu_px = { SUB_SCU(x) >> (tree_type == KVZ_CHROMA_T), SUB_SCU(y) >> (tree_type == KVZ_CHROMA_T) };
const int8_t width = LCU_WIDTH >> depth;
if (cur_cu == NULL) {
cur_cu = LCU_GET_CU_AT_PX(lcu, lcu_px.x, lcu_px.y);
@ -1555,16 +1557,16 @@ void uvg_intra_recon_cu(
const int32_t x2 = x + offset;
const int32_t y2 = y + offset;
uvg_intra_recon_cu(state, x, y, depth + 1, search_data, NULL, lcu);
uvg_intra_recon_cu(state, x2, y, depth + 1, search_data, NULL, lcu);
uvg_intra_recon_cu(state, x, y2, depth + 1, search_data, NULL, lcu);
uvg_intra_recon_cu(state, x2, y2, depth + 1, search_data, NULL, lcu);
uvg_intra_recon_cu(state, x, y, depth + 1, search_data, NULL, lcu, tree_type);
uvg_intra_recon_cu(state, x2, y, depth + 1, search_data, NULL, lcu, tree_type);
uvg_intra_recon_cu(state, x, y2, depth + 1, search_data, NULL, lcu, tree_type);
uvg_intra_recon_cu(state, x2, y2, depth + 1, search_data, NULL, lcu, tree_type);
// Propagate coded block flags from child CUs to parent CU.
uint16_t child_cbfs[3] = {
LCU_GET_CU_AT_PX(lcu, lcu_px.x + offset, lcu_px.y )->cbf,
LCU_GET_CU_AT_PX(lcu, lcu_px.x, lcu_px.y + offset)->cbf,
LCU_GET_CU_AT_PX(lcu, lcu_px.x + offset, lcu_px.y + offset)->cbf,
LCU_GET_CU_AT_PX(lcu, (lcu_px.x + offset) >> (tree_type == KVZ_CHROMA_T), lcu_px.y >> (tree_type == KVZ_CHROMA_T))->cbf,
LCU_GET_CU_AT_PX(lcu, lcu_px.x >> (tree_type == KVZ_CHROMA_T), (lcu_px.y + offset) >> (tree_type == KVZ_CHROMA_T))->cbf,
LCU_GET_CU_AT_PX(lcu, (lcu_px.x + offset) >> (tree_type == KVZ_CHROMA_T), (lcu_px.y + offset) >> (tree_type == KVZ_CHROMA_T))->cbf,
};
if (mode_luma != -1 && depth <= MAX_DEPTH) {

View file

@ -143,6 +143,7 @@ void uvg_intra_recon_cu(
int depth,
intra_search_data_t* search_data,
cu_info_t *cur_cu,
lcu_t *lcu);
lcu_t *lcu,
enum kvz_tree_type tree_type);
int uvg_get_mip_flag_context(int x, int y, int width, int height, const lcu_t* lcu, cu_array_t* const cu_a);

View file

@ -112,7 +112,7 @@ static INLINE void copy_cu_coeffs(int x_local, int y_local, int width, lcu_t *fr
static void work_tree_copy_up(int x_local, int y_local, int depth, lcu_t *work_tree, bool joint, enum
kvz_tree_type tree_type)
{
const int width = tree_type != KVZ_CHROMA_T ? LCU_WIDTH >> depth : LCU_WIDTH_C >> depth;
const int width = LCU_WIDTH >> depth;
copy_cu_info (x_local, y_local, width, &work_tree[depth + 1], &work_tree[depth]);
copy_cu_pixels(x_local, y_local, width, &work_tree[depth + 1], &work_tree[depth], tree_type);
copy_cu_coeffs(x_local, y_local, width, &work_tree[depth + 1], &work_tree[depth], joint, tree_type);
@ -808,8 +808,8 @@ static double search_cu(
lcu_t *const lcu = &work_tree[depth];
int x_local = SUB_SCU_TREE(x, tree_type);
int y_local = SUB_SCU_TREE(y, tree_type);
int x_local = SUB_SCU(x) >> (tree_type == KVZ_CHROMA_T);
int y_local = SUB_SCU(y) >> (tree_type == KVZ_CHROMA_T);
int32_t frame_width = tree_type != KVZ_CHROMA_T ? frame->width : frame->width / 2;
int32_t frame_height = tree_type != KVZ_CHROMA_T ? frame->height: frame->height / 2;
@ -923,13 +923,13 @@ static double search_cu(
double intra_cost = intra_search.cost;
if (intra_cost < cost || tree_type == KVZ_CHROMA_T) {
int8_t intra_mode = intra_search.pred_cu.intra.mode;
if(state->encoder_control->cfg.cclm && tree_type != KVZ_CHROMA_T) {
if(state->encoder_control->cfg.cclm && tree_type == KVZ_BOTH_T) {
intra_search.pred_cu.intra.mode_chroma = -1;
uvg_intra_recon_cu(state,
x, y,
depth, &intra_search,
&intra_search.pred_cu,
lcu);
lcu, tree_type);
downsample_cclm_rec(
state, x, y, cu_width / 2, cu_width / 2, lcu->rec.y, lcu->left_ref.y[64]
@ -945,7 +945,7 @@ static double search_cu(
// really bad for chroma.
intra_search.pred_cu.intra.mode_chroma = intra_search.pred_cu.intra.mode;
if (ctrl->cfg.rdo >= 3 || ctrl->cfg.jccr || ctrl->cfg.lfnst) {
uvg_search_cu_intra_chroma(state, x, y, depth, lcu, &intra_search);
uvg_search_cu_intra_chroma(state, x, y, depth, lcu, &intra_search, tree_type);
if (intra_search.pred_cu.joint_cb_cr == 0) {
intra_search.pred_cu.joint_cb_cr = 4;
@ -962,7 +962,8 @@ static double search_cu(
x, y,
depth, &intra_search,
&intra_search.pred_cu,
lcu);
lcu,
tree_type);
intra_cost += uvg_cu_rd_cost_chroma(state, x_local, y_local, depth, &intra_search.pred_cu, lcu);
intra_search.pred_cu.intra.mode = intra_mode;
intra_search.pred_cu.violates_lfnst_constrained_chroma = false;
@ -987,12 +988,15 @@ static double search_cu(
if ((depth == 4 && (x % 8 == 0 || y % 8 == 0)) || state->encoder_control->chroma_format == UVG_CSP_400 || tree_type == KVZ_LUMA_T) {
intra_search.pred_cu.intra.mode_chroma = -1;
}
if(tree_type == KVZ_CHROMA_T) {
intra_search.pred_cu.intra.mode = -1;
}
lcu_fill_cu_info(lcu, x_local, y_local, cu_width, cu_width, cur_cu);
uvg_intra_recon_cu(state,
x, y,
depth, &intra_search,
NULL,
lcu);
lcu, tree_type);
if(depth == 4 && x % 8 && y % 8) {
intra_search.pred_cu.intra.mode_chroma = cur_cu->intra.mode_chroma;
intra_search.pred_cu.intra.mode = -1;
@ -1000,7 +1004,8 @@ static double search_cu(
x, y,
depth, &intra_search,
NULL,
lcu);
lcu,
tree_type);
}
if (cur_cu->joint_cb_cr == 4) cur_cu->joint_cb_cr = 0;
lcu_fill_cu_info(lcu, x_local, y_local, cu_width, cu_width, cur_cu);
@ -1117,7 +1122,7 @@ static double search_cu(
// Recursively split all the way to max search depth.
if (can_split_cu) {
int half_cu = cu_width / 2;
int half_cu = cu_width >> (tree_type != KVZ_CHROMA_T);
double split_cost = 0.0;
int cbf = cbf_is_set_any(cur_cu->cbf, depth);
cabac_data_t post_seach_cabac;
@ -1132,8 +1137,8 @@ static double search_cu(
if (depth < MAX_DEPTH) {
// Add cost of cu_split_flag.
uvg_write_split_flag(state, &state->search_cabac,
x > 0 ? LCU_GET_CU_AT_PX(lcu, SUB_SCU(x) - 1, SUB_SCU(y)) : NULL,
y > 0 ? LCU_GET_CU_AT_PX(lcu, SUB_SCU(x), SUB_SCU(y) - 1) : NULL,
x > 0 ? LCU_GET_CU_AT_PX(lcu, (SUB_SCU(x) - 1) >> (tree_type == KVZ_CHROMA_T), SUB_SCU(y) >> (tree_type == KVZ_CHROMA_T)) : NULL,
y > 0 ? LCU_GET_CU_AT_PX(lcu, SUB_SCU(x) >> (tree_type == KVZ_CHROMA_T), (SUB_SCU(y) - 1) >> (tree_type == KVZ_CHROMA_T)) : NULL,
1, depth, cu_width, x, y, tree_type,
&split_bits);
}
@ -1161,6 +1166,7 @@ static double search_cu(
// gets used, at least in the most obvious cases, while avoiding any
// searching.
// TODO: Dual tree
if (cur_cu->type == CU_NOTSET && depth < MAX_PU_DEPTH
&& x + cu_width <= frame_width && y + cu_width <= frame_height
&& state->encoder_control->cfg.combine_intra_cus)
@ -1200,7 +1206,8 @@ static double search_cu(
depth,
&proxy,
NULL,
lcu);
lcu,
tree_type);
double mode_bits = calc_mode_bits(state, lcu, cur_cu, x, y, depth) + bits;
cost += mode_bits * state->lambda;
@ -1476,8 +1483,8 @@ void uvg_search_lcu(encoder_state_t * const state, const int x, const int y, con
if(state->frame->slicetype == UVG_SLICE_I && state->encoder_control->cfg.dual_tree) {
cost = search_cu(
state,
x / 2,
y / 2,
x,
y,
0,
work_tree,
KVZ_CHROMA_T);

View file

@ -375,14 +375,14 @@ static double search_intra_trdepth(
if (pred_cu->lfnst_idx > 0 && pred_cu->tr_idx > 0) {
continue;
}
uvg_intra_recon_cu(
state,
x_px,
y_px,
depth,
search_data,
uvg_intra_recon_cu(state,
x_px, y_px,
depth, search_data,
pred_cu,
lcu);
lcu,
KVZ_LUMA_T);
if (trafo != 0 && !cbf_is_set(pred_cu->cbf, depth, COLOR_Y)) continue;
// TODO: Not sure if this should be 0 or 1 but at least seems to work with 1
@ -473,7 +473,8 @@ static double search_intra_trdepth(
depth,
search_data,
pred_cu,
lcu);
lcu,
KVZ_BOTH_T);
best_rd_cost += uvg_cu_rd_cost_chroma(
state,
lcu_px.x,
@ -515,9 +516,21 @@ static double search_intra_trdepth(
}
}
if (best_tr_idx == MTS_SKIP) break;
// Very unlikely that further search is necessary if skip seems best option
} // end lfnst_index loop
}
if(reconstruct_chroma) {
int8_t luma_mode = pred_cu->intra.mode;
pred_cu->intra.mode = -1;
pred_cu->intra.mode_chroma = chroma_mode;
pred_cu->joint_cb_cr= 4; // TODO: Maybe check the jccr mode here also but holy shit is the interface of search_intra_rdo bad currently
uvg_intra_recon_cu(state,
x_px, y_px,
depth, search_data,
pred_cu,
lcu,
KVZ_BOTH_T);
best_rd_cost += uvg_cu_rd_cost_chroma(state, lcu_px.x, lcu_px.y, depth, pred_cu, lcu);
pred_cu->intra.mode = luma_mode;
}
pred_cu->tr_skip = best_tr_idx == MTS_SKIP;
pred_cu->tr_idx = best_tr_idx;
pred_cu->lfnst_idx = best_lfnst_idx;
@ -1407,7 +1420,8 @@ int8_t uvg_search_intra_chroma_rdo(
int8_t num_modes,
lcu_t *const lcu,
intra_search_data_t* chroma_data,
int8_t luma_mode)
int8_t luma_mode,
enum kvz_tree_type tree_type)
{
const bool reconstruct_chroma = (depth != 4) || (x_px & 4 && y_px & 4);
@ -1552,11 +1566,13 @@ int8_t uvg_search_intra_chroma_rdo(
x_px, y_px,
depth, &chroma_data[mode_i],
pred_cu,
lcu);
lcu,
tree_type);
chroma_data[mode_i].cost += uvg_cu_rd_cost_chroma(state, lcu_px.x, lcu_px.y, depth, pred_cu, lcu);
memcpy(&state->search_cabac, &temp_cabac, sizeof(cabac_data_t));
}
}
pred_cu->cr_lfnst_idx = best_lfnst_index;
}
sort_modes(chroma_data, num_modes);
@ -1569,9 +1585,14 @@ int8_t uvg_search_intra_chroma_rdo(
#undef IS_JCCR_MODE
int8_t uvg_search_cu_intra_chroma(encoder_state_t * const state,
const int x_px, const int y_px,
const int depth, lcu_t *lcu, intra_search_data_t *search_data)
int8_t uvg_search_cu_intra_chroma(
encoder_state_t * const state,
const int x_px,
const int y_px,
const int depth,
lcu_t *lcu,
intra_search_data_t *search_data,
enum kvz_tree_type tree_type)
{
const vector2d_t lcu_px = { SUB_SCU(x_px), SUB_SCU(y_px) };
@ -1640,7 +1661,7 @@ int8_t uvg_search_cu_intra_chroma(encoder_state_t * const state,
}
if (num_modes > 1) {
uvg_search_intra_chroma_rdo(state, x_px, y_px, depth, num_modes, lcu, chroma_data, intra_mode);
uvg_search_intra_chroma_rdo(state, x_px, y_px, depth, num_modes, lcu, chroma_data, intra_mode, tree_type);
}
*search_data = chroma_data[0];
return chroma_data[0].pred_cu.intra.mode_chroma;

View file

@ -48,9 +48,14 @@ double uvg_luma_mode_bits(const encoder_state_t *state, const cu_info_t* const c
double uvg_chroma_mode_bits(const encoder_state_t *state,
int8_t chroma_mode, int8_t luma_mode);
int8_t uvg_search_cu_intra_chroma(encoder_state_t * const state,
const int x_px, const int y_px,
const int depth, lcu_t *lcu, intra_search_data_t* best_cclm);
int8_t uvg_search_cu_intra_chroma(
encoder_state_t * const state,
const int x_px,
const int y_px,
const int depth,
lcu_t *lcu,
intra_search_data_t* best_cclm,
enum kvz_tree_type tree_type);
void uvg_search_cu_intra(
encoder_state_t * const state,