[Unity][TVMScript] Optionally hide StructInfo that can be inferred #16356

Lunderberg · 2024-01-05T22:14:33Z

By default, TVMScript prints the struct info of every variable being bound, which can become quite verbose. This commit adds the configuration option show_inferable_type_annotations, which determines whether struct info annotations are shown in cases where they can be inferred.

The show_inferable_type_annotations option defaults to True, preserving the current default behavior.

Hzfengsy · 2024-01-08T12:02:26Z

Thanks, @Lunderberg for the improvement. I like this PR to print in a concise format. But I'm not sure the name of type_annotations is suitable. How about changing to struct_info, i.e. show_inferable_struct_info or show_all_struct_info

Lunderberg · 2024-01-08T14:07:57Z

Thank you, @Hzfengsy, and I had been wondering on the name. Referring to struct info in the name would make it clearer how the option impacts printing of relax functions, but would restrict its usage to only relax functions. Referring to type annotations in the name would make it less clear for relax, but would allow the option to have a similar functionality for TIR.

I don't have any strong preference on the name, and will update show_inferable_struct_info to show_all_struct_info.

tqchen · 2024-01-08T14:44:18Z

Just want to bring up some background on annotations and complications in struct info.

Given the complication of the relax struct info annotation, we cannot guanrantee that the struct info can always be inferred in a reproducible manner from rhs. This is also one reason why we always require the annotation to be printed out.

The primary reason is when the struct info itself can involve complicated bindings

def example(x: R.Tensor):
        s = shape_func(x)
        y: R.Tensor[s] = opaque_fn(x)

Removing the annotation R.Tensor[s] again would cause the deduction of opaque_fn to retrigger, which results in the parsed back function as

def example(x: R.Tensor):
        s = shape_func(x)
        s1 = shape_func(x)
        y: R.Tensor[s1] = opaque_fn(x)

While semantically it is still ok after DCE, however, it is not always desirable. more complicated case can ocurr when the set of annotations involve other set of computations, where some may have common sub expressions (e.g. same shape_func calls being combined after some pass).

The main take-away here is that in order for the parser to be able to fully roundtrippable, we would need to rely on the annotations on the lhs.

While it is ok to toggle off the signature printing, I think we should note that this will no longer guarantee roundtripping.

Lunderberg · 2024-01-08T18:41:36Z

@tqchen Thank you on the details, and agreed that this is primarily for ease of readability. Until and unless the hidden struct info can round-trip in all cases, it shouldn't be used for saving results.

For the opaque functions, I think the current implementation will avoid the issue you described. In order to hide the struct info, three conditions need to hold. (1) The show_all_struct_info config is False, (2) the struct info of the expression can be inferred, and (3) the inferred struct info is structurally equal to the actual struct info. I think the example you gave would fail the third condition, with R.Tensor[s] and R.Tensor[s1] being recognized as different.

I don't see any unit tests that currently use opaque functions. Are there examples somewhere so I could add an appropriate unit test?

tqchen · 2024-01-08T19:58:51Z

Get it, this is mainly for future proofing so no test cases is necessary atm. Thanks for explaination, as long as it defaults to false and we know the caveat i think it is ok

Lunderberg · 2024-01-08T22:03:17Z

Sounds good. The test case would be primarily for future-proofing. Since this can improve readability for many cases, I can imagine a later PR being submitted to change the default, especially if the author is unaware of the potential pitfalls (i.e. like me prior to this thread). A test case for a known case that is difficult to round-trip correctly would catch the error when that hypothetical PR is submitted.

slyubomirsky · 2024-01-09T03:48:24Z

I am very much in favor of having test cases to avoid regression.

Lunderberg · 2024-01-22T19:13:37Z

Updated to target the main branch, and to include unit tests for round-trip of opaque functions, with and without displaying the inferable struct info.

Lunderberg · 2024-01-24T15:31:56Z

CI is now passing on main, including the unit test to validate that the new behavior preserves round-trips of opaque functions, so this PR is ready for (re-)review.

slyubomirsky · 2024-02-06T20:55:55Z

src/script/printer/relax/utils.h

