Merge pull request #693 from pengcheng888/issue/679_2

issue/679 缓存c++的device和python的device的对应关系，以减少互相转换的耗时

python/infinicore/device.py

 from infinicore.lib import _infinicore
 
+_infinicore_2_python_dict = {}
+_python_2_infinicore_dict = {}
+
 
 class device:
     # Public attributes describing the device
@@ -31,8 +34,7 @@ class device:
     def __getattr__(self, name):
         # Lazily construct and cache an attribute.
         # such as, self._underlying .
-        _type, _index = device._to_infinicore_device(self.type, self.index)
-        setattr(self, name, _infinicore.Device(_type, _index))
+        setattr(self, name, device._to_infinicore_device(self.type, self.index))
         return getattr(self, name)
 
     def __repr__(self):
@@ -56,52 +58,14 @@ class device:
         return self.type == other.type and self.index == other.index
 
     @staticmethod
-    def _to_infinicore_device(type, index):
-        all_device_types = tuple(_infinicore.Device.Type.__members__.values())[:-1]
-        all_device_count = tuple(
-            _infinicore.get_device_count(device) for device in all_device_types
-        )
-
-        torch_devices = {
-            torch_type: {
-                infinicore_type: 0
-                for infinicore_type in all_device_types
-                if _TORCH_DEVICE_MAP[infinicore_type] == torch_type
-            }
-            for torch_type in _TORCH_DEVICE_MAP.values()
-        }
-
-        for i, count in enumerate(all_device_count):
-            infinicore_device_type = _infinicore.Device.Type(i)
-            torch_devices[_TORCH_DEVICE_MAP[infinicore_device_type]][
-                infinicore_device_type
-            ] += count
-
-        for infinicore_device_type, infinicore_device_count in torch_devices[
-            type
-        ].items():
-            for i in range(infinicore_device_count):
-                if index == 0:
-                    return infinicore_device_type, i
-
-                index -= 1
+    def _to_infinicore_device(type: str, index: int):
+        return _python_2_infinicore_dict[type][index]
 
     @staticmethod
-    def _from_infinicore_device(infinicore_device):
-        type = _TORCH_DEVICE_MAP[infinicore_device.type]
-
-        base_index = 0
-
-        for infinicore_type, torch_type in _TORCH_DEVICE_MAP.items():
-            if torch_type != type:
-                continue
-
-            if infinicore_type == infinicore_device.type:
-                break
-
-            base_index += _infinicore.get_device_count(infinicore_type)
-
-        return device(type, base_index + infinicore_device.index)
+    def _from_infinicore_device(infinicore_device: _infinicore.Device):
+        return _infinicore_2_python_dict[infinicore_device.type][
+            infinicore_device.index
+        ]
 
 
 _TORCH_DEVICE_MAP = {
@@ -116,3 +80,55 @@ _TORCH_DEVICE_MAP = {
     _infinicore.Device.Type.HYGON: "cuda",
     _infinicore.Device.Type.QY: "cuda",
 }
+
+
+def _initialize_device_relationship(all_device_types, all_device_count):
+    # python_device_type_set: {'cpu', 'musa', 'npu', 'cuda', 'mlu'}
+    python_device_type_set = set([type for type in _TORCH_DEVICE_MAP.values()])
+    python_device_current_index = {type: 0 for type in python_device_type_set}
+
+    infinicore_2_python_dict = {}
+    python_2_infinicore_dict = {}
+    for infinicore_device_type, infinicore_device_count in zip(
+        all_device_types, all_device_count
+    ):
+        if 0 == infinicore_device_count:
+            continue
+
+        # Found one device
+        for infinicore_device_index in range(infinicore_device_count):
+            # Create instantiation objects for C++ devices
+            infinicore_instance = _infinicore.Device(
+                infinicore_device_type, infinicore_device_index
+            )
+
+            # Create instantiation objects for python devices
+            python_device_type = _TORCH_DEVICE_MAP[infinicore_device_type]
+
+            python_device_index = python_device_current_index[python_device_type]
+            python_device_current_index[python_device_type] += 1
+
+            python_instance = device(python_device_type, python_device_index)
+
+            # Cache corresponding relationship
+            if infinicore_2_python_dict.get(infinicore_device_type) is not None:
+                infinicore_2_python_dict[infinicore_device_type].append(python_instance)
+            else:
+                infinicore_2_python_dict[infinicore_device_type] = [python_instance]
+
+            if python_2_infinicore_dict.get(python_device_type) is not None:
+                python_2_infinicore_dict[python_device_type].append(infinicore_instance)
+            else:
+                python_2_infinicore_dict[python_device_type] = [infinicore_instance]
+
+    return infinicore_2_python_dict, python_2_infinicore_dict
+
+
+_all_device_types = tuple(_infinicore.Device.Type.__members__.values())[:-1]
+_all_device_count = tuple(
+    _infinicore.get_device_count(device) for device in _all_device_types
+)
+
+_infinicore_2_python_dict, _python_2_infinicore_dict = _initialize_device_relationship(
+    _all_device_types, _all_device_count
+)

python/infinicore/tensor.py

@@ -14,7 +14,7 @@ from .utils import (
 
 
 class Tensor:
-    # Public attributes describing the device
+    # Public attributes describing the Tensor
     _underlying: _infinicore.Tensor
     _torch_ref: "torch.Tensor"  # noqa: F821
     shape: list[int]

test/infinicore/test.py

 import torch
+from infinicore.lib import _infinicore
 
 import infinicore
 
@@ -168,9 +169,106 @@ def test5_bf16():
     print("误差:", torch_tensor - torch_ans_result)
 
 
+def func6_initialize_device_relationship():
+    from infinicore.device import _initialize_device_relationship
+
+    all_device_types = [
+        _infinicore.Device.Type.CPU,  # 0  "cpu"
+        _infinicore.Device.Type.NVIDIA,  # 1  "cuda"
+        _infinicore.Device.Type.CAMBRICON,  # 2  "mlu"
+        _infinicore.Device.Type.ASCEND,  # 3  "npu"
+        _infinicore.Device.Type.METAX,  # 4  "cuda"
+        _infinicore.Device.Type.MOORE,  # 5  "musa"
+        _infinicore.Device.Type.ILUVATAR,  # 6  "cuda"
+        _infinicore.Device.Type.QY,  # 9  "cuda"
+        _infinicore.Device.Type.KUNLUN,  # 7  "cuda"
+        _infinicore.Device.Type.HYGON,  # 8  "cuda"
+    ]
+    if True:
+        print("\n ---------- 测试 CPU")
+        all_device_count = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+        infinicore_2_python_dict, python_2_infinicore_dict = (
+            _initialize_device_relationship(all_device_types, all_device_count)
+        )
+        print("infinicore_2_python_dict", infinicore_2_python_dict)
+        print("python_2_infinicore_dict: ", python_2_infinicore_dict)
+
+        print("\n ---------- 测试 CPU+NVIDIA")
+        all_device_count = [1, 2, 0, 0, 0, 0, 0, 0, 0, 0]
+        infinicore_2_python_dict, python_2_infinicore_dict = (
+            _initialize_device_relationship(all_device_types, all_device_count)
+        )
+        print("infinicore_2_python_dict", infinicore_2_python_dict)
+        print("python_2_infinicore_dict: ", python_2_infinicore_dict)
+
+        print("\n ---------- 测试 CPU+NVIDIA+HYGON")
+        all_device_count = [1, 2, 0, 0, 0, 0, 0, 0, 0, 2]
+        infinicore_2_python_dict, python_2_infinicore_dict = (
+            _initialize_device_relationship(all_device_types, all_device_count)
+        )
+        print("infinicore_2_python_dict", infinicore_2_python_dict)
+        print("python_2_infinicore_dict: ", python_2_infinicore_dict)
+
+        print("\n ---------- 测试 CPU+NVIDIA+HYGON+CAMBRICON+ASCEND")
+        all_device_count = [1, 2, 2, 2, 0, 0, 0, 0, 0, 2]
+        infinicore_2_python_dict, python_2_infinicore_dict = (
+            _initialize_device_relationship(all_device_types, all_device_count)
+        )
+        print("infinicore_2_python_dict", infinicore_2_python_dict)
+        print("python_2_infinicore_dict: ", python_2_infinicore_dict)
+
+    if True:
+        print("\n ---------- 算子测试 cpu")
+        x = torch.ones((2, 3), dtype=torch.float32, device="cpu")
+        x_infini = infinicore.from_torch(x)
+
+        y = torch.ones((2, 3), dtype=torch.float32, device="cpu")
+        y_infini = infinicore.from_torch(y)
+
+        z_infini = x_infini + y_infini
+        print(z_infini.device)
+        z_infini.debug()
+
+    if True:
+        print("\n ---------- 算子测试 cuda")
+        x = torch.ones((2, 3), dtype=torch.float32, device="cuda:0")
+        x_infini = infinicore.from_torch(x)
+
+        y = torch.ones((2, 3), dtype=torch.float32, device="cuda:0")
+        y_infini = infinicore.from_torch(y)
+
+        z_infini = x_infini + y_infini
+        # print(z_infini.device)
+        # z_infini.debug()
+
+    if False:
+        print("\n ---------- 算子测试 cuda", infinicore.device("cuda", 0)._underlying)
+        x = infinicore.empty(
+            (2, 3), dtype=infinicore.float32, device=infinicore.device("cuda")
+        )
+        y = infinicore.empty(
+            (2, 3), dtype=infinicore.float32, device=infinicore.device("cuda")
+        )
+        z = x + y
+        print(z.device)
+
+    if False:
+        print("\n ---------- 算子测试 MOORE")
+        x = torch.ones((2, 3), dtype=torch.float32, device="musa:1")
+        x_infini = infinicore.from_torch(x)
+
+        y = torch.ones((2, 3), dtype=torch.float32, device="musa:1")
+        y_infini = infinicore.from_torch(y)
+
+        z_infini = x_infini + y_infini
+        print(z_infini.device)
+        z_infini.debug()
+
+
 if __name__ == "__main__":
     test()
     test2()
     test3()
     test4_to()
     test5_bf16()
+    func6_initialize_device_relationship()