Migrate more kernels to ninetoothed.build

Author: voltjia

ops/ninetoothed/kernels/addmm.py

@@ -1,13 +1,17 @@
+import functools
+
 import ninetoothed
 from ninetoothed import Tensor
 
 import ops.ninetoothed.kernels.mm as mm
+from ops.ninetoothed.kernels._common import build
 
 
-def arrangement(input, mat1, mat2, beta, alpha, output):
-    _, _, input_arranged = mm.arrangement(mat1, mat2, input)
-
-    mat1_arranged, mat2_arranged, output_arranged = mm.arrangement(mat1, mat2, output)
+def arrangement(input, mat1, mat2, beta, alpha, output, **block_sizes):
+    _, _, input_arranged = mm.arrangement(mat1, mat2, input, **block_sizes)
+    mat1_arranged, mat2_arranged, output_arranged = mm.arrangement(
+        mat1, mat2, output, **block_sizes
+    )
 
     return input_arranged, mat1_arranged, mat2_arranged, beta, alpha, output_arranged
 
@@ -17,6 +21,76 @@ def application(input, mat1, mat2, beta, alpha, output):
     output = beta * input + alpha * output
 
 
-tensors = (Tensor(2), Tensor(2), Tensor(2), Tensor(0), Tensor(0), Tensor(2))
+def premake(m, n, k, dtype, block_size_m, block_size_n, block_size_k):
+    arrangement_ = functools.partial(
+        arrangement,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        block_size_k=block_size_k,
+    )
+    tensors = (
+        Tensor(shape=(m, n), dtype=dtype),
+        Tensor(shape=(m, k), dtype=dtype),
+        Tensor(shape=(k, n), dtype=dtype),
+        Tensor(0, dtype=ninetoothed.float32),
+        Tensor(0, dtype=ninetoothed.float32),
+        Tensor(shape=(m, n), dtype=dtype),
+    )
+
+    return arrangement_, application, tensors
+
+
+# Compile the square shapes used by the benchmark; other shapes use the generic
+# make-based fallback below.
+configs = tuple(
+    (
+        (),
+        {
+            "m": s,
+            "n": s,
+            "k": s,
+            "dtype": ninetoothed.float16,
+            "block_size_m": bm,
+            "block_size_n": bn,
+            "block_size_k": bk,
+        },
+        {"num_warps": nw, "num_stages": 3},
+    )
+    for s in (128 * i for i in range(2, 33))
+    for bm in (64, 128)
+    for bn in (64, 128)
+    for bk in (32, 64)
+    for nw in (4, 8)
+)
+
+_build_kernel = build(
+    premake,
+    configs,
+    meta_parameters=("block_size_m", "block_size_n", "block_size_k"),
+    kernel_name="addmm",
+)
+
+_BUILD_CONFIGS = frozenset(
+    (kwargs["m"], kwargs["n"], kwargs["k"], kwargs["dtype"])
+    for _, kwargs, _ in configs
+)
+
+_fallback_kernel = ninetoothed.make(
+    arrangement,
+    application,
+    (
+        Tensor(2),
+        Tensor(2),
+        Tensor(2),
+        Tensor(0),
+        Tensor(0),
+        Tensor(2),
+    ),
+)
+
+
+def kernel(input, mat1, mat2, beta, alpha, output, m, n, k, dtype):
+    if (m, n, k, dtype) in _BUILD_CONFIGS:
+        return _build_kernel(input, mat1, mat2, beta, alpha, output, m, n, k, dtype)
 
-kernel = ninetoothed.make(arrangement, application, tensors)
+    return _fallback_kernel(input, mat1, mat2, beta, alpha, output)

ops/ninetoothed/kernels/conv2d.py

@@ -1,10 +1,13 @@
+import functools
+
 import ninetoothed
 from ninetoothed import Tensor
 
 import ops.ninetoothed.kernels.mm as mm
+from ops.ninetoothed.kernels._common import build
 
 
-def arrangement(input, filter, output):
+def arrangement(input, filter, output, **block_sizes):
     input_arranged = input.tile((1, *filter.shape[1:]), strides=(-1, -1, 1, 1))
     input_arranged = input_arranged.squeeze(1)
     input_arranged.dtype = input_arranged.dtype.squeeze(0)
@@ -16,10 +19,90 @@ def arrangement(input, filter, output):
 
     output_arranged = output.permute((0, 2, 3, 1)).flatten(end_dim=3)
 
