feat: add torch.compile triton self-test

Author: Godnight1006

API.md

@@ -15,6 +15,13 @@ torchruntime.install(["torch", "torchvision<0.20"])
 
 On Windows CUDA, Linux ROCm (6.x+), and Linux XPU, this also installs the appropriate Triton package to enable `torch.compile` (`triton-windows`, `pytorch-triton-rocm`, or `pytorch-triton-xpu`).
 
+## Test torch
+Run:
+`python -m torchruntime test`
+
+To specifically verify `torch.compile` / Triton:
+`python -m torchruntime test compile`
+
 ## Get device info
 You can use the device database built into `torchruntime` for your projects:
 ```py

README.md

@@ -44,7 +44,7 @@ torchruntime.configure()
 ```
 
 ### (Optional) Step 3. Test torch
-Run `python -m torchruntime test` to run a set of tests to check whether the installed version of torch is working correctly.
+Run `python -m torchruntime test` to run a set of tests to check whether the installed version of torch is working correctly (including a `torch.compile` / Triton check on CUDA/XPU systems). You can also run `python -m torchruntime test compile` to run only the compile check.
 
 ## Customizing packages
 By default, `python -m torchruntime install` will install the latest available `torch`, `torchvision` and `torchaudio` suitable on the user's platform.

torchruntime/__main__.py

@@ -10,7 +10,7 @@ def print_usage(entry_command: str):
 
 Commands:
     install             Install PyTorch packages
-    test [subcommand]   Run tests (subcommands: all, devices, math, functions)
+    test [subcommand]   Run tests (subcommands: all, import, devices, compile, math, functions)
     --help             Show this help message
 
 Examples:
@@ -20,10 +20,11 @@ def print_usage(entry_command: str):
     {entry_command} install --uv torch>=2.0.0 torchaudio
     {entry_command} install torch==2.1.* torchvision>=0.16.0 torchaudio==2.1.0
 
-    {entry_command} test          # Runs all tests (import, devices, math, functions)
+    {entry_command} test          # Runs all tests (import, devices, compile, math, functions)
     {entry_command} test all      # Same as above
     {entry_command} test import  # Test only import
     {entry_command} test devices  # Test only devices
+    {entry_command} test compile  # Test torch.compile (Triton)
     {entry_command} test math     # Test only math
     {entry_command} test functions # Test only functions
 

torchruntime/utils/torch_test/__init__.py

@@ -1,6 +1,8 @@
+import importlib.util
+import platform
 import time
 
-from ..torch_device_utils import get_installed_torch_platform, get_device_count, get_device_name, get_device
+from ..torch_device_utils import get_device, get_device_count, get_device_name, get_installed_torch_platform
 
 
 def test(subcommand):
@@ -16,7 +18,7 @@ def test(subcommand):
 
 
 def test_all():
-    for fn in (test_import, test_devices, test_math, test_functions):
+    for fn in (test_import, test_devices, test_compile, test_math, test_functions):
         fn()
         print("")
 
@@ -101,3 +103,74 @@ def test_functions():
         t.run_all_tests()
 
     print("--- / FUNCTIONAL TEST ---")
+
+
+def test_compile():
+    print("--- COMPILE TEST ---")
+
+    try:
+        import torch
+    except ImportError:
+        print("torch.compile: SKIPPED (torch not installed)")
+        print("--- / COMPILE TEST ---")
+        return
+
+    if not hasattr(torch, "compile"):
+        print("torch.compile: SKIPPED (requires torch>=2.0)")
+        print("--- / COMPILE TEST ---")
+        return
+
+    torch_platform_name, _ = get_installed_torch_platform()
+    if torch_platform_name not in ("cuda", "xpu"):
+        print(f"torch.compile: SKIPPED (unsupported backend: {torch_platform_name})")
+        print("--- / COMPILE TEST ---")
+        return
+
+    if importlib.util.find_spec("triton") is None:
+        print("triton: NOT INSTALLED")
+    else:
+        print("triton: installed")
+
+    device = get_device(0)
+    print("On torch device:", device)
+
+    def f(x):
+        return x * 2 + 1
+
+    try:
+        compiled_f = torch.compile(f)
+        x = torch.randn((1024,), device=device)
+        y = compiled_f(x)
+        expected = f(x)
+        if not torch.allclose(y, expected):
+            print("torch.compile: FAILED (output mismatch)")
+        else:
+            if torch_platform_name == "cuda":
+                torch.cuda.synchronize()
+            if torch_platform_name == "xpu" and hasattr(torch, "xpu") and hasattr(torch.xpu, "synchronize"):
+                torch.xpu.synchronize()
+            print("torch.compile: PASSED")
+    except Exception as e:
+        print(f"torch.compile: FAILED ({type(e).__name__}: {e})")
+
+        hint = None
+        os_name = platform.system()
+        if torch_platform_name == "cuda" and os_name == "Windows":
+            hint = "pip install triton-windows  (or: python -m torchruntime install)"
+        elif torch_platform_name == "cuda" and os_name == "Linux":
+            if getattr(torch.version, "hip", None):
+                hint = (
+                    "pip install pytorch-triton-rocm --index-url https://download.pytorch.org/whl  "
+                    "(or: python -m torchruntime install)"
+                )
+        elif torch_platform_name == "xpu" and os_name == "Linux":
+            hint = (
+                "pip install pytorch-triton-xpu --index-url https://download.pytorch.org/whl  "
+                "(or: python -m torchruntime install)"
+            )
+
+        if hint:
+            print("If this failed due to Triton, try:")
+            print("  ", hint)
+
+    print("--- / COMPILE TEST ---")