translate bytecode ops

examples/simple_dynamo.py

@@ -42,7 +42,13 @@ def foo(a, b):
     d = a * b
     e = c + d
     f = e - a
-    return f
+    g = f > e
+    h = g < f
+    i = h <= g
+    j = i >= i
+    k = j == i
+    l = j != k
+    return l
 
 
 optimized_foo = paddlefx.optimize(my_compiler)(foo)
@@ -51,3 +57,15 @@ def foo(a, b):
 optimized_res = optimized_foo(in_a, in_b)
 
 np.testing.assert_equal(original_res.numpy(), optimized_res.numpy())
+
+# in_a = paddle.randn([1])
+# def foo(a):
+#     b = not a
+#     return b
+#
+# optimized_foo = paddlefx.optimize(my_compiler)(foo)
+
+# original_res = foo(in_a)
+# optimized_res = optimized_foo(in_a)
+
+# np.testing.assert_equal(original_res.numpy(), optimized_res.numpy())

src/paddlefx/translator.py

@@ -52,7 +52,40 @@ def convert_instruction(i: dis.Instruction):
     )
 
 
-class InstructionTranslatorBase:
+def _binary_constructor(op_name: str):
+    def _binary(self, inst: Instruction):
+        op = getattr(operator, op_name)
+        args = list(reversed([self.stack.pop() for _ in range(2)]))
+        res = self.output.create_node('call_function', op, args, {})
+        self.stack.append(res)
+
+    return _binary
+
+
+class InstructionTranslatorMeta(type):
+    def __new__(cls, *args, **kwargs):
+        # binary op
+        op_mapper = {
+            'add': 'BINARY_ADD',
+            'sub': 'BINARY_SUBTRACT',
+            'mul': 'BINARY_MULTIPLY',
+            'floordiv': 'BINARY_FLOOR_DIVIDE',
+            # NOTE: in fact, paddle doesn't support floor_divide
+            'truediv': 'BINARY_TRUE_DIVIDE',
+            'add': 'BINARY_ADD',
+            'add': 'BINARY_ADD',
+        }
+        inst = type.__new__(cls, *args, **kwargs)
+        for op_name, func_name in op_mapper.items():
+            func = _binary_constructor(op_name)
+            func = types.FunctionType(
+                func.__code__, globals(), None, None, func.__closure__
+            )
+            setattr(inst, func_name, func)
+        return inst
+
+
+class InstructionTranslatorBase(metaclass=InstructionTranslatorMeta):
     def __init__(
         self,
         instructions: list[Instruction],
@@ -103,35 +136,18 @@ def LOAD_FAST(self, inst: Instruction):
     def RETURN_VALUE(self, inst: Instruction):
         self.compile_subgraph()
 
-    def BINARY_ADD(self, inst: Instruction):
-        add = getattr(operator, 'add')
-        args = list(reversed([self.stack.pop() for _ in range(2)]))
-        res = self.output.create_node('call_function', add, args, {})
-        self.stack.append(res)
-
-    def BINARY_SUBTRACT(self, inst: Instruction):
-        sub = getattr(operator, 'sub')
-        args = list(reversed([self.stack.pop() for _ in range(2)]))
-        res = self.output.create_node('call_function', sub, args, {})
-        self.stack.append(res)
-
-    def BINARY_MULTIPLY(self, inst: Instruction):
-        mul = getattr(operator, 'mul')
-        args = list(reversed([self.stack.pop() for _ in range(2)]))
-        res = self.output.create_node('call_function', mul, args, {})
-        self.stack.append(res)
-
-    # NOTE: in fact, paddle doesn't support floor_divide
-    def BINARY_FLOOR_DIVIDE(self, inst: Instruction):
-        floordiv = getattr(operator, 'floordiv')
-        args = list(reversed([self.stack.pop() for _ in range(2)]))
-        res = self.output.create_node('call_function', floordiv, args, {})
-        self.stack.append(res)
-
-    def BINARY_TRUE_DIVIDE(self, inst: Instruction):
-        truediv = getattr(operator, 'truediv')
+    def COMPARE_OP(self, inst: Instruction):
+        op_mapper = {
+            '>': 'gt',
+            '<': 'lt',
+            '>=': 'ge',
+            '<=': 'le',
+            '==': 'eq',
+            '!=': 'ne',
+        }
+        op = getattr(operator, op_mapper[inst.argval])
         args = list(reversed([self.stack.pop() for _ in range(2)]))
-        res = self.output.create_node('call_function', truediv, args, {})
+        res = self.output.create_node('call_function', op, args, {})
         self.stack.append(res)