[Bug][Relax][LLVM] scatter_elements with symbolic H/W extents times out during build[Bug]
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Expected behavior
Compiling a Relax module containing R.scatter_elements with symbolic tensor extents should either finish successfully or fail with a clear compiler error.
The static version of the same computation with concrete shape (6, 8, 6, 6) builds and runs correctly, so the symbolic-shape version should not hang during LLVM codegen.
### Actual behavior
A minimal R.scatter_elements program with symbolic height and width extents times out during build with the official LLVM target-default pipeline.
Observed locally:
static shape (6, 8, 6, 6): builds/runs successfully, output matches
symbolic H/W shape (6, 8, d2, d3): build-only timeout after 30 seconds
A pass trace can finish through AttachGlobalSymbol, so the timeout appears after Relax lowering, during later TIR/LLVM build.
The generated TIR for the symbolic case contains many repeated symbolic product/index expressions involving d2 * d3. In my local trace, the symbolic H/W case had 162 symbolic product occurrences, while the
static control built quickly.
### 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 multiprocessing as mp
import time
import numpy as np
import tvm
from tvm import relax, tir
def build_static_module():
bb = relax.BlockBuilder()
data = relax.Var("data", relax.TensorType((6, 8, 6, 6), "float32"))
idx = relax.Var("idx", relax.TensorType((6, 4, 6, 6), "int64"))
updates = relax.Var("updates", relax.TensorType((6, 4, 6, 6), "float32"))
with bb.function("main", params=[data, idx, updates]):
out = bb.emit(
relax.op.scatter_elements(data, idx, updates, axis=1, reduction="update")
)
bb.emit_func_output(out)
return bb.get()
def build_symbolic_hw_module():
d2 = tir.Var("d2", "int64")
d3 = tir.Var("d3", "int64")
bb = relax.BlockBuilder()
data = relax.Var("data", relax.TensorType((6, 8, d2, d3), "float32"))
idx = relax.Var("idx", relax.TensorType((6, 4, d2, d3), "int64"))
updates = relax.Var("updates", relax.TensorType((6, 4, d2, d3), "float32"))
with bb.function("main", params=[data, idx, updates]):
out = bb.emit(
relax.op.scatter_elements(data, idx, updates, axis=1, reduction="update")
)
bb.emit_func_output(out)
return bb.get()
def make_inputs():
rng = np.random.default_rng(0)
data = rng.normal(size=(6, 8, 6, 6)).astype("float32")
updates = rng.normal(size=(6, 4, 6, 6)).astype("float32")
idx = np.broadcast_to(
np.arange(4, dtype="int64").reshape(1, 4, 1, 1),
(6, 4, 6, 6),
).copy()
return data, idx, updates
def run_static():
target = tvm.target.Target("llvm")
exe = relax.build(
build_static_module(),
target=target,
relax_pipeline=relax.get_default_pipeline(target),
exec_mode="bytecode",
)
vm = relax.VirtualMachine(exe, tvm.cpu())
data, idx, updates = make_inputs()
actual = vm["main"](
tvm.runtime.tensor(data, tvm.cpu()),
tvm.runtime.tensor(idx, tvm.cpu()),
tvm.runtime.tensor(updates, tvm.cpu()),
).numpy()
expected = data.copy()
np.put_along_axis(expected, idx, updates, axis=1)
print("static match:", np.allclose(actual, expected, rtol=1e-5, atol=1e-5))
def build_symbolic_only():
target = tvm.target.Target("llvm")
relax.build(
build_symbolic_hw_module(),
target=target,
relax_pipeline=relax.get_default_pipeline(target),
exec_mode="bytecode",
)
def run_with_timeout(fn, timeout_sec):
proc = mp.Process(target=fn)
start = time.time()
proc.start()
proc.join(timeout_sec)
if proc.is_alive():
proc.terminate()
proc.join()
print("timeout after", timeout_sec, "seconds")
else:
print("finished in", round(time.time() - start, 3), "seconds, exitcode", proc.exitcode)
if __name__ == "__main__":
run_static()
run_with_timeout(build_symbolic_only, timeout_sec=30)
Observed output:
static match: True
timeout after 30 seconds
Additional local controls:
one-op static scatter_elements: ok
one-op batch-only symbolic scatter_elements: ok
one-op height-only symbolic scatter_elements: ok, but slower
one-op width-only symbolic scatter_elements: ok
one-op height+width symbolic scatter_elements: timeout
symbolic gather_elements/take/scatter_nd/slice_scatter controls: build quickly
This suggests the slow path is specific to scatter_elements lowering/codegen with multiple symbolic spatial extents, not a general symbolic-index build problem.
### Triage
- needs-triage
- type: bug
- relax
- llvm
Contributor guide
No contributing guide indexed for this repository
Research direction
Run build_symbolic_only from the reproducer with the default Relax pipeline and LLVM target, then inspect the pass trace after AttachGlobalSymbol and compare symbolic TIR with the static control. The issue is done when the symbolic H/W case no longer hangs during build and either completes successfully or emits a clear compiler error, while the static case and output behavior remain correct.
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
- 52/100