support xccl test

comms/torchcomms/tests/integration/py/AllGatherSingleTest.py

@@ -18,6 +18,8 @@ class AllGatherSingleTest(unittest.TestCase):
 
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/AllGatherTest.py

@@ -18,6 +18,8 @@ class AllGatherTest(unittest.TestCase):
 
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/AllGatherVTest.py

@@ -19,6 +19,8 @@ class AllGatherVTest(unittest.TestCase):
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
     num_replays = 4
 
     def get_wrapper(self):

comms/torchcomms/tests/integration/py/AllReduceTest.py

@@ -21,6 +21,9 @@ class AllReduceTest(unittest.TestCase):
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
+        dtypes = [torch.float, torch.int]
     ops = [ReduceOp.SUM, ReduceOp.MAX]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/AllToAllSingleTest.py

@@ -18,6 +18,8 @@ class AllToAllSingleTest(unittest.TestCase):
 
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/AllToAllTest.py

@@ -18,6 +18,8 @@ class AllToAllTest(unittest.TestCase):
 
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/BroadcastTest.py

@@ -18,6 +18,8 @@ class BroadcastTest(unittest.TestCase):
 
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/GatherTest.py

@@ -18,6 +18,8 @@ class GatherTest(unittest.TestCase):
 
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/ObjColTest.py

@@ -2,6 +2,7 @@
 # pyre-unsafe
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 
+import os
 import unittest
 
 from contextlib import contextmanager
@@ -29,6 +30,8 @@ class ObjColTest(unittest.TestCase):
 
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/ReduceScatterSingleTest.py

@@ -21,6 +21,9 @@ class ReduceScatterSingleTest(unittest.TestCase):
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
+        dtypes = [torch.float, torch.int]
     ops = [ReduceOp.SUM, ReduceOp.MAX]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/ReduceScatterTest.py

@@ -21,6 +21,9 @@ class ReduceScatterTest(unittest.TestCase):
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
+        dtypes = [torch.float, torch.int]
     ops = [ReduceOp.SUM, ReduceOp.MAX]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/ReduceTest.py

@@ -21,6 +21,9 @@ class ReduceTest(unittest.TestCase):
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
+        dtypes = [torch.float, torch.int]
     ops = [ReduceOp.SUM, ReduceOp.MAX]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/ScatterTest.py

@@ -18,6 +18,8 @@ class ScatterTest(unittest.TestCase):
 
     # Class variables for test parameters
     counts = [0, 4, 1024, 1024 * 1024]
+    if os.environ.get("TEST_BACKEND") == "xccl":
+        counts = [4, 1024, 1024 * 1024]
     dtypes = [torch.float, torch.int, torch.int8]
     num_replays = 4
 

comms/torchcomms/tests/integration/py/TorchCommTestHelpers.py

@@ -218,11 +218,13 @@ def get_device(self, backend, rank):
         if device := os.environ.get("TEST_DEVICE"):
             return torch.device(device)
 
-        # Check for CUDA availability and abort if not available
-        if not torch.cuda.is_available():
+        # Check for accelerator availability and abort if not available
+        if not torch.accelerator.is_available():
             return torch.device("cpu")
 
-        device_id = rank % torch.cuda.device_count()
+        device_id = rank % torch.accelerator.device_count()
+        if os.getenv("TEST_BACKEND") == "xccl":
+            return torch.device(f"xpu:{device_id}")
         return torch.device(f"cuda:{device_id}")
 
     def __init__(self, store=None):