Gate partition ML/pruning in fast two pass partitions by yeqing-wu · Pull Request #5317 · AOMediaCodec/avm

yeqing-wu · 2026-08-25T19:32:32Z

This newly activates fast two pass partition for speed=2, while improving speed-quality ratio for speed=3/4. At speed=2 it is turned on only for class A1 (2160p and above). Class A2 (2K) lost coding efficiency at speed=2, so we leave it off there. Speeds 3 and 4 are not changed and stay on for every resolution.

Some of the partition ML/pruning schemes interact poorly with the fast two-pass partition search, so skip them in the dry pass and on trusted wet-pass blocks. The full policy across regimes:

FAST dry pass: disabled. Shape is picked from reduced-tool
RD, and the dry pass is already fast.
FAST wet pass, forced: disabled. Partition is already fixed by the
dry-pass template; the decision is discarded.
FAST wet pass, reopened: enabled. Block is genuinely searched.
CONSERVATIVE / one-pass: enabled. Predicate is inert on this path.

Anchor: commit 7361d54
Speed 4 (cpu-used=4): FG16 CTC (33 frames, class A1 and A2, RA)
Speed 3 (cpu-used=3): FG16 CTC (33 frames, class A1 and A2, RA)
Speed 2 (cpu-used=2): FG16 CTC (33 frames, class A1 and A2, RA)

1) Speed 4 (cpu-used=4)
  +------------+-------+-------+-------+-------+------+------+
  | Class      |     Y |    Cb |    Cr |  wAvg | Enc% | Dec% |
  +------------+-------+-------+-------+-------+------+------+
  | A1         | -0.73 |  0.26 | -0.05 | -0.66 |  106 |  101 |
  | A2         | -0.49 |  0.22 |  0.20 | -0.43 |  105 |  101 |
  | Avg w/o B2 | -0.56 |  0.23 |  0.12 | -0.50 |  105 |  101 |
  +------------+-------+-------+-------+-------+------+------+

2) Speed 3 (cpu-used=3)
  +------------+-------+-------+-------+-------+------+------+
  | Class      |     Y |    Cb |    Cr |  wAvg | Enc% | Dec% |
  +------------+-------+-------+-------+-------+------+------+
  | A1         | -0.75 |  0.08 |  0.17 | -0.68 |  106 |  101 |
  | A2         | -0.47 |  0.18 | -0.15 | -0.43 |  104 |  100 |
  | Avg w/o B2 | -0.55 |  0.15 | -0.05 | -0.51 |  105 |  100 |
  +------------+-------+-------+-------+-------+------+------+

3) Speed 2 (cpu-used=2)
  +------------+-------+-------+-------+-------+------+------+
  | Class      |     Y |    Cb |    Cr |  wAvg | Enc% | Dec% |
  +------------+-------+-------+-------+-------+------+------+
  | A1         | -0.50 | -2.10 | -2.72 | -0.65 |   84 |   87 |
  | A2         |  0.46 |  1.97 |  1.60 |  0.56 |   83 |   95 |
  | Avg w/o B2 |  0.17 |  0.76 |  0.32 |  0.20 |   84 |   93 |
  +------------+-------+-------+-------+-------+------+------+

STATS_CHANGED

urvangjoshi · 2026-08-25T20:40:25Z

@yeqing-wu : I have a general question related to the fast two pass method:
Have you tested to make sure it works as expected when extended partitions (H partitions, for example) are turned on?

Reason for asking:

On top of commit 44be072 (today), if I try to enable extended partitions at all speeds:
- At speed 4, seeing a loss (e.g. 0.43% loss on A5 set) instead of gain (despite 119% enc-time) <-- unexpected
- At speed 2, seeing a gain as expected

Given that fast two pass is enabled for speed >= 3 on that commit, I'm trying to understand something isn't fundamentally broken with this feature in combination with extended partitions.

It's more important to double check this, as you move the feature to speed 2.

Note: On commit 2abcfff (~1 week ago), the behavior was as expected (gain at both speed 2 and 4)

yeqing-wu · 2026-08-25T20:51:23Z

Hi @urvangjoshi, thanks for your information. I haven’t conducted any special tests with the extended partitions turned on. All my tests are based on the default speed features in the av2-enc branch.

