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

ninetoothed/jit.py

@@ -106,13 +106,18 @@ class JIT:
     def _get_tree(self):
         module = ast.parse(inspect.getsource(inspect.getmodule(self.func)))
 
-        _AliasRestorer().visit(module)
         collector = _ImportCollector()
         collector.visit(module)
+
         finder = _FunctionDefFinder(self.func.__name__)
         finder.visit(module)
+        func_def = finder.result
+
+        inliner = _Inliner(self.func.__globals__)
+        inliner.visit(func_def)
+        module.body = collector.imports + inliner.imports + [finder.result]
 
-        return ast.Module(body=collector.imports + [finder.result], type_ignores=[])
+        return _AliasRestorer().visit(module)
 
     def _find_dependencies(self):
         dependencies = set()
@@ -483,13 +488,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 +512,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 +564,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 +577,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
@@ -663,6 +706,211 @@ class Tritonizer(ast.NodeTransformer):
         return node
 
 
+class _Inliner(ast.NodeTransformer):
+    def __init__(self, globals, imports=[]):
+        self._globals = globals
+
+        self._count = 0
+
+        self.imports = imports
+
+    def visit_Expr(self, node):
+        value, stmts = self._inline_expr(node.value)
+        node.value = value
+        node = self.generic_visit(node)
+
+        if stmts:
+            if isinstance(value, ast.Constant) and value.value is None:
+                return stmts
+
+            return stmts + [node]
+
+        return node
+
+    def visit_Assign(self, node):
+        value, stmts = self._inline_expr(node.value)
+        node.value = value
+        node = self.generic_visit(node)
+
+        if stmts:
+            return stmts + [node]
+
+        return node
+
+    def visit_Return(self, node):
+        if node.value:
+            value, stmts = self._inline_expr(node.value)
+            node.value = value
+
+            if stmts:
+                return stmts + [node]
+
+        return node
+
+    def _inline_expr(self, expr):
+        def _inline_list(lst):
+            new_list = []
+            new_stmts = []
+
+            for expr in lst:
+                expr, stmts = self._inline_expr(expr)
+
+                new_list.append(expr)
+                new_stmts.extend(stmts)
+
+            return new_list, new_stmts
+
+        def _inline_field(field):
+            if isinstance(field, ast.AST):
+                return self._inline_expr(field)
+
+            return field, []
+
+        if isinstance(expr, ast.Call):
+            new_expr, new_stmts = self._inline_call(expr)
+
+            if new_expr is not None:
+                return new_expr, new_stmts
+
+        new_stmts = []
+
+        for field, value in ast.iter_fields(expr):
+            if isinstance(value, list):
+                new_value, new_stmts = _inline_list(value)
+            else:
+                new_value, new_stmts = _inline_field(value)
+
+            setattr(expr, field, new_value)
+            new_stmts.extend(new_stmts)
+
+        return expr, new_stmts
+
+    def _inline_call(self, node):
+        class _ParameterReplacer(ast.NodeTransformer):
+            def __init__(self, mapping):
+                self._mapping = mapping
+
+            def visit_Name(self, node):
+                return self._mapping.get(node.id, node)
+
+        class _LocalVariableRenamer(ast.NodeTransformer):
+            def __init__(self, prefix, local_vars):
+                self._prefix = prefix
+
+                self._local_vars = local_vars
+
+            def visit_Name(self, node):
+                if node.id in self._local_vars:
+                    node.id = f"{self._prefix}{node.id}"
+
+                return node
+
+            def visit_arg(self, node):
+                return node
+
+        def _resolve_function(node, globals):
+            if isinstance(node, ast.Name):
+                return globals.get(node.id)
+
+            if isinstance(node, ast.Attribute):
+                obj = _resolve_function(node.value, globals)
+
+                if obj is not None:
+                    return getattr(obj, node.attr, None)
+
+            return None
+
+        def _get_source(func):
+            try:
+                return inspect.getsource(func)
+            except TypeError:
+                return None
+
+        def _find_function_definition(source):
+            finder = _FunctionDefFinder(func.__name__)
+            finder.visit(ast.parse(source))
+
+            return finder.result
+
+        def _find_assigned_names(stmts):
+            class _AssignedNameFinder(ast.NodeVisitor):
+                def __init__(self):
+                    self.result = set()
+
+                def visit_Name(self, node):
+                    if isinstance(node.ctx, ast.Store):
+                        self.result.add(node.id)
+
+            names = set()
+
+            for stmt in stmts:
+                finder = _AssignedNameFinder()
+                finder.visit(stmt)
+                names |= finder.result
+
+            return names
+
+        def _make_temporary():
+            prefix = naming.auto_generate(f"temporary_{self._count}")
+            self._count += 1
+
+            return prefix
+
+        func = _resolve_function(node.func, self._globals)
+
+        if func is None:
+            return None, []
+
+        source = _get_source(func)
+
+        if source is None:
+            return None, []
+
+        func_def = _find_function_definition(source)
+
+        if func_def is None:
+            return None, []
+
+        collector = _ImportCollector()
+        collector.visit(ast.parse(inspect.getsource(inspect.getmodule(func))))
+        self.imports.extend(collector.imports)
+
+        param_names = [arg.arg for arg in func_def.args.args]
+
+        mapping = {param: arg for param, arg in zip(param_names, node.args)}
+        param_replacer = _ParameterReplacer(mapping)
+        body = [param_replacer.visit(stmt) for stmt in func_def.body]
+
+        local_vars = _find_assigned_names(body) - set(param_names)
+        prefix = _make_temporary()
+        local_var_renamer = _LocalVariableRenamer(prefix, local_vars)
+        body = [local_var_renamer.visit(stmt) for stmt in body]
+
+        inlined_body = []
+
+        inliner = _Inliner(func.__globals__)
+
+        for stmt in body:
+            inlined_stmt = inliner.visit(stmt)
+
+            if isinstance(inlined_stmt, list):
+                inlined_body.extend(inlined_stmt)
+            else:
+                inlined_body.append(inlined_stmt)
+
+        if not inlined_body or not isinstance(inlined_body[-1], ast.Return):
+            return ast.Constant(value=None), inlined_body
+
+        ret = inlined_body.pop()
+        temp = _make_temporary()
+        assignment = ast.Assign(
+            targets=[ast.Name(id=temp, ctx=ast.Store())], value=ret.value
+        )
+        inlined_body.append(assignment)
+
+        return ast.Name(id=temp, ctx=ast.Load()), inlined_body
+
+
 class _BinOpSimplifier(ast.NodeTransformer):
     def visit_BinOp(self, node):
         self.generic_visit(node)

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,14 +83,26 @@ 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):
+    def tile(self, tile_shape, strides=None, dilation=None, floor_mode=False):
         """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.
+        :param floor_mode: If ``True``, will use floor division to
+            compute the outer shape.
         :return: A hierarchical tensor.
         """
 
