Merge pull request #41 from InfiniTensor/develop-premake-mechanism

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,14 +1,39 @@
 import functools
 
-import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
 import ntops.kernels.mm as mm
 
 
-def arrangement(input, x, y, beta, alpha, output):
-    _, _, input_arranged = mm.arrangement(x, y, input)
+def arrangement(
+    input,
+    x,
+    y,
+    beta,
+    alpha,
+    output,
+    block_size_m=None,
+    block_size_n=None,
+    block_size_k=None,
+):
+    if block_size_m is None:
+        block_size_m = mm.BLOCK_SIZE_M
+
+    if block_size_n is None:
+        block_size_n = mm.BLOCK_SIZE_N
+
+    if block_size_k is None:
+        block_size_k = mm.BLOCK_SIZE_K
+
+    _, _, input_arranged = mm.arrangement(
+        x,
+        y,
+        input,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        block_size_k=block_size_k,
+    )
 
     x_arranged, y_arranged, output_arranged = mm.arrangement(x, y, output)
 
@@ -21,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,22 +1,32 @@
 import functools
 
-import ninetoothed
 from ninetoothed import Tensor
 
 from ntops.kernels.mm import BLOCK_SIZE_K, BLOCK_SIZE_M, BLOCK_SIZE_N, application
 
 
-def arrangement(input, other, output):
-    output_arranged = output.tile((1, BLOCK_SIZE_M, BLOCK_SIZE_N))
+def arrangement(
+    input, other, output, block_size_m=None, block_size_n=None, block_size_k=None
+):
+    if block_size_m is None:
+        block_size_m = BLOCK_SIZE_M
+
+    if block_size_n is None:
+        block_size_n = BLOCK_SIZE_N
+
+    if block_size_k is None:
+        block_size_k = BLOCK_SIZE_K
+
+    output_arranged = output.tile((1, block_size_m, block_size_n))
     output_arranged.dtype = output_arranged.dtype.squeeze(0)
 
-    input_arranged = input.tile((1, BLOCK_SIZE_M, BLOCK_SIZE_K))
+    input_arranged = input.tile((1, block_size_m, block_size_k))
     input_arranged = input_arranged.tile((1, 1, -1))
     input_arranged = input_arranged.expand((-1, -1, output_arranged.shape[-1]))
     input_arranged.dtype = input_arranged.dtype.squeeze((0, 1))
     input_arranged.dtype.dtype = input_arranged.dtype.dtype.squeeze(0)
 
-    other_arranged = other.tile((1, BLOCK_SIZE_K, BLOCK_SIZE_N))
+    other_arranged = other.tile((1, block_size_k, block_size_n))
     other_arranged = other_arranged.tile((1, -1, 1))
     other_arranged = other_arranged.expand((-1, output_arranged.shape[-2], -1))
     other_arranged.dtype = other_arranged.dtype.squeeze((0, 2))
@@ -25,6 +35,14 @@ def arrangement(input, other, output):
     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