Skip to content

[Bug] Task dep-gen misses WAR edge: reader task not ordered before a later aliasing inout writer #1306

Description

@YunjiQin

Background

In pypto-lib, a fused per-band MLP for models/qwen3/14b/prefill_fwd.py on a2a3 produced wrong results (golden ~10% off) when two bands reused one loop-carried accumulator buffer. Reducing it isolated a runtime task-dependency (WAR) bug: the runtime does not order a later loop iteration's writer task after the current iteration's reader task when they alias the same buffer. Minimal reproducer below fails deterministically.

Reproduce with:

python models/qwen3/14b/war_repro_draft.py -p a2a3 -d <device>

Reproduction environment:

Component Version
pypto-lib 198ab5a (branch: main)
pypto 10d6b128 (branch: main)
simpler c94aa9f3 (pin)
pto-isa 83d01313
CANN 9.0.0
Driver 26.0.rc1

Diagnosis: simpler — task dependency generation misses the write-after-read (WAR) anti-dependency between a pure-reader task and a later aliasing inout writer.

Platform

a2a3 (Ascend 910B/C hardware)

Runtime Variant

All / Unknown

Description

The task-graph dependency generation does not emit a WAR (write-after-read) edge from a reader task to a subsequent writer task that aliases the same tensor across a loop.

Pattern (one create_tensor buffer buf, carried across a pl.range loop):

  • iteration N: writer task takes buf add_inout (RMW) → reader task takes buf add_input (pure read)
  • iteration N+1: writer task takes buf add_inout again

The runtime emits writer(N) → reader(N) (RAW) and writer(N) → writer(N+1) (WAW), but not reader(N) → writer(N+1) (WAR). Because the reader produces no new version of buf, the loop carries the writer's output forward, so iteration N+1's writer depends only on iteration N's writer — never on iteration N's reader. The runtime is then free to run writer(N+1) concurrently with reader(N), overwriting buf while reader(N) is still reading it → data race.

This is visible directly in the generated orchestration (single buf, add_inout in writer, add_input in reader):

// Generated orchestration: war_repro.cpp  (buf = alloc_tensors(...), created ONCE)
for (int64_t b = 0; b < 2; b += 1) {
    PTO2_SCOPE() {
        // Spmd writer_spmd: writer
        L0TaskArgs params_t0;
        params_t0.add_inout(buf);          // writer RMW on buf
        params_t0.add_input(ext_src);
        params_t0.add_scalar(b);
        params_t0.launch_spec.set_block_num(24);
        rt_submit_aiv_task(0, params_t0);

        // Spmd reader_spmd: reader
        L0TaskArgs params_t1;
        params_t1.add_inout(ext_out);
        params_t1.add_input(buf);          // reader pure-read of buf
        params_t1.add_scalar(b);
        params_t1.launch_spec.set_block_num(24);
        rt_submit_aiv_task(1, params_t1);
    }
}

Expected: writer(b=1) (add_inout buf) waits for reader(b=0) (add_input buf) — a WAR edge. Actual: no such edge; writer(b=1) overwrites buf while reader(b=0) reads it.

Steps to Reproduce

Minimal standalone frontend (no matmul; two loop iterations; writer inout + reader pure-read of one carried buffer):

import pypto.language as pl

BANDS = 2
M = 128
N = 8704          # large enough that writer/reader overlap in time
CHUNK = 256
NCHUNKS = N // CHUNK
SPMD = 24


@pl.jit
def war_repro(
    src: pl.Tensor[[BANDS * M, N], pl.FP32],
    out: pl.Out[pl.Tensor[[BANDS * M, N], pl.FP32]],
):
    buf = pl.create_tensor([M, N], dtype=pl.FP32)
    for b in pl.range(BANDS):          # buf loop-carried across bands
        # writer: overwrite buf from src[b]  (inout)
        for wcore in pl.spmd(SPMD, name_hint="writer_spmd"):
            for cc in pl.range(wcore, NCHUNKS, SPMD):
                c0 = cc * CHUNK
                s = pl.slice(src, [M, CHUNK], [b * M, c0])
                buf = pl.assemble(buf, s, [0, c0])
        # reader: copy buf -> out[b]  (pure read)
        for rcore in pl.spmd(SPMD, name_hint="reader_spmd"):
            for cc in pl.range(rcore, NCHUNKS, SPMD):
                c0 = cc * CHUNK
                r = pl.slice(buf, [M, CHUNK], [0, c0])
                out = pl.assemble(out, r, [b * M, c0])
    return out

Harness init: src[band0]=1.0, src[band1]=2.0; golden out = src. Run:

python models/qwen3/14b/war_repro_draft.py -p a2a3 -d <device>

Full repro file: models/qwen3/14b/war_repro_draft.py in pypto-lib.

Expected Behavior

out[band0] == src[band0] == 1.0 and out[band1] == src[band1] == 2.0, i.e. each band's reader sees its own writer's data (WAR ordering enforced).

Actual Behavior

Golden FAILs deterministically:

'out' FAIL  shape=(256, 8704)
Mismatched elements: 392704/2228224  rtol=0.0 atol=0.0
[0] actual=2.0, expected=1.0
[1] actual=2.0, expected=1.0
...

out[band0] contains 2.0 (band1's writer data) in ~35% of its columns — band0's reader read buf after band1's writer overwrote it. Missing WAR edge → race.

Git Commit ID

c94aa9f

CANN Version

9.0.0

Driver Version

26.0.rc1

Host Platform

Linux (aarch64)

Additional Context

Real-kernel manifestation: in models/qwen3/14b/prefill_fwd.py, a fused per-band MLP shared the gate/up/silu accumulator buffers across the two MLP bands to force a band0→band1 software pipeline. gate is inout, silu is a pure input reader of the gate accumulator; band1's gate raced band0's silu exactly as above, corrupting ~10% of logits. Per-band-private buffers (no aliasing) are correct — this repro is the minimized aliasing case.

Related: hw-native-sys/pypto-lib#481 — a sibling WAR-anti-dependency miss, but at the orchestration auto-dep layer (a gather add_input view + writeback add_output view of the same external inout tensor get no WAR edge). This issue is the runtime task-scheduler layer: the orch emits correct add_inout/add_input annotations (see war_repro.cpp above), but the runtime dep-gen does not derive the reader → next-iteration writer WAR edge from them.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions