ninetoothed 0.17.0__py3-none-any.whl → 0.18.0__py3-none-any.whl

ninetoothed/aot.py

@@ -1,9 +1,11 @@
 import ast
 import pathlib
+import re
 import subprocess
 import tempfile
 import uuid
 
+import ninetoothed.naming as naming
 from ninetoothed.dtype import int64
 from ninetoothed.generation import CACHE_DIR, CodeGenerator
 from ninetoothed.tensor import Tensor
@@ -25,8 +27,10 @@ def aot(
 
 def _aot(func, caller, kernel_name, num_warps, num_stages):
     def _find_tensor_by_source_name(tensors, name):
+        name = naming.remove_prefixes(name)
+
         for tensor in tensors:
-            if tensor.source.name == name:
+            if naming.remove_prefixes(tensor.source.name) == name:
                 return tensor
 
     _HEADER_PATH.parent.mkdir(exist_ok=True)
@@ -49,11 +53,15 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
     kernel_func = code_generator.kernel_func
     launch_func = code_generator.launch_func
 
+    param_strings = ["stream"]
     param_types = []
+    constexpr_param_indices = []
 
     for arg in kernel_func.args.args:
         param = arg.arg
 
+        param_strings.append(param)
+
         if match := Tensor.pointer_pattern().fullmatch(param):
             source_name = match.group(0).removesuffix("_pointer")
             tensor = _find_tensor_by_source_name(tensors, source_name)
@@ -64,9 +72,23 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
             param_types.append(int64)
         elif Tensor.stride_pattern().fullmatch(param):
             param_types.append(int64)
+        else:
+            source_name = param
+            tensor = _find_tensor_by_source_name(tensors, source_name)
+            dtype = tensor.source.dtype
+
+            if tensor.constexpr:
+                param_types.append(f"{tensor.value}")
+                constexpr_param_indices.append(len(param_types) - 1)
+            else:
+                param_types.append(dtype)
 
     signature = ", ".join(param_types)
 
+    for index in sorted(set(constexpr_param_indices), reverse=True):
+        param_strings.pop(index + 1)
+        param_types.pop(index)
+
     grid_extractor = _GridExtractor()
     launch_func = grid_extractor.visit(launch_func)
     grid_extractor.visit(code_generator.raw_grid)
@@ -76,21 +98,26 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
         source_file, kernel_name, signature, grid, num_warps, num_stages
     )
 
-    unparser = _Unparser()
+    c_source_file_name = f"{kernel_name}.{signature_hash}.c"
+    c_source_file = output_contents[c_source_file_name]
+
+    c_header_file_name = f"{kernel_name}.{signature_hash}.h"
+    c_header_file = output_contents[c_header_file_name]
+
+    pattern = rf"\({', '.join(rf'(.*) {param}' for param in param_strings)}\)"
+    c_param_type_strings = re.search(pattern, c_header_file).groups()
+
+    unparser = _Unparser(c_param_type_strings)
 
     launch_func_unparsed = unparser.unparse(launch_func)
     launch_func_unparsed = launch_func_unparsed.replace(
         func.__name__, f"{kernel_name}_{signature_hash}"
     )
 
-    c_source_file_name = f"{kernel_name}.{signature_hash}.c"
-    c_source_file = output_contents[c_source_file_name]
     c_source_file = f"{c_source_file}\n{launch_func_unparsed}\n"
     c_source_file = c_source_file.replace("<stdint.h>", f'"{_HEADER_PATH}"')
     output_contents[c_source_file_name] = c_source_file
 
-    c_header_file_name = f"{kernel_name}.{signature_hash}.h"
-    c_header_file = output_contents[c_header_file_name]
     c_header_file = f'{c_header_file}\n#ifdef __cplusplus\nextern "C" {unparser.header};\n#else\n{unparser.header};\n#endif\n'
     c_header_file = c_header_file.replace("<stdint.h>", f'"{_HEADER_PATH}"')
     output_contents[c_header_file_name] = c_header_file
