[Bug] [Relax][BYOC] JSONSerializer fails to handle tir_vars in BYOC for models with symbolic arithmetic
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Expected behavior
The BYOC serialization flow should correctly handle tir_vars during graph partitioning. It should support dynamic shape models with symbolic arithmetic. The compiler should complete the code generation stage successfully.
### Actual behavior
The `JSONSerializer` crashes during the code generation stage. It triggers a `tvm.error.InternalError`. The system expects `TensorType` but receives `relax.ShapeType`. This failure occurs when `FuseOpsByPattern` injects `tir_vars` into the fused function signature.
### Environment
- **OS**: Ubuntu 24.04
- **TVM Version**: TVM main branch (Commit: [551be33ed3026ebde5bfe8399940c42c91373e96])
- **BYOC Backend**: official example_npu
### Steps to reproduce
The following script demonstrates the specific behavior of this issue. `test1` contains symbolic arithmetic (e.g., `n // 2`). It generates `tir_vars` and causes a crash during `RunCodegen`. In contrast, `test2` uses a simple symbolic variable `n`. It proceeds without error. This issue persists across any custom BYOC backend that relies on the `JSONSerializer` infrastructure.
```python
import tvm
from tvm.script.parser import relax as R
from tvm.relax.dpl import is_op, wildcard
@tvm.script.ir_module
class InputModule:
@R.function
def test1(x: R.Tensor((1, 16, "n", 16), dtype="float32"), bias: R.Tensor((1, 16, 8, 8), dtype="float32"),):
with R.dataflow():
# test 1
# (1, 16, n, 16) -> (1, 16, n/2, 8)
lv = R.nn.max_pool2d(x, pool_size=(2, 2), strides=(2, 2), padding=(0, 0), layout="NCHW")
gv = R.add(lv, bias)
R.output(gv)
return gv
@R.function
def test2(x: R.Tensor((1, 16, "n", 16), dtype="float32"), bias: R.Tensor((1, 16, 16, 16), dtype="float32"),):
with R.dataflow():
# test 2
gv = R.add(x, bias)
R.output(gv)
return gv
mod = InputModule
mod.show() # Examine the original IR.
patterns = [("example_npu.add", is_op("relax.add")(wildcard(), wildcard()))]
mod = tvm.relax.transform.FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
mod.show() # Examine the post-pass IR.
mod = tvm.relax.transform.RunCodegen()(mod)
```
### Error Log
```
Traceback (most recent call last):
File "/home/zin/tvm_code/test_BYOC/test_tir_vars.py", line 54, in
mod = tvm.relax.transform.RunCodegen()(mod)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
...
File "/home/zin/tvm_new/src/relax/backend/contrib/example_npu/codegen.cc", line 82, in tvm::relax::contrib::ExampleNPUCompiler(tvm::ffi::Array, tvm::ffi::Map, tvm::ffi::Map)
serializer.serialize(func);
File "/home/zin/tvm_new/src/relax/backend/contrib/example_npu/../codegen_json/codegen_json.h", line 242, in tvm::relax::backend::contrib::JSONSerializer::serialize(tvm::relax::Function)
memo_[param] = AddNode(node_ptr, param);
File "/home/zin/tvm_new/src/relax/backend/contrib/example_npu/../codegen_json/codegen_json.h", line 290, in tvm::relax::backend::contrib::NodeEntries tvm::relax::backend::contrib::JSONSerializer::AddNode(tvm::relax::backend::contrib::JSONGraphObjectPtr, const tvm::relax::Expr&)
TVM_FFI_ICHECK(tensor_ty) << "Expect TensorType, but received: " << ty->GetTypeKey();
tvm.error.InternalError: Check failed: (tensor_ty) is false: Expect TensorType, but received: relax.ShapeType
```
### Triage
* type:bug
* byoc:codegen_json
### Additional Note
I implemented a temporary workaround. I modified the traversal logic in `codegen_json.h`. I also patched `SetInputOutputBuffers` in `json_runtime.h`. However, this fix requires manual adjustments for every specific backend. Developers must manually implement logic to detect `tir_vars`. Each backend's `Run()` function must explicitly identify and ignore these symbolic nodes. This repetition increases implementation complexity for hardware vendors.
#### tvm/src/relax/backend/contrib/codegen_json/codegen_json.h
```C++
void serialize(Function func) {
// First we convert all the parameters into input nodes.
for (const auto& param : func->params) {
auto node_ptr = std::make_shared(param->name_hint(), "input" /* op_type_ */);
if(param->name_hint() == "tir_vars"){ // user add
auto idx = func->params.size() - 2; // user add
memo_[param] = AddNode(node_ptr, func->params[idx]); // user add
} else {
memo_[param] = AddNode(node_ptr, param);
}
}
heads_ = VisitExpr(func->body);
}
```
#### tvm/src/runtime/extra/contrib/json/json_runtime.h
```C++
void SetInputOutputBuffers(const ffi::PackedArgs& args) {
ICHECK_EQ(args.size(), input_var_eid_.size() + outputs_.size())
<< "Found mismatch in the number of provided data entryies and required.";
for (size_t i = 0; i < static_cast(args.size()); i++) {
auto eid = i < input_var_eid_.size() ? input_var_eid_[i]
: EntryID(outputs_[i - input_var_eid_.size()]);
const DLTensor* arg;
if (auto opt_nd = args[i].as()) {
Tensor arr = opt_nd.value();
arg = arr.operator->();
} else if(nodes_[eid].name_ == "tir_vars") { // user add start
auto opt_nd = args[i-1].as();
Tensor arr = opt_nd.value();
arg = arr.operator->(); // user add end
} else {
arg = args[i].cast();
}
```
#### Each backend's `Run()`
```C++
if(nodes_[id].GetOpName() == "tir_vars"){
...
```
Contributor guide
No contributing guide indexed for this repository
Research direction
Run the provided Python reproducer through FuseOpsByPattern and RunCodegen, then inspect tvm/src/relax/backend/contrib/codegen_json/codegen_json.h where JSONSerializer handles function parameters and AddNode. Trace how relax.ShapeType parameters representing tir_vars reach serialization and compare the runtime path in tvm/src/runtime/extra/contrib/json/json_runtime.h. Done means the symbolic-arithmetic case completes code generation without backend-specific workarounds and preserves existing simple-symbolic behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend-api-design, compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100