[Bug] Nested scalar if silently drops outer condition during a2a3 lowering

[sjduan]

Component

Codegen

Description

A nested scalar guard inside an in-core task is silently miscompiled for the
a2a3 target. Given:

if logical_pos >= 0:
    if logical_pos < ROWS:
        value = state[logical_pos : logical_pos + 1, 0:COLS]

the generated PTO drops the outer logical_pos >= 0 predicate and keeps only
logical_pos < ROWS. A negative offset then reaches partition_view, where it
is clamped to row 0, producing silently incorrect output.

The logically equivalent combined predicate lowers correctly:

if logical_pos >= 0 and logical_pos < ROWS:
    value = state[logical_pos : logical_pos + 1, 0:COLS]

Steps to Reproduce

Save the following as `pypto_nested_if_lowering_probe.py` in a `pypto-lib`
checkout so its `golden` runner is importable:


#!/usr/bin/env python3
"""Minimal repro for a PyPTO nested scalar-if lowering bug.

Both kernels should return ``state[0]``. Their only loop iteration has
``logical_pos == -1``, so it must not read state. The nested spelling is
currently miscompiled on a2a3: the outer ``logical_pos >= 0`` guard disappears
and the negative slice is clamped to row 0, producing ``2 * state[0]``. The
logically equivalent combined predicate works.
"""

import pypto.language as pl


ROWS = 2
COLS = 64


@pl.jit.inline
def nested_if_body(
    token_position: pl.Tensor[[1], pl.INT32],
    state: pl.Tensor[[ROWS, COLS], pl.FP32],
    output: pl.Tensor[[1, COLS], pl.FP32],
):
    pooled = pl.create_tensor([1, COLS], dtype=pl.FP32)
    with pl.at(level=pl.Level.CORE_GROUP, name_hint="nested_if_probe") as pool_tid:
        token_pos = pl.cast(pl.read(token_position, [0]), pl.INDEX)
        window_start = token_pos - ROWS + 1
        acc = state[token_pos : token_pos + 1, 0:COLS]
        for state_idx in pl.range(ROWS - 1):
            logical_pos = window_start + state_idx
            value = pl.full([1, COLS], dtype=pl.FP32, value=0.0)
            if logical_pos >= 0:
                if logical_pos < ROWS:
                    value = state[logical_pos : logical_pos + 1, 0:COLS]
            acc = pl.add(acc, value)
        pooled[0:1, 0:COLS] = acc
    with pl.at(level=pl.Level.CORE_GROUP, name_hint="nested_copy", deps=[pool_tid]):
        output[0:1, 0:COLS] = pooled[0:1, 0:COLS]


@pl.jit.inline
def combined_if_body(
    token_position: pl.Tensor[[1], pl.INT32],
    state: pl.Tensor[[ROWS, COLS], pl.FP32],
    output: pl.Tensor[[1, COLS], pl.FP32],
):
    pooled = pl.create_tensor([1, COLS], dtype=pl.FP32)
    with pl.at(level=pl.Level.CORE_GROUP, name_hint="combined_if_probe") as pool_tid:
        token_pos = pl.cast(pl.read(token_position, [0]), pl.INDEX)
        window_start = token_pos - ROWS + 1
        acc = state[token_pos : token_pos + 1, 0:COLS]
        for state_idx in pl.range(ROWS - 1):
            logical_pos = window_start + state_idx
            value = pl.full([1, COLS], dtype=pl.FP32, value=0.0)
            if logical_pos >= 0 and logical_pos < ROWS:
                value = state[logical_pos : logical_pos + 1, 0:COLS]
            acc = pl.add(acc, value)
        pooled[0:1, 0:COLS] = acc
    with pl.at(level=pl.Level.CORE_GROUP, name_hint="combined_copy", deps=[pool_tid]):
        output[0:1, 0:COLS] = pooled[0:1, 0:COLS]


@pl.jit
def nested_if_probe(
    token_position: pl.Tensor[[1], pl.INT32],
    state: pl.Tensor[[ROWS, COLS], pl.FP32],
    output: pl.Out[pl.Tensor[[1, COLS], pl.FP32]],
):
    nested_if_body(token_position, state, output)


@pl.jit
def combined_if_probe(
    token_position: pl.Tensor[[1], pl.INT32],
    state: pl.Tensor[[ROWS, COLS], pl.FP32],
    output: pl.Out[pl.Tensor[[1, COLS], pl.FP32]],
):
    combined_if_body(token_position, state, output)


def build_tensor_specs():
    import torch
    from golden import TensorSpec

    return [
        TensorSpec(
            "token_position",
            [1],
            torch.int32,
            init_value=lambda: torch.tensor([0], dtype=torch.int32),
        ),
        TensorSpec(
            "state",
            [ROWS, COLS],
            torch.float32,
            init_value=lambda: torch.arange(
                1, ROWS * COLS + 1, dtype=torch.float32
            ).reshape(ROWS, COLS),
        ),
        TensorSpec("output", [1, COLS], torch.float32, is_output=True),
    ]


def golden(tensors):
    tensors["output"][:] = tensors["state"][0:1]


if __name__ == "__main__":
    import argparse
    from golden import run_jit

    parser = argparse.ArgumentParser()
    parser.add_argument("-p", "--platform", default="a2a3")
    parser.add_argument("-d", "--device", type=int, default=0)
    parser.add_argument("--case", choices=["nested", "combined", "all"], default="all")
    args = parser.parse_args()

    cases = {"nested": nested_if_probe, "combined": combined_if_probe}
    selected = cases if args.case == "all" else {args.case: cases[args.case]}
    failed = []
    for name, fn in selected.items():
        print(f"--- {name} ---")
        result = run_jit(
            fn=fn,
            specs=build_tensor_specs(),
            golden_fn=golden,
            runtime_cfg={"platform": args.platform, "device_id": args.device},
            rtol=0.0,
            atol=0.0,
        )
        if not result.passed:
            failed.append(name)
            if result.error:
                print(result.error)

    raise SystemExit(1 if failed else 0)


Run on an available Ascend 910B device:


python pypto_nested_if_lowering_probe.py -p a2a3 -d 0 --case all

Expected Behavior

Both variants are semantically equivalent and should return state[0]:

output[0, 0:4] = [1.0, 2.0, 3.0, 4.0]

Both cases should pass exact validation.

Actual Behavior

The nested variant fails all 64 elements and returns twice the expected value:

--- nested ---
output FAIL: mismatched elements 64/64
actual[0, 0:4]   = [2.0, 4.0, 6.0, 8.0]
expected[0, 0:4] = [1.0, 2.0, 3.0, 4.0]

--- combined ---
output PASS

The generated PTO for the nested form contains only the upper-bound check:

%is_lt_rows = arith.cmpi slt, %logical_pos, %c2_index : index
scf.if %is_lt_rows {
  %clamped = arith.maxsi %logical_pos, %c0_index : index
  %view = pto.partition_view %state, offsets = [%clamped, %c0_index], ...
  pto.tload ...
}

There is no logical_pos >= 0 comparison. Since logical_pos == -1, the
remaining upper-bound check succeeds and the slice offset is clamped to row 0.

The combined form lowers correctly:

%is_nonnegative = arith.cmpi sle, %c0_index, %logical_pos : index
%is_lt_rows = arith.cmpi slt, %logical_pos, %c2_index : index
%valid = arith.andi %is_nonnegative, %is_lt_rows : i1
scf.if %valid {
  ...
}

### Git Commit ID

c56951b923e528026b1020087adc257b1279e645

### NPU Kind

Ascend 910C

### Host Platform

Linux (aarch64)

### Additional Context

_No response_