Before I run the tests, let me confirm a few things:

Is my understanding correct that the extended partitions are turned off for speed >= 2 by sf->part_sf.disable_ext_partitions?
How much gain did you observe in speed 2 when you turned on the extended partitions?
Did you remember the gain of the extended partitions on speed >= 3 before applying the fast two-pass partition if you already did that?

Thanks
Yeqing

urvangjoshi · 2026-08-25T20:58:50Z

Hi @urvangjoshi, thanks for your information. I haven’t conducted any special tests with the extended partitions turned on. All my tests are based on the default speed features in the av2-enc branch.

Before I run the tests, let me confirm a few things:

Is my understanding correct that the extended partitions are turned off for speed >= 2 by sf->part_sf.disable_ext_partitions?

Correct. For my tests, I basically moved this speed feature to speed 5 (to test gains on speed 2, 3 and 4).

How much gain did you observe in speed 2 when you turned on the extended partitions?

Did you remember the gain of the extended partitions on speed >= 3 before applying the fast two-pass partition if you already did that?

The coding gains when enabling ext-partitions (on top of commit 2abcfff) were as follows (33 frames RA, all classes without B2):

Speed 4: -0.43% at 118% enc-time
Speed 3: -0.46% at 126% enc-time
Speed 2: -0.45% at 127% enc-time

The ratio isn't good enough to directly make this change, but I'm thinking about possible pruning methods to get some of the gains.

Thanks Yeqing

yeqing-wu · 2026-08-25T21:04:04Z

Got it. I can run some tests with sf->part_sf.disable_ext_partitions = false.

By the way, when part_sf.disable_ext_partitions is set to true, how many partition shapes would be disabled? There’s no comment in the code explaining this, so could you please elaborate on it? Does it disable all 3-ways and 4-ways partitions?

urvangjoshi · 2026-08-25T21:08:27Z

Got it. I can run some tests with sf->part_sf.disable_ext_partitions = false.

By the way, when part_sf.disable_ext_partitions is set to true, how many partition shapes would be disabled? There’s no comment in the code explaining this, so could you please elaborate on it? Does it disable all 3-ways and 4-ways partitions?

This will only enable the HORZ_3 / VERT_3 (that is, H shaped) partitions.
[The 4A/4B partitions are disabled for speed >= 1]

yeqing-wu · 2026-08-27T00:21:54Z

@urvangjoshi, after reading the code of the extended partition, I now understand why disable_ext_partitions = false doesn't help and in fact hurts performance. The same issue affects both the slow and the fast two-pass partition. The reasons are as follows:

Extended partitions are only allowed for block sizes between 16x16 and 64x64 inclusive. For all other block sizes they are forbidden.
In the two-pass partition (both slow and fast), the minimum block size in the dry pass is 16x16. Combined with (1), this means only 64x64 blocks can use extended partitions in the dry pass. And because a dry-pass partition decision
is reused directly in the wet pass whenever the block size is >= 32x32, extended partitions get very few opportunities to be evaluated: only when the parent block is 64x64.
For block sizes <= 32x32 the wet pass re-opens the partition search, so extended partitions can be evaluated there — but those opportunities are also rare.

So extended partitions are exercised very rarely, while enabling them adds block-level signaling cost. The cost of signaling the partition shape isn't offset by the residual savings. This is naturally even worse in the fast
two-pass partition, since it evaluates only a very limited set of basic coding tools in the dry pass.

  // Extended partition is disabled above BLOCK_64X64 to avoid crossing the
  // 64X64 boundary.
  if (bsize > BLOCK_64X64 && bsize <= BLOCK_LARGEST) {
    return false;
  }
  // At bsize <= 8X8, extended partitions will lead to dimension < 2.
  if (bsize <= BLOCK_8X8) {
    return false;
  }

yeqing-wu · 2026-08-27T00:35:19Z

BTW, on top of this PR, set disable_ext_partitions = false, the results are:

Speed 2:

Speed 3:

Speed 4:

urvangjoshi · 2026-08-27T00:47:55Z

@urvangjoshi, after reading the code of the extended partition, I now understand why disable_ext_partitions = false doesn't help and in fact hurts performance. The same issue affects both the slow and the fast two-pass partition. The reasons are as follows:

