[Dispatch] use dispatch replace getitem, setitem, delitem and bool

sot/opcode_translator/executor/opcode_executor.py

@@ -432,7 +432,11 @@ def BINARY_SUBSCR(self, instr):
         container = self.pop()
         assert isinstance(key, VariableBase)
         self._graph.add_global_guarded_variable(key)
-        self.push(container[key.value])
+        self.push(
+            BuiltinVariable(operator.getitem, self._graph, DanglingTracker())(
+                container, key.value
+            )
+        )
 
     # inplace operators
     # paddle variable do not have inplace operators. For example when call `y **= x`, will call var.__pow__

sot/opcode_translator/executor/variable_dispatch.py

@@ -90,12 +90,44 @@
     {},
     lambda var: var.bool(),
 )
+Dispatcher.register(
+    bool,
+    ("ConstantVariable",),
+    {},
+    lambda var: var.bool(),
+)
 Dispatcher.register(
     operator.truth,
     ("ContainerVariable",),
     {},
     lambda var: var.bool(),
 )
+Dispatcher.register(
+    operator.truth,
+    ("ConstantVariable",),
+    {},
+    lambda var: var.bool(),
+)
+
+# getitem
+Dispatcher.register(
+    operator.getitem,
+    (
+        "VariableBase | TensorVariable | ContainerVariable",
+        "int | str | TensorVariable | slice",
+    ),
+    {},
+    lambda var, key: var.getitem(key),
+)
+Dispatcher.register(
+    operator.getitem,
+    (
+        "VariableBase | TensorVariable | ContainerVariable",
+        "ConstantVariable | SliceVariable",
+    ),
+    {},
+    lambda var, key: var.getitem(key.get_value()),
+)
 
 
 # VariableBase

sot/opcode_translator/executor/variables/base.py

@@ -291,9 +291,10 @@ def __repr__(self):
     def __str__(self):
         return self.__repr__()
 
-    def __getitem__(self, item):
-        # TODO: Remove this function after we use builtin dispatcher instead
+    def __getitem__(self, idx):
+        return self.getitem(idx)
 
