Decouple the kernels

Author: voltjia

src/ntops/__init__.py

@@ -1,12 +0,0 @@
-from ntops.abs import abs
-from ntops.add import add
-from ntops.addmm import addmm
-from ntops.bmm import bmm
-from ntops.div import div
-from ntops.exp import exp
-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", "rsqrt"]

src/ntops/abs.py

@@ -0,0 +1,16 @@
+import functools
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, output):
+    output = ntl.abs(input)  # noqa: F841
+
+
+@functools.cache
+def make(ndim):
+    return ninetoothed.make(arrangement, application, (Tensor(ndim), Tensor(ndim)))

src/ntops/add.py

@@ -0,0 +1,17 @@
+import functools
+
+import ninetoothed
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, other, alpha, output):
+    output = input + alpha * other  # noqa: F841
+
+
+@functools.cache
+def make(ndim):
+    tensors = (Tensor(ndim), Tensor(ndim), Tensor(0), Tensor(ndim))
+
+    return ninetoothed.make(arrangement, application, tensors)

src/ntops/exp.py

@@ -2,10 +2,9 @@
 
 import ninetoothed
 import ninetoothed.language as ntl
-import torch
 from ninetoothed import Tensor
 
-import ntops.mm as mm
+import ntops.kernels.mm as mm
 
 
 def arrangement(input, x, y, beta, alpha, output):
@@ -22,22 +21,8 @@ def application(input, x, y, beta, alpha, output):
     output = beta * input + alpha * mm_output
 
 
-def addmm(input, x, y, beta, alpha, output=None):
-    m, _ = x.shape
-    _, n = y.shape
-
-    if output is None:
-        output = torch.empty((m, n), dtype=input.dtype, device=input.device)
-
-    kernel = _make()
-
-    kernel(input, x, y, beta, alpha, output)
-
-    return output
-
-
 @functools.cache
-def _make():
+def make():
     tensors = (Tensor(2), Tensor(2), Tensor(2), Tensor(0), Tensor(0), Tensor(2))
 
     return ninetoothed.make(arrangement, application, tensors)

src/ntops/kernels/__init__.py

@@ -1,10 +1,9 @@
 import functools
 
 import ninetoothed
-import torch
 from ninetoothed import Tensor
 
-from ntops.mm import BLOCK_SIZE_K, BLOCK_SIZE_M, BLOCK_SIZE_N, application
+from ntops.kernels.mm import BLOCK_SIZE_K, BLOCK_SIZE_M, BLOCK_SIZE_N, application
 
 
 def arrangement(input, other, output):
@@ -26,20 +25,6 @@ def arrangement(input, other, output):
     return input_arranged, other_arranged, output_arranged
 
 
-def bmm(input, other, output=None):
-    b, m, _ = input.shape
-    _, _, n = other.shape
-
-    if output is None:
-        output = torch.empty((b, m, n), dtype=input.dtype, device=input.device)
-
-    kernel = _make()
-
-    kernel(input, other, output)
-
-    return output
-
-
 @functools.cache
-def _make():
+def make():
     return ninetoothed.make(arrangement, application, (Tensor(3), Tensor(3), Tensor(3)))

src/ntops/kernels/abs.py

@@ -2,10 +2,9 @@
 
 import ninetoothed
 import ninetoothed.language as ntl
-import torch
 from ninetoothed import Tensor
 
-from ntops import element_wise
+from ntops.kernels.element_wise import arrangement
 
 
 def default_application(input, other, output):
@@ -20,26 +19,15 @@ def floor_application(input, other, output):
     output = ntl.floor(input / other)  # noqa: F841
 
 
-def div(input, other, rounding_mode=None, output=None):
-    if output is None:
-        output = torch.empty_like(input)
-
-    kernel = _make(input.ndim, rounding_mode)
-
-    kernel(input, other, output)
-
-    return output
-
-
 @functools.cache
-def _make(ndim, rounding_mode):
-    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim))
-
+def make(ndim, rounding_mode):
     if rounding_mode == "trunc":
         application = trunc_application
     elif rounding_mode == "floor":
         application = floor_application
     else:
         application = default_application
 
-    return ninetoothed.make(element_wise.arrangement, application, tensors)
+    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim))
+
+    return ninetoothed.make(arrangement, application, tensors)