[Bug] [TIRx] CompactBufferAllocation leaves a stale double_buffer_scope reference and causes InjectDoubleBuffer to fail #20972

Expected behavior

Buffer-rebuilding passes should keep double_buffer_scope attached to the allocated buffer that its body actually accesses. The PrimFunc below should compile through the default CPU TIR pipeline without leaving a dangling annotation reference.

The input IR matches the unmodified shape, dtype, and computation in the official test_double_buffer_transform test in tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py: A is float32[16,32], B is float32[16], and the local cache has 32 elements. The official test applies InjectDoubleBuffer directly and checks the transformation; the reproducer below exercises that same input through the default compilation pipeline.

Actual behavior

Compilation fails in InjectDoubleBuffer:

File "src/s_tir/transform/inject_double_buffer.cc", line 265,
in DoubleBufferInjector::GetRemappedBuffer
    TVM_FFI_ICHECK(stride.defined());

tvm.error.InternalError:
Check failed: (stride.defined()) is false:

A pass-by-pass trace identifies s_tir.CompactBufferAllocation as the first pass after which the annotation no longer refers to the allocated buffer. Its output has the following structure:

cache = T.alloc_buffer((32,))
cache_1 = T.Buffer((32,))
with T.attr(cache_1.data, "double_buffer_scope", 1):
    for j in range(32):
        cache[j] = A[i, j]
    B[i] = T.float32(0.0)
    for j in range(32):
        B[i] = B[i] + cache[j]

The allocation and accesses use cache, while the annotation uses a distinct buffer variable cache_1 with no corresponding allocation. This was checked using both structural comparison of the attribute's data expression and object identity of its underlying buffer variable; it is not just a difference in printed names.

The stale reference survives until InjectDoubleBuffer. That pass obtains the double-buffer half-size from the associated AllocBuffer.shape[0]; the missing association leaves its internal stride undefined. This is not an unspecified input-tensor stride.

The following controls were reproduced in two fresh Python processes:

In both processes, applying FlattenBuffer directly to the original input, followed by InjectDoubleBuffer, succeeds and retains the correct buffer association. Thus the later FlattenBuffer snapshot should not be mistaken for the first point of failure.

For the local control, I captured the actual input immediately before InjectDoubleBuffer and, only in this one-allocation/one-annotation reproducer, rebound the annotation node to the actual allocation's data expression. I then applied InjectDoubleBuffer and compiled its result with further applications of that pass disabled. Random and exactly representable input checks passed, including different initial values of B. No shape, indexing, arithmetic, or tensor strides were changed.

Reassociating the annotation with the allocation eliminates the assertion in this reproducer, supporting the stale-reference diagnosis.

Environment

• TVM: 0.26.dev0+source.8f328e8
• Build-reported source revision: 8f328e802cfe5e41fcc8f5c17e7582b1c28bfce4
• LLVM: 15.0.7
• OS: Ubuntu 22.04.5 LTS, Linux 6.8.0-52-generic, x86_64
• CPU: Intel Core i9-13900K
• Python: 3.11.16
• Target: {"kind": "llvm", "mtriple": "x86_64-pc-linux-gnu"}
• opt_level=2, tir_pipeline="default"
• No disabled passes in the failing default compilation
• No GPU is required. The failure occurs during compilation, before any input arrays are needed.

Steps to reproduce

Save the following as repro.py and run python repro.py in the recorded TVM environment. This standalone script requires only TVM and uses the original official shape.

import tvm
from tvm.script import ir as I, tirx as T

@I.ir_module
class Module:
    @T.prim_func(s_tir=True)
    def main(
        A: T.Buffer((16, 32), "float32"),
        B: T.Buffer((16,), "float32"),
    ):
        for i in range(16):
            cache = T.alloc_buffer((32,), "float32")
            with T.attr(cache.data, "double_buffer_scope", 1):
                for j in range(32):
                    cache[j] = A[i, j]
                B[i] = T.float32(0.0)
                for j in range(32):
                    B[i] = B[i] + cache[j]

target = tvm.target.Target(
    {"kind": "llvm", "mtriple": "x86_64-pc-linux-gnu"}
)
with tvm.transform.PassContext(opt_level=2):
    tvm.compile(Module, target=target, tir_pipeline="default")

The exact script above was also executed separately and produced the expected stride.defined() failure. Changing the width and all corresponding loop/cache extents from 32 to 33 reproduces the same failure.

Triage

• needs-triage