ninetoothed 0.11.1__py3-none-any.whl → 0.13.0__py3-none-any.whl

ninetoothed/jit.py

@@ -483,13 +483,15 @@ class CodeGenerator(ast.NodeTransformer):
         mask = functools.reduce(
             lambda x, y: x & y,
             (
-                offsets[source_dim][target_dim][
-                    type(self)._generate_slices(tensor, target_dim)
-                ]
+                sum(
+                    offsets[source_dim][target_dim][
+                        type(self)._generate_slices(tensor, target_dim)
+                    ]
+                    for target_dim in range(tensor.target.ndim)
+                    if offsets[source_dim][target_dim] != 0
+                )
                 < tensor.source.shape[source_dim]
                 for source_dim in range(tensor.source.ndim)
-                for target_dim in range(tensor.target.ndim)
-                if offsets[source_dim][target_dim] != 0
             ),
         ) & functools.reduce(
             lambda x, y: x & y,
@@ -505,13 +507,21 @@ class CodeGenerator(ast.NodeTransformer):
         return pointers, mask
 
     def _complete_indices(self, tensor, indices):
+        class _NextPowerOfTwoMaker(ast.NodeTransformer):
+            def visit_Name(self, node):
+                name = node.id
+
+                if not naming.is_meta(name):
+                    next_power_of_2_name = naming.make_next_power_of_2(name)
+
+                    return ast.Name(id=next_power_of_2_name, ctx=ast.Load())
+
+                return self.generic_visit(node)
+
         indices = list(self._generate_pid_indices(tensor) + tuple(indices))
 
         for size in tensor.innermost().shape:
-            if Symbol.is_name(size):
-                name = size.node.id
-                if not naming.is_meta(name):
-                    size = naming.make_next_power_of_2(name)
+            size = _NextPowerOfTwoMaker().visit(Symbol(copy.deepcopy(size)).node)
 
             indices.append(call("arange", 0, size))
 
@@ -549,8 +559,10 @@ class CodeGenerator(ast.NodeTransformer):
 
     @staticmethod
     def _generate_offsets(tensor, indices):
-        offsets = collections.defaultdict(
-            lambda: collections.defaultdict(lambda: Symbol(0))
+        raw_offsets = collections.defaultdict(
+            lambda: collections.defaultdict(
+                lambda: collections.defaultdict(lambda: Symbol(0))
+            )
         )
 
         curr = tensor
@@ -560,36 +572,62 @@ class CodeGenerator(ast.NodeTransformer):
             stop = start + curr.ndim
             curr_indices = indices[start:stop]
 
-            for index, stride, source_dim, target_dim in zip(
-                curr_indices, curr.strides, curr.source_dims, curr.target_dims
+            for index, stride, source_dim, target_dim, unflattened_dim in zip(
+                curr_indices,
+                curr.strides,
+                curr.source_dims,
+                curr.target_dims,
+                curr.unflattened_dims,
             ):
-                offsets[source_dim][target_dim] += index * stride
+                raw_offsets[source_dim][target_dim][unflattened_dim] += index * stride
 
             start = stop
             curr = curr.dtype
 
-        for source_dim in tuple(offsets):
-            for target_dim in tuple(offsets[source_dim]):
-                if not isinstance(source_dim, tuple):
-                    continue
+        offsets = collections.defaultdict(
+            lambda: collections.defaultdict(lambda: Symbol(0))
+        )
 
-                unraveled = CodeGenerator._unravel_index(
-                    offsets[source_dim][target_dim],
-                    tuple(tensor.source.shape[dim] for dim in source_dim),
-                )
+        source_strides = tuple(Symbol(stride) for stride in tensor.source.strides)
 
-                for offs, dim in zip(unraveled, source_dim):
-                    offsets[dim][target_dim] = offs
+        unflattened_strides = tuple(
+            Symbol(stride) for stride in tensor.unflattened.strides
+        )
 
-        for source_dim in range(tensor.source.ndim):
-            for target_dim in range(tensor.target.ndim):
-                offsets[source_dim][target_dim] = copy.deepcopy(
-                    offsets[source_dim][target_dim]
+        def _add_unraveled_offsets(raw_offs, source_dim, target_dim, unflattened_dim):
+            if not isinstance(unflattened_dim, tuple):
+                offsets[source_dim][target_dim] += copy.deepcopy(
+                    raw_offs
+                ).find_and_replace(
+                    unflattened_strides, Symbol(1)
+                ) * unflattened_strides[unflattened_dim].find_and_replace(
+                    source_strides, Symbol(1)
                 )
-                offsets[source_dim][target_dim].find_and_replace(
-                    Symbol(tensor.source.strides[source_dim]), Symbol(1)
+
+                return
+
+            unraveled_offs = CodeGenerator._unravel_index(
+                raw_offs,
+                tuple(tensor.unflattened.shape[dim] for dim in unflattened_dim),
+            )
+
+            for raw_offs, source_dim, unflattened_dim in zip(
+                unraveled_offs, source_dim, unflattened_dim
+            ):
+                _add_unraveled_offsets(
+                    raw_offs, source_dim, target_dim, unflattened_dim
                 )
 
+        for source_dim in tuple(raw_offsets):
+            for target_dim in tuple(raw_offsets[source_dim]):
+                for unflattened_dim in tuple(raw_offsets[source_dim][target_dim]):
+                    _add_unraveled_offsets(
+                        raw_offsets[source_dim][target_dim][unflattened_dim],
+                        source_dim,
+                        target_dim,
+                        unflattened_dim,
+                    )
+
         return offsets
 
     @staticmethod

ninetoothed/symbol.py

@@ -140,7 +140,12 @@ class Symbol:
         return ast.unparse(self._node)
 
     def find_and_replace(self, target, replacement):
-        _FindAndReplacer(target.node, replacement.node).visit(self._node)
+        if isinstance(target, tuple):
+            targets = tuple(item.node for item in target)
+        else:
+            targets = (target.node,)
+
+        return Symbol(_FindAndReplacer(targets, replacement.node).visit(self._node))
 
     def names(self):
         class NameCollector(ast.NodeVisitor):
@@ -175,12 +180,14 @@ class Symbol:
 
 
 class _FindAndReplacer(ast.NodeTransformer):
-    def __init__(self, target, replacement):
-        self._target_id = target.id
+    def __init__(self, targets, replacement):
+        self._targets_unparsed = tuple(
+            sorted({ast.unparse(target) for target in targets}, key=len, reverse=True)
+        )
         self._replacement = replacement
 
-    def visit_Name(self, node):
-        if node.id == self._target_id:
+    def visit(self, node):
+        if ast.unparse(node) in self._targets_unparsed:
             return self._replacement
 
-        return self.generic_visit(node)
+        return super().visit(node)

ninetoothed/tensor.py

@@ -3,7 +3,6 @@ import math
 import re
 
 import ninetoothed.naming as naming
-from ninetoothed.language import call
 from ninetoothed.symbol import Symbol
 
 
@@ -38,6 +37,8 @@ class Tensor:
         source_dims=None,
         target=None,
         target_dims=None,
+        unflattened=None,
+        unflattened_dims=None,
     ):
         self.dtype = dtype
 
@@ -82,6 +83,16 @@ class Tensor:
         else:
             self.target_dims = (dim for dim in range(self.target.ndim))
 
+        if unflattened is not None:
+            self.unflattened = unflattened
+        else:
+            self.unflattened = self
+
+        if unflattened_dims is not None:
+            self.unflattened_dims = unflattened_dims
+        else:
+            self.unflattened_dims = (dim for dim in range(self.unflattened.ndim))
+
         type(self).num_instances += 1
 
     def tile(self, tile_shape, strides=None, dilation=None):
@@ -113,8 +124,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
             )
@@ -135,10 +149,14 @@ class Tensor:
                 strides=inner_strides,
                 source=self.source,
                 source_dims=self.source_dims,
+                unflattened=self.unflattened,
+                unflattened_dims=self.unflattened_dims,
             ),
             strides=outer_strides,
             source=self.source,
             source_dims=self.source_dims,
+            unflattened=self.unflattened,
+            unflattened_dims=self.unflattened_dims,
         )
 
     def expand(self, shape):
@@ -162,6 +180,8 @@ class Tensor:
             source=self.source,
             source_dims=self.source_dims,
             target_dims=self.target_dims,
+            unflattened=self.unflattened,
+            unflattened_dims=self.unflattened_dims,
         )
 
     def squeeze(self, dim):
