ninetoothed 0.15.0__py3-none-any.whl → 0.16.0__py3-none-any.whl

ninetoothed/__init__.py

@@ -13,12 +13,13 @@ from ninetoothed.dtype import (
 )
 from ninetoothed.jit import jit
 from ninetoothed.make import make
-from ninetoothed.symbol import Symbol
+from ninetoothed.symbol import Symbol, block_size
 from ninetoothed.tensor import Tensor
 
 __all__ = [
     "Symbol",
     "Tensor",
+    "block_size",
     "float16",
     "float32",
     "float64",

ninetoothed/aot.py

@@ -4,7 +4,7 @@ import subprocess
 import tempfile
 import uuid
 
-from ninetoothed.dtype import int64, uint64
+from ninetoothed.dtype import int64
 from ninetoothed.generation import CACHE_DIR, CodeGenerator
 from ninetoothed.tensor import Tensor
 
@@ -36,7 +36,13 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
 
     code_generator = CodeGenerator()
     source_file = code_generator(
-        func, caller=caller, kernel_name=kernel_name, prettify=False
+        func,
+        caller=caller,
+        kernel_name=kernel_name,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        max_num_configs=None,
+        prettify=False,
     )
 
     tensors = code_generator.tensors
@@ -55,7 +61,7 @@ def _aot(func, caller, kernel_name, num_warps, num_stages):
 
             param_types.append(f"*{dtype}")
         elif Tensor.size_pattern().fullmatch(param):
-            param_types.append(uint64)
+            param_types.append(int64)
         elif Tensor.stride_pattern().fullmatch(param):
             param_types.append(int64)
 

ninetoothed/generation.py

@@ -5,11 +5,20 @@ import functools
 import hashlib
 import inspect
 import itertools
+import json
 import math
+import os
 import pathlib
+import random
+import shutil
 import subprocess
+import tempfile
+import time
+import uuid
 
+import sympy
 import triton
+import triton.language as tl
 
 import ninetoothed.naming as naming
 from ninetoothed.cudaifier import Cudaifier
@@ -19,19 +28,47 @@ from ninetoothed.tensor import Tensor
 from ninetoothed.torchifier import Torchifier
 
 CACHE_DIR = pathlib.Path.home() / ".ninetoothed"
+CACHE_DIR.mkdir(exist_ok=True)
 
 
 class CodeGenerator(ast.NodeTransformer):
     def __init__(self):
         super().__init__()
 
-        self._POWER_OF_TWOS = tuple(2**n for n in range(5, 11))
+        cache_file = CACHE_DIR / "code_generator_cache.json"
 
-        self._MIN_PRODUCT = 2**10
+        log2_min_num_elements = 4
 
-        self._MAX_PRODUCT = 2**20
+        if cache_file.exists():
+            with open(cache_file) as f:
+                cache = json.load(f)
 
-    def __call__(self, func, caller, kernel_name, prettify):
+            log2_max_num_elements = cache["log2_max_num_elements"]
+        else:
+            log2_max_num_elements = _determine_log2_max_num_elements_per_block(
+                log2_min_num_elements
+            )
+
+            cache = {"log2_max_num_elements": log2_max_num_elements}
+
+            with open(cache_file, "w") as f:
+                json.dump(cache, f, indent=4)
+                f.write("\n")
+
+        self._min_num_elements = 2**log2_min_num_elements
+
+        self._max_num_elements = 2**log2_max_num_elements
+
+    def __call__(
+        self,
+        func,
+        caller,
+        kernel_name,
+        num_warps,
+        num_stages,
+        max_num_configs,
+        prettify,
+    ):
         def _get_tree(func):
             module = ast.parse(inspect.getsource(inspect.getmodule(func)))
 
@@ -63,6 +100,12 @@ class CodeGenerator(ast.NodeTransformer):
 
         self._caller = caller
 
+        self._num_wraps = num_warps
+
+        self._num_stages = num_stages
+
+        self._max_num_configs = max_num_configs
+
         self._context = inspect.get_annotations(func)
 
         self._args = list(self._context.values())
@@ -82,6 +125,7 @@ class CodeGenerator(ast.NodeTransformer):
         dependencies = _find_dependencies(func)
         source = "\n\n".join((unparsed, dependencies)).strip()
         source = source.replace(func.__name__, kernel_name)
+        source += "\n"
 
         if prettify:
             for original, simplified in name_collector.simplified_names.items():
@@ -93,9 +137,7 @@ class CodeGenerator(ast.NodeTransformer):
             )
 
         digest = hashlib.sha256(source.encode("utf-8")).hexdigest()
