triton-windows 3.5.0.post21__cp314-cp314-win_amd64.whl

triton/experimental/gluon/language/amd/cdna4/__init__.py

@@ -0,0 +1,48 @@
+from triton.experimental.gluon.language import _core as ttgl
+from ..._core import builtin, float32
+from ..._layouts import DotOperandLayout
+from .._layouts import AMDMFMALayout
+from ..cdna3 import *  # NOQA: F403
+from ..cdna3 import __all__ as __cdna3_all
+from . import async_copy
+
+__all__ = [*__cdna3_all, "async_copy", "mfma_scaled"]
+
+
+@builtin
+def mfma_scaled(a, a_scale, a_format, b, b_scale, b_format, acc, _semantic=None):
+    """
+    AMD Scaled MFMA operation.
+
+    ```
+    c = a * a_scale @ b * b_scale + acc
+    ```
+
+    `a` and `b` use microscaling formats described in
+    "OCP Microscaling Formats (MX) Specification":
+    https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf.
+    Currently supported only on CDNA4 hardware.
+
+    Args:
+        a (tensor): The operand A to be multiplied.
+        a_scale (tensor): Scale factor for operand A.
+        a_format (str): Format of the operand A. Available formats: `e2m1`, `e4m3`, `e5m2`.
+        b (tensor): The operand B to be multiplied.
+        b_scale (tensor): Scale factor for operand B. Available formats: `e2m1`, `e4m3`, `e5m2`.
+        b_format (str): Format of the operand B.
+        acc (tensor): Accumulator tensor.
+    """
+    layout = acc.type.layout
+    assert isinstance(layout, AMDMFMALayout), "Expected layout to be an instance of AMDMFMALayout"
+    assert (isinstance(a.type.layout, DotOperandLayout) and a.type.layout.parent== layout), \
+            "Expected lhs layout to be a DotOperandLayout with parent matching MFMA layout"
+    assert (isinstance(b.type.layout, DotOperandLayout) and b.type.layout.parent == layout), \
+            "Expected rhs layout to be a DotOperandLayout with parent matching MFMA layout"
+
+    assert a_format.value in {"e2m1", "e4m3", "e5m2"}, f"Unsupported lhs_format: {a_format.value}"
+    assert b_format.value in {"e2m1", "e4m3", "e5m2"}, f"Unsupported rhs_format: {b_format.value}"
+
+    tensor = _semantic.dot_scaled(a, a_scale, a_format, b, b_scale, b_format, acc, False, True, True, float32)
+
+    ret_ty = ttgl.distributed_type(tensor.dtype, tensor.shape, layout)
+    return ttgl.tensor(tensor.handle, ret_ty)

triton/experimental/gluon/language/amd/cdna4/async_copy.py

@@ -0,0 +1,151 @@
+from ..._core import ir, builtin, _unwrap_if_constexpr
+from ..._semantic import _check
+from ..._layouts import BlockedLayout, SliceLayout
+from ..cdna3 import _verify_buffer_ops
+
+__all__ = [
+    "global_load_to_shared",
+    "buffer_load_to_shared",
+    "async_wait",
+    "load_shared_relaxed",
+]
+
+
+@builtin
+def global_load_to_shared(dest, ptr, mask=None, other=None, cache_modifier="", _semantic=None):
+    """
+    AMD global load to shared operation. This operation loads data directly
+    from global memory to shared memory without going through registers. It
+    happens asynchronously and requires a subsequent `async_wait` to ensure the
+    data is available in shared memory.
+    Compared to `buffer_load_to_shared`, it requires a tensor pointer which
+    supports 64-bit indexing range for each thread in a block, which gives more
+    flexibility, but at the cost of higher register pressure and no hardware
+    out-of-bound masking support. Prefer to use `buffer_load_to_shared` when
+    possible for better performance.
+
+    The underlying hardware instruction uses separate registers for global
+    memory address for each thread but the same register for local memory
+    address for the whole warp. Therefore, while using this operation
+    the following conditions must be met or lowering to LLVM will fail:
+
+    - For the `ptr` layout, size per thread * bits per element must be 128 or 32.
+      To get ideal performance, it is recommended to use 128 bits per element.
+    - Writes to `dest` must be coalesced.
+    - If `dest` is swizzled, it only can be swizzled within warp boundary.
+
+    Args:
+        dest (shared_memory_descriptor): Destination shared memory descriptor.
+        ptr (pointer tensor): Tensor of pointers to global memory to load from.
+        mask (tensor, optional): Mask tensor for predicated loads. Defaults to None.
+        other (tensor, optional): Tensor providing default values for masked elements. Defaults to None.
+        cache_modifier (str): Cache modifier specifier. Defaults to "".
+    """
+    _check(ptr.type.is_block(), lambda: "expected ptr to be a tensor")
+    _check(isinstance(ptr.type.layout, (BlockedLayout, SliceLayout)),
+           lambda: "expected ptr type layout to be BlockedLayout or SliceLayout")
+    _check(
+        dest.shape == ptr.shape, lambda:
+        f"expected dest shape to match pointer shape but got dest.shape = {dest.shape}, pointer.shape = {ptr.shape}")
+
+    mask = _unwrap_if_constexpr(mask)
+    if mask is not None:
+        ptr, mask = _semantic.broadcast_impl_value(ptr, mask)
+    other = _unwrap_if_constexpr(other)
+    if other is not None:
+        ptr, other = _semantic.broadcast_impl_value(ptr, other)
+
+    cache_modifier = _semantic._str_to_load_cache_modifier(cache_modifier)
+    mask_handle = mask.handle if mask is not None else ir.value()
+    other_handle = other.handle if other is not None else ir.value()
+    _semantic.builder.create_async_copy_global_to_local(dest.handle, ptr.handle, mask_handle, other_handle,
+                                                        cache_modifier, ir.EVICTION_POLICY.NORMAL, False)
+
+
+@builtin
+def buffer_load_to_shared(dest, ptr, offsets, mask=None, other=None, cache_modifier="", _semantic=None):
+    """
+    AMD buffer load to shared operation. Buffer load is similar to global load
+    but it accesses global memory via a scalar base pointer and a tensor of
+    32-bit offsets instead of a tensor of pointers. This operation loads data
+    directly from global memory to shared memory without going through
+    registers. It happens asynchronously and requires a subsequent `async_wait`
+    to ensure the data is available in shared memory.
+    Compared to `global_load_to_shared`, it has better performance and also
+    supports hardware out-of-bound masking. But it strictly requires a
+    32-bit offset instead of a 64-bit tensor pointer.
+
+    The underlying hardware instruction uses separate registers for global
+    memory address for each thread but the same register for local memory
+    address for the whole warp. Therefore, while using this operation
+    the following conditions must be met or lowering to LLVM will fail:
+
+    - For the `offsets` layout, size per thread * bits per element must be 128 or 32.
+      To get ideal performance, it is recommended to use 128 bits per element.
+    - Writes to `dest` must be coalesced.
+    - If `dest` is swizzled, it only can be swizzled within warp boundary.
+
+    Args:
+        dest (shared_memory_descriptor): Destination shared memory descriptor.
+        ptr (pointer to scalar): Global memory scalar base pointer to load from.
+        offsets (tensor): Offsets tensor for the load operation.
+        mask (tensor, optional): Mask tensor for predicated loads. Defaults to None.
+        other (tensor, optional): Tensor providing default values for masked elements. Defaults to None.
+        cache_modifier (str): Cache modifier specifier. Defaults to "".
+    """
+    _check(isinstance(offsets.type.layout, (BlockedLayout, SliceLayout)),
+           lambda: "expected offsets type layout to be BlockedLayout or SliceLayout")
+    _verify_buffer_ops(ptr, offsets, mask, other)
+
+    mask = _unwrap_if_constexpr(mask)
+    if mask is not None:
+        offsets, mask = _semantic.broadcast_impl_value(offsets, mask)
+    other = _unwrap_if_constexpr(other)
+    if other is not None:
+        offsets, other = _semantic.broadcast_impl_value(offsets, other)
+
+    mask = mask.handle if mask is not None else ir.value()
+    other = other.handle if other is not None else ir.value()
+    stride = ir.value()
+    cache_modifier = _semantic._str_to_load_cache_modifier(cache_modifier)
+
+    _semantic.builder.create_buffer_load_to_local(dest.handle, ptr.handle, offsets.handle, mask, other, stride,
+                                                  cache_modifier)
+
+
+@builtin
+def async_wait(num_outstanding=0, _semantic=None):
+    """
+    Wait for outstanding memory operations, this includes normal load like
+    `load` and `buffer_load`, as well as direct load to shared memory
+    like `global_load_to_shared` and `buffer_load_to_shared`.
+    It will block until the number of outstanding memory operations is less than
+    or equal to `num_outstanding`.
+
+    Args:
+        num_outstanding (int): The number of outstanding operations to wait for. Defaults to 0.
+    """
+    num_outstanding = _unwrap_if_constexpr(num_outstanding)
+    _semantic.builder.create_async_wait_group(num_outstanding)
+
+
+@builtin
+def load_shared_relaxed(smem, layout, _semantic=None):
+    """
+    Load a tensor from shared memory with extra hints for the underlying
+    compiler to avoid emitting unnecessary waits before loading from the target
+    shared memory.
+
+    Args:
+        smem (shared_memory_descriptor): Shared memory descriptor to load from.
+        layout (DistributedLayout): The destination layout of the tensor.
+
+    Returns:
+        tensor: A Gluon tensor containing the loaded data.
+    """
+    SYNCED_VIA_WAIT_ATTR_NAME = "ttg.amdgpu.syncedViaAsyncWait"
+
+    layout = _unwrap_if_constexpr(layout)
+    ret = _semantic.shared_load(smem, layout)
+    ret.handle.set_attr(SYNCED_VIA_WAIT_ATTR_NAME, _semantic.builder.get_bool_attr(True))
+    return ret