@@ -190,6 +210,12 @@ class Tensor:
                 for i, target_dim in enumerate(self.target_dims)
                 if i not in dim
             ],
+            unflattened=self.unflattened,
+            unflattened_dims=[
+                unflattened_dim
+                for i, unflattened_dim in enumerate(self.unflattened_dims)
+                if i not in dim
+            ],
         )
 
     def permute(self, dims):
@@ -203,11 +229,13 @@ class Tensor:
         new_shape = [None for _ in range(self.ndim)]
         new_strides = [None for _ in range(self.ndim)]
         new_source_dims = [None for _ in range(self.ndim)]
+        new_unflattened_dims = [None for _ in range(self.ndim)]
 
         for original_dim, permuted_dim in enumerate(dims):
             new_shape[original_dim] = self.shape[permuted_dim]
             new_strides[original_dim] = self.strides[permuted_dim]
             new_source_dims[original_dim] = self.source_dims[permuted_dim]
+            new_unflattened_dims[original_dim] = self.unflattened_dims[permuted_dim]
 
         return type(self)(
             shape=new_shape,
@@ -216,6 +244,8 @@ class Tensor:
             source=self.source,
             source_dims=new_source_dims,
             target_dims=self.target_dims,
+            unflattened=self.unflattened,
+            unflattened_dims=new_unflattened_dims,
         )
 
     def flatten(self, start_dim=None, end_dim=None):
