[Relax][TVMScript] Print ExternFunc struct_info when non-default

src/script/printer/relax/function.cc

@@ -22,6 +22,15 @@ namespace tvm {
 namespace script {
 namespace printer {
 
+static bool HasDefaultExternFuncStructInfo(const relax::ExternFunc& n) {
+  const auto* sinfo = n->struct_info_.as<relax::FuncStructInfoNode>();
+  if (sinfo == nullptr || sinfo->params.defined() || sinfo->purity ||
+      !sinfo->ret->IsInstance<relax::ObjectStructInfoNode>()) {
+    return false;
+  }
+  return true;
+}
+
 bool AtTopLevelFunction(const IRDocsifier& d) {
   // fewer than 2 frames: not in a function at all
   if (d->frames.size() < 2) {
@@ -128,8 +137,12 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
 TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
     .set_dispatch<relax::ExternFunc>(  //
         "", [](relax::ExternFunc n, AccessPath n_p, IRDocsifier d) -> Doc {
-          // TODO(@junrushao): print more information out of extern function.
-          return Relax(d, "ExternFunc")->Call({LiteralDoc::Str(n->global_symbol, n_p)});
+          ffi::Array<ExprDoc> args;
+          args.push_back(LiteralDoc::Str(n->global_symbol, n_p->Attr("global_symbol")));
+          if (!HasDefaultExternFuncStructInfo(n)) {
+            args.push_back(d->AsDoc<ExprDoc>(n->struct_info_, n_p->Attr("struct_info_")));
+          }
+          return Relax(d, "ExternFunc")->Call(args);
         });
 
 TVM_SCRIPT_REPR(relax::FunctionNode, ReprPrintRelax);

tests/python/relax/test_tvmscript_printer_relax.py

@@ -98,6 +98,41 @@ def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
     )
 
 
+def test_extern_func_with_struct_info():
+    obj = IRModule(
+        {
+            "my_ext": relax.ExternFunc(
+                "my_ext",
+                relax.FuncStructInfo([], relax.TensorStructInfo(dtype="float32", ndim=2), purity=True),
+            ),
+        }
+    )
+    _assert_print(
+        obj,
+        """
+# from tvm.script import ir as I
+# from tvm.script import relax as R
+
+@I.ir_module
+class Module:
+    my_ext = R.ExternFunc("my_ext", R.Callable((), R.Tensor(dtype="float32", ndim=2), True))
+""",
+    )
+
+
+def test_extern_func_with_struct_info_roundtrip():
+    mod = IRModule(
+        {
+            "my_ext": relax.ExternFunc(
+                "my_ext",
+                relax.FuncStructInfo([], relax.TensorStructInfo(dtype="float32", ndim=2), purity=True),
+            ),
+        }
+    )
+    roundtrip = tvm.script.from_source(mod.script(verbose_expr=True))
+    tvm.ir.assert_structural_equal(mod, roundtrip)
+
+
 def test_nested_function():
     @I.ir_module
     class NestedFunction: