ninetoothed 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

ninetoothed/jit.py

@@ -5,7 +5,6 @@ import inspect
 import itertools
 import math
 import tempfile
-import textwrap
 
 from ninetoothed.language import attribute, call
 from ninetoothed.symbol import Symbol
@@ -33,8 +32,7 @@ class JIT:
         ):
             return type(self).handles[source_file][source_line]
 
-        source = textwrap.dedent(inspect.getsource(self.func))
-        tree = ast.parse(source)
+        tree = self._get_tree()
 
         CodeGenerator(inspect.get_annotations(self.func)).visit(tree)
         Tritonizer().visit(tree)
@@ -56,15 +54,7 @@ class JIT:
         namespace = {}
         exec(code, namespace)
 
-        class Handle:
-            def __init__(self, kernel, launch):
-                self._kernel = kernel
-                self._launch = launch
-
-            def __call__(self, *args, **kwargs):
-                return self._launch(*args, **kwargs)
-
-        handle = Handle(
+        handle = _Handle(
             namespace[self.func.__name__],
             namespace[f"launch_{self.func.__name__}"],
         )
@@ -73,6 +63,15 @@ class JIT:
 
         return handle
 
+    def _get_tree(self):
+        module = ast.parse(inspect.getsource(inspect.getmodule(self.func)))
+
+        _AliasRestorer().visit(module)
+        finder = _FunctionDefFinder(self.func.__name__)
+        finder.visit(module)
+
+        return ast.Module(body=[finder.result], type_ignores=[])
+
 
 class CodeGenerator(ast.NodeTransformer):
     def __init__(self, context):
@@ -100,6 +99,18 @@ class CodeGenerator(ast.NodeTransformer):
 
         self.generic_visit(node)
 
+        for arg in self._args:
+            if not isinstance(arg, Tensor):
+                continue
+
+            node.body.insert(
+                0,
+                ast.Assign(
+                    targets=[Symbol(f"{arg.name}_ptrs").node],
+                    value=arg.pointers().node,
+                ),
+            )
+
         return node
 
     def visit_arguments(self, node):
@@ -136,12 +147,12 @@ class CodeGenerator(ast.NodeTransformer):
             value = self._context[node.value.id]
 
             if isinstance(value, Tensor):
-                if isinstance(node.slice, ast.Tuple):
-                    indices = value.indices() + tuple(node.slice.elts)
-                else:
-                    indices = value.indices() + (node.slice,)
-                offsets = value.offsets(indices)
-                pointers = value.pointers(offsets)
+                pointers = type(self)._create_pointers(
+                    value,
+                    node.slice.elts
+                    if isinstance(node.slice, ast.Tuple)
+                    else (node.slice,),
+                )
 
                 return call("load", pointers).node
 
@@ -166,7 +177,9 @@ class CodeGenerator(ast.NodeTransformer):
         self.generic_visit(node)
 
         if node.id in self._context and isinstance(node.ctx, ast.Load):
-            return call("load", self._context[node.id].pointers().node).node
+            return call(
+                "load", type(self)._create_pointers(self._context[node.id], ()).node
+            ).node
 
         return node
 
@@ -180,7 +193,7 @@ class CodeGenerator(ast.NodeTransformer):
                 return ast.Expr(
                     call(
                         "store",
-                        self._context[target.id].pointers().node,
+                        type(self)._create_pointers(self._context[target.id], ()).node,
                         node.value,
                     ).node
                 )
@@ -195,13 +208,12 @@ class CodeGenerator(ast.NodeTransformer):
                 if isinstance(value, Tensor):
                     self.generic_visit(node)
 
-                    indices = value.indices() + tuple(
+                    pointers = type(self)._create_pointers(
+                        value,
                         target.slice.elts
                         if isinstance(target.slice, ast.Tuple)
-                        else target.slice
+                        else (target.slice,),
                     )
-                    offsets = value.offsets(indices)
-                    pointers = value.pointers(offsets)
 
                     return ast.Expr(
                         call(
@@ -316,6 +328,14 @@ class CodeGenerator(ast.NodeTransformer):
 
         return ast.parse(f"lambda meta: ({num_elements},)", mode="eval").body
 
+    @staticmethod
+    def _create_pointers(tensor, indices):
+        return Symbol(f"{tensor.name}_ptrs") + tensor.offsets(
+            [0 for _ in range(tensor.ndim())]
+            + list(indices)
+            + [0 for _ in range(tensor.inmost().ndim())]
+        )
+
 
 class Tritonizer(ast.NodeTransformer):
     def visit_Module(self, node):
@@ -329,8 +349,8 @@ class Tritonizer(ast.NodeTransformer):
     def visit_Name(self, node):
         self.generic_visit(node)
 
-        if node.id == "ninetoothed":
-            node.id = "triton"
+        if node.id == "ninetoothed" or "ninetoothed." in node.id:
+            node.id = node.id.replace("ninetoothed", "triton")
 
         return node
 
@@ -348,3 +368,73 @@ class Tritonizer(ast.NodeTransformer):
             )
 
         return node
+
+
+class _Handle:
+    def __init__(self, kernel, launch):
+        self._kernel = kernel
+        self._launch = launch
+
+    def __call__(self, *args, **kwargs):
+        return self._launch(*args, **kwargs)
+
+
+class _AliasRestorer(ast.NodeTransformer):
+    def __init__(self):
+        super().__init__()
+
+        self._aliases = {}
+        self._redefined = set()
+
+    def visit_Import(self, node):
+        for alias in node.names:
+            if alias.asname:
+                self._aliases[alias.asname] = alias.name
+
+        return node
+
+    def visit_ImportFrom(self, node):
+        for alias in node.names:
+            full_name = f"{node.module}.{alias.name}"
+            if alias.asname:
+                self._aliases[alias.asname] = full_name
+
+        return node
+
+    def visit_Assign(self, node):
+        for target in node.targets:
+            if isinstance(target, ast.Name):
+                self._redefined.add(target.id)
+
+        return self.generic_visit(node)
+
+    def visit_FunctionDef(self, node):
+        original_redefined = self._redefined.copy()
+
+        self.generic_visit(node)
+
+        self._redefined = original_redefined
+
+        return node
+
+    def visit_Name(self, node):
+        if node.id in self._redefined:
+            return node
+
+        if node.id in self._aliases:
+            return ast.Name(id=self._aliases[node.id], ctx=node.ctx)
+
+        return node
+
+
+class _FunctionDefFinder(ast.NodeVisitor):
+    def __init__(self, name):
+        self._name = name
+
+        self.result = None
+
+    def visit_FunctionDef(self, node):
+        if node.name == self._name:
+            self.result = node
+
+        self.generic_visit(node)

ninetoothed/symbol.py

@@ -34,24 +34,47 @@ class Symbol:
             self._node.id = type(self)._create_constexpr(self._node.id)
 
     def __add__(self, other):
-        return type(self)(
-            ast.BinOp(left=self._node, op=ast.Add(), right=type(self)(other)._node)
-        )
+        other = type(self)(other)
+
+        if isinstance(self._node, ast.Constant) and self._node.value == 0:
+            return other
+
+        if isinstance(other._node, ast.Constant) and other._node.value == 0:
+            return self
+
+        return type(self)(ast.BinOp(left=self._node, op=ast.Add(), right=other._node))
 
     def __radd__(self, other):
         return self.__add__(other)
 
     def __mul__(self, other):
-        return type(self)(
-            ast.BinOp(left=self._node, op=ast.Mult(), right=type(self)(other)._node)
-        )
+        other = type(self)(other)
+
+        if isinstance(self._node, ast.Constant) and self._node.value == 0:
+            return type(self)(0)
+
+        if isinstance(other._node, ast.Constant) and other._node.value == 0:
+            return type(self)(0)
+
+        if isinstance(self._node, ast.Constant) and self._node.value == 1:
+            return other
+
+        if isinstance(other._node, ast.Constant) and other._node.value == 1:
+            return self
+
+        return type(self)(ast.BinOp(left=self._node, op=ast.Mult(), right=other._node))
 
     def __rmul__(self, other):
         return self.__mul__(other)
 
     def __floordiv__(self, other):
+        other = type(self)(other)
+
+        if isinstance(other._node, ast.Constant) and other._node.value == 1:
+            return self
+
         return type(self)(
-            ast.BinOp(left=self._node, op=ast.FloorDiv(), right=type(self)(other)._node)
+            ast.BinOp(left=self._node, op=ast.FloorDiv(), right=other._node)
         )
 
     def __mod__(self, other):

ninetoothed/tensor.py

@@ -103,11 +103,12 @@ class Tensor:
             indices = self.indices()
 
         if not isinstance(self.dtype, type(self)):
-            if indices:
+            if len(indices) != self.ndim():
                 raise IndexError("Incorrect number of indices.")
 
             return sum(
-                self.stride(idx)
+                indices[idx]
+                * self.stride(idx)
                 * call("arange", 0, self.size(idx))[
                     tuple(slice(None) if i == idx else None for i in range(self.ndim()))
                 ]
@@ -131,8 +132,21 @@ class Tensor:
             indices.append(index // stride)
             index %= stride
 
+        curr = self.dtype
+        while isinstance(curr, type(self)):
+            indices.extend(
+                0 if curr is not self.inmost() else 1 for _ in range(curr.ndim())
+            )
+            curr = curr.dtype
+
         return tuple(indices)
 
+    def inmost(self):
+        if not isinstance(self.dtype, type(self)):
+            return self
+
+        return self.dtype.inmost()
+
     def ndim(self):
         return len(self.shape)
 

{ninetoothed-0.1.1.dist-info → ninetoothed-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: ninetoothed
-Version: 0.1.1
+Version: 0.2.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.2.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+ninetoothed/__init__.py,sha256=T5UJXlC-wbo8JKPbLUNT65Kccp12xP52WFV5FsugETI,147
+ninetoothed/jit.py,sha256=hmUzkFZzsiKLgOHbsN0MAr1G1JCiyQ22cFPtmyZ1OyE,12725
+ninetoothed/language.py,sha256=cSuTgi5OwmLFy-dy_AHGZzRm18wz01ByHQ2vioP1vTg,437
+ninetoothed/symbol.py,sha256=I2Mc9D1w7AYAIQtyAXyDQ-FBqowVZrd-PK-JOt_SpgA,3787
+ninetoothed/tensor.py,sha256=RfwYzdYASkr6usJklESm1n8RoxvYjWnPtCjIfipa2fg,5000
+ninetoothed/torchifier.py,sha256=JmIVQE8r0zr_RLExsRDOGNsMu0F7v6J_o22aWqlw81k,841
+ninetoothed-0.2.0.dist-info/METADATA,sha256=w6qkc2riniG0N4nDUCUkZWF8Eve3j5brBQHIWIEqLXQ,5422
+ninetoothed-0.2.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+ninetoothed-0.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.2.0.dist-info/RECORD,,

ninetoothed-0.1.1.dist-info/RECORD

@@ -1,10 +0,0 @@
-ninetoothed/__init__.py,sha256=T5UJXlC-wbo8JKPbLUNT65Kccp12xP52WFV5FsugETI,147
-ninetoothed/jit.py,sha256=DdRdZ7DhfZwJeS7AcO_RhD9TZcCebKI55V4_6UHs3bo,10523
-ninetoothed/language.py,sha256=cSuTgi5OwmLFy-dy_AHGZzRm18wz01ByHQ2vioP1vTg,437
-ninetoothed/symbol.py,sha256=8BI4ekeLuUdHTEREvMMlAzwrJ93pqiCdSHGc38clBFA,3034
-ninetoothed/tensor.py,sha256=o_HLEuaBzojmbMLnbPGLcw4iqBI34TNdES3YLTagztE,4590
-ninetoothed/torchifier.py,sha256=JmIVQE8r0zr_RLExsRDOGNsMu0F7v6J_o22aWqlw81k,841
-ninetoothed-0.1.1.dist-info/METADATA,sha256=1Nv6Xcz7CrpEUrzAYH93bYVX8GfPtHwzj4yofeaoJro,5422
-ninetoothed-0.1.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-ninetoothed-0.1.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ninetoothed-0.1.1.dist-info/RECORD,,