ninetoothed 0.5.0__tar.gz → 0.7.0__tar.gz

{ninetoothed-0.5.0 → ninetoothed-0.7.0}/PKG-INFO

@@ -1,11 +1,10 @@
 Metadata-Version: 2.3
 Name: ninetoothed
-Version: 0.5.0
+Version: 0.7.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
 Author-email: Jiacheng Huang <huangjiacheng0709@outlook.com>
-License-File: LICENSE
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
@@ -51,6 +50,8 @@ def add_kernel(
 
 In this code, we first define `BLOCK_SIZE`, which is a `Symbol`. You can think of `"BLOCK_SIZE"` as its name. We see that `meta` is set to `True`, indicating to the compiler that it is a meta-parameter and its value can be determined by the compiler. The `Tensor(1)` constructs a one-dimensional tensor (vector), and `Tensor(1).tile((BLOCK_SIZE,))` means we want to create a vector and divide it into blocks of size `BLOCK_SIZE`. Suppose the size of this vector is `8192` and `BLOCK_SIZE` is `1024`, then the vector will be divided into `8` blocks, each of size `1024`.
 
+![Tiling of Vector Addition](docs/source/_static/vecadd-tiling.png)
+
 By using type annotations, we tell the compiler that we will have three tensor parameters, which will be divided into blocks, and `x`, `y`, and `z` are these blocks. It's important to understand that `x`, `y`, and `z` are the blocks, not the tensors themselves. In the function body, `x`, `y`, and `z` are also the blocks. The rest is straightforward (only one line `z = x + y` left, haha), we add each block of `x` and `y` and store it in `z`. Since each block of the parameter tensors undergoes this operation, the addition is completed for the whole tensors as well.
 
 ### Matrix Multiplication
@@ -82,4 +83,10 @@ 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]`.
 
+![Tiling of Matrix Multiplication](docs/source/_static/matmul-tiling.png)
+
 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.
+
+## License
+
+This project is distributed under the Apache-2.0 license. See the included [LICENSE](LICENSE) file for details.

{ninetoothed-0.5.0 → ninetoothed-0.7.0}/README.md

@@ -36,6 +36,8 @@ def add_kernel(
 
 In this code, we first define `BLOCK_SIZE`, which is a `Symbol`. You can think of `"BLOCK_SIZE"` as its name. We see that `meta` is set to `True`, indicating to the compiler that it is a meta-parameter and its value can be determined by the compiler. The `Tensor(1)` constructs a one-dimensional tensor (vector), and `Tensor(1).tile((BLOCK_SIZE,))` means we want to create a vector and divide it into blocks of size `BLOCK_SIZE`. Suppose the size of this vector is `8192` and `BLOCK_SIZE` is `1024`, then the vector will be divided into `8` blocks, each of size `1024`.
 
+![Tiling of Vector Addition](docs/source/_static/vecadd-tiling.png)
+
 By using type annotations, we tell the compiler that we will have three tensor parameters, which will be divided into blocks, and `x`, `y`, and `z` are these blocks. It's important to understand that `x`, `y`, and `z` are the blocks, not the tensors themselves. In the function body, `x`, `y`, and `z` are also the blocks. The rest is straightforward (only one line `z = x + y` left, haha), we add each block of `x` and `y` and store it in `z`. Since each block of the parameter tensors undergoes this operation, the addition is completed for the whole tensors as well.
 
 ### Matrix Multiplication
@@ -67,4 +69,10 @@ 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]`.
 
+![Tiling of Matrix Multiplication](docs/source/_static/matmul-tiling.png)
+
 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.
+
+## License
+
+This project is distributed under the Apache-2.0 license. See the included [LICENSE](LICENSE) file for details.

{ninetoothed-0.5.0 → ninetoothed-0.7.0}/docs/README.zh.md

