[Feature Request][RFC] Redesign reducers as ownership-safe deferred reductions

## Required prerequisites

- [x] I searched the issue tracker and did not find an existing design proposal covering first-class reducer update semantics. This RFC is related to, but broader than, the concrete wrong-code reports linked below.

## Motivation

TileLang needs a way to express **deferred reductions**: contributions may be produced across multiple `T.Parallel` loops and multiple `T.Pipelined` iterations, each thread should accumulate local partials first, and cross-thread communication should happen only once at the end.

This is useful for kernels such as the in-tree GEMV reducer example:

```python
acc = T.alloc_reducer((block_M,), accum_dtype, op="sum", replication="all")
T.clear(acc)

for k_tile in T.Pipelined(...):
# load A_frag and x_frag
for i, j in T.Parallel(block_M, block_N):
acc[i] += A_frag[i, j] * x_frag[j]

T.finalize_reducer(acc)
```

Using `T.reduce_sum` once per tile would introduce a collective per pipeline iteration. Materializing the entire reduction domain and reducing once would be too expensive. A reducer is therefore useful, but its semantic purpose should be **local monoid accumulation followed by one deferred collective**, not a user-selected fully replicated fragment layout.

The current implementation conflates these two concepts. That has produced silent wrong-code, unclear ownership rules, and special cases throughout layout inference and verification.

This RFC proposes making reducer updates first-class operations, defining a compositional correctness lowering, and treating subgroup ownership analysis as an optional optimization rather than a requirement for correctness.

Inspected against `main` at `fef9f760920fa5e81596e08696ee3c6e8e1c18ce`.

## Current model

The public API currently returns an ordinary fragment buffer and attaches reducer metadata to the surrounding block:

```python
acc = T.alloc_reducer(shape, dtype, op="sum", replication="all")
T.fill(acc, 0)
acc[i] += value
T.finalize_reducer(acc)
```

Relevant implementation choices:

