ninetoothed 0.1.0__tar.gz → 0.2.0__tar.gz

ninetoothed-0.2.0/.github/workflows/pytest.yml

@@ -0,0 +1,18 @@
+name: pytest
+on: [push, pull_request]
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+      - name: Test with pytest
+        run: |
+          pytest --doctest-modules --junitxml=junit/test-results.xml --cov=ninetoothed --cov-report=xml --cov-report=html

ninetoothed-0.2.0/PKG-INFO

@@ -0,0 +1,79 @@
+Metadata-Version: 2.3
+Name: ninetoothed
+Version: 0.2.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
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# NineToothed
+
+A domain-specific language (DSL) based on Triton but providing higher-level abstractions.
+
+**Other language versions: [English](README.md), [简体中文](docs/README.zh.md).**
+
+## Installation
+
+We can use `pip` to install `ninetoothed`.
+
+```shell
+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`.
+
+## Usage
+
+Currently, we can use the `Tensor` and `Symbol` classes in the `ninetoothed` package to perform meta-operations like `tile` and `expand` to easily construct kernel functions. Below, we will use these features to create vector addition and matrix multiplication kernel functions.
+
+### Vector Addition
+
+```python
+BLOCK_SIZE = Symbol("BLOCK_SIZE", meta=True)
+
+@ninetoothed.jit
+def add_kernel(
+    x: Tensor(1).tile((BLOCK_SIZE,)),
+    y: Tensor(1).tile((BLOCK_SIZE,)),
+    z: Tensor(1).tile((BLOCK_SIZE,)),
+):
+    z = x + y
+```
+
+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`.
+
+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
+
+```python
+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)
+
+a_tiled = Tensor(2).tile((BLOCK_SIZE_M, BLOCK_SIZE_K)).tile((1, -1))
+b_tiled = Tensor(2).tile((BLOCK_SIZE_K, BLOCK_SIZE_N)).tile((-1, 1))
+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))
+
+@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)
+    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$.
+
+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/README.md

@@ -0,0 +1,65 @@
+# NineToothed
+
+A domain-specific language (DSL) based on Triton but providing higher-level abstractions.
+
+**Other language versions: [English](README.md), [简体中文](docs/README.zh.md).**
+
+## Installation
+
+We can use `pip` to install `ninetoothed`.
+
+```shell
+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`.
+
+## Usage
+
+Currently, we can use the `Tensor` and `Symbol` classes in the `ninetoothed` package to perform meta-operations like `tile` and `expand` to easily construct kernel functions. Below, we will use these features to create vector addition and matrix multiplication kernel functions.
+
+### Vector Addition
+
+```python
+BLOCK_SIZE = Symbol("BLOCK_SIZE", meta=True)
+
+@ninetoothed.jit
+def add_kernel(
+    x: Tensor(1).tile((BLOCK_SIZE,)),
+    y: Tensor(1).tile((BLOCK_SIZE,)),
+    z: Tensor(1).tile((BLOCK_SIZE,)),
+):
+    z = x + y
+```
+
+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`.
+
+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
+
+```python
+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)
+
+a_tiled = Tensor(2).tile((BLOCK_SIZE_M, BLOCK_SIZE_K)).tile((1, -1))
+b_tiled = Tensor(2).tile((BLOCK_SIZE_K, BLOCK_SIZE_N)).tile((-1, 1))
+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))
+
+@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)
+    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$.
+
+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.1.0 → ninetoothed-0.2.0}/docs/README.zh.md

@@ -6,20 +6,13 @@
 
 ## 安装
 
-由于现在 `ninetoothed` 还没有进入到 PyPI，我们需要手动 `git clone` 本仓库，再 `pip install`。
+我们可以使用 `pip` 安装 `ninetoothed`。
 
 ```shell
-git clone https://github.com/InfiniTensor/ninetoothed.git
-pip install ./ninetoothed
+pip install ninetoothed
 ```
 
