ninetoothed 0.17.0__tar.gz → 0.18.0__tar.gz

{ninetoothed-0.17.0 → ninetoothed-0.18.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.17.0
+Version: 0.18.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

{ninetoothed-0.17.0 → ninetoothed-0.18.0}/pyproject.toml

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

{ninetoothed-0.17.0 → ninetoothed-0.18.0}/src/ninetoothed/aot.py

@@ -1,9 +1,11 @@
 import ast
 import pathlib
+import re
 import subprocess
 import tempfile
 import uuid
 
+import ninetoothed.naming as naming
 from ninetoothed.dtype import int64
 from ninetoothed.generation import CACHE_DIR, CodeGenerator
 from ninetoothed.tensor import Tensor
@@ -25,8 +27,10 @@ def aot(
 
 def _aot(func, caller, kernel_name, num_warps, num_stages):
     def _find_tensor_by_source_name(tensors, name):
+        name = naming.remove_prefixes(name)
+
         for tensor in tensors:
-            if tensor.source.name == name:
+            if naming.remove_prefixes(tensor.source.name) == name:
                 return tensor
 
     _HEADER_PATH.parent.mkdir(exist_ok=True)
@@ -49,11 +53,15 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
     kernel_func = code_generator.kernel_func
     launch_func = code_generator.launch_func
 
+    param_strings = ["stream"]
     param_types = []
+    constexpr_param_indices = []
 
     for arg in kernel_func.args.args:
         param = arg.arg
 
+        param_strings.append(param)
+
         if match := Tensor.pointer_pattern().fullmatch(param):
             source_name = match.group(0).removesuffix("_pointer")
             tensor = _find_tensor_by_source_name(tensors, source_name)
@@ -64,9 +72,23 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
             param_types.append(int64)
         elif Tensor.stride_pattern().fullmatch(param):
             param_types.append(int64)
+        else:
+            source_name = param
+            tensor = _find_tensor_by_source_name(tensors, source_name)
+            dtype = tensor.source.dtype
+
+            if tensor.constexpr:
+                param_types.append(f"{tensor.value}")
+                constexpr_param_indices.append(len(param_types) - 1)
+            else:
+                param_types.append(dtype)
 
     signature = ", ".join(param_types)
 
+    for index in sorted(set(constexpr_param_indices), reverse=True):
+        param_strings.pop(index + 1)
+        param_types.pop(index)
+
     grid_extractor = _GridExtractor()
     launch_func = grid_extractor.visit(launch_func)
     grid_extractor.visit(code_generator.raw_grid)
@@ -76,21 +98,26 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
         source_file, kernel_name, signature, grid, num_warps, num_stages
     )
 
-    unparser = _Unparser()
+    c_source_file_name = f"{kernel_name}.{signature_hash}.c"
+    c_source_file = output_contents[c_source_file_name]
+
+    c_header_file_name = f"{kernel_name}.{signature_hash}.h"
+    c_header_file = output_contents[c_header_file_name]
+
+    pattern = rf"\({', '.join(rf'(.*) {param}' for param in param_strings)}\)"
+    c_param_type_strings = re.search(pattern, c_header_file).groups()
+
+    unparser = _Unparser(c_param_type_strings)
 
     launch_func_unparsed = unparser.unparse(launch_func)
     launch_func_unparsed = launch_func_unparsed.replace(
         func.__name__, f"{kernel_name}_{signature_hash}"
     )
 
-    c_source_file_name = f"{kernel_name}.{signature_hash}.c"
-    c_source_file = output_contents[c_source_file_name]
     c_source_file = f"{c_source_file}\n{launch_func_unparsed}\n"
     c_source_file = c_source_file.replace("<stdint.h>", f'"{_HEADER_PATH}"')
     output_contents[c_source_file_name] = c_source_file
 
-    c_header_file_name = f"{kernel_name}.{signature_hash}.h"
-    c_header_file = output_contents[c_header_file_name]
     c_header_file = f'{c_header_file}\n#ifdef __cplusplus\nextern "C" {unparser.header};\n#else\n{unparser.header};\n#endif\n'
     c_header_file = c_header_file.replace("<stdint.h>", f'"{_HEADER_PATH}"')
     output_contents[c_header_file_name] = c_header_file
@@ -120,6 +147,9 @@ _HEADER_PATH = CACHE_DIR / "ninetoothed.h"
 
 
 class _Unparser:
+    def __init__(self, param_types):
+        self._param_types = param_types
+
     def unparse(self, node):
         method_name = "_unparse_" + node.__class__.__name__
 
@@ -137,11 +167,29 @@ class _Unparser:
     def _unparse_Call(self, node):
         call = ast.Call(
             func=node.func,
-            args=[ast.Name(id="stream", ctx=ast.Load())] + node.args,
+            args=[ast.Name(id="stream", ctx=ast.Load())]
+            + [
+                arg
+                for arg in node.args
+                if not isinstance(arg, ast.Name) or not naming.is_constexpr(arg.id)
+            ],
             keywords=[],
         )
 
-        return f"return {self._generic_unparse(call)};"
+        unparsed = f"return {self._generic_unparse(call)};"
+
+        pattern = rf"\((stream), {', '.join(r'([^,]*)' for _ in range(len(self._param_types) - 1))}\)"
+        args = re.search(pattern, unparsed).groups()
+
+        for i, (arg, type) in enumerate(zip(args, self._param_types)):
+            if i != 0 and "." not in arg:
+                new_arg = f"*({type} *){arg}.data"
+            else:
+                new_arg = f"({type}){arg}"
+
+            unparsed = unparsed.replace(arg, new_arg)
+
+        return unparsed
 
     def _unparse_FunctionDef(self, node):
         params = ["NineToothedStream stream"]
@@ -153,13 +201,16 @@ class _Unparser:
         body_lines = []
 
         for stmt in node.body:
+            if isinstance(stmt, ast.Assign):
+                continue
+
             stmt_unparsed = self.unparse(stmt)
 
             if isinstance(stmt, ast.Expr):
                 stmt_unparsed = stmt_unparsed.strip()
 
-                if not stmt_unparsed.endswith(";"):
-                    stmt_unparsed += ";"
+            if not stmt_unparsed.endswith(";"):
+                stmt_unparsed += ";"
 
             body_lines.append("    " + stmt_unparsed)
 

{ninetoothed-0.17.0 → ninetoothed-0.18.0}/src/ninetoothed/generation.py

@@ -289,12 +289,12 @@ class CodeGenerator(ast.NodeTransformer):
         if isinstance(value, Tensor):
             attr = getattr(value, node.attr)
 
-            if node.attr == "dtype" and attr is None:
-                return Symbol(f"{value.source.pointer_string()}.type.element_ty").node
-
             if isinstance(attr, Tensor):
                 return attr
 
+            if node.attr == "dtype":
+                return Symbol(f"{value.source.pointer_string()}.type.element_ty").node
+
             return Symbol(attr).node
 
         self.generic_visit(node)
@@ -500,16 +500,19 @@ class CodeGenerator(ast.NodeTransformer):
             naming.remove_prefixes(param) for param in next_power_of_2_params
         ]
 
