tqchen commented on a change in pull request #7630:
URL: https://github.com/apache/tvm/pull/7630#discussion_r598027818
##########
File path: python/tvm/script/context_maintainer.py
##########
@@ -16,59 +16,179 @@
# under the License.
"""TVM Script Context Maintainer for TIR"""
-from tvm.te import schedule
+from typing import List, Mapping, Union, Optional, Dict, Callable
+import synr
+
+
+import tvm
+from tvm.ir import Span
+from tvm.tir import Var, Buffer, PrimExpr, Stmt, MatchBufferRegion
+from tvm.runtime import Object
+from .node import BufferSlice
+
+
+class BlockInfo:
+ """Information for block and block_realize signature"""
+
+ alloc_buffers: List[Buffer] = []
+ """List[Buffer]: list of tir.alloc_buffer statements in the block
signature"""
+ match_buffers: List[MatchBufferRegion] = []
+ """List[MatchBufferRegion]: list of tir.match_buffer_region statements in
the block signature"""
+ iter_bindings: Mapping[Var, PrimExpr] = {}
+ """Mapping[Var, PrimExpr]: map of block iter var to its values"""
+ reads: Optional[List[BufferSlice]] = None
+ """Optional[List[BufferSlice]]:
+ list of tir.reads statements in the block signature, None for
not-visited"""
+ writes: Optional[List[BufferSlice]] = None
+ """Optional[List[BufferSlice]]:
+ list of tir.writes statements in the block signature, None for
not-visited"""
+ annotations: Optional[Mapping[str, Object]] = None
+ """Optional[Mapping[str, Object]]:
+ list of tir.block_attr statements in the block signature, None for
not-visited"""
+ predicate: Optional[PrimExpr] = None
+ """Optional[PrimExpr]: block realize predicate, None for not-visited"""
+ init: Optional[Stmt] = None
+ """Optional[Stmt]: init part of the block, None for not-visited"""
+
+ def __init__(self):
+ self.alloc_buffers = []
+ self.match_buffers = []
+ self.iter_bindings = {}
+ self.reads = None
+ self.writes = None
+ self.annotations = None
+ self.predicate = None
+ self.init = None
class ContextMaintainer:
- """Maintain all the necessary context info"""
+ """Maintain all the necessary context info
+ Parameters
+ ----------
+ _report_error : Callable[[str, Union[Span, synr.ast.Span]], None]
+ The report error function handle
+ """
- def __init__(self, parser):
+ # scope context
+ node_stack: List[List[synr.ast.Node]] = []
+ """List[List[synr.ast.Node]]: The ast nodes insides the current scope"""
+ block_info_stack: List[BlockInfo] = []
+ """List[BlockInfo]: The block info for the current block scope"""
+ loop_stack: List[List[Var]] = []
+ """List[List[Var]]: List of loop vars inside the current block scope"""
+ symbols: List[Dict[str, Union[Var, Buffer]]] = []
+ """List[Dict[str, Union[Var, Buffer]]]: Symbol map from name to object for
the current scope"""
+
+ # function context
+ func_params: List[Var] = []
+ """List[Var]: The function parameters"""
+ func_buffer_map: Mapping[Var, Buffer] = {}
+ """Mapping[Var, Buffer]: The function buffer map"""
+ func_dict_attr: Mapping[str, Object] = {}
+ """Mapping[str, Object]: The function attrs"""
+ func_var_env_dict: Mapping[Var, str] = {}
+ """Mapping[Var, str]: The map from var to env thread"""
+
+ # parser and analyzer
+ analyzer: tvm.arith.Analyzer = tvm.arith.Analyzer()
+ """tvm.arith.Analyzer: The analyzer for simplifying"""
+ _report_error: Callable[[str, Union[Span, synr.ast.Span]], None]
+ """Callable[[str, Union[Span, synr.ast.Span]], None]: The report error
function handle"""
+
+ def __init__(self, _report_error: Callable[[str, Union[Span,
synr.ast.Span]], None]):
# scope context
- self.node_stack = [] # AST nodes of scopes
- self.symbols = [] # symbols of scopes
+ self.node_stack = []
+ self.block_info_stack = []
+ self.loop_stack = []
+ self.symbols = []
# function context
- self.func_params = [] # parameter list of function
- self.func_buffer_map = {} # buffer_map of function
- self.func_dict_attr = {} # func_attr of function
- self.func_var_env_dict = {} # map from var to env_name
- # parser
- self.parser = parser
-
- def pop_scope(self):
- """Pop the inner most scope"""
- self.symbols.pop()
- self.node_stack.pop()
+ self.func_params = []
+ self.func_buffer_map = {}
+ self.func_dict_attr = {}
+ self.func_var_env_dict = {}
+ # parser and analyzer
+ self._report_error = _report_error
+ self.analyzer = tvm.arith.Analyzer()
- def new_scope(self, nodes=None):
- """Creating a new scope"""
+ def enter_scope(self, nodes: Optional[List[synr.ast.Node]] = None):
+ """Creates a new scope
+
+ Note
+ ----
+ This function is used for normal scopes that do not involve
+ a `with block` scope. Use `enter_block_scope`
+ for block scope cases.
+
+ Parameters
+ ----------
+ nodes : Optional[List[synr.ast.Node]]
+ The synr AST nodes in new scope
+ """
if nodes is None:
nodes = []
self.node_stack.append(list(reversed(nodes)))
self.symbols.append(dict())
- def update_symbol(self, name, symbol):
+ def enter_block_scope(self, nodes: Optional[List[synr.ast.Node]] = None):
+ """Creates a new block scope, the function will call `enter_scope`
implicitly
+ Besides the behaviors of `enter_scope`, it will update loop_stack and
block_info_stack
+ to maintain block info.
+
+ Note
+ ----
+ This function should be used to handle a block scope,
+ aka the blocks that involve a `with block` scope.
Review comment:
Thanks @electriclilies. I agree with all you said in particular wrt to
code readability. We already followed the principle of enforcing heavy
documentation in the case of user facing code and making sure the overall logic
flows well.
Code readability also goes beyond the comments, a lot of efforts needs to be
spent on API naming, intuitive callings and error handling. This PR does a lot
of that, for example:
- E0: Clear API naming: see convert_to_int function in the diff
- E1: Good error handling: call_with_error_reporting that allows accurate
error reporting in most of the cases.
- E2: Type annotations: adding type annotations to a lot of the places(while
our current convention does not require so) to increase readability
- E3: Documentation of most functions
There is of course a tradeoff between the time we spend and amount of
comments to be added and other efforts. On one hand we certainly want to add as
many comments as possible. On the other hand, adding every code blocks may not
be the best way of investing time -- we could spend more time on overall
architectural correctness, the scaffolds(APIs and components) and other
elements that makes the code more readable and maintainable.
Comming back to a related example(e.g. reviewing the quantization code). It
is certainly helpful to add examples about network patterns happened during the
quantization process, values being involved and so on. But that may not be the
most important thing for now, since we can focus on more important issues on
readability and maintainability -- e.g. clarifying the key APIs, make sure they
compose well and so on.
Right now we are prioritizing to add the examples to developer code paths
that are more sutble, like the arithmetic modules and so on. In this
particular case, an example can be suggested and checked in as followup PRs is
more than welcomed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
