chore: import upstream snapshot with attribution

test/auto_parallel/test_cluster_v2.py

@@ -0,0 +1,124 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+import unittest
+
+from paddle.distributed.auto_parallel.static.cluster_v2 import DeviceMesh
+from paddle.framework import core
+
+
+class TestDeviceMesh(unittest.TestCase):
+    def test_device_mesh(self):
+        name = "my_device_mesh"
+        mesh = [[0, 1, 2], [3, 4, 5]]
+        device_mesh = DeviceMesh(name, mesh, dim_names=["x", "y"])
+        device_mesh1 = DeviceMesh("another_mesh", [0, 1, 2, 3])
+        self.assertEqual(device_mesh.name, "my_device_mesh")
+        self.assertEqual(device_mesh.shape, [2, 3])
+        self.assertEqual(device_mesh.device_ids, [0, 1, 2, 3, 4, 5])
+        self.assertEqual(device_mesh.dim_names, ["x", "y"])
+        self.assertEqual(device_mesh.device_type, "UNKNOWN")
+        self.assertEqual(device_mesh.size, 6)
+        self.assertEqual(device_mesh.ndim, 2)
+        self.assertEqual(device_mesh.dim_size(0), 2)
+        self.assertEqual(device_mesh.dim_size(-1), 3)
+        self.assertEqual(device_mesh.dim_size("x"), 2)
+        self.assertEqual(device_mesh.dim_size("y"), 3)
+        self.assertEqual(device_mesh.empty(), False)
+        self.assertEqual(device_mesh.contains(0), True)
+        self.assertEqual(device_mesh.contains(6), False)
+
+        dev0 = core.Device(global_id=0, local_id=0, machine_id=0, type="GPU")
+        dev1 = core.Device(global_id=1, local_id=1, machine_id=0, type="GPU")
+        dev2 = core.Device(global_id=2, local_id=2, machine_id=0, type="GPU")
+        dev3 = core.Device(global_id=3, local_id=0, machine_id=1, type="GPU")
+        dev4 = core.Device(global_id=4, local_id=1, machine_id=1, type="GPU")
+        dev5 = core.Device(global_id=5, local_id=2, machine_id=1, type="GPU")
+        device_mesh.add_device(dev0)
+        device_mesh.add_device(dev1)
+        device_mesh.add_device(dev2)
+        device_mesh.add_device(dev3)
+        device_mesh.add_device(dev4)
+        device_mesh.add_device(dev5)
+        self.assertEqual(device_mesh.device(0), dev0)
+        self.assertEqual(device_mesh.device(1), dev1)
+        self.assertEqual(device_mesh.device(2), dev2)
+        self.assertEqual(device_mesh.device(3), dev3)
+        self.assertEqual(device_mesh.device(4), dev4)
+        self.assertEqual(device_mesh.device(5), dev5)
+
+        link0 = core.Link(source_id=0, target_id=1, type="NVL")
+        link1 = core.Link(source_id=1, target_id=0, type="NVL")
+        link2 = core.Link(source_id=3, target_id=4, type="NVL")
+        link3 = core.Link(source_id=4, target_id=5, type="NVL")
+        device_mesh.add_link(link0)
+        device_mesh.add_link(link1)
+        device_mesh.add_link(link2)
+        device_mesh.add_link(link3)
+        self.assertEqual(device_mesh.link(0, 1), link0)
+        self.assertEqual(device_mesh.link(1, 0), link1)
+        self.assertEqual(device_mesh.link(3, 4), link2)
+        self.assertEqual(device_mesh.link(4, 5), link3)
+
+        self.assertEqual(device_mesh.machine(0).id, 0)
+        self.assertEqual(device_mesh.machine(0).contains(3), False)
+        self.assertEqual(device_mesh.machine(0).device(2), dev2)
+        self.assertEqual(device_mesh.machine(1).link(3, 4), link2)
+        self.assertEqual(
+            device_mesh.machine(0).devices, device_mesh.machine(0).devices
+        )
+        self.assertEqual(
+            device_mesh.machine(0).links, device_mesh.machine(0).links
+        )
+
+        self.assertEqual(device_mesh.device_type, "GPU")
+        self.assertEqual(device_mesh.devices, device_mesh.devices)
+        self.assertEqual(device_mesh.links, device_mesh.links)
+        self.assertEqual(device_mesh.machines, device_mesh.machines)
+        self.assertEqual(device_mesh, device_mesh)
+        self.assertNotEqual(device_mesh, device_mesh1)
+        self.assertEqual(str(device_mesh), str(device_mesh))
+
+    def test_device(self):
+        device = core.Device(global_id=0, local_id=1, machine_id=2, type="GPU")
+        device.capability.sflops = 100
+        device.capability.dflops = 200
+        device.capability.memory = 32
+        device.capability.rate = 2
+        self.assertEqual(device.global_id, 0)
+        self.assertEqual(device.local_id, 1)
+        self.assertEqual(device.machine_id, 2)
+        self.assertEqual(device.type, "GPU")
+        self.assertAlmostEqual(device.capability.sflops, 100)
+        self.assertAlmostEqual(device.capability.dflops, 200)
+        self.assertAlmostEqual(device.capability.memory, 32)
+        self.assertAlmostEqual(device.capability.rate, 2)
+        self.assertEqual(device, device)
+        self.assertEqual(str(device), str(device))
+
+    def test_link(self):
+        link = core.Link(source_id=0, target_id=1, type="NVL")
+        link.capability.bandwidth = 100
+        link.capability.latency = 1
+        self.assertEqual(link.source_id, 0)
+        self.assertEqual(link.target_id, 1)
+        self.assertEqual(link.type, "NVL")
+        self.assertAlmostEqual(link.capability.bandwidth, 100)
+        self.assertAlmostEqual(link.capability.latency, 1)
+        self.assertEqual(link, link)
+        self.assertEqual(str(link), str(link))
+
+
+if __name__ == "__main__":
+    unittest.main()