-    return mm.arrangement(input_arranged, filter_arranged, output_arranged)
+    return mm.arrangement(
+        input_arranged, filter_arranged, output_arranged, **block_sizes
+    )
+
+
+def premake(n, c, h, w, k, r, s, dtype, block_size_m, block_size_n, block_size_k):
+    arrangement_ = functools.partial(
+        arrangement,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        block_size_k=block_size_k,
+    )
+
+    p = h - r + 1
+    q = w - s + 1
+
+    tensors = (
+        Tensor(shape=(n, c, h, w), dtype=dtype),
+        Tensor(shape=(k, c, r, s), dtype=dtype),
+        Tensor(shape=(n, k, p, q), dtype=dtype),
+    )
+
+    return arrangement_, mm.application, tensors
+
+
+# Block sweep approximating the JIT auto-tuner default range. Conv2d's im2col
+# arrangement produces a tall-skinny matmul, so wider bn (up to 256) and
+# longer-pipelined num_stages help match the JIT-tuned kernel.
+configs = tuple(
+    (
+        (),
+        {
+            "n": n,
+            "c": 512,
+            "h": 14,
+            "w": 14,
+            "k": 512,
+            "r": 3,
+            "s": 3,
+            "dtype": ninetoothed.float16,
+            "block_size_m": bm,
+            "block_size_n": bn,
+            "block_size_k": bk,
+        },
+        {"num_warps": 8, "num_stages": ns},
+    )
+    for n in (2, 4, 8, 16, 32, 64, 128, 256, 512, 1024)
+    for bm in (64, 128)
+    for bn in (128, 256)
+    for bk in (32, 64)
+    for ns in (3, 5)
+    if bm * bn <= 32768 and bm * bk <= 32768 and bn * bk <= 32768
+)
+
+_build_kernel = build(
+    premake,
+    configs,
+    meta_parameters=("block_size_m", "block_size_n", "block_size_k"),
+    kernel_name="conv2d",
+)
+
+_BUILD_CONFIGS = frozenset(
+    (
+        kwargs["n"],
+        kwargs["c"],
+        kwargs["h"],
+        kwargs["w"],
+        kwargs["k"],
+        kwargs["r"],
+        kwargs["s"],
+        kwargs["dtype"],
+    )
+    for _, kwargs, _ in configs
+)
+
+_fallback_kernel = ninetoothed.make(
+    arrangement,
+    mm.application,
+    tuple(Tensor(4, shape_options={"constexpr": True}) for _ in range(3)),
+)
 
 
-shape_options = {"constexpr": True}
-tensors = tuple(Tensor(4, shape_options=shape_options) for _ in range(3))
+def kernel(input, filter, output, n, c, h, w, k, r, s, dtype):
+    if (n, c, h, w, k, r, s, dtype) in _BUILD_CONFIGS:
+        return _build_kernel(input, filter, output, n, c, h, w, k, r, s, dtype)
 
-kernel = ninetoothed.make(arrangement, mm.application, tensors)
+    return _fallback_kernel(input, filter, output)

ops/ninetoothed/kernels/max_pool2d.py

@@ -1,15 +1,22 @@
+import functools
+
 import ninetoothed
 import ninetoothed.language as ntl
-from ninetoothed import Symbol, Tensor
+from ninetoothed import Symbol, Tensor, block_size
 
-BLOCK_SIZE = Symbol("BLOCK_SIZE", meta=True)
+from ops.ninetoothed.kernels._common import build
 
-WINDOW_HEIGHT = Symbol("WINDOW_HEIGHT", constexpr=True, upper_bound=16)
-WINDOW_WIDTH = Symbol("WINDOW_WIDTH", constexpr=True, upper_bound=16)
+BLOCK_SIZE = block_size()
 
 
-def arrangement(input, output):
-    input_arranged = input.tile((1, 1, WINDOW_HEIGHT, WINDOW_WIDTH))
+def arrangement(
+    input,
+    output,
+    window_height,
+    window_width,
+    BLOCK_SIZE=BLOCK_SIZE,
+):
+    input_arranged = input.tile((1, 1, window_height, window_width))
     input_arranged = input_arranged.ravel()
     input_arranged = input_arranged.flatten(end_dim=4).flatten(start_dim=1)
     input_arranged = input_arranged.tile((BLOCK_SIZE, -1))
@@ -27,6 +34,76 @@ def application(input, output):
     output = ntl.max(input, axis=1)  # noqa: F841
 
 
-kernel = ninetoothed.make(
-    arrangement, application, (Tensor(4, other=float("-inf")), Tensor(4))
+def premake(window_height, window_width, dtype, block_size_):
+    arrangement_ = functools.partial(
+        arrangement,
+        window_height=window_height,
+        window_width=window_width,
+        BLOCK_SIZE=block_size_,
+    )
+    tensors = (
+        Tensor(4, dtype=dtype, other=float("-inf")),
+        Tensor(4, dtype=dtype),
+    )
+
+    return arrangement_, application, tensors
+
+
+# Compile the benchmarked 3x3 fp16 case; other windows and dtypes use the
+# generic make-based fallback below.
+configs = (
+    (
+        (),
+        {
+            "window_height": 3,
+            "window_width": 3,
+            "dtype": ninetoothed.float16,
+            "block_size_": 256,
+        },
+        {},
+    ),
+)
+
+_build_kernel = build(
+    premake,
+    configs,
+    meta_parameters=("block_size_",),
+    kernel_name="max_pool2d",
+)
+
+_BUILD_CONFIGS = frozenset(
+    (kwargs["window_height"], kwargs["window_width"], kwargs["dtype"])
+    for _, kwargs, _ in configs
+)
+
+_FALLBACK_BLOCK_SIZE = block_size()
+_FALLBACK_WINDOW_HEIGHT = Symbol(
+    "FALLBACK_WINDOW_HEIGHT", constexpr=True, upper_bound=16
 )
+_FALLBACK_WINDOW_WIDTH = Symbol("FALLBACK_WINDOW_WIDTH", constexpr=True, upper_bound=16)
+
+
+def _fallback_arrangement(
+    input,
+    output,
+    BLOCK_SIZE=_FALLBACK_BLOCK_SIZE,
+    WINDOW_HEIGHT=_FALLBACK_WINDOW_HEIGHT,
+    WINDOW_WIDTH=_FALLBACK_WINDOW_WIDTH,
+):
+    return arrangement(input, output, WINDOW_HEIGHT, WINDOW_WIDTH, BLOCK_SIZE)
+
+
+_fallback_kernel = ninetoothed.make(
+    _fallback_arrangement,
+    application,
+    (Tensor(4, other=float("-inf")), Tensor(4)),
+)
+
+
+def kernel(input, output, window_height, window_width, dtype):
+    if (window_height, window_width, dtype) in _BUILD_CONFIGS:
+        return _build_kernel(input, output, window_height, window_width, dtype)
+
+    return _fallback_kernel(
+        input, output, WINDOW_HEIGHT=window_height, WINDOW_WIDTH=window_width
+    )