+        arg_names = [naming.remove_prefixes(arg.source.name) for arg in self._args]
+
+        arg_names += [
+            param
+            for param in non_next_power_of_2_constexpr_params_without_prefixes
+            if not Tensor.size_pattern().fullmatch(param) and param not in arg_names
+        ]
+
         launch = ast.FunctionDef(
             name=self.launch_func_name,
             args=ast.arguments(
                 posonlyargs=[],
-                args=[ast.arg(arg=arg.source.name) for arg in self._args]
-                + [
-                    ast.arg(arg=param)
-                    for param in non_next_power_of_2_constexpr_params_without_prefixes
-                    if not Tensor.size_pattern().fullmatch(param)
-                ],
+                args=[ast.arg(arg=name) for name in arg_names],
                 kwonlyargs=[],
                 defaults=[],
             ),

{ninetoothed-0.17.0 → ninetoothed-0.18.0}/src/ninetoothed/tensor.py

@@ -32,6 +32,8 @@ class Tensor:
         strides=None,
         other=None,
         shape_options=None,
+        constexpr=None,
+        value=None,
         name=None,
         source=None,
         source_dims=None,
@@ -74,6 +76,21 @@ class Tensor:
 
         self.other = other
 
+        if constexpr and self.ndim != 0:
+            raise ValueError(
+                "`constexpr` can only be set for zero-dimensional tensors."
+            )
+
+        self.constexpr = constexpr
+
+        if self.constexpr:
+            self.name = naming.make_constexpr(self.name)
+
+        if not constexpr and value is not None:
+            raise ValueError("`value` can only be set for constexpr tensors.")
+
+        self.value = value
+
         if source is not None:
             self.source = source
         else:

{ninetoothed-0.17.0 → ninetoothed-0.18.0}/tests/test_addmm.py

@@ -9,11 +9,23 @@ from ninetoothed import Tensor
 from tests.skippers import skip_if_cuda_not_available, skip_if_float8_e5m2_not_supported
 
 
-def arrangement(input, mat1, mat2, beta, alpha, output):
-    _, _, input_arranged = matmul.arrangement(mat1, mat2, input)
+def arrangement(
+    input,
+    mat1,
+    mat2,
+    beta,
+    alpha,
+    output,
+    BLOCK_SIZE_M=matmul.BLOCK_SIZE_M,
+    BLOCK_SIZE_N=matmul.BLOCK_SIZE_N,
+    BLOCK_SIZE_K=matmul.BLOCK_SIZE_K,
+):
+    _, _, input_arranged = matmul.arrangement(
+        mat1, mat2, input, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K
+    )
 
     mat1_arranged, mat2_arranged, output_arranged = matmul.arrangement(
-        mat1, mat2, output
+        mat1, mat2, output, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K
     )
 
     return input_arranged, mat1_arranged, mat2_arranged, beta, alpha, output_arranged

{ninetoothed-0.17.0 → ninetoothed-0.18.0}/tests/test_aot.py

@@ -7,6 +7,8 @@ import torch.nn.functional as F
 
 import ninetoothed
 import ninetoothed.generation
+import tests.test_addmm as addmm
+import tests.test_attention as attention
 import tests.test_conv2d as conv2d
 import tests.test_matmul as matmul
 from ninetoothed import Tensor
@@ -19,6 +21,91 @@ class TestCUDA:
     def setup_class(cls):
         torch.manual_seed(0)
 
+    def test_addmm(self):
+        arrangement = functools.partial(
+            addmm.arrangement, BLOCK_SIZE_M=64, BLOCK_SIZE_N=64, BLOCK_SIZE_K=64
+        )
+        application = addmm.application
+        tensors = tuple(
+            Tensor(ndim, dtype=ninetoothed.float16) for ndim in (2, 2, 2, 0, 0, 2)
+        )
+        caller = "cuda"
+        kernel_name = "addmm"
+        output_dir = ninetoothed.generation.CACHE_DIR
+
+        launch_func = _generate_launch_func(
+            arrangement,
+            application,
+            tensors,
+            caller=caller,
+            kernel_name=kernel_name,
+            output_dir=output_dir,
+        )
+
+        shape = (512, 512)
+        dtype = torch.float16
+        device = caller
+
+        input = torch.randn(shape, dtype=dtype, device=device)
+        mat1 = torch.randn(shape, dtype=dtype, device=device)
+        mat2 = torch.randn(shape, dtype=dtype, device=device)
+        beta = torch.randn((), dtype=dtype)
+        alpha = torch.randn((), dtype=dtype)
+        output = torch.empty(
+            (mat1.shape[0], mat2.shape[1]), dtype=mat1.dtype, device=mat1.device
+        )
+
+        _run_launch_func(launch_func, input, mat1, mat2, beta, alpha, output)
+
+        assert torch.allclose(
+            output, torch.addmm(input, mat1, mat2, beta=beta, alpha=alpha), atol=0.075
+        )
+
+    def test_attention(self):
+        emb_dim = 64
+
+        arrangement = functools.partial(
+            attention.arrangement, BLOCK_SIZE_M=64, BLOCK_SIZE_N=64
+        )
+        application = attention.application
+        query_, key_, value_, output_ = tuple(
+            Tensor(4, dtype=ninetoothed.float16) for _ in range(4)
+        )
+        for tensor in (query_, key_, value_, output_):
+            tensor.shape = tensor.shape[:-1] + (emb_dim,)
+        is_causal_ = Tensor(0, constexpr=True, value=1)
+        tensors = (query_, key_, value_, is_causal_, output_)
+        caller = "cuda"
+        kernel_name = "attention"
+        output_dir = ninetoothed.generation.CACHE_DIR
+
+        launch_func = _generate_launch_func(
+            arrangement,
+            application,
+            tensors,
+            caller=caller,
+            kernel_name=kernel_name,
+            output_dir=output_dir,
+        )
+
+        shape = (2, 4, 1024, emb_dim)
+        dtype = torch.float16
+        device = caller
+
+        query = torch.randn(shape, dtype=dtype, device=device)
+        key = torch.randn(shape, dtype=dtype, device=device)
+        value = torch.randn(shape, dtype=dtype, device=device)
+        is_causal = torch.tensor(True)
+        output = torch.empty(shape, dtype=dtype, device=device)
+
+        _run_launch_func(launch_func, query, key, value, is_causal, output)
+
+        assert torch.allclose(
+            output,
+            F.scaled_dot_product_attention(query, key, value, is_causal=True, scale=1),
+            atol=0.01,
+        )
+
     def test_matmul(self):
         arrangement = functools.partial(
             matmul.arrangement, BLOCK_SIZE_M=64, BLOCK_SIZE_N=64, BLOCK_SIZE_K=64

{ninetoothed-0.17.0 → ninetoothed-0.18.0}/tests/test_attention.py

@@ -7,11 +7,13 @@ import ninetoothed.language as ntl
 from ninetoothed import Tensor
 from tests.skippers import skip_if_cuda_not_available
 
+BLOCK_SIZE_M = ninetoothed.block_size(lower_bound=64, upper_bound=128)
+BLOCK_SIZE_N = ninetoothed.block_size(lower_bound=32, upper_bound=64)
 
-def arrangement(q, k, v, o):
-    BLOCK_SIZE_M = ninetoothed.block_size(lower_bound=64, upper_bound=128)
-    BLOCK_SIZE_N = ninetoothed.block_size(lower_bound=32, upper_bound=64)
 
+def arrangement(
+    q, k, v, is_causal, o, BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N=BLOCK_SIZE_N
+):
     def arrange_q_or_o(input):
         arranged = input.tile((1, 1, BLOCK_SIZE_M, -1))
         arranged.dtype = arranged.dtype.squeeze((0, 1))
@@ -31,10 +33,16 @@ def arrangement(q, k, v, o):
 
     q_arranged = arrange_q_or_o(q)
 
-    return q_arranged, arrange_k_or_v(k), arrange_k_or_v(v), arrange_q_or_o(o)
+    return (
+        q_arranged,
+        arrange_k_or_v(k),
+        arrange_k_or_v(v),
+        is_causal,
+        arrange_q_or_o(o),
+    )
 
 
-def application(q, k, v, o):
+def application(q, k, v, is_causal, o):
     q_loaded = (q * 1.44269504089).to(q.dtype)
 
     acc = ntl.zeros((q.shape[-2], q.shape[-1]), dtype=ntl.float32)
@@ -45,6 +53,10 @@ def application(q, k, v, o):
         qk = ntl.dot(q_loaded, ntl.trans(k[i]))
         qk = ntl.where(k[i].offsets(-2) < k.source.shape[-2], qk, float("-inf"))
 
+        if is_causal:
+            mask = q.offsets(-2)[:, None] >= k[i].offsets(-2)[None, :]
+            qk = ntl.where(mask, qk, float("-inf"))
+
         m_ij = ntl.maximum(m_i, ntl.max(qk, 1))
         p = ntl.exp2(qk - m_ij[:, None])
         l_ij = ntl.sum(p, 1)
@@ -58,27 +70,28 @@ def application(q, k, v, o):
     o = acc  # noqa: F841
 
 
-def attention(q, k, v):
+def attention(q, k, v, is_causal=False):
     o = torch.empty_like(q, dtype=v.dtype)
 
-    attention_kernel = ninetoothed.make(
-        arrangement,
-        application,
-        (
-            Tensor(
-                4,
-                shape_options=(
-                    None,
-                    None,
-                    {"constexpr": True},
-                    {"constexpr": True, "upper_bound": 128},
-                ),
-            )
-            for _ in range(4)
-        ),
+    q_, k_, v_, o_ = (
+        Tensor(
+            4,
+            shape_options=(
+                None,
+                None,
+                {"constexpr": True},
+                {"constexpr": True, "upper_bound": 128},
+            ),
+        )
+        for _ in range(4)
     )
+    is_causal_ = Tensor(0, constexpr=True)
 
-    attention_kernel(q, k, v, o)
+    tensors = (q_, k_, v_, is_causal_, o_)
+
+    attention_kernel = ninetoothed.make(arrangement, application, tensors)
+
+    attention_kernel(q, k, v, is_causal, o)
 
     return o
 
@@ -87,6 +100,10 @@ def attention(q, k, v):
 class TestCUDA:
     shapes = ((2, 4, 1024, 64), (2, 4, 1, 64))
 
+    dtypes = (torch.float32, torch.float16)
+
+    is_causal_values = (False, True)
+
     @classmethod
     def setup_class(cls):
         torch.manual_seed(0)
@@ -96,23 +113,22 @@ class TestCUDA:
             for shape in cls.shapes
         }
 
+    @pytest.mark.parametrize("is_causal", is_causal_values)
+    @pytest.mark.parametrize("dtype", dtypes)
     @pytest.mark.parametrize("shape", shapes)
-    def test_fp32(self, shape):
-        q, k, v = (arg.to(torch.float32) for arg in type(self).args[shape])
+    def test(self, shape, dtype, is_causal):
+        q, k, v = (arg.to(dtype) for arg in type(self).args[shape])
 
-        assert torch.allclose(
-            attention(q, k, v),
-            F.scaled_dot_product_attention(q, k, v, scale=1),
-            atol=0.01,
-            rtol=0.01,
-        )
-
-    @pytest.mark.parametrize("shape", shapes)
-    def test_fp16(self, shape):
-        q, k, v = (arg.to(torch.float16) for arg in type(self).args[shape])
+        if dtype == torch.float32:
+            atol = 0.025
+            rtol = 0.025
+        elif dtype == torch.float16:
+            atol = 0.01
+            rtol = 0.01
 
         assert torch.allclose(
-            attention(q, k, v),
-            F.scaled_dot_product_attention(q, k, v, scale=1),
-            atol=0.01,
+            attention(q, k, v, is_causal=is_causal),
+            F.scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=1),
+            atol=atol,
+            rtol=rtol,
         )