triton/experimental/gluon/language/extra/__init__.py

@@ -0,0 +1,3 @@
+from triton.language.extra import libdevice
+
+__all__ = ["libdevice"]

triton/experimental/gluon/language/nvidia/__init__.py

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

triton/experimental/gluon/language/nvidia/ampere/__init__.py

@@ -0,0 +1,3 @@
+from . import async_copy, mbarrier
+
+__all__ = ["async_copy", "mbarrier"]

triton/experimental/gluon/language/nvidia/ampere/async_copy.py

@@ -0,0 +1,74 @@
+from ..._semantic import _check
+from ..._core import _unwrap_if_constexpr, builtin
+from triton._C.libtriton import ir
+
+__all__ = [
+    "async_copy_global_to_shared",
+    "mbarrier_arrive",
+    "commit_group",
+    "wait_group",
+]
+
+
+@builtin
+def async_copy_global_to_shared(smem, pointer, mask=None, cache_modifier="", eviction_policy="", volatile=False,
+                                _semantic=None):
+    """
+    Asynchronously copy elements from global memory to shared memory.
+
+    Args:
+        smem (shared_memory_descriptor): Destination shared memory descriptor.
+        pointer (tensor): Source pointer tensor.
+        mask (tensor, optional): Mask tensor for predicated loads. Defaults to None.
+        cache_modifier (str): Cache modifier specifier. Defaults to "".
+        eviction_policy (str): Eviction policy specifier. Defaults to "".
+        volatile (bool): Whether the load is volatile. Defaults to False.
+    """
+    mask = _unwrap_if_constexpr(mask)
+    cache_modifier = _semantic._str_to_load_cache_modifier(cache_modifier)
+    eviction_policy = _semantic._str_to_eviction_policy(eviction_policy)
+    volatile = _unwrap_if_constexpr(volatile)
+    if mask is not None:
+        pointer, mask = _semantic.broadcast_impl_value(pointer, mask)
+    _check(
+        smem.shape == pointer.shape, lambda:
+        f"expected smem shape to match pointer shape but got smem.shape = {smem.shape}, pointer.shape = {pointer.shape}"
+    )
+    mask_handle = mask.handle if mask is not None else ir.value()
+    _semantic.builder.create_async_copy_global_to_local(smem.handle, pointer.handle, mask_handle, ir.value(),
+                                                        cache_modifier, eviction_policy, volatile)
+
+
+@builtin
+def mbarrier_arrive(mbarrier, increment_count=True, _semantic=None):
+    """
+    Arrive on the mbarrier once all outstanding async copies are complete.
+
+    Args:
+        mbarrier (shared_memory_descriptor): Barrier object to arrive on.
+        increment_count (bool): Whether to increment the arrival count. Defaults to True.
+    """
+    increment_count = _unwrap_if_constexpr(increment_count)
+    _semantic.builder.create_async_copy_mbarrier_arrive(mbarrier.handle, increment_count)
+
+
+@builtin
+def commit_group(_semantic=None):
+    """
+    Commit the current asynchronous copy group.
+
+    This finalizes a set of asynchronous copy operations.
+    """
+    _semantic.builder.create_async_commit_group()
+
+
+@builtin
+def wait_group(num_outstanding=0, _semantic=None):
+    """
+    Wait for outstanding asynchronous copy group operations.
+
+    Args:
+        num_outstanding (int): Wait until `num_outstanding` or less async copy groups in-flight. Defaults to 0.
+    """
+    num_outstanding = _unwrap_if_constexpr(num_outstanding)
+    _semantic.builder.create_async_wait_group(num_outstanding)

