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

    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