ninetoothed 0.11.0__tar.gz → 0.12.0__tar.gz

ninetoothed-0.12.0/.gitattributes

@@ -0,0 +1 @@
+*.png filter=lfs diff=lfs merge=lfs -text

ninetoothed-0.12.0/.github/workflows/sphinx.yml

@@ -0,0 +1,37 @@
+name: "Sphinx: Render docs"
+
+on: push
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
+        lfs: true
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install .
+        pip install -r docs/requirements.txt
+    - name: Build HTML
+      run: make -C docs html
+    - name: Upload artifacts
+      uses: actions/upload-artifact@v4
+      with:
+        name: html-docs
+        path: docs/build/html/
+    - name: Deploy
+      uses: peaceiris/actions-gh-pages@v3
+      if: github.ref == 'refs/heads/master'
+      with:
+        github_token: ${{ secrets.GITHUB_TOKEN }}
+        publish_dir: docs/build/html
+        cname: ninetoothed.org

{ninetoothed-0.11.0 → ninetoothed-0.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.11.0
+Version: 0.12.0
 Summary: A domain-specific language based on Triton but providing higher-level abstraction.
 Project-URL: Homepage, https://github.com/InfiniTensor/ninetoothed
 Project-URL: Issues, https://github.com/InfiniTensor/ninetoothed/issues
@@ -11,6 +11,9 @@ Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Requires-Dist: triton>=3.0.0
+Provides-Extra: visualization
+Requires-Dist: matplotlib>=3.9.0; extra == 'visualization'
+Requires-Dist: numpy>=2.1.0; extra == 'visualization'
 Description-Content-Type: text/markdown
 
 # NineToothed

ninetoothed-0.12.0/docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

ninetoothed-0.12.0/docs/make.bat

@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd

ninetoothed-0.12.0/docs/requirements.txt

@@ -0,0 +1,2 @@
+sphinx
+pydata-sphinx-theme

ninetoothed-0.12.0/docs/source/_static/matmul-tiling.png

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4436ac3d9a7b1362efdbd33cd14ff9be34ceef89d596d4554a4727dd0ef57d5
+size 224240

ninetoothed-0.12.0/docs/source/_static/ninetoothed-logo.png

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2eb81c9d72a5f53d3409a8f56f5bf8ea85145ac788930743a5449692097d4537
+size 17335

ninetoothed-0.12.0/docs/source/_static/vecadd-tiling.png

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2d6725fb932a4d96d1c1e7a4c730e2685d00c94d8677f63512d98182de638a1
+size 44416

ninetoothed-0.12.0/docs/source/code_generation.rst

@@ -0,0 +1,9 @@
+Code Generation
+===============
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+
+   ninetoothed.jit
+   ninetoothed.make

ninetoothed-0.12.0/docs/source/conf.py

@@ -0,0 +1,28 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# For the full list of built-in configuration values, see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+
+project = "NineToothed"
+copyright = "2024, NineToothed Contributors"
+author = "NineToothed Contributors"
+
+# -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
+
+extensions = ["sphinx.ext.autodoc", "sphinx.ext.autosummary"]
+
+templates_path = ["_templates"]
+exclude_patterns = []
+
+
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_theme = "pydata_sphinx_theme"
+html_static_path = ["_static"]
+html_title = "NineToothed"
+html_logo = "_static/ninetoothed-logo.png"

ninetoothed-0.12.0/docs/source/index.rst

@@ -0,0 +1,14 @@
+NineToothed Documentation
+=========================
+
+**NineToothed** is a domain-specific language (DSL) based on Triton, offering higher-level abstractions. Through its tensor-oriented metaprogramming (TOM) model, it empowers developers to write high-performance compute kernels intuitively, without the need to manage low-level details like pointer arithmetic or memory access.
+
+.. note::
+
+   This project is under active development.
+
+.. toctree::
+   :maxdepth: 2
+
+   installation
+   python_api

ninetoothed-0.12.0/docs/source/installation.rst

@@ -0,0 +1,12 @@
+Installation
+============
+
+You can install NineToothed using ``pip``:
+
+.. code-block::
+
+    pip install ninetoothed
+
+To fully leverage its capabilities, you will also need to install a compatible deep learning framework. Currently, NineToothed supports `PyTorch <https://pytorch.org/>`_.
+
+It is generally considered good practice to use a virtual environment when installing packages with pip, though it is optional. You may find this `documentation <https://docs.python.org/3/library/venv.html>`_ helpful.

ninetoothed-0.12.0/docs/source/python_api.rst

@@ -0,0 +1,9 @@
+Python API
+==========
+
+.. toctree::
+   :maxdepth: 1
+
+   code_generation
+   tensor
+   symbol

ninetoothed-0.12.0/docs/source/symbol.rst

@@ -0,0 +1,4 @@
+Symbol
+======
+
+.. autoclass:: ninetoothed.Symbol

ninetoothed-0.12.0/docs/source/tensor.rst

@@ -0,0 +1,18 @@
+Tensor
+======
+
+.. autoclass:: ninetoothed.Tensor
+
+Meta-Operations
+---------------
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+
+   ninetoothed.Tensor.tile
+   ninetoothed.Tensor.expand
+   ninetoothed.Tensor.squeeze
+   ninetoothed.Tensor.permute
+   ninetoothed.Tensor.flatten
+   ninetoothed.Tensor.ravel

{ninetoothed-0.11.0 → ninetoothed-0.12.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ninetoothed"
-version = "0.11.0"
+version = "0.12.0"
 authors = [{ name = "Jiacheng Huang", email = "huangjiacheng0709@outlook.com" }]
 description = "A domain-specific language based on Triton but providing higher-level abstraction."
 readme = "README.md"
@@ -20,6 +20,9 @@ classifiers = [
 Homepage = "https://github.com/InfiniTensor/ninetoothed"
 Issues = "https://github.com/InfiniTensor/ninetoothed/issues"
 
+[project.optional-dependencies]
+visualization = ["matplotlib>=3.9.0", "numpy>=2.1.0"]
+
 [tool.ruff]
 src = [".", "src", "tests"]
 

{ninetoothed-0.11.0 → ninetoothed-0.12.0}/src/ninetoothed/jit.py

@@ -20,6 +20,13 @@ from ninetoothed.torchifier import Torchifier
 
 
 def make(arrangement, application, tensors):
+    """Integrate the arrangement and the application of the tensors.
+
+    :param arrangement: The arrangement of the tensors.
+    :param application: The application of the tensors.
+    :param tensors: The tensors.
+    :return: A handle to the compute kernel.
+    """
     params = inspect.signature(application).parameters
     types = arrangement(*tensors)
     annotations = {param: type for param, type in zip(params, types)}
@@ -28,14 +35,26 @@ def make(arrangement, application, tensors):
     return jit(application)
 
 
-def jit(_func=None, *, _prettify=False):
+def jit(func=None, *, _prettify=False):
+    """A decorator for generating compute kernels.
+
+    :param func: The function to be compiled.
+    :param _prettify: Whether to prettify the generated code.
+    :return: A handle to the compute kernel.
+
+    .. note::
+
+        The ``_prettify`` parameter is experimental, which might break
+        the generated code.
+    """
+
     def wrapper(func):
         return JIT(func, _prettify=_prettify)()
 
-    if _func is None:
+    if func is None:
         return wrapper
 
-    return wrapper(_func)
+    return wrapper(func)
 
 
 class JIT:
@@ -472,6 +491,15 @@ class CodeGenerator(ast.NodeTransformer):
                 for target_dim in range(tensor.target.ndim)
                 if offsets[source_dim][target_dim] != 0
             ),
+        ) & functools.reduce(
+            lambda x, y: x & y,
+            (
+                indices[dim - tensor.innermost().target.ndim][
+                    type(self)._generate_slices(tensor, target_dim)
+                ]
+                < tensor.innermost().target.shape[dim]
+                for dim, target_dim in enumerate(tensor.innermost().target_dims)
+            ),
         )
 
         return pointers, mask

{ninetoothed-0.11.0 → ninetoothed-0.12.0}/src/ninetoothed/symbol.py

@@ -7,6 +7,13 @@ import ninetoothed.naming as naming
 
 
 class Symbol:
+    """A class uesed to represent a symbol.
+
+    :param expr: The expression used to construct the symbol.
+    :param constexpr: Whether the symbol is a constexpr.
+    :param mata: Whether the symbol is a meta.
+    """
+
     def __init__(self, expr, constexpr=None, meta=None):
         if isinstance(expr, type(self)):
             self._node = expr._node

{ninetoothed-0.11.0 → ninetoothed-0.12.0}/src/ninetoothed/tensor.py

@@ -3,11 +3,25 @@ import math
 import re
 
 import ninetoothed.naming as naming
-from ninetoothed.language import call
 from ninetoothed.symbol import Symbol
 
 
 class Tensor:
+    """A class uesed to represent a symbolic tensor.
+
+    :param ndim: The number of dimensions of the tensor.
+    :param shape: The shape of the tensor.
+    :param dtype: The element type of the tensor.
+    :param strides: The strides of the tensor.
+    :param other: The values for out-of-bounds positions.
+    :param constexpr_shape: Whether the sizes are constexpr.
+    :param name: The name of the tensor.
+    :param source: For internal use only.
+    :param source_dims: For internal use only.
+    :param target: For internal use only.
+    :param target_dims: For internal use only.
+    """
+
     num_instances = 0
 
     def __init__(
@@ -70,6 +84,14 @@ class Tensor:
         type(self).num_instances += 1
 
     def tile(self, tile_shape, strides=None, dilation=None):
+        """Tiles the tensor into a hierarchical tensor.
+
+        :param tile_shape: The shape of a tile.
+        :param strides: The interval at which each tile is generated.
+        :param dilation: The spacing between tiles.
+        :return: A hierarchical tensor.
+        """
+
         if strides is None:
             strides = [-1 for _ in tile_shape]
 
@@ -90,8 +112,11 @@ class Tensor:
             if stride == -1:
                 stride = tile_size
 
+            def cdiv(x, y):
+                return (x + y - 1) // y
+
             new_size = (
-                call("cdiv", self_size - spacing * (tile_size - 1) - 1, stride) + 1
+                (cdiv(self_size - spacing * (tile_size - 1) - 1, stride) + 1)
                 if stride != 0
                 else -1
             )
@@ -119,6 +144,12 @@ class Tensor:
         )
 
     def expand(self, shape):
+        """Expands the specified singleton dimensions of the tensor.
+
+        :param shape: The expanded shape.
+        :return: The expanded tensor.
+        """
+
         # TODO: Add error handling.
         return type(self)(
             shape=[
@@ -136,6 +167,12 @@ class Tensor:
         )
 
     def squeeze(self, dim):
+        """Removes the specified singleton dimensions of the tensor.
+
+        :param dim: The dimension(s) to be squeezed.
+        :return: The squeezed tensor.
+        """
+
         if not isinstance(dim, tuple):
             dim = (dim,)
 
@@ -158,6 +195,12 @@ class Tensor:
         )
 
     def permute(self, dims):
+        """Permutes the dimensions of the tensor.
+
+        :param dims: The permuted ordering of the dimensions.
+        :return: The permuted tensor.
+        """
+
         # TODO: Add error handling.
         new_shape = [None for _ in range(self.ndim)]
         new_strides = [None for _ in range(self.ndim)]
@@ -178,6 +221,16 @@ class Tensor:
         )
 
     def flatten(self, start_dim=None, end_dim=None):
