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