(tvm) branch main updated: [TVMScript] Fix error reporting inside Macro func (#16967)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 9cfebca136 [TVMScript] Fix error reporting inside Macro func (#16967)
9cfebca136 is described below

commit 9cfebca136a6dd58e59deeb19690d37cc6e9426a
Author: Siyuan Feng <hzfen...@sjtu.edu.cn>
AuthorDate: Sun May 5 21:51:53 2024 +0800

    [TVMScript] Fix error reporting inside Macro func (#16967)
---
 python/tvm/script/parser/core/parser.py | 53 +++++++++++++++++++++++----------
 1 file changed, 38 insertions(+), 15 deletions(-)

diff --git a/python/tvm/script/parser/core/parser.py 
b/python/tvm/script/parser/core/parser.py
index b41a05689d..0ecf669566 100644
--- a/python/tvm/script/parser/core/parser.py
+++ b/python/tvm/script/parser/core/parser.py
@@ -145,26 +145,27 @@ class ScriptMacro(abc.ABC):
         local_vars = param_binding.arguments
         parser = self._find_parser_def()
 
-        if self.hygienic:
-            saved_var_table = parser.var_table
-            parser.var_table = VarTable()
+        with parser.with_diag_source(self.source):
+            if self.hygienic:
+                saved_var_table = parser.var_table
+                parser.var_table = VarTable()
 
-            with parser.var_table.with_frame():
-                for k, v in self.closure_vars.items():
-                    parser.var_table.add(k, v)
-                for k, v in local_vars.items():
-                    parser.var_table.add(k, v)
+                with parser.var_table.with_frame():
+                    for k, v in self.closure_vars.items():
+                        parser.var_table.add(k, v)
+                    for k, v in local_vars.items():
+                        parser.var_table.add(k, v)
 
-                parse_result = self.parse_macro(parser)
+                    parse_result = self.parse_macro(parser)
 
-            parser.var_table = saved_var_table
+                parser.var_table = saved_var_table
 
-        else:
-            with parser.var_table.with_frame():
-                for k, v in local_vars.items():
-                    parser.var_table.add(k, v)
+            else:
+                with parser.var_table.with_frame():
+                    for k, v in local_vars.items():
+                        parser.var_table.add(k, v)
 
-                parse_result = self.parse_macro(parser)
+                    parse_result = self.parse_macro(parser)
 
         return parse_result
 
@@ -415,6 +416,28 @@ class Parser(doc.NodeVisitor):
 
         return _deferred(pop_token)
 
+    def with_diag_source(self, source: Source):
+        """Add a new source as with statement.
+
+        Parameters
+        ----------
+        source : Source
+            The source for diagnostics.
+
+        Returns
+        -------
+        res : Any
+            The context with new source.
+        """
+
+        last_diag = self.diag
+        self.diag = Diagnostics(source)
+
+        def pop_source():
+            self.diag = last_diag
+
+        return _deferred(pop_source)
+
     def eval_expr(
         self,
         node: Union[doc.Expression, doc.expr],