[BugFix] Support keep-dim destinations in fragment reduce

[li-ruinan]

[BugFix] Support keep-dim destinations in fragment reduce

Summary

Fixes #2530.

T.reduce_sum / T.reduce_max / etc. crash at compile time when the
destination buffer keeps the reduced axis with extent 1 (e.g. reducing
(64, 64) along dim=1 into a (64, 1) fragment):

Check failed: vars.size() == InputDim() (2 vs. 1)

This shape is part of the documented frontend contract — tilelang/language/reduce_op.py:50-56
explicitly accepts both [X, Y] and [X, 1, Y] output shapes and builds
expected_shapes for both — so this is a backend gap, not user misuse. The fix
teaches both the layout-inference and lowering stages to produce a reducer
layout whose rank matches the keep-dim destination.

Root Cause

ComputeReducerLayout unconditionally drops the reduced axis when building
the reducer fragment, so it always returns a layout of rank n-1 for a rank-n
source. Two places consume it, and both assume the destination is rank n-1.

Stage 1 — layout inference (src/op/reduce.cc)

ReduceOpNode::InferLayout binds that rank-n-1 layout straight to the
destination buffer. When dst has no layout yet it takes an early return,
which skips the rank check further down in the same function — the one place that
would have caught the mismatch. A rank-1 layout gets bound to a rank-2 buffer,
and nothing complains until a much later Forward call trips the invariant guard
at src/layout/layout.cc:929:

ICHECK_EQ(vars.size(), InputDim())

That is where the opaque (2 vs. 1) surfaces. The reported location is the
detection point, not the introduction point.

Stage 2 — lowering (src/backend/common/op/reduce.h)

Fixing stage 1 alone only relocates the crash. The lowering path derives its own
red_layout from the same helper, then feeds one shared dst_vars array to
two different layouts:

auto dst_indices = dst_layout->Forward(dst_vars.Map(...));  // needs dst_dim args
auto red_indices = red_layout->Forward(dst_vars.Map(...));  // needs src_dim args

dst_vars has exactly one length, so red_layout and dst_layout must agree on
rank or one of the two calls fails the same assertion. On top of that,
ICHECK_EQ(src_dim, dst_dim + 1) rejects keep-dim outright (2 vs. 3), and
src_vars.insert(begin() + op.dim, reduce_iv) would push src_vars to rank
n+1 because dst_vars already carries the dummy axis.

Note that LowerLocal (src/backend/common/op/reduce.h:531) already handles
keep-dim correctly for the local scope — it reads op.dst->shape.size() and
branches on dst_dim == src_dim. Only the fragment path was missing it, which
is why the bug is scope-specific.

Changes

ComputeReducerLayout — new bool keepdim = false parameter (both copies:
src/op/reduce.cc:134 and src/backend/common/op/reduce.h:63, which are
byte-for-byte duplicates apart from static vs inline)

• reducer_shape: Set(dim, 1) instead of erase(dim), keeping the axis at
  extent 1 so InputDim() matches the destination rank.
• fwd: InputPlaceholders(n) + Set(dim, FloorMod(rep, ...)) instead of
  InputPlaceholders(n-1) + insert. Non-keep-dim deliberately offsets
  placeholder numbering from position; keep-dim needs them aligned one-to-one,
  and the extent-1 dummy coordinate contributes nothing to the thread map, so
  overwriting it is correct.
• Replication handling (reducer_rep_extent, CondenseReplicateVar,
  BindThreadRange) is untouched: keep-dim and non-keep-dim describe the same
  physical data with the same thread mapping, differing only in how many logical
  coordinates they accept.
• Default argument keeps every existing caller behaviourally unchanged.
• Cross-referencing comments added to both copies so they do not drift apart.

Stage 1 (src/op/reduce.cc:235)

• Derive keepdim from dst->shape.size() == src_layout->InputDim() and
  is_one(dst->shape[this->dim]). The destination buffer is the only carrier
  of this intent(keepdim) — ReduceOpNode has no keepdim field, and the function
  never read dst->shape before.
