tqchen commented on code in PR #20004:
URL: https://github.com/apache/tvm/pull/20004#discussion_r3577334433
##########
python/tvm/relax/script/parser/parser.py:
##########
@@ -168,30 +213,25 @@ def is_recursive(node: doc.FunctionDef) -> bool:
def collect_symbolic_var_from_prelude(
- self: Parser, node: doc.FunctionDef, symbolic_vars: dict[str, tirx.Var]
-) -> dict[str, tirx.Var]:
+ self: Parser, node: doc.FunctionDef, symbolic_vars: dict[str, tvm.ir.Var]
+) -> dict[str, tvm.ir.Var]:
prelude_vars = {}
for stmt in node.body:
- if isinstance(stmt, doc.Assign) and all(
- isinstance(target, doc.Name) and target.id in symbolic_vars for
target in stmt.targets
- ):
- values = self.eval_expr(stmt.value)
-
- try:
- iter(values)
- except TypeError:
- values = [values]
-
- assert len(stmt.targets) == len(values)
- for target, value in zip(stmt.targets, values):
- name = target.id
- prelude_vars[name] = value
+ if isinstance(stmt, doc.Assign) and len(stmt.targets) == 1:
+ declarations =
collect_symbolic_var_declaration_nodes(stmt.targets[0], stmt.value)
+ for name, value_node in declarations.items():
+ if name not in symbolic_vars:
+ continue
+ declared_var = self.eval_expr(value_node)
+ if tvm.ir.is_prim_var(declared_var) and declared_var.ty ==
symbolic_vars[name].ty:
Review Comment:
Thanks for flagging this. I verified the behavior directly: this Python
expression is not pointer equality. `tvm.ir.Type.__eq__` calls
`tvm_ffi.structural_equal`, and `PrimType.__eq__` delegates to it for
non-string operands. As a stronger check, I constructed two `int64` Var types
on separate OS threads, bypassing the thread-local PrimType cache:
`lhs.ty.same_as(rhs.ty)` was `False`, while `lhs.ty == rhs.ty` and
`structural_equal(lhs.ty, rhs.ty)` were both `True`; `int64 == int32` was
`False`. Both runtime operands were `tvm.ir.PrimType`. Therefore the existing
comparison already has the intended structural semantics; replacing it with an
explicit `structural_equal` call would be behavior-equivalent, so I kept the
parser unchanged. At exact updated head `b79d07391687`, the focused
parser/type/ExprFunctor gate passes 25/25.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
