apache / apache/tvm

[Bug][Relax] Default LLVM pipeline crashes in FoldConstant when reshape target is a runtime Shape

Open
#20,193 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 handle a valid reshape whose target shape is runtime shape metadata.

In this pattern:

shape_of(source) -> shape_to_tensor -> tensor_to_shape -> reshape(source, runtime_shape)

the reshape target is a runtime Shape, not a static ShapeExpr. The compiler should either preserve it, lower it correctly, or skip constant-fold legalization for this call. It should not abort inside
FoldConstant.

### Actual behavior

The same module builds and runs correctly with the generic/default pipeline, but the official LLVM target-default pipeline fails inside FoldConstant.

Observed locally:

generic/default: match
official llvm target-default: AssertionError
first bad pass: FoldConstant
manual skip FoldConstant: match
ShapeExpr reshape control: match

The failure is:

AssertionError
assert isinstance(tgt_shape, ShapeExpr)

This appears to happen because FoldConstant invokes the reshape legalization callback. That legalization path assumes a reshape target Var is bound to ShapeExpr, but in this program the target is bound to
R.tensor_to_shape(...), which is a runtime Shape.

### Environment

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

### Steps to reproduce

import numpy as np
import tvm
from tvm import relax, tir

def build_module():
bb = relax.BlockBuilder()

m = tir.Var("m", "int64")
d0 = tir.Var("d0", "int64")
d2 = tir.Var("d2", "int64")
d3 = tir.Var("d3", "int64")

x = relax.Var("x", relax.TensorType((m, 8, 8, 8), "float32"))
begin = relax.Var("begin", relax.TensorType((4,), "int64"))
end = relax.Var("end", relax.TensorType((4,), "int64"))
strides = relax.Var("strides", relax.TensorType((4,), "int64"))

with bb.function("main", params=[x, begin, end, strides]):
with bb.dataflow():
sliced = bb.emit(relax.op.dynamic_strided_slice(x, begin, end, strides))
source = bb.match_cast(
sliced,
relax.TensorType((d0, 8, d2, d3), "float32"),
)

shape0 = bb.emit(relax.op.shape_of(source))
shape_tensor = bb.emit(relax.op.shape_to_tensor(shape0))
shape_back = bb.emit(relax.op.tensor_to_shape(shape_tensor))

view0 = bb.emit(relax.op.reshape(source, shape_back))
view = bb.match_cast(
view0,
relax.TensorType((d0, 8, d2, d3), "float32"),
)

out = bb.emit(relax.op.add(source, view))
gv = bb.emit_output(out)

bb.emit_func_output(gv)

return bb.get()

def make_inputs():
x = np.random.default_rng(0).normal(size=(6, 8, 8, 8)).astype("float32")
begin = np.array([0, 0, 1, 1], dtype="int64")
end = np.array([6, 8, 7, 7], dtype="int64")
strides = np.array([1, 1, 1, 1], dtype="int64")
return [x, begin, end, strides]

def run(pipeline):
mod = build_module()
target = tvm.target.Target("llvm")

exe = relax.build(
mod,
target=target,
relax_pipeline=pipeline,
exec_mode="compiled",
)
vm = relax.VirtualMachine(exe, tvm.cpu())

inputs = make_inputs()
tvm_inputs = [tvm.runtime.tensor(x, tvm.cpu()) for x in inputs]
actual = vm["main"](*tvm_inputs).numpy()

x, begin, end, strides = inputs
source = x[
begin[0] : end[0] : strides[0],
begin[1] : end[1] : strides[1],
begin[2] : end[2] : strides[2],
begin[3] : end[3] : strides[3],
]
expected = source + source

print("match:", np.allclose(actual, expected, rtol=1e-5, atol=1e-5))
print("max_abs:", np.max(np.abs(actual.astype("float64") - expected.astype("float64"))))

target = tvm.target.Target("llvm")

print("generic/default")
run("default")

print("llvm target-default")
# This fails locally in FoldConstant with:
# AssertionError: assert isinstance(tgt_shape, ShapeExpr)
run(relax.get_default_pipeline(target))

Observed output:

generic/default
match: True
max_abs: 0.0

llvm target-default
AssertionError

A pass-prefix trace localizes the first failing pass to FoldConstant:

DispatchSampling: ok
DispatchSortScan: ok
LegalizeOps: ok
AnnotateTIROpPattern: ok
FoldConstant: AssertionError

A ShapeExpr control, where the reshape target is the known symbolic shape tuple instead of tensor_to_shape(shape_to_tensor(shape_of(source))), builds and runs correctly under the same target-default pipeline.

### Triage

- needs-triage
- type: bug
- relax

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the provided Python reproducer and trace the LLVM target-default pipeline through FoldConstant, focusing on the reshape legalization callback and its ShapeExpr assertion. Compare the runtime Shape case with the ShapeExpr control and the manual FoldConstant skip. Done means the LLVM pipeline builds and runs the runtime-shape reshape without aborting, while preserving the expected result.

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
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.