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

[chen-zin]

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: [551be33])
• 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.

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 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

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

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*>();
      }

Each backend's Run()

if(nodes_[id].GetOpName() == "tir_vars"){
...

[tlopex]

Thanks for reporting it! would have a look into it