Extended partitions are only allowed for block sizes between 16x16 and 64x64 inclusive. For all other block sizes they are forbidden.

In the two-pass partition (both slow and fast), the minimum block size in the dry pass is 16x16. Combined with (1), this means only 64x64 blocks can use extended partitions in the dry pass. And because a dry-pass partition decision
is reused directly in the wet pass whenever the block size is >= 32x32, extended partitions get very few opportunities to be evaluated: only when the parent block is 64x64.

For block sizes <= 32x32 the wet pass re-opens the partition search, so extended partitions can be evaluated there — but those opportunities are also rare.

So extended partitions are exercised very rarely, while enabling them adds block-level signaling cost. The cost of signaling the partition shape isn't offset by the residual savings. This is naturally even worse in the fast two-pass partition, since it evaluates only a very limited set of basic coding tools in the dry pass.
  // Extended partition is disabled above BLOCK_64X64 to avoid crossing the
  // 64X64 boundary.
  if (bsize > BLOCK_64X64 && bsize <= BLOCK_LARGEST) {
    return false;
  }
  // At bsize <= 8X8, extended partitions will lead to dimension < 2.
  if (bsize <= BLOCK_8X8) {
    return false;
  }

Thanks for digging through the details.
So, looks like 2-pass and extended partitions (H partitions) are incompatible.
It will be even worse for uneven 4-way in fact.

Anyway, we'll have to keep this in mind going forward.
Basically, for speed >= 2, to enable extended partitions, we'll have to disable / redesign the 2 pass partitioning.

Can you add a comment for this in the code -- wherever you feel appropriate?

yeqing-wu · 2026-08-27T01:48:44Z

@urvangjoshi sure, add these notes above the definition of disable_ext_partitions.
If you can achieve a good tradeoff for the extended partitions, I'm happy to work with you on making it work well in the two-pass partition flow too. I have a few ideas, though I'm not sure if they can work.

// Notes: Setting disable_ext_partitions = false (i.e. enabling extended
// partitions) for speed >= 2 doesn't improve coding efficiency and in fact
// hurts performance. It interacts poorly with the two-pass partition, both
// slow and fast. The reasons are as follows:
//
// 1. Extended partitions are only allowed for block sizes between 16x16 and
//    64x64 inclusive. For all other block sizes they are forbidden.
// 2. In the two-pass partition (both slow and fast), the minimum block size
//    in the dry pass is 16x16. Combined with (1), this means only 64x64
//    blocks can use extended partitions in the dry pass. And because a
//    dry-pass partition decision is reused directly in the wet pass whenever
//    the block size is >= 32x32, extended partitions get very few
//    opportunities to be evaluated: only when the parent block is 64x64.
// 3. For block sizes <= 32x32 the wet pass re-opens the partition search, so
//    extended partitions can be evaluated there, but those opportunities are
//    also rare.
//
// So extended partitions are exercised very rarely once the two-pass
// partition is enabled, while still adding block-level signaling cost. The
// cost of signaling the partition shape isn't offset by the residual
// savings. This is naturally even worse in the fast two-pass partition,
// since it evaluates only a very limited set of basic coding tools in the
// dry pass.
bool disable_ext_partitions;

yeqing-wu · 2026-08-27T16:52:17Z

Hi All, if you have time, could you please help review this PR?

urvangjoshi · 2026-08-27T21:17:33Z

@yeqing-wu : Looking at the speed 2 numbers:

 3) Speed 2 (cpu-used=2)
  +------------+-------+-------+-------+-------+------+------+
  | Class      |     Y |    Cb |    Cr |  wAvg | Enc% | Dec% |
  +------------+-------+-------+-------+-------+------+------+
  | A1         | -0.50 | -2.10 | -2.72 | -0.65 |   84 |   87 |
  | A2         |  0.46 |  1.97 |  1.60 |  0.56 |   83 |   95 |
  | Avg w/o B2 |  0.17 |  0.76 |  0.32 |  0.20 |   84 |   93 |
  +------------+-------+-------+-------+-------+------+------+

The ratio for A2 is ~30, which seems borderline.
Resolution gain/loss swings are a bit concerning to me.

