apache / apache/tvm

[Bug] [Relax][BYOC] JSONSerializer fails to handle tir_vars in BYOC for models with symbolic arithmetic

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.