Skip to content

[Frontend] Gathering the same tensor twice returns zeros: both gathers alias one scratchpad and CSE folds the loads #297

Description

@YWHyuk

Summary

Gathering the same tensor twice with different indices in one kernel silently returns wrong
results on npu. No error, no warning — the output is just zeros.

Minimal repro

import torch

N = 16
torch.manual_seed(0)
x  = torch.randn(N)
i0 = torch.randint(0, N, (N,), dtype=torch.int64)
i1 = torch.randint(0, N, (N,), dtype=torch.int64)

def f(t, a, b):
    return t[a] - t[b]

golden = f(x, i0, i1)
opt = torch.compile(dynamic=False)(f)
got = opt(x.to("npu:0"), i0.to("npu:0"), i1.to("npu:0")).cpu()

print(golden[:4])   # tensor([ 0.7562,  1.2457,  0.7839, -1.8444])
print(got[:4])      # tensor([0., 0., 0., 0.])   <-- wrong

Observed (functional Spike pass on, systolic_ws_128x128_c2_simple_noc_tpuv3.yml):

golden : [0.7561537027359009, 1.2456812858581543, 0.7838619947433472, -1.8443694114685059]
npu    : [0.0, 0.0, 0.0, 0.0]
allclose: False      npu all-zero: True

Why

Both gathers alias one scratchpad, and the emitted loads are textually identical, so CSE folds them
into a single value. a - b becomes a - a:

"togsim.transfer"(..., %spad1, ..., %alloc1) {indirect = true}   // gather #1 -> spad1
"togsim.transfer"(..., %spad1, ..., %alloc3) {indirect = true}   // gather #2 -> SAME spad1, overwrites
affine.for %compute_idx = 0 to 2 step 2 {
    %tmp3 = affine.vector_load %spad1[...]
    %tmp7 = arith.subf %tmp3, %tmp3                              // A - A  ==  0

allocate_sram_buffer keys the scratchpad by (dram_name, str(raw_index)), but load() has already
moved the indirect part of the index into offset_desc, so x[i0] and x[i1] produce the same key.
The index buffer identity never reaches the cache key.

Full analysis in #296 (comment). This is the same root cause as #296, but a different — and more
dangerous — symptom: #296 crashes with invalid MLIR, this one returns wrong data silently.

Severity

Any model that gathers one tensor twice is affected, e.g. embedding lookups with two index sets,
x[i] - x[j], and bilinear resampling. There is no diagnostic; the result is simply wrong.

Environment

  • branch feature/togsim-cpp-trace
  • pytorchsim_functional_mode: 1 (Spike) — the value above is what the functional simulator produces.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions