chore: import upstream snapshot with attribution

python/tvm/script/parser/core/dispatch.py

@@ -0,0 +1,157 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Parser dispatching infrastructure"""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from .doc import AST
+
+if TYPE_CHECKING:
+    from .parser import Parser
+
+
+ParseMethod = Callable[["Parser", AST], None]
+ParseVTable: dict[tuple[str, str], ParseMethod] = {}
+
+OpMethod = Callable[..., Any]
+OpVTable: dict[tuple[type, AST, int], OpMethod] = {}
+
+
+def register(token: str, type_name: str):
+    """Register a method for a dispatch token and type name.
+
+    Parameters
+    ----------
+    token : str
+        The token for IR, e.g., T for TIR and R for Relax.
+
+    type_name : str
+        The type name of AST node, e.g., FunctionDef, With, For.
+
+    Returns
+    -------
+    func : callable
+        The function to register dispatched method of parsing
+        corresponding token and AST node type.
+    """
+
+    def func(method: ParseMethod):
+        """Register a method in parser virtual table.
+
+        Parameters
+        ----------
+        method : ParseMethod
+            The dispatched method to be registered in parser virtual table.
+        """
+        ParseVTable[(token, type_name)] = method
+
+    return func
+
+
+def get(
+    token: str,
+    type_name: str,
+    default: ParseMethod | None = None,
+) -> ParseMethod | None:
+    """Get a registered method for a dispatch token and type name,
+    or return a default method if no registered methods with this dispatch token and type name.
+
+    Parameters
+    ----------
+    token : str
+        The token for IR, e.g., T for TIR and R for Relax.
+
+    type_name : str
+        The type name of AST node, e.g., FunctionDef, With, For.
+
+    default : Optional[ParseMethod]
+        The default method when no registered methods with this dispatch token and type name.
+
+    Returns
+    -------
+    func : Optional[ParseMethod]
+        The dispatched method of parsing corresponding token and AST node type.
+    """
+    return ParseVTable.get((token, type_name), default)
+
+
+def register_op(operand_type: type, op_node_type: AST, operand_index: int):
+    """Register a method for a operand type, AST operator node and operand index.
+
+    Parameters
+    ----------
+    operand_type : Type
+        The type of operands, e.g., tirx.Expr, tirx.IterVar.
+
+    op_node_type : AST
+        The doc AST operator node type, e.g., doc.Add, doc.Eq.
+
+    operand_index : int
+        The operand index, i.e., 0 for left operand and 1 for right operand.
+
+    Returns
+    -------
+    func : callable
+        The function to register dispatched method of parsing
+        corresponding a operand type, AST operator node and operand index.
+    """
+
+    def func(method: OpMethod):
+        """Register a method in parser operator virtual table.
+
+        Parameters
+        ----------
+        method : ParseMethod
+            The dispatched method to be registered in parser operator virtual table.
+        """
+        OpVTable[(operand_type, op_node_type, operand_index)] = method
+
+    return func
+
+
+def get_op(
+    operand_type: type,
+    op_node_type: type,
+    operand_index: int,
+    default: OpMethod | None = None,
+) -> OpMethod | None:
+    """Register a method for a operand type, AST operator node and operand index.
+
+    Parameters
+    ----------
+    operand_type : Type
+        The type of operands, e.g., tirx.Expr, tirx.IterVar.
+
+    op_node_type : AST
+        The doc AST operator node type, e.g., doc.Add, doc.Eq.
+
+    operand_index : int
+        The operand index, i.e., 0 for left operand and 1 for right operand.
+
+
+    default : Optional[OpMethod]
+        The default method when no registered methods with this operand type,
+        AST operator node and operand index.
+
+    Returns
+    -------
+    func : Optional[OpMethod]
+        The function to register dispatched method of parsing
+        corresponding a operand type, AST operator node and operand index.
+    """
+    return OpVTable.get((operand_type, op_node_type, operand_index), default)