ninetoothed 0.11.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

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/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/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/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.11.0.dist-info → ninetoothed-0.12.0.dist-info}/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,12 @@
+ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
+ninetoothed/jit.py,sha256=U3Nen5vyx69ulW7_hnRuATW86Ag9NgVgd3U02NVB20c,24430
+ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
+ninetoothed/symbol.py,sha256=mN96tp-2eUxbiNfxuxtKWNSxOSdYqlcmpY2MYQ-FiEg,4993
+ninetoothed/tensor.py,sha256=E63sq3jh7ZLiLwFYTtavztKEZx7kRX-UVa2ZXSP2X0s,12008
+ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
+ninetoothed/visualization.py,sha256=VPPh__Bral_Z9hKj9D4UOo8HvRFQidCWSe9cS-D5QfY,3351
+ninetoothed-0.12.0.dist-info/METADATA,sha256=gb5zeAxwYQRm-dFNmdX3uDIn_U-abEZ-VU9bSBBxioA,7198
+ninetoothed-0.12.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ninetoothed-0.12.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.12.0.dist-info/RECORD,,

ninetoothed-0.11.0.dist-info/RECORD

@@ -1,11 +0,0 @@
-ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
-ninetoothed/jit.py,sha256=0LeDBpSYFgPx4hatP_ZsvElsj0d9d552OKRc__L1Jvc,23460
-ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
-ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
-ninetoothed/symbol.py,sha256=rZ5nXtn-U1Nw0BBRJ-kfrwmX_zCbAi76un-Z2QFaoZc,4773
-ninetoothed/tensor.py,sha256=ehgbHrxL1gc9iOI9rGUp5k-nI_daLmCfOdLI9hE-GLw,9756
-ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
-ninetoothed-0.11.0.dist-info/METADATA,sha256=NlC4oj1R7gNoCdA1AxYc_7UWY3uZPv3LeOq9Jo-9Q8w,7055
-ninetoothed-0.11.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ninetoothed-0.11.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ninetoothed-0.11.0.dist-info/RECORD,,