ninetoothed 0.3.0__py3-none-any.whl → 0.5.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)
@@ -178,7 +178,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)
 
@@ -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/tensor.py

@@ -24,8 +24,8 @@ class Tensor:
         self.name = f"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
 
@@ -191,6 +191,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.3.0.dist-info → ninetoothed-0.5.0.dist-info}/METADATA

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

@@ -0,0 +1,10 @@
+ninetoothed/__init__.py,sha256=T5UJXlC-wbo8JKPbLUNT65Kccp12xP52WFV5FsugETI,147
+ninetoothed/jit.py,sha256=q-TRGF81rUEwV1TGDrew3ijwvzCWenR8EejZbYteZSI,16188
+ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+ninetoothed/symbol.py,sha256=8Wg-JQPkVv9mMIxB1Rj4SHzOytHXPgHLkuK0BEFPDkc,5243
+ninetoothed/tensor.py,sha256=UO79yYwHMfdqv32Ww2mtcl-ki1C9zInC0vBNwDtzlHU,6575
+ninetoothed/torchifier.py,sha256=8M2PDwyFIfVypX6Z-Vt_bGbsCPqxqKnftL0rXeh9bOM,911
+ninetoothed-0.5.0.dist-info/METADATA,sha256=ObwfQtwBk3x90adbQfiSo5wK11qUG9f4NdmunyjC--0,6785
+ninetoothed-0.5.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+ninetoothed-0.5.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.5.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,,