[Bug][TIR] Array declaration of expr occurred multiple times

[Cookiee235]

Actual behavior

Traceback (most recent call last):
  File "/data/qshenaf/remote_pc/TirFuzz/bugs/topi.nn.simplify_0.py", line 21, in <module>
    ref_output = compile_mod(sch.mod, np_inputs, op_output.shape, op_output.dtype, opt_level=0)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/remote_pc/TirFuzz/bugs/topi.nn.simplify_0.py", line 8, in compile_mod
    ref_mod = tvm.build(mod, target='llvm')
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/python/tvm/driver/build_module.py", line 59, in build
    return tvm.tir.build(mod, target, pipeline)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/python/tvm/tir/build.py", line 173, in build
    mod = pipeline(mod)
          ^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/python/tvm/ir/transform.py", line 238, in __call__
    return _ffi_transform_api.RunPass(self, mod)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "tvm/_ffi/_cython/./packed_func.pxi", line 339, in tvm._ffi._cy3.core.PackedFuncBase.__call__
  File "tvm/_ffi/_cython/./packed_func.pxi", line 270, in tvm._ffi._cy3.core.FuncCall
  File "tvm/_ffi/_cython/./packed_func.pxi", line 259, in tvm._ffi._cy3.core.FuncCall3
  File "tvm/_ffi/_cython/./base.pxi", line 185, in tvm._ffi._cy3.core.CHECK_CALL
  File "/data/qshenaf/envs/tvm/python/tvm/_ffi/base.py", line 468, in raise_last_ffi_error
    raise py_err
  File "tvm/_ffi/_cython/./packed_func.pxi", line 56, in tvm._ffi._cy3.core.tvm_callback
  File "/data/qshenaf/envs/tvm/python/tvm/tir/pipeline.py", line 122, in _pipeline
    mod = tvm.ir.transform.Sequential(passes)(mod)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/python/tvm/ir/transform.py", line 238, in __call__
    return _ffi_transform_api.RunPass(self, mod)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "tvm/_ffi/_cython/./packed_func.pxi", line 339, in tvm._ffi._cy3.core.PackedFuncBase.__call__
  File "tvm/_ffi/_cython/./packed_func.pxi", line 270, in tvm._ffi._cy3.core.FuncCall
  File "tvm/_ffi/_cython/./packed_func.pxi", line 259, in tvm._ffi._cy3.core.FuncCall3
  File "tvm/_ffi/_cython/./base.pxi", line 185, in tvm._ffi._cy3.core.CHECK_CALL
  File "/data/qshenaf/envs/tvm/src/tir/ir/transform.cc", line 121, in tvm::tir::transform::PrimFuncPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
    func = pass_func(std::move(func), mod, pass_ctx);
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/tir/transforms/storage_rewrite.cc", line 1759, in unpack_call<tvm::tir::PrimFunc, 3, tvm::tir::transform::StorageRewrite()::<lambda(tvm::tir::PrimFunc, tvm::IRModule, tvm::transform::PassContext)> >
    false);
^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/tir/transforms/storage_rewrite.cc", line 1716, in tvm::tir::PointerValueTypeRewrite(tvm::tir::PrimFunc, bool, bool, bool, bool, bool, bool, bool, bool)
    rewrite_scalar_read_to_vector_shuffle);
^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/tir/transforms/storage_rewrite.cc", line 1159, in tvm::tir::VectorTypeAccessChecker::VectorTypeAccessChecker(tvm::runtime::Array<tvm::tir::Var, void> const&, tvm::runtime::Map<tvm::tir::Var, tvm::tir::Buffer, void,void> const&, bool, bool)
    OnArrayDeclaration(buffer_var, dtype, extent, BufferVarInfo::kPrimFuncParam);
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/tir/transforms/storage_rewrite.cc", line 1253, in tvm::tir::VectorTypeAccessChecker::OnArrayDeclaration(tvm::tir::Var, tvm::runtime::DataType, tvm::PrimExpr, tvm::tir::BufferVarInfo::DeclarationLocation)
    ICHECK(info_map_.find(buffer.get()) == info_map_.end())
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
tvm.error.InternalError: Traceback (most recent call last):
  4: tvm::tir::transform::PrimFuncPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
        at /data/qshenaf/envs/tvm/src/tir/ir/transform.cc:121
  3: unpack_call<tvm::tir::PrimFunc, 3, tvm::tir::transform::StorageRewrite()::<lambda(tvm::tir::PrimFunc, tvm::IRModule, tvm::transform::PassContext)> >
        at /data/qshenaf/envs/tvm/src/tir/transforms/storage_rewrite.cc:1759
  2: tvm::tir::PointerValueTypeRewrite(tvm::tir::PrimFunc, bool, bool, bool, bool, bool, bool, bool, bool)
        at /data/qshenaf/envs/tvm/src/tir/transforms/storage_rewrite.cc:1716
  1: tvm::tir::VectorTypeAccessChecker::VectorTypeAccessChecker(tvm::runtime::Array<tvm::tir::Var, void> const&, tvm::runtime::Map<tvm::tir::Var, tvm::tir::Buffer, void, void> const&, bool, bool)
        at /data/qshenaf/envs/tvm/src/tir/transforms/storage_rewrite.cc:1159
  0: tvm::tir::VectorTypeAccessChecker::OnArrayDeclaration(tvm::tir::Var, tvm::runtime::DataType, tvm::PrimExpr, tvm::tir::BufferVarInfo::DeclarationLocation)
        at /data/qshenaf/envs/tvm/src/tir/transforms/storage_rewrite.cc:1253
  File "/data/qshenaf/envs/tvm/src/tir/transforms/storage_rewrite.cc", line 1253
InternalError: Check failed: (info_map_.find(buffer.get()) == info_map_.end()) is false: Array declaration of expr occurred multiple times.

Environment

tvm: 0.20.dev0

Steps to reproduce

import tvm
from tvm import te, topi, tir
import numpy as np


expr = te.placeholder([2, 3], dtype='float32', name='expr')
op_output = topi.nn.simplify(expr)
sch = tir.Schedule(te.create_prim_func([expr, op_output]))
tvm.build(sch.mod)

Triage

• needs-triage

cc @Hzfengsy @junrushao @quic-sanirudh @shingjan

[jikechao]

It seems the topi.nn.simplify generated an incorrect IR shown below:

@I.ir_module
class Module:
    @T.prim_func
    def main(var_expr: T.handle, var_expr_1: T.handle):
        T.func_attr({"tir.noalias": T.bool(True)})
        expr = T.match_buffer(var_expr, (2, 3))
        expr = T.match_buffer(var_expr_1, (2, 3), data=expr.data)  # invlid due to duplicate
        T.evaluate(0)