This is an automated email from the ASF dual-hosted git repository.
lunderberg pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d5fab9e4fb [TVMScript] Use environment variable TVM_BLACK_FORMAT for
.show() (#15762)
d5fab9e4fb is described below
commit d5fab9e4fb3b6f06a07b7c4a887d1f7db9ed19a1
Author: Eric Lunderberg <[email protected]>
AuthorDate: Tue Sep 26 15:06:53 2023 -0500
[TVMScript] Use environment variable TVM_BLACK_FORMAT for .show() (#15762)
Prior to this commit, the default behavior of the `black_format`
argument in TVMScript printing has changed back and forth, based on
conflicting user preferences. This commit allows the default to be
specified by each using using the `TVM_BLACK_FORMAT` environment
variable. If unspecified in a `obj.show()` method call, this
environment variable is used to determine the default.
---
python/tvm/runtime/script_printer.py | 30 +++++++++++++++++++++++++++---
python/tvm/tir/schedule/schedule.py | 25 ++++++++++++-------------
python/tvm/tir/schedule/trace.py | 9 ++++++++-
3 files changed, 47 insertions(+), 17 deletions(-)
diff --git a/python/tvm/runtime/script_printer.py
b/python/tvm/runtime/script_printer.py
index e7846c0680..dff5429111 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/runtime/script_printer.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Configuration of TVMScript printer"""
+import os
from typing import Dict, List, Optional, Sequence
from tvm._ffi import get_global_func, register_object
@@ -199,7 +200,7 @@ class Scriptable:
def show(
self,
style: Optional[str] = None,
- black_format: bool = False,
+ black_format: Optional[bool] = None,
*,
name: Optional[str] = None,
show_meta: bool = False,
@@ -226,8 +227,26 @@ class Scriptable:
style : str, optional
Pygmentize printing style, auto-detected if None. See
`tvm.script.highlight.cprint` for more details.
- black_format: bool
- If true, use the formatter Black to format the TVMScript
+
+ black_format: Optional[bool]
+
+ If true, use the formatter Black to format the TVMScript.
+ If false, do not apply the auto-formatter.
+
+ If None (default), determine the behavior based on the
+ environment variable "TVM_BLACK_FORMAT". If this
+ environment variable is unset, set to the empty string, or
+ set to the integer zero, black auto-formatting will be
+ disabled. If the environment variable is set to a
+ non-zero integer, black auto-formatting will be enabled.
+
+ Note that the "TVM_BLACK_FORMAT" environment variable only
+ applies to the `.show()` method, and not the underlying
+ `.script()` method. The `.show()` method is intended for
+ human-readable output based on individual user
+ preferences, while the `.script()` method is intended to
+ provided a consistent output regardless of environment.
+
name : Optional[str] = None
The name of the object
show_meta : bool = False
@@ -263,11 +282,16 @@ class Scriptable:
Object to be underlined
obj_to_annotate : Optional[Dict[Object, str]] = None
Object to be annotated
+
"""
from tvm.script.highlight import ( # pylint:
disable=import-outside-toplevel
cprint,
)
+ if black_format is None:
+ env = os.environ.get("TVM_BLACK_FORMAT")
+ black_format = env and int(env)
+
cprint(
self.script(
name=name,
diff --git a/python/tvm/tir/schedule/schedule.py
b/python/tvm/tir/schedule/schedule.py
index 85d1f19bba..5df3a486ca 100644
--- a/python/tvm/tir/schedule/schedule.py
+++ b/python/tvm/tir/schedule/schedule.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""The TensorIR schedule class"""
+import inspect
from typing import Callable, Dict, List, Optional, Tuple, Union
from tvm._ffi import register_object as _register_object
@@ -268,26 +269,24 @@ class Schedule(Object):
"""
return _ffi_api.ScheduleForkSeed(self) # type: ignore # pylint:
disable=no-member
- def show(self, style: Optional[str] = None, black_format: bool = False) ->
None:
+ def show(self, *args, **kwargs) -> None:
"""A sugar for print highlighted TVM script.
- Parameters
- ----------
- style : str, optional
-
- Pygmentize printing style, auto-detected if None. See
- `tvm.script.highlight.cprint` for more details.
-
- black_format: bool
-
- If true, use the formatter Black to format the TVMScript
+ All parameters are forwarded to the underlying `Module.show`
+ and `Trace.show` methods.
"""
mod = self.mod
if mod is not None:
- mod.show(style=style, black_format=black_format)
+ mod.show(*args, **kwargs)
+
trace = self.trace
if trace is not None:
- trace.show(style=style, black_format=black_format)
+ # Trace.show only supports the style and black_format arguments
+ param_binding = inspect.signature(mod.show).bind(*args, **kwargs)
+ param_binding.apply_defaults()
+ bound_args = param_binding.arguments
+
+ trace.show(style=bound_args["style"],
black_format=bound_args["black_format"])
########## Lookup ##########
diff --git a/python/tvm/tir/schedule/trace.py b/python/tvm/tir/schedule/trace.py
index e317304993..cb8d5ce997 100644
--- a/python/tvm/tir/schedule/trace.py
+++ b/python/tvm/tir/schedule/trace.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""An execution trace of a scheduling program"""
+import os
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
from tvm._ffi import register_object as _register_object
@@ -274,10 +275,16 @@ class Trace(Object):
black_format: bool
- If true, use the formatter Black to format the TVMScript
+ If true, use the formatter Black to format the TVMScript.
+ If None, determine based on the "TVM_BLACK_FORMAT" environment
+ variable.
"""
from tvm.script.highlight import ( # pylint:
disable=import-outside-toplevel
cprint,
)
+ if black_format is None:
+ env = os.environ.get("TVM_BLACK_FORMAT")
+ black_format = bool(env and int(env))
+
cprint(str(self), style=style, black_format=black_format)
