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

triton/language/__init__.py

@@ -6,6 +6,7 @@ from . import extra
 from .standard import (
     argmax,
     argmin,
+    bitonic_merge,
     cdiv,
     cumprod,
     cumsum,
@@ -14,11 +15,13 @@ from .standard import (
     max,
     min,
     ravel,
+    reduce_or,
     sigmoid,
     softmax,
     sort,
     sum,
     swizzle2d,
+    topk,
     xor_sum,
     zeros,
     zeros_like,
@@ -26,16 +29,17 @@ from .standard import (
 from .core import (
     PropagateNan,
     TRITON_MAX_TENSOR_NUMEL,
-    _experimental_descriptor_load,
-    _experimental_descriptor_store,
-    _experimental_make_tensor_descriptor,
-    _experimental_reinterpret_tensor_descriptor,
-    _experimental_tensor_descriptor,
+    load_tensor_descriptor,
+    store_tensor_descriptor,
+    make_tensor_descriptor,
+    tensor_descriptor,
+    tensor_descriptor_type,
     add,
     advance,
     arange,
     associative_scan,
     assume,
+    async_task,
     atomic_add,
     atomic_and,
     atomic_cas,
@@ -51,8 +55,10 @@ from .core import (
     cat,
     cast,
     clamp,
+    condition,
     const,
     constexpr,
+    constexpr_type,
     debug_barrier,
     device_assert,
     device_print,
@@ -80,6 +86,7 @@ from .core import (
     join,
     load,
     make_block_ptr,
+    map_elementwise,
     max_constancy,
     max_contiguous,
     maximum,
@@ -89,7 +96,6 @@ from .core import (
     permute,
     pi32_t,
     pointer_type,
-    nv_tma_desc_type,
     program_id,
     range,
     reduce,
@@ -126,15 +132,15 @@ from .random import (
     randn4x,
     uint_to_uniform_float,
 )
+from . import target_info
 
 __all__ = [
     "PropagateNan",
     "TRITON_MAX_TENSOR_NUMEL",
-    "_experimental_descriptor_load",
-    "_experimental_descriptor_store",
-    "_experimental_make_tensor_descriptor",
-    "_experimental_reinterpret_tensor_descriptor",
-    "_experimental_tensor_descriptor",
+    "load_tensor_descriptor",
+    "store_tensor_descriptor",
+    "make_tensor_descriptor",
+    "tensor_descriptor",
     "abs",
     "add",
     "advance",
@@ -143,6 +149,7 @@ __all__ = [
     "argmin",
     "associative_scan",
     "assume",
+    "async_task",
     "atomic_add",
     "atomic_and",
     "atomic_cas",
@@ -152,6 +159,7 @@ __all__ = [
     "atomic_xchg",
     "atomic_xor",
     "bfloat16",
+    "bitonic_merge",
     "block_type",
     "broadcast",
     "broadcast_to",
@@ -160,8 +168,10 @@ __all__ = [
     "cdiv",
     "ceil",
     "clamp",
+    "condition",
     "const",
     "constexpr",
+    "constexpr_type",
     "cos",
     "cumprod",
     "cumsum",
@@ -204,6 +214,7 @@ __all__ = [
     "log",
     "log2",
     "make_block_ptr",
+    "map_elementwise",
     "math",
     "max",
     "max_constancy",
@@ -219,7 +230,6 @@ __all__ = [
     "philox_impl",
     "pi32_t",
     "pointer_type",
-    "nv_tma_desc_type",
     "program_id",
     "rand",
     "rand4x",
@@ -230,6 +240,7 @@ __all__ = [
     "range",
     "ravel",
     "reduce",
+    "reduce_or",
     "reshape",
     "rsqrt",
     "slice",
@@ -246,7 +257,9 @@ __all__ = [
     "store",
     "sum",
     "swizzle2d",
+    "target_info",
     "tensor",
+    "topk",
     "trans",
     "tuple",
     "uint16",
@@ -264,12 +277,12 @@ __all__ = [
 ]
 
 
-def str_to_ty(name):
+def str_to_ty(name, c):
     from builtins import tuple
 
     if isinstance(name, tuple):
         fields = type(name).__dict__.get("_fields", None)
-        return tuple_type([str_to_ty(x) for x in name], fields)
+        return tuple_type([str_to_ty(x, c) for x in name], fields)
 
     if name[0] == "*":
         name = name[1:]
@@ -277,14 +290,32 @@ def str_to_ty(name):
         if name[0] == "k":
             name = name[1:]
             const = True
-        ty = str_to_ty(name)
+        ty = str_to_ty(name, c)
         return pointer_type(element_ty=ty, const=const)
 
-    if name == "nvTmaDesc":
-        return nv_tma_desc_type()
+    if name.startswith("tensordesc"):
+        inner = name.split("<")[1].rstrip(">")
+        dtype, rest = inner.split("[", maxsplit=1)
+        block_shape, rest = rest.split("]", maxsplit=1)
+        block_shape = [int(s.strip()) for s in block_shape.rstrip("]").split(",")]
+        layout = rest.lstrip(",")
+        is_gluon = len(layout)
+        dtype = str_to_ty(dtype, None)
+        ndim = len(block_shape)
+        shape_type = tuple_type([int32] * ndim)
+        # FIXME: Last dim stride should be constexpr(1)
+        stride_type = tuple_type(([int64] * ndim))
+        block = block_type(dtype, block_shape)
+        if is_gluon:
+            from triton.experimental.gluon.language._layouts import NVMMASharedLayout
+            from triton.experimental.gluon.language.nvidia.hopper.tma import tensor_descriptor_type as gluon_tensor_descriptor_type
+            layout = eval(layout, dict(NVMMASharedLayout=NVMMASharedLayout))
+            assert isinstance(layout, NVMMASharedLayout)
+            return gluon_tensor_descriptor_type(block, shape_type, stride_type, layout)
+        return tensor_descriptor_type(block, shape_type, stride_type)
 
-    if name == "constexpr":
-        return constexpr
+    if name.startswith("constexpr"):
+        return constexpr_type(c)
 
     tys = {
         "fp8e4nv": float8e4nv,