[dynamo] Optimize COMPARE_OP (pytorch#122039)

Improves `benchmarks/dynamo/microbenchmarks/dynamo_microbenchmarks.py`
from 5.6 to 5.1s.

Pull Request resolved: pytorch#122039
Approved by: https://github.com/Skylion007, https://github.com/anijain2305

Author: jansel

torch/_dynamo/symbolic_convert.py

@@ -100,8 +100,8 @@
 )
 from .variables.nn_module import NNModuleVariable
 from .variables.tensor import (
+    supported_comparison_ops,
     supported_const_comparison_ops,
-    supported_tensor_comparison_ops,
     SymNodeVariable,
     TensorVariable,
 )
@@ -880,8 +880,7 @@ def pop(self) -> VariableTracker:
         return self.stack.pop()
 
     def popn(self, n: int) -> List[VariableTracker]:
-        assert n >= 0
-        return list(reversed([self.pop() for _ in range(n)]))
+        return [*reversed([self.pop() for _ in range(n)])]
 
     def LOAD_FAST(self, inst):
         name = inst.argval
@@ -1203,59 +1202,47 @@ def FOR_ITER(self, inst):
     def COMPARE_OP(self, inst):
         left, right = self.popn(2)
         op = inst.argval
-        supported_any = dict(
-            itertools.chain(
-                supported_tensor_comparison_ops.items(),
-                supported_const_comparison_ops.items(),
-            )
-        )
-        if (
-            isinstance(
-                left,
-                (
-                    TensorVariable,
-                    SymNodeVariable,
-                    NNModuleVariable,
-                    BaseListVariable,
-                    UserDefinedVariable,
-                    BaseUserFunctionVariable,
-                    ConstDictVariable,
-                ),
-            )
-            and isinstance(right, ConstantVariable)
-            and right.value is None
-            and op in supported_const_comparison_ops
-        ):
-            # <non-None> is None
-            self.push(
-                ConstantVariable.create(
-                    supported_const_comparison_ops[op](object(), right.value)
-                )
-            )
-
-        elif (
-            left.is_python_constant()
-            and right.is_python_constant()
-            and op in supported_any
-        ):
-            # constant fold
-            self.push(
-                ConstantVariable.create(
-                    supported_any[op](
-                        left.as_python_constant(), right.as_python_constant()
-                    ),
-                )
-            )
-        elif op in ("in", "not in"):
+        if op == "in" or op == "not in":
             self.push(right.call_method(self, "__contains__", [left], {}))
             if op == "not in":
                 self.UNARY_NOT(inst)
-        else:
-            self.push(
-                BuiltinVariable(supported_any[op]).call_function(
-                    self, [left, right], {}
+            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], {}
             )
+        )
 
     def GET_ITER(self, inst):
         self.call_function(BuiltinVariable(iter), [self.pop()], {})

torch/_dynamo/variables/builtin.py

@@ -1526,8 +1526,8 @@ def _comparison(self, tx, left, right):
         )
         from .lists import SizeVariable
         from .tensor import (
-            supported_const_comparison_ops,
-            supported_tensor_comparison_ops,
+            supported_const_comparison_op_values,
+            supported_tensor_comparison_op_values,
         )
 
         op = self.fn
@@ -1540,7 +1540,7 @@ def _unimplemented():
                 isinstance(x, (NNModuleVariable, ConstantVariable))
                 for x in [left, right]
             )
-            and op in supported_const_comparison_ops.values()
+            and op in supported_const_comparison_op_values
         ):
             left = (
                 tx.output.get_submodule(left.module_key)
@@ -1555,7 +1555,7 @@ def _unimplemented():
             return ConstantVariable.create(op(left, right))
 
         if isinstance(left, UserFunctionVariable):
-            if op not in supported_const_comparison_ops.values():
+            if op not in supported_const_comparison_op_values:
                 _unimplemented()
             if not isinstance(right, UserFunctionVariable):
                 _unimplemented()
@@ -1594,7 +1594,7 @@ def _unimplemented():
                 else:
                     return ConstantVariable.create(not is_result)
 
-            if op not in supported_tensor_comparison_ops.values():
+            if op not in supported_tensor_comparison_op_values:
                 _unimplemented()
             if (
                 isinstance(left, TensorVariable)
@@ -1618,7 +1618,7 @@ def _unimplemented():
             )
 
         if isinstance(left, SymNodeVariable) or isinstance(right, SymNodeVariable):
-            if op not in supported_tensor_comparison_ops.values():
+            if op not in supported_tensor_comparison_op_values:
                 _unimplemented()
 
             proxy = tx.output.create_proxy(

torch/_dynamo/variables/tensor.py

@@ -57,6 +57,7 @@
 from .constant import ConstantVariable
 from .lists import SizeVariable
 
+# Ops that allow tensor <op> tensor
 supported_tensor_comparison_ops = {
     ">": operator.gt,
     "<": operator.lt,
@@ -65,12 +66,23 @@
     "==": operator.eq,
     "!=": operator.ne,
 }
+# Ops that allow tensor <op> None
 supported_const_comparison_ops = {
     "is": operator.is_,
     "is not": operator.is_not,
     "==": operator.eq,
     "!=": operator.ne,
 }
+supported_comparison_ops = {
+    **supported_tensor_comparison_ops,
+    **supported_const_comparison_ops,
+}
+supported_tensor_comparison_op_values = dict.fromkeys(
+    supported_tensor_comparison_ops.values()
+)
+supported_const_comparison_op_values = dict.fromkeys(
+    supported_const_comparison_ops.values()
+)
 
 
 class TensorVariable(VariableTracker):