ninetoothed 0.1.0__tar.gz → 0.1.1__tar.gz

ninetoothed-0.1.1/.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.1.1/.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

ninetoothed-0.1.1/.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

ninetoothed-0.1.1/.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Unknown virtualenv at \\wsl.localhost\Ubuntu\home\voltman\ninetoothed\venv" project-jdk-type="Python SDK" />
+</project>

ninetoothed-0.1.1/.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/ninetoothed.iml" filepath="$PROJECT_DIR$/.idea/ninetoothed.iml" />
+    </modules>
+  </component>
+</project>

ninetoothed-0.1.1/.idea/ninetoothed.iml

@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$">
+      <excludeFolder url="file://$MODULE_DIR$/.venv" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

ninetoothed-0.1.1/.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

ninetoothed-0.1.1/.idea/workspace.xml

@@ -0,0 +1,47 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ChangeListManager">
+    <list default="true" id="b7a10642-6e58-4644-a069-aad20bd5add0" name="Changes" comment="" />
+    <option name="SHOW_DIALOG" value="false" />
+    <option name="HIGHLIGHT_CONFLICTS" value="true" />
+    <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
+    <option name="LAST_RESOLUTION" value="IGNORE" />
+  </component>
+  <component name="Git.Settings">
+    <option name="RECENT_GIT_ROOT_PATH" value="$PROJECT_DIR$" />
+  </component>
+  <component name="ProjectColorInfo"><![CDATA[{
+  "customColor": "",
+  "associatedIndex": 8
+}]]></component>
+  <component name="ProjectId" id="2jzxGjCO1rbfcApmCQlTd14207t" />
+  <component name="ProjectViewState">
+    <option name="hideEmptyMiddlePackages" value="true" />
+    <option name="showLibraryContents" value="true" />
+  </component>
+  <component name="PropertiesComponent"><![CDATA[{
+  "keyToString": {
+    "RunOnceActivity.ShowReadmeOnStart": "true",
+    "git-widget-placeholder": "dev",
+    "settings.editor.selected.configurable": "advanced.settings"
+  }
+}]]></component>
+  <component name="SharedIndexes">
+    <attachedChunks>
+      <set>
+        <option value="bundled-python-sdk-975db3bf15a3-31b6be0877a2-com.jetbrains.pycharm.community.sharedIndexes.bundled-PC-241.18034.82" />
+      </set>
+    </attachedChunks>
+  </component>
+  <component name="SpellCheckerSettings" RuntimeDictionaries="0" Folders="0" CustomDictionaries="0" DefaultDictionary="application-level" UseSingleDictionary="true" transferred="true" />
+  <component name="TaskManager">
+    <task active="true" id="Default" summary="Default task">
+      <changelist id="b7a10642-6e58-4644-a069-aad20bd5add0" name="Changes" comment="" />
+      <created>1722403613879</created>
+      <option name="number" value="Default" />
+      <option name="presentableId" value="Default" />
+      <updated>1722403613879</updated>
+    </task>
+    <servers />
+  </component>
+</project>

ninetoothed-0.1.1/PKG-INFO

@@ -0,0 +1,79 @@
+Metadata-Version: 2.3
+Name: ninetoothed
+Version: 0.1.1
+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.1.1/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.1.1}/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.1.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ninetoothed"
-version = "0.1.0"
+version = "0.1.1"
 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.1.1/requirements.txt

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

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

@@ -1,4 +1,5 @@
 import ast
+import collections
 import functools
 import inspect
 import itertools
@@ -12,6 +13,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]
+
+        source = textwrap.dedent(inspect.getsource(self.func))
+        tree = ast.parse(source)
+
+        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)
+
+        class Handle:
+            def __init__(self, kernel, launch):
+                self._kernel = kernel
+                self._launch = launch
+
+            def __call__(self, *args, **kwargs):
+                return self._launch(*args, **kwargs)
+
+        handle = Handle(
+            namespace[self.func.__name__],
+            namespace[f"launch_{self.func.__name__}"],
+        )
+
+        type(self).handles[source_file][source_line] = handle
+
+        return handle
+
+
 class CodeGenerator(ast.NodeTransformer):
     def __init__(self, context):
         super().__init__()
@@ -286,34 +348,3 @@ class Tritonizer(ast.NodeTransformer):
             )
 
         return node
-
-
-def jit(func):
-    source = textwrap.dedent(inspect.getsource(func))
-    tree = ast.parse(source)
-
-    CodeGenerator(func.__annotations__).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)
-
-    class Handle:
-        def __init__(self, kernel, launch):
-            self._kernel = kernel
-            self._launch = launch
-
-        def __call__(self, *args, **kwargs):
-            return self._launch(*args, **kwargs)
-
-    return Handle(namespace[func.__name__], namespace[f"launch_{func.__name__}"])

{ninetoothed-0.1.0 → ninetoothed-0.1.1}/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)

ninetoothed-0.1.1/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.1.1/tests/test_add.py

@@ -0,0 +1,41 @@
+import ninetoothed
+import torch
+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.1.1/tests/test_matmul.py

@@ -0,0 +1,72 @@
+import ninetoothed
+import torch
+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 = ninetoothed.language.zeros(
+            output.shape, dtype=ninetoothed.language.float32
+        )
+        for k in range(lhs.shape[1]):
+            accumulator = ninetoothed.language.dot(lhs[0, k], rhs[k, 0], accumulator)
+        output = accumulator.to(ninetoothed.language.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).**