[isp] Implement coefficient encoding for isp splits. Make get_split_dim non static, it was needed elsewhere after all.

This commit is contained in:
siivonek 2022-08-17 16:42:22 +03:00
parent 27a735ab5a
commit 6a9ca94d25
2 changed files with 40 additions and 19 deletions

View file

@ -525,7 +525,8 @@ static void encode_transform_unit(
int depth,
bool only_chroma,
lcu_coeff_t* coeff,
enum uvg_tree_type tree_type)
enum uvg_tree_type tree_type,
bool last_split)
{
assert(depth >= 1 && depth <= MAX_PU_DEPTH);
@ -586,7 +587,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) {
if ((chroma_cbf_set || joint_chroma) && last_split) {
//Need to drop const to get lfnst constraints
encode_chroma_tu(state, x, y, depth, width_c, (cu_info_t*)cur_pu, &scan_idx, coeff, joint_chroma, tree_type);
}
@ -611,7 +612,8 @@ static void encode_transform_coeff(
uint8_t parent_coeff_v,
bool only_chroma,
lcu_coeff_t* coeff,
enum uvg_tree_type tree_type)
enum uvg_tree_type tree_type,
bool last_split) // Always true except when writing sub partition coeffs (ISP)
{
cabac_data_t * const cabac = &state->cabac;
//const encoder_control_t *const ctrl = state->encoder_control;
@ -642,8 +644,6 @@ static void encode_transform_coeff(
int8_t split = (LCU_WIDTH >> depth > TR_MAX_WIDTH);
const int cb_flag_y = tree_type != UVG_CHROMA_T ? cbf_is_set(cur_pu->cbf, depth, COLOR_Y) : 0;
const int cb_flag_u = tree_type != UVG_LUMA_T ?( cur_pu->joint_cb_cr ? (cur_pu->joint_cb_cr >> 1) & 1 : cbf_is_set(cur_cu->cbf, depth, COLOR_U)) : 0;
const int cb_flag_v = tree_type != UVG_LUMA_T ? (cur_pu->joint_cb_cr ? cur_pu->joint_cb_cr & 1 : cbf_is_set(cur_cu->cbf, depth, COLOR_V)) : 0;
@ -671,7 +671,7 @@ static void encode_transform_coeff(
// - they have already been signaled to 0 previously
// When they are not present they are inferred to be 0, except for size 4
// when the flags from previous level are used.
if (state->encoder_control->chroma_format != UVG_CSP_400 && (depth != 4 || only_chroma) && tree_type != UVG_LUMA_T) {
if (state->encoder_control->chroma_format != UVG_CSP_400 && (depth != 4 || only_chroma) && tree_type != UVG_LUMA_T && last_split) {
if (!split) {
if (true) {
@ -690,10 +690,10 @@ static void encode_transform_coeff(
uint8_t offset = LCU_WIDTH >> (depth + 1);
int x2 = x + offset;
int y2 = y + offset;
encode_transform_coeff(state, x, y, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type);
encode_transform_coeff(state, x2, y, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type);
encode_transform_coeff(state, x, y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type);
encode_transform_coeff(state, x2, y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type);
encode_transform_coeff(state, x, y, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true);
encode_transform_coeff(state, x2, y, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true);
encode_transform_coeff(state, x, y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true);
encode_transform_coeff(state, x2, y2, depth + 1, tr_depth + 1, cb_flag_u, cb_flag_v, only_chroma, coeff, tree_type, true);
return;
}
@ -737,12 +737,13 @@ static void encode_transform_coeff(
|| (cb_flag_u && cb_flag_v))
&& (depth != 4 || only_chroma || tree_type == UVG_CHROMA_T)
&& state->encoder_control->cfg.jccr
&& last_split
) {
assert(cur_pu->joint_cb_cr < 4 && "JointCbCr is in search state.");
cabac->cur_ctx = &cabac->ctx.joint_cb_cr[cb_flag_u * 2 + cb_flag_v - 1];
CABAC_BIN(cabac, cur_pu->joint_cb_cr != 0, "tu_joint_cbcr_residual_flag");
}
encode_transform_unit(state, x, y, depth, only_chroma, coeff, tree_type);
encode_transform_unit(state, x, y, depth, only_chroma, coeff, tree_type, last_split);
}
}
@ -1671,7 +1672,7 @@ void uvg_encode_coding_tree(
// Code (possible) coeffs to bitstream
if (cbf) {
encode_transform_coeff(state, x, y, depth, 0, 0, 0, 0, coeff, tree_type);
encode_transform_coeff(state, x, y, depth, 0, 0, 0, 0, coeff, tree_type, true);
}
encode_mts_idx(state, cabac, cur_cu);
@ -1688,7 +1689,25 @@ void uvg_encode_coding_tree(
}
if (tree_type != UVG_CHROMA_T) {
encode_transform_coeff(state, x, y, depth, 0, 0, 0, 0, coeff, tree_type);
// Cycle through sub partitions if ISP enabled.
// 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 part_dim = cu_width;
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;
// Check if last split to write chroma
bool last_split = (part + part_dim) == limit;
encode_transform_coeff(state, part_x, part_y, depth, 0, 0, 0, 0, coeff, tree_type, last_split);
}
}
if (tree_type != UVG_CHROMA_T) {
@ -1706,7 +1725,7 @@ void uvg_encode_coding_tree(
tmp->violates_lfnst_constrained_luma = false;
tmp->violates_lfnst_constrained_chroma = false;
tmp->lfnst_last_scan_pos = false;
encode_transform_coeff(state, x, y, depth, 0, 0, 0, 1, coeff, tree_type);
encode_transform_coeff(state, x, y, depth, 0, 0, 0, 1, coeff, tree_type, true);
// Write LFNST only once for single tree structure
encode_lfnst_idx(state, cabac, tmp, x, y, depth, cu_width, cu_height, tree_type, COLOR_UV);
}

View file

@ -1481,8 +1481,10 @@ const cu_info_t* uvg_get_co_located_luma_cu(
* \param height Block height.
* \param split_type Horizontal or vertical split.
*/
static int get_isp_split_dim(const int width, const int height, const int split_type)
int uvg_get_isp_split_dim(const int width, const int height, const int split_type)
{
assert(split_type != ISP_MODE_NO_ISP && "Cannot calculate split dimension if no split type is set. Make sure this function is not called in this case.");
bool divide_in_rows = split_type == SPLIT_TYPE_HOR;
int split_dim_size, non_split_dim_size, partition_size, div_shift = 2;
@ -1666,9 +1668,9 @@ 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 = get_isp_split_dim(width, height, split_type);
int part_dim = uvg_get_isp_split_dim(width, height, split_type);
int limit = split_type == ISP_MODE_HOR ? height : width;
for (int part = 0; part < limit; part + part_dim) {
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 ? part_dim : width;
@ -1749,8 +1751,8 @@ bool uvg_can_use_isp_with_lfnst(const int width, const int height, const int isp
return false;
}
const int tu_width = (isp_split_type == ISP_MODE_HOR) ? width : get_isp_split_dim(width, height, SPLIT_TYPE_VER);
const int tu_height = (isp_split_type == ISP_MODE_HOR) ? get_isp_split_dim(width, height, SPLIT_TYPE_HOR) : height;
const int tu_width = (isp_split_type == ISP_MODE_HOR) ? width : uvg_get_isp_split_dim(width, height, SPLIT_TYPE_VER);
const int tu_height = (isp_split_type == ISP_MODE_HOR) ? uvg_get_isp_split_dim(width, height, SPLIT_TYPE_HOR) : height;
if (!(tu_width >= TR_MIN_WIDTH && tu_height >= TR_MIN_WIDTH))
{