ninetoothed 0.16.0__tar.gz → 0.17.0__tar.gz

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/.github/workflows/pytest.yml

@@ -12,6 +12,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
+          pip install -e .[all]
           pip install -r requirements.txt
       - name: Test with pytest
         run: |

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.16.0
+Version: 0.17.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

ninetoothed-0.17.0/docs/source/basics.rst

@@ -0,0 +1,254 @@
+The Basics
+==========
+
+Symbols
+-------
+
+The concept of **symbols** is similar to what is described in the `SymPy tutorial <https://docs.sympy.org/latest/tutorials/intro-tutorial/intro.html>`_. Symbols do not store actual numerical values; instead, they represent symbolic names or symbolic expressions. This allows for performing symbolic mathematical operations.
+
+In NineToothed, you can create a symbol using the ``Symbol`` class. For example, in the code below, we first create two symbols named ``BLOCK_SIZE_M`` and ``BLOCK_SIZE_N``, and then perform a multiplication operation on them:
+
+.. code-block::
+
+    >>> from ninetoothed import Symbol
+    >>> BLOCK_SIZE_M = Symbol("BLOCK_SIZE_M")
+    >>> BLOCK_SIZE_M
+    BLOCK_SIZE_M
+    >>> BLOCK_SIZE_N = Symbol("BLOCK_SIZE_N")
+    >>> BLOCK_SIZE_N
+    BLOCK_SIZE_N
+    >>> BLOCK_SIZE_M * BLOCK_SIZE_N
+    BLOCK_SIZE_M * BLOCK_SIZE_N
+
+Symbolic Tensors
+----------------
+
+Similar to many deep learning frameworks, tensors are a core concept in NineToothed. However, tensors in NineToothed differ slightly from those in other frameworks—they do not store actual data. Instead, they store symbolic expressions in member variables such as ``shape`` and ``strides``. For this reason, we refer to them as **symbolic tensors**.
+
+In NineToothed, you can create a tensor using the ``Tensor`` class. As shown in the example below, ``Tensor(2)`` creates a 2-dimensional tensor—essentially, a matrix. Note that the ``shape`` member contains symbolic expressions rather than concrete values:
+
+.. code-block::
+
+    >>> from ninetoothed import Tensor
+    >>> x = Tensor(2)
+    >>> x.shape
+    (ninetoothed_tensor_0_size_0, ninetoothed_tensor_0_size_1)
+
+Tensor-Oriented Metaprogramming
+-------------------------------
+
+Thanks to symbolic tensors, we can perform certain compile-time operations on tensors in NineToothed. These operations are called **meta-operations**, such as ``tile``, ``expand``, ``squeeze``, ``permute``, and so on.
+
+For example, in the following code, we apply the ``tile`` operation to ``x``, which divides ``x`` into blocks of shape ``(BLOCK_SIZE_M, BLOCK_SIZE_N)``:
+
+.. code-block::
+
+    >>> x_tiled = x.tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
+    >>> x_tiled.shape
+    ((ninetoothed_tensor_0_size_0 - (BLOCK_SIZE_M - 1) - 1 + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + 1, (ninetoothed_tensor_0_size_1 - (BLOCK_SIZE_N - 1) - 1 + BLOCK_SIZE_N - 1) // BLOCK_SIZE_N + 1)
+    >>> x_tiled.dtype.shape
+    (BLOCK_SIZE_M, BLOCK_SIZE_N)
+
+We notice that the ``dtype`` of ``x_tiled`` also has a ``shape`` attribute. This is because tensors in NineToothed can be nested—that is, the elements of a tensor can themselves be tensors.
+
+In other words, during the ``tile`` operation, we create a two-level tensor: Each element of the outer tensor is itself an inner tensor. To make this easier to understand, let's walk through a numerical example:
+
+.. code-block::
+
+    >>> BLOCK_SIZE_M = 2
+    >>> BLOCK_SIZE_N = 2
+    >>> x = Tensor(shape=(4, 8))
+    >>> x_tiled = x.tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
+    >>> x_tiled.shape
+    (2, 4)
+    >>> x_tiled.dtype.shape
+    (2, 2)
+
+As shown in the figure below, we've tiled the original tensor ``x`` of shape ``(4, 8)`` into blocks of shape ``(2, 2)`` (the inner tensors), resulting in a total of ``(2, 4)`` such blocks (the outer tensor):
+
+.. image:: generated/x-tiled.png
+
+Arrange-and-Apply Paradigm
+--------------------------
+
+After introducing meta-operations, you should be able to understand what operations can be performed on individual tensors at compile time. A series of such operations is referred to as **arrangement**. However, this is not enough, as we also need to establish the relationship between multiple parameter tensors.
+
+Such relationships are managed by the NineToothed compiler: The NineToothed compiler launches programs based on the shape of the outermost tensors of the arranged parameter tensors and maps the second outermost tensors to these programs.
+
+We can understand this concept using a simple arrangement function:
+
+.. code-block:: python
+
+    def arrangement(x, y, z, BLOCK_SIZE=ninetoothed.block_size()):
+        return x.tile((BLOCK_SIZE,)), y.tile((BLOCK_SIZE,)), z.tile((BLOCK_SIZE,))
+
+In this function, we apply the ``tile`` operation to the vectors ``x``, ``y``, and ``z`` to divide each vector into blocks of size ``BLOCK_SIZE``. For example, if each vector's length is ``16`` and ``BLOCK_SIZE`` is ``2``, each vector can be divided into ``8`` blocks, with each block having a length of ``2``. The arranged ``x``, ``y``, and ``z`` would then look as follows:
+
+.. image:: generated/x-arranged.png
+
+.. image:: generated/y-arranged.png
+
+.. image:: generated/z-arranged.png
+
+Based on this arrangement, the NineToothed compiler will launch ``8`` programs and map the elements of the outermost tensors of the arranged ``x``, ``y``, and ``z`` (i.e., the second outermost tensors) to these ``8`` programs.
+
+Now that we have these mappings, we can launch the programs accordingly. However, we are still one step away from fully implementing an algorithm, because we have not defined what each program should do. In other words, we need to define how to apply the arranged tensors. In NineToothed, this can be done by defining an **application** function.
+
+For example, to define vector addition, we can create the following application function:
+
+.. code-block:: python
+
+    def application(x, y, z):
+        z = x + y
+
+The logic of the code is simple: It adds ``x`` and ``y`` and stores the result in ``z``. However, it is important to note that the parameters of the application function are the elements of the outermost tensors of the arranged parameter tensors (i.e., the second outermost tensors), not the tensors themselves. That is, based on the above assumptions, ``x``, ``y``, and ``z`` here represent blocks of length ``2``, not the original tensors of length ``16``.
+
+At this point, we have defined both the arrangement and application functions. The remaining task is to integrate them into a compute kernel. In NineToothed, we can use ``ninetoothed.make`` to achieve this:
+
+.. code-block:: python
+
+    kernel = ninetoothed.make(arrangement, application, (Tensor(1), Tensor(1), Tensor(1)))
+
+This code means that we want to arrange three 1-dimensional tensors (vectors) according to the ``arrangement`` function, and apply the arranged tensors using the ``application`` function to form a compute kernel ``kernel``. The paradigm of constructing a compute kernel this way is called the **arrange-and-apply paradigm**.
+
+We can invoke ``kernel`` as follows:
+
+.. code-block:: python
+
+    import torch
+
+    dtype = torch.float16
+    device = "cuda"
+
+    x = torch.tensor((1, 2, 3), dtype=dtype, device=device)
+    y = torch.tensor((4, 5, 6), dtype=dtype, device=device)
+
+    z = torch.empty_like(x)
+    kernel(x, y, z)
+
+    reference = torch.tensor((5, 7, 9), dtype=dtype, device=device)
+    assert torch.allclose(z, reference)
+
+As shown, when we call ``kernel``, we do not provide an actual value for ``BLOCK_SIZE``. This is because when constructing ``BLOCK_SIZE``, we used ``ninetoothed.block_size``, which represents that we want to use the configurations generated by the NineToothed compiler for auto-tuning. If we want to provide a value manually (for example, during debugging), we can directly assign a specific value as follows:
+
+.. code-block:: python
+
+    def arrangement(x, y, z, BLOCK_SIZE=1024):
+        return x.tile((BLOCK_SIZE,)), y.tile((BLOCK_SIZE,)), z.tile((BLOCK_SIZE,))
+
+Indexing and Iteration
+----------------------
+
+Through the vector addition example, we got a brief understanding of how to develop compute kernels using NineToothed. In that example, the parameter tensors were arranged into two-level tensors. However, tensors in NineToothed are not limited to just two levels—they can be three-level or even more. Only the outermost level of an arranged tensor is used to launch programs. In other words, tensors with more than two levels are hierarchical and can be indexed and iterated over within the application function.
+
+Now, let's implement a matrix multiplication kernel to better understand indexing and iteration in NineToothed, as well as deepen our understanding of the arrange-and-apply paradigm.
+
+Before we begin implementation, we first need to understand the algorithm we want to realize. Here's a diagram for the algorithm:
+
+.. image:: generated/tiled-matrix-multiplication.png
+
+In simple terms, we tile three matrices into blocks. For each block in :math:`C`, we need to iterate over the corresponding row of blocks of :math:`A` and the corresponding column of blocks of :math:`B`. Then, for each iteration, we need to perform a small matrix multiplication between the blocks of :math:`A` and :math:`B`, and accumulate the result into the block of :math:`C`.
+
+With this algorithm in mind, let's begin the implementation, starting with the arrangement phase:
+
+.. code-block:: python
+
+    BLOCK_SIZE_M = ninetoothed.block_size()
+    BLOCK_SIZE_N = ninetoothed.block_size()
+    BLOCK_SIZE_K = ninetoothed.block_size()
+
+
+    def arrangement(input, other, output):
+        output_arranged = output.tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
+
+        input_arranged = input.tile((BLOCK_SIZE_M, BLOCK_SIZE_K))
+        input_arranged = input_arranged.tile((1, -1))
+        input_arranged = input_arranged.expand((-1, output_arranged.shape[1]))
+        input_arranged.dtype = input_arranged.dtype.squeeze(0)
+
+        other_arranged = other.tile((BLOCK_SIZE_K, BLOCK_SIZE_N))
+        other_arranged = other_arranged.tile((-1, 1))
+        other_arranged = other_arranged.expand((output_arranged.shape[0], -1))
+        other_arranged.dtype = other_arranged.dtype.squeeze(1)
+
+        return input_arranged, other_arranged, output_arranged
+
+In this code, we first define the symbols ``BLOCK_SIZE_M``, ``BLOCK_SIZE_N``, and ``BLOCK_SIZE_K``, which represent the shapes of the blocks. We then tile ``output`` into blocks of shape ``(BLOCK_SIZE_M, BLOCK_SIZE_N)``, ``input`` into ``(BLOCK_SIZE_M, BLOCK_SIZE_K)``, and ``other`` into ``(BLOCK_SIZE_K, BLOCK_SIZE_N)``:
+
+.. image:: generated/input-arranged-0.png
+
+.. image:: generated/other-arranged-0.png
+
+.. image:: generated/output-arranged-0.png
+
+We notice that simple arrangement is not enough for matrix multiplication. According to the diagram, each block in ``output`` corresponds to a row of blocks in ``input`` and a column of blocks in ``other``. So we need to further tile ``input`` row-wise and ``other`` column-wise:
+
+.. image:: generated/input-arranged-1.png
+
+.. image:: generated/other-arranged-1.png
+
+But we're still not done. Remember how the NineToothed compiler establishes the relationship between multiple parameter tensors?
+
+    The NineToothed compiler launches programs based on the shape of the outermost tensors of the arranged parameter tensors and maps the second outermost tensors to these programs.
+
+Why is this important? Because it implies a crucial rule: The outermost tensors of the arranged parameter tensors must have the same shape.
+
+Currently, the shapes of the outermost tensors of the arranged parameter tensors are ``(4, 1)``, ``(1, 4)``, and ``(4, 4)``—clearly inconsistent. This suggests that the arrangement is incorrect or incomplete. From the diagram, we know we need to align each row of blocks of ``input`` with each column of blocks of ``other``. We can achieve this via ``expand``, horizontally expanding ``input`` and vertically expanding ``other`` to match the shape of ``output``:
+
+.. image:: generated/input-arranged-2.png
+
+.. image:: generated/other-arranged-2.png
+
+Now, the outermost tensors of the arranged parameter tensors have matching shapes. Technically, arrangement is complete and we could proceed to write the application function. However, we notice that the row of blocks of ``input`` and the column of blocks of ``other`` are two-dimensional, and their shapes are of the form ``(1, ...)`` and ``(..., 1)`` respectively. In other words, if we do not perform additional operations, the way to index the row of blocks and the column of blocks would be ``input[0, k]`` and ``other[k, 0]``. If we want to find the range of ``k`` based on ``input``, we would need to use ``input.shape[1]``. But we know that dimensions of size ``1`` can be safely removed here. That's why we add ``squeeze``:
+
+.. image:: generated/input-arranged-3.png
+
+.. image:: generated/other-arranged-3.png
+
+With this, we can now index the row of blocks and the column of blocks with ``input[k]`` and ``other[k]``, and use ``input.shape[0]`` to determine the range of ``k``.
+
+At this point, the arrangement phase is complete. The final arrangement result is:
+
+.. image:: generated/input-arranged-3.png
+
+.. image:: generated/other-arranged-3.png
+
+.. image:: generated/output-arranged-0.png
+
+Now let's look at the application function:
+
+.. code-block:: python
+
+    def application(input, other, output):
+        accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
+
+        for k in range(input.shape[0]):
+            accumulator += ntl.dot(input[k], other[k])
+
+        output = accumulator
+
+Within the function body, we first define an ``accumulator`` to accumulate intermediate results. Then, we iterate over the row of blocks of ``input`` and the column of blocks of ``other``, and accumulate the results of small matrix multiplications into the ``accumulator``. Finally, we write the ``accumulator`` into the corresponding block of ``output``. Since this happens on each block of ``output``, the overall matrix multiplication is completed.
+
+Just like in the vector addition example, after defining ``arrangement`` and ``application``, we can integrate them using ``ninetoothed.make`` to build a kernel:
+
+.. code-block:: python
+
+    kernel = ninetoothed.make(arrangement, application, (Tensor(2), Tensor(2), Tensor(2)))
+
+The kernel can be invoked like this:
+
+.. code-block:: python
+
+    import torch
+
+    dtype = torch.float16
+    device = "cuda"
+
+    input = torch.tensor(((1, 2), (3, 4)), dtype=dtype, device=device)
+    other = torch.tensor(((5, 6), (7, 8)), dtype=dtype, device=device)
+
+    output = torch.empty((input.shape[0], other.shape[1]), dtype=dtype, device=device)
+    kernel(input, other, output)
+
+    reference = torch.tensor(((19, 22), (43, 50)), dtype=dtype, device=device)
+    assert torch.allclose(output, reference)

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/docs/source/conf.py

@@ -3,6 +3,18 @@
 # For the full list of built-in configuration values, see the documentation:
 # https://www.sphinx-doc.org/en/master/usage/configuration.html
 
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("."))
+
+from visualize import (
+    visualize_add,
+    visualize_mm,
+    visualize_tiled_matrix_multiplication,
+    visualize_x_tiled,
+)
+
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
@@ -36,3 +48,13 @@ html_theme_options = {
         }
     ]
 }
+
+os.makedirs("generated", exist_ok=True)
+
+visualize_x_tiled(4, 8, 2, 2)
+
+visualize_add(16, 2)
+
+visualize_tiled_matrix_multiplication(8, 8, 8, 2, 2, 2)
+
+visualize_mm(8, 8, 8, 2, 2, 2)

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/docs/source/index.rst

@@ -11,5 +11,6 @@ NineToothed Documentation
    :maxdepth: 2
 
    installation
+   basics
    python_api
    visualization

ninetoothed-0.17.0/docs/source/visualize.py

@@ -0,0 +1,189 @@
+import matplotlib.pyplot as plt
+
+from ninetoothed import Tensor
+from ninetoothed.visualization import (
+    _prepare_figure_and_axes,
+    _visualize_tensor,
+    visualize,
+)
+
+
+def visualize_x_tiled(m, n, block_size_m, block_size_n):
+    BLOCK_SIZE_M = block_size_m
+    BLOCK_SIZE_N = block_size_n
+
+    x = Tensor(shape=(m, n))
+    x_tiled = x.tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
+
+    visualize(x_tiled, color="C0", save_path="generated/x-tiled.png")
+
+
+def visualize_add(size, block_size):
+    BLOCK_SIZE = block_size
+
+    x = Tensor(shape=(size,))
+    y = Tensor(shape=(size,))
+    z = Tensor(shape=(size,))
+
+    x_arranged = x.tile((BLOCK_SIZE,))
+    visualize(x_arranged, color="C0", save_path="generated/x-arranged.png")
+
+    y_arranged = y.tile((BLOCK_SIZE,))
+    visualize(y_arranged, color="C1", save_path="generated/y-arranged.png")
+
+    z_arranged = z.tile((BLOCK_SIZE,))
+    visualize(z_arranged, color="C2", save_path="generated/z-arranged.png")
+
+
+def visualize_tiled_matrix_multiplication(
+    m, n, k, block_size_m, block_size_n, block_size_k
+):
+    BLOCK_SIZE_M = block_size_m
+    BLOCK_SIZE_N = block_size_n
+    BLOCK_SIZE_K = block_size_k
+
+    a = Tensor(shape=(m, k))
+    b = Tensor(shape=(k, n))
+    c = Tensor(shape=(m, n))
+
+    a_tiled = a.tile((BLOCK_SIZE_M, BLOCK_SIZE_K))
+    b_tiled = b.tile((BLOCK_SIZE_K, BLOCK_SIZE_N))
+    c_tiled = c.tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
+
+    a_tile = a_tiled.innermost()
+    b_tile = b_tiled.innermost()
+    c_tile = c_tiled.innermost()
+
+    color_0 = "#1f77b4"
+    color_1 = "#ff7f0e"
+    color_2 = "#2ca02c"
+
+    a_color = _lighten_color(color_0, 20)
+    b_color = _lighten_color(color_1, 20)
+    c_color = _lighten_color(color_2, 20)
+
+    a_tile_base_color = color_0
+    b_tile_base_color = color_1
+    c_tile_base_color = color_2
+
+    def _visualize_matrices(ax):
+        a_verts, a_max_pos_x, a_max_pos_y = _visualize_tensor(ax, a, 0, 0, a_color)
+        b_verts, b_max_pos_x, b_max_pos_y = _visualize_tensor(
+            ax, b, a_max_pos_x + 2, a_max_pos_y + 2, b_color
+        )
+        c_verts, _, _ = _visualize_tensor(ax, c, 0, a_max_pos_y + 2, c_color)
+
+        a_min_pos_x, a_min_pos_y = a_verts[0][0]
+        b_min_pos_x, b_min_pos_y = b_verts[0][0]
+        c_min_pos_x, c_min_pos_y = c_verts[0][0]
+
+        percentage = 30
+
+        for i in range(0, a_tiled.shape[1]):
+            y_offset = i * a_tile.shape[1]
+
+            x = a_min_pos_x + (a_tiled.shape[0] - 2) * a_tile.shape[0]
+            y = a_min_pos_y + y_offset
+
+            darkened_color = _darken_color(
+                a_tile_base_color, (i + 1) / a_tiled.shape[1] * percentage
+            )
+
+            _visualize_tensor(ax, a_tile, x, y, darkened_color)
+
+        for i in range(0, b_tiled.shape[0]):
+            x_offset = i * b_tile.shape[0]
+
+            x = b_min_pos_x + x_offset
+            y = b_min_pos_y + b_tile.shape[1]
+
+            darkened_color = _darken_color(
+                b_tile_base_color,
+                (b_tiled.shape[0] - i) / b_tiled.shape[0] * percentage,
+            )
+
+            _visualize_tensor(ax, b_tile, x, y, darkened_color)
+
+        _visualize_tensor(
+            ax,
+            c_tile,
+            c_min_pos_x + (a_tiled.shape[0] - 2) * a_tile.shape[0],
+            c_min_pos_y + b_tile.shape[1],
+            _darken_color(c_tile_base_color, percentage),
+        )
+
+        return b_max_pos_x, b_max_pos_y
+
+    max_pos_x, max_pos_y = _visualize_matrices(plt.gca())
+
+    width = max_pos_y + 1
+    height = max_pos_x + 1
+
+    _, ax = _prepare_figure_and_axes(width, height)
+
+    _visualize_matrices(ax)
+
+    save_path = "generated/tiled-matrix-multiplication.png"
+    plt.savefig(save_path, transparent=True, bbox_inches="tight", pad_inches=0)
+
+    plt.close()
+
+
+def visualize_mm(m, n, k, block_size_m, block_size_n, block_size_k):
+    BLOCK_SIZE_M = block_size_m
+    BLOCK_SIZE_N = block_size_n
+    BLOCK_SIZE_K = block_size_k
+
+    input = Tensor(shape=(m, k))
+    other = Tensor(shape=(k, n))
+    output = Tensor(shape=(m, n))
+
+    output_arranged = output.tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
+    visualize(output_arranged, color="C2", save_path="generated/output-arranged-0.png")
+
+    input_arranged = input.tile((BLOCK_SIZE_M, BLOCK_SIZE_K))
+    visualize(input_arranged, color="C0", save_path="generated/input-arranged-0.png")
+    input_arranged = input_arranged.tile((1, -1))
+    visualize(input_arranged, color="C0", save_path="generated/input-arranged-1.png")
+    input_arranged = input_arranged.expand((-1, output_arranged.shape[1]))
+    visualize(input_arranged, color="C0", save_path="generated/input-arranged-2.png")
+    input_arranged.dtype = input_arranged.dtype.squeeze(0)
+    visualize(input_arranged, color="C0", save_path="generated/input-arranged-3.png")
+
+    other_arranged = other.tile((BLOCK_SIZE_K, BLOCK_SIZE_N))
+    visualize(other_arranged, color="C1", save_path="generated/other-arranged-0.png")
+    other_arranged = other_arranged.tile((-1, 1))
+    visualize(other_arranged, color="C1", save_path="generated/other-arranged-1.png")
+    other_arranged = other_arranged.expand((output_arranged.shape[0], -1))
+    visualize(other_arranged, color="C1", save_path="generated/other-arranged-2.png")
+    other_arranged.dtype = other_arranged.dtype.squeeze(1)
+    visualize(other_arranged, color="C1", save_path="generated/other-arranged-3.png")
+
+
+def _darken_color(hex_color, percentage):
+    rgb = _hex_to_rgb(hex_color)
+    factor = (100 - percentage) / 100
+    darkened_rgb = tuple(int(c * factor) for c in rgb)
+
+    return _rgb_to_hex(darkened_rgb)
+
+
+def _lighten_color(hex_color, percentage):
+    rgb = _hex_to_rgb(hex_color)
+    factor = percentage / 100
+    lightened_rgb = tuple(int(c + (255 - c) * factor) for c in rgb)
+
+    return _rgb_to_hex(lightened_rgb)
+
+
+def _hex_to_rgb(hex_color):
+    hex_color = hex_color.lstrip("#")
+
+    if len(hex_color) == 3:
+        hex_color = "".join(c * 2 for c in hex_color)
+
+    return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4))
+
+
+def _rgb_to_hex(rgb_color):
+    return "#{:02x}{:02x}{:02x}".format(*rgb_color)

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ninetoothed"
-version = "0.16.0"
+version = "0.17.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"

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/src/ninetoothed/aot.py

