ninetoothed 0.10.0__tar.gz → 0.11.1__tar.gz

ninetoothed-0.11.1/.gitattributes

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

ninetoothed-0.11.1/.github/workflows/publish-to-pypi.yml

@@ -0,0 +1,92 @@
+name: Publish Python 🐍 distribution 📦 to PyPI
+
+on: push
+
+jobs:
+  build:
+    name: Build distribution 📦
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.x"
+    - name: Install pypa/build
+      run: >-
+        python3 -m
+        pip install
+        build
+        --user
+    - name: Build a binary wheel and a source tarball
+      run: python3 -m build
+    - name: Store the distribution packages
+      uses: actions/upload-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+
+  publish-to-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    if: startsWith(github.ref, 'refs/tags/')
+    needs:
+    - build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/ninetoothed
+    permissions:
+      id-token: write
+
+    steps:
+    - name: Download all the dists
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Publish distribution 📦 to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+
+  github-release:
+    name: >-
+      Sign the Python 🐍 distribution 📦 with Sigstore
+      and upload them to GitHub Release
+    needs:
+    - publish-to-pypi
+    runs-on: ubuntu-latest
+
+    permissions:
+      contents: write
+      id-token: write
+
+    steps:
+    - name: Download all the dists
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Sign the dists with Sigstore
+      uses: sigstore/gh-action-sigstore-python@v3.0.0
+      with:
+        inputs: >-
+          ./dist/*.tar.gz
+          ./dist/*.whl
+    - name: Create GitHub Release
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      run: >-
+        gh release create
+        "$GITHUB_REF_NAME"
+        --repo "$GITHUB_REPOSITORY"
+        --notes ""
+    - name: Upload artifact signatures to GitHub Release
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      run: >-
+        gh release upload
+        "$GITHUB_REF_NAME" dist/**
+        --repo "$GITHUB_REPOSITORY"

ninetoothed-0.11.1/.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.10.0 → ninetoothed-0.11.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.10.0
+Version: 0.11.1
 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

ninetoothed-0.11.1/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.11.1/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.11.1/docs/requirements.txt

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

ninetoothed-0.11.1/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.11.1/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.11.1/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.11.1/docs/source/code_generation.rst

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

ninetoothed-0.11.1/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.11.1/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.11.1/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.11.1/docs/source/python_api.rst

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

ninetoothed-0.11.1/docs/source/symbol.rst

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

ninetoothed-0.11.1/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.10.0 → ninetoothed-0.11.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ninetoothed"
-version = "0.10.0"
+version = "0.11.1"
 authors = [{ name = "Jiacheng Huang", email = "huangjiacheng0709@outlook.com" }]
 description = "A domain-specific language based on Triton but providing higher-level abstraction."
 readme = "README.md"

{ninetoothed-0.10.0 → ninetoothed-0.11.1}/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.10.0 → ninetoothed-0.11.1}/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.10.0 → ninetoothed-0.11.1}/src/ninetoothed/tensor.py

@@ -8,6 +8,21 @@ 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 +85,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]
 
@@ -119,6 +142,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=[
@@ -132,9 +161,16 @@ class Tensor:
             ],
             source=self.source,
             source_dims=self.source_dims,
+            target_dims=self.target_dims,
         )
 
     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,)
 
@@ -149,9 +185,20 @@ class Tensor:
                 for i, source_dim in enumerate(self.source_dims)
                 if i not in dim
             ],
+            target_dims=[
+                target_dim
+                for i, target_dim in enumerate(self.target_dims)
+                if i not in dim
+            ],
         )
 
     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)]
@@ -168,9 +215,20 @@ class Tensor:
             strides=new_strides,
             source=self.source,
             source_dims=new_source_dims,
+            target_dims=self.target_dims,
         )
 
     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
@@ -197,15 +255,36 @@ class Tensor:
             leading_source_dims + (flattening_source_dims,) + trailing_source_dims
         )
 
+        leading_target_dims = self.target_dims[:start_dim]
+        flattening_target_dims = self.target_dims[start_dim:end_dim]
+        trailing_target_dims = self.target_dims[end_dim:]
+
+        new_target_dims = (
+            leading_target_dims + (flattening_target_dims[-1],) + trailing_target_dims
+        )
+
         return type(self)(
             shape=new_shape,
             dtype=self.dtype,
             strides=new_strides,
             source=self.source,
             source_dims=new_source_dims,
+            target_dims=new_target_dims,
         )
 
     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.11.1/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.11.1/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.10.0 → ninetoothed-0.11.1}/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