Add unit test for round-trip of opaque function

Author: Lunderberg

src/script/printer/relax/utils.h

@@ -84,14 +84,16 @@ inline Optional<ExprDoc> StructInfoAsAnn(const relax::Var& v, const ObjectPath&
   if (!v->struct_info_.defined()) {
     return NullOpt;
   }
+  bool attempt_to_hide_struct_info = !d->cfg->show_all_struct_info;
+
   if (const auto* call = rhs.as<relax::CallNode>()) {
     static const Op& call_tir_op = Op::Get("relax.call_tir");
     static const Op& call_dps_packed_op = Op::Get("relax.call_dps_packed");
     if (call->op.same_as(call_tir_op) || call->op.same_as(call_dps_packed_op)) {
-      return NullOpt;
+      attempt_to_hide_struct_info = true;
     }
   }
-  if (!d->cfg->show_all_struct_info) {
+  if (attempt_to_hide_struct_info) {
     Optional<relax::StructInfo> inferred_sinfo = NullOpt;
     if (auto opt = rhs.as<relax::Call>()) {
       auto call = opt.value();

tests/python/tvmscript/test_tvmscript_roundtrip.py

@@ -21,7 +21,7 @@
 import tvm
 import tvm.testing
 from tvm import tir
-from tvm.script import tir as T, ir as I
+from tvm.script import tir as T, ir as I, relax as R
 
 import numpy as np
 
@@ -3996,6 +3996,24 @@ def func():
         yield make_ir_generator(op, arg)
 
 
+def relax_extern_func():
+    @R.function
+    def func(A: R.Tensor([10, 20], "float32")):
+        func = R.ExternFunc("dummy_func")
+
+        B: R.Tensor([10, 20], "float32") = R.call_dps_packed(
+            func, [A], out_sinfo=R.Tensor([10, 20], "float32")
+        )
+
+        C: R.Tensor(ndim=2, dtype="float32") = R.call_dps_packed(
+            func, [B], out_sinfo=R.Tensor([10, 20], "float32")
+        )
+
+        return C
+
+    return func
+
+
 ir_generator = tvm.testing.parameter(
     launch_env_thread,
     opt_gemm_normalize,
@@ -4081,13 +4099,35 @@ def func():
     *op_of_literal(),
 )
 
+relax_ir_generator = tvm.testing.parameter(
+    relax_extern_func,
+)
+
+show_all_relax_struct_info = tvm.testing.parameter(
+    by_dict={
+        "show_all_struct_info": True,
+        "hide_inferable_struct_info": False,
+    }
+)
+
 
 def test_roundtrip(ir_generator):
     original = ir_generator()
     after_roundtrip = tvm.script.from_source(original.script(show_meta=True))
     tvm.ir.assert_structural_equal(original, after_roundtrip, True)
 
 
+def test_relax_roundtrip(relax_ir_generator, show_all_relax_struct_info):
+    original = relax_ir_generator()
+    after_roundtrip = tvm.script.from_source(
+        original.script(
+            show_meta=True,
+            show_all_struct_info=show_all_relax_struct_info,
+        )
+    )
+    tvm.ir.assert_structural_equal(original, after_roundtrip, True)
+
+
 def test_return_none_no_trailing_type():
     func = return_none()
     script = func.script()