Skip to content

Commit 0176f43

Browse files
tqchentqchen
authored
[Script] Make transpiler call helper keywords explicit (#20448)
The transpiler call helpers currently mix builder arguments and source-location metadata in variadic keyword dictionaries. Make the builder keyword mapping optional and explicit, with dedicated keyword-only parameters for `span`, `name_span`, and `value_span`. Update callers to separate builder arguments from source ranges while preserving their evaluation order and leaving caller-owned mappings unchanged. This is a readability cleanup.
1 parent c607484 commit 0176f43

1 file changed

Lines changed: 50 additions & 20 deletions

File tree

‎python/tvm/script/parser/transpile.py‎

Lines changed: 50 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -287,33 +287,60 @@ def _raise_error(self, node: ast.AST, message: str) -> NoReturn:
287287
)
288288

289289
def _call(
290-
self, namespace: str, member: str, args: list[ast.expr], node: ast.AST, **keywords: ast.expr
290+
self,
291+
namespace: str,
292+
member: str,
293+
args: list[ast.expr],
294+
node: ast.AST,
295+
*,
296+
keywords: Mapping[str, ast.expr] | None = None,
297+
span: ast.expr | None = None,
298+
name_span: ast.expr | None = None,
299+
value_span: ast.expr | None = None,
291300
) -> ast.Call:
292301
"""Build a generated operation with its source range and named arguments."""
293-
if not self.module.track_span:
294-
keywords.pop("span", None)
295-
keywords.pop("name_span", None)
296-
keywords.pop("value_span", None)
302+
arguments = [ast.keyword(key, value) for key, value in keywords.items()] if keywords else []
303+
if self.module.track_span:
304+
if span is not None:
305+
arguments.append(ast.keyword("span", span))
306+
if name_span is not None:
307+
arguments.append(ast.keyword("name_span", name_span))
308+
if value_span is not None:
309+
arguments.append(ast.keyword("value_span", value_span))
297310
return ast.copy_location(
298311
ast.Call(
299312
ast.Attribute(ast.Name(namespace, ast.Load()), member, ast.Load()),
300313
args,
301-
[ast.keyword(key, value) for key, value in keywords.items()],
314+
arguments,
302315
),
303316
node,
304317
)
305318

306319
def _call_dialect(
307-
self, member: str, args: list[ast.expr], node: ast.AST, **keywords: ast.expr
320+
self,
321+
member: str,
322+
args: list[ast.expr],
323+
node: ast.AST,
324+
*,
325+
keywords: Mapping[str, ast.expr] | None = None,
326+
span: ast.expr | None = None,
327+
name_span: ast.expr | None = None,
328+
value_span: ast.expr | None = None,
308329
) -> ast.Call:
309-
return self._call(
330+
# Dialect calls evaluate their span, name span, builder arguments, then value span.
331+
call = self._call(
310332
self.function.dialect_prefix,
311333
member,
312334
args,
313335
node,
314-
span=self.module.make_span_expr(node),
315-
**keywords,
336+
span=self.module.make_span_expr(node) if span is None else span,
337+
name_span=name_span,
316338
)
339+
if keywords:
340+
call.keywords.extend(ast.keyword(key, value) for key, value in keywords.items())
341+
if self.module.track_span and value_span is not None:
342+
call.keywords.append(ast.keyword("value_span", value_span))
343+
return call
317344

318345
def _attach_span(self, value: ast.expr, node: ast.AST) -> ast.expr:
319346
if (
@@ -974,7 +1001,7 @@ def _bind(
9741001
[value],
9751002
statement,
9761003
name_span=self.module.make_span_expr(target),
977-
**keywords,
1004+
keywords=keywords,
9781005
)
9791006
elif kind == "module_alias" and not frame_value:
9801007
# -------------------- Pattern --------------------
@@ -1014,18 +1041,19 @@ def _bind(
10141041
# Builder:
10151042
# y = X.bind_(value, name="y", span=target_span, value_span=rhs_span)
10161043
# -------------------------------------------------
1017-
if value_span is not None:
1018-
keywords["value_span"] = value_span
10191044
value = (
10201045
self._call_dialect(
10211046
"bind_",
10221047
[value],
10231048
statement,
10241049
name_span=self.module.make_span_expr(target),
1025-
**keywords,
1050+
keywords=keywords,
1051+
value_span=value_span,
10261052
)
10271053
if frame_value
1028-
else self._call_dialect("bind_", [value], target, **keywords)
1054+
else self._call_dialect(
1055+
"bind_", [value], target, keywords=keywords, value_span=value_span
1056+
)
10291057
)
10301058
if frame_value:
10311059
# Binding an entered frame originates at the source as-target;
@@ -1136,9 +1164,11 @@ def visit_Assign(self, node: ast.Assign) -> ast.stmt | list[ast.stmt]:
11361164
"setitem_",
11371165
[],
11381166
node,
1139-
value=self._rewrite_assignment_value(node.value, ordinary=ordinary),
1140-
target=self.visit(target.value),
1141-
key=self._rewrite_index(target.slice),
1167+
keywords={
1168+
"value": self._rewrite_assignment_value(node.value, ordinary=ordinary),
1169+
"target": self.visit(target.value),
1170+
"key": self._rewrite_index(target.slice),
1171+
},
11421172
span=self.module.make_span_expr(node),
11431173
)
11441174
),
@@ -1464,7 +1494,7 @@ def visit_For(self, node: ast.For) -> ast.For | ast.With | list[ast.stmt]:
14641494
)
14651495
else:
14661496
self._raise_error(node.target, "Loop targets must be names or a flat tuple of names")
1467-
context = self._call_dialect("for_", [iterable], node, names=names)
1497+
context = self._call_dialect("for_", [iterable], node, keywords={"names": names})
14681498
# The generated iteration check originates at the source iterable, not
14691499
# the final body line. Its native frame span still covers the whole loop.
14701500
ast.copy_location(context, node.iter)
@@ -1838,7 +1868,7 @@ def _create_symbol_declarations(
18381868
"resolve_type_var_",
18391869
[ast.Constant(parameter.name)],
18401870
parameter,
1841-
dtype=ast.Constant(dtype),
1871+
keywords={"dtype": ast.Constant(dtype)},
18421872
),
18431873
parameter,
18441874
)

0 commit comments

Comments
 (0)