[GitHub] [tvm] junrushao1994 commented on a change in pull request #7630: [TensorIR] TVMScript Parser/Printer

[GitBox <gi...@apache.org>]

junrushao1994 commented on a change in pull request #7630:
URL: https://github.com/apache/tvm/pull/7630#discussion_r598086482



##########
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:
       I don't think any of us disagree with the extreme importance of comprehensive docs, but we need think a bit deeper what kind docs really help.
   
   I would like to further specifically reiterate my understanding of categorization of helpful docs:
   - D1. User-facing. This helps a user, who is possibly not familiar with the implementation, to understand how to use some functionalities of TVM;
   - D2. Design doc. This provides the high-level overview of how the codebase is structured and why it is designed like this;
   - D3. Private API doc. Document the private APIs that are not directly user-facing.
   
   D1 directly helps a user to better understand how to make things work and how it works. It is clear that D1 is desirable.
   
   D2 helps developers to understand the design philosophy, how the codebase is structured, etc, so that more people could help better maintain the codebase. Without D2, people are unable to understand the key concepts - that is why Tianqi redirects us to the RFC, so that we could all understand what a "Block" is, etc. Without D2, no matter how many words the developer use to document private APIs, it is still painful to understand some data structures.
   
   On D3, we always insist that APIs to be at least somewhat documented, so that maintainers could get a brief sense what will be going on if we call an API. Assuming we have good D1 and D2, we could substantially lower the steep learning curve and makes understanding D3 much easier - the prerequisite is that maintainers should read D1 and D2 first, in our specific case, the RFC and related materials.
   
   I think everybody totally agrees with Lily's words, and trust me, nobody wants bring trouble to future maintainers :-) With D1 and D2 ready in place (after M1s and M2s merged), it will be much easier to understand the design philosophy. This is what we are doing in Ansor upstreaming too - we upstream many tutorials after the codebase is fully functioning.
   
   It is totally understandable that reviewers are feeling frustrated when not understanding the design philosophy, and that is what RFCs are for, aren't they :-) Next time, we would love to see such frustration converts to clear questions and answers. For design philosophy-related questions, I would love to redirect everybody to the RFC from the very beginning, and we shouldn't debate outside the RFC about topics like "why two scopes".
   
   Then should we include the RFC text into the inlined docs? I think it is debatable, and I prefer not to replicate. The basic reason is that we will have D1 and D2 in the end, and maintaining two copies is quite error-prone. An easier solution for maintenance is to add links to D1 and D2 on critical data structures.




-- 
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:
users@infra.apache.org