• Guard with src_layout->InputDim() > 1: a rank-1 source already degenerates to
  [1] via the reducer_shape.empty() fallback and would otherwise be
  misclassified.
• Add ICHECK(is_one(dst->shape[this->dim])) for same-rank destinations
  (src/op/reduce.cc:226), before the keepdim decision. A destination of the
  same rank as the source is only meaningful as a keep-dim reduction, so a
  non-unit reduced axis is a defect in the extent, and this reports it as one.
  Without it such a buffer would take keepdim == false, build a rank-n-1
  layout, and fail the rank check below with "reducer layout rank does not match
  the destination buffer rank" — a message that does not describe the actual
  problem. tilelang/language/reduce_op.py:51-57 already rejects this shape at
  the frontend (expected shapes are [64] or [64, 1]), and every
  _REDUCE_OP_KEY construction site sits behind that check, so the new ICHECK is
  defensive: it keeps the C++ layer self-consistent without assuming that
  frontend guard survives future refactoring.
• Add ICHECK_EQ(reducer_layout->InputDim(), dst->shape.size()) on the free
  binding path (src/op/reduce.cc:242), before the early return, so any future
  rank divergence fails here with a readable message instead of sliding into
  layout.cc. The pre-existing rank check at :250 only covers the
  already-bound path.

Stage 2 (src/backend/common/op/reduce.h:807)

• Same keepdim criterion — including the is_one(op.dst->shape[op.dim]) extent
  check — so both stages cannot disagree. Stage 1 runs first and already
  diagnoses a same-rank destination with a non-unit reduced axis, so lowering
  never sees one.
• New keep-dim branch: ICHECK_EQ(src_dim, dst_dim) plus
  ICHECK(is_one(dst_layout->InputShape()[op.dim])), replacing the
  unconditional ICHECK_EQ(src_dim, dst_dim + 1).
• src_vars: Set(op.dim, reduce_iv) instead of insert under keep-dim.
  src_vars starts as a copy of dst_vars, which already holds the extent-1
  dummy at op.dim; replacing that slot keeps src_vars at source rank, while
  inserting would overshoot to n+1. This mirrors what LowerLocal already does.

Deliberately unchanged

• src/layout/layout.cc:929 — a correct invariant guard. Relaxing it would turn a
  compile error into silent numerical corruption.
• tilelang/language/reduce_op.py — the frontend contract is right; the backend
  just had not honoured it.
• LowerLocal — already correct, and used as the reference for the fragment path.

Tested

• repro.py from the issue: both the keep-dim (64, 1) case and the
  non-keep-dim (64,) control compile and match the torch reference.
• Extra variants, all matching torch: keep-dim on dim=0, reduce_max /
  reduce_min keep-dim, 3D (32, 16, 64) reducing the middle axis into
  (32, 1, 64), plus non-keep-dim controls for each.
• New parametrized regression test test_reduce_keepdim in
  testing/python/language/test_tilelang_language_reduce.py — 8 cases covering
  sum/max/min, dim=0 and dim=1, 2D and 3D, and both clear=True and
  clear=False. The accumulate (clear=False) path is the one that exercises
  red_indices against clear_buffer directly, so it is covered for sum, max and
  min. Each case asserts the output shape as well as the values, so a silent
  rank regression fails too.
• Frontend rejection of the invalid same-rank shape confirmed empirically: a
  (64, 64) source reduced along dim=1 into a (64, 64) destination raises
  ValueError: Invalid reduce output shape ... expected shapes are [64] or [64, 1]
  before reaching the backend.
• test_tilelang_language_reduce.py re-run with TILELANG_DISABLE_CACHE=1 after
  adding the extent check: 0 failed.
• Reduce suites plus testing/python/layout — test_tilelang_language_reduce.py,
  test_tilelang_language_warp_reduce.py,
  test_tilelang_language_reduce_maxmin_nan.py, test_math_bitwise_reduce.py,
  layout/: 0 failed.
• testing/python/issue: 0 failed.
• Both backends recompile clean: the shared header is included by
  src/cuda/op/reduce.cc and src/rocm/op/reduce.cc.