@@ -31,7 +31,7 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
 
     _HEADER_PATH.parent.mkdir(exist_ok=True)
 
-    if not _HEADER_PATH.exists():
+    if not _HEADER_PATH.exists() or _HEADER_PATH.read_text() != _HEADER_CONTENT:
         _HEADER_PATH.write_text(_HEADER_CONTENT)
 
     code_generator = CodeGenerator()
@@ -91,20 +91,29 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
 
     c_header_file_name = f"{kernel_name}.{signature_hash}.h"
     c_header_file = output_contents[c_header_file_name]
-    c_header_file = f"{c_header_file}\n{unparser.header};\n"
+    c_header_file = f'{c_header_file}\n#ifdef __cplusplus\nextern "C" {unparser.header};\n#else\n{unparser.header};\n#endif\n'
     c_header_file = c_header_file.replace("<stdint.h>", f'"{_HEADER_PATH}"')
     output_contents[c_header_file_name] = c_header_file
 
     return output_contents
 
 
-_HEADER_CONTENT = """#include <stdint.h>
+_HEADER_CONTENT = """#ifndef NINETOOTHED_H
+#define NINETOOTHED_H
+
+#include <stdint.h>
 
 typedef struct {
-    uintptr_t data;
+    void *data;
     uint64_t *shape;
     int64_t *strides;
 } NineToothedTensor;
