1 // Copyright (c) 2017-2021, The rav1e contributors. All rights reserved 2 // 3 // This source code is subject to the terms of the BSD 2 Clause License and 4 // the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License 5 // was not distributed with this source code in the LICENSE file, you can 6 // obtain it at www.aomedia.org/license/software. If the Alliance for Open 7 // Media Patent License 1.0 was not distributed with this source code in the 8 // PATENTS file, you can obtain it at www.aomedia.org/license/patent. 9 10 use super::cdf_context::ContextWriter; 11 12 use super::*; 13 14 // Generates 4 bit field in which each bit set to 1 represents 15 // a blocksize partition 1111 means we split 64x64, 32x32, 16x16 16 // and 8x8. 1000 means we just split the 64x64 to 32x32 17 pub static partition_context_lookup: [[u8; 2]; BlockSize::BLOCK_SIZES_ALL] = [ 18 [31, 31], // 4X4 - {0b11111, 0b11111} 19 [31, 30], // 4X8 - {0b11111, 0b11110} 20 [30, 31], // 8X4 - {0b11110, 0b11111} 21 [30, 30], // 8X8 - {0b11110, 0b11110} 22 [30, 28], // 8X16 - {0b11110, 0b11100} 23 [28, 30], // 16X8 - {0b11100, 0b11110} 24 [28, 28], // 16X16 - {0b11100, 0b11100} 25 [28, 24], // 16X32 - {0b11100, 0b11000} 26 [24, 28], // 32X16 - {0b11000, 0b11100} 27 [24, 24], // 32X32 - {0b11000, 0b11000} 28 [24, 16], // 32X64 - {0b11000, 0b10000} 29 [16, 24], // 64X32 - {0b10000, 0b11000} 30 [16, 16], // 64X64 - {0b10000, 0b10000} 31 [16, 0], // 64X128- {0b10000, 0b00000} 32 [0, 16], // 128X64- {0b00000, 0b10000} 33 [0, 0], // 128X128-{0b00000, 0b00000} 34 [31, 28], // 4X16 - {0b11111, 0b11100} 35 [28, 31], // 16X4 - {0b11100, 0b11111} 36 [30, 24], // 8X32 - {0b11110, 0b11000} 37 [24, 30], // 32X8 - {0b11000, 0b11110} 38 [28, 16], // 16X64 - {0b11100, 0b10000} 39 [16, 28], // 64X16 - {0b10000, 0b11100} 40 ]; 41 42 pub const CFL_JOINT_SIGNS: usize = 8; 43 pub const CFL_ALPHA_CONTEXTS: usize = 6; 44 pub const CFL_ALPHABET_SIZE: usize = 16; 45 46 pub const PARTITION_PLOFFSET: usize = 4; 47 pub const PARTITION_BLOCK_SIZES: usize = 4 + 1; 48 const PARTITION_CONTEXTS_PRIMARY: usize = 49 PARTITION_BLOCK_SIZES * PARTITION_PLOFFSET; 50 pub const PARTITION_CONTEXTS: usize = PARTITION_CONTEXTS_PRIMARY; 51 pub const PARTITION_TYPES: usize = 4; 52 pub const EXT_PARTITION_TYPES: usize = 10; 53 54 pub const SKIP_CONTEXTS: usize = 3; 55 pub const SKIP_MODE_CONTEXTS: usize = 3; 56 57 // partition contexts are at 8x8 granularity, as it is not possible to 58 // split 4x4 blocks any further than that 59 pub const PARTITION_CONTEXT_GRANULARITY: usize = 8; 60 pub const PARTITION_CONTEXT_MAX_WIDTH: usize = 61 MAX_TILE_WIDTH / PARTITION_CONTEXT_GRANULARITY; 62 63 #[derive(Copy, Clone, Debug, PartialEq)] 64 pub enum CFLSign { 65 CFL_SIGN_ZERO = 0, 66 CFL_SIGN_NEG = 1, 67 CFL_SIGN_POS = 2, 68 } 69 70 impl CFLSign { from_alpha(a: i16) -> CFLSign71 pub fn from_alpha(a: i16) -> CFLSign { 72 [CFL_SIGN_NEG, CFL_SIGN_ZERO, CFL_SIGN_POS][(a.signum() + 1) as usize] 73 } 74 } 75 76 use crate::context::CFLSign::*; 77 78 const CFL_SIGNS: usize = 3; 79 static cfl_sign_value: [i16; CFL_SIGNS] = [0, -1, 1]; 80 81 #[derive(Copy, Clone, Debug)] 82 pub struct CFLParams { 83 pub sign: [CFLSign; 2], 84 pub scale: [u8; 2], 85 } 86 87 impl Default for CFLParams { 88 #[inline] default() -> Self89 fn default() -> Self { 90 Self { sign: [CFL_SIGN_NEG, CFL_SIGN_ZERO], scale: [1, 0] } 91 } 92 } 93 94 impl CFLParams { 95 #[inline] joint_sign(self) -> u3296 pub fn joint_sign(self) -> u32 { 97 assert!(self.sign[0] != CFL_SIGN_ZERO || self.sign[1] != CFL_SIGN_ZERO); 98 (self.sign[0] as u32) * (CFL_SIGNS as u32) + (self.sign[1] as u32) - 1 99 } 100 #[inline] context(self, uv: usize) -> usize101 pub fn context(self, uv: usize) -> usize { 102 assert!(self.sign[uv] != CFL_SIGN_ZERO); 103 (self.sign[uv] as usize - 1) * CFL_SIGNS + (self.sign[1 - uv] as usize) 104 } 105 #[inline] index(self, uv: usize) -> u32106 pub fn index(self, uv: usize) -> u32 { 107 assert!(self.sign[uv] != CFL_SIGN_ZERO && self.scale[uv] != 0); 108 (self.scale[uv] - 1) as u32 109 } 110 #[inline] alpha(self, uv: usize) -> i16111 pub fn alpha(self, uv: usize) -> i16 { 112 cfl_sign_value[self.sign[uv] as usize] * (self.scale[uv] as i16) 113 } 114 #[inline] from_alpha(u: i16, v: i16) -> CFLParams115 pub fn from_alpha(u: i16, v: i16) -> CFLParams { 116 CFLParams { 117 sign: [CFLSign::from_alpha(u), CFLSign::from_alpha(v)], 118 scale: [u.abs() as u8, v.abs() as u8], 119 } 120 } 121 } 122 123 #[cfg(test)] 124 mod test { 125 #[test] cdf_map()126 fn cdf_map() { 127 use super::*; 128 129 let cdf = CDFContext::new(8); 130 let cdf_map = FieldMap { map: cdf.build_map() }; 131 let f = &cdf.partition_cdf[2]; 132 cdf_map.lookup(f.as_ptr() as usize); 133 } 134 135 use super::CFLSign; 136 use super::CFLSign::*; 137 138 static cfl_alpha_signs: [[CFLSign; 2]; 8] = [ 139 [CFL_SIGN_ZERO, CFL_SIGN_NEG], 140 [CFL_SIGN_ZERO, CFL_SIGN_POS], 141 [CFL_SIGN_NEG, CFL_SIGN_ZERO], 142 [CFL_SIGN_NEG, CFL_SIGN_NEG], 143 [CFL_SIGN_NEG, CFL_SIGN_POS], 144 [CFL_SIGN_POS, CFL_SIGN_ZERO], 145 [CFL_SIGN_POS, CFL_SIGN_NEG], 146 [CFL_SIGN_POS, CFL_SIGN_POS], 147 ]; 148 149 static cfl_context: [[usize; 8]; 2] = 150 [[0, 0, 0, 1, 2, 3, 4, 5], [0, 3, 0, 1, 4, 0, 2, 5]]; 151 152 #[test] cfl_joint_sign()153 fn cfl_joint_sign() { 154 use super::*; 155 156 let mut cfl = CFLParams::default(); 157 for (joint_sign, &signs) in cfl_alpha_signs.iter().enumerate() { 158 cfl.sign = signs; 159 assert!(cfl.joint_sign() as usize == joint_sign); 160 for uv in 0..2 { 161 if signs[uv] != CFL_SIGN_ZERO { 162 assert!(cfl.context(uv) == cfl_context[uv][joint_sign]); 163 } 164 } 165 } 166 } 167 } 168 169 impl<'a> ContextWriter<'a> { partition_gather_horz_alike( out: &mut [u16; 2], cdf_in: &[u16], _bsize: BlockSize, )170 fn partition_gather_horz_alike( 171 out: &mut [u16; 2], cdf_in: &[u16], _bsize: BlockSize, 172 ) { 173 out[0] = 32768; 174 out[0] -= ContextWriter::cdf_element_prob( 175 cdf_in, 176 PartitionType::PARTITION_HORZ as usize, 177 ); 178 out[0] -= ContextWriter::cdf_element_prob( 179 cdf_in, 180 PartitionType::PARTITION_SPLIT as usize, 181 ); 182 out[0] -= ContextWriter::cdf_element_prob( 183 cdf_in, 184 PartitionType::PARTITION_HORZ_A as usize, 185 ); 186 out[0] -= ContextWriter::cdf_element_prob( 187 cdf_in, 188 PartitionType::PARTITION_HORZ_B as usize, 189 ); 190 out[0] -= ContextWriter::cdf_element_prob( 191 cdf_in, 192 PartitionType::PARTITION_VERT_A as usize, 193 ); 194 out[0] -= ContextWriter::cdf_element_prob( 195 cdf_in, 196 PartitionType::PARTITION_HORZ_4 as usize, 197 ); 198 out[0] = 32768 - out[0]; 199 out[1] = 0; 200 } 201 partition_gather_vert_alike( out: &mut [u16; 2], cdf_in: &[u16], _bsize: BlockSize, )202 fn partition_gather_vert_alike( 203 out: &mut [u16; 2], cdf_in: &[u16], _bsize: BlockSize, 204 ) { 205 out[0] = 32768; 206 out[0] -= ContextWriter::cdf_element_prob( 207 cdf_in, 208 PartitionType::PARTITION_VERT as usize, 209 ); 210 out[0] -= ContextWriter::cdf_element_prob( 211 cdf_in, 212 PartitionType::PARTITION_SPLIT as usize, 213 ); 214 out[0] -= ContextWriter::cdf_element_prob( 215 cdf_in, 216 PartitionType::PARTITION_HORZ_A as usize, 217 ); 218 out[0] -= ContextWriter::cdf_element_prob( 219 cdf_in, 220 PartitionType::PARTITION_VERT_A as usize, 221 ); 222 out[0] -= ContextWriter::cdf_element_prob( 223 cdf_in, 224 PartitionType::PARTITION_VERT_B as usize, 225 ); 226 out[0] -= ContextWriter::cdf_element_prob( 227 cdf_in, 228 PartitionType::PARTITION_VERT_4 as usize, 229 ); 230 out[0] = 32768 - out[0]; 231 out[1] = 0; 232 } 233 234 #[inline] write_skip<W: Writer>( &mut self, w: &mut W, bo: TileBlockOffset, skip: bool, )235 pub fn write_skip<W: Writer>( 236 &mut self, w: &mut W, bo: TileBlockOffset, skip: bool, 237 ) { 238 let ctx = self.bc.skip_context(bo); 239 let cdf = &mut self.fc.skip_cdfs[ctx]; 240 symbol_with_update!(self, w, skip as u32, cdf, 2); 241 } 242 get_segment_pred(&self, bo: TileBlockOffset) -> (u8, u8)243 pub fn get_segment_pred(&self, bo: TileBlockOffset) -> (u8, u8) { 244 let mut prev_ul = -1; 245 let mut prev_u = -1; 246 let mut prev_l = -1; 247 if bo.0.x > 0 && bo.0.y > 0 { 248 prev_ul = self.bc.blocks.above_left_of(bo).segmentation_idx as i8; 249 } 250 if bo.0.y > 0 { 251 prev_u = self.bc.blocks.above_of(bo).segmentation_idx as i8; 252 } 253 if bo.0.x > 0 { 254 prev_l = self.bc.blocks.left_of(bo).segmentation_idx as i8; 255 } 256 257 /* Pick CDF index based on number of matching/out-of-bounds segment IDs. */ 258 let cdf_index: u8; 259 if prev_ul < 0 || prev_u < 0 || prev_l < 0 { 260 /* Edge case */ 261 cdf_index = 0; 262 } else if (prev_ul == prev_u) && (prev_ul == prev_l) { 263 cdf_index = 2; 264 } else if (prev_ul == prev_u) || (prev_ul == prev_l) || (prev_u == prev_l) 265 { 266 cdf_index = 1; 267 } else { 268 cdf_index = 0; 269 } 270 271 /* If 2 or more are identical returns that as predictor, otherwise prev_l. */ 272 let r: i8; 273 if prev_u == -1 { 274 /* edge case */ 275 r = if prev_l == -1 { 0 } else { prev_l }; 276 } else if prev_l == -1 { 277 /* edge case */ 278 r = prev_u; 279 } else { 280 r = if prev_ul == prev_u { prev_u } else { prev_l }; 281 } 282 (r as u8, cdf_index) 283 } 284 write_cfl_alphas<W: Writer>(&mut self, w: &mut W, cfl: CFLParams)285 pub fn write_cfl_alphas<W: Writer>(&mut self, w: &mut W, cfl: CFLParams) { 286 symbol_with_update!(self, w, cfl.joint_sign(), &mut self.fc.cfl_sign_cdf); 287 for uv in 0..2 { 288 if cfl.sign[uv] != CFL_SIGN_ZERO { 289 symbol_with_update!( 290 self, 291 w, 292 cfl.index(uv), 293 &mut self.fc.cfl_alpha_cdf[cfl.context(uv)] 294 ); 295 } 296 } 297 } 298 write_partition( &mut self, w: &mut impl Writer, bo: TileBlockOffset, p: PartitionType, bsize: BlockSize, )299 pub fn write_partition( 300 &mut self, w: &mut impl Writer, bo: TileBlockOffset, p: PartitionType, 301 bsize: BlockSize, 302 ) { 303 debug_assert!(bsize.is_sqr()); 304 assert!(bsize >= BlockSize::BLOCK_8X8); 305 let hbs = bsize.width_mi() / 2; 306 let has_cols = (bo.0.x + hbs) < self.bc.blocks.cols(); 307 let has_rows = (bo.0.y + hbs) < self.bc.blocks.rows(); 308 let ctx = self.bc.partition_plane_context(bo, bsize); 309 assert!(ctx < PARTITION_CONTEXTS); 310 311 if !has_rows && !has_cols { 312 return; 313 } 314 315 if has_rows && has_cols { 316 if ctx < PARTITION_TYPES { 317 let cdf = &mut self.fc.partition_w8_cdf[ctx]; 318 symbol_with_update!(self, w, p as u32, cdf, 4); 319 } else if ctx < 4 * PARTITION_TYPES { 320 let cdf = &mut self.fc.partition_cdf[ctx - PARTITION_TYPES]; 321 symbol_with_update!(self, w, p as u32, cdf); 322 } else { 323 let cdf = &mut self.fc.partition_w128_cdf[ctx - 4 * PARTITION_TYPES]; 324 symbol_with_update!(self, w, p as u32, cdf); 325 } 326 } else if !has_rows && has_cols { 327 assert!( 328 p == PartitionType::PARTITION_SPLIT 329 || p == PartitionType::PARTITION_HORZ 330 ); 331 assert!(bsize > BlockSize::BLOCK_8X8); 332 let mut cdf = [0u16; 2]; 333 if ctx < PARTITION_TYPES { 334 let partition_cdf = &mut self.fc.partition_w8_cdf[ctx]; 335 ContextWriter::partition_gather_vert_alike( 336 &mut cdf, 337 partition_cdf, 338 bsize, 339 ); 340 } else if ctx < 4 * PARTITION_TYPES { 341 let partition_cdf = &mut self.fc.partition_cdf[ctx - PARTITION_TYPES]; 342 ContextWriter::partition_gather_vert_alike( 343 &mut cdf, 344 partition_cdf, 345 bsize, 346 ); 347 } else { 348 let partition_cdf = 349 &mut self.fc.partition_w128_cdf[ctx - 4 * PARTITION_TYPES]; 350 ContextWriter::partition_gather_vert_alike( 351 &mut cdf, 352 partition_cdf, 353 bsize, 354 ); 355 } 356 w.symbol((p == PartitionType::PARTITION_SPLIT) as u32, &cdf); 357 } else { 358 assert!( 359 p == PartitionType::PARTITION_SPLIT 360 || p == PartitionType::PARTITION_VERT 361 ); 362 assert!(bsize > BlockSize::BLOCK_8X8); 363 let mut cdf = [0u16; 2]; 364 if ctx < PARTITION_TYPES { 365 let partition_cdf = &mut self.fc.partition_w8_cdf[ctx]; 366 ContextWriter::partition_gather_horz_alike( 367 &mut cdf, 368 partition_cdf, 369 bsize, 370 ); 371 } else if ctx < 4 * PARTITION_TYPES { 372 let partition_cdf = &mut self.fc.partition_cdf[ctx - PARTITION_TYPES]; 373 ContextWriter::partition_gather_horz_alike( 374 &mut cdf, 375 partition_cdf, 376 bsize, 377 ); 378 } else { 379 let partition_cdf = 380 &mut self.fc.partition_w128_cdf[ctx - 4 * PARTITION_TYPES]; 381 ContextWriter::partition_gather_horz_alike( 382 &mut cdf, 383 partition_cdf, 384 bsize, 385 ); 386 } 387 w.symbol((p == PartitionType::PARTITION_SPLIT) as u32, &cdf); 388 } 389 } 390 neg_interleave(x: i32, r: i32, max: i32) -> i32391 fn neg_interleave(x: i32, r: i32, max: i32) -> i32 { 392 assert!(x < max); 393 if r == 0 { 394 return x; 395 } else if r >= (max - 1) { 396 return -x + max - 1; 397 } 398 let diff = x - r; 399 if 2 * r < max { 400 if diff.abs() <= r { 401 if diff > 0 { 402 return (diff << 1) - 1; 403 } else { 404 return (-diff) << 1; 405 } 406 } 407 x 408 } else { 409 if diff.abs() < (max - r) { 410 if diff > 0 { 411 return (diff << 1) - 1; 412 } else { 413 return (-diff) << 1; 414 } 415 } 416 (max - x) - 1 417 } 418 } 419 write_segmentation<W: Writer>( &mut self, w: &mut W, bo: TileBlockOffset, bsize: BlockSize, skip: bool, last_active_segid: u8, )420 pub fn write_segmentation<W: Writer>( 421 &mut self, w: &mut W, bo: TileBlockOffset, bsize: BlockSize, skip: bool, 422 last_active_segid: u8, 423 ) { 424 let (pred, cdf_index) = self.get_segment_pred(bo); 425 if skip { 426 self.bc.blocks.set_segmentation_idx(bo, bsize, pred); 427 return; 428 } 429 let seg_idx = self.bc.blocks[bo].segmentation_idx; 430 let coded_id = Self::neg_interleave( 431 seg_idx as i32, 432 pred as i32, 433 (last_active_segid + 1) as i32, 434 ); 435 symbol_with_update!( 436 self, 437 w, 438 coded_id as u32, 439 &mut self.fc.spatial_segmentation_cdfs[cdf_index as usize] 440 ); 441 } 442 } 443 444 impl<'a> BlockContext<'a> { partition_plane_context( &self, bo: TileBlockOffset, bsize: BlockSize, ) -> usize445 pub fn partition_plane_context( 446 &self, bo: TileBlockOffset, bsize: BlockSize, 447 ) -> usize { 448 // TODO: this should be way simpler without sub8x8 449 let above_ctx = self.above_partition_context[bo.0.x >> 1]; 450 let left_ctx = self.left_partition_context[bo.y_in_sb() >> 1]; 451 let bsl = bsize.width_log2() - BLOCK_8X8.width_log2(); 452 let above = (above_ctx >> bsl) & 1; 453 let left = (left_ctx >> bsl) & 1; 454 455 assert!(bsize.is_sqr()); 456 457 (left * 2 + above) as usize + bsl as usize * PARTITION_PLOFFSET 458 } 459 reset_skip_context( &mut self, bo: TileBlockOffset, bsize: BlockSize, xdec: usize, ydec: usize, cs: ChromaSampling, )460 pub fn reset_skip_context( 461 &mut self, bo: TileBlockOffset, bsize: BlockSize, xdec: usize, 462 ydec: usize, cs: ChromaSampling, 463 ) { 464 let num_planes = if cs == ChromaSampling::Cs400 { 1 } else { 3 }; 465 let nplanes = if bsize >= BLOCK_8X8 { 466 num_planes 467 } else { 468 1 + (num_planes - 1) * has_chroma(bo, bsize, xdec, ydec, cs) as usize 469 }; 470 471 for plane in 0..nplanes { 472 let xdec2 = if plane == 0 { 0 } else { xdec }; 473 let ydec2 = if plane == 0 { 0 } else { ydec }; 474 475 let plane_bsize = 476 if plane == 0 { bsize } else { bsize.subsampled_size(xdec2, ydec2) }; 477 let bw = plane_bsize.width_mi(); 478 let bh = plane_bsize.height_mi(); 479 480 for above in 481 &mut self.above_coeff_context[plane][(bo.0.x >> xdec2)..][..bw] 482 { 483 *above = 0; 484 } 485 486 let bo_y = bo.y_in_sb(); 487 for left in &mut self.left_coeff_context[plane][(bo_y >> ydec2)..][..bh] 488 { 489 *left = 0; 490 } 491 } 492 } 493 skip_context(&self, bo: TileBlockOffset) -> usize494 pub fn skip_context(&self, bo: TileBlockOffset) -> usize { 495 let above_skip = bo.0.y > 0 && self.blocks.above_of(bo).skip; 496 let left_skip = bo.0.x > 0 && self.blocks.left_of(bo).skip; 497 above_skip as usize + left_skip as usize 498 } 499 update_partition_context( &mut self, bo: TileBlockOffset, subsize: BlockSize, bsize: BlockSize, )500 pub fn update_partition_context( 501 &mut self, bo: TileBlockOffset, subsize: BlockSize, bsize: BlockSize, 502 ) { 503 assert!(bsize.is_sqr()); 504 505 let bw = bsize.width_mi(); 506 let bh = bsize.height_mi(); 507 508 let above_ctx = 509 &mut self.above_partition_context[bo.0.x >> 1..(bo.0.x + bw) >> 1]; 510 let left_ctx = &mut self.left_partition_context 511 [bo.y_in_sb() >> 1..(bo.y_in_sb() + bh) >> 1]; 512 513 // update the partition context at the end notes. set partition bits 514 // of block sizes larger than the current one to be one, and partition 515 // bits of smaller block sizes to be zero. 516 for above in &mut above_ctx[..bw >> 1] { 517 *above = partition_context_lookup[subsize as usize][0]; 518 } 519 520 for left in &mut left_ctx[..bh >> 1] { 521 *left = partition_context_lookup[subsize as usize][1]; 522 } 523 } 524 } 525