Remove formatter argument

Signed-off-by: Ganesan Ramalingam <[email protected]>

Author: gramalingam

onnxscript/_internal/analysis.py

@@ -58,7 +58,7 @@ def __init__(
         self._constant_if_condition: dict[ast.If, bool] = {}
         if globals:
             self._compute_constant_if_conditions(fun, globals)
-        self.do_liveness_analysis(fun, formatter)
+        self.do_liveness_analysis(fun)
 
     def _compute_constant_if_conditions(
         self, fun: ast.FunctionDef, globals: dict[str, Any]
@@ -70,7 +70,7 @@ def _compute_constant_if_conditions(
         conditions. The value of such conditions is determined from the outer-scope.
         """
 
-        assigned_vars = self.assigned_vars(fun.body, self._formatter)
+        assigned_vars = self.assigned_vars(fun.body)
         for node in ast.walk(fun):
             if isinstance(node, ast.If):
                 if isinstance(node.test, ast.Name):
@@ -90,17 +90,15 @@ def constant_if_condition(self, if_stmt: ast.If) -> Optional[bool]:
         """
         return self._constant_if_condition.get(if_stmt, None)  # type: ignore[return-value]
 
-    def assigned_vars(
-        self, stmt: ast.stmt | list[ast.stmt], formatter: sourceinfo.Formatter
-    ) -> Set[str]:
+    def assigned_vars(self, stmt: ast.stmt | list[ast.stmt]) -> Set[str]:
         """Return the set of all variables that may be assigned to in an execution of input stmt
         or sequence of statements.
         """
 
         def assigned_in_block(block: Sequence[ast.stmt]) -> Set[str]:
             result: set[Any] = set()
             for s in block:
-                result = result | self.assigned_vars(s, formatter)
+                result = result | self.assigned_vars(s)
             return result
 
         if isinstance(stmt, ast.Assign):
@@ -118,7 +116,7 @@ def assigned_in_block(block: Sequence[ast.stmt]) -> Set[str]:
             else:
                 return assigned_in_block(stmt.orelse)
         if isinstance(stmt, ast.For):
-            return assigned_in_block(stmt.body) | {_get_loop_var(stmt, formatter)}
+            return assigned_in_block(stmt.body) | {_get_loop_var(stmt, self._formatter)}
         if isinstance(stmt, ast.While):
             return assigned_in_block(stmt.body)
         if isinstance(stmt, list):
@@ -133,10 +131,10 @@ def assigned_in_block(block: Sequence[ast.stmt]) -> Set[str]:
             return set()
         if ast_utils.is_doc_string(stmt):
             return set()
-        error_message = formatter(stmt, f"Unsupported statement type {type(stmt)!r}.")
+        error_message = self._formatter(stmt, f"Unsupported statement type {type(stmt)!r}.")
         raise ValueError(error_message)
 
-    def do_liveness_analysis(self, fun: ast.FunctionDef, formatter: sourceinfo.Formatter):
+    def do_liveness_analysis(self, fun: ast.FunctionDef):
         """Perform liveness analysis of the given function-ast. The results of the
         analysis are stored directly with each statement-ast `s` as attributes `s.live_in`
         and `s.live_out`.
@@ -171,7 +169,7 @@ def visitBlock(block: Sequence[ast.stmt], live_out: Set[str]) -> Set[str]:
                 else:
                     return visitBlock(stmt.orelse, live_out)
             if isinstance(stmt, ast.For):
-                p_loop_var = _get_loop_var(stmt, formatter)
+                p_loop_var = _get_loop_var(stmt, self._formatter)
                 prev = None
                 curr = live_out
                 while curr != prev:
@@ -198,14 +196,16 @@ def visitBlock(block: Sequence[ast.stmt], live_out: Set[str]) -> Set[str]:
                 return live_out
             if ast_utils.is_print_call(stmt):
                 return live_out
-            raise ValueError(formatter(stmt, f"Unsupported statement type {type(stmt)!r}."))
+            raise ValueError(
+                self._formatter(stmt, f"Unsupported statement type {type(stmt)!r}.")
+            )
 
         assert isinstance(fun, ast.FunctionDef)
         live: set[Any] = set()
         for s in reversed(fun.body):
             live = visit(s, live)
 
-    def exposed_uses(self, stmts: Sequence[ast.stmt], formatter: sourceinfo.Formatter):
+    def exposed_uses(self, stmts: Sequence[ast.stmt]):
         """Return the set of variables that are used before being defined by given block.
         In essence, this identifies the "inputs" to a given code-block.
         For example, consider the following code-block:
@@ -251,7 +251,7 @@ def visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
             if isinstance(stmt, ast.For):
                 # Analysis assumes loop may execute zero times. Results can be improved
                 # for loops that execute at least once.
-                loop_var_set = {_get_loop_var(stmt, formatter)}
+                loop_var_set = {_get_loop_var(stmt, self._formatter)}
                 used_after_loop = live_out.difference(loop_var_set)
                 used_inside_loop = visitBlock(stmt.body, set()).difference(loop_var_set)
                 used_in_loop_header = _used_vars(stmt.iter)
@@ -269,13 +269,15 @@ def visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
             if isinstance(stmt, ast.FunctionDef):
                 if stmt.name in live_out:
                     live_out.remove(stmt.name)
-                    live_out = live_out | self.outer_scope_variables(stmt, formatter)
+                    live_out = live_out | self.outer_scope_variables(stmt)
                 return live_out
-            raise ValueError(formatter(stmt, f"Unsupported statement type {type(stmt)!r}."))
+            raise ValueError(
+                self._formatter(stmt, f"Unsupported statement type {type(stmt)!r}.")
+            )
 
         return visitBlock(stmts, set())
 
-    def outer_scope_variables(self, fun: ast.FunctionDef, formatter: sourceinfo.Formatter):
+    def outer_scope_variables(self, fun: ast.FunctionDef):
         """Return the set of outer-scope variables used in a nested function.
 
         Args:
@@ -286,6 +288,6 @@ def outer_scope_variables(self, fun: ast.FunctionDef, formatter: sourceinfo.Form
             A set of variable names (strings).
         """
         assert isinstance(fun, ast.FunctionDef)
-        used_vars_ = self.exposed_uses(fun.body, formatter)
+        used_vars_ = self.exposed_uses(fun.body)
         inputs = [x.arg for x in fun.args.args]
         return used_vars_.difference(inputs)

onnxscript/_internal/analysis_test.py

@@ -114,7 +114,7 @@ class TestExposedUses(unittest.TestCase):
     def assertUses(self, f, expected):
         source, parse_tree = ast_utils.get_src_and_ast(f)
         analyzer = analysis.AstAnalyzer(parse_tree, formatter(source))
-        result = analyzer.exposed_uses(parse_tree.body, formatter(source))
+        result = analyzer.exposed_uses(parse_tree.body)
         self.assertEqual(result, set(expected))
 
     def test_basic(self):
@@ -192,7 +192,7 @@ class TestAssignedVarAnalysis(unittest.TestCase):
     def assert_assigned_vars(self, f, expected: set[str]):
         source, parse_tree = ast_utils.get_src_and_ast(f)
         analyzer = analysis.AstAnalyzer(parse_tree, formatter(source))
-        result = analyzer.assigned_vars(parse_tree.body, formatter(source))
+        result = analyzer.assigned_vars(parse_tree.body)
         self.assertEqual(result, expected)
 
     def test_basic_defs(self):

onnxscript/converter.py

@@ -1108,11 +1108,9 @@ def _translate_if_stmt(self, stmt: ast.If) -> None:
                 self._translate_stmt(s)
             return
         if hasattr(stmt, "live_out"):
-            live_defs = list(
-                stmt.live_out.intersection(self.analyzer.assigned_vars(stmt, self._message))
-            )
+            live_defs = list(stmt.live_out.intersection(self.analyzer.assigned_vars(stmt)))
         else:
-            live_defs = list(self.analyzer.assigned_vars(stmt, self._message))
+            live_defs = list(self.analyzer.assigned_vars(stmt))
         test = self._translate_expr(stmt.test, "cond").name
         lineno = self._source_of(stmt).lineno
         thenGraph, sub_fct_then = self._translate_block(
@@ -1192,8 +1190,8 @@ def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
         else:
             self.fail(loop_stmt, f"Unexpected loop type {type(loop_stmt)!r}.")
         # analyze loop body
-        exposed_uses = self.analyzer.exposed_uses(loop_stmt.body, self._message)
-        vars_def_in_loop = self.analyzer.assigned_vars(loop_stmt.body, self._message)
+        exposed_uses = self.analyzer.exposed_uses(loop_stmt.body)
+        vars_def_in_loop = self.analyzer.assigned_vars(loop_stmt.body)
         loop_state_vars = vars_def_in_loop.intersection(exposed_uses | loop_stmt.live_out)
         scan_outputs = set()  # TODO
         outputs = list(loop_state_vars | scan_outputs)
@@ -1380,7 +1378,7 @@ def _translate_nested_function_def(self, fn: ast.FunctionDef) -> None:
         self._enter_scope(fn.name, fn)
         self._translate_function_def_common(fn)
         function_ir = self._exit_scope()
-        outer_scope_vars = self.analyzer.outer_scope_variables(fn, self._message)
+        outer_scope_vars = self.analyzer.outer_scope_variables(fn)
         function_ir.outer_scope_variables = [
             (var, self._lookup(var, self._source_of(fn))) for var in outer_scope_vars
         ]