[refactor] Remove the old AST builder from python frontend (#3527)

* Remove old code from kernel_impl.py

* remove old builder from transformer.py

* fix format

* fix pylint

* delete old builder

* format

* fix pylint

python/taichi/lang/__init__.py

@@ -396,7 +396,6 @@ def __init__(self):
         self.gdb_trigger = False
         self.excepthook = False
         self.experimental_real_function = False
-        self.experimental_ast_refactor = True
 
 
 def prepare_sandbox():
@@ -489,7 +488,6 @@ def init(arch=None,
     env_spec.add('gdb_trigger')
     env_spec.add('excepthook')
     env_spec.add('experimental_real_function')
-    env_spec.add('experimental_ast_refactor')
 
     # compiler configurations (ti.cfg):
     for key in dir(ti.cfg):
@@ -519,8 +517,6 @@ def init(arch=None,
         impl.get_runtime().print_preprocessed = spec_cfg.print_preprocessed
         impl.get_runtime().experimental_real_function = \
             spec_cfg.experimental_real_function
-        impl.get_runtime(
-        ).experimental_ast_refactor = spec_cfg.experimental_ast_refactor
         ti.set_logging_level(spec_cfg.log_level.lower())
         if spec_cfg.excepthook:
             # TODO(#1405): add a way to restore old excepthook

python/taichi/lang/ast/symbol_resolver.py

@@ -1,6 +1,5 @@
 """Provides helpers to resolve AST nodes."""
 import ast
-import inspect
 
 
 class ASTResolver:
@@ -52,22 +51,3 @@ def resolve_to(node, wanted, scope):
                 return False
         # The name ``scope`` here could be a bit confusing
         return scope is wanted
-
-
-class ModuleResolver:
-    @staticmethod
-    def get_module_name(function, wanted):
-        """Get the user defined module name of the given
-        ``wanted`` module object from the golbal variables of a function.
-
-        Args:
-            function (Callable): A function to get the global variables from.
-            wanted (Module): The wanted module object.
-        """
-        if not inspect.isfunction(function):
-            return None
-
-        for name, module in function.__globals__.items():
-            if module == wanted:
-                return name
-        return None

python/taichi/lang/ast/transformer.py

@@ -2,11 +2,10 @@
 
 import astor
 from taichi.lang import impl
-from taichi.lang.ast.symbol_resolver import ASTResolver, ModuleResolver
-from taichi.lang.ast_builder_utils import BuilderContext, IRBuilderContext
+from taichi.lang.ast.symbol_resolver import ASTResolver
+from taichi.lang.ast_builder_utils import IRBuilderContext
 from taichi.lang.exception import TaichiSyntaxError
 from taichi.lang.ir_builder import IRBuilder
-from taichi.lang.stmt_builder import build_stmt
 
 import taichi as ti
 
@@ -25,7 +24,6 @@ def __init__(self,
         self.arg_features = arg_features
         self.pass_checks = ASTTransformerChecks(func=func,
                                                 global_vars=global_vars)
-        self.rename_module = ASTTransformerUnifyModule(func=func)
         self.global_vars = global_vars
 
     @staticmethod
@@ -37,47 +35,33 @@ def print_ast(tree, title=None):
         print(astor.to_source(tree.body[0], indent_with='    '), flush=True)
 
     def visit(self, tree, *arguments):
-        if impl.get_runtime().experimental_ast_refactor:
-            self.print_ast(tree, 'Initial AST')
-            ctx = IRBuilderContext(
-                func=self.func,
-                excluded_parameters=self.excluded_parameters,
-                is_kernel=self.is_kernel,
-                arg_features=self.arg_features,
-                global_vars=self.global_vars,
-                argument_data=arguments)
-            # Convert Python AST to Python code that generates Taichi C++ AST.
-
-            tree = IRBuilder()(ctx, tree)
-            ast.fix_missing_locations(tree)
-            self.print_ast(tree, 'Preprocessed')
-            self.pass_checks.visit(tree)  # does not modify the AST
-            return ctx.return_data
         self.print_ast(tree, 'Initial AST')
-        self.rename_module.visit(tree)
-        self.print_ast(tree, 'AST with module renamed')
-        ctx = BuilderContext(func=self.func,
-                             excluded_parameters=self.excluded_parameters,
-                             is_kernel=self.is_kernel,
-                             arg_features=self.arg_features)
+        ctx = IRBuilderContext(func=self.func,
+                               excluded_parameters=self.excluded_parameters,
+                               is_kernel=self.is_kernel,
+                               arg_features=self.arg_features,
+                               global_vars=self.global_vars,
+                               argument_data=arguments)
         # Convert Python AST to Python code that generates Taichi C++ AST.
-        tree = build_stmt(ctx, tree)
+
+        tree = IRBuilder()(ctx, tree)
         ast.fix_missing_locations(tree)
         self.print_ast(tree, 'Preprocessed')
         self.pass_checks.visit(tree)  # does not modify the AST
-
-        return None
+        return ctx.return_data
 
 
-class ASTTransformerBase(ast.NodeTransformer):
-    def __init__(self, func):
+# Performs checks at the Python AST level. Does not modify the AST.
+class ASTTransformerChecks(ast.NodeTransformer):
+    def __init__(self, func, global_vars):
         super().__init__()
         self.func = func
+        self.has_return = False
+        self.in_static_if = False
+        self.globals = global_vars
 
     @staticmethod
     def get_decorator(global_vars, node):
-        if not impl.get_runtime().experimental_ast_refactor:
-            global_vars = globals()
         if not isinstance(node, ast.Call):
             return ''
         for wanted, name in [
@@ -89,34 +73,6 @@ def get_decorator(global_vars, node):
                 return name
         return ''
 
-
-class ASTTransformerUnifyModule(ast.NodeTransformer):
-    """Rename the module alias to `ti`.
-    module func calls like `<module alias>.<func-name>` will be renamed to
-    `ti.<func-name>`
-    """
-    def __init__(self, func):
-        super().__init__()
-        self.func = func
-        # Get the module alias from the global symbols table of the given func.
-        self.custom_module_name = ModuleResolver.get_module_name(func.func, ti)
-        self.default_module_name = "ti"
-
-    def visit_Name(self, node):
-        # Get the id of the ast.Name, rename if it equals to self.module_name.
-        if node.id == self.custom_module_name:
-            node.id = self.default_module_name
-        return node
-
-
-# Performs checks at the Python AST level. Does not modify the AST.
-class ASTTransformerChecks(ASTTransformerBase):
-    def __init__(self, func, global_vars):
-        super().__init__(func)
-        self.has_return = False
-        self.in_static_if = False
-        self.globals = global_vars
-
     def visit_If(self, node):
         node.test = self.visit(node.test)
 

python/taichi/lang/ast_builder_utils.py

@@ -1,4 +1,3 @@
-import ast
 from enum import Enum
 
 from taichi.lang.exception import TaichiSyntaxError
@@ -17,81 +16,6 @@ def __call__(self, ctx, node):
         return method(ctx, node)
 
 
-def parse_stmt(stmt):
-    return ast.parse(stmt).body[0]
-
-
-def parse_expr(expr):
-    return ast.parse(expr).body[0].value
-
-
-class ScopeGuard:
-    def __init__(self, scopes, stmt_block=None):
-        self.scopes = scopes
-        self.stmt_block = stmt_block
-
-    def __enter__(self):
-        self.scopes.append([])
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        local = self.scopes[-1]
-        if self.stmt_block is not None:
-            for var in reversed(local):
-                stmt = parse_stmt('del var')
-                stmt.targets[0].id = var
-                self.stmt_block.append(stmt)
-        self.scopes.pop()
-
-
-class BuilderContext:
-    def __init__(self,
-                 excluded_parameters=(),
-                 is_kernel=True,
-                 func=None,
-                 arg_features=None):
-        self.func = func
-        self.local_scopes = []
-        self.control_scopes = []
-        self.excluded_parameters = excluded_parameters
-        self.is_kernel = is_kernel
-        self.arg_features = arg_features
-        self.returns = None
-
-    # e.g.: FunctionDef, Module, Global
-    def variable_scope_guard(self, *args):
-        return ScopeGuard(self.local_scopes, *args)
-
-    # e.g.: For, While
-    def control_scope_guard(self):
-        return ScopeGuard(self.control_scopes)
-
-    def current_scope(self):
-        return self.local_scopes[-1]
-
-    def current_control_scope(self):
-        return self.control_scopes[-1]
-
-    def is_var_declared(self, name):
-        for s in self.local_scopes:
-            if name in s:
-                return True
-        return False
-
-    def is_creation(self, name):
-        return not self.is_var_declared(name)
-
-    def create_variable(self, name):
-        assert name not in self.current_scope(
-        ), "Recreating variables is not allowed"
-        self.current_scope().append(name)
-
-    def check_loop_var(self, loop_var):
-        if self.is_var_declared(loop_var):
-            raise TaichiSyntaxError(
-                f"Variable '{loop_var}' is already declared in the outer scope and cannot be used as loop variable"
-            )
-
-
 class VariableScopeGuard:
     def __init__(self, scopes, stmt_block=None):
         self.scopes = scopes