-        cache_dir = CACHE_DIR
-        cache_dir.mkdir(exist_ok=True)
-        cache_file = cache_dir / f"{digest}.py"
+        cache_file = CACHE_DIR / f"{digest}.py"
 
         if not cache_file.exists():
             with open(cache_file, "w", encoding="utf-8") as f:
@@ -110,11 +152,15 @@ class CodeGenerator(ast.NodeTransformer):
     def visit_Module(self, node):
         self.generic_visit(node)
 
-        func_with_auto_tuning = f"{Symbol(self._autotune)}({self._func_def.name})"
+        if self._autotune is not None:
+            func_with_auto_tuning = f"{Symbol(self._autotune)}({self._func_def.name})"
+
+            node.body.append(
+                ast.parse(
+                    f"{self._func_name_with_auto_tuning} = {func_with_auto_tuning}"
+                )
+            )
 
-        node.body.append(
-            ast.parse(f"{self._func_name_with_auto_tuning} = {func_with_auto_tuning}")
-        )
         node.body.append(self._launch)
 
         return node
@@ -136,8 +182,13 @@ class CodeGenerator(ast.NodeTransformer):
     def visit_arguments(self, node):
         self.generic_visit(node)
 
-        names_of_args = [arg.names() - {"ninetoothed"} for arg in self._args]
-        names = functools.reduce(lambda x, y: x | y, names_of_args)
+        symbols = {
+            name.node.id: name
+            for arg in self._args
+            for name in arg.names()
+            if name != "ninetoothed"
+        }
+        names = symbols.keys()
         meta_names = {name for name in names if naming.is_meta(name)}
         non_meta_names = {name for name in names if name not in meta_names}
         non_meta_names |= {
@@ -146,6 +197,8 @@ class CodeGenerator(ast.NodeTransformer):
             if naming.is_constexpr(name)
         }
 
+        self._symbols = symbols
+
         non_meta_names = sorted(non_meta_names)
         meta_names = sorted(meta_names)
 
@@ -160,6 +213,12 @@ class CodeGenerator(ast.NodeTransformer):
         ]
 
         self._autotune = self._generate_autotune(non_meta_names, meta_names)
+
+        if self._autotune is not None:
+            self._func_name = self._func_name_with_auto_tuning
+        else:
+            self._func_name = self._func_def.name
+
         self._func_def.decorator_list = [Symbol("triton.jit").node]
 
         self._launch = self._generate_launch(non_meta_names, meta_names)
@@ -243,12 +302,69 @@ class CodeGenerator(ast.NodeTransformer):
         return isinstance(node, ast.Name) and node.id in self._context
 
     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"]
+        inequalities = True
+
+        for arg in self._args:
+            if arg.ndim == 0:
+                continue
+
+            num_elements = sympy.simplify(str(math.prod(arg.innermost().shape)))
+
+            inequalities &= num_elements <= self._max_num_elements
+            inequalities &= num_elements >= self._min_num_elements
+
+        values_of_meta_params = []
+
+        for param in meta:
+            symbol = self._symbols[param]
+
+            values = range(symbol.lower_bound, symbol.upper_bound + 1)
+
+            if symbol.power_of_two:
+                values = tuple(value for value in values if value & (value - 1) == 0)
+            else:
+                values = tuple(values)
+
+            values_of_meta_params.append(values)
+
+        max_values_of_non_meta_params = {}
+
+        for free_symbol in inequalities.free_symbols:
+            symbol_str = str(free_symbol)
+
+            if symbol_str in meta:
+                continue
+
+            symbol = self._symbols[symbol_str]
+
+            max_values_of_non_meta_params[symbol_str] = symbol.upper_bound
+
+        block_size_configs = []
+
+        for values in itertools.product(*values_of_meta_params):
+            config = {param: value for param, value in zip(meta, values)}
+
+            if sympy.logic.simplify_logic(
+                inequalities.subs(config | max_values_of_non_meta_params)
+            ):
+                block_size_configs.append(config)
+
+        if isinstance(self._num_wraps, collections.abc.Iterable):
+            num_warps_configs = self._num_wraps
+        else:
+            num_warps_configs = (self._num_wraps,)
+
+        if isinstance(self._num_stages, collections.abc.Iterable):
+            num_stages_configs = self._num_stages
+        else:
+            num_stages_configs = (self._num_stages,)
 