+        """Flattens the specified dimensions of the tensor.
+
+        See :func:`ravel` for the differences between :func:`flatten`
+        and :func:`ravel`.
+
+        :param start_dim: The first dimension to flatten.
+        :param end_dim: The dimension after the last to flatten.
+        :return: The flattened tensor.
+        """
+
         # TODO: Add error handling.
         if start_dim is None:
             start_dim = 0
@@ -222,6 +275,18 @@ class Tensor:
         )
 
     def ravel(self):
+        """Flattens the hierarchy of the tensor.
+
+        :func:`ravel` differs from :func:`flatten`, which only flattens
+        dimensions at a single level. For example, consider a tensor
+        with two levels: the first level has a shape of ``(N, P, Q)``,
+        and the second level has a shape of ``(C, R, S)``. After
+        applying :func:`ravel`, the resulting tensor will have a single
+        flattened level with a shape of ``(N, P, Q, C, R, S)``.
+
+        :return: The raveled tensor.
+        """
+
         # TODO: Add error handling.
         new_shape = []
         new_strides = []

ninetoothed-0.12.0/src/ninetoothed/visualization.py

@@ -0,0 +1,122 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from mpl_toolkits.axes_grid1 import Divider, Size
+
+
+def visualize(tensor, color=None, save_path=None):
+    outline_width = 0.1
+    plt.rcParams["lines.linewidth"] = 72 * outline_width
+
+    if color is None:
+        color = f"C{visualize.count}"
+
+    _, max_pos_x, max_pos_y = _visualize_tensor(plt.gca(), tensor, 0, 0, color)
+
+    width = max_pos_y + 1
+    height = max_pos_x + 1
+
+    fig = plt.figure(figsize=(width + outline_width, height + outline_width))
+
+    h = (Size.Fixed(0), Size.Fixed(width + outline_width))
+    v = (Size.Fixed(0), Size.Fixed(height + outline_width))
+
+    divider = Divider(fig, (0, 0, 1, 1), h, v, aspect=False)
+
+    ax = fig.add_axes(
+        divider.get_position(), axes_locator=divider.new_locator(nx=1, ny=1)
+    )
+
+    ax.set_aspect("equal")
+    ax.invert_yaxis()
+
+    plt.axis("off")
+
+    half_outline_width = outline_width / 2
+    plt.xlim((-half_outline_width, width + half_outline_width))
+    plt.ylim((-half_outline_width, height + half_outline_width))
+
+    _visualize_tensor(ax, tensor, 0, 0, color)
+
+    plt.savefig(save_path, transparent=True, bbox_inches="tight", pad_inches=0)
+
+    plt.close()
+
+    visualize.count += 1
+
+
+visualize.count = 0
+
+
+def _visualize_tensor(ax, tensor, x, y, color, level_spacing=4):
+    verts = _visualize_level(ax, tensor, x, y, color)
+
+    if tensor.dtype is None:
+        return verts, verts[1][1][0], verts[1][1][1]
+
+    next_x, next_y = verts[0][1]
+    next_y += level_spacing + 1
+
+    next_verts, max_pos_x, max_pos_y = _visualize_tensor(
+        ax, tensor.dtype, next_x, next_y, color
+    )
+
+    conn_verts = _verts_of_rect(1, level_spacing, next_x, next_y - level_spacing)
+    conn_verts = [list(vert) for vert in conn_verts]
+    conn_verts[2][0] += next_verts[1][0][0]
+
+    pos_y, pos_x = zip(*conn_verts)
+    pos_x = pos_x + (pos_x[0],)
+    pos_y = pos_y + (pos_y[0],)
+
+    ax.plot(pos_x[1:3], pos_y[1:3], "k--")
+    ax.plot(pos_x[3:5], pos_y[3:5], "k--")
+
+    max_pos_x = max(max_pos_x, verts[1][1][0])
+    max_pos_y = max(max_pos_y, verts[1][1][1])
+
+    return verts, max_pos_x, max_pos_y
+
+
+def _visualize_level(ax, level, x, y, color):
+    offsets = [1 for _ in range(level.ndim)]
+
+    for dim in range(-3, -level.ndim - 1, -1):
+        offsets[dim] = offsets[dim + 2] * level.shape[dim + 2] + 1
+
+    indices = np.indices(level.shape)
+    flattened_indices = np.stack(
+        [indices[i].flatten() for i in range(level.ndim)], axis=-1
+    )
+
+    max_pos_x = x
+    max_pos_y = y
+
+    for indices in flattened_indices:
+        pos = [x, y]
+
+        for dim, index in enumerate(indices):
+            pos[(level.ndim - dim) % 2] += index * offsets[dim]
+
+        max_pos_x = max(max_pos_x, pos[0])
+        max_pos_y = max(max_pos_y, pos[1])
+
+        _visualize_unit_square(ax, pos[1], pos[0], color)
+
+    verts = (((x, y), (x, max_pos_y)), ((max_pos_x, y), (max_pos_x, max_pos_y)))
+
+    return verts
+
+
+def _visualize_unit_square(ax, x, y, color):
+    _visualize_rect(ax, 1, 1, x, y, color)
+
+
+def _visualize_rect(ax, width, height, x, y, color):
+    pos_x, pos_y = zip(*_verts_of_rect(width, height, x, y))
+
+    ax.fill(pos_x, pos_y, color)
+    ax.plot(pos_x + (pos_x[0],), pos_y + (pos_y[0],), "k")
+
+
+def _verts_of_rect(width, height, x, y):
+    return ((x, y), (x + width, y), (x + width, y + height), (x, y + height))

