From f8423d8e718cc306a8d29d63b0bdaafe00a02a7d Mon Sep 17 00:00:00 2001
From: "Han, Chao" <chao1.han@intel.com>
Date: Fri, 7 Nov 2025 10:27:40 +0800
Subject: [PATCH] support xccl test

---
 .../tests/integration/py/AllGatherSingleTest.py           | 2 ++
 comms/torchcomms/tests/integration/py/AllGatherTest.py    | 2 ++
 comms/torchcomms/tests/integration/py/AllGatherVTest.py   | 2 ++
 comms/torchcomms/tests/integration/py/AllReduceTest.py    | 3 +++
 .../torchcomms/tests/integration/py/AllToAllSingleTest.py | 2 ++
 comms/torchcomms/tests/integration/py/AllToAllTest.py     | 2 ++
 comms/torchcomms/tests/integration/py/BroadcastTest.py    | 2 ++
 comms/torchcomms/tests/integration/py/GatherTest.py       | 2 ++
 comms/torchcomms/tests/integration/py/ObjColTest.py       | 3 +++
 .../tests/integration/py/ReduceScatterSingleTest.py       | 3 +++
 .../torchcomms/tests/integration/py/ReduceScatterTest.py  | 3 +++
 comms/torchcomms/tests/integration/py/ReduceTest.py       | 3 +++
 comms/torchcomms/tests/integration/py/ScatterTest.py      | 2 ++
 .../tests/integration/py/TorchCommTestHelpers.py          | 8 +++++---
 14 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/comms/torchcomms/tests/integration/py/AllGatherSingleTest.py b/comms/torchcomms/tests/integration/py/AllGatherSingleTest.py
index 428f0e9f1e..ab3fb89e61 100644
--- a/comms/torchcomms/tests/integration/py/AllGatherSingleTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/AllGatherTest.py b/comms/torchcomms/tests/integration/py/AllGatherTest.py
index 553ca4904a..a5cc1fb7f5 100644
--- a/comms/torchcomms/tests/integration/py/AllGatherTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/AllGatherVTest.py b/comms/torchcomms/tests/integration/py/AllGatherVTest.py
index f971144d55..eda57a385c 100644
--- a/comms/torchcomms/tests/integration/py/AllGatherVTest.py
+++ b/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):
diff --git a/comms/torchcomms/tests/integration/py/AllReduceTest.py b/comms/torchcomms/tests/integration/py/AllReduceTest.py
index b6125e9e5a..78d4a0c526 100644
--- a/comms/torchcomms/tests/integration/py/AllReduceTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/AllToAllSingleTest.py b/comms/torchcomms/tests/integration/py/AllToAllSingleTest.py
index 5d06af449e..241e7e3b40 100644
--- a/comms/torchcomms/tests/integration/py/AllToAllSingleTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/AllToAllTest.py b/comms/torchcomms/tests/integration/py/AllToAllTest.py
index fd39660ea7..ecdee379f4 100644
--- a/comms/torchcomms/tests/integration/py/AllToAllTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/BroadcastTest.py b/comms/torchcomms/tests/integration/py/BroadcastTest.py
index 9b1b5bbc69..521983a53c 100644
--- a/comms/torchcomms/tests/integration/py/BroadcastTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/GatherTest.py b/comms/torchcomms/tests/integration/py/GatherTest.py
index a3ab139a60..acaa7d30e8 100644
--- a/comms/torchcomms/tests/integration/py/GatherTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/ObjColTest.py b/comms/torchcomms/tests/integration/py/ObjColTest.py
index 28e6e897d8..30028874b0 100644
--- a/comms/torchcomms/tests/integration/py/ObjColTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/ReduceScatterSingleTest.py b/comms/torchcomms/tests/integration/py/ReduceScatterSingleTest.py
index 5e15f6208b..dc1725235a 100644
--- a/comms/torchcomms/tests/integration/py/ReduceScatterSingleTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/ReduceScatterTest.py b/comms/torchcomms/tests/integration/py/ReduceScatterTest.py
index daf5941e9a..8069228a7d 100644
--- a/comms/torchcomms/tests/integration/py/ReduceScatterTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/ReduceTest.py b/comms/torchcomms/tests/integration/py/ReduceTest.py
index b058822c61..69b892a8f1 100644
--- a/comms/torchcomms/tests/integration/py/ReduceTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/ScatterTest.py b/comms/torchcomms/tests/integration/py/ScatterTest.py
index 09f0593d59..8b3ff0bb08 100644
--- a/comms/torchcomms/tests/integration/py/ScatterTest.py
+++ b/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
 
diff --git a/comms/torchcomms/tests/integration/py/TorchCommTestHelpers.py b/comms/torchcomms/tests/integration/py/TorchCommTestHelpers.py
index b6a955c215..44da2f0555 100644
--- a/comms/torchcomms/tests/integration/py/TorchCommTestHelpers.py
+++ b/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):