[Bug][Relax] Default LLVM pipeline can leave high-level Relax ops after FoldConstant resolves dynamic reshape shape
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Expected behavior
The default LLVM Relax pipeline should not pass high-level Relax tensor ops such as R.add to VM codegen.
In the program below, a reshape target shape is computed from bound tensor parameters:
bound int64 tensors -> tensor_to_shape -> reshape(data, shape) -> add
After constant folding resolves the reshape target to a static shape, the following R.add should be legalized before VM codegen.
### Actual behavior
The official LLVM target-default pipeline fails during build because a high-level R.add remains in the final module.
Observed locally:
official target-default build: error
residual high-level ops after official pipeline: R.add = 1
FoldConstant-first then default build: ok, max_abs = 0.0
LegalizeOps after FoldConstant contrast: ok, max_abs = 0.0
The issue appears to be a pass-order gap:
LegalizeOps runs before FoldConstant
FoldConstant later resolves the reshape shape
no later LegalizeOps pass legalizes the consumer op
VMCodeGen receives residual R.add and fails
### Environment
OS: Linux x86_64
Python: 3.10.12
TVM version: 0.26.dev1
TVM commit: 5a8dae4d95c55c8fec9246a607a28c3ff54ffe05
Target: llvm
Relax VM exec_mode: bytecode
### Steps to reproduce
import numpy as np
import tvm
from tvm import relax
def make_module():
bb = relax.BlockBuilder()
data = relax.Var("data", relax.TensorType((256,), "float32"))
c0 = relax.Var("c0", relax.TensorType((2,), "int64"))
c1 = relax.Var("c1", relax.TensorType((2,), "int64"))
bias = relax.Var("bias", relax.TensorType((16, 16), "float32"))
with bb.function("main", params=[data, c0, c1, bias]):
with bb.dataflow():
doubled = bb.emit(relax.op.add(c0, c0))
target_shape = bb.emit(relax.op.multiply(doubled, c1))
shape = bb.emit(relax.op.tensor_to_shape(target_shape))
reshaped = bb.emit(relax.op.reshape(data, shape))
out = bb.emit(relax.op.add(reshaped, bias))
gv = bb.emit_output(out)
bb.emit_func_output(gv)
mod = bb.get()
return relax.transform.BindParams(
"main",
{
"c0": tvm.runtime.tensor(np.array([8, 8], dtype="int64")),
"c1": tvm.runtime.tensor(np.array([1, 1], dtype="int64")),
"bias": tvm.runtime.tensor(np.ones((16, 16), dtype="float32")),
},
)(mod)
mod = make_module()
target = tvm.target.Target("llvm")
# This fails locally because a high-level R.add remains after the pipeline.
exe = relax.build(
mod,
target=target,
relax_pipeline=relax.get_default_pipeline(target),
exec_mode="bytecode",
)
vm = relax.VirtualMachine(exe, tvm.cpu())
data = tvm.runtime.tensor(np.arange(256, dtype="float32"), tvm.cpu())
out = vm["main"](data).numpy()
print(out)
Observed error:
CodeGenVM cannot handle this intrinsic now
A diagnostic contrast supports the pass-order diagnosis. If FoldConstant is applied before the default build, or if LegalizeOps is run again after FoldConstant, the same program builds and runs correctly with
max_abs=0.0.
This is not specific to R.add. I reproduced the same pattern with multiple consumers including:
add
subtract
multiply
divide
maximum
minimum
permute_dims
concat
broadcast_to
collapse_sum_to
### Triage
- needs-triage
- type: bug
- relax
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by running the provided Python reproducer with relax.get_default_pipeline(target), then inspect the pass ordering around LegalizeOps and FoldConstant. The fix is done when the default LLVM Relax pipeline legalizes the listed consumer ops after dynamic reshape folding, builds successfully, and produces the expected output without residual high-level Relax ops.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100