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

ninetoothed/aot.py

@@ -31,7 +31,7 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
 
     _HEADER_PATH.parent.mkdir(exist_ok=True)
 
-    if not _HEADER_PATH.exists():
+    if not _HEADER_PATH.exists() or _HEADER_PATH.read_text() != _HEADER_CONTENT:
         _HEADER_PATH.write_text(_HEADER_CONTENT)
 
     code_generator = CodeGenerator()
@@ -91,20 +91,29 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
 
     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{unparser.header};\n"
+    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
 
     return output_contents
 
 
-_HEADER_CONTENT = """#include <stdint.h>
+_HEADER_CONTENT = """#ifndef NINETOOTHED_H
+#define NINETOOTHED_H
+
+#include <stdint.h>
 
 typedef struct {
-    uintptr_t data;
+    void *data;
     uint64_t *shape;
     int64_t *strides;
 } NineToothedTensor;
+
+typedef void *NineToothedStream;
+
+typedef int NineToothedResult;
+
+#endif // NINETOOTHED_H
 """
 
 _HEADER_PATH = CACHE_DIR / "ninetoothed.h"
@@ -135,9 +144,9 @@ class _Unparser:
         return f"return {self._generic_unparse(call)};"
 
     def _unparse_FunctionDef(self, node):
-        params = ["CUstream stream"]
+        params = ["NineToothedStream stream"]
         params += [f"NineToothedTensor {arg.arg}" for arg in node.args.args]
-        header = f"CUresult {node.name}({', '.join(params)})"
+        header = f"NineToothedResult {node.name}({', '.join(params)})"
 
         self.header = header
 

ninetoothed/generation.py

@@ -19,6 +19,7 @@ import uuid
 import sympy
 import triton
 import triton.language as tl
+from triton.language.extra import libdevice
 
 import ninetoothed.naming as naming
 from ninetoothed.cudaifier import Cudaifier
@@ -225,6 +226,41 @@ class CodeGenerator(ast.NodeTransformer):
 
         return node
 
+    def visit_Call(self, node):
+        def _offsets(tensor, dim=None):
+            if dim is None:
+                return tensor._last_generated_overall_offsets.node
+
+            offsets = tensor._last_generated_offsets
+
+            if dim < 0:
+                dim += tensor.source.ndim
+
+            return sum(
+                offsets[dim][target_dim] for target_dim in range(tensor.target.ndim)
+            ).node
+
+        func = node.func
+        args = node.args
+
+        if isinstance(func, ast.Attribute):
+            if func.attr == "offsets":
+                value = func.value
+
+                if self._in_context(value):
+                    tensor = self._context[value.id]
+                elif isinstance(value, ast.Subscript) and self._in_context(value.value):
+                    tensor = self._context[value.value.id]
+
+                self.visit(value)
+
+                # TODO: Add error handling.
+                return _offsets(tensor, ast.literal_eval(args[0]) if args else None)
+
+        self.generic_visit(node)
+
+        return node
+
     def visit_Subscript(self, node):
         if self._in_context(node.value) and isinstance(node.ctx, ast.Load):
             value = self._context[node.value.id]
@@ -242,13 +278,24 @@ class CodeGenerator(ast.NodeTransformer):
         return node
 
     def visit_Attribute(self, node):
-        if self._in_context(node.value):
-            value = self._context[node.value.id]
+        value = node.value
 
-            if isinstance(value, Tensor):
-                inner = value.dtype
+        if isinstance(value, ast.Attribute):
+            value = self.visit_Attribute(value)
+
+        if self._in_context(value):
+            value = self._context[value.id].dtype
+
+        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
 
-                return Symbol(getattr(inner, node.attr)).node
+            if isinstance(attr, Tensor):
+                return attr
+
+            return Symbol(attr).node
 
         self.generic_visit(node)
 
@@ -560,6 +607,8 @@ class CodeGenerator(ast.NodeTransformer):
         indices = self._complete_indices(tensor, indices)
         offsets = type(self)._generate_offsets(tensor, indices)
 
+        tensor._last_generated_offsets = offsets
+
         for source_dim in range(tensor.source.ndim):
             for target_dim in range(tensor.target.ndim):
                 if target_dim not in invariant_target_dims:
@@ -584,7 +633,7 @@ class CodeGenerator(ast.NodeTransformer):
                     * tensor.source.strides[source_dim]
                 )
 
-        pointers = name_for_pointers + sum(
+        overall_offsets = sum(
             offsets[source_dim][target_dim][
                 type(self)._generate_slices(tensor, target_dim)
             ]
@@ -594,6 +643,10 @@ class CodeGenerator(ast.NodeTransformer):
             if target_dim not in invariant_target_dims
             and offsets[source_dim][target_dim] != 0
         )
+
+        tensor._last_generated_overall_offsets = overall_offsets
+
+        pointers = name_for_pointers + overall_offsets
         mask = functools.reduce(
             lambda x, y: x & y,
             (
@@ -980,6 +1033,9 @@ class _Inliner(ast.NodeTransformer):
         if func_def is None:
             return None, []
 
+        if inspect.getmodule(func) is libdevice:
+            return None, []
+
         collector = _ImportCollector()
         collector.visit(ast.parse(inspect.getsource(inspect.getmodule(func))))
         self.imports.extend(collector.imports)

ninetoothed/language.py

@@ -1,7 +1,11 @@
 import ast
 
+from triton.language.extra import libdevice
+
 from ninetoothed.symbol import Symbol
 
+__all__ = ["libdevice"]
+
 LANGUAGE = "ninetoothed.language"
 
 

ninetoothed/visualization.py

@@ -10,8 +10,6 @@ def visualize(tensor, color=None, save_path=None):
     :param color: The color to be used for visualization.
     :param save_path: The path where the visualization should be saved.
     """