@@ -120,6 +147,9 @@ _HEADER_PATH = CACHE_DIR / "ninetoothed.h"
 
 
 class _Unparser:
+    def __init__(self, param_types):
+        self._param_types = param_types
+
     def unparse(self, node):
         method_name = "_unparse_" + node.__class__.__name__
 
@@ -137,11 +167,29 @@ class _Unparser:
     def _unparse_Call(self, node):
         call = ast.Call(
             func=node.func,
-            args=[ast.Name(id="stream", ctx=ast.Load())] + node.args,
+            args=[ast.Name(id="stream", ctx=ast.Load())]
+            + [
+                arg
+                for arg in node.args
+                if not isinstance(arg, ast.Name) or not naming.is_constexpr(arg.id)
+            ],
             keywords=[],
         )
 
-        return f"return {self._generic_unparse(call)};"
+        unparsed = f"return {self._generic_unparse(call)};"
+
+        pattern = rf"\((stream), {', '.join(r'([^,]*)' for _ in range(len(self._param_types) - 1))}\)"
+        args = re.search(pattern, unparsed).groups()
+
+        for i, (arg, type) in enumerate(zip(args, self._param_types)):
+            if i != 0 and "." not in arg:
+                new_arg = f"*({type} *){arg}.data"
+            else:
+                new_arg = f"({type}){arg}"
+
+            unparsed = unparsed.replace(arg, new_arg)
+
+        return unparsed
 
     def _unparse_FunctionDef(self, node):
         params = ["NineToothedStream stream"]
@@ -153,13 +201,16 @@ class _Unparser:
         body_lines = []
 
         for stmt in node.body:
+            if isinstance(stmt, ast.Assign):
+                continue
+
             stmt_unparsed = self.unparse(stmt)
 
             if isinstance(stmt, ast.Expr):
                 stmt_unparsed = stmt_unparsed.strip()
 
-                if not stmt_unparsed.endswith(";"):
-                    stmt_unparsed += ";"
+            if not stmt_unparsed.endswith(";"):
+                stmt_unparsed += ";"
 
             body_lines.append("    " + stmt_unparsed)
 

ninetoothed/generation.py

@@ -289,12 +289,12 @@ class CodeGenerator(ast.NodeTransformer):
         if isinstance(value, Tensor):
             attr = getattr(value, node.attr)
 
-            if node.attr == "dtype" and attr is None:
-                return Symbol(f"{value.source.pointer_string()}.type.element_ty").node
-
             if isinstance(attr, Tensor):
                 return attr
 
+            if node.attr == "dtype":
+                return Symbol(f"{value.source.pointer_string()}.type.element_ty").node
+
             return Symbol(attr).node
 
         self.generic_visit(node)
@@ -500,16 +500,19 @@ class CodeGenerator(ast.NodeTransformer):
             naming.remove_prefixes(param) for param in next_power_of_2_params
         ]
 