Follow-up (not in this PR)

ComputeReducerLayout and its helper InputPlaceholders exist as identical
copies in src/op/reduce.cc and src/backend/common/op/reduce.h, in different
namespaces (tvm::tl vs tvm::tl::backend::reduce). Merging them would need
src/op/reduce.cc to depend on src/backend/, or a new shared header — a
layering change too broad for a bugfix. Both copies now carry comments pointing
at each other; a dedicated refactor PR could unify them if necessary.

Summary

• Added keep-dimension reducer layout inference and lowering for fragment reductions.
• Preserves reduced axes as extent-1 dimensions and validates destination rank/layout compatibility.
• Keeps existing non-keep-dimension behavior unchanged.
• Added coverage for sum, max, min, 2D/3D inputs, multiple dimensions, and clear/accumulate paths against PyTorch keepdim=True results.

C++ style / lint notes

• The PR changes C++ reduction implementation and API declarations.
• Review docs/developer_guide/cpp_style.md and the warning-only “C++ API Style Audit” CI step for any advisory findings.
• No correctness or build issues are indicated; TLCPP003/TLCPP004 warnings should remain non-blocking unless they identify a concrete API or maintainability risk.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

Keep-dimension reduction support

Layer / File(s)	Summary
Reducer layout inference src/op/reduce.cc, src/backend/common/op/reduce.h	ComputeReducerLayout now preserves the reduced axis as extent 1 when requested. ReduceOpNode::InferLayout validates destination rank and reduced-axis extent.
Reduction lowering and validation src/backend/common/op/reduce.h, testing/python/language/test_tilelang_language_reduce.py	Lowering propagates keep-dimension state, validates destination shape, preserves source-rank indexing, and tests sum, max, and min reductions.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: leiwang1999, siriusneo

Sequence Diagram(s)

sequenceDiagram
  participant ReduceOpNode
  participant ComputeReducerLayout
  participant ReduceLowerer
  participant KeepdimTest
  ReduceOpNode->>ComputeReducerLayout: derive layout with keepdim
  ComputeReducerLayout-->>ReduceOpNode: return rank-preserving reducer layout
  ReduceLowerer->>ComputeReducerLayout: construct lowering layout with keepdim
  ReduceLowerer->>ReduceLowerer: validate destination rank and extent 1
  ReduceLowerer->>KeepdimTest: execute reduction kernel
  KeepdimTest-->>ReduceLowerer: compare shape and values with keepdim reference

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 15.38% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The changes address issue [#2530] by supporting keep-dimension layout inference, lowering, backend layout generation, and regression coverage.
Out of Scope Changes check	✅ Passed	All code and test changes directly support keep-dimension fragment reductions and the linked issue objectives.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: fixing fragment reductions that support keep-dimension destinations.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/op/reduce.cc`:
- Around line 218-224: Require the destination extent at the reduced axis to be
1 before enabling keep-dim in the reducer-layout selection: update the predicate
using dst->shape[this->dim] in src/op/reduce.cc lines 218-224, and mirror it
with op.dst->shape[op.dim] in src/backend/common/op/reduce.h lines 802-808.
Ensure same-rank destinations with a non-unit reduced axis are rejected clearly
rather than treated as keep-dim.

In `@testing/python/language/test_tilelang_language_reduce.py`:
- Around line 659-700: Extend REDUCE_KEEPDIM_CASES and test_reduce_keepdim to
include non-clearing min cases. In the accumulator initialization, use positive
infinity for min and negative infinity for max; update the non-sum reduction
dispatch so op == "min" invokes T.reduce_min rather than T.reduce_max, while
preserving sum and clearing behavior.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 990138ef-c02d-445b-b3fa-c710983af178

📥 Commits

Reviewing files that changed from the base of the PR and between 8f34abf and 494967c.

📒 Files selected for processing (3)

• src/backend/common/op/reduce.h
• src/op/reduce.cc
• testing/python/language/test_tilelang_language_reduce.py