+      attempt_to_hide_struct_info = true;
+    }
+  }
+  if (attempt_to_hide_struct_info) {


Hm, I wonder if this is more complicated than what we would really want. Perhaps a syntactic rule for when to write StructInfo might make more sense? E.g., have it for op calls but not tuple indices, etc.? I'm not sure what situations are sufficiently non-obvious to require an annotation if a user wants to hide them. It's not wrong, but it seems a little unusual for the TVMScript printer to be running type inference itself.

My worry with having separate rules for it is that those rules could diverge from the actual struct inference. Even when the syntactic rules allow the struct info to be hidden, it would still need to compare against the actual struct info to see if the annotations provide different information. That comparison would require StructInfo to compare against, and so this function would effectively be another implementation of InferStructInfo, but one which could get out of sync with the canonical version.

By implementing it in terms of the struct inference, the same inference rules are used for both parsing and printing. Just as parsing can fill in missing information by calling the type inference, the printing can check if information can be omitted by calling the type inference.

It seems a little like overkill, but we are running the normalizer between every pass anyway so I guess it's not so bad.

e: And you're right, it's futureproof too.

Yeah, probably overkill, but I tend to be more okay with overkill for non-default behavior. Good point on the comparison with normalizing on each pass.

slyubomirsky

As a futureproof way of condensing the amount of StructInfo printed, this seems fine given the discussion above.

Lunderberg · 2024-02-08T02:01:46Z

Sounds good, and thank you for the review! Before merging, I'm going to rebase onto main, since the CI results look a bit out of date.

By default, TVMScript prints the struct info of every variable being bound, which can become quite verbose. This commit adds the configuration option `show_inferable_type_annotations`, which determines whether struct info annotations are shown in cases where they can be inferred. The `show_inferable_type_annotations` option defaults to `True`, preserving the current default behavior.

…16356

Lunderberg · 2024-02-09T20:58:44Z

Merged from main to PR branch, to resolve CI breakage that was fixed in #16546.

MasterJH5574 · 2024-02-14T15:23:07Z

Ooops looks like the Python decorator here causes clang warning when compiling every .cc file

tvm/include/tvm/node/script_printer.h

Lines 89 to 91 in 2b813ec

    
              *     # func.show(show_all_struct_info=True) 
        
              *     @R.function 
        
              *     def func(

[ 82%] Building CXX object CMakeFiles/tvm_objs.dir/src/relay/transforms/flatten_atrous_conv.cc.o
In file included from /home/ruihang/Workspace/tvm/src/relay/transforms/annotate_texture_storage.cc:36:
In file included from /home/ruihang/Workspace/tvm/include/tvm/relay/attrs/nn.h:27:
In file included from /home/ruihang/Workspace/tvm/include/tvm/ir/attrs.h:48:
In file included from /home/ruihang/Workspace/tvm/include/tvm/ir/expr.h:27:
In file included from /home/ruihang/Workspace/tvm/include/tvm/ir/source_map.h:26:
In file included from /home/ruihang/Workspace/tvm/include/tvm/node/node.h:38:
In file included from /home/ruihang/Workspace/tvm/include/tvm/node/repr_printer.h:27:
/home/ruihang/Workspace/tvm/include/tvm/node/script_printer.h:90:10: warning: unknown command tag name [-Wdocumentation-unknown-command]
   *     @R.function
         ^~
1 warning generated.
[ 82%] Building CXX object CMakeFiles/tvm_objs.dir/src/relay/transforms/fold_constant.cc.o
In file included from /home/ruihang/Workspace/tvm/src/relay/collage/sub_graph.cc:25:
In file included from /home/ruihang/Workspace/tvm/src/relay/collage/./sub_graph.h:28:
In file included from /home/ruihang/Workspace/tvm/include/tvm/ir/transform.h:59:
In file included from /home/ruihang/Workspace/tvm/include/tvm/ir/diagnostic.h:29:
In file included from /home/ruihang/Workspace/tvm/include/tvm/ir/module.h:27:
In file included from /home/ruihang/Workspace/tvm/include/tvm/ir/adt.h:30:
In file included from /home/ruihang/Workspace/tvm/include/tvm/ir/expr.h:27:
In file included from /home/ruihang/Workspace/tvm/include/tvm/ir/source_map.h:26:
In file included from /home/ruihang/Workspace/tvm/include/tvm/node/node.h:38:
In file included from /home/ruihang/Workspace/tvm/include/tvm/node/repr_printer.h:27:
/home/ruihang/Workspace/tvm/include/tvm/node/script_printer.h:90:10: warning: unknown command tag name [-Wdocumentation-unknown-command]
   *     @R.function
         ^~
1 warning generated.

I wonder if we want to find a work around here?

Hzfengsy · 2024-02-14T15:24:11Z

Fixed in #16571

Lunderberg requested a review from junrushao January 5, 2024 22:14

Lunderberg force-pushed the unity_tvmscript_optional_hiding_of_sinfo branch from a87931a to 6e60cc3 Compare January 22, 2024 19:10

Lunderberg changed the base branch from unity to main January 22, 2024 19:11

Lunderberg force-pushed the unity_tvmscript_optional_hiding_of_sinfo branch from 6e60cc3 to b702e13 Compare January 22, 2024 20:57

slyubomirsky reviewed Feb 6, 2024

View reviewed changes

slyubomirsky approved these changes Feb 6, 2024

View reviewed changes

Lunderberg added 3 commits February 8, 2024 02:02

Rename show_inferable_type_annotations to show_all_struct_info

75282b6

Add unit test for round-trip of opaque function

f11725f

Lunderberg force-pushed the unity_tvmscript_optional_hiding_of_sinfo branch from b702e13 to f11725f Compare February 8, 2024 02:02

Merge branch 'main' into unity_tvmscript_optional_hiding_of_sinfo_pr_…

6ebeca3

…16356

Lunderberg merged commit bf43947 into apache:main Feb 12, 2024
18 checks passed

Lunderberg deleted the unity_tvmscript_optional_hiding_of_sinfo branch February 12, 2024 18:06

ysh329 mentioned this pull request Apr 21, 2024

[Release] v0.16.0 Release Candidate Notes #16911

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Unity][TVMScript] Optionally hide StructInfo that can be inferred #16356

[Unity][TVMScript] Optionally hide StructInfo that can be inferred #16356

Lunderberg commented Jan 5, 2024

Hzfengsy commented Jan 8, 2024

Lunderberg commented Jan 8, 2024

tqchen commented Jan 8, 2024

Lunderberg commented Jan 8, 2024

tqchen commented Jan 8, 2024

Lunderberg commented Jan 8, 2024

slyubomirsky commented Jan 9, 2024

Lunderberg commented Jan 22, 2024

Lunderberg commented Jan 24, 2024

slyubomirsky Feb 6, 2024

Lunderberg Feb 6, 2024

slyubomirsky Feb 6, 2024 •

edited

Loading

Lunderberg Feb 8, 2024

slyubomirsky left a comment

Lunderberg commented Feb 8, 2024

Lunderberg commented Feb 9, 2024

MasterJH5574 commented Feb 14, 2024

Hzfengsy commented Feb 14, 2024

[Unity][TVMScript] Optionally hide StructInfo that can be inferred #16356

[Unity][TVMScript] Optionally hide StructInfo that can be inferred #16356

Conversation

Lunderberg commented Jan 5, 2024

Hzfengsy commented Jan 8, 2024

Lunderberg commented Jan 8, 2024

tqchen commented Jan 8, 2024

Lunderberg commented Jan 8, 2024

tqchen commented Jan 8, 2024

Lunderberg commented Jan 8, 2024

slyubomirsky commented Jan 9, 2024

Lunderberg commented Jan 22, 2024

Lunderberg commented Jan 24, 2024

slyubomirsky Feb 6, 2024

Choose a reason for hiding this comment

Lunderberg Feb 6, 2024

Choose a reason for hiding this comment

slyubomirsky Feb 6, 2024 • edited Loading

Choose a reason for hiding this comment

Lunderberg Feb 8, 2024

Choose a reason for hiding this comment

slyubomirsky left a comment

Choose a reason for hiding this comment

Lunderberg commented Feb 8, 2024

Lunderberg commented Feb 9, 2024

MasterJH5574 commented Feb 14, 2024

Hzfengsy commented Feb 14, 2024

slyubomirsky Feb 6, 2024 •

edited

Loading