+
+typedef void *NineToothedStream;
+
+typedef int NineToothedResult;
+
+#endif // NINETOOTHED_H
 """
 
 _HEADER_PATH = CACHE_DIR / "ninetoothed.h"
@@ -135,9 +144,9 @@ class _Unparser:
         return f"return {self._generic_unparse(call)};"
 
     def _unparse_FunctionDef(self, node):
-        params = ["CUstream stream"]
+        params = ["NineToothedStream stream"]
         params += [f"NineToothedTensor {arg.arg}" for arg in node.args.args]
-        header = f"CUresult {node.name}({', '.join(params)})"
+        header = f"NineToothedResult {node.name}({', '.join(params)})"
 
         self.header = header
 

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/src/ninetoothed/generation.py

@@ -19,6 +19,7 @@ import uuid
 import sympy
 import triton
 import triton.language as tl
+from triton.language.extra import libdevice
 
 import ninetoothed.naming as naming
 from ninetoothed.cudaifier import Cudaifier
@@ -225,6 +226,41 @@ class CodeGenerator(ast.NodeTransformer):
 
         return node
 
+    def visit_Call(self, node):
+        def _offsets(tensor, dim=None):
+            if dim is None:
+                return tensor._last_generated_overall_offsets.node
+
+            offsets = tensor._last_generated_offsets
+
+            if dim < 0:
+                dim += tensor.source.ndim
+
+            return sum(
+                offsets[dim][target_dim] for target_dim in range(tensor.target.ndim)
+            ).node
+
+        func = node.func
+        args = node.args
+
+        if isinstance(func, ast.Attribute):
+            if func.attr == "offsets":
+                value = func.value
+
+                if self._in_context(value):
+                    tensor = self._context[value.id]
+                elif isinstance(value, ast.Subscript) and self._in_context(value.value):
+                    tensor = self._context[value.value.id]
+
+                self.visit(value)
+
+                # TODO: Add error handling.
+                return _offsets(tensor, ast.literal_eval(args[0]) if args else None)
+
+        self.generic_visit(node)
+
+        return node
+
     def visit_Subscript(self, node):
         if self._in_context(node.value) and isinstance(node.ctx, ast.Load):
             value = self._context[node.value.id]
@@ -242,13 +278,24 @@ class CodeGenerator(ast.NodeTransformer):
         return node
 
     def visit_Attribute(self, node):
-        if self._in_context(node.value):
-            value = self._context[node.value.id]
+        value = node.value
 
-            if isinstance(value, Tensor):
-                inner = value.dtype
+        if isinstance(value, ast.Attribute):
+            value = self.visit_Attribute(value)
+
+        if self._in_context(value):
+            value = self._context[value.id].dtype
+
+        if isinstance(value, Tensor):
+            attr = getattr(value, node.attr)
+
+            if node.attr == "dtype" and attr is None:
+                return Symbol(f"{value.source.pointer_string()}.type.element_ty").node
 
-                return Symbol(getattr(inner, node.attr)).node
+            if isinstance(attr, Tensor):
+                return attr
+
+            return Symbol(attr).node
 
         self.generic_visit(node)
 
@@ -560,6 +607,8 @@ class CodeGenerator(ast.NodeTransformer):
         indices = self._complete_indices(tensor, indices)
         offsets = type(self)._generate_offsets(tensor, indices)
 
+        tensor._last_generated_offsets = offsets
+
         for source_dim in range(tensor.source.ndim):
             for target_dim in range(tensor.target.ndim):
                 if target_dim not in invariant_target_dims:
@@ -584,7 +633,7 @@ class CodeGenerator(ast.NodeTransformer):
                     * tensor.source.strides[source_dim]
                 )
 
-        pointers = name_for_pointers + sum(
+        overall_offsets = sum(
             offsets[source_dim][target_dim][
                 type(self)._generate_slices(tensor, target_dim)
             ]
@@ -594,6 +643,10 @@ class CodeGenerator(ast.NodeTransformer):
             if target_dim not in invariant_target_dims
             and offsets[source_dim][target_dim] != 0
         )
+
+        tensor._last_generated_overall_offsets = overall_offsets
+
+        pointers = name_for_pointers + overall_offsets
         mask = functools.reduce(
             lambda x, y: x & y,
             (
@@ -980,6 +1033,9 @@ class _Inliner(ast.NodeTransformer):
         if func_def is None:
             return None, []
 
+        if inspect.getmodule(func) is libdevice:
+            return None, []
+
         collector = _ImportCollector()
         collector.visit(ast.parse(inspect.getsource(inspect.getmodule(func))))
         self.imports.extend(collector.imports)

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/src/ninetoothed/language.py

@@ -1,7 +1,11 @@
 import ast
 
+from triton.language.extra import libdevice
+
 from ninetoothed.symbol import Symbol
 
+__all__ = ["libdevice"]
+
 LANGUAGE = "ninetoothed.language"
 
 

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/src/ninetoothed/visualization.py

@@ -10,8 +10,6 @@ def visualize(tensor, color=None, save_path=None):
     :param color: The color to be used for visualization.
     :param save_path: The path where the visualization should be saved.
     """
