chen-zin opened a new issue, #20037:
URL: https://github.com/apache/tvm/issues/20037
### 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 <module>
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::relax::Function,
void>, tvm::ffi::Map<tvm::ffi::String, tvm::ffi::Any, void>,
tvm::ffi::Map<tvm::relax::Constant, tvm::ffi::String, void>)
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 have a temporary workaround. I modify the traversal logic in
codegen_json.h. I also patch SetInputOutputBuffers in json_runtime.h. However,
this fix requires manual adjustments for every specific backend.
#### 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<JSONGraphNode>(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<size_t>(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>()) {
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>();
Tensor arr = opt_nd.value();
arg = arr.operator->(); // user add end
} else {
arg = args[i].cast<DLTensor*>();
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