triton/experimental/gluon/language/nvidia/ampere/mbarrier.py

@@ -0,0 +1,80 @@
+from triton.experimental.gluon.language._layouts import SwizzledSharedLayout
+from triton.experimental.gluon.language._core import builtin, _unwrap_if_constexpr
+
+__all__ = ["arrive", "init", "invalidate", "MBarrierLayout", "wait"]
+
+
+class MBarrierLayout(SwizzledSharedLayout):
+    """
+    Layout for mbarrier synchronization in Ampere and later architectures.
+
+    Args:
+        ctas_per_cga (int): CTAs per CGA grouping. Defaults to 1.
+        cta_split_num (int): CTA split factor. Defaults to 1.
+    """
+
+    def __init__(self, ctas_per_cga: int = 1, cta_split_num: int = 1):
+        super().__init__(
+            vec=1,
+            per_phase=1,
+            max_phase=1,
+            order=[0],
+            ctas_per_cga=[ctas_per_cga],
+            cta_split_num=[cta_split_num],
+            cta_order=[0],
+        )
+
+
+@builtin
+def init(mbarrier, count, _semantic=None):
+    """
+    Initialize an mbarrier with a specified count.
+
+    Args:
+        mbarrier (shared_memory_descriptor): The barrier object to initialize.
+        count (int): The initial count for the barrier.
+    """
+    count = _unwrap_if_constexpr(count)
+    _semantic.builder.create_mbarrier_init(mbarrier.handle, count)
+
+
+@builtin
+def invalidate(mbarrier, _semantic=None):
+    """
+    Invalidate an mbarrier, resetting its state.
+
+    Args:
+        mbarrier (shared_memory_descriptor): The barrier object to invalidate.
+    """
+    _semantic.builder.create_mbarrier_inval(mbarrier.handle)
+
+
+@builtin
+def wait(mbarrier, phase, pred=True, deps=(), _semantic=None):
+    """
+    Wait until the mbarrier object completes its current phase.
+
+    Args:
+        mbarrier (shared_memory_descriptor): The barrier object to wait on.
+        phase (int): The phase index to wait for.
+        pred (bool): Predicate. Operation is skipped if predicate is False. Defaults to True.
+        deps (Sequence[shared_memory_descriptor]): Dependent allocations barrier is waiting on. Used to track liveness of dependent allocations. Defaults to ().
+    """
+    phase = _semantic.to_tensor(phase)
+    pred = _semantic.to_tensor(pred)
+    deps = [x.handle for x in deps]
+    _semantic.builder.create_mbarrier_wait(mbarrier.handle, phase.handle, pred.handle, deps)
+
+
+@builtin
+def arrive(mbarrier, *, pred=True, _semantic=None):
+    """
+    Arrive on an mbarrier, signaling that a thread has reached the barrier.
+
+    Args:
+        mbarrier (shared_memory_descriptor): The barrier object to arrive on.
+        pred (bool): Predicate. Operation is skipped if predicate is False. Defaults to True.
+    """
+    count = 1
+    pred = _semantic.to_tensor(pred)
+    _semantic.builder.create_mbarrier_arrive(mbarrier.handle, count, pred.handle)