[Bug] `CompactBufferAllocation` leaves `double_buffer_scope` pointing at an unallocated buffer → `InjectDoubleBuffer` ICHECK `stride.defined()`

[YuanchengJiang]

The following code:

import tvm
from tvm.script import tirx as T

@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")
        T.attr(cache.data, "double_buffer_scope", 1)
        for j in range(32):
            cache[j] = A[i, j]
        B[i] = 0.0
        for j in range(32):
            B[i] = B[i] + cache[j]

tvm.compile(main, target="llvm")

Resulted in this output:

Traceback (most recent call last):
  File "min.py", line 15, in <module>
  File "python/tvm/driver/build_module.py", line 111, in compile
  File "src/s_tir/transform/inject_double_buffer.cc", line 440, in operator()
  File "src/s_tir/transform/inject_double_buffer.cc", line 164, in tvm::s_tir::DoubleBufferInjector::Inject(tvm::tirx::Stmt)
  File "src/s_tir/transform/inject_double_buffer.cc", line 200, in tvm::s_tir::DoubleBufferInjector::Mutate_(tvm::tirx::ForNode const*, tvm::ffi::InplaceMode)
  File "src/s_tir/transform/inject_double_buffer.cc", line 169, in tvm::s_tir::DoubleBufferInjector::Mutate_(tvm::tirx::AttrStmtNode const*, tvm::ffi::InplaceMode)
  File "src/s_tir/transform/inject_double_buffer.cc", line 395, in tvm::s_tir::DoubleBufferInjector::MakeProducer(tvm::tirx::AttrStmtNode const*, tvm::ffi::InplaceMode)
  File "src/s_tir/transform/inject_double_buffer.cc", line 327, in tvm::tirx::BufferVar tvm::s_tir::DoubleBufferInjector::GetRemappedBuffer(tvm::tirx::BufferVar, tvm::PrimExpr)
tvm.error.InternalError: Check failed: (stride.defined()) is false: 

To reproduce: python3 min.py

The IR is the Before module of test_double_buffer_transform in tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py; InjectDoubleBuffer alone handles it fine. The crash comes from s_tir.transform.CompactBufferAllocation, which runs earlier in the default pipeline. It rewrites the double_buffer_scope AttrStmt to reference a fresh buffer (cache_1), while the AllocBuffer and every access still use cache:

for i in range(16):
    cache = T.alloc_buffer((32,))
    cache_1 = T.Buffer((32,))          # never allocated
    with T.attr(cache_1.data, "double_buffer_scope", 1):
        for j in range(32):
            cache[j] = A[i, j]
        ...

InjectDoubleBuffer then finds no allocation for the attr's buffer, so its stride is never set. (tirx.analysis.verify_well_formed also accepts the module after CompactBufferAllocation, despite the undefined buffer.) Standalone:

mod = tvm.IRModule({"main": main})
print(tvm.s_tir.transform.CompactBufferAllocation()(mod))

TVM commit: e269315c90e3a061c9e1c77b370ce883b1b223f4

This bug was found by fusion-fuzz

Triage

• [tir:transform]

[anbv29]

Hi! I’d like to create a minimal reproducer for this fusion regression and use it to guide a narrowly scoped fix and test. Is there a preferred Relax test location?