triton-windows 3.2.0.post11__cp312-cp312-win_amd64.whl

triton/language/extra/__init__.py

@@ -0,0 +1,26 @@
+import pkgutil
+from importlib.util import module_from_spec
+from sys import modules
+
+_backends = []
+for module_finder, module_name, is_pkg in pkgutil.iter_modules(
+        __path__,
+        prefix=__name__ + ".",
+):
+    # skip .py files (like libdevice.py)
+    if not is_pkg:
+        continue
+
+    # import backends (like cuda and hip) that are included during setup.py
+    spec = module_finder.find_spec(module_name)
+    if spec is None or spec.loader is None:
+        continue
+    module = module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    _backends.append(module_name)
+    modules[module_name] = module
+
+__all__ = _backends
+
+del _backends

triton/language/extra/cuda/__init__.py

@@ -0,0 +1,13 @@
+from . import libdevice
+
+from .utils import (globaltimer, num_threads, num_warps, smid, convert_custom_float8_sm70, convert_custom_float8_sm80)
+
+from ._experimental_tma import *  # noqa: F403
+from ._experimental_tma import __all__ as _tma_all
+
+__all__ = [
+    "libdevice", "globaltimer", "num_threads", "num_warps", "smid", "convert_custom_float8_sm70",
+    "convert_custom_float8_sm80", *_tma_all
+]
+
+del _tma_all

triton/language/extra/cuda/_experimental_tma.py

@@ -0,0 +1,108 @@
+from typing import Sequence
+
+from triton.language import core
+from triton.language import semantic
+from triton._C.libtriton import ir
+
+__all__ = [
+    "experimental_device_tensormap_create1d",
+    "experimental_device_tensormap_create2d",
+    "experimental_tensormap_fenceproxy_acquire",
+]
+
+
+def _determine_elem_type(element_ty: core.dtype):
+    if element_ty.primitive_bitwidth == 8:
+        return 0
+    elif element_ty.primitive_bitwidth == 16:
+        return 1
+    elif element_ty.primitive_bitwidth == 32:
+        return 2
+    else:
+        raise ValueError("element_ty must be a primitive of size 1, 2, or 4 bytes but got")
+
+
+@core.builtin
+def experimental_device_tensormap_create1d(
+    desc_ptr: core.tensor,
+    global_address: core.tensor,
+    load_size: core.tensor,
+    global_size: core.tensor,
+    element_ty: core.dtype,
+    _builder: ir.builder,
+):
+    load_size = core._constexpr_to_value(load_size)
+    global_size = semantic.to_tensor(global_size, _builder)
+    element_ty = core._constexpr_to_value(element_ty)
+    element_stride = [core.full([], 1, core.int32, _builder=_builder)]
+
+    semantic.tensormap_create(
+        desc_ptr=desc_ptr,
+        global_address=global_address,
+        box_dim=[semantic.to_tensor(load_size, _builder)],
+        global_dim=[global_size],
+        global_stride=[],
+        element_stride=element_stride,
+        elem_type=_determine_elem_type(element_ty),
+        interleave_layout=0,
+        swizzle_mode=0,
+        fill_mode=0,
+        builder=_builder,
+    )
+
+
+@core.builtin
+def experimental_device_tensormap_create2d(
+    desc_ptr: core.tensor,
+    global_address: core.tensor,
+    load_size: Sequence[core.constexpr],
+    global_size: Sequence[core.tensor],
+    element_ty: core.dtype,
+    _builder: ir.builder,
+):
+    assert len(load_size) == 2
+    assert len(global_size) == 2
+    load_size = [core._constexpr_to_value(x) for x in load_size]
+    global_size = [semantic.to_tensor(x, _builder) for x in global_size]
+
+    element_size = element_ty.primitive_bitwidth // 8
+    element_size_t = core.full([], element_size, core.int64, _builder=_builder)
+    global_stride = semantic.mul(element_size_t, global_size[-1], True, _builder)
+    # Undocumented, but global_stride seems to be divided by 16
+    global_stride = semantic.ashr(global_stride, semantic.to_tensor(4, _builder), _builder)
+
+    contig_dim_size_in_bytes = element_size * load_size[-1]
+    if contig_dim_size_in_bytes > 128:
+        load_size[-1] = 128 // element_size
+
+    elem_stride = core.full([], 1, core.int32, _builder=_builder)
+
+    semantic.tensormap_create(
+        desc_ptr=desc_ptr,
+        global_address=global_address,
+        box_dim=[semantic.to_tensor(x, _builder) for x in load_size[::-1]],
+        global_dim=global_size[::-1],
+        global_stride=[global_stride],
+        element_stride=[elem_stride, elem_stride],
+        elem_type=_determine_elem_type(element_ty),
+        interleave_layout=0,
+        swizzle_mode=_determine_swizzle_mode_2d(contig_dim_size_in_bytes, load_size),
+        fill_mode=0,
+        builder=_builder,
+    )
+
+
+def _determine_swizzle_mode_2d(contig_dim_size_in_bytes, load_size):
+    if contig_dim_size_in_bytes >= 128:
+        return 3
+    elif contig_dim_size_in_bytes >= 64:
+        return 2
+    elif contig_dim_size_in_bytes >= 32:
+        return 1
+    else:
+        raise ValueError("block size too small")
+
+
+@core.builtin
+def experimental_tensormap_fenceproxy_acquire(desc_ptr: core.tensor, _builder: ir.builder):
+    semantic.tensormap_fenceproxy_acquire(desc_ptr, _builder)