[Bug] InternalError "Check failed: indices.size() == 1 (2 vs. 1): CodeGenLLVM requires all buffers to be flat 1-d buffers"

[Thrsu]

When building a Relax module with a TIR function that uses axis separators in the buffer layout, TVM throws an InternalError during the compilation process. The error message suggests that TVM's LLVM code generation requires all buffers to be flat 1-dimensional, but the provided buffer has more than one index.

The issue appears to be related to how the axis_separators attribute is handled in match_buffer, which results in a mismatch during LLVM code generation.

Steps to reproduce

import tvm
from tvm import relax

from tvm.script import ir as I
from tvm.script import tir as T
from tvm.script import relax as R

@I.ir_module
class Module:
    @T.prim_func(private=True)
    def te_layout_transform_axis_separator(x: T.Buffer((T.int64(16),), "float32"), var_te_layout_transform_axis_separator: T.handle):
        T.func_attr({"tir.noalias": T.bool(True)})
        te_layout_transform_axis_separator = T.match_buffer(var_te_layout_transform_axis_separator, (T.int64(4), T.int64(4)), axis_separators=[1])
        # with T.block("root"):
        for self in range(T.int64(16)):
            with T.block("te_layout_transform_axis_separator"):
                v_self = T.axis.spatial(T.int64(16), self)
                T.reads(x[v_self])
                T.writes(te_layout_transform_axis_separator[v_self // T.int64(4), v_self % T.int64(4)])
                te_layout_transform_axis_separator[v_self // T.int64(4), v_self % T.int64(4)] = x[v_self]

    @R.function
    def main(x: R.Tensor((16,), dtype="float32")) -> R.Tensor((4, 4), dtype="float32"):
        cls = Module
        with R.dataflow():
            lv = R.call_tir(cls.te_layout_transform_axis_separator, (x,), out_sinfo=R.Tensor((4, 4), dtype="float32"))
            R.output(lv)
        return lv

mod = Module
ex = relax.build(mod, target='llvm')

Actual behavior

Traceback (most recent call last):
  File "path/to/script.py", line 31, in <module>
    ex = relax.build(mod, target='llvm')
  ...
tvm.error.InternalError: Traceback (most recent call last):
  ...
  File "/software/tvm/src/target/llvm/codegen_llvm.cc", line 960
InternalError: Check failed: indices.size() == 1 (2 vs. 1) : CodeGenLLVM requires all buffers to be flat 1-d buffers.

Could you please help confirm if this is a bug in TVM or an issue with my usage?