Gunse11er opened a new pull request, #20116:
URL: https://github.com/apache/tvm/pull/20116
# PR title
`[BugFix][TE] Initialize nested reductions at the outermost reduction scope`
# PR body
## Problem
Fixes #20105.
`te.create_prim_func` may place reduction iterators at different nested
block levels when their domains have different dependencies. The previous
lowering initialized the same accumulator in both the leaf reduction block and
every parent block that introduced a reduction iterator. When reduction axes
spanned multiple levels, the inner `T.init` reset partial results produced by
the outer level.
For adaptive average pooling from `3x4` to `2x2`, the incorrect lowering
contained two `T.init` regions and produced:
```text
actual = [[6.5, 7.5],
[8.5, 9.5]]
expected = [[12.5, 14.5],
[16.5, 18.5]]
max absolute difference = 9.0
```
The first output is `(11 + 15) / 4 = 6.5` instead of `(10 + 11 + 14 + 15) /
4 = 12.5`, showing that the inner initialization discarded the first reduction
slice.
## Root cause
`GenerateStmtFromCompute` unconditionally generated an initialization for a
leaf reduction block. It also generated an initialization for every parent
scope containing a reduction iterator. These conditions overlap when reduction
iterators are first defined at different nesting levels, causing duplicate
initialization of one logical reduction.
## Fix
Record the outermost scope that first defines a commutative reduction
iterator during the existing axis-definition pass, and generate the reduction
initialization only at that scope. Nested blocks below it continue reading and
updating the same accumulator without reinitializing it.
This is a general TE lowering fix; it does not special-case adaptive pooling
or Relax.
## Regression coverage
The new numerical regression compiles and executes adaptive average pooling
for both mixed-level orders:
- `3x4 -> 2x2`, where the height reduction has a dependent extent.
- `4x3 -> 2x2`, where the width reduction has a dependent extent.
The same test file was run against an independent build of the base commit
and against the fixed build:
```text
base bb9bc20a: 2 failed in 6.99s
fixed bb9bc20a: 2 passed in 7.77s
```
The base failures were numerical mismatches with maximum absolute
differences of `9.0` and `8.75`. Both fixed cases have zero error.
The generated-IR and numerical trigger matrix is:
| Input -> output | Base `T.init` | Base max diff | Fixed `T.init` | Fixed
max diff |
| --- | ---: | ---: | ---: | ---: |
| `4x4 -> 2x2` | 1 | 0.0 | 1 | 0.0 |
| `3x4 -> 2x2` | 2 | 9.0 | 1 | 0.0 |
| `4x3 -> 2x2` | 2 | 8.75 | 1 | 0.0 |
| `3x3 -> 2x2` | 1 | 0.0 | 1 | 0.0 |
This also confirms that the change preserves cases whose reduction iterators
already occupy a single level.
## Testing
Focused regression:
```bash
python -m pytest -q \
tests/python/te/test_te_create_primfunc.py::test_adaptive_pooling_mixed_reduction_levels
```
Relevant TE and Relax suites:
```bash
python -m pytest -q \
tests/python/te/test_te_create_primfunc.py \
tests/python/relax/test_transform_legalize_ops_nn.py
```
```text
base: 3 failed, 111 passed, 3 skipped, 15 warnings in 7.78s
fixed: 114 passed, 3 skipped, 15 warnings in 8.00s
```
The three base-only failures are exactly the two new numerical cases and the
updated nested-reduction structural expectation. The warnings are identical
between both runs and pre-existing.
The original Relax/LLVM reproducer from #20105 was also run end to end:
```text
base: max absolute difference = 9.0
fixed: max absolute difference = 0.0
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
