Tile encoding

[rom1v]

(description updated on 16 april 2019)

This PR implements tile encoding (#631).

1. The (many) first commits introduce tiling structures, which allow to expose simultaneous tiled regions of the whole frame data.
2. Following commits use the tiling structures where necessary in the whole codebase.
3. Then, command line arguments are added, the encoder encodes tiles separately, and separate tiles are written to the bitstream.
4. Finally, parallelization is enabled (spoiler: 1 line of code).

Context

Encoding a frame first involves frame-wise accesses (initialization, etc.), then tile-wise accesses (to encode tiles in parallel), then frame-wise accesses using the results of tile-encoding (deblocking, cdef, …):

                                \
      +----------------+         |
      |                |         |
      |                |         |  Frame-wise accesses
      |                |          >
      |                |         |   - FrameState<T>
      |                |         |   - Frame<T>
      +----------------+         |   - Plane<T>
                                /    - ...

              ||   tiling views
              \/
                                \
  +---+  +---+  +---+  +---+     |
  |   |  |   |  |   |  |   |     |  Tile encoding (possibly in parallel)
  +---+  +---+  +---+  +---+     |
                                 |
  +---+  +---+  +---+  +---+     |  Tile-wise accesses
  |   |  |   |  |   |  |   |      >
  +---+  +---+  +---+  +---+     |   - TileStateMut<'_, T>
                                 |   - TileMut<'_, T>
  +---+  +---+  +---+  +---+     |   - PlaneRegionMut<'_, T>
  |   |  |   |  |   |  |   |     |
  +---+  +---+  +---+  +---+     |
                                /

              ||   vanishing of tiling views
              \/
                                \
      +----------------+         |
      |                |         |
      |                |         |  Frame-wise accesses
      |                |          >
      |                |         |  (deblocking, CDEF, ...)
      |                |         |
      +----------------+         |
                                /

Tiling

As you know, in Rust, it is not sufficient not to read/write the same memory from several threads, it must be impossible to write (safe) code that could do it. More precisely, a mutable reference may not alias any other reference to the same memory.

That's the reason why, as a preliminary step, I replaced accesses using the whole plane as a raw slice in addition to the stride information by PlaneSlice (#1035) and PlaneMutSlice (#1043).

But Plane(Mut)Slice still borrows the whole plane slice, so it does not, in itself, solves the problem.

There are several structures to be tiled, which form a tree:

 +- FrameState → TileState
 |  +- Frame → Tile
 |  |  +- Plane → PlaneRegion 
 |  +  RestorationState → TileRestorationState
 |  |  +- RestorationPlane → TileRestorationPlane
 |  |     +- FrameRestorationUnits → TileRestorationUnits
 |  +  FrameMotionVectors → TileMotionVectors
 +- FrameBlocks → TileBlocks

Most of them exist both in const and mutable version (e.g. PlaneRegion and PlaneRegionMut).

Tiling structures

PlaneRegion

This is a view of bounded region of a Plane. It is similar to PlaneSlice, except that it does not borrow the whole underlying raw slice. That way, it is possible to get several non-overlapping regions simultaneously.

In the end, we should probably merge it with PlaneSlice, but it requires more work because some frame-wise code still uses PlaneSlice in the code base.

It is possible to retrieve a subregion of a region (which may not exceed its parent). In theory, a subregion is defined by a rectangle (for example: x, y, width, height), but in practice, we need more flexibility. For example, we often need to retrieve a region from an offset, using the same bottom-right corner as its parent without providing width and height.

For that purpose, I propose a specific Area structure (actually, a Rust enum) to describe subregion bounds. Here are some usage examples:

let region = plane.region(Area::Rect { x: 32, y: 32, width: 512, height: 512 });

// the area is relative to the parent region
let subregion = region.subregion(Area::StartingAt { x: 128, y: 128 });
// it is equivalent to
let subregion = region.subregion(Area::Rect { x: 128, y: 128, width: 384, height: 384 });
// or
let subregion = plane.region(Area:: Rect { x: 160, y: 160, width: 384, height: 384 });

Retrieving a subregion from a BlockOffset is so common accross the code base that I decided to expose it directly:

let bo = BlockOffset { x: 2, y: 3 };
let subregion = region.subregion(Area::BlockStartingAt { bo });

Like Plane(Mut)Slice, it provides operator[] and iterators over its rows:

let row5 = &region[5];
let value = region[3][4];
for row in region.rows_iter() {
    let _first_four_values = &row[..4];
}

The mutable versions of the structure (PlaneRegionMut) and methods are also provided.

Tile

A Tile is a view of 3 colocated plane regions (Tile is to a PlaneRegion as a Frame is to a Plane).

The mutable version (TileMut) is also provided.

TileState

The way the FrameState fields are mapped in TileState depends on how they are accessed tile-wise and frame-wise.

Some fields (like qc) are only used during tile-encoding, so they are only stored in TileState.

Some other fields (like input or segmentation) are not written tile-wise, so they just reference the matching field in FrameState.

Some others (like rec) are written tile-wise, but must be accessible frame-wise once the tile views vanish (e.g. for deblocking).

It contains 2 tiled views: TileRestorationState and a vector of TileMotionVectorsMut (a tiled view of FrameMotionVectors).

This structure is only provided as mutable (TileStateMut). A const version is not necessary, and would require to instantiate a const version of all its embedded tiled views.

TileBlocks

TileBlocks is a tiled view of FrameBlocks. It exposes the blocks associated to the tile.

The mutable version (TileBlocksMut) is also provided.

Splitting into tiles

A TilingInfo structure computes all the details about tiling from the frame width and height and the (log2 of the) number of tile columns and rows. The details are accessible for initializing data or writing into the bitstream.

It provides an iterator over tiles (yielding one TileStateMut and one TileBlocksMut for each tile).

Frame offsets vs tile offsets

In encode_tile(), super-block, block and plane offsets are expressed relative to the tile. The tiling views expose its data relative to the tile:

• plane_region[y][x] is pixel (x, y) relative to the plane region,
• tile_blocks[boy][box] contains the Block at (box, boy) relative to the tile,
• …

TileStateMut exposes some references to frame-level data stored in FrameState:

• input is a reference to the whole frame,
• input_hres and input_qres are references to the whole planes.

When accessing these frame-level data, tile offsets are converted to frame offsets, for example by:

let frame_bo = ts.to_frame_block_offset(bo);

Current state

It works.

Need more tests and reviews.

Usage

Pass the requested log2 number of tiles, with --tile-cols-log2 and --tile-rows-log2. For example, to request 2x2 tiles:

rav1e video.y4m -o video.ivf --tile-cols-log2 1 --tile-rows-log2 1

Currently, the number of tiles is passed in log2 (like in libaom, even if the aomenc options are called --tile-columns and --tile-rows), to avoid any confusion. Maybe we could find a correct user-friendly option later.

Note that the actual number of tiles may be smaller (e.g. if the image size has fewer super-blocks).

[coveralls]

Coverage increased (+6.8%) to 82.605% when pulling b52795e on rom1v:tiling into b9fef7c on xiph:master.

[rom1v]

I added #[derive(Copy)] to BlockOffset, and used BlockOffset by value in the Area enum. Using a reference would add a lifetime parameter (Area<'_>), which would be overkill here.

But since I added Copy, clippy warns about BlockOffset everywhere:

error: this argument is passed by reference, but would be more efficient if passed by value
    --> src/partition.rs:1234:19
     |
1234 | pub fn has_tr(bo: &BlockOffset, bsize: BlockSize) -> bool {
     |                   ^^^^^^^^^^^^ help: consider passing by value instead: `BlockOffset`
     |
     = note: `-D clippy::trivially-copy-pass-by-ref` implied by `-D warnings`
     = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#trivially_copy_pass_by_ref

I think we should replace all &BlockOffset to BlockOffset in function parameters. Do you agree that I do that on master?

[rom1v]

At least, we can see the tiles

EDIT: another one with wrong chroma:

and with wrong luma:

[rom1v]

Here is my first working tiled (2×2) video encoded with rav1e:
bbb_tiled3.mp4 7919c917bd8769ef29ef108553deec55caee88becf1cb65b8acb1a649dd89ef6
bbb_tiled3.ivf 73fba3cbd3beb4b811f05a8d2daf5ba5169a29299a2d1d44d60ad8b5ad922116

From this version: https://github.com/rom1v/rav1e/commits/tiling.100

ffmpeg -i big_buck_bunny_720p.mp4 -ss 10 -t 10 bbb.y4m
target/release/rav1e bbb.y4m --tile-rows-log2 1 --tile-cols-log2 1 -o bbb_tiled3.ivf
ffmpeg -i bbb_tiled3.ivf -c:v copy bbb_tiled3.mp4

[rom1v]

1 tile vs 4 tiles: https://beta.arewecompressedyet.com/?job=%402019-04-05T21%3A29%3A59.430Z_ref_1_tile&job=%402019-04-05T21%3A28%3A52.445Z_4_tiles

Of course, there is a cost in quality (SSIM, PSNR…) because it loses the possibility to exploit some redundancy across tiles in the same frame. For now, the current version only saves CDF from tile 0 (it should choose the bigger tile in bytes instead), and always store tile sizes on 4 bytes. It can (and will) be improved.

The encoding time is worse with tiling on AWCY because I think that it uses only 1 core per instance.

[rom1v]

Unfortunately (but as expected), the tiling structures are not a zero cost abstraction. They add an overhead in encoding time between 1~3%.

Concretely, if we compare the version compiled from master branch with the version compiled from this pull request, and encode a video without tiling (1 tile), the latter takes a bit more time to encode. Said otherwise, we pay a (small) cost for tile encoding even if we don't use it.

As an example, on my laptop, an encoding takes 3mn33,245 on master and 3mn35,722 on tiling (it's just an example, but the difference is quite stable).

You can compare encoding times on AWCY for 1 tile: https://beta.arewecompressedyet.com/?job=master-70005e353aa8ce21e3ecd257c927f71d4012a117&job=%402019-04-05T21%3A29%3A59.430Z_ref_1_tile

Maybe some work could be done to minimize this overhead (for example using more #[inline(always)]), but I think the overhead cannot be removed (within a memory-safe language).

EDIT: now there is no overhead (even a negative overhead with more inlines than on master) 🎉

[ycho]

I think it is ready to land now.
Thank you very much for the great work!

As I observed, BD Rate change by this PR is:
At speed 0, low_latency=true (Sine low_latency=false seems not working correctly)
1 tile -> 2x2=4 tiles:
AWCY link

PSNR	PSNR Cb	PSNR Cr	PSNR HVS	SSIM	MS SSIM	CIEDE 2000
2.8320	1.7719	2.2092	2.7641	2.8205	2.7920	2.4722

1 tile -> 4x2 (col x row) tiles:
AWCY link

PSNR	PSNR Cb	PSNR Cr	PSNR HVS	SSIM	MS SSIM	CIEDE 2000
5.2285	3.8110	4.2109	5.0781	5.2115	5.1621	4.7286

[rom1v]

(Since low_latency=false seems not working correctly)

Actually, in absolute, --low_latency=false gives better objective quality results. But relatively, tiling cause less loss with --low_latency=true: AWCY.

I don't know why.

[ycho]

Actually, in absolute, --low_latency=false gives better objective quality results.

Sure, far better (> 10%), that is why we wanted to have pyramid and/or frame-reordering.

[rom1v]

I published a blog post about this feature: https://blog.rom1v.com/2019/04/implementing-tile-encoding-in-rav1e/