Meanwhile, the benefit for speed >= 3 is clear.
So, I'd suggest keeping the speed feature for speed >= 3 for now, until we can improve it for speed 2.

[Note that speed 2 is already much worse than speed 1, btw].

Feel free to chime in @leolzhao @yunqingwang1

yeqing-wu · 2026-08-27T21:30:56Z

@yeqing-wu : Looking at the speed 2 numbers:
 3) Speed 2 (cpu-used=2)
  +------------+-------+-------+-------+-------+------+------+
  | Class      |     Y |    Cb |    Cr |  wAvg | Enc% | Dec% |
  +------------+-------+-------+-------+-------+------+------+
  | A1         | -0.50 | -2.10 | -2.72 | -0.65 |   84 |   87 |
  | A2         |  0.46 |  1.97 |  1.60 |  0.56 |   83 |   95 |
  | Avg w/o B2 |  0.17 |  0.76 |  0.32 |  0.20 |   84 |   93 |
  +------------+-------+-------+-------+-------+------+------+
The ratio for A2 is ~30, which seems borderline. Resolution gain/loss swings are a bit concerning to me.

Meanwhile, the benefit for speed >= 3 is clear. So, I'd suggest keeping the speed feature for speed >= 3 for now, until we can improve it for speed 2.

[Note that speed 2 is already much worse than speed 1, btw].

Feel free to chime in @leolzhao @yunqingwang1

@leolzhao @yunqingwang1 @jingninghan, what are your thoughts? The overall tradeoff is quite good, with a coding loss of 0.2% and a significant reduction in encoding time of 16%. This results in a ratio of 80x, which is much higher than the required minimum bar 30x ratio.

If you really have concerns about the trade-off of class A2 in speed 2, another solution is to apply this only to class A1 for speed 2. It offers a -0.65% coding gain and still achieves a 16% reduction in encoding time, which are both significant advantages for class A1 in speed 2.

Thanks
Yeqing

urvangjoshi · 2026-08-27T21:36:43Z

If you really have concerns about the trade-off of class A2 in speed 2, another solution is to apply this only to class A1 for speed 2. It offers a -0.65% coding gain and still achieves a 16% reduction in encoding time, which are both significant advantages for class A1 in speed 2.

Enabling for A1 only at speed 2 sounds good!

yeqing-wu · 2026-08-27T21:44:16Z

Let’s wait for others’ comments. If they’re okay with only enabling class A1 in speed 2, I’m fine with that. My only concern is that having so many customized settings for different classes would make the encoder extremely complex, but this is just my personal opinion.

leolzhao · 2026-08-28T16:24:23Z

@yeqing-wu : Looking at the speed 2 numbers:
 3) Speed 2 (cpu-used=2)
  +------------+-------+-------+-------+-------+------+------+
  | Class      |     Y |    Cb |    Cr |  wAvg | Enc% | Dec% |
  +------------+-------+-------+-------+-------+------+------+
  | A1         | -0.50 | -2.10 | -2.72 | -0.65 |   84 |   87 |
  | A2         |  0.46 |  1.97 |  1.60 |  0.56 |   83 |   95 |
  | Avg w/o B2 |  0.17 |  0.76 |  0.32 |  0.20 |   84 |   93 |
  +------------+-------+-------+-------+-------+------+------+
The ratio for A2 is ~30, which seems borderline. Resolution gain/loss swings are a bit concerning to me.

Meanwhile, the benefit for speed >= 3 is clear. So, I'd suggest keeping the speed feature for speed >= 3 for now, until we can improve it for speed 2.

[Note that speed 2 is already much worse than speed 1, btw].

Feel free to chime in @leolzhao @yunqingwang1

I agree with Urvang that we can turn on this speed feature for speed >=3 for now. If the trade off for speed = 2 is further improved, then we can turn it on.

yeqing-wu · 2026-08-28T16:36:30Z