1. [`alloc_reducer`](https://github.com/tile-ai/tilelang/blob/fef9f760920fa5e81596e08696ee3c6e8e1c18ce/tilelang/language/allocate.py#L283-L319) allocates a `local.fragment` buffer and records string metadata for `op` and `replication`.
2. [`LayoutReducer`](https://github.com/tile-ai/tilelang/blob/fef9f760920fa5e81596e08696ee3c6e8e1c18ce/src/transform/layout_reducer.cc#L187-L253) maps `replication="all"` to `Fragment::FullyReplicated(...)`. Reducer store validation is still a TODO.
3. [`ParallelOp`](https://github.com/tile-ai/tilelang/blob/fef9f760920fa5e81596e08696ee3c6e8e1c18ce/src/op/parallel.cc#L432-L468) skips all-replicated reducers as layout sources, and [skips their normal fragment ownership validation](https://github.com/tile-ai/tilelang/blob/fef9f760920fa5e81596e08696ee3c6e8e1c18ce/src/op/parallel.cc#L650-L700).
4. [`VerifyParallelLoop`](https://github.com/tile-ai/tilelang/blob/fef9f760920fa5e81596e08696ee3c6e8e1c18ce/src/transform/verify_parallel_loop.cc#L57-L62) exempts reducer stores from the ordinary race check.
5. [`FinalizeReducerLowerer`](https://github.com/tile-ai/tilelang/blob/fef9f760920fa5e81596e08696ee3c6e8e1c18ce/src/backend/common/op/finalize_reducer.h#L40-L56) obtains `reducing_threads` directly from the reducer fragment's `ReplicateExtent()`.

This works only when every physical thread-local value is an independent partial that should participate in the final collective. A fragment's final replication layout does not establish that property.

## Minimal wrong-code example

```python
BD = 8
THREADS = 128

with T.Kernel(1, threads=THREADS):
x_frag = T.alloc_fragment((BD,), T.float32)
T.copy(x, x_frag)

total = T.alloc_reducer((1,), T.float32, op="sum", replication="all")
T.fill(total, 0.0)
for j in T.Parallel(BD):
total[0] += x_frag[j]
T.finalize_reducer(total)
```

The inferred layouts are:

```text
total: thread = rep, replicate = 128
x_frag: thread = rep*8 + j, replicate = 16
```

The generated accumulation and finalization are equivalent to:

```cpp
x_frag[0] = x[threadIdx.x & 7];
total[0] += x_frag[0];
total[0] = tl::AllReduce>::run(total[0], workspace);
```

Each logical `x[j]` contribution is inserted by 16 replica threads. The result is therefore:

```text
sum(thread=0..127, x[thread % 8]) = 16 * sum(x)
```

For `x = [1, ..., 8]`, the kernel returns `576` instead of `36`.

In contrast, `T.reduce_sum(x_frag, total)` sees the source thread expression and correctly emits:

```cpp
tl::AllReduce>::run(...)
```

The two `128` values in the incorrect call also have different meanings: the first is the reduction span, while `NamedBarrier<128>` is the number of threads reaching a shared barrier. The latter can remain the participant count even when independent 8-thread reduction groups are used.

This is the small, exact version of the failure reported in #2408.

## Problems with the current abstraction

### 1. A logical reduction update is represented as an arbitrary `BufferStore`

The compiler knows the declared reducer op, but the update itself is just general TIR:

```python
acc[i] += value
acc[i] = T.max(acc[i], value)
acc[i] = unrelated_expression
```

There is no first-class IR node identifying the logical output index and contribution. The implementation currently does not verify that stores conform to the declared op. This makes it difficult to distinguish a duplicated pure fragment assignment from a duplicated non-idempotent reduction contribution.

It also forces reducer buffers to be exempted from ordinary ownership/race validation instead of giving reducer updates their own legality rule.

### 2. Final-result replication is conflated with partial ownership

`Fragment::FullyReplicated` normally describes a value available in every participating thread. During a reducer region, however, those slots are intentionally incoherent partial states until finalization.

More importantly, a final layout with `ReplicateExtent() == 128` cannot distinguish between:

- 128 independent partials that require one 128-thread reduction;
- 16 equivalent replica groups, each containing an 8-thread reduction;
- 8 groups of 16 threads;
- multiple update sites with different groupings.

The final reducer layout has already lost the contribution ownership information. `FinalizeReducer` cannot reconstruct it from the result layout alone.

### 3. A single saved ownership plan is not compositional

It is tempting to preserve `ThreadReduceStep` or `ReduceOwnershipPlan` from the accumulation loop and pass it to finalization. That is sound for a closed `T.reduce`, which has one source layout and explicit reduction axes. It is not generally a property of a reducer region:

```python
for j in T.Parallel(8): # 8-thread groups, 16 replicas
T.reducer_update(acc, 0, a[j])

for k in T.Parallel(16): # 16-thread groups, 8 replicas
T.reducer_update(acc, 0, b[k])
```

After both update sites have accumulated into the same thread-local state, neither `AllReduce<8>` nor `AllReduce<16>` is generally sufficient. Plans can also differ in scale, offset, thread range, output-index projection, or predicates.

Combining arbitrary plans would require region-level effect and group-equivalence analysis. Correctness should not depend on solving that optimization problem.

### 4. Identity, seed, and lifetime are not explicit enough

`T.fill` currently acts as the start marker and initializer. The docs require users to choose the proper identity, but this is not the same as a general initial value:

- every thread-local partial must start from the monoid identity;
- a user seed such as `initial=1` for sum must be applied exactly once, not once per thread;
- the reducer should not be readable before finalization;
- updates after finalization, duplicate finalization, unmatched fill/finalize, and divergent finalization should be rejected clearly.

The current buffer type does not distinguish partial state from a finalized result, and pairing is inferred by scanning calls between `T.fill` and `T.finalize_reducer`.

### 5. Physical policy leaks into the logical API

`replication="all"` / `"none"` is chosen when the logical reducer is allocated, even though partial storage layout and final-result distribution are different decisions.

Similarly, `batch` is exposed as a semantic-looking `finalize_reducer` argument even though it should only select an equivalent code-generation strategy. This has also created a separate wrong-code path in #2623.

### 6. Participant sets are implicit

The collective participant range may differ from the raw kernel thread extent under producer/consumer warp specialization. The reducer should operate over an explicit, compiler-known participant set, including its offset, instead of assuming that all relevant workers start at thread zero. #2346 demonstrates this class of problem.

### 7. Extensibility and testing are fragmented

The reducer op enum is separate from the richer `T.reduce` implementation and currently supports only sum/max/min (#1016). Numerical coverage is also incomplete: some batched paths have codegen tests but are deliberately excluded from correctness tests.

## Proposed semantic model

A reducer is an opaque **deferred monoid reduction**, not a readable fragment buffer.

For a reducer `R` with combine operation `⊕` and identity `e`:

1. allocation starts a reduction epoch with all physical partials initialized to `e`;
2. `reducer_update(R, index, value)` contributes `value` exactly once to logical output `index`;
3. update sites may appear in multiple loops, pipeline iterations, and compatible control-flow regions;
4. finalization combines all partials over the compiler-defined participant set;
5. only the finalized result is readable;
6. physical replication, subgroup shape, batching, vectorization, and workspace strategy do not change these semantics.

Floating-point sum may be reassociated in the same way as existing GPU reductions; bitwise reproducibility across different layouts is not a goal.

## Proposed API

### Recommended v2 surface

```python
acc = T.alloc_reducer((M,), T.float32, op="sum")

for k_tile in T.Pipelined(...):
for i, j in T.Parallel(M, K_TILE):
T.reducer_update(acc, (i,), A_frag[i, j] * x_frag[j])

result = T.alloc_fragment((M,), T.float32)
T.finalize_reducer(acc, result)
T.copy(result, out)
```

Properties:

- `alloc_reducer` no longer accepts a physical replication policy.
- The reducer op and identity belong to the reducer definition.
- `reducer_update` carries only the reducer handle, logical result indices, and contribution.
- `finalize_reducer` consumes the partial state and writes a normal destination buffer whose layout is inferred or explicitly annotated through existing layout mechanisms.
- A destination may eventually be fragment/shared/global, allowing finalization directly into the next consumer when legal.

For an incremental migration, finalization may remain in-place initially:

```python
T.finalize_reducer(acc)
# acc becomes readable after this point
```

but a separate destination is cleaner because the partial and finalized values have different semantics and potentially different layouts.

### Update operation

Conceptually:

```text
partial[index] = combine(partial[index], contribution)
```

Examples:

```python
sum_acc = T.alloc_reducer(..., op="sum")
T.reducer_update(sum_acc, i, value)

max_acc = T.alloc_reducer(..., op="max")
T.reducer_update(max_acc, i, T.abs(value))
```

`abs_sum` and `abs_max` do not require separate combine operations; `abs` is a transformation of the contribution. Subtraction is not a reduction operator, but `acc -= value` may be frontend sugar for a sum update with `-value`.

The frontend may temporarily recognize legacy forms such as `acc[i] += value` and canonicalize them to `ReducerUpdateOp`, but the core IR should not rely indefinitely on matching arbitrary `BufferStore` expressions.

### Identity versus seed

The monoid identity should be initialized implicitly:

| op | identity |
|---|---|
| sum | `0` |
| product | `1` |
| max | lowest representable value |
| min | highest representable value |
| bitwise-or/xor | `0` |
| bitwise-and | all bits set |

If an initial logical seed is desired, it should be a separate API/property and must be combined exactly once after or during finalization. It must not replace the per-thread identity fill.

## Proposed IR

The first implementation can keep an underlying allocation buffer for storage, while exposing it as an opaque reducer handle to verification and high-level transforms.

### `ReducerInfo`

```text
ReducerInfo {
shape
dtype
combine_op
identity
scope / participant domain
}
```

The combine op should reuse a common reduction-op registry shared with `T.reduce` rather than maintaining a separate fixed enum and separate backend naming table.

### `ReducerUpdateOp`

```text
ReducerUpdateOp {
reducer
logical_indices
contribution
}
```

This is an effectful high-level op. It is not an ordinary buffer read/write and should remain visible through layout inference so `ParallelOp` can reason about the logical iteration and its replica coordinate.

### `FinalizeReducerOp`

```text
FinalizeReducerOp {
reducer
destination
optional_codegen_hints
}
```

Codegen hints such as batch width must not affect semantics. An unsupported hint should fall back to scalar lowering rather than silently changing which values are reduced.

### State verification

Add a dedicated reducer verifier with a simple state machine:

```text
Allocated/Active --update*--> Active --finalize--> Finalized
```

It should reject:

- reducer loads before finalization;
- arbitrary stores that did not canonicalize to `ReducerUpdateOp`;
- update after finalization;
- missing or duplicate finalization;
- out-of-bounds logical output indices;
- incompatible contribution dtype;
- finalize under thread-divergent control flow;
- reducer escape/alias patterns the compiler cannot analyze;
- unsupported non-associative combine operations.

## Correctness lowering: canonical unique contributions

Correctness should use a local, compositional rule:

> Every dynamic logical `ReducerUpdateOp` contributes exactly once to its logical output across the participant set.

When an update occurs inside a replicated `T.Parallel` layout, `ParallelOp` still knows the loop layout and can derive the loop's replica coordinate. The canonical lowering should guard **only the reducer update** with a representative condition such as `rep == 0`:

```cpp
// Ordinary replicated fragment work remains unchanged.
y_frag[...] = f(x_frag[...]);

// Only the reducer contribution is deduplicated.
if (loop_replica == 0) {
reducer_partial[...] = combine(reducer_partial[...], contribution);
}
```

It is not safe to guard the entire `T.Parallel` body because an ordinary replicated destination fragment may require all replicas to be populated.

For the minimal `BD=8`, `threads=128` example, eight representative threads contribute and the remaining partial slots retain the sum identity. A participant-wide `AllReduce<128>` then produces the correct result. With multiple update sites and different layouts, each site independently contributes once; the final participant-wide collective remains correct without merging ownership plans.

Finalization must remain outside the representative predicate so every thread required by the selected barrier policy reaches the collective.

The participant set must come from the current execution domain (including warp-specialized consumer offset/range), not be reconstructed from the reducer result layout.

This baseline may not always be optimal, but it is sound and provides a fallback for every supported reducer region.

## Optional optimization: compatible subgroup signatures

Subgroup finalization should be an optimization fast path, not the semantic foundation.

After layouts are known, each update site may produce a normalized signature such as:

```text
ThreadGroupSignature {
participant_range
group_partition / thread map
reduction_steps (extent, scale, lower factor)
logical-output projection
required uniform predicate
}
```

The compiler may avoid representative-only updates and emit subgroup `AllReduce` only when all update sites for an epoch have provably compatible signatures. A conservative first version should require exact normalized equality rather than attempting to merge plans.

Examples:

- one update site with 16 identical 8-thread groups: emit `AllReduce<..., 8, ...>` independently in every group;
- several update sites with the exact same grouping and output projection: same optimization is legal;
- sites with widths 8 and 16, different scales, different participant ranges, or unproven predicates: use canonical unique contributions plus participant-wide finalization.

This reuses the ownership machinery developed for [`T.reduce`](https://github.com/tile-ai/tilelang/blob/fef9f760920fa5e81596e08696ee3c6e8e1c18ce/src/backend/common/op/reduce.h#L386-L433) where applicable, without claiming that a `ReduceOwnershipPlan` from one update site describes an arbitrary region.

## Layout responsibilities

The design should distinguish three layouts/domains:

1. **Logical reducer shape**: user-visible output indices.
2. **Partial storage layout**: an internal implementation choice; initially this may conservatively allocate a slot per participant and logical output as the current all-replicated path does.
3. **Final result layout**: the normal destination layout required by downstream consumers.

The initial implementation can prioritize correctness and reuse current partial storage. Later work may infer a sparse/segmented partial layout to reduce register pressure. That optimization should not be exposed as `replication=` on the logical reducer.

## Alternatives considered

### Keep arbitrary buffer syntax and pattern-match stores

This minimizes frontend changes, but it remains difficult to validate, gives poor diagnostics, and continues to make reducer semantics depend on recognizing expression shapes. It is reasonable only as a compatibility canonicalization step.

### Carry one `ReduceOwnershipPlan` to finalization

Efficient for a single closed update site, but not compositional across multiple layouts or predicates. This should be an optional fast path with compatibility checks, not the correctness path.

### Divide sum results by the replica count

This only works for a subset of sum cases, does not generalize to arbitrary monoids, and fails when update sites have different multiplicities. It also changes floating-point behavior unnecessarily.

### Guard the entire replicated loop

This can leave ordinary replicated fragment destinations uninitialized. The representative predicate must apply specifically to reducer updates unless layout inference globally and safely de-replicates every affected fragment consumer.

### Eliminate reducers and call `T.reduce` per tile

This loses the main benefit of deferred reduction by introducing collectives inside streaming/pipelined loops. A reducer abstraction remains valuable; the proposal changes its semantic foundation.

## Migration plan

### Phase 0: stop silent wrong-code

- Fix or conservatively reject replicated contributions covered by #2408.
- Fix the independent batched-finalize defect in #2623, or temporarily fall back to scalar finalization.
- Add numerical tests before adding more reducer codegen optimizations.

### Phase 1: introduce first-class updates and verification

- Add the Python `T.reducer_update` API and a registered `ReducerUpdateOp`.
- Add reducer state/lifetime verification.
- Canonicalize supported legacy sum/max/min store patterns to the new op.
- Reject unsupported reducer stores with a source-span diagnostic.
- Keep the current in-place finalize API initially if necessary.

### Phase 2: canonical ownership-safe lowering

- Derive update-specific representative predicates from `ParallelOp` layouts.
- Apply the predicate only to reducer updates.
- Use explicit participant ranges for identity initialization, final collectives, barriers, and workspace addressing.
- Remove reducer-specific blanket exemptions from normal fragment ownership logic where the new op makes them unnecessary.

### Phase 3: separate result layout and add conservative fast paths

- Support an explicit finalize destination.
- Move result replication/distribution to destination layout inference.
- Add exact-signature subgroup reduction.
- Select scalar/batched/vectorized finalization automatically and share implementation with `T.reduce`.

### Phase 4: deprecate the legacy surface

- Deprecate `replication=` on `alloc_reducer`.
- Deprecate `T.fill(reducer, identity)` in favor of implicit identity initialization.
- Eventually stop treating arbitrary reducer `BufferStore` as a supported core IR form.

## Test plan / acceptance criteria

### Numerical correctness

- `BD=8`, `threads=128`, sum of `[1, ..., 8]` returns `36`, not `576`.
- #2408 non-dividing GEMV tile widths are correct.
- Two update sites with incompatible group widths (for example 8 and 16) fall back and produce the correct result.
- Multiple compatible update sites remain correct and may use the subgroup fast path.
- A loop containing both an ordinary replicated fragment write and a reducer update initializes all ordinary fragment replicas while counting each reducer contribution once.
- Predicated updates count exactly the logically enabled contributions.
- Streaming/pipelined GEMV performs one final collective and matches the reference.
- Sum/max/min across scalar and multi-element reducer shapes.
- Batched and scalar codegen produce identical numerical results for every output element.

### Participant domains

- Full CTA.
- Partial/offset consumer ranges under warp specialization.
- Warp-only and cross-warp collectives.
- CUDA and ROCm implementations where supported.

### Diagnostics

- Read before finalize.
- Update after finalize.
- Missing/double finalize.
- Arbitrary or op-mismatched store.
- Non-uniform finalize.
- Invalid index/dtype.
- Unsupported alias or escape.

### Codegen invariants

- Barrier participant count is derived independently from subgroup reduction width.
- Final collectives are outside update representative predicates.
- Unsupported batching/vectorization hints fall back without changing semantics.
- The fallback path works without a compatible `ThreadGroupSignature`.

## Related issues and history

- #2408: all-replicated reducer contributions can be counted multiple times.
- #2623: batched finalization reduces only the first batch of output elements.
- #2346: reducer participant range/offset interacts incorrectly with warp specialization.
- #1016: reducer op coverage and finalize vectorization.
- #1560: reducer buffers are exempted from ordinary `T.Parallel` verification.
- #2053: motivation for efficient fragment/global reductions.
- #757: introduced `alloc_reducer` to separate local accumulation from inter-thread reduction.
- #1976: introduced batched AllReduce optimization.

## Open questions

1. Should v2 finalization require a separate destination immediately, or retain in-place finalization during migration?
2. Should final-result distribution be expressed entirely by destination layout, or should there be a small semantic policy such as `broadcast="all"`?
3. What is the minimum useful built-in monoid set? Should custom scalar/structured monoids be a later RFC?
4. Should reducers support explicit `reset`/multiple epochs, or require one allocation per epoch initially?
5. What conservative partial storage layout is acceptable for the correctness implementation before sparse/segmented inference is added?
6. How long should legacy `acc[i] += value` pattern canonicalization be supported?

The main decision requested by this RFC is whether TileLang agrees with the following semantic boundary:

> `ReducerUpdateOp` defines logical contribution multiplicity; layouts select storage and communication, but must not change how many times a logical contribution enters the reduction.