ninetoothed 0.12.0__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

@@ -37,6 +37,8 @@ class Tensor:
         source_dims=None,
         target=None,
         target_dims=None,
+        unflattened=None,
+        unflattened_dims=None,
     ):
         self.dtype = dtype
 
@@ -81,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):
@@ -137,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):
@@ -164,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):
@@ -192,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):
@@ -205,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,
@@ -218,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):
@@ -265,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,
@@ -272,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):
@@ -290,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
 
@@ -304,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):
@@ -385,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

@@ -4,6 +4,12 @@ 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
 

{ninetoothed-0.12.0.dist-info → ninetoothed-0.13.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.12.0
+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,9 @@ 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'

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

@@ -1,12 +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=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,,