Refactor kernels to use premake instead of make

Author: voltjia

src/ntops/kernels/abs.py

@@ -1,6 +1,5 @@
 import functools
 
-import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
@@ -11,6 +10,9 @@ def application(input, output):
     output = ntl.abs(input)  # noqa: F841
 
 
-@functools.cache
-def make(ndim):
-    return ninetoothed.make(arrangement, application, (Tensor(ndim), Tensor(ndim)))
+def premake(ndim, dtype=None, block_size=None):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
+
+    tensors = (Tensor(ndim, dtype=dtype), Tensor(ndim, dtype=dtype))
+
+    return arrangement_, application, tensors

src/ntops/kernels/add.py

@@ -1,6 +1,5 @@
 import functools
 
-import ninetoothed
 from ninetoothed import Tensor
 
 from ntops.kernels.element_wise import arrangement
@@ -10,8 +9,14 @@ 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))
+def premake(ndim, dtype=None, block_size=None):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
 
-    return ninetoothed.make(arrangement, application, tensors)
+    tensors = (
+        Tensor(ndim, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+        Tensor(0, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+    )
+
+    return arrangement_, application, tensors

src/ntops/kernels/addmm.py

@@ -1,6 +1,5 @@
 import functools
 
-import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
@@ -47,8 +46,21 @@ def application(input, x, y, beta, alpha, output):
     output = beta * input + alpha * mm_output
 
 
-@functools.cache
-def make():
-    tensors = (Tensor(2), Tensor(2), Tensor(2), Tensor(0), Tensor(0), Tensor(2))
+def premake(dtype=None, block_size_m=None, block_size_n=None, block_size_k=None):
+    arrangement_ = functools.partial(
+        arrangement,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        block_size_k=block_size_k,
+    )
+
+    tensors = (
+        Tensor(2, dtype=dtype),
+        Tensor(2, dtype=dtype),
+        Tensor(2, dtype=dtype),
+        Tensor(0, dtype=dtype),
+        Tensor(0, dtype=dtype),
+        Tensor(2, dtype=dtype),
+    )
 
-    return ninetoothed.make(arrangement, application, tensors)
+    return arrangement_, application, tensors

src/ntops/kernels/bitwise_and.py

@@ -1,6 +1,5 @@
 import functools
 
-import ninetoothed
 from ninetoothed import Tensor
 
 from ntops.kernels.element_wise import arrangement
@@ -10,8 +9,13 @@ def application(input, other, output):
     output = input & other  # noqa: F841
 
 
-@functools.cache
-def make(ndim):
-    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim))
+def premake(ndim, dtype=None, block_size=None):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
 
-    return ninetoothed.make(arrangement, application, tensors)
+    tensors = (
+        Tensor(ndim, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+    )
+
+    return arrangement_, application, tensors

src/ntops/kernels/bitwise_not.py

@@ -1,6 +1,5 @@
 import functools
 
-import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
@@ -15,10 +14,11 @@ def logical_application(input, output):
     output = ntl.where(input, False, True)  # noqa: F841
 
 
-@functools.cache
-def make(ndim, logical=False):
-    tensors = (Tensor(ndim), Tensor(ndim))
+def premake(ndim, logical=False, dtype=None, block_size=None):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
 
     application = logical_application if logical else bitwise_application
 
-    return ninetoothed.make(arrangement, application, tensors)
+    tensors = (Tensor(ndim, dtype=dtype), Tensor(ndim, dtype=dtype))
+
+    return arrangement_, application, tensors

src/ntops/kernels/bitwise_or.py

@@ -1,6 +1,5 @@
 import functools
 
-import ninetoothed
 from ninetoothed import Tensor
 
 from ntops.kernels.element_wise import arrangement
@@ -10,8 +9,13 @@ def application(input, other, output):
     output = input | other  # noqa: F841
 
 
-@functools.cache
-def make(ndim):
-    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim))
+def premake(ndim, dtype=None, block_size=None):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
 
-    return ninetoothed.make(arrangement, application, tensors)
+    tensors = (
+        Tensor(ndim, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+    )
+
+    return arrangement_, application, tensors

src/ntops/kernels/bmm.py

@@ -1,6 +1,5 @@
 import functools
 
-import ninetoothed
 from ninetoothed import Tensor
 
 from ntops.kernels.mm import BLOCK_SIZE_K, BLOCK_SIZE_M, BLOCK_SIZE_N, application
@@ -36,6 +35,14 @@ def arrangement(
     return input_arranged, other_arranged, output_arranged
 
 
-@functools.cache
-def make():
-    return ninetoothed.make(arrangement, application, (Tensor(3), Tensor(3), Tensor(3)))
+def premake(dtype=None, block_size_m=None, block_size_n=None, block_size_k=None):
+    arrangement_ = functools.partial(
+        arrangement,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        block_size_k=block_size_k,
+    )
+
+    tensors = (Tensor(3, dtype=dtype), Tensor(3, dtype=dtype), Tensor(3, dtype=dtype))
+
+    return arrangement_, application, tensors

src/ntops/kernels/clamp.py

@@ -1,6 +1,5 @@
 import functools
 
-import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
@@ -11,8 +10,14 @@ def application(input, min_val, max_val, output):
     output = ntl.clamp(input, min_val, max_val)  # noqa: F841
 
 
-@functools.cache
-def make(ndim):
-    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim), Tensor(ndim))
+def premake(ndim, dtype=None, block_size=None):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
 
-    return ninetoothed.make(arrangement, application, tensors)
+    tensors = (
+        Tensor(ndim, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+    )
+
+    return arrangement_, application, tensors

src/ntops/kernels/cos.py

@@ -1,6 +1,5 @@
 import functools
 
-import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
@@ -11,6 +10,9 @@ def application(input, output):
     output = ntl.cos(input)  # noqa: F841
 
 
-@functools.cache
-def make(ndim):
-    return ninetoothed.make(arrangement, application, (Tensor(ndim), Tensor(ndim)))
+def premake(ndim, dtype=None, block_size=None):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
+
+    tensors = (Tensor(ndim, dtype=dtype), Tensor(ndim, dtype=dtype))
+
+    return arrangement_, application, tensors

src/ntops/kernels/div.py

@@ -1,6 +1,5 @@
 import functools
 
-import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
@@ -19,15 +18,20 @@ def floor_application(input, other, output):
     output = ntl.floor(input / other)  # noqa: F841
 
 
-@functools.cache
-def make(ndim, rounding_mode):
+def premake(ndim, rounding_mode, dtype=None, block_size=None):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
+
     if rounding_mode == "trunc":
         application = trunc_application
     elif rounding_mode == "floor":
         application = floor_application
     else:
         application = default_application
 
-    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim))
+    tensors = (
+        Tensor(ndim, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+        Tensor(ndim, dtype=dtype),
+    )
 
-    return ninetoothed.make(arrangement, application, tensors)
+    return arrangement_, application, tensors