[Unity][MSC] Avoid depending on trivial bindings in Relax intermediate

python/tvm/contrib/msc/core/codegen/codegen.py

@@ -172,4 +172,18 @@ def relay_to_relax(
  def _bind_weights(mod: tvm.IRModule, folder: msc_utils.MSCDirectory) -> tvm.IRModule:
  return BindParams("main", weights)(mod)
 
- return codegen.load(inputs, post_load=_bind_weights)
+ mod = codegen.load(inputs, post_load=_bind_weights)
+
+ mod = tvm.ir.transform.Sequential(
+ [
+ # The canonicalization of relax variable bindings is not required
+ # for correctness. It does, however, remove trivial `x = y`
+ # bindings, preventing test cases from depending on their
+ # presence.
+ tvm.relax.transform.CanonicalizeBindings(),
+ tvm.relax.transform.ConvertToDataflow(min_size=1),
+ ],
+ name="tvm.contrib.msc.core.codegen.relay_to_relax_postproc",
+ )(mod)
+
+ return mod

python/tvm/contrib/msc/framework/tvm/codegen/codegen.py

@@ -71,4 +71,18 @@ def _bind_weights(mod: tvm.IRModule, folder: msc_utils.MSCDirectory) -> tvm.IRMo
  return mod
 
  codegen = CodeGen(graph, _ffi_api.GetRelaxSources, codegen_config, print_config, build_folder)
- return codegen.load(inputs, pre_load=_save_weights, post_load=_bind_weights)
+ mod = codegen.load(inputs, pre_load=_save_weights, post_load=_bind_weights)
+
+ mod = tvm.ir.transform.Sequential(
+ [
+ # The canonicalization of relax variable bindings is not required
+ # for correctness. It does, however, remove trivial `x = y`
+ # bindings, preventing test cases from depending on their
+ # presence.
+ tvm.relax.transform.CanonicalizeBindings(),
+ tvm.relax.transform.ConvertToDataflow(min_size=1),
+ ],
+ name="tvm.contrib.msc.framework.tvm.codegen.to_relax_postproc",
+ )(mod)
+
+ return mod

src/contrib/msc/framework/tvm/codegen.cc

@@ -71,9 +71,7 @@ void RelaxCodeGen::CodeGenGraph() {
  continue;
  }
  int scope_level = CompareScope(node);
- if (scope_level == 1) {
- stack_.scope_start("block_builder.dataflow()");
- } else if (scope_level == -1) {
+ if (scope_level == -1) {
  stack_.scope_end();
  }
  CodeGenNode(node, config()->use_tools);
@@ -83,13 +81,11 @@ void RelaxCodeGen::CodeGenGraph() {
  for (size_t i = 0; i < scopes().size() - 1; i++) {
  stack_.scope_end();
  }
- } else if (scopes().size() == 0) {
- // start dataflow scope for non-scope graph
- stack_.scope_start("block_builder.dataflow()");
  }
  // mark outputs
  stack_.comment("Emit the outputs");
  Array<String> idx_exits;
+
  for (const auto& e : graph()->GetExits()) {
  const auto& idx_exit = IdxNodeBase(e) + (config()->use_tools ? "_exit" : "");
  if (config()->use_tools) {
@@ -104,10 +100,9 @@ void RelaxCodeGen::CodeGenGraph() {
  stack_.call_arg(DocUtils::ToStr(e->name + "_exit"), "name_hint");
  }
  }
- stack_.func_call("emit_output", idx_exit, "block_builder").call_arg(idx_exit);
  idx_exits.push_back(idx_exit);
  }
- stack_.scope_end();
+
  if (config()->use_tools) {
  stack_.func_call("msc_tools.execute_step", "output").call_arg(DocUtils::ToStr("after_build"));
  if (idx_exits.size() == 1) {

src/contrib/msc/framework/tvm/relax_opcode.cc

@@ -36,9 +36,6 @@ const Array<Doc> RelaxOpCode::GetDocs() {
  if (node()->optype == "input" || node()->optype == "constant" || node()->optype == "shape") {
  emit_var = false;
  }
- if (node()->optype == "tuple" && node()->children.size() == 0) {
- emit_var = false;
- }
  if (emit_var) {
  const auto& name = config()->explicit_name ? node()->name : "";
  BuilderEmit(IdxNode(), name);