Function _visit_class_def

python/tvm/script/parser/ir/parser.py:36–83

The class definition visiting method for ir module. Parameters ---------- self : Parser The visiting parser. node : doc.ClassDef The doc AST class definition node.

(self: Parser, node: doc.ClassDef)

@dispatch.register(token="ir", type_name="ClassDef")
def _visit_class_def(self: Parser, node: doc.ClassDef) -> None:
    """The class definition visiting method for ir module.

    Parameters
    ----------
    self : Parser
        The visiting parser.

    node : doc.ClassDef
        The doc AST class definition node.
    """

    with self.var_table.with_frame():
        with I.ir_module():
            # Step 0. Add the class name to the var table
            fake_module = ModuleWithGlobalVars()
            self.var_table.add(node.name, fake_module)

            # Step 1: Check if this class inherits from BasePyModule
            is_base_py_module = _check_base_py_module_inheritance(node)
            if is_base_py_module:
                # Store this information in the IRModule for later use
                I.module_attrs({"base_py_module": True})
                # Set the parser context to allow Python functions
                self.set_class_context(node.name, True)
            else:
                # Set the parser context to disallow Python functions
                self.set_class_context(node.name, False)

            # Step 2. Visit non-function stmts, including but not limited to
            # 1. `I.module_attrs`
            # 2. `I.module_global_infos`
            with self.with_dispatch_token("ir"):
                for stmt in node.body:
                    if not isinstance(stmt, doc.FunctionDef):
                        self.visit(stmt)

            # Step 3. Visit function stmts to declare the global vars
            for stmt in node.body:
                if isinstance(stmt, doc.FunctionDef):
                    global_var = self.visit_tvm_declare_function(stmt)
                    fake_module.__setattr__(stmt.name, global_var)

            # Step 4. Visit and parse the functions
            with self.with_dispatch_token("ir"):
                for stmt in node.body:
                    if isinstance(stmt, doc.FunctionDef):
                        self.visit(stmt)


@dispatch.register(token="ir", type_name="Assign")