triton-windows 3.3.1.post21__cp311-cp311-win_amd64.whl → 3.4.0.post21__cp311-cp311-win_amd64.whl

triton/__init__.py

@@ -1,5 +1,5 @@
 """isort:skip_file"""
-__version__ = '3.3.1'
+__version__ = '3.4.0'
 
 # ---------------------------------------
 # Note: import order is significant here.
@@ -26,6 +26,8 @@ from . import language
 from . import testing
 from . import tools
 
+must_use_result = language.core.must_use_result
+
 __all__ = [
     "autotune",
     "cdiv",
@@ -39,6 +41,7 @@ __all__ = [
     "KernelInterface",
     "language",
     "MockTensor",
+    "must_use_result",
     "next_power_of_2",
     "OutOfResources",
     "reinterpret",

triton/_filecheck.py

@@ -0,0 +1,87 @@
+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._C.libtriton import ir
+
+# ===-----------------------------------------------------------------------===#
+# filecheck_test
+# ===-----------------------------------------------------------------------===#
+
+# Stub target for testing the frontend.
+stub_target = GPUTarget("cuda", 100, 32)
+stub_backend = make_backend(stub_target)
+
+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],
+                                    stderr=subprocess.STDOUT)
+        except subprocess.CalledProcessError as error:
+            decoded = error.output.decode('unicode_escape')
+            raise ValueError(decoded)
+
+
+def run_parser(kernel_fn):
+    sigkeys = [x.name for x in kernel_fn.params]
+    sigvals = [f"arg{i}" for i in range(len(sigkeys))]
+    signature = {k: v for (k, v) in zip(sigkeys, sigvals)}
+    source_cls = GluonASTSource if kernel_fn.is_gluon() else ASTSource
+    src = source_cls(fn=kernel_fn, signature=signature)
+
+    context = ir.context()
+    ir.load_dialects(context)
+    stub_backend.load_dialects(context)
+
+    extra_options = src.parse_options()
+    options = stub_backend.parse_options(dict(**extra_options))
+    codegen_fns = stub_backend.get_codegen_implementation(options)
+    module_map = stub_backend.get_module_map()
+    module = src.make_ir(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):
+
+    def test_fn():
+        run_filecheck_test(fn)
+
+    return test_fn

triton/_internal_testing.py

@@ -4,7 +4,7 @@ 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
 import pytest
 
 from numpy.random import RandomState
@@ -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():
@@ -46,29 +47,29 @@ def is_hip():
     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_gfx12():
+    target = get_current_target()
+    print(target.arch)
+    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 is_xpu():
@@ -161,7 +162,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 +177,13 @@ 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

triton/_utils.py

@@ -1,35 +1,124 @@
+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):
+    assert len(path) != 0
     prev = iterable if len(path) == 1 else get_iterable_path(iterable, path[:-1])
-    prev[path[-1]] = val
+    prev[path[-1]] = val  # type: ignore[index]
 
 
-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()