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

ninetoothed/jit.py

@@ -1,9 +1,11 @@
 import ast
 import collections
 import functools
+import importlib.util
 import inspect
 import itertools
 import math
+import sys
 import tempfile
 
 import triton
@@ -41,24 +43,20 @@ class JIT:
         ast.fix_missing_locations(tree)
 
         unparsed = ast.unparse(tree).replace("None:", ":").replace(":None", ":")
+        dependencies = self._find_dependencies()
+        source = "\n\n".join((unparsed, dependencies)).strip()
 
         with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as temp_file:
-            temp_file.write(unparsed.encode("utf-8"))
+            temp_file.write(source.encode("utf-8"))
             temp_file_name = temp_file.name
 
-        with open(temp_file_name, "r") as temp_file:
-            code = compile(
-                source=temp_file.read(),
-                filename=temp_file_name,
-                mode="exec",
-            )
-
-        namespace = {}
-        exec(code, namespace)
+        module = type(self)._import_from_path(temp_file_name, temp_file_name)
+        module_vars = vars(module)
 
         handle = _Handle(
-            namespace[self.func.__name__],
-            namespace[f"launch_{self.func.__name__}"],
+            module_vars[self.func.__name__],
+            module_vars[f"launch_{self.func.__name__}"],
+            source,
         )
 
         type(self).handles[source_file][source_line] = handle
@@ -74,6 +72,24 @@ class JIT:
 
         return ast.Module(body=[finder.result], type_ignores=[])
 
+    def _find_dependencies(self):
+        dependencies = set()
+
+        for obj in self.func.__globals__.values():
+            if isinstance(obj, triton.runtime.JITFunction):
+                dependencies.add(obj.src)
+
+        return "\n".join(f"@triton.jit\n{dependency}" for dependency in dependencies)
+
+    @staticmethod
+    def _import_from_path(module_name, file_path):
+        spec = importlib.util.spec_from_file_location(module_name, file_path)
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        spec.loader.exec_module(module)
+
+        return module
+
 
 class CodeGenerator(ast.NodeTransformer):
     def __init__(self, context):
@@ -102,7 +118,7 @@ class CodeGenerator(ast.NodeTransformer):
         self.generic_visit(node)
 
         for arg in self._args:
-            if not isinstance(arg, Tensor):
+            if not isinstance(arg, Tensor) or arg.ndim == 0:
                 continue
 
             offsets = arg.offsets()
@@ -178,7 +194,7 @@ class CodeGenerator(ast.NodeTransformer):
             if isinstance(value, Tensor):
                 inner = value.dtype
 
-                return Symbol(inner.__dict__[node.attr]).node
+                return Symbol(getattr(inner, node.attr)).node
 
         self.generic_visit(node)
 
@@ -355,6 +371,9 @@ class CodeGenerator(ast.NodeTransformer):
 
     @staticmethod
     def _generate_load(tensor, intermediate_indices=()):
+        if tensor.ndim == 0:
+            return Symbol(tensor.original.name).node
+
         pointers, mask = CodeGenerator._generate_pointers_and_mask(
             tensor, intermediate_indices
         )
@@ -459,9 +478,10 @@ class Tritonizer(ast.NodeTransformer):
 
 
 class _Handle:
-    def __init__(self, kernel, launch):
+    def __init__(self, kernel, launch, source):
         self._kernel = kernel
         self._launch = launch
+        self._source = source
 
     def __call__(self, *args, **kwargs):
         return self._launch(*args, **kwargs)

ninetoothed/tensor.py

@@ -21,11 +21,11 @@ class Tensor:
 
         self.dtype = dtype
 
-        self.name = f"tensor_{type(self).num_instances}"
+        self.name = f"_ninetoothed_tensor_{type(self).num_instances}"
 
         if ndim is not None:
-            self.shape = [Symbol(self.size_string(i)) for i in range(ndim)]
-            self.strides = [Symbol(self.stride_string(i)) for i in range(ndim)]
+            self.shape = (Symbol(self.size_string(i)) for i in range(ndim))
+            self.strides = (Symbol(self.stride_string(i)) for i in range(ndim))
         else:
             self.shape = shape
 
@@ -103,6 +103,9 @@ class Tensor:
         )
 
     def names(self):