-成功运行完以上两个命令之后，`ninetoothed` 就被安装好了。需要注意的是，我们目前还没有加入任何的依赖管理，所以如果想要运行 Triton 核函数，那么最起码应该再安装 Triton。
-
-```shell
-pip install triton
-```
-
-其余包可以根据需要自行安装，如 `torch`、`matplotlib`、`pandas` 等。
+成功运行完以上两个命令之后，`ninetoothed` 就被安装好了。但是除了 `ninetoothed` 的本体之外，如果我们想要真正发挥它的作用，至少还需要安装 `triton` 和一个 `ninetoothed` 所支持的深度学习框架。以尝试为目的的话，我们推荐安装 `triton` 和 `torch`。
 
 ## 使用
 
@@ -67,6 +60,6 @@ def matmul_kernel(a: a_tiled, b: b_tiled, c: c_tiled):
     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$ 的列块。
 
-对应好了分块，后续的部分就简单多了。在函数体当中，我们定义了一个 `accumulator`，用于累加中间结果，之后就遍历了对应好的 A 的行块和 B 的列块，并且把他们相乘的结果累加到了 `accumulator` 当中，最后再将 `accumulator` 放到了对应的 C 的分块当中。由于参数张量被分成的每一块都被执行了这样的操作，因此即便对于整体而言，乘法也被完成了。
+对应好了分块，后续的部分就简单多了。在函数体当中，我们定义了一个 `accumulator`，用于累加中间结果，之后就遍历了对应好的 $A$ 的行块和 $B$ 的列块，并且把他们相乘的结果累加到了 `accumulator` 当中，最后再将 `accumulator` 放到了对应的 $C$ 的分块当中。由于参数张量被分成的每一块都被执行了这样的操作，因此即便对于整体而言，乘法也被完成了。

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

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ninetoothed"
-version = "0.1.0"
+version = "0.2.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"
@@ -18,3 +18,9 @@ classifiers = [
 [project.urls]
 Homepage = "https://github.com/InfiniTensor/ninetoothed"
 Issues = "https://github.com/InfiniTensor/ninetoothed/issues"
+
+[tool.ruff]
+src = [".", "src", "tests"]
+
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F", "I"]

ninetoothed-0.2.0/requirements.txt

@@ -0,0 +1,7 @@
+matplotlib
+pandas
+pytest
+pytest-cov
+ruff
+torch
+triton

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

@@ -1,10 +1,10 @@
 import ast
+import collections
 import functools
 import inspect
 import itertools
 import math
 import tempfile
-import textwrap
 
 from ninetoothed.language import attribute, call
 from ninetoothed.symbol import Symbol
@@ -12,6 +12,67 @@ from ninetoothed.tensor import Tensor
 from ninetoothed.torchifier import Torchifier
 
 
+def jit(func):
+    return JIT(func)()
+
+
+class JIT:
+    handles = collections.defaultdict(dict)
+
+    def __init__(self, func):
+        self.func = func
+
+    def __call__(self):
+        source_file = inspect.getsourcefile(self.func)
+        source_line = inspect.getsourcelines(self.func)[1]
+
+        if (
+            source_file in type(self).handles
+            and source_line in type(self).handles[source_file]
+        ):
+            return type(self).handles[source_file][source_line]
+
+        tree = self._get_tree()
+
+        CodeGenerator(inspect.get_annotations(self.func)).visit(tree)
+        Tritonizer().visit(tree)
+        ast.fix_missing_locations(tree)
+
+        unparsed = ast.unparse(tree).replace("None:", ":").replace(":None", ":")
+
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as temp_file:
+            temp_file.write(unparsed.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)
+
+        handle = _Handle(
+            namespace[self.func.__name__],
+            namespace[f"launch_{self.func.__name__}"],
+        )
+
+        type(self).handles[source_file][source_line] = handle
+
+        return handle
+
+    def _get_tree(self):
+        module = ast.parse(inspect.getsource(inspect.getmodule(self.func)))
+
+        _AliasRestorer().visit(module)
+        finder = _FunctionDefFinder(self.func.__name__)
+        finder.visit(module)
+
+        return ast.Module(body=[finder.result], type_ignores=[])
+
+
 class CodeGenerator(ast.NodeTransformer):
     def __init__(self, context):
         super().__init__()
@@ -38,6 +99,18 @@ class CodeGenerator(ast.NodeTransformer):
 
         self.generic_visit(node)
 
+        for arg in self._args:
+            if not isinstance(arg, Tensor):
+                continue
+
+            node.body.insert(
+                0,
+                ast.Assign(
+                    targets=[Symbol(f"{arg.name}_ptrs").node],
+                    value=arg.pointers().node,
+                ),
+            )
+
         return node
 
     def visit_arguments(self, node):
@@ -74,12 +147,12 @@ class CodeGenerator(ast.NodeTransformer):
             value = self._context[node.value.id]
 
             if isinstance(value, Tensor):
-                if isinstance(node.slice, ast.Tuple):
-                    indices = value.indices() + tuple(node.slice.elts)
-                else:
-                    indices = value.indices() + (node.slice,)
-                offsets = value.offsets(indices)
-                pointers = value.pointers(offsets)
+                pointers = type(self)._create_pointers(
+                    value,
+                    node.slice.elts
+                    if isinstance(node.slice, ast.Tuple)
+                    else (node.slice,),
+                )
 
                 return call("load", pointers).node
 
@@ -104,7 +177,9 @@ class CodeGenerator(ast.NodeTransformer):
         self.generic_visit(node)
 
         if node.id in self._context and isinstance(node.ctx, ast.Load):
-            return call("load", self._context[node.id].pointers().node).node
+            return call(
+                "load", type(self)._create_pointers(self._context[node.id], ()).node
+            ).node
 
         return node
 
@@ -118,7 +193,7 @@ class CodeGenerator(ast.NodeTransformer):
                 return ast.Expr(
                     call(
                         "store",
-                        self._context[target.id].pointers().node,
+                        type(self)._create_pointers(self._context[target.id], ()).node,
                         node.value,
                     ).node
                 )
@@ -133,13 +208,12 @@ class CodeGenerator(ast.NodeTransformer):
                 if isinstance(value, Tensor):
                     self.generic_visit(node)
 
-                    indices = value.indices() + tuple(
+                    pointers = type(self)._create_pointers(
+                        value,
                         target.slice.elts
                         if isinstance(target.slice, ast.Tuple)
-                        else target.slice
+                        else (target.slice,),
                     )
-                    offsets = value.offsets(indices)
-                    pointers = value.pointers(offsets)
 
                     return ast.Expr(
                         call(
@@ -254,6 +328,14 @@ 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())]
+        )
+
 
 class Tritonizer(ast.NodeTransformer):
     def visit_Module(self, node):
@@ -267,8 +349,8 @@ class Tritonizer(ast.NodeTransformer):
     def visit_Name(self, node):
         self.generic_visit(node)
 
-        if node.id == "ninetoothed":
-            node.id = "triton"
+        if node.id == "ninetoothed" or "ninetoothed." in node.id:
+            node.id = node.id.replace("ninetoothed", "triton")
 
         return node
 
@@ -288,32 +370,71 @@ class Tritonizer(ast.NodeTransformer):
         return node
 
 
-def jit(func):
-    source = textwrap.dedent(inspect.getsource(func))
-    tree = ast.parse(source)
+class _Handle:
+    def __init__(self, kernel, launch):
+        self._kernel = kernel
+        self._launch = launch
+
+    def __call__(self, *args, **kwargs):
+        return self._launch(*args, **kwargs)
+
+
+class _AliasRestorer(ast.NodeTransformer):
+    def __init__(self):
+        super().__init__()
+
+        self._aliases = {}
+        self._redefined = set()
+
+    def visit_Import(self, node):
+        for alias in node.names:
+            if alias.asname:
+                self._aliases[alias.asname] = alias.name
+
+        return node
+
+    def visit_ImportFrom(self, node):
+        for alias in node.names:
+            full_name = f"{node.module}.{alias.name}"
+            if alias.asname:
+                self._aliases[alias.asname] = full_name
+
+        return node
 
-    CodeGenerator(func.__annotations__).visit(tree)
-    Tritonizer().visit(tree)
-    ast.fix_missing_locations(tree)
+    def visit_Assign(self, node):
+        for target in node.targets:
+            if isinstance(target, ast.Name):
+                self._redefined.add(target.id)
+
+        return self.generic_visit(node)
+
+    def visit_FunctionDef(self, node):
+        original_redefined = self._redefined.copy()
+
+        self.generic_visit(node)
+
+        self._redefined = original_redefined
 
-    unparsed = ast.unparse(tree).replace("None:", ":").replace(":None", ":")
+        return node
+
+    def visit_Name(self, node):
+        if node.id in self._redefined:
+            return node
 
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as temp_file:
-        temp_file.write(unparsed.encode("utf-8"))
-        temp_file_name = temp_file.name
+        if node.id in self._aliases:
+            return ast.Name(id=self._aliases[node.id], ctx=node.ctx)
 
-    with open(temp_file_name, "r") as temp_file:
-        code = compile(source=temp_file.read(), filename=temp_file_name, mode="exec")
+        return node
 
-    namespace = {}
-    exec(code, namespace)
 
-    class Handle:
-        def __init__(self, kernel, launch):
-            self._kernel = kernel
-            self._launch = launch
+class _FunctionDefFinder(ast.NodeVisitor):
+    def __init__(self, name):
+        self._name = name
 
-        def __call__(self, *args, **kwargs):
-            return self._launch(*args, **kwargs)
+        self.result = None
 
-    return Handle(namespace[func.__name__], namespace[f"launch_{func.__name__}"])
+    def visit_FunctionDef(self, node):
+        if node.name == self._name:
+            self.result = node
+
+        self.generic_visit(node)

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

@@ -34,24 +34,47 @@ class Symbol:
             self._node.id = type(self)._create_constexpr(self._node.id)
 
     def __add__(self, other):
-        return type(self)(
-            ast.BinOp(left=self._node, op=ast.Add(), right=type(self)(other)._node)
-        )
+        other = type(self)(other)
+
+        if isinstance(self._node, ast.Constant) and self._node.value == 0:
+            return other
+
+        if isinstance(other._node, ast.Constant) and other._node.value == 0:
+            return self
+
+        return type(self)(ast.BinOp(left=self._node, op=ast.Add(), right=other._node))
 
     def __radd__(self, other):
         return self.__add__(other)
 
     def __mul__(self, other):
-        return type(self)(
-            ast.BinOp(left=self._node, op=ast.Mult(), right=type(self)(other)._node)
-        )
+        other = type(self)(other)
+
+        if isinstance(self._node, ast.Constant) and self._node.value == 0:
+            return type(self)(0)
+
+        if isinstance(other._node, ast.Constant) and other._node.value == 0:
+            return type(self)(0)
+
+        if isinstance(self._node, ast.Constant) and self._node.value == 1:
+            return other
+
+        if isinstance(other._node, ast.Constant) and other._node.value == 1:
+            return self
+
+        return type(self)(ast.BinOp(left=self._node, op=ast.Mult(), right=other._node))
 
     def __rmul__(self, other):
         return self.__mul__(other)
 
     def __floordiv__(self, other):
+        other = type(self)(other)
+
+        if isinstance(other._node, ast.Constant) and other._node.value == 1:
+            return self
+
         return type(self)(
-            ast.BinOp(left=self._node, op=ast.FloorDiv(), right=type(self)(other)._node)
+            ast.BinOp(left=self._node, op=ast.FloorDiv(), right=other._node)
         )
 
     def __mod__(self, other):

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

@@ -46,7 +46,7 @@ class Tensor:
             new_size = call("cdiv", size, tile_size)
             outer_shape.append(new_size)
 
-            new_stride = call("cdiv", stride * size, (new_size * tile_stride))
+            new_stride = stride * tile_size // tile_stride
             outer_strides.append(new_stride)
 
             inner_shape.append(tile_size)
@@ -103,11 +103,12 @@ class Tensor:
             indices = self.indices()
 
         if not isinstance(self.dtype, type(self)):
-            if indices:
+            if len(indices) != self.ndim():
                 raise IndexError("Incorrect number of indices.")
 
             return sum(
-                self.stride(idx)
+                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()))
                 ]
@@ -131,8 +132,21 @@ class Tensor:
             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())
+            )
+            curr = curr.dtype
+
         return tuple(indices)
 
+    def inmost(self):
+        if not isinstance(self.dtype, type(self)):
+            return self
+
+        return self.dtype.inmost()
+
     def ndim(self):
         return len(self.shape)
 

ninetoothed-0.2.0/tests/skippers.py

@@ -0,0 +1,16 @@
+import pytest
+import torch
+
+
+def skip_if_cuda_not_available(func):
+    return pytest.mark.skipif(
+        not torch.cuda.is_available(),
+        reason="CUDA not available",
+    )(func)
+
+
+def skip_if_float8_e5m2_not_supported(func):
+    return pytest.mark.skipif(
+        not hasattr(torch, "float8_e5m2"),
+        reason="`float8_e5m2` not supported",
+    )(func)

ninetoothed-0.2.0/tests/test_add.py

@@ -0,0 +1,41 @@
+import torch
+
+import ninetoothed
+from ninetoothed import Symbol, Tensor
+from tests.skippers import skip_if_cuda_not_available
+
+
+def add(lhs, rhs):
+    BLOCK_SIZE = Symbol("BLOCK_SIZE", meta=True)
+
+    @ninetoothed.jit
+    def add_kernel(
+        lhs: Tensor(1).tile((BLOCK_SIZE,)),
+        rhs: Tensor(1).tile((BLOCK_SIZE,)),
+        output: Tensor(1).tile((BLOCK_SIZE,)),
+    ):
+        output = lhs + rhs  # noqa: F841
+
+    output = torch.empty_like(lhs)
+
+    add_kernel(lhs, rhs, output)
+
+    return output
+
+
+@skip_if_cuda_not_available
+class TestCUDA:
+    @classmethod
+    def setup_class(cls):
+        torch.manual_seed(0)
+
+        size = 98432
+
+        cls.lhs = torch.rand(size, device="cuda")
+        cls.rhs = torch.rand(size, device="cuda")
+
+    def test_fp32(self):
+        lhs = type(self).lhs.to(torch.float32)
+        rhs = type(self).rhs.to(torch.float32)
+
+        assert torch.allclose(add(lhs, rhs), lhs + rhs)

