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