ninetoothed 0.8.0__py3-none-any.whl → 0.10.0__py3-none-any.whl

ninetoothed/jit.py

@@ -39,23 +39,12 @@ def jit(_func=None, *, _prettify=False):
 
 
 class JIT:
-    handles = collections.defaultdict(dict)
-
     def __init__(self, func, _prettify=False):
         self.func = func
 
         self._prettify = _prettify
 
     def __call__(self):
-        source_file = inspect.getsourcefile(self.func)
-        source_line = inspect.getsourcelines(self.func)[1]
-
-        if (
-            source_file in type(self).handles
-            and source_line in type(self).handles[source_file]
-        ):
-            return type(self).handles[source_file][source_line]
-
         tree = self._get_tree()
 
         CodeGenerator(inspect.get_annotations(self.func)).visit(tree)
@@ -93,8 +82,6 @@ class JIT:
             source,
         )
 
-        type(self).handles[source_file][source_line] = handle
-
         return handle
 
     def _get_tree(self):
@@ -350,6 +337,7 @@ class CodeGenerator(ast.NodeTransformer):
                 + [
                     ast.arg(arg=param)
                     for param in non_next_power_of_2_constexpr_params_without_prefixes
+                    if not Tensor.size_pattern().fullmatch(param)
                 ],
                 kwonlyargs=[],
                 defaults=[],
@@ -489,9 +477,9 @@ class CodeGenerator(ast.NodeTransformer):
         return pointers, mask
 
     def _complete_indices(self, tensor, indices):
-        indices = list(self._generate_pid_indices(tensor) + indices)
+        indices = list(self._generate_pid_indices(tensor) + tuple(indices))
 
-        for size in tensor.inmost().shape:
+        for size in tensor.innermost().shape:
             if Symbol.is_name(size):
                 name = size.node.id
                 if not naming.is_meta(name):
@@ -526,7 +514,10 @@ class CodeGenerator(ast.NodeTransformer):
 
     @staticmethod
     def _generate_slices(tensor, dim):
-        return tuple(slice(None) if i == dim else None for i in range(tensor.ndim))
+        return tuple(
+            slice(None) if target_dim == dim else None
+            for target_dim in tensor.innermost().target_dims
+        )
 
     @staticmethod
     def _generate_offsets(tensor, indices):

ninetoothed/naming.py

@@ -1,6 +1,10 @@
 import re
 
 
+def auto_generate(name):
+    return f"ninetoothed_{name}"
+
+
 def make_constexpr(name):
     return _add_prefix(name, _CONSTEXPR)
 

ninetoothed/tensor.py

@@ -2,6 +2,7 @@ import itertools
 import math
 import re
 
+import ninetoothed.naming as naming
 from ninetoothed.language import call
 from ninetoothed.symbol import Symbol
 
@@ -16,6 +17,7 @@ class Tensor:
         dtype=None,
         strides=None,
         other=None,
+        constexpr_shape=None,
         name=None,
         source=None,
         source_dims=None,
@@ -27,10 +29,13 @@ class Tensor:
         if name is not None:
             self.name = name
         else:
-            self.name = f"_ninetoothed_tensor_{type(self).num_instances}"
+            self.name = naming.auto_generate(f"tensor_{type(self).num_instances}")
 
         if ndim is not None:
-            self.shape = (Symbol(self.size_string(i)) for i in range(ndim))
+            self.shape = (
+                Symbol(self.size_string(i), constexpr=constexpr_shape)
+                for i in range(ndim)
+            )
             self.strides = (Symbol(self.stride_string(i)) for i in range(ndim))
         else:
             self.shape = shape
@@ -107,14 +112,10 @@ class Tensor:
                 strides=inner_strides,
                 source=self.source,
                 source_dims=self.source_dims,
-                target=self.target,
-                target_dims=self.target_dims,
             ),
             strides=outer_strides,
             source=self.source,
             source_dims=self.source_dims,
-            target=self.target,
-            target_dims=self.target_dims,
         )
 
     def expand(self, shape):
@@ -131,23 +132,22 @@ class Tensor:
             ],
             source=self.source,
             source_dims=self.source_dims,
-            target=self.target,
-            target_dims=self.target_dims,
         )
 
     def squeeze(self, dim):