@@ -36,6 +36,8 @@ def add_kernel(
 
 在这段代码当中，我们首先定义了 `BLOCK_SIZE`，它是一个 `Symbol`，我们可以把 `"BLOCK_SIZE"` 理解成它的名字。我们可以看到 `meta` 被设成了 `True`，这是在告诉编译器，它是一个元参数，可以由编译器决定它的取值。之后出现的 `Tensor(1)` 则是在构造一个一维的张量（向量），`Tensor(1).tile((BLOCK_SIZE,))` 的意思就是说，我们想要构造一个向量，并且把它分成大小为 `BLOCK_SIZE` 的块。假设这个向量的大小为 `8192`，而 `BLOCK_SIZE` 是 `1024`，那么这个向量就会被分成 `8` 块，每一块的大小都是 `1024`。
 
+![向量加法的 Tiling](source/_static/vecadd-tiling.png)
+
 我们通过类型标注的方式，告诉了编译器，我们将会有三个参数张量，并且每个参数张量，都会被按照这样的方式分块，而 `x`、`y`、`z` 就是被分成的块。这一点很重要，我们要意识到，`x`、`y`、`z` 是被分成的块，而不是被分块的张量本身，并且函数体当中的 `x`、`y`、`z` 也都是被分成的块。剩下的就很好理解了（也就剩下 `z = x + y` 一行了，哈哈哈），我们把每一块 `x` 和 `y` 相加，放到了 `z` 中，由于参数张量被分成的每一块都被执行了这样的操作，因此即便对于整体而言，加法也被完成了。
 
 ### 矩阵乘法
@@ -67,4 +69,10 @@ def matmul_kernel(a: a_tiled, b: b_tiled, c: c_tiled):
 
 对于矩阵乘法来说，我们也有三个参数张量，但是分块的方式肯定比向量加法要复杂一些。我们将三个矩阵分别记作 $A$、$B$、$C$，其中 $A$ 和 $B$ 为输入，$C$ 为输出。其中 $C$ 的分块操作很简单，我们只需要按照行和列，将其分成大小为 `(BLOCK_SIZE_M, BLOCK_SIZE_N)` 的块即可，这样只要每个这样的块都算出了结果，整个 $C$ 也就都算出了结果。那么该如何分 $A$ 和 $B$ 呢？答案是再引入一个元参数 `BLOCK_SIZE_K`，这样我们就可以把 $A$ 分成 `(BLOCK_SIZE_M, BLOCK_SIZE_K)` 大小的块，把 $B$ 分成 `(BLOCK_SIZE_K, BLOCK_SIZE_N)` 的块。但是对于矩阵乘法，$A$ 和 $B$ 并不是块块对应，而是需要对应 $A$ 的每一行和 $B$ 的每一列，所以我们还需要继续 `tile`，把 $A$ 和 $B$ 进一步分成以行为单位和以列为单位的块。到目前为止，我们有了一堆 $A$ 的行块和 $B$ 的列块，但是对于每一个 $A$ 的行块，我们都要对应 $B$ 的每一个列块。这个时候，我们就需要进行 `expand` 了，我们把 $A$ 的行块沿着列 `expand` 成 $C$ 的列数那么多列，把 $B$ 的列块沿着行 `expand` 成 $C$ 的行数那么多行。这样，我们就成功地将 $A$、$B$、$C$ 三者都分好了块，并且对于每一个 $C$ 的块，我们都有对应好的 $A$ 的行块和 $B$ 的列块。其实我们的元操作到此为止，已经能够编写出核函数了，但是我们发现，刚才所提到的行块和列块所在的层级，是二维的，而且大小是 `(1, ...)` 和 `(..., 1)` 这样的形式。也就是说，如果不进行其它操作，那么我们访问行块和列块的方式就得是 `a[0, k]` 和 `b[k, 0]`，如果我们想要依靠 `a` 找到 `k` 的范围，那就得是 `a.shape[1]`。但是我们知道，大小为 `1` 的维度，其实完全可以被去掉，这就是为什么我们加了两行 `squeeze`，其中的 `dtype` 是数据类型的意思，在 PyTorch 中一般可以是某些整数类型或者浮点类型之类的，比如 `torch.float32`，但是由于九齿当中可以进行 `tile` 等元操作，所以 `dtype` 也可以是 `Tensor`。也就是说，在九齿当中，存在着“存储张量的张量”这样的概念。总而言之，这两行就是对最外层张量所存储的下一层的张量进行操作，把大小为 `1` 的维度去掉了，这样，我们在访问行块和列块时就可以使用 `a[k]` 和 `b[k]`，找 `k` 的范围时也可以使用 `a.shape[0]` 了。
 
+![矩阵乘法的 Tiling](source/_static/matmul-tiling.png)
+
 对应好了分块，后续的部分就简单多了。在函数体当中，我们定义了一个 `accumulator`，用于累加中间结果，之后就遍历了对应好的 $A$ 的行块和 $B$ 的列块，并且把他们相乘的结果累加到了 `accumulator` 当中，最后再将 `accumulator` 放到了对应的 $C$ 的分块当中。由于参数张量被分成的每一块都被执行了这样的操作，因此即便对于整体而言，乘法也被完成了。
+
+## 许可证
+
+本项目采用 Apache-2.0 许可证发布。详情请参见随附的 [LICENSE](LICENSE) 文件。

{ninetoothed-0.5.0 → ninetoothed-0.7.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ninetoothed"
-version = "0.5.0"
+version = "0.7.0"
 authors = [{ name = "Jiacheng Huang", email = "huangjiacheng0709@outlook.com" }]
 description = "A domain-specific language based on Triton but providing higher-level abstraction."
 readme = "README.md"

{ninetoothed-0.5.0 → ninetoothed-0.7.0}/src/ninetoothed/jit.py

@@ -1,29 +1,41 @@
 import ast
 import collections
 import functools
+import importlib.util
 import inspect
 import itertools
 import math
+import subprocess
+import sys
 import tempfile
 
 import triton
 
+import ninetoothed.naming as naming
 from ninetoothed.language import attribute, call
 from ninetoothed.symbol import Symbol
 from ninetoothed.tensor import Tensor
 from ninetoothed.torchifier import Torchifier
 
 
-def jit(func):
-    return JIT(func)()
+def jit(_func=None, *, _prettify=False):
+    def wrapper(func):
+        return JIT(func, _prettify=_prettify)()
+
+    if _func is None:
+        return wrapper
+
+    return wrapper(_func)
 
 
 class JIT:
     handles = collections.defaultdict(dict)
 
-    def __init__(self, func):
+    def __init__(self, func, _prettify=False):
         self.func = func
 
+        self._prettify = _prettify
+
     def __call__(self):
         source_file = inspect.getsourcefile(self.func)
         source_line = inspect.getsourcelines(self.func)[1]
@@ -38,27 +50,37 @@ class JIT:
 
         CodeGenerator(inspect.get_annotations(self.func)).visit(tree)
         Tritonizer().visit(tree)
+        _BinOpSimplifier().visit(tree)
         ast.fix_missing_locations(tree)
 
+        if self._prettify:
+            name_collector = _SimplifiedNameCollector()
+            name_collector.visit(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_name = temp_file.name
+        if self._prettify:
+            for original, simplified in name_collector.simplified_names.items():
+                if simplified not in name_collector.simplified_names:
+                    source = source.replace(original, simplified)
 
-        with open(temp_file_name, "r") as temp_file:
-            code = compile(
-                source=temp_file.read(),
-                filename=temp_file_name,
-                mode="exec",
+            source = subprocess.check_output(
+                ["ruff", "format", "-"], input=source, encoding="utf-8"
             )
 
-        namespace = {}
-        exec(code, namespace)
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as temp_file:
+            temp_file.write(source.encode("utf-8"))
+            temp_file_name = temp_file.name
+
+        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
@@ -69,10 +91,30 @@ class JIT:
         module = ast.parse(inspect.getsource(inspect.getmodule(self.func)))
 
         _AliasRestorer().visit(module)
+        collector = _ImportCollector()
+        collector.visit(module)
         finder = _FunctionDefFinder(self.func.__name__)
         finder.visit(module)
 
-        return ast.Module(body=[finder.result], type_ignores=[])
+        return ast.Module(body=collector.imports + [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):
@@ -102,7 +144,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()
@@ -131,12 +173,17 @@ class CodeGenerator(ast.NodeTransformer):
 
         names_of_args = [arg.names() - {"ninetoothed"} for arg in self._args]
         names = functools.reduce(lambda x, y: x | y, names_of_args)
-        meta_names = {name for name in names if Symbol.is_meta(name)}
+        meta_names = {name for name in names if naming.is_meta(name)}
         non_meta_names = {name for name in names if name not in meta_names}
+        non_meta_names |= {
+            naming.make_next_power_of_2(name)
+            for name in non_meta_names
+            if naming.is_constexpr(name)
+        }
 
         node.args = [
             ast.arg(arg=name)
-            if not Symbol.is_constexpr(name)
+            if not naming.is_constexpr(name)
             else ast.arg(arg=name, annotation=attribute("constexpr").node)
             for name in non_meta_names
         ] + [
@@ -287,9 +334,20 @@ 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
+        non_next_power_of_2_constexpr_params = [
+            param
+            for param in params
+            if naming.is_constexpr(param) and not naming.is_next_power_of_2(param)
+        ]
+        non_next_power_of_2_constexpr_params_without_prefixes = [
+            naming.remove_prefixes(param)
+            for param in non_next_power_of_2_constexpr_params
+        ]
+        next_power_of_2_params = [
+            param for param in params if naming.is_next_power_of_2(param)
+        ]
+        next_power_of_2_params_without_prefixes = [
+            naming.remove_prefixes(param) for param in next_power_of_2_params
         ]
 
         launch = ast.FunctionDef(
@@ -297,17 +355,33 @@ class CodeGenerator(ast.NodeTransformer):
             args=ast.arguments(
                 posonlyargs=[],
                 args=[ast.arg(arg=arg.original.name) for arg in self._args]
-                + [ast.arg(arg=param) for param in constexpr_params_without_prefixes],
+                + [
+                    ast.arg(arg=param)
+                    for param in non_next_power_of_2_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()),
+                    value=ast.Name(id=param_without_prefixes, ctx=ast.Load()),
+                )
+                for param, param_without_prefixes in zip(
+                    non_next_power_of_2_constexpr_params,
+                    non_next_power_of_2_constexpr_params_without_prefixes,
+                )
+            ]
+            + [
+                ast.Assign(
+                    targets=[ast.Name(id=param, ctx=ast.Store())],
+                    value=Symbol(
+                        f"triton.next_power_of_2({param_without_prefixes})"
+                    ).node,
                 )
-                for param, param_without_prefix in zip(
-                    constexpr_params, constexpr_params_without_prefixes
+                for param, param_without_prefixes in zip(
+                    next_power_of_2_params,
+                    next_power_of_2_params_without_prefixes,
                 )
             ]
             + [
@@ -355,6 +429,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
         )
@@ -458,10 +535,41 @@ class Tritonizer(ast.NodeTransformer):
         return node
 
 
+class _BinOpSimplifier(ast.NodeTransformer):
+    def visit_BinOp(self, node):
+        self.generic_visit(node)
+
+        if isinstance(node.op, ast.Mult):
+            left = Symbol(node.left)
+            right = Symbol(node.right)
+
+            if left == 0 or right == 0:
+                return Symbol(0).node
+
+            if left == 1:
+                return node.right
+
+            if right == 1:
+                return node.left
+
+        return node
+
+
+class _SimplifiedNameCollector(ast.NodeVisitor):
+    def __init__(self):
+        self.simplified_names = {}
+
+    def visit_Name(self, node):
+        self.generic_visit(node)
+
+        self.simplified_names[node.id] = naming.remove_prefixes(node.id)
+
+
 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)
@@ -515,6 +623,23 @@ class _AliasRestorer(ast.NodeTransformer):
         return node
 
 
+class _ImportCollector(ast.NodeVisitor):
+    def __init__(self):
+        super().__init__()
+
+        self.imports = []
+
+    def visit_Import(self, node):
+        self.imports.append(node)
+
+        self.generic_visit(node)
+
+    def visit_ImportFrom(self, node):
+        self.imports.append(node)
+
+        self.generic_visit(node)
+
+
 class _FunctionDefFinder(ast.NodeVisitor):
     def __init__(self, name):
         self._name = name

ninetoothed-0.7.0/src/ninetoothed/naming.py

@@ -0,0 +1,50 @@
+import re
+
+
+def make_constexpr(name):
+    return _add_prefix(name, _CONSTEXPR)
+
+
+def make_meta(name):
+    return _add_prefix(name, _META)
+
+
+def make_next_power_of_2(name):
+    return _add_prefix(name, _NEXT_POWER_OF_2)
+
+
+def is_constexpr(name):
+    return _CONSTEXPR in _find_prefixes(name) or is_meta(name)
+
+
+def is_meta(name):
+    return _META in _find_prefixes(name)
+
+
+def is_next_power_of_2(name):
+    return _NEXT_POWER_OF_2 in _find_prefixes(name)
+
+
+def remove_prefixes(name):
+    return _PREFIX_PATTERN.sub("", name)
+
+
+_CONSTEXPR = "constexpr"
+
+_META = "meta"
+
+_NEXT_POWER_OF_2 = "next_power_of_2"
+
+_PREFIX_PATTERN = re.compile(r"ninetoothed_((?!_).*?)_prefix_")
+
+
+def _add_prefix(name, string):
+    return f"{_make_prefix(string)}{name}"
+
+
+def _make_prefix(string):
+    return f"ninetoothed_{string}_prefix_"
+
+
+def _find_prefixes(name):
+    return set(_PREFIX_PATTERN.findall(name))

{ninetoothed-0.5.0 → ninetoothed-0.7.0}/src/ninetoothed/symbol.py

@@ -1,7 +1,10 @@
 import ast
 import inspect
+import numbers
 import types
 
+import ninetoothed.naming as naming
+
 
 class Symbol:
     def __init__(self, expr, constexpr=None, meta=None):
@@ -28,18 +31,31 @@ class Symbol:
             if constexpr is False:
                 raise ValueError("Non-constexpr meta symbol is not supported.")
 
-            self._node.id = type(self)._create_meta(self._node.id)
+            self._node.id = naming.make_meta(self._node.id)
 
         if constexpr:
-            self._node.id = type(self)._create_constexpr(self._node.id)
+            self._node.id = naming.make_constexpr(self._node.id)
+
+    def __eq__(self, other):
+        if isinstance(self._node, ast.Constant):
+            if isinstance(other, Symbol) and isinstance(other._node, ast.Constant):
+                return self._node.value == other._node.value
+
+            if isinstance(other, numbers.Number):
+                return self._node.value == other
+
+        return False
+
+    def __hash__(self):
+        return id(self)
 
     def __add__(self, other):
         other = type(self)(other)
 
-        if isinstance(self._node, ast.Constant) and self._node.value == 0:
+        if self == 0:
             return other
 
-        if isinstance(other._node, ast.Constant) and other._node.value == 0:
+        if other == 0:
             return self
 
         return type(self)(ast.BinOp(left=self._node, op=ast.Add(), right=other._node))
@@ -47,19 +63,30 @@ class Symbol:
     def __radd__(self, other):
         return self.__add__(other)
 
-    def __mul__(self, other):
+    def __sub__(self, other):
         other = type(self)(other)
 
-        if isinstance(self._node, ast.Constant) and self._node.value == 0:
-            return type(self)(0)
+        if self == 0:
+            return -other
 
-        if isinstance(other._node, ast.Constant) and other._node.value == 0:
+        if other == 0:
+            return self
+
+        return type(self)(ast.BinOp(left=self._node, op=ast.Sub(), right=other._node))
+
+    def __rsub__(self, other):
+        return type(self)(other).__sub__(self)
+
+    def __mul__(self, other):
+        other = type(self)(other)
+
+        if self == 0 or other == 0:
             return type(self)(0)
 
-        if isinstance(self._node, ast.Constant) and self._node.value == 1:
+        if self == 1:
             return other
 
-        if isinstance(other._node, ast.Constant) and other._node.value == 1:
+        if other == 1:
             return self
 
         return type(self)(ast.BinOp(left=self._node, op=ast.Mult(), right=other._node))
@@ -136,40 +163,8 @@ class Symbol:
         return SliceSimplifier().visit(self._node)
 
     @staticmethod
-    def is_constexpr(name):
-        return name.startswith(Symbol._constexpr_prefix()) or Symbol.is_meta(name)
-
-    @staticmethod
-    def is_meta(name):
-        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"{Symbol._constexpr_prefix()}{name}"
-
-    @staticmethod
-    def _create_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_"
+    def is_name(object):
+        return isinstance(object, Symbol) and isinstance(object.node, ast.Name)
 
 
 class _FindAndReplacer(ast.NodeTransformer):

{ninetoothed-0.5.0 → ninetoothed-0.7.0}/src/ninetoothed/tensor.py

@@ -1,6 +1,7 @@
 import itertools
 import re
 
+import ninetoothed.naming as naming
 from ninetoothed.language import call
 from ninetoothed.symbol import Symbol
 
@@ -15,13 +16,15 @@ class Tensor:
         dtype=None,
         strides=None,
         other=None,
+        name=None,
         original=None,
     ):
-        type(self).num_instances += 1
-
         self.dtype = dtype
 
-        self.name = f"tensor_{type(self).num_instances}"
+        if name is not None:
+            self.name = name
+        else:
+            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))
@@ -41,29 +44,41 @@ class Tensor:
         else:
             self.original = self
 
-    def tile(self, tile_shape, tile_strides=None):
-        if tile_strides is None:
-            tile_strides = [1 for _ in tile_shape]
+        type(self).num_instances += 1
+
+    def tile(self, tile_shape, strides=None, dilation=None):
+        if strides is None:
+            strides = [-1 for _ in tile_shape]
+
+        if dilation is None:
+            dilation = [1 for _ in tile_shape]
 
         outer_shape = []
         outer_strides = []
         inner_shape = []
         inner_strides = []
 
-        for size, stride, tile_size, tile_stride in zip(
-            self.shape, self.strides, tile_shape, tile_strides
+        for self_size, self_stride, tile_size, stride, spacing in zip(
+            self.shape, self.strides, tile_shape, strides, dilation
         ):
             if tile_size == -1:
-                tile_size = size
+                tile_size = self_size
+
+            if stride == -1:
+                stride = tile_size
 
-            new_size = call("cdiv", size, tile_size)
+            new_size = (
+                call("cdiv", self_size - spacing * (tile_size - 1) - 1, stride) + 1
+                if stride != 0
+                else -1
+            )
             outer_shape.append(new_size)
 
-            new_stride = stride * tile_size // tile_stride
+            new_stride = self_stride * stride // spacing
             outer_strides.append(new_stride)
 
             inner_shape.append(tile_size)
-            next_stride = stride * tile_stride
+            next_stride = self_stride * spacing
             inner_strides.append(next_stride)
 
         return type(self)(
@@ -103,6 +118,9 @@ class Tensor:
         )
 
     def names(self):
+        if self.ndim == 0:
+            return {self.original.name}
+
         return (
             {self.original.pointer_string()}
             | {
@@ -112,6 +130,7 @@ class Tensor:
                 for name in value.names()
             }
             | (self.dtype.names() if isinstance(self.dtype, type(self)) else set())
+            | (self.original.names() if self.original is not self else set())
         )
 
     def offsets(self, indices=None):
@@ -160,7 +179,14 @@ class Tensor:
 
         if isinstance(curr, type(self)):
             for dim in range(curr.ndim):
-                indices.append(call("arange", 0, curr.shape[dim]))
+                size = curr.shape[dim]
+
+                if Symbol.is_name(size):
+                    name = size.node.id
+                    if not naming.is_meta(name):
+                        size = naming.make_next_power_of_2(name)
+
+                indices.append(call("arange", 0, size))
 
         return tuple(indices)
 
@@ -237,7 +263,7 @@ class Tensor:
     def _calculate_default_strides(shape):
         strides = [1]
 
-        for size in shape[1:]:
+        for size in reversed(shape[1:]):
             strides.append(size * strides[-1])
 
         return reversed(strides)

ninetoothed-0.7.0/tests/test_addmm.py

@@ -0,0 +1,104 @@
+import random
+
+import torch
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Symbol, Tensor
+from tests.skippers import skip_if_cuda_not_available, skip_if_float8_e5m2_not_supported
+
+
+def addmm(input, mat1, mat2, beta=1, alpha=1):
+    BLOCK_SIZE_M = Symbol("BLOCK_SIZE_M", meta=True)
+    BLOCK_SIZE_N = Symbol("BLOCK_SIZE_N", meta=True)
+    BLOCK_SIZE_K = Symbol("BLOCK_SIZE_K", meta=True)
+
+    input_tiled = Tensor(2).tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
+
+    output_tiled = Tensor(2).tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
+
+    mat1_tiled = (
+        Tensor(2)
+        .tile((BLOCK_SIZE_M, BLOCK_SIZE_K))
+        .tile((1, -1))
+        .expand((-1, output_tiled.shape[1]))
+    )
+    mat1_tiled.dtype = mat1_tiled.dtype.squeeze(0)
+
+    mat2_tiled = (
+        Tensor(2)
+        .tile((BLOCK_SIZE_K, BLOCK_SIZE_N))
+        .tile((-1, 1))
+        .expand((output_tiled.shape[0], -1))
+    )
+    mat2_tiled.dtype = mat2_tiled.dtype.squeeze(1)
+
+    @ninetoothed.jit
+    def addmm_kernel(
+        input: input_tiled,
+        mat1: mat1_tiled,
+        mat2: mat2_tiled,
+        beta: Tensor(0),
+        alpha: Tensor(0),
+        output: output_tiled,
+    ):
+        accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
+        for k in range(mat1.shape[0]):
+            accumulator += ntl.dot(mat1[k], mat2[k])
+        output = beta * input + alpha * accumulator.to(ntl.float16)
+
+    output = torch.empty(
+        (mat1.shape[0], mat2.shape[1]), device=mat1.device, dtype=torch.float16
+    )
+
+    addmm_kernel(input, mat1, mat2, beta, alpha, output)
+
+    return output
+
+
+@skip_if_cuda_not_available
+class TestCUDA:
+    @classmethod
+    def setup_class(cls):
+        torch.manual_seed(0)
+
+        shape = (512, 512)
+
+        cls.input = torch.randn(shape, device="cuda")
+        cls.mat1 = torch.randn(shape, device="cuda")
+        cls.mat2 = torch.randn(shape, device="cuda")
+        cls.beta = random.uniform(0, 1)
+        cls.alpha = random.uniform(0, 1)
+
+    def test_fp16(self):
+        input = type(self).input.to(torch.float16)
+        mat1 = type(self).mat1.to(torch.float16)
+        mat2 = type(self).mat2.to(torch.float16)
+        beta = type(self).beta
+        alpha = type(self).alpha
+
+        assert torch.allclose(
+            addmm(input, mat1, mat2, beta=beta, alpha=alpha),
+            torch.addmm(input, mat1, mat2, beta=beta, alpha=alpha),
+            atol=0.075,
+        )
+
+    @skip_if_float8_e5m2_not_supported
+    def test_fp8(self):
+        input = type(self).input.to(torch.float8_e5m2)
+        mat1 = type(self).mat1.to(torch.float8_e5m2)
+        mat2 = type(self).mat2.T.to(torch.float8_e5m2)
+        beta = type(self).beta
+        alpha = type(self).alpha
+
+        assert torch.allclose(
+            addmm(input, mat1, mat2, beta=beta, alpha=alpha),
+            torch.addmm(
+                input.to(torch.float16),
+                mat1.to(torch.float16),
+                mat2.to(torch.float16),
+                beta=beta,
+                alpha=alpha,
+            ),
+            atol=0.125,
+        )

ninetoothed-0.7.0/tests/test_naming.py

@@ -0,0 +1,51 @@
+import ninetoothed.naming as naming
+
+
+def test_make_constexpr():
+    assert naming.make_constexpr(_NAME) == f"ninetoothed_constexpr_prefix_{_NAME}"
+
+
+def test_make_meta():
+    assert naming.make_meta(_NAME) == f"ninetoothed_meta_prefix_{_NAME}"
+
+
+def test_make_next_power_of_2():
+    assert (
+        naming.make_next_power_of_2(_NAME)
+        == f"ninetoothed_next_power_of_2_prefix_{_NAME}"
+    )
+    assert (
+        naming.make_next_power_of_2(naming.make_constexpr(_NAME))
+        == f"ninetoothed_next_power_of_2_prefix_ninetoothed_constexpr_prefix_{_NAME}"
+    )
+
+
+def test_is_constexpr():
+    assert naming.is_constexpr(naming.make_constexpr(_NAME))
+    assert naming.is_constexpr(naming.make_meta(_NAME))
+
+
+def test_is_meta():
+    assert naming.is_meta(naming.make_meta(_NAME))
+
+
+def test_is_next_power_of_2():
+    assert naming.is_next_power_of_2(naming.make_next_power_of_2(_NAME))
+    assert naming.is_next_power_of_2(
+        naming.make_next_power_of_2(naming.make_constexpr(_NAME))
+    )
+
+
+def test_remove_prefixes():
+    assert naming.remove_prefixes(naming.make_constexpr(_NAME)) == _NAME
+    assert naming.remove_prefixes(naming.make_meta(_NAME)) == _NAME
+    assert naming.remove_prefixes(naming.make_next_power_of_2(_NAME)) == _NAME
+    assert (
+        naming.remove_prefixes(
+            naming.make_next_power_of_2(naming.make_constexpr(_NAME))
+        )
+        == _NAME
+    )
+
+
+_NAME = "ninetoothed_name"

{ninetoothed-0.5.0 → ninetoothed-0.7.0}/tests/test_softmax.py

@@ -1,5 +1,4 @@
 import torch
-import triton
 
 import ninetoothed
 import ninetoothed.language as ntl
@@ -22,7 +21,7 @@ def softmax(input):
 
     output = torch.empty_like(input)
 
-    softmax_kernel(input, output, BLOCK_SIZE=triton.next_power_of_2(input.shape[-1]))
+    softmax_kernel(input, output, BLOCK_SIZE=input.shape[-1])
 
     return output