ninetoothed 0.8.0__py3-none-any.whl → 0.9.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,7 +477,7 @@ 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:
             if Symbol.is_name(size):

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

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.9.0.dist-info}/METADATA

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

@@ -0,0 +1,11 @@
+ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
+ninetoothed/jit.py,sha256=B3q32ksKTRr7I4jLcoDXjwEx7A_Awz9DXGEmIkrtoBc,23393
+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=OU6lVjzKU614mk3EN1AAgTDradbvJkyl22AXwdhxcfs,9577
+ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
+ninetoothed-0.9.0.dist-info/METADATA,sha256=LKtSbgc_mWKJ4L_8BYiRKWjoV-JKjr4hefhrkdPmrHs,7054
+ninetoothed-0.9.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ninetoothed-0.9.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.9.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,,