ninetoothed-0.12.0/tests/test_attention.py

@@ -0,0 +1,92 @@
+import torch
+import torch.nn.functional as F
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Symbol, Tensor
+from tests.skippers import skip_if_cuda_not_available
+
+
+def arrangement(q, k, v, o):
+    BLOCK_SIZE_M = Symbol("BLOCK_SIZE_M", constexpr=True)
+    BLOCK_SIZE_N = Symbol("BLOCK_SIZE_N", constexpr=True)
+
+    def arrange_q_or_o(input):
+        arranged = input.tile((1, 1, BLOCK_SIZE_M, -1))
+        arranged.dtype = arranged.dtype.squeeze((0, 1))
+
+        return arranged
+
+    def arrange_k_or_v(input):
+        arranged = (
+            input.tile((1, 1, BLOCK_SIZE_N, -1))
+            .tile((1, 1, -1, -1))
+            .expand((-1, -1, q_arranged.shape[-2], -1))
+        )
+        arranged.dtype = arranged.dtype.squeeze((0, 1, 3))
+        arranged.dtype.dtype = arranged.dtype.dtype.squeeze((0, 1))
+
+        return arranged
+
+    q_arranged = arrange_q_or_o(q)
+
+    return q_arranged, arrange_k_or_v(k), arrange_k_or_v(v), arrange_q_or_o(o)
+
+
+def application(q, k, v, o):
+    q_loaded = (q * 1.44269504089).to(ntl.float16)
+
+    acc = ntl.zeros((q.shape[-2], q.shape[-1]), dtype=ntl.float32)
+    l_i = ntl.full((q.shape[-2],), 1, dtype=ntl.float32)
+    m_i = ntl.full((q.shape[-2],), float("-inf"), dtype=ntl.float32)
+
+    for i in range(k.shape[0]):
+        qk = ntl.dot(q_loaded, ntl.trans(k[i]))
+
+        m_ij = ntl.maximum(m_i, ntl.max(qk, 1))
+        p = ntl.exp2(qk - m_ij[:, None])
+        l_ij = ntl.sum(p, 1)
+
+        alpha = ntl.exp2(m_i - m_ij)
+        acc = acc * alpha[:, None] + ntl.dot(p.to(ntl.float16), v[i])
+        m_i = m_ij
+        l_i = l_i * alpha + l_ij
+
+    acc /= l_i[:, None]
+    o = acc  # noqa: F841
+
+
+def attention(q, k, v):
+    o = torch.empty_like(q, dtype=v.dtype)
+
+    attention_kernel = ninetoothed.make(
+        arrangement, application, (Tensor(4, constexpr_shape=True) for _ in range(4))
+    )
+
+    attention_kernel(q, k, v, o, BLOCK_SIZE_M=128, BLOCK_SIZE_N=64)
+
+    return o
+
+
+@skip_if_cuda_not_available
+class TestCUDA:
+    @classmethod
+    def setup_class(cls):
+        torch.manual_seed(0)
+
+        shape = (2, 4, 1024, 64)
+
+        cls.q = torch.randn(shape, device="cuda")
+        cls.k = torch.randn(shape, device="cuda")
+        cls.v = torch.randn(shape, device="cuda")
+
+    def test_fp16(self):
+        q = type(self).q.to(torch.float16)
+        k = type(self).k.to(torch.float16)
+        v = type(self).v.to(torch.float16)
+
+        assert torch.allclose(
+            attention(q, k, v),
+            F.scaled_dot_product_attention(q, k, v, scale=1),
+            atol=0.01,
+        )

ninetoothed-0.12.0/tests/test_conv2d.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn.functional as F
+
+import ninetoothed
+import tests.test_matmul as matmul
+from ninetoothed import Tensor
+from tests.skippers import skip_if_cuda_not_available
+
+
+def arrangement(input, filter, output):
+    input_tiled = input.tile((1, *filter.shape[1:]), strides=(-1, -1, 1, 1))
+    input_squeezed = input_tiled.squeeze(1)
+    input_squeezed.dtype = input_squeezed.dtype.squeeze(0)
+    input_raveled = input_squeezed.ravel()
+    input_flattened = input_raveled.flatten(end_dim=3).flatten(start_dim=1)
+
+    filter_flattened = filter.flatten(start_dim=1)
+    filter_permuted = filter_flattened.permute((1, 0))
+
+    output_flattened = output.permute((0, 2, 3, 1)).flatten(end_dim=3)
+
+    return matmul.arrangement(input_flattened, filter_permuted, output_flattened)
+
+
+def conv2d(input, filter):
+    n, _, h, w = input.shape
+    k, _, r, s = filter.shape
+    p = h - r + 1
+    q = w - s + 1
+
+    output = torch.empty((n, k, p, q), device=input.device, dtype=input.dtype)
+
+    conv2d_kernel = ninetoothed.make(
+        arrangement,
+        matmul.application,
+        (Tensor(4), Tensor(4, constexpr_shape=True), Tensor(4)),
+    )
+
+    conv2d_kernel(input, filter, output)
+
+    return output
+
+
+@skip_if_cuda_not_available
+class TestCUDA:
+    @classmethod
+    def setup_class(cls):
+        torch.manual_seed(0)
+
+        n, c, h, w = 4, 64, 16, 16
+        k, _, r, s = 512, c, 3, 3
+
+        cls.input = torch.randn(n, c, h, w, device="cuda")
+        cls.filter = torch.randn(k, c, r, s, device="cuda")
+
+    def test_fp16(self):
+        input = type(self).input.to(torch.float16)
+        filter = type(self).filter.to(torch.float16)
+
+        assert torch.allclose(
+            conv2d(input, filter), F.conv2d(input, filter), atol=0.001, rtol=0.001
+        )