@yeqing-wu : Looking at the speed 2 numbers:
 3) Speed 2 (cpu-used=2)
  +------------+-------+-------+-------+-------+------+------+
  | Class      |     Y |    Cb |    Cr |  wAvg | Enc% | Dec% |
  +------------+-------+-------+-------+-------+------+------+
  | A1         | -0.50 | -2.10 | -2.72 | -0.65 |   84 |   87 |
  | A2         |  0.46 |  1.97 |  1.60 |  0.56 |   83 |   95 |
  | Avg w/o B2 |  0.17 |  0.76 |  0.32 |  0.20 |   84 |   93 |
  +------------+-------+-------+-------+-------+------+------+
The ratio for A2 is ~30, which seems borderline. Resolution gain/loss swings are a bit concerning to me.
Meanwhile, the benefit for speed >= 3 is clear. So, I'd suggest keeping the speed feature for speed >= 3 for now, until we can improve it for speed 2.
[Note that speed 2 is already much worse than speed 1, btw].
Feel free to chime in @leolzhao @yunqingwang1
I agree with Urvang that we can turn on this speed feature for speed >=3 for now. If the trade off for speed = 2 is further improved, then we can turn it on.

Okay, then I’ll turn it on only for class A1 in speed 2, as it offers a coding gain of -0.65% and saves 16% of encoding time. For class A2 in speed 2, given the concerns about the tradeoff, I’ll disable it and save it for future improvements.

yeqing-wu · 2026-08-28T16:52:48Z

Done. Add the following constraints:

  if (oxcf->speed == 2) {
    const int is_2160p_or_larger =
        AVMMIN(oxcf->frm_dim_cfg.width, oxcf->frm_dim_cfg.height) >= 2160;
    return is_2160p_or_larger;
  }

@urvangjoshi and @leolzhao, could you please help review it?

urvangjoshi · 2026-08-28T16:55:39Z

Done. Add the following constraints:
  if (oxcf->speed == 2) {
    const int is_2160p_or_larger =
        AVMMIN(oxcf->frm_dim_cfg.width, oxcf->frm_dim_cfg.height) >= 2160;
    return is_2160p_or_larger;
  }
@urvangjoshi and @leolzhao, could you please help review it?

Looks good. Pls update the PR description as well.

yeqing-wu · 2026-08-28T16:55:42Z

Also updated the commit message to reflect this constraint on speed 2.

yeqing-wu · 2026-08-28T16:58:38Z

@urvangjoshi, I have a question. Is there a way to automatically sync the PR conversation message with the change in the commit message? If not, it would be tedious to manually modify the PR conversation message with the commit message.

urvangjoshi · 2026-08-28T17:06:00Z

@urvangjoshi, I have a question. Is there a way to automatically sync the PR conversation message with the change in the commit message? If not, it would be tedious to manually modify the PR conversation message with the commit message.

I don't believe so.

yeqing-wu · 2026-08-28T21:52:19Z

The pipeline tests are queueing now, do I need to do anything to proceed them?

urvangjoshi · 2026-08-28T21:54:55Z

The pipeline tests are queueing now, do I need to do anything to proceed them?

We have to just wait unfortunately (you may have seen my email about 4 runners causing the queueing).

