apache / apache/tvm

[Bug][Relax] Default LLVM pipeline can leave high-level Relax ops after FoldConstant resolves dynamic reshape shape

Open
#20,192 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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.