This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 3dea168fe2 [BugFix][TE] Initialize nested reductions at the outermost
reduction scope (#20116)
3dea168fe2 is described below
commit 3dea168fe20d6ea5c0d42de2b3fb4ad021fe5cf1
Author: Gunse11er <[email protected]>
AuthorDate: Wed Aug 19 09:46:23 2026 +0800
[BugFix][TE] Initialize nested reductions at the outermost reduction scope
(#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
```
---
src/te/operation/create_primfunc.cc | 17 +++++++++++------
tests/python/te/test_te_create_primfunc.py | 29 ++++++++++++++++++++++++++---
2 files changed, 37 insertions(+), 9 deletions(-)
diff --git a/src/te/operation/create_primfunc.cc
b/src/te/operation/create_primfunc.cc
index 7a8c4ae80b..daa1b712eb 100644
--- a/src/te/operation/create_primfunc.cc
+++ b/src/te/operation/create_primfunc.cc
@@ -488,6 +488,8 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp&
compute_op, CreateFuncInfo* in
TVM_FFI_ICHECK(!axes_levels.empty());
std::vector<NestedScopeInfo> scopes;
scopes.reserve(axes_levels.size());
+ // Initialize a nested reduction at its outermost reduction level.
+ size_t reduction_init_scope = axes_levels.size() - 1;
std::unordered_set<Var> defined_axes;
for (size_t i = 0; i < axes_levels.size(); ++i) {
NestedScopeInfo cur_scope;
@@ -498,6 +500,9 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp&
compute_op, CreateFuncInfo* in
bool first_times_define =
std::find(axes_levels[i].begin(), axes_levels[i].end(), axis) !=
axes_levels[i].end();
if (first_times_define) {
+ if (axis->iter_type == IterVarType::kCommReduce) {
+ reduction_init_scope = std::min(reduction_init_scope, i);
+ }
Var loop_var = Var(axis->var->name, index_type);
Var block_var("v_" + axis->var->name, index_type);
PrimExpr min = axis->dom->min;
@@ -541,9 +546,13 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp&
compute_op, CreateFuncInfo* in
auto leaf = scopes.back();
ffi::Map<ffi::String, ffi::Any> annotations =
GenerateBlockAnnotations(compute_op, info);
const ReduceNode* reduce = compute_op->body[0].as<ReduceNode>();
+
if (reduce) {
PrimExpr expr_body = compute_op->body[0];
- Stmt init = GenerateInitStmt(leaf.store_indices, buffers, reduce,
leaf.axes_remap, info);
+ ffi::Optional<Stmt> init{std::nullopt};
+ if (reduction_init_scope == scopes.size() - 1) {
+ init = GenerateInitStmt(leaf.store_indices, buffers, reduce,
leaf.axes_remap, info);
+ }
Stmt body =
GenerateBodyStmt(leaf.store_indices, buffers, leaf.axes_remap,
expr_body, info, analyzer);
seq_stmt.push_back(SBlockRealize(/*iter_values=*/leaf.bindings,
@@ -592,11 +601,7 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp&
compute_op, CreateFuncInfo* in
const auto& block_iters = cur.block_iters;
ffi::Optional<Stmt> init{std::nullopt};
- if (reduce && std::any_of(block_iters.begin(), block_iters.end(),
[](const IterVar& iter) {
- return iter->iter_type == IterVarType::kCommReduce;
- })) {
- // if the reduce axis defined in non-leaf scopes, the nested block is
also
- // a reduction block, thus we should also insert init stmt in the
parent level.
+ if (reduce && i - 1 == reduction_init_scope) {
init = GenerateInitStmt(cur.store_indices, buffers, reduce,
cur.axes_remap, info);
}
diff --git a/tests/python/te/test_te_create_primfunc.py
b/tests/python/te/test_te_create_primfunc.py
index 38eef4f795..9a879e6701 100644
--- a/tests/python/te/test_te_create_primfunc.py
+++ b/tests/python/te/test_te_create_primfunc.py
@@ -955,6 +955,28 @@ def test_adaptive_pooling_window():
_check_workload(te_workload, tir_workload)
[email protected](
+ ("input_shape", "expected"),
+ [
+ ((3, 4), [[12.5, 14.5], [16.5, 18.5]]),
+ ((4, 3), [[12.0, 13.0], [18.0, 19.0]]),
+ ],
+)
+def test_adaptive_pooling_mixed_reduction_levels(input_shape, expected):
+ data = te.placeholder((1, 1, *input_shape), "float32", "data")
+ output = topi.nn.adaptive_pool(data, [2, 2], pool_type="avg")
+ prim_func = te.create_prim_func([data, output])
+ compiled = tvm.compile(prim_func)
+
+ input_data = np.arange(10, 10 + np.prod(input_shape),
dtype="float32").reshape(
+ 1, 1, *input_shape
+ )
+ actual = tvm.runtime.tensor(np.empty((1, 1, 2, 2), dtype="float32"))
+ compiled(tvm.runtime.tensor(input_data), actual)
+
+ tvm.testing.assert_allclose(actual.numpy()[0, 0], np.array(expected,
dtype="float32"))
+
+
def test_global_pool():
# fix the issue-17938
data = te.placeholder((1, 1, 32, 32), dtype="int8", name="data")
@@ -991,10 +1013,11 @@ def test_nested_reduce_domain_dependency():
v_i2_2 = T.axis.spatial((v_i2_1, v_i2_1 + 1),
v_i2_1)
v_rv_1 = T.axis.reduce((v_rv, v_rv + 1), v_rv)
v_rv_2 = T.axis.reduce(v_rv, rv_1)
- T.reads(x[v_i0_2, v_i1_2, v_i2_2, v_rv_1,
v_rv_2])
+ T.reads(
+ compute[v_i0_2, v_i1_2, v_i2_2],
+ x[v_i0_2, v_i1_2, v_i2_2, v_rv_1, v_rv_2],
+ )
T.writes(compute[v_i0_2, v_i1_2, v_i2_2])
- with T.init():
- compute[v_i0_2, v_i1_2, v_i2_2] =
T.float32(0.0)
compute[v_i0_2, v_i1_2, v_i2_2] = (
compute[v_i0_2, v_i1_2, v_i2_2]
+ x[v_i0_2, v_i1_2, v_i2_2, v_rv_1, v_rv_2]