This newly activates fast two pass partition for speed=2, while improving speed-quality ratio for speed=3/4. At speed=2 it is turned on only for class A1 (2160p and above). Class A2 (2K) lost coding efficiency at speed=2, so we leave it off there. Speeds 3 and 4 are not changed and stay on for every resolution. Some of the partition ML/pruning schemes interact poorly with the fast two-pass partition search, so skip them in the dry pass and on trusted wet-pass blocks. The full policy across regimes: * FAST dry pass: disabled. Shape is picked from reduced-tool RD, and the dry pass is already fast. * FAST wet pass, forced: disabled. Partition is already fixed by the dry-pass template; the decision is discarded. * FAST wet pass, reopened: enabled. Block is genuinely searched. * CONSERVATIVE / one-pass: enabled. Predicate is inert on this path. Anchor: commit 7361d54 Speed 4 (cpu-used=4): FG16 CTC (33 frames, class A1 and A2, RA) Speed 3 (cpu-used=3): FG16 CTC (33 frames, class A1 and A2, RA) Speed 2 (cpu-used=2): FG16 CTC (33 frames, class A1 and A2, RA) ``` 1) Speed 4 (cpu-used=4) +------------+-------+-------+-------+-------+------+------+ | Class | Y | Cb | Cr | wAvg | Enc% | Dec% | +------------+-------+-------+-------+-------+------+------+ | A1 | -0.73 | 0.26 | -0.05 | -0.66 | 106 | 101 | | A2 | -0.49 | 0.22 | 0.20 | -0.43 | 105 | 101 | | Avg w/o B2 | -0.56 | 0.23 | 0.12 | -0.50 | 105 | 101 | +------------+-------+-------+-------+-------+------+------+ 2) Speed 3 (cpu-used=3) +------------+-------+-------+-------+-------+------+------+ | Class | Y | Cb | Cr | wAvg | Enc% | Dec% | +------------+-------+-------+-------+-------+------+------+ | A1 | -0.75 | 0.08 | 0.17 | -0.68 | 106 | 101 | | A2 | -0.47 | 0.18 | -0.15 | -0.43 | 104 | 100 | | Avg w/o B2 | -0.55 | 0.15 | -0.05 | -0.51 | 105 | 100 | +------------+-------+-------+-------+-------+------+------+ 3) Speed 2 (cpu-used=2) +------------+-------+-------+-------+-------+------+------+ | Class | Y | Cb | Cr | wAvg | Enc% | Dec% | +------------+-------+-------+-------+-------+------+------+ | A1 | -0.50 | -2.10 | -2.72 | -0.65 | 84 | 87 | | A2 | 0.46 | 1.97 | 1.60 | 0.56 | 83 | 95 | | Avg w/o B2 | 0.17 | 0.76 | 0.32 | 0.20 | 84 | 93 | +------------+-------+-------+-------+-------+------+------+ ``` STATS_CHANGED

aom-bot added the cla-signed label Aug 25, 2026

yeqing-wu requested review from jingninghan, leolzhao, urvangjoshi and yunqingwang1 August 25, 2026 19:33

yeqing-wu force-pushed the sb_tpp_ml_pruning_strategy_submission_merge branch from 29862b1 to f2ccd3e Compare August 27, 2026 01:45

yeqing-wu force-pushed the sb_tpp_ml_pruning_strategy_submission_merge branch from f2ccd3e to f8ae855 Compare August 28, 2026 16:51

urvangjoshi approved these changes Aug 28, 2026

View reviewed changes

urvangjoshi force-pushed the sb_tpp_ml_pruning_strategy_submission_merge branch from f8ae855 to a51fc80 Compare August 28, 2026 23:10

urvangjoshi merged commit 22008a2 into AOMediaCodec:av2-enc Aug 28, 2026
8 checks passed

Uh oh!

Conversation

yeqing-wu commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

urvangjoshi commented Aug 25, 2026

Uh oh!

yeqing-wu commented Aug 25, 2026

Uh oh!

urvangjoshi commented Aug 25, 2026

Uh oh!

yeqing-wu commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

urvangjoshi commented Aug 25, 2026

Uh oh!

yeqing-wu commented Aug 27, 2026

Uh oh!

yeqing-wu commented Aug 27, 2026

Uh oh!

urvangjoshi commented Aug 27, 2026

Uh oh!

yeqing-wu commented Aug 27, 2026

Uh oh!

yeqing-wu commented Aug 27, 2026

Uh oh!

urvangjoshi commented Aug 27, 2026

Uh oh!

yeqing-wu commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

urvangjoshi commented Aug 27, 2026

Uh oh!

yeqing-wu commented Aug 27, 2026

Uh oh!

leolzhao commented Aug 28, 2026

Uh oh!

yeqing-wu commented Aug 28, 2026

Uh oh!

yeqing-wu commented Aug 28, 2026

Uh oh!

urvangjoshi commented Aug 28, 2026

Uh oh!

yeqing-wu commented Aug 28, 2026

Uh oh!

yeqing-wu commented Aug 28, 2026

Uh oh!

urvangjoshi commented Aug 28, 2026

Uh oh!

yeqing-wu commented Aug 28, 2026

Uh oh!

urvangjoshi commented Aug 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yeqing-wu commented Aug 25, 2026 •

edited

Loading

yeqing-wu commented Aug 25, 2026 •

edited

Loading

yeqing-wu commented Aug 27, 2026 •

edited

Loading