-    outline_width = 0.1
-    plt.rcParams["lines.linewidth"] = 72 * outline_width
 
     if color is None:
         color = f"C{visualize.count}"
@@ -21,6 +19,24 @@ def visualize(tensor, color=None, save_path=None):
     width = max_pos_y + 1
     height = max_pos_x + 1
 
+    _, ax = _prepare_figure_and_axes(width, height)
+
+    _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 _prepare_figure_and_axes(width, height):
+    outline_width = 0.1
+    plt.rcParams["lines.linewidth"] = 72 * outline_width
+
     fig = plt.figure(figsize=(width + outline_width, height + outline_width))
 
     h = (Size.Fixed(0), Size.Fixed(width + outline_width))
@@ -41,16 +57,7 @@ def visualize(tensor, color=None, save_path=None):
     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
+    return fig, ax
 
 
 def _visualize_tensor(ax, tensor, x, y, color, level_spacing=4):

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/tests/test_addmm.py

@@ -12,11 +12,11 @@ from tests.skippers import skip_if_cuda_not_available, skip_if_float8_e5m2_not_s
 def arrangement(input, mat1, mat2, beta, alpha, output):
     _, _, input_arranged = matmul.arrangement(mat1, mat2, input)
 
-    mat1_arrange, mat2_arranged, output_arranged = matmul.arrangement(
+    mat1_arranged, mat2_arranged, output_arranged = matmul.arrangement(
         mat1, mat2, output
     )
 
-    return input_arranged, mat1_arrange, mat2_arranged, beta, alpha, output_arranged
+    return input_arranged, mat1_arranged, mat2_arranged, beta, alpha, output_arranged
 
 
 def application(input, mat1, mat2, beta, alpha, output):

{ninetoothed-0.16.0 → ninetoothed-0.17.0}/tests/test_attention.py

@@ -1,3 +1,4 @@
+import pytest
 import torch
 import torch.nn.functional as F
 
@@ -34,7 +35,7 @@ def arrangement(q, k, v, o):
 
 
 def application(q, k, v, o):
-    q_loaded = (q * 1.44269504089).to(ntl.float16)
+    q_loaded = (q * 1.44269504089).to(q.dtype)
 
     acc = ntl.zeros((q.shape[-2], q.shape[-1]), dtype=ntl.float32)
     l_i = ntl.full((q.shape[-2],), 1, dtype=ntl.float32)
@@ -42,13 +43,14 @@ def application(q, k, v, o):
 
     for i in range(k.shape[0]):
         qk = ntl.dot(q_loaded, ntl.trans(k[i]))
+        qk = ntl.where(k[i].offsets(-2) < k.source.shape[-2], qk, float("-inf"))
 
         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])