+        if self.ndim == 0:
+            return {self.original.name}
+
         return (
             {self.original.pointer_string()}
             | {
@@ -191,6 +194,22 @@ class Tensor:
 
         return self.strides[dim]
 
+    @property
+    def shape(self):
+        return self._shape
+
+    @shape.setter
+    def shape(self, value):
+        self._shape = tuple(value)
+
+    @property
+    def strides(self):
+        return self._strides
+
+    @strides.setter
+    def strides(self, value):
+        self._strides = tuple(value)
+
     @property
     def ndim(self):
         return len(self.shape)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: ninetoothed
-Version: 0.4.0
+Version: 0.6.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
@@ -15,6 +15,8 @@ Description-Content-Type: text/markdown
 
 # NineToothed
 
+![NineToothed Logo](docs/source/_static/ninetoothed-logo.png)
+
 A domain-specific language (DSL) based on Triton but providing higher-level abstractions.
 
 **Other language versions: [English](README.md), [简体中文](docs/README.zh.md).**
@@ -80,4 +82,4 @@ def matmul_kernel(a: a_tiled, b: b_tiled, c: c_tiled):
 
 For matrix multiplication, we also have three tensor parameters, but the tiling method is more complex than vector addition. We denote the three matrices as $A$, $B$, and $C$, where $A$ and $B$ are inputs, and $C$ is the output. Tiling $C$ is simple; we just need to divide it into blocks of size `(BLOCK_SIZE_M, BLOCK_SIZE_N)` by rows and columns. Once each block computes its result, the entire $C$ is computed. However, how should we tile $A$ and $B$? The answer is to introduce another meta-parameter `BLOCK_SIZE_K`. This way, we can divide $A$ into blocks of size `(BLOCK_SIZE_M, BLOCK_SIZE_K)` and $B$ into blocks of size `(BLOCK_SIZE_K, BLOCK_SIZE_N)`. However, for matrix multiplication, $A$ and $B$ do not correspond block by block; each row of $A$ needs to correspond to each column of $B$. Therefore, we need to further `tile` $A$ and $B$ by rows and columns, respectively. Up to this point, we have a set of row blocks of $A$ and column blocks of $B$. However, each row block of $A$ must correspond to every column block of $B$. This is where `expand` comes in. We `expand` the row blocks of $A$ along the columns to the number of columns of $C$ and the column blocks of $B$ along the rows to the number of rows of $C$. This way, we successfully tile $A$, $B$, and $C$. In fact, our meta-operations up to this point have already enabled us to write kernel functions. However, we notice that the levels where the row blocks and column blocks reside, which we mentioned earlier, are two-dimensional, and their sizes are of the forms `(1, ...)` and `(..., 1)`. This means that if no other operations are performed, the way we access row blocks and column blocks would have to be `a[0, k]` and `b[k, 0]`. If we want to use `a` to find the range of `k`, we would need to use `a.shape[1]`, but we know that dimensions of size `1` can actually be removed completely. This is why we added two lines of `squeeze`. The `dtype` refers to the data type, which in PyTorch can generally be some integer or floating-point type, such as `torch.float32`. However, since meta-operations like `tile` can be performed in NineToothed, `dtype` can also be a `Tensor`. In other words, there is a concept of "tensors that store tensors" in NineToothed. In summary, these two lines perform operations on the tensors stored in the outmost tensor, removing the dimensions of size `1`. This way, when we access the row and column blocks, we can use `a[k]` and `b[k]`, and when finding the range of `k`, we can use `a.shape[0]`.
 
-With tiling done, the rest is simple. In the function body, we define an `accumulator` to accumulate intermediate results. We then iterate through the corresponding row blocks of $A$ and column blocks of B, multiplying them and accumulating the results in `accumulator`. Finally, we place the `accumulator` in the corresponding block of $C$. Since each block of the parameter tensors undergoes this operation, the multiplication is completed for the whole tensors as well.
+With tiling done, the rest is simple. In the function body, we define an `accumulator` to accumulate intermediate results. We then iterate through the corresponding row blocks of $A$ and column blocks of $B$, multiplying them and accumulating the results in `accumulator`. Finally, we place the `accumulator` in the corresponding block of $C$. Since each block of the parameter tensors undergoes this operation, the multiplication is completed for the whole tensors as well.

ninetoothed-0.6.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+ninetoothed/__init__.py,sha256=T5UJXlC-wbo8JKPbLUNT65Kccp12xP52WFV5FsugETI,147
+ninetoothed/jit.py,sha256=5gNp4HixCkural_Ns3DxwT4LL3OUcG0ECj4NLjb-EYk,16959
+ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+ninetoothed/symbol.py,sha256=8Wg-JQPkVv9mMIxB1Rj4SHzOytHXPgHLkuK0BEFPDkc,5243
+ninetoothed/tensor.py,sha256=L-9LhwnM4uRtRvj3tqrzerUijEfKeTQvFBcmS1hQilI,6656
+ninetoothed/torchifier.py,sha256=8M2PDwyFIfVypX6Z-Vt_bGbsCPqxqKnftL0rXeh9bOM,911
+ninetoothed-0.6.0.dist-info/METADATA,sha256=zvY4nvKt7R8kWDYrGnApem_C07trLgOj1-7zXPfqD9U,6785
+ninetoothed-0.6.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+ninetoothed-0.6.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.6.0.dist-info/RECORD,,

ninetoothed-0.4.0.dist-info/RECORD

@@ -1,10 +0,0 @@
-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,,