Add test_sycl_extension from test_cpp_extensions_aot for XPU

.github/workflows/_windows_ut.yml

@@ -249,6 +249,22 @@ jobs:
 
           if not exist "%GITHUB_WORKSPACE%\ut_log" mkdir "%GITHUB_WORKSPACE%\ut_log"
           copy op_extended.xml %GITHUB_WORKSPACE%\ut_log /Y
+      - name: Run CPP Extensions AOT SYCL UT
+        if: contains(inputs.ut, 'test_xpu') || github.event_name == 'schedule'
+        shell: cmd
+        continue-on-error: true
+        run: |
+          call "C:\ProgramData\miniforge3\Scripts\activate.bat"
+          call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat"
+          call "C:\Program Files (x86)\Intel\oneAPI\compiler\latest\env\vars.bat"
+          call "C:\Program Files (x86)\Intel\oneAPI\ocloc\latest\env\vars.bat"
+          call "C:\Program Files (x86)\Intel\oneAPI\pti\latest\env\vars.bat"
+          call conda activate windows_ci
+          cd ../pytorch/third_party/torch-xpu-ops/test/xpu/
+          python -m pytest test_cpp_extensions_aot_xpu.py -v --junit-xml=test_cpp_extensions_aot_xpu.xml
+
+          if not exist "%GITHUB_WORKSPACE%\ut_log" mkdir "%GITHUB_WORKSPACE%\ut_log"
+          copy test_cpp_extensions_aot_xpu.xml %GITHUB_WORKSPACE%\ut_log /Y
       - name: Run Test XPU UT
         if: contains(inputs.ut, 'test_xpu') || github.event_name == 'schedule'
         shell: cmd

test/xpu/test_cpp_extensions_aot_xpu.py

@@ -0,0 +1,119 @@
+# Owner(s): ["module: cpp-extensions"]
+
+import os
+import shutil
+import subprocess
+import sys
+import unittest
+
+import torch
+import torch.testing._internal.common_utils as common
+from torch.testing._internal.common_utils import TEST_XPU
+
+# Resolved in setUpClass; points to pytorch/test/ in the CI layout.
+_test_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../../test")
+_cpp_extensions_dir = os.path.join(_test_dir, "cpp_extensions")
+
+
+def _build_cpp_extensions():
+    """Build AOT cpp extensions from pytorch/test/cpp_extensions/.
+
+    Mirrors the build logic from pytorch/test/run_test.py's
+    _test_cpp_extensions_aot helper, but only builds the parts needed
+    for the SYCL extension test.
+    """
+    build_dir = os.path.join(_cpp_extensions_dir, "build")
+    if os.path.exists(build_dir):
+        shutil.rmtree(build_dir)
+
+    shell_env = os.environ.copy()
+    shell_env["USE_NINJA"] = "1"
+    if sys.platform == "win32":
+        shell_env["DISTUTILS_USE_SDK"] = "1"
+
+    install_cmd = [
+        sys.executable,
+        "-m",
+        "pip",
+        "install",
+        "--no-build-isolation",
+        ".",
+        "--root",
+        "./install",
+    ]
+
+    return_code = subprocess.call(install_cmd, cwd=_cpp_extensions_dir, env=shell_env)
+    if return_code != 0:
+        raise RuntimeError(
+            f"Failed to build cpp extensions (exit code {return_code}). "
+            f"Build dir: {_cpp_extensions_dir}"
+        )
+
+    # Add the installed packages to sys.path so they can be imported
+    install_directories = []
+    for root, directories, _ in os.walk(os.path.join(_cpp_extensions_dir, "install")):
+        for directory in directories:
+            if "-packages" in directory:
+                install_directories.append(os.path.join(root, directory))
+
+    for d in install_directories:
+        if d not in sys.path:
+            sys.path.insert(0, d)
+
+
+def _check_sycl_extension_available():
+    try:
+        import torch_test_cpp_extension.sycl  # noqa: F401
+
+        return True
+    except ModuleNotFoundError as e:
+        # Only treat the extension module itself being absent as "not available".
+        # Re-raise other missing-module errors (e.g. a broken dependency) so the
+        # underlying failure is visible instead of being silently masked.
+        if getattr(e, "name", None) in (
+            "torch_test_cpp_extension",
+            "torch_test_cpp_extension.sycl",
+        ):
+            return False
+        raise
+
+
+@torch.testing._internal.common_utils.markDynamoStrictTest
+class TestCppExtensionAOT(common.TestCase):
+    """Tests ahead-of-time SYCL cpp extensions on XPU."""
+
+    _sycl_extension_built = False
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        if not TEST_XPU:
+            return
+        # Defer sys.path mutation until we actually need it
+        if _test_dir not in sys.path:
+            sys.path.insert(0, _test_dir)
+        if _check_sycl_extension_available():
+            cls._sycl_extension_built = True
+            return
+        # Build the AOT cpp extensions; let failures surface as test errors
+        _build_cpp_extensions()
+        cls._sycl_extension_built = _check_sycl_extension_available()
+
+    @unittest.skipIf(not TEST_XPU, "XPU not found")
+    def test_sycl_extension(self):
+        if not self._sycl_extension_built:
+            self.skipTest("torch_test_cpp_extension.sycl not built")
+
+        import torch_test_cpp_extension.sycl as sycl_extension
+
+        x = torch.zeros(100, device="xpu", dtype=torch.float32)
+        y = torch.zeros(100, device="xpu", dtype=torch.float32)
+
+        z = sycl_extension.sigmoid_add(x, y).cpu()
+
+        # 2 * sigmoid(0) = 2 * 0.5 = 1
+        self.assertEqual(z, torch.ones_like(z))
+
+
+if __name__ == "__main__":
+    common.run_tests()