[UX] highlight tvm script

[ganler]

See tlc-pack/relax#185

cc: @YuchenJin @junrushao1994 @Hzfengsy

[YuchenJin]

Thanks @ganler, LGTM!

[tkonolige]

Would it be possible to use ansi colors instead of hardcoding specific rgb values. That way the printed code would match the user's terminal colors?

Here is with rgb colors:

And here is with ansi colors. Notice how it matches my prompt and other text in my terminal.

This is the theme I used:

        Keyword: "bold ansimagenta",
        Keyword.Namespace: "ansicyan",
        Keyword.Type: "ansibrightblue",
        Name.Function: "bold ansiyellow",
        Name.Class: "bold ansiblue",
        Name.Decorator: "italic ansibrightmagenta",
        String: "ansibrightyellow",
        Number: "ansibrightgreen",
        Operator: "ansired",
        Operator.Word: "ansibrightcyan",
        Comment: "italic ansigray",

[ganler]

@tkonolige Thanks for the suggestions.

To put some background: The reason why I use these magic RGB values is to match exact themes in jupyter notebook (or vscode) so they will look compatible. In jupyter notebook, the colors will be displayed as what the RGB values exactly are since it uses HTML to render things. While for terminals, I guess it will map things to nearest 3/4/8-bit colors.

That said I fully agree that we want the theme look compatible for both notebook users and terminal users. My proposal is to: 1) first set them as different theme styles (including the theme you suggested that we have 3 tvm-builtin themes where 2 for jupyter notebook environemnts and 1 for terminal env); 2) we set the default style to a magic string, say "auto", and then we detect if the user is doing experiements in a notebook (if notebook -> "light"; if terminal -> ansi colors).

[areusch]

@ganler thanks for this PR! would be great to add highlighting to TVMScript.

on the python deps: unfortunately we're still working through a PR to generate a virtualenv from gen_requirements.txt, so you'll also need to add it to docker/install/ubuntu_install_python_package.sh for now.

regarding the actual PR, some thoughts: in general it'd be good to avoid adding dependencies to the TVM core unless they're clearly beneficial to a large set of code paths. the approach here seems to be to lex the stringified TVMScript and then highlight based on tokens. I wonder if it would be possible to instead modify the TVMScript printer? The benefit of doing this is that then, you could automatically determine whether you wanted to highlight TVMScript by checking whether it was being operator<< to a TTY. otherwise, i wonder if it might be better to put this in contrib (and not add pygments to the core tvm dep list) or create a setup.py "extra" to distinguish between dependencies needed just to run inference and dependencies needed for work on the compiler.

[areusch]

cc @yelite @junrushao1994 who might have ideas whether the printer should handle this

[ganler]

Thanks @areusch for the suggestions!

I agree that we want to avoid dependencies if possible. As you said, one way is to implement the functionality from scratch.

I wonder if it would be possible to instead modify the TVMScript printer?

It is possible with some effort to build a highlighting functionality from scratch: 1) lexer (*); 2) generating ansi escape code; and 3) make things configurable.

*With a lexer, we can decouple things as directly modifying operator<< in printer makes printer implementation complicated: when extending printing for some type, we hope that the developer only needs to focus on the content without the fear of what kind of colors to use. It is also hard to directly modifying operator<< as the printer is a bit coarse-grained (statement level) that it can contain multiple token types for one statement type:

tvm/src/printer/tvmscript_printer.cc

Line 883 in 4a94a94

doc << tir_prefix_ << ".let(" << Print(op->var) << ", " << Print(op->value) << ", "

in general it'd be good to avoid adding dependencies to the TVM core unless they're clearly beneficial to a large set of code paths.

Agreed. I think we should not put this under "core". Meanwhile, "Pygments" is actually a very fundamental and reliable library in Python community that has been actively developed for years. For example, latest pypi can use rich (which depends on "Pygments") to display the progress bars.

Summing up the bits above, I suggest that we still use Pygments to hightlight TVM script but we don't put it as a required dependency. Instead we can put some log to suggest users to install "Pygments" if they have not done it so far or simply put it in extras_require.

[yelite]

I think it's reasonable to let 3rd party library to handle the highlighting. It will take non-trivial effort to implement highlight in the TVMScript printer by ourselves. This also applies to other nice-to-have features like code formatting. Users might want the printed TVMScript to be free from very long lines and we can use black to format the output.

IMO all those extra dependencies should be optional and the show function should just fallback to plain printing if those dependencies are not installed and then emit a warning message.

[ganler]

Some updates:

Drop "core" dep, fallback as plain text and warn thanks to @yidawang and @areusch

If Pygments is not installed, we print plain text and post a warning:

ANSI colors for terminal users thanks to @tkonolige

If the environment is not notebook (i.e., terminal or something), we use a style filled with ansi colors (the colors are tuned to close to the default theme of notebook).

Dark (iterm2):

Light (macOS terminal):

[junrushao]

LGTM