@@ -112,11 +126,21 @@ class Tensor:
             if stride == -1:
                 stride = tile_size
 
-            def cdiv(x, y):
+            def _div(x, y, floor_mode=False):
+                if floor_mode:
+                    return x // y
+
                 return (x + y - 1) // y
 
             new_size = (
-                (cdiv(self_size - spacing * (tile_size - 1) - 1, stride) + 1)
+                (
+                    _div(
+                        self_size - spacing * (tile_size - 1) - 1,
+                        stride,
+                        floor_mode=floor_mode,
+                    )
+                    + 1
+                )
                 if stride != 0
                 else -1
             )
@@ -137,10 +161,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 +192,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 +222,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 +241,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 +256,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 +305,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 +322,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 +342,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 +358,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 +441,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.14.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.12.0
+Version: 0.14.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.14.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
+ninetoothed/jit.py,sha256=wsj9RmaQ1dKPJXtLL4zcWpinwh-7km7-wA7pexy7vq4,31675
+ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
+ninetoothed/symbol.py,sha256=UpGmx_jvaDtowADnp1DwYC3fvBXSiaMiYpU-ewkVo50,5261
+ninetoothed/tensor.py,sha256=W1XY8_vaYmszX4lIWuas-ZKGbbdEZU7Z5h1A4FBXDXg,14358
+ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
+ninetoothed/visualization.py,sha256=IZ7iTT4dl5_JFbO-WfSWPFWpgkyPr4nylwhSZVy8gss,3601
+ninetoothed-0.14.0.dist-info/METADATA,sha256=XAI4ibJdwJDCPCpoOeFOT0XWEDQGUZi6OvJwyjgcEC8,7311
+ninetoothed-0.14.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ninetoothed-0.14.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.14.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,,