Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[TVMScript] Improvements tvm.script.highlight #13438

Merged
merged 2 commits into from
Nov 23, 2022

Conversation

Lunderberg
Copy link
Contributor

  • Automatically use "black" formatter if available.

  • Allow overrides of pygmentize style based on environment variable TVM_PYGMENTIZE_STYLE.

@tvm-bot
Copy link
Collaborator

tvm-bot commented Nov 18, 2022

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

  • No users to tag found in teams: tvmscript See #10317 for details
  • Built docs for commit 46cd938 can be found here.

Generated by tvm-bot

@Lunderberg
Copy link
Contributor Author

This PR started from finding that the dracula style shows up best in my terminal, and wanted to have a way to enable it as the default pygmentize for my environment.

- Automatically use "black" formatter if available.

- Allow overrides of pygmentize style based on environment variable
  `TVM_PYGMENTIZE_STYLE`.
@Lunderberg Lunderberg force-pushed the tvmscript_highlight_improvements branch from a992917 to 1bde12c Compare November 18, 2022 21:53
@junrushao
Copy link
Member

junrushao commented Nov 19, 2022

Using black for formatting is a great idea! On the other hand, sometimes we want to toggle it off for huge IRs, for example, the winograd TIR has two extremely long lines which we dont really need to read much into: https://gist.github.com/zxybazh/0019023b0796a55c9fadb12593852f84. Would love to get some thoughts from you which one is better: formatting by default or not?

@Lunderberg
Copy link
Contributor Author

Oof, yeah, looks like that one becomes pretty heavily indented with black, with several sections starting at indentations around line 60-80. It looks like it is specifically the nested function calls (here the chain of T.select), each of which adds indentation.

I think I'd still lean toward having the formatting enabled by default, since I think it improves readability on average, but it would be good to have a way to disable it when debugging cases where the readability is worsened. Maybe having the argument default be set by an environment variable, so that the formatting could be enabled/disabled across all calls as needed.

@Lunderberg
Copy link
Contributor Author

Hmm, I'm also wondering on the representation of T.select. Both it and T.if_then_else are useful to use in nested chains, which are what causes the poor formatting. If a nested sequence of N select/if_then_else were sugared as a sequence of 2*N + 1 arguments, we could avoid having a lot of the structures that get poorly formatted. So the winograd TIR example, line 37 would be formatted into ~40 flat lines, instead of the ~80 heavily indented lines that it currently formats into.

B[i, j] = T.Select(
    i % 4 == 3 and j % 4 == 3,
    T.float32(1),
    i % 4 == 3 and j % 4 == 2,
    T.float32(0),
    i % 4 == 3 and j % 4 == 1,
    T.float32(0),
    i % 4 == 3 and j % 4 == 0,
    T.float32(0),
    i % 4 == 2 and j % 4 == 3,
    T.float32(0),
    i % 4 == 2 and j % 4 == 2,
    T.float32(1),
    i % 4 == 2 and j % 4 == 1,
    T.float32(1),
    i % 4 == 2 and j % 4 == 0,
    T.float32(-1),
    i % 4 == 1 and j % 4 == 3,
    T.float32(-1),
    i % 4 == 1 and j % 4 == 2,
    T.float32(1),
    i % 4 == 1 and j % 4 == 1,
    T.float32(-1),
    i % 4 == 1 and j % 4 == 0,
    T.float32(0),
    i % 4 == 0 and j % 4 == 3,
    T.float32(0),
    i % 4 == 0 and j % 4 == 2,
    T.float32(0),
    i % 4 == 0 and j % 4 == 1,
    T.float32(0),
    i % 4 == 0 and j % 4 == 0,
    T.float32(1),
    T.float32(0),
)

@junrushao
Copy link
Member

Some background: The reason that winograd conv2d has such a long chained T.Select is that it's used to represent A, B, G matrices, and by unrolling the corresponding loops in our scheduling, the T.Select is constant-folded into scalar constants. Therefore there is no performance degradation if scheduled properly.

On sugaring long chained T.Select: The usecases of sugaring T.Select, if it's only limited to constant matrices in winograd conv, then it's probably less common and I would say it may not be desirable because it requires users to understand a new syntax.

My two-cents: as long as we have a flag to turn it off, in .show() methods, I don't have a strong opinion if it defaults to on or off :-) Alternatively, when the IR is much longer (e.g. +20 extra lines after formatting), we could print out a doc, saying you may turn formatting off by tweaking the flag, etc.

@Lunderberg
Copy link
Contributor Author

Good point on the T.Select optimizing out during unrolling, so it wouldn't be something with semantic meaning early on. That makes sense on having the flag to disable it the formatting. Looks like I initially added black_format=True as an argument to cprint, but not to the methods that call into it. All the obj.show() methods now also have a black_format=True argument.

@Lunderberg Lunderberg force-pushed the tvmscript_highlight_improvements branch from 891a9f4 to 46cd938 Compare November 22, 2022 20:42
Copy link
Member

@junrushao junrushao left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@Lunderberg Lunderberg merged commit e2fc4d7 into apache:main Nov 23, 2022
@Lunderberg Lunderberg deleted the tvmscript_highlight_improvements branch November 23, 2022 02:52
xinetzone pushed a commit to daobook/tvm that referenced this pull request Nov 25, 2022
* [TVMScript] Improvements tvm.script.highlight

- Automatically use "black" formatter if available.

- Allow overrides of pygmentize style based on environment variable
  `TVM_PYGMENTIZE_STYLE`.

- Forwarded `black_format` argument from `show` method to `cprint`
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants