ninetoothed 0.2.0__tar.gz → 0.4.0__tar.gz

{ninetoothed-0.2.0 → ninetoothed-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: ninetoothed
-Version: 0.2.0
+Version: 0.4.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
@@ -10,6 +10,7 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
+Requires-Dist: triton>=3.0.0
 Description-Content-Type: text/markdown
 
 # NineToothed
@@ -26,7 +27,7 @@ We can use `pip` to install `ninetoothed`.
 pip install ninetoothed
 ```
 
-After successfully running the above command, `ninetoothed` will be installed. However, to fully utilize its capabilities, you also need to install `triton` and a deep learning framework supported by `ninetoothed`. For trial purposes, we recommend installing `triton` and `torch`.
+After successfully running the above command, `ninetoothed` will be installed. However, to fully utilize its capabilities, you also need to install a deep learning framework supported by `ninetoothed`. For trial purposes, we recommend installing `torch`.
 
 ## Usage
 
@@ -64,16 +65,19 @@ c_tiled = Tensor(2).tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
 a_tiled = a_tiled.expand((-1, c_tiled.shape[1]))
 b_tiled = b_tiled.expand((c_tiled.shape[0], -1))
 
+a_tiled.dtype = a_tiled.dtype.squeeze(0)
+b_tiled.dtype = b_tiled.dtype.squeeze(1)
+
 @ninetoothed.jit
 def matmul_kernel(a: a_tiled, b: b_tiled, c: c_tiled):
     accumulator = ninetoothed.language.zeros(
         c.shape, dtype=ninetoothed.language.float32
     )
-    for k in range(a.shape[1]):
-        accumulator = ninetoothed.language.dot(a[0, k], b[k, 0], accumulator)
+    for k in range(a.shape[0]):
+        accumulator += ninetoothed.language.dot(a[k], b[k])
     c = accumulator.to(ninetoothed.language.float16)
 ```
 
-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$.
+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.

{ninetoothed-0.2.0 → ninetoothed-0.4.0}/README.md

@@ -12,7 +12,7 @@ We can use `pip` to install `ninetoothed`.
 pip install ninetoothed
 ```
 
-After successfully running the above command, `ninetoothed` will be installed. However, to fully utilize its capabilities, you also need to install `triton` and a deep learning framework supported by `ninetoothed`. For trial purposes, we recommend installing `triton` and `torch`.
+After successfully running the above command, `ninetoothed` will be installed. However, to fully utilize its capabilities, you also need to install a deep learning framework supported by `ninetoothed`. For trial purposes, we recommend installing `torch`.
 
 ## Usage
 
@@ -50,16 +50,19 @@ c_tiled = Tensor(2).tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
 a_tiled = a_tiled.expand((-1, c_tiled.shape[1]))
 b_tiled = b_tiled.expand((c_tiled.shape[0], -1))
 
+a_tiled.dtype = a_tiled.dtype.squeeze(0)
+b_tiled.dtype = b_tiled.dtype.squeeze(1)
+
 @ninetoothed.jit
 def matmul_kernel(a: a_tiled, b: b_tiled, c: c_tiled):
     accumulator = ninetoothed.language.zeros(
         c.shape, dtype=ninetoothed.language.float32
     )
-    for k in range(a.shape[1]):
-        accumulator = ninetoothed.language.dot(a[0, k], b[k, 0], accumulator)
+    for k in range(a.shape[0]):
+        accumulator += ninetoothed.language.dot(a[k], b[k])
     c = accumulator.to(ninetoothed.language.float16)
 ```
 
-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$.
+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.

{ninetoothed-0.2.0 → ninetoothed-0.4.0}/docs/README.zh.md

@@ -12,7 +12,7 @@
 pip install ninetoothed
 ```
 
-成功运行完以上两个命令之后，`ninetoothed` 就被安装好了。但是除了 `ninetoothed` 的本体之外，如果我们想要真正发挥它的作用，至少还需要安装 `triton` 和一个 `ninetoothed` 所支持的深度学习框架。以尝试为目的的话，我们推荐安装 `triton` 和 `torch`。
+成功运行完以上两个命令之后，`ninetoothed` 就被安装好了。但是除了 `ninetoothed` 的本体之外，如果我们想要真正发挥它的作用，至少还需要安装一个 `ninetoothed` 所支持的深度学习框架。以尝试为目的的话，我们推荐安装 `torch`。
 
 ## 使用
 
@@ -50,16 +50,19 @@ c_tiled = Tensor(2).tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
 a_tiled = a_tiled.expand((-1, c_tiled.shape[1]))
 b_tiled = b_tiled.expand((c_tiled.shape[0], -1))
 
+a_tiled.dtype = a_tiled.dtype.squeeze(0)
+b_tiled.dtype = b_tiled.dtype.squeeze(1)
+
 @ninetoothed.jit
 def matmul_kernel(a: a_tiled, b: b_tiled, c: c_tiled):
     accumulator = ninetoothed.language.zeros(
         c.shape, dtype=ninetoothed.language.float32
     )
-    for k in range(a.shape[1]):
-        accumulator = ninetoothed.language.dot(a[0, k], b[k, 0], accumulator)
+    for k in range(a.shape[0]):
+        accumulator += ninetoothed.language.dot(a[k], b[k])
     c = accumulator.to(ninetoothed.language.float16)
 ```
 
-对于矩阵乘法来说，我们也有三个参数张量，但是分块的方式肯定比向量加法要复杂一些。我们将三个矩阵分别记作 $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$ 的列块。
+对于矩阵乘法来说，我们也有三个参数张量，但是分块的方式肯定比向量加法要复杂一些。我们将三个矩阵分别记作 $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]` 了。
 
 对应好了分块，后续的部分就简单多了。在函数体当中，我们定义了一个 `accumulator`，用于累加中间结果，之后就遍历了对应好的 $A$ 的行块和 $B$ 的列块，并且把他们相乘的结果累加到了 `accumulator` 当中，最后再将 `accumulator` 放到了对应的 $C$ 的分块当中。由于参数张量被分成的每一块都被执行了这样的操作，因此即便对于整体而言，乘法也被完成了。

{ninetoothed-0.2.0 → ninetoothed-0.4.0}/pyproject.toml

@@ -4,10 +4,11 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ninetoothed"
-version = "0.2.0"
+version = "0.4.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"
+dependencies = ["triton>=3.0.0"]
 requires-python = ">=3.10"
 classifiers = [
     "Programming Language :: Python :: 3",

{ninetoothed-0.2.0 → ninetoothed-0.4.0}/src/ninetoothed/jit.py

@@ -6,6 +6,8 @@ import itertools
 import math
 import tempfile
 
+import triton
+
 from ninetoothed.language import attribute, call
 from ninetoothed.symbol import Symbol
 from ninetoothed.tensor import Tensor
@@ -103,13 +105,24 @@ class CodeGenerator(ast.NodeTransformer):
             if not isinstance(arg, Tensor):
                 continue
 
-            node.body.insert(
-                0,
-                ast.Assign(
-                    targets=[Symbol(f"{arg.name}_ptrs").node],
-                    value=arg.pointers().node,
-                ),
-            )
+            offsets = arg.offsets()
+
+            initializations = {
+                type(self)._name_for_offsets(arg, dim): offs
+                for dim, offs in enumerate(offsets)
+            } | {
+                type(self)._name_for_pointers(arg): arg.original.pointer_string()
+                + sum(
+                    type(self)._name_for_offsets(arg, dim)[
+                        type(self)._generate_slices(arg, dim)
+                    ]
+                    * stride
+                    for dim, stride in enumerate(arg.original.strides)
+                )
+            }
+
+            for target, value in reversed(initializations.items()):
+                node.body.insert(0, ast.Assign(targets=[target.node], value=value.node))
 
         return node
 
@@ -124,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)
@@ -147,15 +160,13 @@ class CodeGenerator(ast.NodeTransformer):
             value = self._context[node.value.id]
 
             if isinstance(value, Tensor):
-                pointers = type(self)._create_pointers(
+                return type(self)._generate_load(
                     value,
-                    node.slice.elts
+                    intermediate_indices=node.slice.elts
                     if isinstance(node.slice, ast.Tuple)
                     else (node.slice,),
                 )
 
-                return call("load", pointers).node
-
         self.generic_visit(node)
 
         return node
@@ -177,9 +188,7 @@ class CodeGenerator(ast.NodeTransformer):
         self.generic_visit(node)
 
         if node.id in self._context and isinstance(node.ctx, ast.Load):
-            return call(
-                "load", type(self)._create_pointers(self._context[node.id], ()).node
-            ).node
+            return type(self)._generate_load(self._context[node.id])
 
         return node
 
@@ -191,11 +200,7 @@ class CodeGenerator(ast.NodeTransformer):
                 self.generic_visit(node)
 
                 return ast.Expr(
-                    call(
-                        "store",
-                        type(self)._create_pointers(self._context[target.id], ()).node,
-                        node.value,
-                    ).node
+                    type(self)._generate_store(self._context[target.id], node.value)
                 )
             elif (
                 isinstance(target, ast.Subscript)
@@ -208,19 +213,14 @@ class CodeGenerator(ast.NodeTransformer):
                 if isinstance(value, Tensor):
                     self.generic_visit(node)
 
-                    pointers = type(self)._create_pointers(
-                        value,
-                        target.slice.elts
-                        if isinstance(target.slice, ast.Tuple)
-                        else (target.slice,),
-                    )
-
                     return ast.Expr(
-                        call(
-                            "store",
-                            pointers.node,
+                        type(self)._generate_store(
+                            value,
                             node.value,
-                        ).node
+                            intermediate_indices=target.slice.elts
+                            if isinstance(target.slice, ast.Tuple)
+                            else (target.slice,),
+                        )
                     )
 
         self.generic_visit(node)
@@ -228,6 +228,13 @@ class CodeGenerator(ast.NodeTransformer):
         return node
 
     def _generate_autotune(self, params, meta):
+        device = triton.runtime.driver.active.get_current_device()
+        properties = triton.runtime.driver.active.utils.get_device_properties(device)
+        max_shared_mem = properties["max_shared_mem"]
+
+        num_warps = 8
+        num_stages = max_shared_mem // 2**15
+
         configs = [
             ast.Call(
                 func=ast.Attribute(
@@ -241,7 +248,10 @@ class CodeGenerator(ast.NodeTransformer):
                         values=[ast.Constant(value=value) for value in values],
                     )
                 ],
-                keywords=[],
+                keywords=[
+                    ast.keyword(arg="num_warps", value=ast.Constant(value=num_warps)),
+                    ast.keyword(arg="num_stages", value=ast.Constant(value=num_stages)),
+                ],
             )
             for values in itertools.product(self._POWER_OF_TWOS, repeat=len(meta))
             if self._MIN_PRODUCT <= math.prod(values) <= self._MAX_PRODUCT
@@ -268,7 +278,7 @@ class CodeGenerator(ast.NodeTransformer):
                         elts=[
                             ast.Constant(value=param)
                             for param in params
-                            if not Tensor.is_pointer(param)
+                            if not Tensor.pointer_pattern().fullmatch(param)
                         ],
                         ctx=ast.Load(),
                     ),
@@ -277,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.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(
@@ -329,13 +354,76 @@ class CodeGenerator(ast.NodeTransformer):
         return ast.parse(f"lambda meta: ({num_elements},)", mode="eval").body
 
     @staticmethod
-    def _create_pointers(tensor, indices):
-        return Symbol(f"{tensor.name}_ptrs") + tensor.offsets(
-            [0 for _ in range(tensor.ndim())]
-            + list(indices)
-            + [0 for _ in range(tensor.inmost().ndim())]
+    def _generate_load(tensor, intermediate_indices=()):
+        pointers, mask = CodeGenerator._generate_pointers_and_mask(
+            tensor, intermediate_indices
+        )
+        other = CodeGenerator._generate_other(tensor)
+
+        return call("load", pointers, mask=mask, other=other).node
+
+    @staticmethod
+    def _generate_store(tensor, value, intermediate_indices=()):
+        pointers, mask = CodeGenerator._generate_pointers_and_mask(
+            tensor, intermediate_indices
+        )
+
+        return call("store", pointers, value, mask=mask).node
+
+    @staticmethod
+    def _generate_pointers_and_mask(tensor, intermediate_indices):
+        intermediate_offsets = CodeGenerator._generate_intermediate_offsets(
+            tensor, intermediate_indices
+        )
+        offsets = [
+            CodeGenerator._name_for_offsets(tensor, dim) + intermediate_offsets[dim]
+            for dim in range(tensor.original.ndim)
+        ]
+        pointers = CodeGenerator._name_for_pointers(tensor) + sum(
+            map(lambda x, y: x * y, intermediate_offsets, tensor.original.strides)
+        )
+        mask = functools.reduce(
+            lambda x, y: x & y,
+            (
+                offs[CodeGenerator._generate_slices(tensor, dim)] < size
+                for dim, (offs, size) in enumerate(zip(offsets, tensor.original.shape))
+            ),
+        )
+
+        return pointers, mask
+
+    @staticmethod
+    def _generate_other(tensor):
+        other = tensor.original.other
+
+        if isinstance(other, float) and not math.isfinite(other):
+            return f"float('{other}')"
+
+        return other
+
+    @staticmethod
+    def _generate_slices(tensor, dim):
+        return tuple(slice(None) if i == dim else None for i in range(tensor.ndim))
+
+    @staticmethod
+    def _generate_intermediate_offsets(tensor, intermediate_indices):
+        return tuple(
+            offs
+            for offs in tensor.offsets(
+                [0 for _ in range(tensor.ndim)]
+                + list(intermediate_indices)
+                + [0 for _ in range(tensor.inmost().ndim)]
+            )
         )
 
+    @staticmethod
+    def _name_for_pointers(tensor):
+        return Symbol(f"{tensor.original.name}_pointers")
+
+    @staticmethod
+    def _name_for_offsets(tensor, dim):
+        return Symbol(f"{tensor.original.name}_offsets_{dim}")
+
 
 class Tritonizer(ast.NodeTransformer):
     def visit_Module(self, node):

{ninetoothed-0.2.0 → ninetoothed-0.4.0}/src/ninetoothed/language.py

@@ -10,7 +10,10 @@ def call(func, *args, **kwargs):
         ast.Call(
             func=attribute(func).node,
             args=[Symbol(arg).node for arg in args],
-            keywords=[(kwarg, Symbol(kwargs[kwarg]).node) for kwarg in kwargs],
+            keywords=[
+                ast.keyword(arg=kwarg, value=Symbol(kwargs[kwarg]).node)
+                for kwarg in kwargs
+            ],
         )
     )
 

{ninetoothed-0.2.0 → ninetoothed-0.4.0}/src/ninetoothed/symbol.py

@@ -78,16 +78,36 @@ class Symbol:
         )
 
     def __mod__(self, other):
+        other = type(self)(other)
+
+        return type(self)(ast.BinOp(left=self._node, op=ast.Mod(), right=other._node))
+
+    def __lt__(self, other):
+        other = type(self)(other)
+
         return type(self)(
-            ast.BinOp(left=self._node, op=ast.Mod(), right=type(self)(other)._node)
+            ast.Compare(left=self._node, ops=[ast.Lt()], comparators=[other._node])
         )
 
+    def __and__(self, other):
+        other = type(self)(other)
+
+        return type(self)(
+            ast.BinOp(left=self._node, op=ast.BitAnd(), right=other._node)
+        )
+
+    def __rand__(self, other):
+        return self.__and__(other)
+
     def __getitem__(self, key):
         return type(self)(ast.Subscript(value=self._node, slice=type(self)(key)._node))
 
     def __repr__(self):
         return ast.unparse(self._node)
 
+    def find_and_replace(self, target, replacement):
+        _FindAndReplacer(target.node, replacement.node).visit(self._node)
+
     def names(self):
         class NameCollector(ast.NodeVisitor):
             def __init__(self):
@@ -117,16 +137,48 @@ 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):
+    def __init__(self, target, replacement):
+        self._target_id = target.id
+        self._replacement = replacement
+
+    def visit_Name(self, node):
+        if node.id == self._target_id:
+            return self._replacement
+
+        return self.generic_visit(node)

{ninetoothed-0.2.0 → ninetoothed-0.4.0}/src/ninetoothed/tensor.py

@@ -1,4 +1,5 @@
 import itertools
+import re
 
 from ninetoothed.language import call
 from ninetoothed.symbol import Symbol
@@ -7,19 +8,24 @@ from ninetoothed.symbol import Symbol
 class Tensor:
     num_instances = 0
 
-    def __init__(self, ndim=None, shape=None, dtype=None, strides=None, name=None):
+    def __init__(
+        self,
+        ndim=None,
+        shape=None,
+        dtype=None,
+        strides=None,
+        other=None,
+        original=None,
+    ):
         type(self).num_instances += 1
 
         self.dtype = dtype
 
-        if name is not None:
-            self.name = name
-        else:
-            self.name = f"tensor_{type(self).num_instances}"
+        self.name = f"tensor_{type(self).num_instances}"
 
         if ndim is not None:
-            self.shape = [Symbol(f"{self.name}_size_{i}") for i in range(ndim)]
-            self.strides = [Symbol(f"{self.name}_stride_{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
 
@@ -28,6 +34,13 @@ class Tensor:
             else:
                 self.strides = self._calculate_default_strides(shape)
 
+        self.other = other
+
+        if original is not None:
+            self.original = original
+        else:
+            self.original = self
+
     def tile(self, tile_shape, tile_strides=None):
         if tile_strides is None:
             tile_strides = [1 for _ in tile_shape]
@@ -59,10 +72,10 @@ class Tensor:
                 shape=inner_shape,
                 dtype=self.dtype,
                 strides=inner_strides,
-                name=self.name,
+                original=self.original,
             ),
             strides=outer_strides,
-            name=self.name,
+            original=self.original,
         )
 
     def expand(self, shape):
@@ -77,12 +90,21 @@ class Tensor:
                 stride if new_size == -1 else 0
                 for new_size, stride in zip(shape, self.strides)
             ],
-            name=self.name,
+            original=self.original,
+        )
+
+    def squeeze(self, dim):
+        # TODO: Add error handling.
+        return type(self)(
+            shape=[size for i, size in enumerate(self.shape) if dim != i],
+            dtype=self.dtype,
+            strides=[stride for i, stride in enumerate(self.strides) if dim != i],
+            original=self.original,
         )
 
     def names(self):
         return (
-            {self._pointer()}
+            {self.original.pointer_string()}
             | {
                 name
                 for value in itertools.chain(self.shape, self.strides)
@@ -92,35 +114,31 @@ class Tensor:
             | (self.dtype.names() if isinstance(self.dtype, type(self)) else set())
         )
 
-    def pointers(self, offsets=None):
-        if offsets is None:
-            offsets = self.offsets()
-
-        return self._pointer() + offsets
-
     def offsets(self, indices=None):
         if indices is None:
             indices = self.indices()
 
-        if not isinstance(self.dtype, type(self)):
-            if len(indices) != self.ndim():
-                raise IndexError("Incorrect number of indices.")
+        offsets = [[] for _ in range(self.original.ndim)]
 
-            return sum(
-                indices[idx]
-                * self.stride(idx)
-                * call("arange", 0, self.size(idx))[
-                    tuple(slice(None) if i == idx else None for i in range(self.ndim()))
-                ]
-                for idx in range(self.ndim())
-            )
+        curr = self
+        start = 0
 
-        outer_indices = indices[: self.ndim()]
-        inner_indices = indices[self.ndim() :]
+        while isinstance(curr, type(self)):
+            stop = start + curr.ndim
+            curr_indices = indices[start:stop]
 
-        return sum(
-            index * stride for index, stride in zip(outer_indices, self.strides)
-        ) + self.dtype.offsets(inner_indices)
+            for index, stride in zip(curr_indices, curr.strides):
+                for dim in self._dims_of(stride):
+                    offsets[dim].append(index * stride)
+
+            start = stop
+            curr = curr.dtype
+
+        for dim in range(self.original.ndim):
+            offsets[dim] = sum(offsets[dim])
+            offsets[dim].find_and_replace(Symbol(self.original.strides[dim]), Symbol(1))
+
+        return offsets
 
     def indices(self, index=None):
         if index is None:
@@ -128,17 +146,22 @@ class Tensor:
 
         indices = []
 
-        for stride in type(self)(shape=self.shape, name=self.name).strides:
+        for stride in type(self)(shape=self.shape, original=self.original).strides:
             indices.append(index // stride)
             index %= stride
 
         curr = self.dtype
-        while isinstance(curr, type(self)):
-            indices.extend(
-                0 if curr is not self.inmost() else 1 for _ in range(curr.ndim())
-            )
+
+        while isinstance(curr.dtype, type(self)):
+            for _ in range(curr.ndim):
+                indices.append(0)
+
             curr = curr.dtype
 
+        if isinstance(curr, type(self)):
+            for dim in range(curr.ndim):
+                indices.append(call("arange", 0, curr.shape[dim]))
+
         return tuple(indices)
 
     def inmost(self):
@@ -147,8 +170,14 @@ class Tensor:
 
         return self.dtype.inmost()
 
-    def ndim(self):
-        return len(self.shape)
+    def pointer_string(self):
+        return f"{self.name}_pointer"
+
+    def size_string(self, dim):
+        return f"{self.name}_size_{dim}"
+
+    def stride_string(self, dim):
+        return f"{self.name}_stride_{dim}"
 
     def size(self, dim=None):
         if dim is None:
@@ -162,12 +191,31 @@ class Tensor:
 
         return self.strides[dim]
 
+    @property
+    def ndim(self):
+        return len(self.shape)
+
     @staticmethod
-    def is_pointer(name):
-        return name.endswith("_ptr")
+    def pointer_pattern():
+        return re.compile(rf"({_identifier_pattern_raw_string()})_(pointer)")
 
-    def _pointer(self):
-        return f"{self.name}_ptr"
+    @staticmethod
+    def size_pattern():
+        return re.compile(rf"({_identifier_pattern_raw_string()})_(size)_(.+)")
+
+    @staticmethod
+    def stride_pattern():
+        return re.compile(rf"({_identifier_pattern_raw_string()})_(stride)_(.+)")
+
+    def _dims_of(self, stride):
+        dims = set()
+        names = stride.names() if isinstance(stride, Symbol) else {stride}
+
+        for dim, original_stride in enumerate(self.original.strides):
+            if str(original_stride) in names:
+                dims.add(dim)
+
+        return dims
 
     @staticmethod
     def _calculate_default_strides(shape):
@@ -177,3 +225,7 @@ class Tensor:
             strides.append(size * strides[-1])
 
         return reversed(strides)
+
+
+def _identifier_pattern_raw_string():
+    return r"[a-zA-Z_][a-zA-Z0-9_]*"

ninetoothed-0.4.0/src/ninetoothed/torchifier.py

@@ -0,0 +1,38 @@
+import ast
+
+from ninetoothed.tensor import Tensor
+
+
+class Torchifier(ast.NodeTransformer):
+    def visit_Name(self, node):
+        self.generic_visit(node)
+
+        source = node.id
+
+        def repl(match):
+            return f"{match.group(1)}"
+
+        source = Tensor.pointer_pattern().sub(repl, source)
+
+        def repl(match):
+            return f"{match.group(1)}.{match.group(2)}({match.group(3)})"
+
+        source = Tensor.size_pattern().sub(repl, source)
+        source = Tensor.stride_pattern().sub(repl, source)
+
+        if source != node.id:
+            return ast.parse(source, mode="eval").body
+
+        return node
+
+    def visit_Attribute(self, node):
+        self.generic_visit(node)
+
+        if (
+            isinstance(node.value, ast.Name)
+            and node.value.id == "ninetoothed"
+            and node.attr == "language"
+        ):
+            return node.value
+
+        return node

{ninetoothed-0.2.0 → ninetoothed-0.4.0}/tests/test_matmul.py

@@ -19,18 +19,21 @@ def matmul(lhs, rhs):
         .tile((1, -1))
         .expand((-1, output_tiled.shape[1]))
     )
+    lhs_tiled.dtype = lhs_tiled.dtype.squeeze(0)
+
     rhs_tiled = (
         Tensor(2)
         .tile((BLOCK_SIZE_K, BLOCK_SIZE_N))
         .tile((-1, 1))
         .expand((output_tiled.shape[0], -1))
     )
+    rhs_tiled.dtype = rhs_tiled.dtype.squeeze(1)
 
     @ninetoothed.jit
     def matmul_kernel(lhs: lhs_tiled, rhs: rhs_tiled, output: output_tiled):
         accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
-        for k in range(lhs.shape[1]):
-            accumulator = ntl.dot(lhs[0, k], rhs[k, 0], accumulator)
+        for k in range(lhs.shape[0]):
+            accumulator += ntl.dot(lhs[k], rhs[k])
         output = accumulator.to(ntl.float16)
 
     output = torch.empty(

ninetoothed-0.4.0/tests/test_softmax.py

@@ -0,0 +1,41 @@
+import torch
+import triton
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Symbol, Tensor
+from tests.skippers import skip_if_cuda_not_available
+
+
+def softmax(input):
+    BLOCK_SIZE = Symbol("BLOCK_SIZE", constexpr=True)
+
+    @ninetoothed.jit
+    def softmax_kernel(
+        input_row: Tensor(2, other=float("-inf")).tile((1, BLOCK_SIZE)),
+        output_row: Tensor(2).tile((1, BLOCK_SIZE)),
+    ):
+        row_minus_max = input_row - ntl.max(input_row)
+        numerator = ntl.exp(row_minus_max)
+        denominator = ntl.sum(numerator)
+        output_row = numerator / denominator  # noqa: F841
+
+    output = torch.empty_like(input)
+
+    softmax_kernel(input, output, BLOCK_SIZE=triton.next_power_of_2(input.shape[-1]))
+
+    return output
+
+
+@skip_if_cuda_not_available
+class TestCUDA:
+    @classmethod
+    def setup_class(cls):
+        torch.manual_seed(0)
+
+        cls.input = torch.randn(1823, 781, device="cuda")
+
+    def test_fp32(self):
+        input = type(self).input.to(torch.float32)
+
+        assert torch.allclose(softmax(input), torch.softmax(input, axis=-1))

ninetoothed-0.2.0/src/ninetoothed/torchifier.py

@@ -1,34 +0,0 @@
-import ast
-import re
-
-
-class Torchifier(ast.NodeTransformer):
-    def visit_Name(self, node):
-        self.generic_visit(node)
-
-        pattern = re.compile(r"([a-zA-Z_][a-zA-Z0-9_]*)_(size|stride)_(.+)")
-
-        node.id = node.id.replace("_ptr", "")
-
-        if re.fullmatch(pattern, node.id):
-            return ast.parse(
-                pattern.sub(
-                    lambda match: f"{match.group(1)}.{match.group(2)}({match.group(3)})",
-                    node.id,
-                ),
-                mode="eval",
-            ).body
-
-        return node
-
-    def visit_Attribute(self, node):
-        self.generic_visit(node)
-
-        if (
-            isinstance(node.value, ast.Name)
-            and node.value.id == "ninetoothed"
-            and node.attr == "language"
-        ):
-            return node.value
-
-        return node