+        acc = acc * alpha[:, None] + ntl.dot(p.to(v[i].dtype), v[i])
         m_i = m_ij
         l_i = l_i * alpha + l_ij
 
@@ -83,20 +85,31 @@ def attention(q, k, v):
 
 @skip_if_cuda_not_available
 class TestCUDA:
+    shapes = ((2, 4, 1024, 64), (2, 4, 1, 64))
+
     @classmethod
     def setup_class(cls):
         torch.manual_seed(0)
 
-        shape = (2, 4, 1024, 64)
+        cls.args = {
+            shape: tuple(torch.randn(shape, device="cuda") for _ in range(3))
+            for shape in cls.shapes
+        }
+
+    @pytest.mark.parametrize("shape", shapes)
+    def test_fp32(self, shape):
+        q, k, v = (arg.to(torch.float32) for arg in type(self).args[shape])
 
-        cls.q = torch.randn(shape, device="cuda")
-        cls.k = torch.randn(shape, device="cuda")
-        cls.v = torch.randn(shape, device="cuda")
+        assert torch.allclose(
+            attention(q, k, v),
+            F.scaled_dot_product_attention(q, k, v, scale=1),
+            atol=0.01,
+            rtol=0.01,
+        )
 
-    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)
+    @pytest.mark.parametrize("shape", shapes)
+    def test_fp16(self, shape):
+        q, k, v = (arg.to(torch.float16) for arg in type(self).args[shape])
 
         assert torch.allclose(
             attention(q, k, v),

ninetoothed-0.17.0/tests/test_dropout.py

@@ -0,0 +1,54 @@
+import random
+
+import torch
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Tensor, block_size
+from tests.skippers import skip_if_cuda_not_available
+
+
+def arrangement(input, p, seed, output, BLOCK_SIZE=block_size()):
+    return input.tile((BLOCK_SIZE,)), p, seed, output.tile((BLOCK_SIZE,))
+
+
+def application(input, p, seed, output):
+    output = ntl.where(ntl.rand(seed, input.offsets()) > p, input / (1 - p), 0)  # noqa: F841
+
+
+def dropout(input, p=0.5):
+    seed = random.randrange(0, 2**31)
+    output = torch.empty_like(input)
+
+    tensors = (Tensor(1), Tensor(0), Tensor(0), Tensor(1))
+    dropout_kernel = ninetoothed.make(arrangement, application, tensors)
+
+    dropout_kernel(input, p, seed, output)
+
+    return output
+
+
+@skip_if_cuda_not_available
+class TestCUDA:
+    @classmethod
+    def setup_class(cls):
+        random.seed(0)
+        torch.manual_seed(0)
+
+        size = 349
+
+        cls.input = torch.randn(size, device="cuda")
+
+    def test_fp16(self):
+        input = type(self).input.to(torch.float16)
+        p = 0.3
+
+        output = dropout(input, p=p)
+
+        assert input.shape == output.shape
+
+        non_zero_ratio = output.nonzero().numel() / input.numel()
+
+        assert abs(non_zero_ratio - (1 - p)) < 0.05
+
+        assert torch.allclose(output[output != 0], input[output != 0] / (1 - p))

ninetoothed-0.17.0/tests/test_pow.py

@@ -0,0 +1,50 @@
+import torch
+
+import ninetoothed
+from ninetoothed import Tensor, block_size
+from ninetoothed.language import libdevice
+from tests.skippers import skip_if_cuda_not_available
+
+
+def arrangement(input, exponent, output, BLOCK_SIZE=block_size()):
+    return (
+        input.tile((BLOCK_SIZE,)),
+        exponent.tile((BLOCK_SIZE,)),
+        output.tile((BLOCK_SIZE,)),
+    )
+
+
+def application(input, exponent, output):
+    output = libdevice.pow(input, exponent)  # noqa: F841
+
+
+def pow(input, exponent):
+    output = torch.empty_like(input)
+
+    pow_kernel = ninetoothed.make(
+        arrangement, application, (Tensor(1), Tensor(1), Tensor(1))
+    )
+
+    pow_kernel(input, exponent, output)
+
+    return output
+
+
+@skip_if_cuda_not_available
+class TestCUDA:
+    @classmethod
+    def setup_class(cls):
+        torch.manual_seed(0)
+
+        size = 44925
+
+        cls.input = torch.randn(size, device="cuda")
+        cls.exponent = torch.randn(size, device="cuda")
+
+    def test_fp32(self):
+        input = type(self).input.to(torch.float32)
+        exponent = type(self).exponent.to(torch.float32)
+
+        assert torch.allclose(
+            pow(input, exponent), torch.pow(input, exponent), equal_nan=True
+        )