@@ -263,6 +293,16 @@ class Tensor:
             leading_target_dims + (flattening_target_dims[-1],) + trailing_target_dims
         )
 
+        leading_unflattened_dims = self.unflattened_dims[:start_dim]
+        flattening_unflattened_dims = self.unflattened_dims[start_dim:end_dim]
+        trailing_unflattened_dims = self.unflattened_dims[end_dim:]
+
+        new_unflattened_dims = (
+            leading_unflattened_dims
+            + (flattening_unflattened_dims,)
+            + trailing_unflattened_dims
+        )
+
         return type(self)(
             shape=new_shape,
             dtype=self.dtype,
@@ -270,6 +310,8 @@ class Tensor:
             source=self.source,
             source_dims=new_source_dims,
             target_dims=new_target_dims,
+            unflattened=self.unflattened,
+            unflattened_dims=new_unflattened_dims,
         )
 
     def ravel(self):
@@ -288,12 +330,14 @@ class Tensor:
         # TODO: Add error handling.
         new_shape = []
         new_strides = []
+        new_source_dims = []
 
         curr = self
 
         while isinstance(curr, type(self)):
             new_shape.extend(curr.shape)
             new_strides.extend(curr.strides)
+            new_source_dims.extend(curr.source_dims)
 
             curr = curr.dtype
 
@@ -302,6 +346,8 @@ class Tensor:
             strides=new_strides,
             other=self.source.other,
             name=self.source.name,
+            source=self.source,
+            source_dims=new_source_dims,
         )
 
     def names(self):
@@ -383,6 +429,14 @@ class Tensor:
     def target_dims(self, value):
         self._target_dims = tuple(value)
 
+    @property
+    def unflattened_dims(self):
+        return self._unflattened_dims
+
+    @unflattened_dims.setter
+    def unflattened_dims(self, value):
+        self._unflattened_dims = tuple(value)
+
     @staticmethod
     def pointer_pattern():
         return re.compile(rf"({_identifier_pattern_raw_string()})_(pointer)")

ninetoothed/visualization.py

@@ -0,0 +1,128 @@
+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):
+    """Visualize a tensor as a structured grid representation.
+
+    :param tensor: The tensor to be visualized.
+    :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}"
+
+    _, 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.1.dist-info → ninetoothed-0.13.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.11.1
+Version: 0.13.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,12 @@ Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Requires-Dist: triton>=3.0.0
+Provides-Extra: all
+Requires-Dist: matplotlib>=3.9.0; extra == 'all'
+Requires-Dist: numpy>=2.1.0; extra == 'all'
+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.13.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
+ninetoothed/jit.py,sha256=fuMfb4gEvz78n-7--GS1Ud14g33SLT-vI1kqt9BPAIQ,25759
+ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
+ninetoothed/symbol.py,sha256=UpGmx_jvaDtowADnp1DwYC3fvBXSiaMiYpU-ewkVo50,5261
+ninetoothed/tensor.py,sha256=Q7WPigNyGOgP5lYYac39pF6zlFbCyXYrICPgGOuuyr4,13976
+ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
+ninetoothed/visualization.py,sha256=IZ7iTT4dl5_JFbO-WfSWPFWpgkyPr4nylwhSZVy8gss,3601
+ninetoothed-0.13.0.dist-info/METADATA,sha256=Vp_huL1YVgNAa9q6C0J5hZnCR7SRyMFRALGD21ikrjc,7311
+ninetoothed-0.13.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ninetoothed-0.13.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.13.0.dist-info/RECORD,,

ninetoothed-0.11.1.dist-info/RECORD

@@ -1,11 +0,0 @@
-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=Pgl3t08qNDJrjbNMLltEIwMu19vnKwMUncOq4aedTjY,11983
-ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
-ninetoothed-0.11.1.dist-info/METADATA,sha256=DsvP5HloDDgEiHUtPEu8blUVYCBr3_VXYILUikfOE-Y,7055
-ninetoothed-0.11.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ninetoothed-0.11.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ninetoothed-0.11.1.dist-info/RECORD,,