triton-windows 3.3.0.post19__cp312-cp312-win_amd64.whl → 3.4.0.post20__cp312-cp312-win_amd64.whl

triton/experimental/gluon/language/nvidia/hopper/tma.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+from typing import List, Tuple, TYPE_CHECKING
+from dataclasses import dataclass
+import triton.experimental.gluon.language._core as ttgl
+from triton.experimental.gluon.language._layouts import NVMMASharedLayout
+from triton.experimental.gluon.language._core import builtin, _unwrap_if_constexpr
+
+if TYPE_CHECKING:
+    from triton._C import ir
+
+__all__ = ["async_copy_global_to_shared", "async_copy_shared_to_global", "store_wait"]
+
+
+@dataclass(eq=True)
+class tensor_descriptor_type:
+    block_type: ttgl.block_type
+    shape_type: ttgl.tuple_type
+    strides_type: ttgl.tuple_type
+    layout: NVMMASharedLayout
+
+    def __str__(self) -> str:
+        return f"tensor_descriptor<{self.block_type}, {self.layout}>"
+
+    def _unflatten_ir(self, handles: List[ir.value], cursor: int) -> Tuple[tensor_descriptor, int]:
+        handle = handles[cursor]
+        cursor += 1
+        shape, cursor = self.shape_type._unflatten_ir(handles, cursor)
+        strides, cursor = self.strides_type._unflatten_ir(handles, cursor)
+        value = tensor_descriptor(handle, shape, strides, self.block_type, layout=self.layout)
+        return value, cursor
+
+    def _flatten_ir_types(self, builder: ir.builder, out: List[ir.type]) -> None:
+        is_signed = self.block_type.element_ty.is_int_signed()
+        ty = builder.get_tensor_descriptor_layout_type(
+            self.block_type.to_ir(builder),
+            is_signed,
+            self.layout._to_ir(builder),
+        )
+        out.append(ty)
+        self.shape_type._flatten_ir_types(builder, out)
+        self.strides_type._flatten_ir_types(builder, out)
+
+    def mangle(self) -> str:
+        return f"TD{self.block_type.mangle}_{self.layout.mangle()}TD"
+
+
+class tensor_descriptor:
+
+    def __init__(self, handle, shape: List[ttgl.tensor], strides: List[ttgl.tensor], block_type: ttgl.block_type,
+                 layout: NVMMASharedLayout):
+        self.handle = handle
+        self.shape = ttgl.tuple(shape)
+        self.strides = ttgl.tuple(strides)
+        self.type = tensor_descriptor_type(block_type, shape_type=self.shape.type, strides_type=self.strides.type,
+                                           layout=layout)
+
+    def _flatten_ir(self, handles: List[ir.value]) -> None:
+        handles.append(self.handle)
+        self.shape._flatten_ir(handles)
+        self.strides._flatten_ir(handles)
+
+    @property
+    def block_type(self):
+        return self.type.block_type
+
+    @property
+    def block_shape(self):
+        return self.type.block_type.shape
+
+    @property
+    def dtype(self):
+        return self.type.block_type.element_ty
+
+    @property
+    def layout(self):
+        return self.type.layout
+
+
+@builtin
+def async_copy_global_to_shared(tensor_desc, coord, barrier, result, pred=True, _semantic=None):
+    coord = _semantic._convert_to_ir_values(coord, require_i64=False)
+    pred = _semantic.to_tensor(pred)
+    _semantic.builder.create_async_tma_copy_global_to_local(tensor_desc.handle, coord, barrier.handle, result.handle,
+                                                            pred.handle)
+
+
+@builtin
+def async_copy_shared_to_global(tensor_desc, coord, src, _semantic=None):
+    coord = _semantic._convert_to_ir_values(coord, require_i64=False)
+    _semantic.builder.create_async_tma_copy_local_to_global(tensor_desc.handle, coord, src.handle)
+
+
+@builtin
+def store_wait(pendings, _semantic=None):
+    pendings = _unwrap_if_constexpr(pendings)
+    _semantic.builder.create_async_tma_store_wait(pendings)

triton/experimental/gluon/nvidia/__init__.py

@@ -0,0 +1,4 @@
+from . import hopper
+from . import blackwell
+
+__all__ = ["hopper", "blackwell"]

triton/experimental/gluon/nvidia/blackwell.py

@@ -0,0 +1,3 @@
+from .hopper import TensorDescriptor
+
+__all__ = ["TensorDescriptor"]

triton/experimental/gluon/nvidia/hopper.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from typing import List, Any
+from triton._utils import validate_block_shape, canonicalize_dtype, get_primitive_bitwidth
+from triton.experimental.gluon.language._layouts import NVMMASharedLayout
+
+__all__ = ["TensorDescriptor"]
+
+
+@dataclass
+class TensorDescriptor:
+    base: Any
+    shape: List[int]
+    strides: List[int]
+    block_shape: List[int]
+    layout: NVMMASharedLayout
+
+    def __post_init__(self):
+        rank = len(self.shape)
+        assert len(self.strides) == rank, f"rank mismatch: {self}"
+        assert len(self.block_shape) == rank, f"rank mismatch: {self}"
+        assert rank > 0, "rank must not be zero"
+        assert rank <= 5, "rank cannot be more than 5"
+        assert self.base.data_ptr() % 16 == 0, "base must be 16-byte aligned"
+        validate_block_shape(self.block_shape)
+        dtype_str = canonicalize_dtype(self.base.dtype)
+        elem_bytes = get_primitive_bitwidth(dtype_str) // 8
+        for stride in self.strides[:-1]:
+            assert (stride * elem_bytes) % 16 == 0, "strides must be 16-byte aligned"
+        assert self.strides[-1] == 1, "Last dimension must be contiguous"
+        assert isinstance(self.layout, NVMMASharedLayout), "Layout must be NVMMASharedLayout"
+
+    @staticmethod
+    def from_tensor(tensor: Any, block_shape: List[int], layout: NVMMASharedLayout):
+        return TensorDescriptor(
+            tensor,
+            tensor.shape,
+            tensor.stride(),
+            block_shape,
+            layout,
+        )

triton/knobs.py

