[FIX][TIRx] Remap typed buffer expressions during specialization

[jinhongyii]

Motivation

TIRx buffers are represented by typed variables. When PrimFunc.specialize remaps a buffer, every expression that refers to its BufferVar must be remapped as well.

PrimFuncSpecializer::VisitExpr_(VarNode*) returned variables that were not in the scalar specialization map unchanged. This bypassed the base StmtExprMutator, which is responsible for recognizing and remapping typed buffer variables. As a result, an alias buffer could retain the original data variable after specialization instead of referring to the specialized buffer.

Example

Consider a buffer alias that shares A.data:

@T.prim_func(private=True, s_tir=True)
def before(A_handle: T.handle, n: T.int32):
    A = T.match_buffer(A_handle, (n,), "int32")
    A_alias = T.decl_buffer((n,), "int32", data=A.data)
    A_alias[n - 1] = 42

after = before.specialize({before.params[1]: 8})

The specialized function should still preserve the alias:

A = T.match_buffer(A_handle, (8,), "int32")
A_alias = T.decl_buffer((8,), "int32", data=A.data)
A_alias[7] = 42

Before this PR, A_alias.data could retain the pre-specialization BufferVar. After this PR, it is remapped to the specialized A.data.

What changed

• Delegate unmatched VarNode expressions to the base StmtExprMutator.
• Add a regression test that specializes a symbolic extent and checks that the alias relationship is preserved.

Tests

• Fresh default CMake/Ninja build
• python -m pytest tests/python/tirx-base/test_tir_specialize.py -q (11 passed)
• pre-commit run --files src/tirx/ir/specialize.cc tests/python/tirx-base/test_tir_specialize.py

Ported from mlc-ai/tvm@14bf144.

[jinhongyii]

@tvm-bot rerun

[jinhongyii]

@tvm-bot rerun