-        num_warps = 8
-        num_stages = max_shared_mem // 2**15
+        compiler_configs = tuple(
+            {"num_warps": num_warps, "num_stages": num_stages}
+            for num_warps, num_stages in itertools.product(
+                num_warps_configs, num_stages_configs
+            )
+        )
 
         configs = [
             ast.Call(
@@ -259,19 +375,38 @@ class CodeGenerator(ast.NodeTransformer):
                 ),
                 args=[
                     ast.Dict(
-                        keys=[ast.Constant(value=param) for param in meta],
-                        values=[ast.Constant(value=value) for value in values],
+                        keys=[
+                            ast.Constant(value=param)
+                            for param in block_size_config.keys()
+                        ],
+                        values=[
+                            ast.Constant(value=value)
+                            for value in block_size_config.values()
+                        ],
                     )
                 ],
                 keywords=[
-                    ast.keyword(arg="num_warps", value=ast.Constant(value=num_warps)),
-                    ast.keyword(arg="num_stages", value=ast.Constant(value=num_stages)),
+                    ast.keyword(
+                        arg="num_warps",
+                        value=ast.Constant(value=compiler_config["num_warps"]),
+                    ),
+                    ast.keyword(
+                        arg="num_stages",
+                        value=ast.Constant(value=compiler_config["num_stages"]),
+                    ),
                 ],
             )
-            for values in itertools.product(self._POWER_OF_TWOS, repeat=len(meta))
-            if self._MIN_PRODUCT <= math.prod(values) <= self._MAX_PRODUCT
+            for block_size_config, compiler_config in itertools.product(
+                block_size_configs, compiler_configs
+            )
         ]
 
+        if self._max_num_configs is not None and len(configs) > self._max_num_configs:
+            configs = random.sample(configs, k=self._max_num_configs)
+
+        if not configs:
+            return None
+
         return ast.Call(
             func=ast.Attribute(
                 value=ast.Name(id="ninetoothed", ctx=ast.Load()),
@@ -357,9 +492,7 @@ class CodeGenerator(ast.NodeTransformer):
                 ast.Expr(
                     ast.Call(
                         func=ast.Subscript(
-                            value=ast.Name(
-                                id=self._func_name_with_auto_tuning, ctx=ast.Load()
-                            ),
+                            value=ast.Name(id=self._func_name, ctx=ast.Load()),
                             slice=self._generate_grid(),
                             ctx=ast.Load(),
                         ),
@@ -993,3 +1126,82 @@ class _FunctionDefFinder(ast.NodeVisitor):
             self.result = node
 
         self.generic_visit(node)
+
+
+def _determine_log2_max_num_elements_per_block(
+    min_exponent, max_exponent=30, num_iterations=3
+):
+    _profile_pseudo_add_kernel(1)
+
+    for n in range(min_exponent, max_exponent + 1):
+        elapsed_time = 0
+
+        for _ in range(num_iterations):
+            elapsed_time += _profile_pseudo_add_kernel(2**n)
+
+        average_elapsed_time = elapsed_time / num_iterations
+
+        if average_elapsed_time >= 1:
+            return n - 1
+
+
+def _profile_pseudo_add_kernel(block_size):
+    cache_dir = triton.runtime.cache.default_cache_dir()
+    os.makedirs(cache_dir, exist_ok=True)
+
+    with tempfile.TemporaryDirectory() as backup_dir:
+        backup_path = os.path.join(backup_dir, str(uuid.uuid4()))
+
+        if os.path.exists(backup_path):
+            shutil.rmtree(backup_path)
+
+        shutil.move(cache_dir, backup_path)
+
+        try:
+            start_time = time.time()
+
+            _run_pseudo_add_kernel(block_size)
+
+            end_time = time.time()
+
+            elapsed_time = end_time - start_time
+        finally:
+            if os.path.exists(cache_dir):
+                shutil.rmtree(cache_dir)
+
+            shutil.move(backup_path, cache_dir)
+
+        return elapsed_time
+
+
+def _run_pseudo_add_kernel(block_size):
+    @triton.jit
+    def kernel(a_ptr, b_ptr, c_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
+        pid = tl.program_id(0)
+
+        offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        mask = offs < num_elements
+
+        a = tl.load(a_ptr + offs, mask=mask)
+        b = tl.load(b_ptr + offs, mask=mask)
+
+        c = a + b
+
+        tl.store(c_ptr + offs, c, mask=mask)
+
+    num_elements = 0
+    shape = (num_elements,)
+    dtype = tl.float32
+
+    a = Tensor(shape=shape, dtype=dtype)
+    b = Tensor(shape=shape, dtype=dtype)
+    c = Tensor(shape=shape, dtype=dtype)
+
+    def data_ptr():
+        return 0
+
+    a.data_ptr = data_ptr
+    b.data_ptr = data_ptr
+    c.data_ptr = data_ptr
+
+    kernel[(1,)](a, b, c, num_elements, block_size)

ninetoothed/jit.py

@@ -4,12 +4,25 @@ import sys
 from ninetoothed.generation import CodeGenerator
 
 
-def jit(func=None, *, caller="torch", kernel_name=None, _prettify=False):
+def jit(
+    func=None,
+    *,
+    caller="torch",
+    kernel_name=None,
+    num_warps=None,
+    num_stages=None,
+    max_num_configs=None,
+    _prettify=False,
+):
     """A decorator for generating compute kernels.
 
     :param func: The function to be compiled.
     :param caller: Who will call the compute kernel.
     :param kernel_name: The name for the generated kernel.
+    :param num_warps: The number of warps to use.
+    :param num_stages: The number of pipeline stages.
+    :param max_num_configs: The maximum number of auto-tuning
+        configurations to use.
     :param _prettify: Whether to prettify the generated code.
     :return: A handle to the compute kernel.
 
@@ -20,7 +33,15 @@ def jit(func=None, *, caller="torch", kernel_name=None, _prettify=False):
     """
 
     def wrapper(func):
-        return JIT(func, caller=caller, kernel_name=kernel_name, _prettify=_prettify)()
+        return JIT(
+            func,
+            caller=caller,
+            kernel_name=kernel_name,
+            num_warps=num_warps,
+            num_stages=num_stages,
+            max_num_configs=max_num_configs,
+            _prettify=_prettify,
+        )()
 
     if func is None:
         return wrapper
@@ -29,7 +50,16 @@ def jit(func=None, *, caller="torch", kernel_name=None, _prettify=False):
 
 
 class JIT:
-    def __init__(self, func, caller, kernel_name, _prettify=False):
+    def __init__(
+        self,
+        func,
+        caller,
+        kernel_name,
+        num_warps,
+        num_stages,
+        max_num_configs,
+        _prettify=False,
+    ):
         self.func = func
 
         self._caller = caller
@@ -39,12 +69,24 @@ class JIT:
         else:
             self._kernel_name = func.__name__
 
+        self._num_warps = num_warps
+
+        self._num_stages = num_stages
+
+        self._max_num_configs = max_num_configs
+
         self._prettify = _prettify
 
     def __call__(self):
         code_generator = CodeGenerator()
         source_file = code_generator(
-            self.func, self._caller, self._kernel_name, self._prettify
+            self.func,
+            self._caller,
+            self._kernel_name,
+            self._num_warps,
+            self._num_stages,
+            self._max_num_configs,
+            self._prettify,
         )
         module = type(self)._import_from_path(source_file, source_file)
         module_vars = vars(module)

ninetoothed/make.py

@@ -2,6 +2,7 @@ import inspect
 
 from ninetoothed.aot import aot
 from ninetoothed.jit import jit
+from ninetoothed.utils import calculate_default_configs
 
 
 def make(
@@ -11,8 +12,9 @@ def make(
     caller="torch",
     kernel_name=None,
     output_dir=None,
-    num_warps=4,
-    num_stages=3,
+    num_warps=None,
+    num_stages=None,
+    max_num_configs=None,
 ):
     """Integrate the arrangement and the application of the tensors.
 
@@ -24,16 +26,33 @@ def make(
     :param output_dir: The directory to store the generated files.
     :param num_warps: The number of warps to use.
     :param num_stages: The number of pipeline stages.
+    :param max_num_configs: The maximum number of auto-tuning
+        configurations to use.
     :return: A handle to the compute kernel.
     """
 
+    default_num_warps, default_num_stages = calculate_default_configs()
+
+    if num_warps is None:
+        num_warps = default_num_warps
+
+    if num_stages is None:
+        num_stages = default_num_stages
+
     params = inspect.signature(application).parameters
     types = arrangement(*tensors)
     annotations = {param: type for param, type in zip(params, types)}
     application.__annotations__ = annotations
 
     if caller == "torch":
-        return jit(application, caller=caller, kernel_name=kernel_name)
+        return jit(
+            application,
+            caller=caller,
+            kernel_name=kernel_name,
+            num_warps=num_warps,
+            num_stages=num_stages,
+            max_num_configs=max_num_configs,
+        )
 
     return aot(
         application,

ninetoothed/symbol.py

@@ -12,9 +12,20 @@ class Symbol:
     :param expr: The expression used to construct the symbol.
     :param constexpr: Whether the symbol is a constexpr.
     :param mata: Whether the symbol is a meta.
+    :param lower_bound: The minimum value for the symbol's range.
+    :param upper_bound: The maximum value for the symbol's range.
+    :param power_of_two: Whether the value should be a power of two.
     """
 
-    def __init__(self, expr, constexpr=None, meta=None):
+    def __init__(
+        self,
+        expr,
+        constexpr=None,
+        meta=None,
+        lower_bound=None,
+        upper_bound=None,
+        power_of_two=None,
+    ):
         if isinstance(expr, type(self)):
             self._node = expr._node
             return
@@ -43,6 +54,40 @@ class Symbol:
         if constexpr:
             self._node.id = naming.make_constexpr(self._node.id)
 
+        self._node.symbol = self
+
+        DEFAULT_LOWER_BOUND_FOR_META_SYMBOLS = 2**5
+        DEFAULT_UPPER_BOUND_FOR_META_SYMBOLS = 2**10
+        DEFAULT_POWER_OF_TWO_FOR_META_SYMBOLS = True
+
+        DEFAULT_LOWER_BOUND_FOR_NON_META_CONSTEXPR_SYMBOLS = 1
+        DEFAULT_UPPER_BOUND_FOR_NON_META_CONSTEXPR_SYMBOLS = 2**20
+        DEFAULT_POWER_OF_TWO_FOR_NON_META_CONSTEXPR_SYMBOLS = False
+
+        if lower_bound is not None:
+            self.lower_bound = lower_bound
+        else:
+            if meta:
+                self.lower_bound = DEFAULT_LOWER_BOUND_FOR_META_SYMBOLS
+            elif constexpr:
+                self.lower_bound = DEFAULT_LOWER_BOUND_FOR_NON_META_CONSTEXPR_SYMBOLS
+
+        if upper_bound is not None:
+            self.upper_bound = upper_bound
+        else:
+            if meta:
+                self.upper_bound = DEFAULT_UPPER_BOUND_FOR_META_SYMBOLS
+            elif constexpr:
+                self.upper_bound = DEFAULT_UPPER_BOUND_FOR_NON_META_CONSTEXPR_SYMBOLS
+
+        if power_of_two is not None:
+            self.power_of_two = power_of_two
+        else:
+            if meta:
+                self.power_of_two = DEFAULT_POWER_OF_TWO_FOR_META_SYMBOLS
+            elif constexpr:
+                self.power_of_two = DEFAULT_POWER_OF_TWO_FOR_NON_META_CONSTEXPR_SYMBOLS
+
     def __eq__(self, other):
         if isinstance(self._node, ast.Constant):
             if isinstance(other, Symbol) and isinstance(other._node, ast.Constant):
@@ -155,7 +200,7 @@ class Symbol:
             def visit_Name(self, node):
                 self.generic_visit(node)
 
-                self.names.add(node.id)
+                self.names.add(node.symbol)
 
         name_collector = NameCollector()
 
@@ -179,6 +224,24 @@ class Symbol:
         return isinstance(object, Symbol) and isinstance(object.node, ast.Name)
 
 
+def block_size(lower_bound=None, upper_bound=None):
+    """Create a block size symbol that serves as a meta-parameter.
+
+    :param lower_bound: The lower bound for the block size's range.
+    :param upper_bound: The upper bound for the block size's range.
+    :return: A block size symbol that serves as a meta-parameter.
+    """
+
+    name = naming.auto_generate(f"BLOCK_SIZE_{block_size._num_block_sizes}")
+
+    block_size._num_block_sizes += 1
+
+    return Symbol(name, meta=True, lower_bound=lower_bound, upper_bound=upper_bound)
+
+
+block_size._num_block_sizes = 0
+
+
 class _FindAndReplacer(ast.NodeTransformer):
     def __init__(self, targets, replacement):
         self._targets_unparsed = tuple(

ninetoothed/tensor.py

@@ -14,7 +14,7 @@ class Tensor:
     :param dtype: The element type of the tensor.
     :param strides: The strides of the tensor.
     :param other: The values for out-of-bounds positions.
-    :param constexpr_shape: Whether the sizes are constexpr.
+    :param shape_options: The options for configuring shape symbols.
     :param name: The name of the tensor.
     :param source: For internal use only.
     :param source_dims: For internal use only.
@@ -31,7 +31,7 @@ class Tensor:
         dtype=None,
         strides=None,
         other=None,
-        constexpr_shape=None,
+        shape_options=None,
         name=None,
         source=None,
         source_dims=None,
@@ -48,9 +48,20 @@ class Tensor:
             self.name = naming.auto_generate(f"tensor_{type(self).num_instances}")
 
         if ndim is not None:
+            if shape_options is None:
+                shape_options = tuple({} for _ in range(ndim))
+
+            if isinstance(shape_options, dict):
+                shape_options = tuple(shape_options for _ in range(ndim))
+
+            shape_options = tuple(
+                size_options if size_options is not None else {}
+                for size_options in shape_options
+            )
+
             self.shape = (
-                Symbol(self.size_string(i), constexpr=constexpr_shape)
-                for i in range(ndim)
+                Symbol(self.size_string(i), **size_options)
+                for i, size_options in zip(range(ndim), shape_options)
             )
             self.strides = (Symbol(self.stride_string(i)) for i in range(ndim))
         else:
@@ -146,7 +157,7 @@ class Tensor:
             )
             outer_shape.append(new_size)
 
-            new_stride = self_stride * stride // spacing
+            new_stride = self_stride * stride
             outer_strides.append(new_stride)
 
             inner_shape.append(tile_size)
@@ -364,10 +375,10 @@ class Tensor:
 
     def names(self):
         if self.ndim == 0:
-            return {self.source.name}
+            return {Symbol(self.source.name)}
 
         return (
-            {self.source.pointer_string()}
+            {Symbol(self.source.pointer_string())}
             | {
                 name
                 for value in itertools.chain(self.shape, self.strides)

ninetoothed/utils.py

@@ -0,0 +1,12 @@
+import triton
+
+
+def calculate_default_configs():
+    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
+
+    return num_warps, num_stages

ninetoothed/visualization.py

@@ -118,10 +118,16 @@ def _visualize_unit_square(ax, x, y, color):
 
 
 def _visualize_rect(ax, width, height, x, y, color):
-    pos_x, pos_y = zip(*_verts_of_rect(width, height, x, y))
-
-    ax.fill(pos_x, pos_y, color)
-    ax.plot(pos_x + (pos_x[0],), pos_y + (pos_y[0],), "k")
+    ax.add_patch(
+        plt.Rectangle(
+            (x, y),
+            width,
+            height,
+            edgecolor="k",
+            facecolor=color,
+            linewidth=plt.rcParams["lines.linewidth"],
+        )
+    )
 
 
 def _verts_of_rect(width, height, x, y):

{ninetoothed-0.15.0.dist-info → ninetoothed-0.16.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.15.0
+Version: 0.16.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: sympy>=1.13.0
 Requires-Dist: triton>=3.0.0
 Provides-Extra: all
 Requires-Dist: matplotlib>=3.9.0; extra == 'all'

ninetoothed-0.16.0.dist-info/RECORD

@@ -0,0 +1,18 @@
+ninetoothed/__init__.py,sha256=F2bxRNhzcGdtADA8RehTuf-QK0xnxno8kxvr6H2L5Tg,552
+ninetoothed/aot.py,sha256=8ZCLtnsign14YvY7SXX5ASidhuUAhPwppTXUJNkQup4,6243
+ninetoothed/cudaifier.py,sha256=5ylMr1q0B9NwbeXkpCu3o2nMGpDfh65nAQ0Az_qMQuI,877
+ninetoothed/dtype.py,sha256=-0iBleay5gYA4wtT3l17QjCesr7g26M6CSfhNJdI3k4,165
+ninetoothed/generation.py,sha256=VIqSyZT4yHxY_a2QPmWW6jjALv3e1mohDqdRQBRYsAo,36462
+ninetoothed/jit.py,sha256=CpeSkO_zUe9DwtTJ2K2H7Bwpx-FvIHfrgzOcEosfpek,2946
+ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+ninetoothed/make.py,sha256=fQKuRJL7HC2iGTAN323mlIWXz9Z3jotIoN68ur29Qlw,1834
+ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
+ninetoothed/symbol.py,sha256=lJo3NL2-T7tKbKjb6MCRLMemN94mqS3bIiG943P0Mbo,7454
+ninetoothed/tensor.py,sha256=gQEzHTcXqZVBFLc2YRfXTKxjxPWMxWN7fNl2BCfJwMs,14782
+ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
+ninetoothed/utils.py,sha256=mtRXABBVPnlgd2n1REh9oB3s_5bUsKhd3iwu3oJ5DSQ,338
+ninetoothed/visualization.py,sha256=zlMH-0WplaboePGzcbpcj4UovpX0k2r4SysSPsNS4r4,3674
+ninetoothed-0.16.0.dist-info/METADATA,sha256=nkq3iImebtmcEs-bZq2zfF2_QxrZD9IWky1S86OnUMA,7340
+ninetoothed-0.16.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ninetoothed-0.16.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.16.0.dist-info/RECORD,,

ninetoothed-0.15.0.dist-info/RECORD

@@ -1,17 +0,0 @@
-ninetoothed/__init__.py,sha256=zGaZiUzwJZ2jfwLxp7lT8ll_V5ngP5QYrfVbapftbCY,522
-ninetoothed/aot.py,sha256=1hzl4-6MqscB4tDMqmLCOTlyzsYkbY20EnmDgHO8hU4,6137
-ninetoothed/cudaifier.py,sha256=5ylMr1q0B9NwbeXkpCu3o2nMGpDfh65nAQ0Az_qMQuI,877
-ninetoothed/dtype.py,sha256=-0iBleay5gYA4wtT3l17QjCesr7g26M6CSfhNJdI3k4,165
-ninetoothed/generation.py,sha256=QHrK7DOuJo5wEV-5HAqu_e-suuD4TPPNmCzrUMRYF2w,30940
-ninetoothed/jit.py,sha256=0MFbFIODtw-bxuOC7WByxiVtQMeyvZkoDxvfAZ9rIFQ,2120
-ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
-ninetoothed/make.py,sha256=wRr3JwGt5E2OCquq_nzBZljdW-AJPOqH49cM08gwl4A,1287
-ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
-ninetoothed/symbol.py,sha256=UpGmx_jvaDtowADnp1DwYC3fvBXSiaMiYpU-ewkVo50,5261
-ninetoothed/tensor.py,sha256=W1XY8_vaYmszX4lIWuas-ZKGbbdEZU7Z5h1A4FBXDXg,14358
-ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
-ninetoothed/visualization.py,sha256=IZ7iTT4dl5_JFbO-WfSWPFWpgkyPr4nylwhSZVy8gss,3601
-ninetoothed-0.15.0.dist-info/METADATA,sha256=UmND-TBDf7vrdii9dhiOZmiZBzrcG8-xEniThTSikqM,7311
-ninetoothed-0.15.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ninetoothed-0.15.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ninetoothed-0.15.0.dist-info/RECORD,,