-    outline_width = 0.1
-    plt.rcParams["lines.linewidth"] = 72 * outline_width
 
     if color is None:
         color = f"C{visualize.count}"
@@ -21,6 +19,24 @@ def visualize(tensor, color=None, save_path=None):
     width = max_pos_y + 1
     height = max_pos_x + 1
 
+    _, ax = _prepare_figure_and_axes(width, height)
+
+    _visualize_tensor(ax, tensor, 0, 0, color)
+
+    plt.savefig(save_path, transparent=True, bbox_inches="tight", pad_inches=0)
+
+    plt.close()
+
+    visualize.count += 1
+
+
+visualize.count = 0
+
+
+def _prepare_figure_and_axes(width, height):
+    outline_width = 0.1
+    plt.rcParams["lines.linewidth"] = 72 * outline_width
+
     fig = plt.figure(figsize=(width + outline_width, height + outline_width))
 
     h = (Size.Fixed(0), Size.Fixed(width + outline_width))
@@ -41,16 +57,7 @@ def visualize(tensor, color=None, save_path=None):
     plt.xlim((-half_outline_width, width + half_outline_width))
     plt.ylim((-half_outline_width, height + half_outline_width))
 
-    _visualize_tensor(ax, tensor, 0, 0, color)
-
-    plt.savefig(save_path, transparent=True, bbox_inches="tight", pad_inches=0)
-
-    plt.close()
-
-    visualize.count += 1
-
-
-visualize.count = 0
+    return fig, ax
 
 
 def _visualize_tensor(ax, tensor, x, y, color, level_spacing=4):

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

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

@@ -1,18 +1,18 @@
 ninetoothed/__init__.py,sha256=F2bxRNhzcGdtADA8RehTuf-QK0xnxno8kxvr6H2L5Tg,552
-ninetoothed/aot.py,sha256=8ZCLtnsign14YvY7SXX5ASidhuUAhPwppTXUJNkQup4,6243
+ninetoothed/aot.py,sha256=b7ykTC5roe_xg3NkZv6VyInBrEiNRwjpixCULUPRuEg,6506
 ninetoothed/cudaifier.py,sha256=5ylMr1q0B9NwbeXkpCu3o2nMGpDfh65nAQ0Az_qMQuI,877
 ninetoothed/dtype.py,sha256=-0iBleay5gYA4wtT3l17QjCesr7g26M6CSfhNJdI3k4,165
-ninetoothed/generation.py,sha256=VIqSyZT4yHxY_a2QPmWW6jjALv3e1mohDqdRQBRYsAo,36462
+ninetoothed/generation.py,sha256=wf8BL-x0PR6rG-9OSpgIZi8LtsIdFbqRUFiQFE5FIno,38107
 ninetoothed/jit.py,sha256=CpeSkO_zUe9DwtTJ2K2H7Bwpx-FvIHfrgzOcEosfpek,2946
-ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+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/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
 ninetoothed/utils.py,sha256=mtRXABBVPnlgd2n1REh9oB3s_5bUsKhd3iwu3oJ5DSQ,338
-ninetoothed/visualization.py,sha256=zlMH-0WplaboePGzcbpcj4UovpX0k2r4SysSPsNS4r4,3674
-ninetoothed-0.16.0.dist-info/METADATA,sha256=nkq3iImebtmcEs-bZq2zfF2_QxrZD9IWky1S86OnUMA,7340
-ninetoothed-0.16.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ninetoothed-0.16.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ninetoothed-0.16.0.dist-info/RECORD,,
+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,,