+        if not isinstance(dim, tuple):
+            dim = (dim,)
+
         # TODO: Add error handling.
         return type(self)(
-            shape=[size for i, size in enumerate(self.shape) if dim != i],
+            shape=[size for i, size in enumerate(self.shape) if i not in dim],
             dtype=self.dtype,
-            strides=[stride for i, stride in enumerate(self.strides) if dim != i],
+            strides=[stride for i, stride in enumerate(self.strides) if i not in dim],
             source=self.source,
             source_dims=[
-                source_dim for i, source_dim in enumerate(self.source_dims) if dim != i
-            ],
-            target=self.target,
-            target_dims=[
-                target_dim for i, target_dim in enumerate(self.target_dims) if dim != i
+                source_dim
+                for i, source_dim in enumerate(self.source_dims)
+                if i not in dim
             ],
         )
 
@@ -168,8 +168,6 @@ class Tensor:
             strides=new_strides,
             source=self.source,
             source_dims=new_source_dims,
-            target=self.target,
-            target_dims=self.target_dims,
         )
 
     def flatten(self, start_dim=None, end_dim=None):
@@ -205,8 +203,6 @@ class Tensor:
             strides=new_strides,
             source=self.source,
             source_dims=new_source_dims,
-            target=self.target,
-            target_dims=self.target_dims,
         )
 
     def ravel(self):
@@ -245,11 +241,11 @@ class Tensor:
             | (self.source.names() if self.source is not self else set())
         )
 
-    def inmost(self):
+    def innermost(self):
         if not isinstance(self.dtype, type(self)):
             return self
 
-        return self.dtype.inmost()
+        return self.dtype.innermost()
 
     def pointer_string(self):
         return f"{self.name}_pointer"

ninetoothed/torchifier.py

@@ -1,5 +1,6 @@
 import ast
 
+import ninetoothed.naming as naming
 from ninetoothed.tensor import Tensor
 
 
@@ -9,6 +10,9 @@ class Torchifier(ast.NodeTransformer):
 
         source = node.id
 
+        if naming.is_constexpr(source):
+            return node
+
         def repl(match):
             return f"{match.group(1)}"
 

{ninetoothed-0.8.0.dist-info → ninetoothed-0.10.0.dist-info}/METADATA

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

@@ -0,0 +1,11 @@
+ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
+ninetoothed/jit.py,sha256=0LeDBpSYFgPx4hatP_ZsvElsj0d9d552OKRc__L1Jvc,23460
+ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
+ninetoothed/symbol.py,sha256=rZ5nXtn-U1Nw0BBRJ-kfrwmX_zCbAi76un-Z2QFaoZc,4773
+ninetoothed/tensor.py,sha256=LGS9wYmPKckZjIEXsMHdclTmVhzkfrc3avOPiCQY1tU,9153
+ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
+ninetoothed-0.10.0.dist-info/METADATA,sha256=nQWkQ--AceNN3DoD-kP9_aykZYVw8LOfqf2iO63v1Ek,7055
+ninetoothed-0.10.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ninetoothed-0.10.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.10.0.dist-info/RECORD,,

ninetoothed-0.8.0.dist-info/RECORD

@@ -1,11 +0,0 @@
-ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
-ninetoothed/jit.py,sha256=z70hQEsogfQu0cLxq5m3cOsWsVANcMRJaVv5di9vk1c,23741
-ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
-ninetoothed/naming.py,sha256=3FBnC-S3dAZRcBcob9SrcVpVEYE5IXRacwkCiA3vIGU,891
-ninetoothed/symbol.py,sha256=rZ5nXtn-U1Nw0BBRJ-kfrwmX_zCbAi76un-Z2QFaoZc,4773
-ninetoothed/tensor.py,sha256=_jM0tVgqIwZd3MJJsGVTaLCsSxpPO8JfF4qkMShhQvQ,9429
-ninetoothed/torchifier.py,sha256=8M2PDwyFIfVypX6Z-Vt_bGbsCPqxqKnftL0rXeh9bOM,911
-ninetoothed-0.8.0.dist-info/METADATA,sha256=gPWYhTBH5EdeOyGnArZIEw82aFmoQchD6pxtLi6LGMA,7054
-ninetoothed-0.8.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ninetoothed-0.8.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ninetoothed-0.8.0.dist-info/RECORD,,