nihav-itu/src/codecs/h264/dsp/mc/mod.rs

   1 use nihav_core::frame::*;
   2 use nihav_codec_support::codecs::MV;
   3 use nihav_codec_support::codecs::blockdsp::*;
   4
   5 macro_rules! module_selector {
   6     ($( ($cond:meta, $module:ident) ),*) => {
   7         module_selector!(list; r#false; $(($cond, $module)),*);
   8     };
   9     (list; $nocond:meta; ($ccar:meta, $carmod:ident), $(($condcdr:meta, $cdrmod:ident)),*) => {
  10         module_selector!(single; $nocond; $ccar; $carmod);
  11         module_selector!(list; any($nocond, $ccar); $(($condcdr, $cdrmod)),*);
  12     };
  13     (list; $nocond:meta; ($yescond:meta, $module:ident)) => {
  14         module_selector!(single; $nocond; $yescond; $module);
  15     };
  16     (list; $_:meta; ) => {};
  17     (single; $nocond:meta; $yescond:meta; $module:ident) => {
  18         #[cfg(all(not($nocond), $yescond))]
  19         mod $module;
  20         #[cfg(all(not($nocond), $yescond))]
  21         use $module::*;
  22     };
  23 }
  24
  25 module_selector! (
  26     (all(feature = "simd", target_arch = "x86_64"), x86),
  27     (debug_assertions, debug),
  28     (not(debug_assertions), release)
  29 );
  30
  31 type MCFunc = fn (dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, h: usize);
  32
  33 fn clip_u8(val: i16) -> u8 { val.max(0).min(255) as u8 }
  34
  35 trait RegisterSIMD {
  36     fn register_simd(&mut self);
  37 }
  38
  39 #[allow(clippy::type_complexity)]
  40 pub struct H264MC {
  41     avg_buf:    NAVideoBufferRef<u8>,
  42     pub put_block_weighted:     [fn (dst: &mut [u8], stride: usize, src: &[u8], h: usize, wparams: [i8; 3]); 4],
  43     pub put_block_weighted2:    [fn (dst: &mut [u8], stride: usize, src0: &[u8], src1: &[u8], h: usize, wparams: [i8; 5]); 4],
  44     pub chroma_interp:          [fn (dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, dx: u16, dy: u16, h: usize); 3],
  45     avg:        [fn (dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bh: usize); 4],
  46
  47     width:      usize,
  48     height:     usize,
  49 }
  50
  51 impl H264MC {
  52     pub fn new(avg_buf: NAVideoBufferRef<u8>) -> Self {
  53         let mut obj = Self {
  54             avg_buf,
  55             put_block_weighted:     [put_blk_w_2, put_blk_w_4, put_blk_w_8, put_blk_w_16],
  56             put_block_weighted2:    [put_blk_w2_2, put_blk_w2_4, put_blk_w2_8, put_blk_w2_16],
  57             chroma_interp:          [chroma_interp_2, chroma_interp_4, chroma_interp_8],
  58             avg:                    [avg_2, avg_4, avg_8, avg_16],
  59             width: 0, height: 0,
  60         };
  61         obj.register_simd();
  62         obj
  63     }
  64     pub fn set_dimensions(&mut self, width: usize, height: usize) {
  65         self.width  = width;
  66         self.height = height;
  67     }
  68     pub fn do_mc(&mut self, frm: &mut NASimpleVideoFrame<u8>, refpic: NAVideoBufferRef<u8>, xpos: usize, ypos: usize, w: usize, h: usize, mv: MV) {
  69         let mut ebuf = [0u8; 22 * 22];
  70         let mvx = mv.x >> 2;
  71         let mvy = mv.y >> 2;
  72         let mode = ((mv.x & 3) + (mv.y & 3) * 4) as usize;
  73         let pre  = if mode != 0 { 2isize } else { 0 };
  74         let post = if mode != 0 { 3isize } else { 0 };
  75         let (yw, yh) = (self.width, self.height);
  76         let src = refpic.get_data();
  77         let systride = refpic.get_stride(0);
  78         let src_x = (xpos as isize) + (mvx as isize);
  79         let src_y = (ypos as isize) + (mvy as isize);
  80         let (ysrc, ystride) = if (src_x - pre < 0) || (src_x + (w as isize) + post > (yw as isize)) || (src_y - pre < 0) || (src_y + (h as isize) + post > (yh as isize)) {
  81                 let add = (pre + post) as usize;
  82                 edge_emu(&refpic, src_x - pre, src_y - pre, w + add, h + add, &mut ebuf, 22, 0, 0);
  83                 (&ebuf[..], 22)
  84             } else {
  85                 (&src[refpic.get_offset(0) + ((src_x - pre) as usize) + ((src_y - pre) as usize) * systride..], systride)
  86             };
  87         let wmode = match w {
  88                 4 => 0,
  89                 8 => 1,
  90                 _ => 2,
  91             };
  92         (H264_LUMA_INTERP[wmode][mode])(&mut frm.data[frm.offset[0] + xpos + ypos * frm.stride[0]..], frm.stride[0], ysrc, ystride, h);
  93
  94         let (cw, ch) = (self.width >> 1, self.height >> 1);
  95         let mvx = mv.x >> 3;
  96         let mvy = mv.y >> 3;
  97         let dx = (mv.x & 7) as u16;
  98         let dy = (mv.y & 7) as u16;
  99         let src_x = ((xpos >> 1) as isize) + (mvx as isize);
 100         let src_y = ((ypos >> 1) as isize) + (mvy as isize);
 101         let suoff = refpic.get_offset(1);
 102         let svoff = refpic.get_offset(2);
 103         let sustride = refpic.get_stride(1);
 104         let svstride = refpic.get_stride(2);
 105         let cbw = w / 2;
 106         let cbh = h / 2;
 107         let (csrc, cstride) = if (src_x < 0) || (src_x + (cbw as isize) + 1 > (cw as isize)) || (src_y < 0) || (src_y + (cbh as isize) + 1 > (ch as isize)) {
 108                 edge_emu(&refpic, src_x, src_y, cbw+1, cbh+1, &mut ebuf,      18, 1, 4);
 109                 edge_emu(&refpic, src_x, src_y, cbw+1, cbh+1, &mut ebuf[9..], 18, 2, 4);
 110                 ([&ebuf, &ebuf[9..]], [18, 18])
 111             } else {
 112                 ([&src[suoff + (src_x as usize) + (src_y as usize) * sustride..],
 113                  &src[svoff + (src_x as usize) + (src_y as usize) * svstride..]],
 114                  [sustride, svstride])
 115             };
 116         for chroma in 1..3 {
 117             let off = frm.offset[chroma] + xpos / 2 + (ypos / 2) * frm.stride[chroma];
 118             (self.chroma_interp[wmode])(&mut frm.data[off..], frm.stride[chroma], csrc[chroma - 1], cstride[chroma - 1], dx, dy, cbh);
 119         }
 120     }
 121
 122     pub fn mc_blocks(&mut self, ydst: &mut [u8], udst: &mut [u8], vdst: &mut [u8], refpic: NAVideoBufferRef<u8>, xpos: usize, ypos: usize, w: usize, h: usize, mv: MV) {
 123         let mode = ((mv.x & 3) + (mv.y & 3) * 4) as usize;
 124
 125         let pre  = if mode != 0 { 2 } else { 0 };
 126         let post = if mode != 0 { 3 } else { 0 };
 127         let (width, height) = (self.width, self.height);
 128         let sx = (xpos as isize) + ((mv.x >> 2) as isize);
 129         let sy = (ypos as isize) + ((mv.y >> 2) as isize);
 130
 131         const EBUF_STRIDE: usize = 32;
 132         let mut ebuf = [0u8; EBUF_STRIDE * (16 + 2 + 3)];
 133
 134         let wmode = match w {
 135                 4 => 0,
 136                 8 => 1,
 137                 _ => 2,
 138             };
 139         if (sx - pre < 0) || (sx + (w as isize) + post > (width as isize)) ||
 140            (sy - pre < 0) || (sy + (h as isize) + post > (height as isize)) {
 141             let edge = (pre + post) as usize;
 142             edge_emu(&refpic, sx - pre, sy - pre, w + edge, h + edge,
 143                      &mut ebuf, EBUF_STRIDE, 0, 0);
 144             (H264_LUMA_INTERP[wmode][mode])(ydst, 16, &ebuf, EBUF_STRIDE, h);
 145         } else {
 146             let sstride = refpic.get_stride(0);
 147             let soff    = refpic.get_offset(0);
 148             let sdta    = refpic.get_data();
 149             let sbuf: &[u8] = sdta.as_slice();
 150             let saddr = soff + ((sx - pre) as usize) + ((sy - pre) as usize) * sstride;
 151             (H264_LUMA_INTERP[wmode][mode])(ydst, 16, &sbuf[saddr..], sstride, h);
 152         }
 153
 154         let (cw, ch) = (self.width >> 1, self.height >> 1);
 155         let mvx = mv.x >> 3;
 156         let mvy = mv.y >> 3;
 157         let dx = (mv.x & 7) as u16;
 158         let dy = (mv.y & 7) as u16;
 159         let src_x = ((xpos >> 1) as isize) + (mvx as isize);
 160         let src_y = ((ypos >> 1) as isize) + (mvy as isize);
 161         let suoff = refpic.get_offset(1);
 162         let svoff = refpic.get_offset(2);
 163         let sustride = refpic.get_stride(1);
 164         let svstride = refpic.get_stride(2);
 165         let src = refpic.get_data();
 166         let cbw = w / 2;
 167         let cbh = h / 2;
 168         let (csrc, cstride) = if (src_x < 0) || (src_x + (cbw as isize) + 1 > (cw as isize)) || (src_y < 0) || (src_y + (cbh as isize) + 1 > (ch as isize)) {
 169                 edge_emu(&refpic, src_x, src_y, cbw+1, cbh+1, &mut ebuf,      18, 1, 4);
 170                 edge_emu(&refpic, src_x, src_y, cbw+1, cbh+1, &mut ebuf[9..], 18, 2, 4);
 171                 ([&ebuf, &ebuf[9..]], [18, 18])
 172             } else {
 173                 ([&src[suoff + (src_x as usize) + (src_y as usize) * sustride..],
 174                  &src[svoff + (src_x as usize) + (src_y as usize) * svstride..]],
 175                  [sustride, svstride])
 176             };
 177         (self.chroma_interp[wmode])(udst, 16, csrc[0], cstride[0], dx, dy, cbh);
 178         (self.chroma_interp[wmode])(vdst, 16, csrc[1], cstride[1], dx, dy, cbh);
 179     }
 180
 181     pub fn do_mc_avg(&mut self, frm: &mut NASimpleVideoFrame<u8>, refpic: NAVideoBufferRef<u8>, xpos: usize, ypos: usize, w: usize, h: usize, mv: MV) {
 182         let mut abuf = self.avg_buf.clone();
 183         let mut afrm = NASimpleVideoFrame::from_video_buf(&mut abuf).unwrap();
 184         let amv = MV { x: mv.x + (xpos as i16) * 4, y: mv.y + (ypos as i16) * 4 };
 185         self.do_mc(&mut afrm, refpic, 0, 0, w, h, amv);
 186         let wsize = match w {
 187                 2 => 0,
 188                 4 => 1,
 189                 8 => 2,
 190                 _ => 3,
 191             };
 192         let src = self.avg_buf.get_data();
 193         for comp in 0..3 {
 194             let shift = if comp == 0 { 0 } else { 1 };
 195             let sstride = self.avg_buf.get_stride(comp);
 196             let soff = self.avg_buf.get_offset(comp);
 197             (self.avg[wsize - shift])(&mut frm.data[frm.offset[comp] + (xpos >> shift) + (ypos >> shift) * frm.stride[comp]..], frm.stride[comp], &src[soff..], sstride, h >> shift);
 198         }
 199     }
 200
 201     pub fn gray_block(&mut self, frm: &mut NASimpleVideoFrame<u8>, x: usize, y: usize, w: usize, h: usize) {
 202         let yoff = frm.offset[0] + x + y * frm.stride[0];
 203         let coff = [frm.offset[1] + x / 2 + y / 2 * frm.stride[1],
 204                     frm.offset[2] + x / 2 + y / 2 * frm.stride[2]];
 205         for row in frm.data[yoff..].chunks_mut(frm.stride[0]).take(h) {
 206             for el in row[..w].iter_mut() {
 207                 *el = 128;
 208             }
 209         }
 210         for chroma in 0..2 {
 211             for row in frm.data[coff[chroma]..].chunks_mut(frm.stride[chroma + 1]).take(h / 2) {
 212                 for el in row[..w / 2].iter_mut() {
 213                     *el = 128;
 214                 }
 215             }
 216         }
 217     }
 218 }
 219
 220 fn avg(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bw: usize, bh: usize) {
 221     for (dline, sline) in dst.chunks_mut(dstride).zip(src.chunks(sstride)).take(bh) {
 222         for (dst, src) in dline.iter_mut().zip(sline.iter()).take(bw) {
 223             *dst = ((u16::from(*dst) + u16::from(*src) + 1) >> 1) as u8;
 224         }
 225     }
 226 }
 227
 228 fn avg_2(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bh: usize) {
 229     let _ = src[sstride + 1];
 230     let _ = dst[dstride + 1];
 231     dst[0]           = ((u16::from(dst[0])           + u16::from(src[0])           + 1) >> 1) as u8;
 232     dst[1]           = ((u16::from(dst[1])           + u16::from(src[1])           + 1) >> 1) as u8;
 233     dst[dstride]     = ((u16::from(dst[dstride])     + u16::from(src[sstride])     + 1) >> 1) as u8;
 234     dst[dstride + 1] = ((u16::from(dst[dstride + 1]) + u16::from(src[sstride + 1]) + 1) >> 1) as u8;
 235     if bh == 4 {
 236         let _ = src[sstride * 3 + 1];
 237         let _ = dst[dstride * 3 + 1];
 238         dst[dstride * 2]     = ((u16::from(dst[dstride * 2])     + u16::from(src[sstride * 2])     + 1) >> 1) as u8;
 239         dst[dstride * 2 + 1] = ((u16::from(dst[dstride * 2 + 1]) + u16::from(src[sstride * 2 + 1]) + 1) >> 1) as u8;
 240         dst[dstride * 3]     = ((u16::from(dst[dstride * 3])     + u16::from(src[sstride * 3])     + 1) >> 1) as u8;
 241         dst[dstride * 3 + 1] = ((u16::from(dst[dstride * 3 + 1]) + u16::from(src[sstride * 3 + 1]) + 1) >> 1) as u8;
 242     }
 243 }
 244 fn avg_4(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bh: usize) {
 245     avg(dst, dstride, src, sstride, 4, bh);
 246 }
 247 fn avg_8(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bh: usize) {
 248     avg(dst, dstride, src, sstride, 8, bh);
 249 }
 250 fn avg_16(dst: &mut [u8], dstride: usize, src: &[u8], sstride: usize, bh: usize) {
 251     avg(dst, dstride, src, sstride, 16, bh);
 252 }
 253
 254 fn put_block_weighted(dst: &mut [u8], stride: usize, src: &[u8], w: usize, h: usize, wparams: [i8; 3]) {
 255     let weight = i16::from(wparams[0]);
 256     let offset = i16::from(wparams[1]);
 257     let wshift = wparams[2] as u8;
 258     let bias = (1 << wshift) >> 1;
 259
 260     for (drow, srow) in dst.chunks_mut(stride).zip(src.chunks(16)).take(h) {
 261         for (dst, &src) in drow[..w].iter_mut().zip(srow.iter()) {
 262             *dst = clip_u8(((i16::from(src) * weight + bias) >> wshift) + offset);
 263         }
 264     }
 265 }
 266
 267 fn put_blk_w_2(dst: &mut [u8], stride: usize, src: &[u8], h: usize, wparams: [i8; 3]) {
 268     put_block_weighted(dst, stride, src, 2, h, wparams);
 269 }
 270 fn put_blk_w_4(dst: &mut [u8], stride: usize, src: &[u8], h: usize, wparams: [i8; 3]) {
 271     put_block_weighted(dst, stride, src, 4, h, wparams);
 272 }
 273 fn put_blk_w_8(dst: &mut [u8], stride: usize, src: &[u8], h: usize, wparams: [i8; 3]) {
 274     put_block_weighted(dst, stride, src, 8, h, wparams);
 275 }
 276 fn put_blk_w_16(dst: &mut [u8], stride: usize, src: &[u8], h: usize, wparams: [i8; 3]) {
 277     put_block_weighted(dst, stride, src, 16, h, wparams);
 278 }
 279
 280 fn put_block_weighted2(dst: &mut [u8], stride: usize, src0: &[u8], src1: &[u8], w: usize, h: usize, wparams: [i8; 5]) {
 281     let weight0 = i16::from(wparams[0]);
 282     let offset0 = i16::from(wparams[1]);
 283     let weight1 = i16::from(wparams[2]);
 284     let offset1 = i16::from(wparams[3]);
 285     let wshift = (wparams[4] as u8) + 1;
 286     let offset = (offset0 + offset1 + 1) >> 1;
 287     let bias = (1 << wshift) >> 1;
 288
 289     for (drow, (srow0, srow1)) in dst.chunks_mut(stride).zip(src0.chunks(16).zip(src1.chunks(16))).take(h) {
 290         for (dst, (&src0, &src1)) in drow[..w].iter_mut().zip(srow0.iter().zip(srow1.iter())) {
 291             *dst = clip_u8(((i16::from(src0) * weight0 + i16::from(src1) * weight1 + bias) >> wshift) + offset);
 292         }
 293     }
 294 }
 295
 296 fn put_blk_w2_2(dst: &mut [u8], stride: usize, src0: &[u8], src1: &[u8], h: usize, wparams: [i8; 5]) {
 297     let weight0 = i16::from(wparams[0]);
 298     let offset0 = i16::from(wparams[1]);
 299     let weight1 = i16::from(wparams[2]);
 300     let offset1 = i16::from(wparams[3]);
 301     let wshift = (wparams[4] as u8) + 1;
 302     let offset = (offset0 + offset1 + 1) >> 1;
 303     let bias = (1 << wshift) >> 1;
 304
 305     let _ = src0[16 + 1];
 306     let _ = src1[16 + 1];
 307     let _ = dst[stride + 1];
 308     dst[0]          = clip_u8(((i16::from(src0[ 0]) * weight0 + i16::from(src1[ 0]) * weight1 + bias) >> wshift) + offset);
 309     dst[1]          = clip_u8(((i16::from(src0[ 1]) * weight0 + i16::from(src1[ 1]) * weight1 + bias) >> wshift) + offset);
 310     dst[stride]     = clip_u8(((i16::from(src0[16]) * weight0 + i16::from(src1[16]) * weight1 + bias) >> wshift) + offset);
 311     dst[stride + 1] = clip_u8(((i16::from(src0[17]) * weight0 + i16::from(src1[17]) * weight1 + bias) >> wshift) + offset);
 312     if h == 4 {
 313         let _ = src0[16 * 3 + 1];
 314         let _ = src1[16 * 3 + 1];
 315         let _ = dst[stride * 3 + 1];
 316         dst[stride * 2]     = clip_u8(((i16::from(src0[32]) * weight0 + i16::from(src1[32]) * weight1 + bias) >> wshift) + offset);
 317         dst[stride * 2 + 1] = clip_u8(((i16::from(src0[33]) * weight0 + i16::from(src1[33]) * weight1 + bias) >> wshift) + offset);
 318         dst[stride * 3]     = clip_u8(((i16::from(src0[48]) * weight0 + i16::from(src1[48]) * weight1 + bias) >> wshift) + offset);
 319         dst[stride * 3 + 1] = clip_u8(((i16::from(src0[49]) * weight0 + i16::from(src1[49]) * weight1 + bias) >> wshift) + offset);
 320     }
 321 }
 322 fn put_blk_w2_4(dst: &mut [u8], stride: usize, src0: &[u8], src1: &[u8], h: usize, wparams: [i8; 5]) {
 323     put_block_weighted2(dst, stride, src0, src1, 4, h, wparams);
 324 }
 325 fn put_blk_w2_8(dst: &mut [u8], stride: usize, src0: &[u8], src1: &[u8], h: usize, wparams: [i8; 5]) {
 326     put_block_weighted2(dst, stride, src0, src1, 8, h, wparams);
 327 }
 328 fn put_blk_w2_16(dst: &mut [u8], stride: usize, src0: &[u8], src1: &[u8], h: usize, wparams: [i8; 5]) {
 329     put_block_weighted2(dst, stride, src0, src1, 16, h, wparams);
 330 }