Merge pull request #6 from InfiniTensor/develop-rsqrt

Add `rsqrt` operator

Author: voltjia

src/ntops/__init__.py

@@ -7,5 +7,6 @@
 from ntops.gelu import gelu
 from ntops.mm import mm
 from ntops.mul import mul
+from ntops.rsqrt import rsqrt
 
-__all__ = ["abs", "add", "addmm", "bmm", "div", "exp", "gelu", "mm", "mul"]
+__all__ = ["abs", "add", "addmm", "bmm", "div", "exp", "gelu", "mm", "mul", "rsqrt"]

src/ntops/rsqrt.py

@@ -0,0 +1,32 @@
+import functools
+
+import ninetoothed
+import ninetoothed.language as ntl
+import torch
+from ninetoothed import Tensor
+
+from ntops import element_wise
+
+
+def application(input, output):
+    output = ntl.rsqrt(ntl.cast(input, ntl.float32))  # noqa: F841
+
+
+def rsqrt(input, output=None):
+    if output is None:
+        output = torch.empty_like(input)
+
+    kernel = _make(input.ndim)
+
+    kernel(input, output)
+
+    return output
+
+
+@functools.cache
+def _make(ndim):
+    return ninetoothed.make(
+        element_wise.arrangement,
+        application,
+        (Tensor(ndim), Tensor(ndim)),
+    )

tests/test_rsqrt.py

@@ -0,0 +1,21 @@
+import pytest
+import torch
+
+import ntops
+from tests.skippers import skip_if_cuda_not_available
+from tests.utils import generate_arguments
+
+
+@skip_if_cuda_not_available
+@pytest.mark.parametrize(*generate_arguments())
+def test_cuda(shape, dtype, atol, rtol):
+    device = "cuda"
+
+    input = torch.randn(shape, dtype=dtype, device=device)
+
+    ninetoothed_output = ntops.rsqrt(input)
+    reference_output = torch.rsqrt(input)
+
+    assert torch.allclose(
+        ninetoothed_output, reference_output, atol=atol, rtol=rtol, equal_nan=True
+    )