[Bug][Relax] TopologicalSort moves reshape before the MatchCast defining its symbolic dimensions

Start with src/relax/transform/topological_sort.cc, especially BindingOrderCollector, where the issue reports that symbolic dimensions used by ShapeExpr do not create dependency edges from their MatchCast definitions. Run the standalone reproducer in the issue and check whether TopologicalSort keeps the transformed module well formed. Done means the reported Var n is not defined failure is gone while the source module still produces the expected output.

Expected behavior

TopologicalSort() should preserve dependencies introduced by MatchCast, including uses of its symbolic dimensions inside a ShapeExpr. A well-formed, executable dataflow module should remain well-formed after sorting.

Actual behavior

The default pass moves reshape(y, [n, m]) before the MatchCast that defines n and m. The source passes check_well_formed(check_ty=True) and its LLVM VM output is [[1, 2], [3, 4]]. The result fails the same checker:

ValueError: Var n is not defined.

The relevant reordered bindings are:

y = R.add(x, R.const(1.0, "float32"))
out = R.reshape(y, R.shape([n, m]))
dims = R.tensor_to_shape(shape)
bound = R.match_cast(dims, R.Shape([n, m]))

All supplied tensors have positive, concrete dimensions. This is an invalid-IR transformation, not a numerical-tolerance or ONNX-import issue. It reproduces on three freshly parsed copies of the minimal source; an ordinary static Add module remains well-formed under the same pass.

Environment

• Ubuntu 24.04.4 LTS, x86_64; Python 3.12.3; LLVM 18.1.3, CPU only.
• TVM 0.28.dev0, source revision 27576ce2d1b2cf5212ad46d5710e594bdcbb5df0.
• Matching TVM FFI revision df463f9bd269867fc7f4d578d62d80cc2a11d004.
• The script below was executed against that frozen build. It imports neither ONNX nor GraphOpt and needs no model files.
• Current main was checked at e345029cedf0be76a087fcad9c89a97bed4f7204. The affected pass source file is byte-for-byte identical to the tested revision. I have not built or executed current main.

Steps to reproduce

Run the complete standalone script with the TVM build above and NumPy. Source execution and its output check complete before the final well-formedness check raises the error.

"""Standalone LLVM CPU reproducer; no ONNX or GraphOpt dependency."""
import numpy as np
import tvm
from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R
n = I.dynamic("n", dtype="int64")
m = I.dynamic("m", dtype="int64")
@I.ir_module
class Module:
    @R.function
    def main(x: R.Tensor((2, 2), "float32"), shape: R.Tensor((2,), "int64")) -> R.Tensor(dtype="float32", ndim=2):
        with R.dataflow():
            y = R.add(x, R.const(1.0, "float32"))
            dims = R.tensor_to_shape(shape)
            bound = R.match_cast(dims, R.Shape([n, m]))
            out = R.reshape(y, R.shape([n, m]))
            R.output(out)
        return out


assert relax.analysis.check_well_formed(Module, check_ty=True)
before = Module.script(show_meta=True)
executable = relax.build(Module, target=tvm.target.Target("llvm", host="llvm"), relax_pipeline="default")
vm = relax.VirtualMachine(executable, tvm.cpu())
inputs = [np.arange(4, dtype="float32").reshape(2, 2), np.array([2, 2], dtype="int64")]
expected = np.array([[1, 2], [3, 4]], dtype="float32")
actual = vm["main"](*[tvm.runtime.tensor(x) for x in inputs]).numpy()
np.testing.assert_array_equal(actual, expected)
assert Module.script(show_meta=True) == before
print("TVM:", tvm.__version__, flush=True)
print("Source LLVM output:", actual.tolist(), flush=True)
optimized = relax.transform.TopologicalSort()(Module)
print("Target well formed:", bool(relax.analysis.check_well_formed(optimized, check_ty=True)), flush=True)
relax.analysis.well_formed(optimized, check_ty=True)

Root-cause investigation and related work

The dependency collector records edges for Relax variable references, but the ShapeExpr([n, m]) uses do not create an edge back to the MatchCast defining those primitive variables. See BindingOrderCollector at the inspected main revision. Tracking symbolic definitions and their uses, or retaining a conservative ordering around shape-defining bindings, should preserve this dependency.

I searched open and closed issues/PRs for TopologicalSort; the pass-introduction PR #16697 is related background, but I did not find an existing report for this symbolic-dependency failure.

Triage

• needs-triage
• Relax transformation correctness: well-formed input becomes invalid IR.

Report preparation used Codex; the included script was executed in the stated environment.