Merge pull request #13 from denix56/triton

Gate torch.compile in tests for windows

Author: denix56

test/test_fps.py

@@ -4,14 +4,24 @@
 import torch
 from torch import Tensor
 from torch_cluster import fps
-from torch_cluster.testing import devices, grad_dtypes, tensor
+from torch_cluster.testing import (
+    devices,
+    grad_dtypes,
+    has_compiler,
+    tensor,
+)
 
 
-@torch.compile
-def fps2(x: Tensor, ratio: Tensor) -> Tensor:
+def fps2_impl(x: Tensor, ratio: Tensor) -> Tensor:
     return fps(x, None, ratio, False)
 
 
+if has_compiler():
+    fps2 = torch.compile(fps2_impl)
+else:
+    fps2 = fps2_impl
+
+
 @pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
 def test_fps(dtype, device):
     x = tensor([

test/test_grid.py

@@ -3,7 +3,7 @@
 import pytest
 import torch
 from torch_cluster import grid_cluster
-from torch_cluster.testing import devices, dtypes, tensor
+from torch_cluster.testing import devices, dtypes, has_compiler, tensor
 
 tests = [{
     'pos': [2, 6],
@@ -39,5 +39,6 @@ def test_grid_cluster(test, dtype, device):
     cluster = grid_cluster(pos, size, start, end)
     assert cluster.tolist() == test['cluster']
 
-    jit = torch.compile(grid_cluster)
-    assert torch.equal(jit(pos, size, start, end), cluster)
+    if has_compiler():
+        jit = torch.compile(grid_cluster)
+        assert torch.equal(jit(pos, size, start, end), cluster)

test/test_knn.py

@@ -5,7 +5,13 @@
 import scipy.spatial
 import torch
 from torch_cluster import knn, knn_graph
-from torch_cluster.testing import devices, grad_dtypes, tensor, triton_wrap
+from torch_cluster.testing import (
+    devices,
+    grad_dtypes,
+    has_compiler,
+    tensor,
+    triton_wrap,
+)
 
 HAS_CUDA = torch.cuda.is_available()
 HAS_TRITON = importlib.util.find_spec('triton') is not None
@@ -41,9 +47,10 @@ def test_knn(dtype, device, use_triton):
     edge_index = knn(x, y, 2, use_triton=use_triton)
     assert to_set(edge_index) == {(0, 2), (0, 3), (1, 0), (1, 1)}
 
-    jit = torch.compile(knn)
-    edge_index = jit(x, y, 2, use_triton=use_triton)
-    assert to_set(edge_index) == {(0, 2), (0, 3), (1, 0), (1, 1)}
+    if has_compiler():
+        jit = torch.compile(knn)
+        edge_index = jit(x, y, 2, use_triton=use_triton)
+        assert to_set(edge_index) == {(0, 2), (0, 3), (1, 0), (1, 1)}
 
     edge_index = knn(x, y, 2, batch_x, batch_y, use_triton=use_triton)
     assert to_set(edge_index) == {(0, 2), (0, 3), (1, 4), (1, 5)}
@@ -113,18 +120,19 @@ def test_knn_graph(dtype, device, use_triton):
         (2, 3),
     }
 
-    jit = torch.compile(knn_graph)
-    edge_index = jit(x, k=2, flow='source_to_target')
-    assert to_set(edge_index) == {
-        (1, 0),
-        (3, 0),
-        (0, 1),
-        (2, 1),
-        (1, 2),
-        (3, 2),
-        (0, 3),
-        (2, 3),
-    }
+    if has_compiler():
+        jit = torch.compile(knn_graph)
+        edge_index = jit(x, k=2, flow='source_to_target')
+        assert to_set(edge_index) == {
+            (1, 0),
+            (3, 0),
+            (0, 1),
+            (2, 1),
+            (1, 2),
+            (3, 2),
+            (0, 3),
+            (2, 3),
+        }
 
 
 @pytest.mark.parametrize(

test/test_radius.py

@@ -4,7 +4,12 @@
 import scipy.spatial
 import torch
 from torch_cluster import radius, radius_graph
-from torch_cluster.testing import devices, floating_dtypes, tensor
+from torch_cluster.testing import (
+    devices,
+    floating_dtypes,
+    has_compiler,
+    tensor,
+)
 
 
 def to_set(edge_index):
@@ -44,10 +49,19 @@ def test_radius(dtype, device):
     assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 1),
                                       (1, 2), (1, 5), (1, 6)])
 
-    jit = torch.compile(radius)
-    edge_index = jit(x, y, 2, max_num_neighbors=4)
-    assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 1),
-                                      (1, 2), (1, 5), (1, 6)])
+    if has_compiler():
+        jit = torch.compile(radius)
+        edge_index = jit(x, y, 2, max_num_neighbors=4)
+        assert to_set(edge_index) == set([
+            (0, 0),
+            (0, 1),
+            (0, 2),
+            (0, 3),
+            (1, 1),
+            (1, 2),
+            (1, 5),
+            (1, 6),
+        ])
 
     edge_index = radius(x, y, 2, batch_x, batch_y, max_num_neighbors=4)
     assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 5),
@@ -78,10 +92,19 @@ def test_radius_graph(dtype, device):
     assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
                                       (3, 2), (0, 3), (2, 3)])
 
-    jit = torch.compile(radius_graph)
-    edge_index = jit(x, r=2.5, flow='source_to_target')
-    assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
-                                      (3, 2), (0, 3), (2, 3)])
+    if has_compiler():
+        jit = torch.compile(radius_graph)
+        edge_index = jit(x, r=2.5, flow='source_to_target')
+        assert to_set(edge_index) == set([
+            (1, 0),
+            (3, 0),
+            (0, 1),
+            (2, 1),
+            (1, 2),
+            (3, 2),
+            (0, 3),
+            (2, 3),
+        ])
 
     edge_index = radius_graph(x, r=100, flow='source_to_target',
                               max_num_neighbors=1)

test/test_rw.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 from torch_cluster import random_walk
-from torch_cluster.testing import devices, tensor
+from torch_cluster.testing import devices, has_compiler, tensor
 
 
 @pytest.mark.parametrize('device', devices)
@@ -31,8 +31,12 @@ def test_rw_small(device):
     out = random_walk(row, col, start, walk_length, num_nodes=3)
     assert out.tolist() == [[0, 1, 0, 1, 0], [1, 0, 1, 0, 1], [2, 2, 2, 2, 2]]
 
-    jit = torch.compile(random_walk)
-    assert torch.equal(jit(row, col, start, walk_length, num_nodes=3), out)
+    if has_compiler():
+        jit = torch.compile(random_walk)
+        assert torch.equal(
+            jit(row, col, start, walk_length, num_nodes=3),
+            out,
+        )
 
 
 @pytest.mark.parametrize('device', devices)

torch_cluster/testing.py

@@ -1,4 +1,6 @@
 from typing import Any
+import platform
+import shutil
 
 import torch
 
@@ -27,3 +29,9 @@ def triton_wrap(dt_device_seq):
         for dt, device in dt_device_seq
         for use_triton in ([False, True] if device.type == 'cuda' else [False])
     ]
+
+
+def has_compiler() -> bool:
+    if platform.system().lower() != 'windows':
+        return True
+    return shutil.which('cl') is not None