ninetoothed 0.13.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()
@@ -701,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/tensor.py

@@ -95,12 +95,14 @@ class Tensor:
 
         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.
         """
 
@@ -124,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
             )

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

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

{ninetoothed-0.13.0.dist-info → ninetoothed-0.14.0.dist-info}/RECORD

@@ -1,12 +1,12 @@
 ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
-ninetoothed/jit.py,sha256=fuMfb4gEvz78n-7--GS1Ud14g33SLT-vI1kqt9BPAIQ,25759
+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=Q7WPigNyGOgP5lYYac39pF6zlFbCyXYrICPgGOuuyr4,13976
+ninetoothed/tensor.py,sha256=W1XY8_vaYmszX4lIWuas-ZKGbbdEZU7Z5h1A4FBXDXg,14358
 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.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,,