+        arg_names = [naming.remove_prefixes(arg.source.name) for arg in self._args]
+
+        arg_names += [
+            param
+            for param in non_next_power_of_2_constexpr_params_without_prefixes
+            if not Tensor.size_pattern().fullmatch(param) and param not in arg_names
+        ]
+
         launch = ast.FunctionDef(
             name=self.launch_func_name,
             args=ast.arguments(
                 posonlyargs=[],
-                args=[ast.arg(arg=arg.source.name) for arg in self._args]
-                + [
-                    ast.arg(arg=param)
-                    for param in non_next_power_of_2_constexpr_params_without_prefixes
-                    if not Tensor.size_pattern().fullmatch(param)
-                ],
+                args=[ast.arg(arg=name) for name in arg_names],
                 kwonlyargs=[],
                 defaults=[],
             ),

ninetoothed/tensor.py

@@ -32,6 +32,8 @@ class Tensor:
         strides=None,
         other=None,
         shape_options=None,
+        constexpr=None,
+        value=None,
         name=None,
         source=None,
         source_dims=None,
@@ -74,6 +76,21 @@ class Tensor:
 
         self.other = other
 
+        if constexpr and self.ndim != 0:
+            raise ValueError(
+                "`constexpr` can only be set for zero-dimensional tensors."
+            )
+
+        self.constexpr = constexpr
+
+        if self.constexpr:
+            self.name = naming.make_constexpr(self.name)
+
+        if not constexpr and value is not None:
+            raise ValueError("`value` can only be set for constexpr tensors.")
+
+        self.value = value
+
         if source is not None:
             self.source = source
         else:

{ninetoothed-0.17.0.dist-info → ninetoothed-0.18.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.17.0
+Version: 0.18.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.17.0.dist-info → ninetoothed-0.18.0.dist-info}/RECORD

@@ -1,18 +1,18 @@
 ninetoothed/__init__.py,sha256=F2bxRNhzcGdtADA8RehTuf-QK0xnxno8kxvr6H2L5Tg,552
-ninetoothed/aot.py,sha256=b7ykTC5roe_xg3NkZv6VyInBrEiNRwjpixCULUPRuEg,6506
+ninetoothed/aot.py,sha256=VLPFRNZgq82DumuVMi36_qptM5nkORzmhbP4uPa559Q,8173
 ninetoothed/cudaifier.py,sha256=5ylMr1q0B9NwbeXkpCu3o2nMGpDfh65nAQ0Az_qMQuI,877
 ninetoothed/dtype.py,sha256=-0iBleay5gYA4wtT3l17QjCesr7g26M6CSfhNJdI3k4,165
-ninetoothed/generation.py,sha256=wf8BL-x0PR6rG-9OSpgIZi8LtsIdFbqRUFiQFE5FIno,38107
+ninetoothed/generation.py,sha256=zbqRWvpa-1q44WuZV9S13DDAxvi4dai2AJ47ihjODsM,38150
 ninetoothed/jit.py,sha256=CpeSkO_zUe9DwtTJ2K2H7Bwpx-FvIHfrgzOcEosfpek,2946
 ninetoothed/language.py,sha256=ERiA4dpwiow2AT2xFeFWYg1KqlnBo6xxPGp8VZrP0Lk,574
 ninetoothed/make.py,sha256=fQKuRJL7HC2iGTAN323mlIWXz9Z3jotIoN68ur29Qlw,1834
 ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
 ninetoothed/symbol.py,sha256=lJo3NL2-T7tKbKjb6MCRLMemN94mqS3bIiG943P0Mbo,7454
-ninetoothed/tensor.py,sha256=gQEzHTcXqZVBFLc2YRfXTKxjxPWMxWN7fNl2BCfJwMs,14782
+ninetoothed/tensor.py,sha256=lK8s5-l5cqhM9FCWXMjTle9vA1Nass_92tvuHY8H3OM,15265
 ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
 ninetoothed/utils.py,sha256=mtRXABBVPnlgd2n1REh9oB3s_5bUsKhd3iwu3oJ5DSQ,338
 ninetoothed/visualization.py,sha256=oc3cA5qqT66_RoAs5D681SCxR5E5wgFwk95ZefdSfZU,3794
-ninetoothed-0.17.0.dist-info/METADATA,sha256=_V2M45nT4Yin-zs7hq5-yHlN6KwV5_zcA8afwXP8S-Q,7340
-ninetoothed-0.17.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ninetoothed-0.17.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ninetoothed-0.17.0.dist-info/RECORD,,
+ninetoothed-0.18.0.dist-info/METADATA,sha256=X4TvwcjVuB40X4jmCsRMee8Auj2mGyORYsbk81fd-G0,7340
+ninetoothed-0.18.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ninetoothed-0.18.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.18.0.dist-info/RECORD,,