benchmark
cdc4ae03 - mimic tp_richcompare handling (#182759)

Commit
38 days ago
mimic tp_richcompare handling (#182759) Summary: Add comparison operator support to torch.compile by implementing CPython's do_richcompare / tp_richcompare dispatch on Dynamo's VariableTracker hierarchy. Each VT gets a richcompare_impl slot (analogous to tp_richcompare), and generic_richcompare implements the do_richcompare 4-step algorithm (subclass priority, forward, reflected, fallback) directly using these slots. COMPARE_OP routes through generic_richcompare; call_method("__eq__") calls richcompare_impl directly, matching CPython's distinction between a == b and a.__eq__(b). The base richcompare_impl unconditionally graph-breaks, forcing every VT to have an explicit override. Graph breaks inside user comparison methods propagate to COMPARE_OP (which has break_graph_if_unsupported) and run eagerly. UDOV's richcompare_impl disables nested graph breaks on the resolved funcvar so the InliningInstructionTranslator does not split mid-method. Other changes: - Unified allow_c_slot API replaces allow_c_hash for registering C extension type slots as safe to call at trace time. - FakeIdVariable generally unconditionally graph-break (compile-time-only values are unsound for comparison). The exception is if FakeIdVariable came from the same kind of call (e.g. we can compare 2 FakeIdVariables that both come from `id()`, or both come from `hash()`). - TorchScriptObjectVariable gets identity-based richcompare_impl. - StreamVariable comparison matches CPython (never returns NotImplemented). - Comparison polyfills (cmp_eq etc.) removed — generic_richcompare uses VT logic directly. The architecture is documented in object_protocol.py with CPython permalink references. Fixes since initial submission: 1. UNPACK_EX producing tuple instead of list (symbolic_convert.py) — Python's *rest in pattern matching and unpacking must produce a list. Changed TupleVariable(vals_list) to ListVariable(vals_list, mutation_type=ValueMutationNew()). Fixed test_patma_188, test_patma_189. 2. FunctoolsPartialVariable.args returning list instead of tuple (functions.py) — var_getattr("args") returned ListVariable but partial.args is a tuple. Changed to TupleVariable. Fixed TestPartialC.test_attributes. 3. FakeIdVariable comparison regression (constant.py, object_protocol.py) — FakeIdVariable.richcompare_impl unconditionally graph-broke. Added FakeValueKind enum (ID vs HASH) so same-kind eq/ne comparisons resolve at compile time. The hash creation site passes FakeValueKind.HASH. Fixed test_userlist.test_init. 4. GetAttrVariable.richcompare_impl couldn't resolve attrs on constants (misc.py) — When var_getattr returned an unresolvable GetAttrVariable, we graph-broke even if the object was a python constant. Added fallback using VariableTracker.build(tx, getattr(obj, name)). Initially used ConstantVariable.create which broke on non-literal types like modules; changed to VariableTracker.build. Fixed test_descrdoc, test_module_parameter. 5. TracebackVariable identity comparison (object_protocol.py) — Two different TracebackVariable instances couldn't be identity-compared. Added TracebackVariable to the "objects created during tracing" list in vt_identity_compare. Fixed test_accepts_traceback. 6. UserDefinedDictVariable skipping MRO walk for Python-level comparison overrides (user_defined.py) — richcompare_impl delegated directly to ConstDictVariable, missing Python-level overrides like Counter.__lt__. Added check for Python methods in the subclass MRO before falling through to base dict. Fixed TestCounter.test_lt. 7. MethodWrapperVariable missing richcompare_impl (functions.py) — HEAD moved ConstantMethodWrapperVariable to functions.py as MethodWrapperVariable but our richcompare_impl was on the old class. Added it to the new class during merge conflict resolution. X-link: https://github.com/pytorch/pytorch/pull/182759 Approved by: https://github.com/anijain2305, https://github.com/guilhermeleobas Reviewed By: wdvr Differential Revision: D106241883 fbshipit-source-id: 7b8455e94cbf313c0d895eb61c37d526391d4364
Author
Committer
Parents
Loading