apache / apache/tvm

[Bug] Inconsistency Module Structure in Relax Transform and Build Failure with InlinePrivateFunctions()

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

Description

When applying the relax.transform.InlinePrivateFunctions() optimization to a Relax module using both Sequential and direct application methods, the resulting module structures are inconsistent. Additionally, when using relax.build() after applying the transformation directly, a build failure occurs with the following internal error:
```
InternalError: Check failed: (slot->value_computed) is false: PrimExpr T.int64(4) * n * m in function I.GlobalVar("main") has not been computed.
```

### Expected behavior
The module structures generated by applying relax.transform.InlinePrivateFunctions() using Sequential or direct application should be consistent, and the module should compile successfully without any internal errors when using relax.build().

### Actual behavior

- The structures of the module differ between the two methods of applying the transformation.
- When applying the transformation and using relax.build(), an internal error occurs, indicating that a computation involving n and m is not computed.

### Steps to reproduce

```python
import tvm
from tvm import relax
import numpy as np
from tvm.script import ir as I
from tvm.script import tir as T
from tvm.script import relax as R

@I.ir_module
class Module:
@T.prim_func(private=True)
def add(var_x2: T.handle, var_y2: T.handle, var_T_add: T.handle):
T.func_attr({"tir.noalias": T.bool(True)})
n, m = T.int64(), T.int64()
x2 = T.match_buffer(var_x2, (n, m))
y2 = T.match_buffer(var_y2, (n, m))
T_add = T.match_buffer(var_T_add, (n, m))
for ax0, ax1 in T.grid(n, m):
with T.block("T_add"):
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
T.reads(x2[v_ax0, v_ax1], y2[v_ax0, v_ax1])
T.writes(T_add[v_ax0, v_ax1])
T_add[v_ax0, v_ax1] = x2[v_ax0, v_ax1] + y2[v_ax0, v_ax1]

@R.function(private=True)
def main_inner(x2: R.Tensor(("n", "m"), dtype="float32"), y2: R.Tensor(("n", "m"), dtype="float32")) -> R.Tensor(("n", "m"), dtype="float32"):
n = T.int64()
m = T.int64()
cls = Module
sum_inner = R.call_tir(cls.add, (x2, y2), out_sinfo=R.Tensor((n, m), dtype="float32"))
return sum_inner

@R.function
def main(x1: R.Tensor((10, 5), dtype="float32"), y1: R.Tensor((10, 5), dtype="float32")) -> R.Tensor((10, 5), dtype="float32"):
cls = Module
sum_main: R.Tensor((10, 5), dtype="float32") = cls.main_inner(x1, y1)
return sum_main

mod = Module

mod_seq = tvm.transform.Sequential([relax.transform.InlinePrivateFunctions(),])(mod)
mod = relax.transform.InlinePrivateFunctions()(mod)
#tvm.ir.assert_structural_equal(mod_seq, mod)
print(mod["main"].body.blocks[0].bindings[0].value.sinfo_args[0].shape.values[0])
with tvm.transform.PassContext(opt_level=4):
ex = relax.build(mod, target='llvm')
```

Could you please help confirm if this is a bug in TVM or an issue with my usage?

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the supplied Python reproducer and compare the results of relax.transform.InlinePrivateFunctions() through Sequential and direct application, including the structural-equality check. Then run relax.build(mod, target="llvm") under the shown PassContext and trace the reported uncomputed shape expression. Done means both transformation paths produce consistent module structures and the transformed module builds without the InternalError.

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
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.