ninetoothed 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

ninetoothed/jit.py

@@ -137,7 +137,7 @@ class CodeGenerator(ast.NodeTransformer):
         node.args = [
             ast.arg(arg=name)
             if not Symbol.is_constexpr(name)
-            else ast.arg(arg=name, annotation=attribute("constexpr"))
+            else ast.arg(arg=name, annotation=attribute("constexpr").node)
             for name in non_meta_names
         ] + [
             ast.arg(arg=name, annotation=attribute("constexpr").node)
@@ -287,15 +287,30 @@ class CodeGenerator(ast.NodeTransformer):
         )
 
     def _generate_launch(self, params, meta):
+        constexpr_params = [param for param in params if Symbol.is_constexpr(param)]
+        constexpr_params_without_prefixes = [
+            Symbol.remove_prefix(param) for param in constexpr_params
+        ]
+
         launch = ast.FunctionDef(
             name=f"launch_{self._func_def.name}",
             args=ast.arguments(
                 posonlyargs=[],
-                args=[ast.arg(arg.original.name) for arg in self._args],
+                args=[ast.arg(arg=arg.original.name) for arg in self._args]
+                + [ast.arg(arg=param) for param in constexpr_params_without_prefixes],
                 kwonlyargs=[],
                 defaults=[],
             ),
             body=[
+                ast.Assign(
+                    targets=[ast.Name(id=param, ctx=ast.Store())],
+                    value=ast.Name(id=param_without_prefix, ctx=ast.Load()),
+                )
+                for param, param_without_prefix in zip(
+                    constexpr_params, constexpr_params_without_prefixes
+                )
+            ]
+            + [
                 ast.Expr(
                     ast.Call(
                         func=ast.Subscript(

ninetoothed/symbol.py

@@ -137,19 +137,39 @@ class Symbol:
 
     @staticmethod
     def is_constexpr(name):
-        return name.startswith("_ninetoothed_constexpr_") or Symbol.is_meta(name)
+        return name.startswith(Symbol._constexpr_prefix()) or Symbol.is_meta(name)
 
     @staticmethod
     def is_meta(name):
-        return name.startswith("_ninetoothed_meta_")
+        return name.startswith(Symbol._meta_prefix())
+
+    @staticmethod
+    def remove_prefix(name):
+        if name.startswith(Symbol._constexpr_prefix()):
+            return name.removeprefix(Symbol._constexpr_prefix())
+
+        if name.startswith(Symbol._meta_prefix()):
+            return name.removeprefix(Symbol._meta_prefix())
 
     @staticmethod
     def _create_constexpr(name):
-        return f"_ninetoothed_constexpr_{name}"
+        return f"{Symbol._constexpr_prefix()}{name}"
 
     @staticmethod
     def _create_meta(name):
-        return f"_ninetoothed_meta_{name}"
+        return f"{Symbol._meta_prefix()}{name}"
+
+    @staticmethod
+    def _constexpr_prefix():
+        return f"{Symbol._ninetoothed_prefix()}constexpr_"
+
+    @staticmethod
+    def _meta_prefix():
+        return f"{Symbol._ninetoothed_prefix()}meta_"
+
+    @staticmethod
+    def _ninetoothed_prefix():
+        return "_ninetoothed_"
 
 
 class _FindAndReplacer(ast.NodeTransformer):

{ninetoothed-0.3.0.dist-info → ninetoothed-0.4.0.dist-info}/METADATA

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

@@ -0,0 +1,10 @@
+ninetoothed/__init__.py,sha256=T5UJXlC-wbo8JKPbLUNT65Kccp12xP52WFV5FsugETI,147
+ninetoothed/jit.py,sha256=ECjaHcrVNj1SBxoMdzjGi5iDp3rtv2jUiHjvK0eU6Cs,16188
+ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+ninetoothed/symbol.py,sha256=8Wg-JQPkVv9mMIxB1Rj4SHzOytHXPgHLkuK0BEFPDkc,5243
+ninetoothed/tensor.py,sha256=_DrjOJ-pBvEbSNUvUoYJduLQXmuKgNcqhe4xUDMVoZw,6275
+ninetoothed/torchifier.py,sha256=8M2PDwyFIfVypX6Z-Vt_bGbsCPqxqKnftL0rXeh9bOM,911
+ninetoothed-0.4.0.dist-info/METADATA,sha256=Wgg0CP-j8VkiJWMpyOLOL7C1kVLkeF4OoZD6eyZsgLQ,6720
+ninetoothed-0.4.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+ninetoothed-0.4.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.4.0.dist-info/RECORD,,

ninetoothed-0.3.0.dist-info/RECORD

@@ -1,10 +0,0 @@
-ninetoothed/__init__.py,sha256=T5UJXlC-wbo8JKPbLUNT65Kccp12xP52WFV5FsugETI,147
-ninetoothed/jit.py,sha256=nhjZRi8_kcjWZX0eOrnxLlzJfVg5vn12f9oi0Er2ABE,15515
-ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
-ninetoothed/symbol.py,sha256=Bd54qcI8KQAX0JRE_wPXycswtdSofhZ6Rr5MtZcv9fo,4665
-ninetoothed/tensor.py,sha256=_DrjOJ-pBvEbSNUvUoYJduLQXmuKgNcqhe4xUDMVoZw,6275
-ninetoothed/torchifier.py,sha256=8M2PDwyFIfVypX6Z-Vt_bGbsCPqxqKnftL0rXeh9bOM,911
-ninetoothed-0.3.0.dist-info/METADATA,sha256=CqdtfdV0eHzSwxJmFpD2IG5d4WTc6RDlpqMZue4Ml2Q,6720
-ninetoothed-0.3.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-ninetoothed-0.3.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ninetoothed-0.3.0.dist-info/RECORD,,