Skip to content

Optimize TCQ with block-level AVX2 dispatch - #5329

Open
chengchen-google wants to merge 1 commit into
AOMediaCodec:av2-encfrom
chengchen-google:0825/0001_tcq_block_level_avx2_fusion
Open

Optimize TCQ with block-level AVX2 dispatch#5329
chengchen-google wants to merge 1 commit into
AOMediaCodec:av2-encfrom
chengchen-google:0825/0001_tcq_block_level_avx2_fusion

Conversation

@chengchen-google

Copy link
Copy Markdown
Contributor

Previously, TCQ evaluated coefficients along anti-diagonal scan lines in a portable C loop (trellis_loop_diagonal_st8), invoking individual fine-grained RTCD SIMD functions per coefficient.

This fine-grained dispatch structure incurred significant overhead:

  • Repeated RTCD function calls and indirect branches on every coeff.
  • Register spills to memory structures across subroutine boundaries.
  • Optimization barriers preventing compiler vector scheduling across the full diagonal loop.

This patch introduces a coarse-grained, block-level RTCD entry point:
av2_trellis_loop_diagonal_st8()
with AVX2 specialization (av2_trellis_loop_diagonal_st8_avx2).

Updates are now executed contiguously within AVX2 registers without exiting to scalar C code per coefficient.

Unit tests are added to compare the the new and original implementation.

We can see that the TCQ function becomes 3-10% faster.

=======================================================================
TX Size   Pure C (us)   Base AVX2 (us)   Patch AVX2 (us)  vs Base AVX2
=======================================================================
4x4       110313        45646            44360           1.03x (+2.8%)
8x8       369747        102636           96537           1.06x (+5.9%)
16x16     343618        81354            74201           1.10x (+8.8%)
32x32     268794        60375            54646           1.10x (+9.5%)
4x8       195707        64939            61781           1.05x (+4.9%)
8x4       201683        65227            61872           1.05x (+5.1%)
8x16      175484        44150            40669           1.09x (+7.9%)
16x8      180465        44455            40943           1.09x (+7.9%)
16x32     134924        30975            28011           1.11x (+9.6%)
32x16     134311        30946            28192           1.10x (+8.9%)
4x16      364581        103343           96494           1.07x (+6.6%)
16x4      390933        104326           97702           1.07x (+6.3%)
8x32      343005        81586            74379           1.10x (+8.8%)
32x8      352547        81916            74820           1.09x (+8.7%)
=======================================================================

@chengchen-google
chengchen-google force-pushed the 0825/0001_tcq_block_level_avx2_fusion branch from 0d6bae6 to 034fae7 Compare August 26, 2026 22:44
Comment thread test/trellis_test.cc

@jjustiss-apple jjustiss-apple left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, just a minor unit test maintenance question.

@chengchen-google
chengchen-google force-pushed the 0825/0001_tcq_block_level_avx2_fusion branch from 034fae7 to c4c1a8a Compare August 27, 2026 18:38
@urvangjoshi

Copy link
Copy Markdown
Contributor

@chengchen-google : This is the unit test that is failing:

[ RUN      ] AVX2/TcqLoopDiagonalSt8Test.RandomValues/0
../../test/trellis_test.cc:1117: Failure
Expected equality of these values:
  rate_ref
    Which is: -1099013650
  rate_tst
    Which is: 193634

[  FAILED  ] AVX2/TcqLoopDiagonalSt8Test.RandomValues/0, where GetParam() = bit_depth:0 function:0x59d01f66a390 function:0x59d01fafb590 (0 ms)
[----------] 1 test from AVX2/TcqLoopDiagonalSt8Test (0 ms total)

For reference, this section of wiki page explains how to find the failing test: https://github.com/AOMediaCodec/avm/wiki/Reproducing-CI-Test-Failures-Locally#finding-failing-tests

Previously, TCQ evaluated coefficients along anti-diagonal scan lines
in a portable C loop (trellis_loop_diagonal_st8), invoking individual
fine-grained RTCD SIMD functions per coefficient.

This fine-grained dispatch structure incurred significant overhead:
- Repeated RTCD function calls and indirect branches on every coeff.
- Register spills to memory structures across subroutine boundaries.
- Optimization barriers preventing compiler vector scheduling across
  the full diagonal loop.

This patch introduces a coarse-grained, block-level RTCD entry point:
  av2_trellis_loop_diagonal_st8()
with AVX2 specialization (av2_trellis_loop_diagonal_st8_avx2).

Updates are now executed contiguously within AVX2 registers without
exiting to scalar C code per coefficient.

Unit tests are added to compare the the new and original implementation.

We can see that the TCQ function becomes 3-10% faster.

=====================================================================
TX Size   Pure C (us)   Base AVX2 (us)   Patch AVX2 (us)  vs Base AVX2
---------------------------------------------------------------
4x4       110313        45646            44360           1.03x (+2.8%)
8x8       369747        102636           96537           1.06x (+5.9%)
16x16     343618        81354            74201           1.10x (+8.8%)
32x32     268794        60375            54646           1.10x (+9.5%)
4x8       195707        64939            61781           1.05x (+4.9%)
8x4       201683        65227            61872           1.05x (+5.1%)
8x16      175484        44150            40669           1.09x (+7.9%)
16x8      180465        44455            40943           1.09x (+7.9%)
16x32     134924        30975            28011           1.11x (+9.6%)
32x16     134311        30946            28192           1.10x (+8.9%)
4x16      364581        103343           96494           1.07x (+6.6%)
16x4      390933        104326           97702           1.07x (+6.3%)
8x32      343005        81586            74379           1.10x (+8.8%)
32x8      352547        81916            74820           1.09x (+8.7%)
=====================================================================

Change-Id: Ib14c7e995e37e4c0af0750af50c78db3d2adb3b5
@chengchen-google
chengchen-google force-pushed the 0825/0001_tcq_block_level_avx2_fusion branch from c4c1a8a to 7405cad Compare August 28, 2026 18:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants