triton-windows 3.3.1.post19__cp311-cp311-win_amd64.whl → 3.5.0.post21__cp311-cp311-win_amd64.whl

triton/__init__.py

@@ -1,5 +1,5 @@
 """isort:skip_file"""
-__version__ = '3.3.1'
+__version__ = '3.5.0'
 
 # ---------------------------------------
 # Note: import order is significant here.
@@ -17,7 +17,8 @@ from .runtime import (
     InterpreterError,
     MockTensor,
 )
-from .runtime.jit import jit
+from .runtime.jit import constexpr_function, jit
+from .runtime._async_compile import AsyncCompileMode, FutureKernel
 from .compiler import compile, CompilationError
 from .errors import TritonError
 from .runtime._allocation import set_allocator
@@ -26,12 +27,17 @@ from . import language
 from . import testing
 from . import tools
 
+must_use_result = language.core.must_use_result
+
 __all__ = [
+    "AsyncCompileMode",
     "autotune",
     "cdiv",
     "CompilationError",
     "compile",
     "Config",
+    "constexpr_function",
+    "FutureKernel",
     "heuristics",
     "InterpreterError",
     "jit",
@@ -39,6 +45,7 @@ __all__ = [
     "KernelInterface",
     "language",
     "MockTensor",
+    "must_use_result",
     "next_power_of_2",
     "OutOfResources",
     "reinterpret",
@@ -56,10 +63,12 @@ __all__ = [
 # -------------------------------------
 
 
+@constexpr_function
 def cdiv(x: int, y: int):
     return (x + y - 1) // y
 
 
+@constexpr_function
 def next_power_of_2(n: int):
     """Return the smallest power of 2 greater than or equal to n"""
     n -= 1

triton/_filecheck.py

@@ -0,0 +1,97 @@
+import functools
+import os
+import inspect
+import subprocess
+import tempfile
+
+import triton
+from triton.compiler import ASTSource, make_backend
+from triton.backends.compiler import GPUTarget
+from triton.experimental.gluon._runtime import GluonASTSource
+from triton.runtime.jit import create_function_from_signature
+from triton._C.libtriton import ir
+
+# ===-----------------------------------------------------------------------===#
+# filecheck_test
+# ===-----------------------------------------------------------------------===#
+
+# Stub target for testing the frontend.
+stub_target = GPUTarget("cuda", 100, 32)
+
+triton_dir = os.path.dirname(__file__)
+filecheck_path = os.path.join(triton_dir, "FileCheck")
+
+
+class MatchError(ValueError):
+
+    def __init__(self, message, module_str):
+        super().__init__(message)
+        self.module_str = module_str
+
+    def __str__(self):
+        return f"{super().__str__()}\n{self.module_str}"
+
+
+def run_filecheck(name, module_str, check_template):
+    with tempfile.TemporaryDirectory() as tempdir:
+        temp_module = os.path.join(tempdir, "module")
+        with open(temp_module, "w") as temp:
+            temp.write(module_str)
+
+        temp_expected = os.path.join(tempdir, "expected")
+        with open(temp_expected, "w") as temp:
+            temp.write(check_template)
+
+        try:
+            subprocess.check_output(
+                [filecheck_path, temp_expected, "--input-file", temp_module, "--dump-input-context=50"],
+                stderr=subprocess.STDOUT)
+        except subprocess.CalledProcessError as error:
+            decoded = error.output.decode('unicode_escape')
+            raise ValueError(decoded)
+
+
+def run_parser(kernel_fn, args=(), kwargs={}, target=stub_target):
+    if "sanitize_overflow" not in kwargs:
+        kwargs = dict(kwargs)
+        kwargs["sanitize_overflow"] = False
+    backend = make_backend(target)
+    binder = create_function_from_signature(
+        kernel_fn.signature,
+        kernel_fn.params,
+        backend,
+    )
+
+    bound_args, specialization, options = binder(*args, **kwargs)
+    options, signature, constexprs, attrs = kernel_fn._pack_args(backend, kwargs, bound_args, specialization, options)
+    source_cls = GluonASTSource if kernel_fn.is_gluon() else ASTSource
+    src = source_cls(kernel_fn, signature, constexprs, attrs)
+
+    context = ir.context()
+    ir.load_dialects(context)
+    backend.load_dialects(context)
+
+    codegen_fns = backend.get_codegen_implementation(options)
+    module_map = backend.get_module_map()
+    module = src.make_ir(target, options, codegen_fns, module_map, context)
+    assert module.verify()
+    return module
+
+
+def run_filecheck_test(kernel_fn):
+    assert isinstance(kernel_fn, triton.runtime.JITFunction)
+    check_template = inspect.getsource(kernel_fn.fn)
+    if check_template is None:
+        raise ValueError("kernel function must have a docstring with FileCheck template")
+    mlir_module = run_parser(kernel_fn)
+
+    run_filecheck("placeholder", mlir_module.str_nodebug(), check_template)
+
+
+def filecheck_test(fn):
+
+    @functools.wraps(fn)
+    def test_fn():
+        run_filecheck_test(fn)
+
+    return test_fn

triton/_internal_testing.py

@@ -4,11 +4,11 @@ import numpy as np
 import torch
 import triton
 import triton.language as tl
-from triton.backends.nvidia.compiler import _path_to_binary
+from triton import knobs
+from typing import Optional, Set, Union
 import pytest
 
 from numpy.random import RandomState
-from typing import Optional, Union
 from triton.runtime.jit import TensorWrapper, reinterpret, type_canonicalisation_dict
 
 int_dtypes = ['int8', 'int16', 'int32', 'int64']
@@ -20,6 +20,7 @@ dtypes = integral_dtypes + float_dtypes
 dtypes_with_bfloat16 = dtypes + ['bfloat16']
 torch_float8_dtypes = ['float8_e4m3fn', 'float8_e5m2']
 torch_dtypes = ['bool'] + int_dtypes + ['uint8'] + float_dtypes + ['bfloat16']
+tma_dtypes = sorted(set(dtypes_with_bfloat16) - {"int64", "uint64", "float64"})
 
 
 def is_interpreter():
@@ -37,38 +38,58 @@ def is_cuda():
     return False if target is None else target.backend == "cuda"
 
 
-def is_hopper():
+def is_ampere_or_newer():
+    return is_cuda() and torch.cuda.get_device_capability()[0] >= 8
+
+
+def is_blackwell():
+    return is_cuda() and torch.cuda.get_device_capability()[0] == 10
+
+
+def is_hopper_or_newer():
     return is_cuda() and torch.cuda.get_device_capability()[0] >= 9
 
 
+def is_hopper():
+    return is_cuda() and torch.cuda.get_device_capability()[0] == 9
+
+
 def is_hip():
     target = get_current_target()
     return False if target is None else target.backend == "hip"
 
 
-def is_hip_mi200():
+def is_hip_cdna2():
     target = get_current_target()
-    if target is None or target.backend != 'hip':
-        return False
-    return target.arch == 'gfx90a'
+    return target is not None and target.backend == 'hip' and target.arch == 'gfx90a'
 
 
-def is_hip_mi300():
+def is_hip_cdna3():
     target = get_current_target()
-    if target is None or target.backend != 'hip':
-        return False
-    return target.arch in ('gfx940', 'gfx941', 'gfx942')
+    return target is not None and target.backend == 'hip' and target.arch == 'gfx942'
 
 
-def is_hip_mi350():
+def is_hip_cdna4():
     target = get_current_target()
-    if target is None or target.backend != 'hip':
-        return False
-    return target.arch in ('gfx950')
+    return target is not None and target.backend == 'hip' and target.arch == 'gfx950'
+
+
+def is_hip_gfx11():
+    target = get_current_target()
+    return target is not None and target.backend == 'hip' and 'gfx11' in target.arch
+
+
+def is_hip_gfx12():
+    target = get_current_target()
+    return target is not None and target.backend == 'hip' and 'gfx12' in target.arch
 
 
 def is_hip_cdna():
-    return is_hip_mi200() or is_hip_mi300() or is_hip_mi350()
+    return is_hip_cdna2() or is_hip_cdna3() or is_hip_cdna4()
+
+
+def get_hip_lds_size():
+    return 163840 if is_hip_cdna4() else 65536
 
 
 def is_xpu():
@@ -131,7 +152,7 @@ def to_triton(x: np.ndarray, device, dst_type=None) -> Union[TensorWrapper, torc
 
 
 def str_to_triton_dtype(x: str) -> tl.dtype:
-    return tl.str_to_ty(type_canonicalisation_dict[x])
+    return tl.str_to_ty(type_canonicalisation_dict[x], None)
 
 
 def torch_dtype_name(dtype) -> str:
@@ -161,7 +182,7 @@ def supports_tma(byval_only=False):
         return True
     if not is_cuda():
         return False
-    _, cuda_version = _path_to_binary("ptxas")
+    cuda_version = knobs.nvidia.ptxas.version
     min_cuda_version = (12, 0) if byval_only else (12, 3)
     cuda_version_tuple = tuple(map(int, cuda_version.split(".")))
     assert len(cuda_version_tuple) == 2, cuda_version_tuple
@@ -176,3 +197,59 @@ def tma_skip_msg(byval_only=False):
 
 
 requires_tma = pytest.mark.skipif(not supports_tma(), reason=tma_skip_msg())
+
+
+def default_alloc_fn(size: int, align: int, _):
+    return torch.empty(size, dtype=torch.int8, device="cuda")
+
+
+def unwrap_tensor(t: Union[torch.Tensor, triton.runtime.jit.TensorWrapper]) -> torch.Tensor:
+    if isinstance(t, triton.runtime.jit.TensorWrapper):
+        return t.base
+    return t
+
+
+def _fresh_knobs_impl(skipped_attr: Optional[Set[str]] = None):
+    from triton import knobs
+
+    if skipped_attr is None:
+        skipped_attr = set()
+
+    monkeypatch = pytest.MonkeyPatch()
+
+    knobs_map = {
+        name: knobset
+        for name, knobset in knobs.__dict__.items()
+        if isinstance(knobset, knobs.base_knobs) and knobset != knobs.base_knobs and name not in skipped_attr
+    }
+
+    # We store which variables we need to unset below in finally because
+    # monkeypatch doesn't appear to reset variables that were never set
+    # before the monkeypatch.delenv call below.
+    env_to_unset = []
+    prev_propagate_env = knobs.propagate_env
+
+    def fresh_function():
+        nonlocal env_to_unset
+        for name, knobset in knobs_map.items():
+            setattr(knobs, name, knobset.copy().reset())
+            for knob in knobset.knob_descriptors.values():
+                if knob.key in os.environ:
+                    monkeypatch.delenv(knob.key, raising=False)
+                else:
+                    env_to_unset.append(knob.key)
+        knobs.propagate_env = True
+        return knobs
+
+    def reset_function():
+        for name, knobset in knobs_map.items():
+            setattr(knobs, name, knobset)
+        # `undo` should be placed before `del os.environ`
+        # Otherwise, it may restore environment variables that monkeypatch deleted
+        monkeypatch.undo()
+        for k in env_to_unset:
+            if k in os.environ:
+                del os.environ[k]
+        knobs.propagate_env = prev_propagate_env
+
+    return fresh_function, reset_function

triton/_utils.py

@@ -1,35 +1,126 @@
+from __future__ import annotations
+
 from functools import reduce
+from typing import Any, Callable, TYPE_CHECKING, Union, List, Dict
+
+if TYPE_CHECKING:
+    from .language import core
+    IterableType = Union[list[Any], tuple[Any, ...], core.tuple, core.tuple_type]
+    ObjPath = tuple[int, ...]
 
+TRITON_MAX_TENSOR_NUMEL = 1048576
 
-def get_iterable_path(iterable, path):
-    return reduce(lambda a, idx: a[idx], path, iterable)
 
+def get_iterable_path(iterable: IterableType, path: ObjPath) -> Any:
+    return reduce(lambda a, idx: a[idx], path, iterable)  # type: ignore[index]
 
-def set_iterable_path(iterable, path, val):
+
+def set_iterable_path(iterable: IterableType, path: tuple[int, ...], val: Any):
+    from .language import core
+    assert len(path) != 0
     prev = iterable if len(path) == 1 else get_iterable_path(iterable, path[:-1])
-    prev[path[-1]] = val
+    assert isinstance(prev, core.tuple)
+    prev._setitem(path[-1], val)
 
 
-def find_paths_if(iterable, pred):
+def find_paths_if(iterable: Union[IterableType, Any], pred: Callable[[ObjPath, Any], bool]) -> list[ObjPath]:
     from .language import core
-    is_iterable = lambda x: isinstance(x, (list, tuple, core.tuple, core.tuple_type))
-    ret = dict()
+    is_iterable: Callable[[Any], bool] = lambda x: isinstance(x, (list, tuple, core.tuple, core.tuple_type))
+    # We need to use dict so that ordering is maintained, while set doesn't guarantee order
+    ret: dict[ObjPath, None] = {}
 
-    def _impl(current, path):
-        path = (path[0], ) if len(path) == 1 else tuple(path)
+    def _impl(path: tuple[int, ...], current: Any):
         if is_iterable(current):
             for idx, item in enumerate(current):
-                _impl(item, path + (idx, ))
+                _impl((*path, idx), item)
         elif pred(path, current):
-            if len(path) == 1:
-                ret[(path[0], )] = None
-            else:
-                ret[tuple(path)] = None
-
-    if is_iterable(iterable):
-        _impl(iterable, [])
-    elif pred(list(), iterable):
-        ret = {tuple(): None}
-    else:
-        ret = dict()
+            ret[path] = None
+
+    _impl((), iterable)
+
     return list(ret.keys())
+
+
+def is_power_of_two(x):
+    return (x & (x - 1)) == 0
+
+
+def validate_block_shape(shape: List[int]):
+    numel = 1
+    for i, d in enumerate(shape):
+        if not isinstance(d, int):
+            raise TypeError(f"Shape element {i} must have type `constexpr[int]`, got `constexpr[{type(d)}]")
+        if not is_power_of_two(d):
+            raise ValueError(f"Shape element {i} must be a power of 2")
+        numel *= d
+
+    if numel > TRITON_MAX_TENSOR_NUMEL:
+        raise ValueError(f"numel ({numel}) exceeds triton maximum tensor numel ({TRITON_MAX_TENSOR_NUMEL})")
+    return numel
+
+
+type_canonicalisation_dict = {
+    # we canonicalise all bools to be unsigned:
+    "bool": "u1",
+    "int1": "u1",
+    "uint1": "u1",
+    "i1": "u1",
+    # floating-point dtypes:
+    "float8e4nv": "fp8e4nv",
+    "float8e5": "fp8e5",
+    "float8e4b15": "fp8e4b15",
+    "float8_e4m3fn": "fp8e4nv",
+    "float8e4b8": "fp8e4b8",
+    "float8_e4m3fnuz": "fp8e4b8",
+    "float8_e5m2": "fp8e5",
+    "float8e5b16": "fp8e5b16",
+    "float8_e5m2fnuz": "fp8e5b16",
+    "half": "fp16",
+    "float16": "fp16",
+    "bfloat16": "bf16",
+    "float": "fp32",
+    "float32": "fp32",
+    "double": "fp64",
+    "float64": "fp64",
+    # signed integers:
+    "int8": "i8",
+    "int16": "i16",
+    "int": "i32",
+    "int32": "i32",
+    "int64": "i64",
+    # unsigned integers:
+    "uint8": "u8",
+    "uint16": "u16",
+    "uint32": "u32",
+    "uint64": "u64",
+    "void": "void",
+}
+
+for v in list(type_canonicalisation_dict.values()):
+    type_canonicalisation_dict[v] = v
+
+
+def canonicalize_dtype(dtype):
+    dtype_str = str(dtype).split(".")[-1]
+    return type_canonicalisation_dict[dtype_str]
+
+
+BITWIDTH_DICT: Dict[str, int] = {
+    **{f"u{n}": n
+       for n in (1, 8, 16, 32, 64)},
+    **{f"i{n}": n
+       for n in (1, 8, 16, 32, 64)},
+    **{f"fp{n}": n
+       for n in (16, 32, 64)},
+    **{f"fp8{suffix}": 8
+       for suffix in ("e4nv", "e4b15", "e4b8", "e5", "e5b16")},
+    "bf16": 16,
+    "void": 0,
+}
+
+for k, v in type_canonicalisation_dict.items():
+    BITWIDTH_DICT[k] = BITWIDTH_DICT[v]
+
+
+def get_primitive_bitwidth(dtype: str) -> int:
+    return BITWIDTH_DICT[dtype]

triton/backends/__init__.py

@@ -1,20 +1,22 @@
-import os
-import importlib.util
+import importlib
 import inspect
+import sys
 from dataclasses import dataclass
+from typing import Type, TypeVar, Union
+from types import ModuleType
 from .driver import DriverBase
 from .compiler import BaseBackend
 
+if sys.version_info >= (3, 10):
+    from importlib.metadata import entry_points
+else:
+    from importlib_metadata import entry_points
 
-def _load_module(name, path):
-    spec = importlib.util.spec_from_file_location(name, path)
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module
+T = TypeVar("T", bound=Union[BaseBackend, DriverBase])
 
 
-def _find_concrete_subclasses(module, base_class):
-    ret = []
+def _find_concrete_subclasses(module: ModuleType, base_class: Type[T]) -> Type[T]:
+    ret: list[Type[T]] = []
     for attr_name in dir(module):
         attr = getattr(module, attr_name)
         if isinstance(attr, type) and issubclass(attr, base_class) and not inspect.isabstract(attr):
@@ -28,23 +30,18 @@ def _find_concrete_subclasses(module, base_class):
 
 @dataclass(frozen=True)
 class Backend:
-    compiler: BaseBackend = None
-    driver: DriverBase = None
+    compiler: Type[BaseBackend]
+    driver: Type[DriverBase]
 
 
-def _discover_backends():
+def _discover_backends() -> dict[str, Backend]:
     backends = dict()
-    root = os.path.dirname(__file__)
-    for name in os.listdir(root):
-        if not os.path.isdir(os.path.join(root, name)):
-            continue
-        if name.startswith('__'):
-            continue
-        compiler = _load_module(name, os.path.join(root, name, 'compiler.py'))
-        driver = _load_module(name, os.path.join(root, name, 'driver.py'))
-        backends[name] = Backend(_find_concrete_subclasses(compiler, BaseBackend),
-                                 _find_concrete_subclasses(driver, DriverBase))
+    for ep in entry_points().select(group="triton.backends"):
+        compiler = importlib.import_module(f"{ep.value}.compiler")
+        driver = importlib.import_module(f"{ep.value}.driver")
+        backends[ep.name] = Backend(_find_concrete_subclasses(compiler, BaseBackend),  # type: ignore
+                                    _find_concrete_subclasses(driver, DriverBase))  # type: ignore
     return backends
 
 
-backends = _discover_backends()
+backends: dict[str, Backend] = _discover_backends()