@@ -0,0 +1,481 @@
+from __future__ import annotations
+
+import importlib
+import os
+import re
+import subprocess
+import sysconfig
+
+from dataclasses import dataclass
+from contextlib import contextmanager
+from typing import cast, Any, Callable, Generator, Generic, Optional, Protocol, Type, TypeVar, TypedDict, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    from .runtime.cache import CacheManager, RemoteCacheBackend
+    from .runtime.jit import JitFunctionInfo, KernelParam
+    from .compiler.compiler import ASTSource, LazyDict, IRSource
+
+
+class Env:
+    pass
+
+
+env = Env()
+
+propagate_env: bool = True
+
+
+def getenv(key: str) -> Optional[str]:
+    res = os.getenv(key)
+    return res.strip() if res is not None else res
+
+
+def setenv(key: str, value: Optional[str]) -> None:
+    if not propagate_env:
+        return
+
+    if value is not None:
+        os.environ[key] = value
+    elif key in os.environ:
+        del os.environ[key]
+
+
+def toenv(val: Any) -> Union[None, tuple[Optional[str]]]:
+    if val is None:
+        return (None, )
+
+    t = type(val)
+    if t is bool:
+        return ("1" if val else "0", )
+
+    if t is str:
+        return (val, )
+
+    if t is int:
+        return (str(val), )
+
+    return None
+
+
+# There's an asymmetry here so that e.g. env_nvidia_tool can be specified with a
+# a string but return an NvidiaTool.
+SetType = TypeVar("SetType")
+GetType = TypeVar("GetType")
+
+
+class env_base(Generic[SetType, GetType]):
+
+    def __init__(self, key: str, default: Union[SetType, Callable[[], SetType]]) -> None:
+        self.key = key
+        self.default: Callable[[], SetType] = default if callable(default) else lambda: default
+
+    def __set_name__(self, objclass: Type[object], name: str) -> None:
+        self.name = name
+
+    def __get__(self, obj: Optional[object], objclass: Optional[Type[object]]) -> GetType:
+        if obj is None:
+            raise AttributeError(f"Cannot access {type(self)} on non-instance")
+
+        if self.name in obj.__dict__:
+            return self.transform(obj.__dict__[self.name])
+        else:
+            return self.get()
+
+    @property
+    def env_val(self) -> str | None:
+        return getenv(self.key)
+
+    def get(self) -> GetType:
+        env = self.env_val
+        return self.transform(self.default() if env is None else self.from_env(env))
+
+    def __set__(self, obj: object, value: Union[SetType, Env]) -> None:
+        if isinstance(value, Env):
+            obj.__dict__.pop(self.name, None)
+        else:
+            obj.__dict__[self.name] = value
+            if env_val := toenv(value):
+                setenv(self.key, env_val[0])
+
+    def __delete__(self, obj: object) -> None:
+        obj.__dict__.pop(self.name, None)
+
+    def transform(self, val: SetType) -> GetType:
+        # See comment about GetType/SetType in their definition above. Only needed
+        # if GetType != SetType.
+        return cast(GetType, val)
+
+    def from_env(self, val: str) -> SetType:
+        raise NotImplementedError()
+
+
+class env_str(env_base[str, str]):
+
+    def from_env(self, val: str) -> str:
+        return val
+
+
+class env_bool(env_base[bool, bool]):
+
+    def __init__(self, key: str, default: Union[bool, Callable[[], bool]] = False) -> None:
+        super().__init__(key, default)
+
+    def from_env(self, val: str) -> bool:
+        return val.lower() in ("1", "true", "yes", "on", "y")
+
+
+class env_int(env_base[int, int]):
+
+    def __init__(self, key: str, default: Union[int, Callable[[], int]] = 0) -> None:
+        super().__init__(key, default)
+
+    def from_env(self, val: str) -> int:
+        try:
+            return int(val)
+        except ValueError as exc:
+            raise RuntimeError(f"Unable to use {self.key}={val}: expected int") from exc
+
+
+class env_opt_base(Generic[GetType, SetType], env_base[Optional[GetType], Optional[SetType]]):
+
+    def __init__(self, key: str) -> None:
+        super().__init__(key, None)
+
+
+ClassType = TypeVar("ClassType")
+
+
+class env_class(Generic[ClassType], env_opt_base[Type[ClassType], Type[ClassType]]):
+
+    def __init__(self, key: str, type: str) -> None:
+        super().__init__(key)
+        # We can't pass the type directly to avoid import cycles
+        self.type = type
+
+    def from_env(self, val: str) -> Type[ClassType]:
+        comps = val.split(":", 1)
+        if len(comps) != 2:
+            raise RuntimeError(f"Unable to read {self.key}: '{val}' isn't of the form MODULE:CLASS")
+        cls = getattr(importlib.import_module(comps[0]), comps[1])
+
+        if not any((c.__name__ == self.type for c in cls.mro())):
+            raise RuntimeError(f"Unable to use '{val}' from {self.key}: not of type '{self.type}'")
+
+        return cast(Type[ClassType], cls)
+
+
+@dataclass
+class NvidiaTool:
+    path: str
+    version: str
+
+    @staticmethod
+    def from_path(path: str) -> Optional[NvidiaTool]:
+        try:
+            result = subprocess.check_output([path, "--version"], stderr=subprocess.STDOUT)
+            if result is None:
+                return None
+            version = re.search(r".*release (\d+\.\d+).*", result.decode("utf-8"), flags=re.MULTILINE)
+            if version is None:
+                return None
+            return NvidiaTool(path, version.group(1))
+        except subprocess.CalledProcessError:
+            return None
+
+
+def find_nvidia_tool(binary: str) -> str:
+    path = os.path.join(
+        os.path.dirname(__file__),
+        "backends",
+        "nvidia",
+        "bin",
+        binary,
+    )
+    if os.access(path, os.X_OK):
+        return path
+
+    if os.name == "nt":
+        from triton.windows_utils import find_cuda
+        cuda_bin_path, _, _ = find_cuda()
+        if cuda_bin_path:
+            path = os.path.join(cuda_bin_path, binary)
+            if os.access(path, os.X_OK):
+                return path
+
+    return ""
+
+
+class env_nvidia_tool(env_base[str, NvidiaTool]):
+
+    def __init__(self, binary: str) -> None:
+        binary += sysconfig.get_config_var("EXE")
+        self.binary = binary
+        super().__init__(f"TRITON_{binary.upper()}_PATH", lambda: find_nvidia_tool(self.binary))
+
+    def transform(self, path: str) -> NvidiaTool:
+        paths = [
+            path,
+            # We still add default as fallback in case the pointed binary isn't
+            # accessible.
+            self.default(),
+        ]
+        for path in paths:
+            if not path or not os.access(path, os.X_OK):
+                continue
+            if tool := NvidiaTool.from_path(path):
+                return tool
+
+        raise RuntimeError(f"Cannot find {self.binary}")
+
+    def from_env(self, val: str) -> str:
+        return val
+
+
+# Separate classes so that types are correct
+class env_opt_str(env_opt_base[str, str], env_str):
+    pass
+
+
+class env_opt_bool(env_opt_base[bool, bool], env_bool):
+    pass
+
+
+@dataclass(frozen=True)
+class CompileTimes:
+    """
+    Model holding timing information for an invocation of the compiler.
+
+    All times in microseconds.
+    """
+
+    # Duration of make_ir
+    ir_initialization: int
+
+    # Ordered mapping from lowering stage to duration spent in that stage.
+    # Keyed by stage extension, e.g. ttir, ttgir
+    lowering_stages: list[tuple[str, int]]
+
+    # Duration of saving artifacts/metadata to cache
+    store_results: int
+
+    @property
+    def total_lowering(self) -> int:
+        return sum((stage[1] for stage in self.lowering_stages))
+
+    @property
+    def total(self) -> int:
+        return self.ir_initialization + self.total_lowering + self.store_results
+
+
+class CompilationListener(Protocol):
+
+    def __call__(self, *, src: Union[ASTSource, IRSource], metadata: dict[str, Any], metadata_group: dict[str, str],
+                 times: CompileTimes, cache_hit: bool) -> None:
+        ...
+
+
+knobs_type = TypeVar("knobs_type", bound='base_knobs')
+
+
+class base_knobs:
+
+    @property
+    def knob_descriptors(self) -> dict[str, env_base]:
+        return {
+            k: v
+            # data descriptors live on the class object
+            for k, v in type(self).__dict__.items()
+            if isinstance(v, env_base)
+        }
+
+    @property
+    def knobs(self) -> dict[str, Any]:
+        return {k: getattr(self, k) for k in self.knob_descriptors.keys()}
+
+    def copy(self: knobs_type) -> knobs_type:
+        res = type(self)()
+        res.__dict__.update(self.__dict__)
+        return res
+
+    def reset(self: knobs_type) -> knobs_type:
+        for knob in self.knob_descriptors.keys():
+            delattr(self, knob)
+        return self
+
+    @contextmanager
+    def scope(self) -> Generator[None, None, None]:
+        try:
+            initial_env = {knob.key: knob.env_val for knob in self.knob_descriptors.values()}
+            orig = dict(self.__dict__)
+            yield
+        finally:
+            self.__dict__.clear()
+            self.__dict__.update(orig)
+
+            for k, v in initial_env.items():
+                if v is not None:
+                    os.environ[k] = v
+                elif k in os.environ:
+                    del os.environ[k]
+
+
+class BuildImpl(Protocol):
+
+    def __call__(self, name: str, src: str, srcdir: str, library_dirs: list[str], include_dirs: list[str],
+                 libraries: list[str], /) -> str:
+        ...
+
+
+class build_knobs(base_knobs):
+    """Configuration controlling how the native compiler is invoked"""
+    cc: env_opt_str = env_opt_str("CC")
+
+    cudacrt_path: env_opt_str = env_opt_str("TRITON_CUDACRT_PATH")
+    cudart_path: env_opt_str = env_opt_str("TRITON_CUDART_PATH")
+
+    impl: Optional[BuildImpl] = None
+
+    @property
+    def backend_dirs(self) -> set[str]:
+        return {path for path in (self.cudacrt_path, self.cudart_path) if path is not None}
+
+
+class redis_knobs(base_knobs):
+    key_format: env_str = env_str("TRITON_REDIS_KEY_FORMAT", "triton:{key}:{filename}")
+    host: env_str = env_str("TRITON_REDIS_HOST", "localhost")
+    port: env_int = env_int("TRITON_REDIS_PORT", 6379)
+
+
+cache: cache_knobs
+
+
+class cache_knobs(base_knobs):
+    home_dir: env_str = env_str("TRITON_HOME", lambda: os.path.expanduser("~/"))
+
+    dump_dir: env_str = env_str("TRITON_DUMP_DIR", lambda: cache.get_triton_dir("dump"))
+    override_dir: env_str = env_str("TRITON_OVERRIDE_DIR", lambda: cache.get_triton_dir("override"))
+    dir: env_str = env_str("TRITON_CACHE_DIR", lambda: cache.get_triton_dir("cache"))
+
+    manager_class: env_class[CacheManager] = env_class("TRITON_CACHE_MANAGER", "CacheManager")
+    remote_manager_class: env_class[RemoteCacheBackend] = env_class("TRITON_REMOTE_CACHE_BACKEND", "RemoteCacheBackend")
+
+    def get_triton_dir(self, dirname: str) -> str:
+        return os.path.join(self.home_dir, ".triton", dirname)
+
+
+class compilation_knobs(base_knobs):
+    override: env_bool = env_bool("TRITON_KERNEL_OVERRIDE")
+    dump_ir: env_bool = env_bool("TRITON_KERNEL_DUMP")
+    store_binary_only: env_bool = env_bool("TRITON_STORE_BINARY_ONLY")
+    always_compile: env_bool = env_bool("TRITON_ALWAYS_COMPILE")
+    # TODO: Use enum to constrain / 'typecheck' the values
+    use_ir_loc: env_opt_str = env_opt_str("USE_IR_LOC")
+    enable_asan: env_bool = env_bool("TRITON_ENABLE_ASAN")
+    disable_line_info: env_bool = env_bool("TRITON_DISABLE_LINE_INFO")
+    front_end_debugging: env_bool = env_bool("TRITON_FRONT_END_DEBUGGING")
+    allow_non_constexpr_globals: env_bool = env_bool("TRITON_ALLOW_NON_CONSTEXPR_GLOBALS")
+    listener: Union[CompilationListener, None] = None
+
+
+class autotuning_knobs(base_knobs):
+    cache: env_bool = env_bool("TRITON_CACHE_AUTOTUNING")
+    print: env_bool = env_bool("TRITON_PRINT_AUTOTUNING")
+
+
+class LaunchHook(Protocol):
+
+    def __call__(self, metadata: LazyDict) -> None:
+        ...
+
+
+# This is of the form [attr_name, attr_val]
+# TODO: Use tuple instead of list for better typing.
+KernelAttr = list[Union[str, int]]
+
+
+class JITHookCompileInfo(TypedDict):
+    key: str
+    signature: dict[KernelParam, str]
+    device: int
+    constants: None
+    num_warps: int
+    num_ctas: int
+    num_stages: int
+    enable_fp_fusion: bool
+    launch_cooperative_grid: bool
+    extern_libs: tuple[tuple[str, str], ...]
+    configs: list[dict[tuple[int, ...], list[KernelAttr]]]
+    specialization_data: str
+    is_warmup: bool
+
+
+class JITHook(Protocol):
+
+    def __call__(self, *, key: str, repr: str, fn: JitFunctionInfo, compile: JITHookCompileInfo, is_manual_warmup: bool,
+                 already_compiled: bool) -> Optional[bool]:
+        ...
+
+
+class runtime_knobs(base_knobs):
+    interpret: env_bool = env_bool("TRITON_INTERPRET")
+    debug: env_bool = env_bool("TRITON_DEBUG")
+    override_arch: env_opt_str = env_opt_str("TRITON_OVERRIDE_ARCH")
+
+    launch_enter_hook: Optional[LaunchHook] = None
+    launch_exit_hook: Optional[LaunchHook] = None
+
+    # Hook for inspecting compiled functions and modules
+    jit_cache_hook: Optional[JITHook] = None
+    # Hook to signal that a kernel is done compiling and inspect compiled function.
+    # jit_cache_hook will always be called before compilation and jit_post_compile_hook after.
+    jit_post_compile_hook: Optional[JITHook] = None
+
+
+class language_knobs(base_knobs):
+    fp32_default: env_opt_str = env_opt_str("TRITON_F32_DEFAULT")
+    default_fp_fusion: env_bool = env_bool("TRITON_DEFAULT_FP_FUSION", True)
+
+
+class nvidia_knobs(base_knobs):
+    cuobjdump: env_nvidia_tool = env_nvidia_tool("cuobjdump")
+    nvdisasm: env_nvidia_tool = env_nvidia_tool("nvdisasm")
+    ptxas: env_nvidia_tool = env_nvidia_tool("ptxas")
+
+    dump_nvptx: env_bool = env_bool("NVPTX_ENABLE_DUMP")
+    disable_ptxas_opt: env_bool = env_bool("DISABLE_PTXAS_OPT")
+    mock_ptx_version: env_opt_str = env_opt_str("TRITON_MOCK_PTX_VERSION")
+
+    libdevice_path: env_opt_str = env_opt_str("TRITON_LIBDEVICE_PATH")
+    libcuda_path: env_opt_str = env_opt_str("TRITON_LIBCUDA_PATH")
+
+
+class amd_knobs(base_knobs):
+    use_buffer_ops: env_bool = env_bool("AMDGCN_USE_BUFFER_OPS")
+    dump_amdgcn: env_bool = env_bool("AMDGCN_ENABLE_DUMP")
+    libhip_path: env_opt_str = env_opt_str("TRITON_LIBHIP_PATH")
+    lld_path: env_opt_str = env_opt_str("TRITON_HIP_LLD_PATH")
+
+    # We use strs so that we can have a default value based on other runtime info
+    use_block_pingpong: env_opt_bool = env_opt_bool("TRITON_HIP_USE_BLOCK_PINGPONG")
+    use_in_thread_transpose: env_opt_bool = env_opt_bool("TRITON_HIP_USE_IN_THREAD_TRANSPOSE")
+
+    global_prefetch: env_int = env_int("TRITON_HIP_GLOBAL_PREFETCH")
+    local_prefetch: env_int = env_int("TRITON_HIP_LOCAL_PREFETCH")
+    use_async_copy: env_bool = env_bool("TRITON_HIP_USE_ASYNC_COPY")
+    scalarize_packed_fops: env_bool = env_bool("AMDGCN_SCALARIZE_PACKED_FOPS")
+
+
+class proton_knobs(base_knobs):
+    cupti_dir: env_opt_str = env_opt_str("TRITON_CUPTI_LIB_PATH")
+
+
+build = build_knobs()
+redis = redis_knobs()
+cache = cache_knobs()
+compilation = compilation_knobs()
+autotuning = autotuning_knobs()
+runtime = runtime_knobs()
+language = language_knobs()
+nvidia = nvidia_knobs()
+amd = amd_knobs()
+proton = proton_knobs()