ninetoothed-0.2.0/tests/test_matmul.py

@@ -0,0 +1,71 @@
+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 matmul(lhs, rhs):
+    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)
+
+    output_tiled = Tensor(2).tile((BLOCK_SIZE_M, BLOCK_SIZE_N))
+
+    lhs_tiled = (
+        Tensor(2)
+        .tile((BLOCK_SIZE_M, BLOCK_SIZE_K))
+        .tile((1, -1))
+        .expand((-1, output_tiled.shape[1]))
+    )
+    rhs_tiled = (
+        Tensor(2)
+        .tile((BLOCK_SIZE_K, BLOCK_SIZE_N))
+        .tile((-1, 1))
+        .expand((output_tiled.shape[0], -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)
+        output = accumulator.to(ntl.float16)
+
+    output = torch.empty(
+        (lhs.shape[0], rhs.shape[1]), device=lhs.device, dtype=torch.float16
+    )
+
+    matmul_kernel(lhs, rhs, output)
+
+    return output
+
+
+@skip_if_cuda_not_available
+class TestCUDA:
+    @classmethod
+    def setup_class(cls):
+        torch.manual_seed(0)
+
+        shape = (512, 512)
+
+        cls.lhs = torch.randn(shape, device="cuda")
+        cls.rhs = torch.randn(shape, device="cuda")
+
+    def test_fp16(self):
+        lhs = type(self).lhs.to(torch.float16)
+        rhs = type(self).rhs.to(torch.float16)
+
+        assert torch.allclose(matmul(lhs, rhs), torch.matmul(lhs, rhs))
+
+    @skip_if_float8_e5m2_not_supported
+    def test_fp8(self):
+        lhs = type(self).lhs.to(torch.float8_e5m2)
+        rhs = type(self).rhs.T.to(torch.float8_e5m2)
+
+        assert torch.allclose(
+            matmul(lhs, rhs),
+            torch.matmul(lhs.to(torch.float16), rhs.to(torch.float16)),
+            atol=0.125,
+        )

ninetoothed-0.1.0/PKG-INFO

@@ -1,19 +0,0 @@
-Metadata-Version: 2.3
-Name: ninetoothed
-Version: 0.1.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
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-
-# Nine-Toothed
-
-A domain-specific language based on Triton but providing higher-level abstraction.
-
-**Read this in other languages: [English](README.md), [简体中文](docs/README.zh.md).**

ninetoothed-0.1.0/README.md

@@ -1,5 +0,0 @@
-# Nine-Toothed
-
-A domain-specific language based on Triton but providing higher-level abstraction.
-
-**Read this in other languages: [English](README.md), [简体中文](docs/README.zh.md).**