[dynamo] Refactor COMPARE_OP and comparison builtins (pytorch#122043)

This removes the duplicate handling of comparison ops between symbolic_convert and bultin and refactors the handling to use the binop infrastructure.  This change regresses overheads a bit, but this is fixed in the next PR.

New test skips are variants of `type(e) is np.ndarray` previously falling back to eager.

Pull Request resolved: pytorch#122043
Approved by: https://github.com/anijain2305
ghstack dependencies: pytorch#122039

Author: jansel

benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_inference.csv

@@ -86,7 +86,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0
 
 
 
-detectron2_fcos_r_50_fpn,pass,35
+detectron2_fcos_r_50_fpn,pass,94
 
 
 

benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv

@@ -54,47 +54,47 @@ densenet121,pass,0
 
 
 
-detectron2_fasterrcnn_r_101_c4,pass,51
+detectron2_fasterrcnn_r_101_c4,pass,164
 
 
 
-detectron2_fasterrcnn_r_101_dc5,pass,51
+detectron2_fasterrcnn_r_101_dc5,pass,163
 
 
 
-detectron2_fasterrcnn_r_101_fpn,pass,55
+detectron2_fasterrcnn_r_101_fpn,pass,172
 
 
 
-detectron2_fasterrcnn_r_50_c4,pass,51
+detectron2_fasterrcnn_r_50_c4,pass,113
 
 
 
-detectron2_fasterrcnn_r_50_dc5,pass,51
+detectron2_fasterrcnn_r_50_dc5,pass,112
 
 
 
-detectron2_fasterrcnn_r_50_fpn,pass,55
+detectron2_fasterrcnn_r_50_fpn,pass,121
 
 
 
-detectron2_fcos_r_50_fpn,pass,38
+detectron2_fcos_r_50_fpn,pass,97
 
 
 
-detectron2_maskrcnn_r_101_c4,fail_accuracy,66
+detectron2_maskrcnn_r_101_c4,pass,182
 
 
 
-detectron2_maskrcnn_r_101_fpn,pass,73
+detectron2_maskrcnn_r_101_fpn,pass,192
 
 
 
-detectron2_maskrcnn_r_50_c4,pass,66
+detectron2_maskrcnn_r_50_c4,pass,131
 
 
 
-detectron2_maskrcnn_r_50_fpn,pass,73
+detectron2_maskrcnn_r_50_fpn,pass,141
 
 
 

benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_inference.csv

@@ -86,7 +86,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0
 
 
 
-detectron2_fcos_r_50_fpn,pass,35
+detectron2_fcos_r_50_fpn,pass,94
 
 
 

benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_torchbench_inference.csv

@@ -54,7 +54,7 @@ densenet121,pass,0
 
 
 
-detectron2_fcos_r_50_fpn,pass,38
+detectron2_fcos_r_50_fpn,pass,97
 
 
 

benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv

@@ -86,7 +86,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0
 
 
 
-detectron2_fcos_r_50_fpn,pass,36
+detectron2_fcos_r_50_fpn,pass,95
 
 
 

benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_inference.csv

@@ -86,7 +86,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0
 
 
 
-detectron2_fcos_r_50_fpn,pass,35
+detectron2_fcos_r_50_fpn,pass,94
 
 
 

benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv

@@ -86,7 +86,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0
 
 
 
-detectron2_fcos_r_50_fpn,pass,36
+detectron2_fcos_r_50_fpn,pass,95
 
 
 

test/dynamo/test_functions.py

@@ -1459,6 +1459,26 @@ def test_partials_udf_arg(x):
         par_mul = functools.partial(udf_mul, torch.ones(10, 10))
         return par_mul(x)
 
+    @make_test
+    def test_list_add_then_mutate(x):
+        my_list = [1, x]
+        y = x / 4.0
+        my_list = my_list + [x / 2.0, 4]
+        my_list.append(y)
+        return sum(my_list)
+
+    @make_test
+    def test_list_expand_lhs(x):
+        return sum(4 * [x])
+
+    @make_test
+    def test_in_not_in(x):
+        mylist = [1, 2, 3, 4, 5, x]
+        myotherlist = [1, 2, 3, 4, 5]
+        assert 3 in mylist
+        assert 6 not in myotherlist
+        return sum(mylist)
+
     @make_test
     def test_partials_udf_kwarg(x):
         par_mul = functools.partial(udf_mul, y=torch.ones(10, 10))

test/dynamo_expected_failures/TestFX.test_pytree_concrete test/dynamo_skips/TestHistogramdd.test_bins_array

@@ -99,24 +99,29 @@
     UnknownVariable,
 )
 from .variables.nn_module import NNModuleVariable
-from .variables.tensor import (
-    supported_comparison_ops,
-    supported_const_comparison_ops,
-    SymNodeVariable,
-    TensorVariable,
-)
+from .variables.tensor import supported_comparison_ops, SymNodeVariable, TensorVariable
 from .variables.user_defined import (
     RemovableHandleVariable,
     UserDefinedClassVariable,
     UserDefinedObjectVariable,
-    UserDefinedVariable,
 )
 
 log = logging.getLogger(__name__)
 graph_break_log = torch._logging.getArtifactLogger(__name__, "graph_breaks")
 trace_call_log = torch._logging.getArtifactLogger(__name__, "trace_call")
 trace_source_log = torch._logging.getArtifactLogger(__name__, "trace_source")
 tls = threading.local()
+compare_op_handlers: Dict[str, Any] = {
+    k: BuiltinVariable(v).call_function for k, v in supported_comparison_ops.items()
+}
+handle_contains = BuiltinVariable(operator.contains).call_function
+handle_not = BuiltinVariable(operator.not_).call_function
+compare_op_handlers["in"] = lambda tx, args, _: handle_contains(
+    tx, [*reversed(args)], {}
+)
+compare_op_handlers["not in"] = lambda tx, args, _: handle_not(
+    tx, [handle_contains(tx, [*reversed(args)], {})], {}
+)
 
 
 @dataclasses.dataclass
@@ -1200,49 +1205,7 @@ def FOR_ITER(self, inst):
             unimplemented(f"FOR_ITER {typestr(it)}")
 
     def COMPARE_OP(self, inst):
-        left, right = self.popn(2)
-        op = inst.argval
-        if op == "in" or op == "not in":
-            self.push(right.call_method(self, "__contains__", [left], {}))
-            if op == "not in":
-                self.UNARY_NOT(inst)
-            return
-
-        if right.is_python_constant():
-            if left.is_python_constant():
-                # constant fold
-                return self.push(
-                    ConstantVariable(
-                        supported_comparison_ops[op](
-                            left.as_python_constant(), right.as_python_constant()
-                        ),
-                    )
-                )
-            elif (
-                op in supported_const_comparison_ops
-                and right.as_python_constant() is None
-                and isinstance(
-                    left,
-                    (
-                        TensorVariable,
-                        SymNodeVariable,
-                        NNModuleVariable,
-                        BaseListVariable,
-                        UserDefinedVariable,
-                        BaseUserFunctionVariable,
-                        ConstDictVariable,
-                    ),
-                )
-            ):
-                # <non-None> is None
-                return self.push(
-                    ConstantVariable(supported_const_comparison_ops[op](object(), None))
-                )
-        self.push(
-            BuiltinVariable(supported_comparison_ops[op]).call_function(
-                self, [left, right], {}
-            )
-        )
+        self.push(compare_op_handlers[inst.argval](self, self.popn(2), {}))
 
     def GET_ITER(self, inst):
         self.call_function(BuiltinVariable(iter), [self.pop()], {})