From de8c22c9d52a65b809c928f76937bfb775e7dfbd Mon Sep 17 00:00:00 2001 From: Gunse11er <872562242@qq.com> Date: Tue, 11 Aug 2026 17:25:19 +0800 Subject: [PATCH] [BugFix][TE] Initialize nested reductions at the outermost reduction scope Nested reductions could initialize the same accumulator at multiple dependency levels, discarding partial results. Track the outermost reduction scope and emit initialization only once. Add numerical regressions for both mixed-level reduction orders. --- 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 7a8c4ae80b18..daa1b712eba8 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 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 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 annotations = GenerateBlockAnnotations(compute_op, info); const ReduceNode* reduce = compute_op->body[0].as(); + if (reduce) { PrimExpr expr_body = compute_op->body[0]; - Stmt init = GenerateInitStmt(leaf.store_indices, buffers, reduce, leaf.axes_remap, info); + ffi::Optional 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 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 38eef4f79561..9a879e6701b4 100644 --- a/tests/python/te/test_te_create_primfunc.py +++ b/tests/python/te/test_te_create_primfunc.py @@ -955,6 +955,28 @@ def te_workload(): _check_workload(te_workload, tir_workload) +@pytest.mark.parametrize( + ("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 tir_workload( 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]