{ninetoothed-0.11.0 → ninetoothed-0.12.0}/tests/test_matmul.py

@@ -6,40 +6,46 @@ from ninetoothed import Symbol, Tensor
 from tests.skippers import skip_if_cuda_not_available, skip_if_float8_e5m2_not_supported
 
 
-def matmul(lhs, rhs):
+def arrangement(lhs, rhs, output):
     BLOCK_SIZE_M = Symbol("BLOCK_SIZE_M", meta=True)
     BLOCK_SIZE_N = Symbol("BLOCK_SIZE_N", meta=True)
     BLOCK_SIZE_K = Symbol("BLOCK_SIZE_K", meta=True)
 
-    output_tiled = Tensor(2).tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
+    output_tiled = output.tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
 
     lhs_tiled = (
-        Tensor(2)
-        .tile((BLOCK_SIZE_M, BLOCK_SIZE_K))
+        lhs.tile((BLOCK_SIZE_M, BLOCK_SIZE_K))
         .tile((1, -1))
         .expand((-1, output_tiled.shape[1]))
     )
     lhs_tiled.dtype = lhs_tiled.dtype.squeeze(0)
 
     rhs_tiled = (
-        Tensor(2)
-        .tile((BLOCK_SIZE_K, BLOCK_SIZE_N))
+        rhs.tile((BLOCK_SIZE_K, BLOCK_SIZE_N))
         .tile((-1, 1))
         .expand((output_tiled.shape[0], -1))
     )
     rhs_tiled.dtype = rhs_tiled.dtype.squeeze(1)
 
-    @ninetoothed.jit
-    def matmul_kernel(lhs: lhs_tiled, rhs: rhs_tiled, output: output_tiled):
-        accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
-        for k in range(lhs.shape[0]):
-            accumulator += ntl.dot(lhs[k], rhs[k])
-        output = accumulator.to(ntl.float16)
+    return lhs_tiled, rhs_tiled, output_tiled
+
+
+def application(lhs, rhs, output):
+    accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
+    for k in range(lhs.shape[0]):
+        accumulator += ntl.dot(lhs[k], rhs[k])
+    output = accumulator.to(ntl.float16)
+
 
+def matmul(lhs, rhs):
     output = torch.empty(
         (lhs.shape[0], rhs.shape[1]), device=lhs.device, dtype=torch.float16
     )
 
+    matmul_kernel = ninetoothed.make(
+        arrangement, application, (Tensor(2), Tensor(2), Tensor(2))
+    )
+
     matmul_kernel(lhs, rhs, output)
 
     return output