pyg-team · rusty1s · Oct 22, 2022 · Oct 22, 2022 · Oct 22, 2022 · Oct 22, 2022
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
 ### Added
+- Added `triton` dependency ([#133](https://github.com/pyg-team/pyg-lib/pull/133))
 - Enable `pytest` testing ([#132](https://github.com/pyg-team/pyg-lib/pull/132))
 - Added C++-based autograd and TorchScript support for `segment_matmul` ([#120](https://github.com/pyg-team/pyg-lib/pull/120), [#122](https://github.com/pyg-team/pyg-lib/pull/122))
 - Allow overriding `time` for seed nodes via `seed_time` in `neighbor_sample` ([#118](https://github.com/pyg-team/pyg-lib/pull/118))

diff --git a/pyg_lib/testing.py b/pyg_lib/testing.py
@@ -18,6 +18,14 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def onlyCUDA(func: Callable) -> Callable:
+    import pytest
+    return pytest.mark.skipif(
+        not torch.cuda.is_available(),
+        reason="CUDA not available",
+    )(func)
+
+
 def withCUDA(func: Callable) -> Callable:
     def wrapper(*args, **kwargs):
         func(*args, device=torch.device('cpu'), **kwargs)

diff --git a/setup.cfg b/setup.cfg
@@ -18,6 +18,9 @@ test=pytest
 [tool:pytest]
 addopts=--capture=no
 
+[flake8]
+ignore=E731
+
 [isort]
 multi_line_output=3
 include_trailing_comma=True

diff --git a/setup.py b/setup.py
@@ -70,7 +70,9 @@ def build_extension(self, ext):
                               cwd=self.build_temp)
 
 
-install_requires = []
+install_requires = [
+    'triton',
+]
 
 test_requires = [
     'pytest',

diff --git a/test/test_triton.py b/test/test_triton.py
@@ -0,0 +1,37 @@
+import torch
+import triton
+import triton.language as tl
+from torch import Tensor
+
+from pyg_lib.testing import onlyCUDA
+
+
+@triton.jit
+def add_kernel(x_ptr, y_ptr, out_ptr, numel, **meta):
+    pid = tl.program_id(axis=0)
+    block_start = pid * meta['BLOCK_SIZE']
+
+    offsets = block_start + tl.arange(0, meta['BLOCK_SIZE'])
+    mask = offsets < numel
+
+    x = tl.load(x_ptr + offsets, mask=mask)
+    y = tl.load(y_ptr + offsets, mask=mask)
+
+    output = x + y
+
+    tl.store(out_ptr + offsets, output, mask=mask)
+
+
+def add(x: Tensor, y: Tensor) -> Tensor:
+    out = torch.empty_like(x)
+    assert x.is_cuda and y.is_cuda and out.is_cuda
+    grid = lambda meta: (triton.cdiv(x.numel(), meta['BLOCK_SIZE']), )
+    add_kernel[grid](x, y, out, x.numel(), BLOCK_SIZE=1024)
+    return out
+
+
+@onlyCUDA
+def test_triton():
+    x = torch.rand(100, device='cuda')
+    y = torch.rand(100, device='cuda')
+    assert torch.allclose(x + y, add(x, y))