+    def getitem(self, item):
         class_var = VariableFactory.from_value(
             self.get_value().__class__,
             self.graph,
@@ -331,9 +332,6 @@ def __call__(self, *args, **kwargs):
         output = fn_var(*args, **kwargs)
         return output
 
-    def getitem(self, *args, **kwargs):
-        pass
-
     @VariableFactory.register_from_value()
     def from_value(
         value: Any,

sot/opcode_translator/executor/variables/basic.py

@@ -79,6 +79,11 @@ def main_info(self) -> dict[str, Any]:
     def __bool__(self) -> bool:
         return bool(self.value)
 
+    def bool(self):
+        return VariableFactory.from_value(
+            bool(self), self.graph, DummyTracker([self])
+        )
+
     @VariableFactory.register_from_value()
     def from_value(value: Any, graph: FunctionGraph | None, tracker: Tracker):
         if isinstance(value, ConstTypes):
@@ -167,7 +172,7 @@ def main_info(self) -> dict[str, Any]:
             "var_name": self.var_name,
         }
 
-    def __getitem__(self, key):
+    def getitem(self, key):
         return self.graph.call_tensor_method(
             '__getitem__',
             self,

sot/opcode_translator/executor/variables/container.py

@@ -34,7 +34,7 @@ def len(self):
             len(self), self.graph, DummyTracker([self])
         )
 
-    def __bool__(self):
+    def __bool__(self) -> bool:
         return len(self) > 0
 
     def bool(self):
@@ -78,7 +78,7 @@ def __init__(
         self.value = val_list
 
     def get_value(self):
-        return [self[i].get_value() for i in range(len(self))]
+        return [self[idx].get_value() for idx in range(len(self))]
 
     def _reconstruct(self, codegen: PyCodeGen):
         size = len(self)
@@ -102,7 +102,7 @@ def main_info(self) -> dict[str, Any]:
     def __len__(self):
         return len(self.value)
 
-    def __getitem__(self, key):
+    def getitem(self, key):
         '''
         we need to make sure that:
             before an inplace change happens to ListVariable,
@@ -188,17 +188,17 @@ def from_value(value: Any, graph: FunctionGraph | None, tracker: Tracker):
 class TupleVariable(ContainerVariable):
     def __init__(
         self,
-        val_tuple: list[VariableBase],
+        val_tuple: list[VariableBase] | tuple[VariableBase],
         graph: FunctionGraph,
         tracker: Tracker,
     ):
         super().__init__(tracker)
         self.graph = graph
         # exactly it is a list (need replace item with VariableBase)
-        self.value = list(val_tuple)
+        self.value = val_tuple
 
     def get_value(self):
-        return tuple(self[i].get_value() for i in range(len(self)))
+        return tuple(self[idx].get_value() for idx in range(len(self)))
 
     def _reconstruct(self, codegen: PyCodeGen):
         size = len(self)
@@ -222,7 +222,7 @@ def main_info(self) -> dict[str, Any]:
     def __len__(self):
         return len(self.value)
 
-    def __getitem__(self, key):
+    def getitem(self, key):
         if isinstance(key, VariableBase):
             raise InnerError(
                 f"[{self.__class__.__name__}]: recieved {key} as key."
@@ -312,7 +312,7 @@ def main_info(self) -> dict[str, Any]:
     def __len__(self):
         return len(self.value)
 
-    def __getitem__(self, key):
+    def getitem(self, key):
         if isinstance(key, VariableBase):
             raise InnerError(
                 f"[{self.__class__.__name__}]: recieved {key} as key."

tests/test_03_tuple.py

@@ -3,6 +3,7 @@
 # BINARY_SUBSCR
 
 
+import operator
 import unittest
 
 from test_case_base import TestCaseBase
@@ -15,10 +16,24 @@ def foo(x: int, y: paddle.Tensor):
     return x[1] + 1
 
 
+def foo1(x: int, y: paddle.Tensor):
+    x = (x, y, 1)
+    return operator.getitem(x, slice(0, 2))
+
+
+def foo2(x: int, y: paddle.Tensor):
+    z = (x, y, 3, 4)
+    return z[0:5:1]
+
+
 class TestExecutor(TestCaseBase):
     def test_simple(self):
         self.assert_results(foo, 1, paddle.to_tensor(2))
+        self.assert_results(foo1, 1, paddle.to_tensor(2))
+        self.assert_results(foo2, 1, paddle.to_tensor(2))
 
 
 if __name__ == "__main__":
     unittest.main()
+    # a = foo2(1, paddle.to_tensor(2))
+    # print(a)

tests/test_04_list.py

@@ -1,7 +1,9 @@
 # New Supported Instructions:
 # BUILD_LIST (new)
+# BINARY_SUBSCR
 
 
+import operator
 import unittest
 
 from test_case_base import TestCaseBase
@@ -14,9 +16,15 @@ def foo(x: int, y: paddle.Tensor):
     return x[1] + 1
 
 
+def list_getitem(x: int, y: paddle.Tensor):
+    z = [x, y]
+    return operator.getitem(z, 1) + 1
+
+
 class TestExecutor(TestCaseBase):
     def test_simple(self):
         self.assert_results(foo, 1, paddle.to_tensor(2))
+        self.assert_results(list_getitem, 1, paddle.to_tensor(2))
 
 
 if __name__ == "__main__":

tests/test_05_dict.py

@@ -2,6 +2,7 @@
 # BUILD_MAP (new)
 # BUILD_CONST_KEY_MAP (new)
 
+import operator
 import unittest
 
 from test_case_base import TestCaseBase
@@ -25,6 +26,11 @@ def dict_set_item(x: int, y: paddle.Tensor):
     return z[1]
 
 
+def dict_get_item(x: int, y: paddle.Tensor):
+    z = {1: y, 2: y + 1}
+    return operator.getitem(z, 1)
+
+
 class TestExecutor(TestCaseBase):
     def test_build_map(self):
         self.assert_results(build_map, 1, paddle.to_tensor(2))
@@ -35,6 +41,9 @@ def test_build_const_key_map(self):
     def test_dict_set_item(self):
         self.assert_results(dict_set_item, 1, paddle.to_tensor(2))
 
+    def test_dict_get_item(self):
+        self.assert_results(dict_get_item, 1, paddle.to_tensor(2))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_15_slice.py

@@ -9,22 +9,34 @@
 import paddle
 
 
-def build_slice(x: list, y: paddle.Tensor):
+def build_list_slice(x: list, y: paddle.Tensor):
     x[2:4] = [0, 1]
     return x[0] + y
 
 
-def build_slice_with_step(x: list, y: paddle.Tensor):
+def build_list_slice_with_step(x: list, y: paddle.Tensor):
     x[1:5:2] = [0, 1]
     return x[0] + y
 
 
+def build_tuple_slice(x: list, y: paddle.Tensor):
+    x[2:4] = (0, 1)
+    return x[0] + y
+
+
+def build_tuple_slice_with_step(x: list, y: paddle.Tensor):
+    x[1:5:2] = (0, 1)
+    return x[0] + y
+
+
 class TestExecutor(TestCaseBase):
     def test_simple(self):
         x = list(range(10))
         y = paddle.arange(10)
-        self.assert_results(build_slice, x, y)
-        self.assert_results(build_slice_with_step, x, y)
+        self.assert_results(build_list_slice, x, y)
+        self.assert_results(build_list_slice_with_step, x, y)
+        self.assert_results(build_tuple_slice, x, y)
+        self.assert_results(build_tuple_slice_with_step, x, y)
 
 
 if __name__ == "__main__":