apache / apache/tvm

[Bug][Relax] FuseOps can reorder call_tir before symbolic shape binding used in out_ty

Open
#20,195 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

### Expected behavior

FuseOps should preserve symbolic shape variable dependencies in Relax dataflow blocks.

If a call_tir output type uses a symbolic shape variable that is defined by an earlier R.match_cast(..., R.Shape([...])), then FuseOps should not move that call_tir before the match_cast.

### Actual behavior

For a well-formed Relax module, FuseOps can move:

second_copy = R.call_tir(cls.copy_2d, (x,), out_ty=R.Tensor((s1, 8), dtype="float32"))

before:

second_shape_bound = R.match_cast(second_shape, R.Shape([s1, 8]))

After FuseOps, the module is not well-formed:

Symbolic Var s1 is not defined.

The same input builds and runs correctly with relax_pipeline="default". The failure appears in the official LLVM target-default pass prefix when FuseOps is applied.

This looks like FuseOps misses shape-expression dependencies that appear only in call_tir output struct info / out_ty.

### Environment

OS: Linux x86_64
Python: 3.10.12
TVM version: 0.26.dev1
TVM commit: 5a8dae4d95c55c8fec9246a607a28c3ff54ffe05
Target: llvm

### Steps to reproduce

import numpy as np
import tvm
from tvm import relax, tirx
from tvm.relax.backend.cpu_generic import pipeline as cpu_pipeline
from tvm.script import ir as I
from tvm.script import tirx as T

@I.ir_module(s_tir=True)
class ExplicitShapeCallTIRFuncs:
@T.prim_func(private=True, s_tir=True)
def make_shape(var_x: T.handle, out: T.Buffer((T.int64(2),), "int64")):
T.func_attr({"op_pattern": 0, "tirx.noalias": True})
m = T.int64()
_ = T.match_buffer(var_x, (m, T.int64(8)))
out[T.int64(0)] = m
out[T.int64(1)] = T.int64(8)

@T.prim_func(private=True, s_tir=True)
def copy_2d(var_x: T.handle, var_out: T.handle):
T.func_attr({"op_pattern": 8, "tirx.noalias": True})
m = T.int64()
_ = T.match_buffer(var_x, (m, T.int64(8)))
n = T.int64()
out = T.match_buffer(var_out, (n, T.int64(8)))
for i, j in T.grid(n, T.int64(8)):
out[i, j] = T.float32(1.0)

def emit_shape_bound_copy(bb, x, shape_gv, copy_gv, shape_var, name_prefix):
shape_tensor = bb.emit(
relax.call_tir(shape_gv, (x,), out_ty=relax.TensorType((2,), "int64")),
name_hint=f"{name_prefix}_shape_tensor",
)
shape = bb.emit(relax.op.tensor_to_shape(shape_tensor), name_hint=f"{name_prefix}_shape")
bb.match_cast(
shape,
relax.ShapeType([shape_var, 8]),
name_hint=f"{name_prefix}_shape_bound",
)
return bb.emit(
relax.call_tir(copy_gv, (x,), out_ty=relax.TensorType((shape_var, 8), "float32")),
name_hint=f"{name_prefix}_copy",
)

def build_module():
base_mod = ExplicitShapeCallTIRFuncs
shape_gv = base_mod.get_global_var("make_shape")
copy_gv = base_mod.get_global_var("copy_2d")
bb = relax.BlockBuilder(base_mod)

m = tirx.Var("m", "int64")
d0 = tirx.Var("d0", "int64")
s0 = tirx.Var("s0", "int64")
s1 = tirx.Var("s1", "int64")
x = relax.Var("x", relax.TensorType((m, 8), "float32"))

with bb.function("main", params=[x]):
with bb.dataflow():
src = emit_shape_bound_copy(bb, x, shape_gv, copy_gv, s0, "first")
source = bb.match_cast(src, relax.TensorType((d0, 8), "float32"))
left0 = bb.match_cast(source, relax.TensorType((d0, 8), "float32"))
keep = bb.emit(relax.Tuple([source, left0]))
left = bb.emit(relax.TupleGetItem(keep, 1))

rhs = emit_shape_bound_copy(bb, x, shape_gv, copy_gv, s1, "second")
same = bb.match_cast(rhs, relax.TensorType((d0, 8), "float32"))
right = bb.emit(relax.op.subtract(same, relax.const(0.375, "float32")))
out = bb.emit(relax.op.add(left, right))
gv = bb.emit_output(out)

bb.emit_func_output(gv)

return bb.get()

def official_llvm_passes_to_fuseops():
target = tvm.target.Target("llvm")
with target:
passes = (
cpu_pipeline.library_dispatch_passes(target)
+ cpu_pipeline.legalize_passes(target)
+ cpu_pipeline.dataflow_lower_passes(target)
+ cpu_pipeline.finalize_passes(target)
)

for pass_obj in passes:
yield pass_obj.info.name, pass_obj
if pass_obj.info.name == "FuseOps":
break

mod = build_module()
relax.analysis.well_formed(mod)

# Generic/default succeeds and matches the expected value.
target = tvm.target.Target("llvm")
exe = relax.build(mod, target=target, relax_pipeline="default", exec_mode="compiled")
vm = relax.VirtualMachine(exe, tvm.cpu())

x_np = np.linspace(-1.0, 1.0, 6 * 8, dtype="float32").reshape(6, 8)
actual = vm["main"](tvm.runtime.tensor(x_np, tvm.cpu())).numpy()
expected = np.ones((6, 8), dtype="float32") + (
np.ones((6, 8), dtype="float32") - np.float32(0.375)
)
np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5)
print("generic/default: ok")

# Official LLVM target-default prefix becomes non-well-formed after FuseOps.
current = mod
with target:
for name, pass_obj in official_llvm_passes_to_fuseops():
current = pass_obj(current)
print("after", name)
relax.analysis.well_formed(current)

Observed output:

generic/default: ok
after DispatchSampling
after DispatchSortScan
after LegalizeOps
after AnnotateTIROpPattern
after FoldConstant
after FuseOps
ValueError: Symbolic Var s1 is not defined.

Before FuseOps, the second shape bridge is ordered correctly:

second_shape_tensor = R.call_tir(cls.make_shape, (x,), out_ty=R.Tensor((2,), dtype="int64"))
second_shape = R.tensor_to_shape(second_shape_tensor)
second_shape_bound = R.match_cast(second_shape, R.Shape([s1, 8]))
second_copy = R.call_tir(cls.copy_2d, (x,), out_ty=R.Tensor((s1, 8), dtype="float32"))

After FuseOps, the call_tir using s1 appears before s1 is defined:

second_copy = R.call_tir(cls.copy_2d, (x,), out_ty=R.Tensor((s1, 8), dtype="float32"))
same = R.match_cast(second_copy, R.Tensor((d0, 8), dtype="float32"))
...
second_shape_bound = R.match_cast(second_shape, R.Shape([s1, 8]))

### Triage

- needs-triage
- type: bug
- relax

Contributor guide

No contributing guide indexed for this repository

Research direction

Run the provided Python reproducer and inspect the official LLVM target-default pass sequence through FuseOps, especially how call_tir out_ty and R.match_cast shape dependencies are handled. Use relax.analysis.well_formed after FuseOps to verify that symbolic shape variables remain defined before use; done means the module stays well-formed and the existing default-pipeline behavior is preserved.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.