# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from nvshmem.bindings.device.cute import *
from nvshmem.bindings.device.cute import (
    barrier as _nvshmem_barrier,
    barrier_block as _nvshmem_barrier_block,
    barrier_warp as _nvshmem_barrier_warp,
)

import nvshmem.core
import cutlass
from cutlass import cute
from cutlass.base_dsl.ast_helpers import const_expr
from cutlass.base_dsl.typing import cast as cute_cast

__all__ = [
    "sync_block",
    "sync_warp",
    "sync",
    "sync_all",
    "sync_all_block",
    "sync_all_warp",
    "barrier",
    "barrier_block",
    "barrier_warp",
    "barrier_all",
    "barrier_all_block",
    "barrier_all_warp",
    "reduce",
    "reduce_block",
    "reduce_warp",
    "reducescatter",
    "reducescatter_block",
    "reducescatter_warp",
    "fcollect",
    "fcollect_block",
    "fcollect_warp",
    "broadcast",
    "broadcast_block",
    "broadcast_warp",
    "alltoall",
    "alltoall_block",
    "alltoall_warp",
]


@cute.jit
def _resolve_ptr(arg):
    return arg.iterator


@cute.jit
def _size_of(obj):
    return cute.size(obj)


@cute.jit
def _resolve_dtype(dst):
    return dst.element_type


@cute.jit
def _resolve_team(team):
    return cute_cast(team, cutlass.Int32)


@cute.jit
def _resolve_root(root):
    return cute_cast(root, cutlass.Int32)


@cute.jit
def _resolve_nelems_dst(dst):
    return cute_cast(_size_of(dst), cutlass.Uint64)


@cute.jit
def _resolve_nelems_src(src):
    return cute_cast(_size_of(src), cutlass.Uint64)


@cute.jit
def _resolve_nelems_min(dst, src):
    dst_nelems = _size_of(dst)
    src_nelems = _size_of(src)
    return cute_cast(dst_nelems if dst_nelems < src_nelems else src_nelems, cutlass.Uint64)


@cute.jit
def _resolve_nelems_alltoall(src, team):
    src_nelems = cute_cast(_size_of(src), cutlass.Uint64)
    team_size = cute_cast(team_n_pes(team), cutlass.Uint64)
    return cute_cast(src_nelems // team_size, cutlass.Uint64)


# sync variations
{% for scope in sync_config.scope %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
@cute.jit
def sync{{ scope.suffix }}(team):
    """
    Executes a {{ _scope_name }}-level synchronization across all PEs in ``team``.{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function.{% endif %}

    This is a lightweight synchronization point that guarantees all PEs in the team have
    reached it before any PE proceeds.  It does not provide memory-ordering or
    memory-visibility guarantees; use ``barrier`` when a memory fence is also required.
    All PEs in the team must call this function before any PE can proceed past it.

    Args:
        - ``team`` (``int``): NVSHMEM team handle identifying the set of PEs to synchronize.
          Use ``nvshmem.core.Teams.TEAM_WORLD`` to synchronize all PEs.

    Note:
        All PEs in ``team`` must call ``sync{{ scope.suffix }}`` with the same ``team`` argument.
        Use ``sync_all{{ scope.suffix }}`` to synchronize across all PEs without specifying a team.
    """
    team = _resolve_team(team)
    return team_sync{{ scope.suffix }}(team)
{% endfor %}

# sync_all variations
{% for scope in sync_all_config.scope %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
@cute.jit
def sync_all{{ scope.suffix }}():
    """
    Executes a {{ _scope_name }}-level synchronization across all PEs in the NVSHMEM runtime
    (equivalent to ``sync{{ scope.suffix }}(TEAM_WORLD)``).{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function.{% endif %}

    This is a convenience wrapper around ``sync{{ scope.suffix }}`` that automatically uses
    ``TEAM_WORLD`` as the team, covering all PEs participating in the NVSHMEM job.

    Note:
        All PEs must call ``sync_all{{ scope.suffix }}`` before any PE can proceed past it.
    """
    team = _resolve_team(nvshmem.core.Teams.TEAM_WORLD)
    return team_sync{{ scope.suffix }}(team)
{% endfor %}

# barrier variations
{% for scope in barrier_config.scope %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
@cute.jit
def barrier{{ scope.suffix }}(team):
    """
    Executes a {{ _scope_name }}-level barrier across all PEs in ``team``.{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function.{% endif %}

    A barrier combines synchronization with a full memory fence, ensuring that all
    outstanding NVSHMEM memory operations (puts, gets, atomics) issued before the barrier
    are complete and visible before any PE in the team proceeds past the barrier.

    Args:
        - ``team`` (``int``): NVSHMEM team handle identifying the set of PEs to barrier on.
          Use ``nvshmem.core.Teams.TEAM_WORLD`` to barrier across all PEs.

    Note:
        All PEs in ``team`` must call ``barrier{{ scope.suffix }}`` with the same ``team`` argument.
        Use ``barrier_all{{ scope.suffix }}`` to barrier across all PEs without specifying a team.
        ``barrier`` provides stronger ordering guarantees than ``sync``.
    """
    team = _resolve_team(team)
    return _nvshmem_barrier{{ scope.suffix }}(team)
{% endfor %}

# barrier_all variations
{% for scope in barrier_all_config.scope %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
@cute.jit
def barrier_all{{ scope.suffix }}():
    """
    Executes a {{ _scope_name }}-level barrier across all PEs in the NVSHMEM runtime
    (equivalent to ``barrier{{ scope.suffix }}(TEAM_WORLD)``).{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function.{% endif %}

    This is a convenience wrapper around ``barrier{{ scope.suffix }}`` that automatically uses
    ``TEAM_WORLD`` as the team. It combines synchronization with a full memory fence,
    ensuring all outstanding NVSHMEM memory operations are visible before proceeding.

    Note:
        All PEs must call ``barrier_all{{ scope.suffix }}`` before any PE can proceed past it.
        ``barrier_all`` provides stronger ordering guarantees than ``sync_all``.
    """
    team = _resolve_team(nvshmem.core.Teams.TEAM_WORLD)
    return _nvshmem_barrier{{ scope.suffix }}(team)
{% endfor %}

# reduce variations
{% for scope in reduce_config.scope %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
@cute.jit
def reduce{{ scope.suffix }}(team, dst, src, op):
    """
    Performs a {{ _scope_name }}-scoped all-reduce from ``src`` to ``dst`` across all PEs in ``team``.{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function with the same arguments.{% endif %}

    Each PE contributes ``size(dst)`` elements from ``src``, and the result of applying
    the reduction operator ``op`` element-wise across all PEs is written to ``dst`` on
    every PE in the team (all-reduce semantics).

    Args:
        - ``team`` (``int``): NVSHMEM team handle identifying the participating PEs.
        - ``dst``: CuTe tensor view pointing to the symmetric destination array.
          Must be a symmetric (NVSHMEM-allocated) tensor. The number of elements reduced
          is ``size(dst)``.
        - ``src``: CuTe tensor view pointing to the symmetric source array.
          Must be a symmetric (NVSHMEM-allocated) tensor with at least ``size(dst)`` elements.
        - ``op`` (``str``): Reduction operator string. Supported operators for numeric types:
          ``"sum"``, ``"prod"``, ``"min"``, ``"max"``. Additional bitwise operators for
          integral types: ``"and"``, ``"or"``, ``"xor"``.

    Note:
        All PEs in ``team`` must call ``reduce{{ scope.suffix }}`` before any PE can proceed past it.
        The element count is taken from ``dst``. Passing an unsupported op/dtype combination
        raises ``RuntimeError`` at JIT compile time.
    """
    team = _resolve_team(team)
    dst_ptr = _resolve_ptr(dst)
    src_ptr = _resolve_ptr(src)
    nelem = _resolve_nelems_dst(dst)
    dtype = _resolve_dtype(dst)
{% for op_literal in reduce_config.reduction_op_literal %}
    {{ "if" if loop.first else "elif" }} const_expr(op == "{{ op_literal.name }}"):
{% for dtype in reduce_config.all_dtypes %}
        {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
            return {{ dtype.nvshmem_type_name }}_{{ op_literal.name }}_reduce{{ scope.suffix }}(team, dst_ptr, src_ptr, nelem)
{% endfor %}
{% endfor %}
{% for op_literal in reduce_config.bitwise_op_literal %}
    {{ "elif" if reduce_config.reduction_op_literal | length > 0 or not loop.first else "if" }} const_expr(op == "{{ op_literal.name }}"):
{% for dtype in reduce_config.integral_dtypes %}
        {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
            return {{ dtype.nvshmem_type_name }}_{{ op_literal.name }}_reduce{{ scope.suffix }}(team, dst_ptr, src_ptr, nelem)
{% endfor %}
{% endfor %}
    raise RuntimeError(f"Unsupported CuTe reduce op/dtype combination: op={op}, dtype={dtype}")
{% endfor %}

# reducescatter variations
{% for scope in reducescatter_config.scope %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
@cute.jit
def reducescatter{{ scope.suffix }}(team, dst, src, op):
    """
    Performs a {{ _scope_name }}-scoped reduce-scatter from ``src`` to ``dst`` across all PEs in ``team``.{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function with the same arguments.{% endif %}

    In a reduce-scatter, each PE contributes elements from ``src``, and the result of
    applying the reduction operator element-wise across all PEs is divided into equal
    portions, each portion written to ``dst`` on a different PE (scatter semantics).

    Args:
        - ``team`` (``int``): NVSHMEM team handle identifying the participating PEs.
        - ``dst``: CuTe tensor view pointing to the symmetric destination array.
          Must be a symmetric (NVSHMEM-allocated) tensor. The number of output elements
          per PE is ``size(dst)``. The ``src`` array should have at least
          ``size(dst) * team_n_pes(team)`` elements.
        - ``src``: CuTe tensor view pointing to the symmetric source array.
          Must be a symmetric (NVSHMEM-allocated) tensor.
        - ``op`` (``str``): Reduction operator string. Supported operators for numeric types:
          ``"sum"``, ``"prod"``, ``"min"``, ``"max"``. Additional bitwise operators for
          integral types: ``"and"``, ``"or"``, ``"xor"``.

    Note:
        All PEs in ``team`` must call ``reducescatter{{ scope.suffix }}`` before any PE can proceed past it.
        The element count is taken from ``dst``. Passing an unsupported op/dtype combination
        raises ``RuntimeError`` at JIT compile time.
    """
    team = _resolve_team(team)
    dst_ptr = _resolve_ptr(dst)
    src_ptr = _resolve_ptr(src)
    nelem = _resolve_nelems_dst(dst)
    dtype = _resolve_dtype(dst)
{% for op_literal in reducescatter_config.reduction_op_literal %}
    {{ "if" if loop.first else "elif" }} const_expr(op == "{{ op_literal.name }}"):
{% for dtype in reducescatter_config.all_dtypes %}
        {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
            return {{ dtype.nvshmem_type_name }}_{{ op_literal.name }}_reducescatter{{ scope.suffix }}(team, dst_ptr, src_ptr, nelem)
{% endfor %}
{% endfor %}
{% for op_literal in reducescatter_config.bitwise_op_literal %}
    {{ "elif" if reducescatter_config.reduction_op_literal | length > 0 or not loop.first else "if" }} const_expr(op == "{{ op_literal.name }}"):
{% for dtype in reducescatter_config.integral_dtypes %}
        {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
            return {{ dtype.nvshmem_type_name }}_{{ op_literal.name }}_reducescatter{{ scope.suffix }}(team, dst_ptr, src_ptr, nelem)
{% endfor %}
{% endfor %}
    raise RuntimeError(f"Unsupported CuTe reducescatter op/dtype combination: op={op}, dtype={dtype}")
{% endfor %}

# fcollect variations
{% for scope in fcollect_config.scope %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
@cute.jit
def fcollect{{ scope.suffix }}(team, dst, src):
    """
    Performs a {{ _scope_name }}-scoped fcollect (all-gather) from ``src`` to ``dst`` across all PEs in ``team``.{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function with the same arguments.{% endif %}

    Each PE contributes ``size(src)`` elements, and the concatenated result from all PEs
    is written to ``dst`` on every PE. The ``dst`` array must be large enough to hold
    contributions from all PEs: ``size(dst) >= size(src) * team_n_pes(team)``.

    Args:
        - ``team`` (``int``): NVSHMEM team handle identifying the participating PEs.
        - ``dst``: CuTe tensor view pointing to the symmetric destination array on this PE.
          Must be a symmetric (NVSHMEM-allocated) tensor with at least
          ``size(src) * team_n_pes(team)`` elements.
        - ``src``: CuTe tensor view pointing to the symmetric source array on this PE.
          Must be a symmetric (NVSHMEM-allocated) tensor. The number of elements
          contributed per PE is ``size(src)``.

    Note:
        All PEs in ``team`` must call ``fcollect{{ scope.suffix }}`` before any PE can proceed past it.
        The element count per PE is taken from ``src``.
        Passing an unsupported dtype raises ``RuntimeError`` at JIT compile time.
    """
    team = _resolve_team(team)
    dst_ptr = _resolve_ptr(dst)
    src_ptr = _resolve_ptr(src)
    nelem = _resolve_nelems_src(src)
    dtype = _resolve_dtype(dst)
{% for dtype in fcollect_config.all_dtypes %}
    {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
        return {{ dtype.nvshmem_type_name }}_fcollect{{ scope.suffix }}(team, dst_ptr, src_ptr, nelem)
{% endfor %}
    raise RuntimeError(f"Unsupported CuTe dtype for fcollect: {dtype}")
{% endfor %}

# broadcast variations
{% for scope in broadcast_config.scope %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
@cute.jit
def broadcast{{ scope.suffix }}(team, dst, src, root=0):
    """
    Performs a {{ _scope_name }}-scoped broadcast from ``src`` on ``root`` PE to ``dst`` on all PEs in ``team``.{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function with the same arguments.{% endif %}

    The root PE broadcasts the contents of its ``src`` array to the ``dst`` array on
    every PE in the team (including the root itself). Non-root PEs ignore their own ``src``.

    Args:
        - ``team`` (``int``): NVSHMEM team handle identifying the participating PEs.
        - ``dst``: CuTe tensor view pointing to the symmetric destination array.
          Must be a symmetric (NVSHMEM-allocated) tensor on all PEs.
        - ``src``: CuTe tensor view pointing to the symmetric source array on the root PE.
          Must be a symmetric (NVSHMEM-allocated) tensor. Only the root PE's ``src`` is used.
        - ``root`` (``int``, optional): PE rank within ``team`` that serves as the broadcast source.
          Defaults to ``0`` (the first PE in the team).

    Note:
        All PEs in ``team`` must call ``broadcast{{ scope.suffix }}`` before any PE can proceed past it.
        The number of elements broadcast is ``min(size(dst), size(src))``.
        Passing an unsupported dtype raises ``RuntimeError`` at JIT compile time.
    """
    team = _resolve_team(team)
    dst_ptr = _resolve_ptr(dst)
    src_ptr = _resolve_ptr(src)
    nelem = _resolve_nelems_min(dst, src)
    root = _resolve_root(root)
    dtype = _resolve_dtype(dst)
{% for dtype in broadcast_config.all_dtypes %}
    {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
        return {{ dtype.nvshmem_type_name }}_broadcast{{ scope.suffix }}(team, dst_ptr, src_ptr, nelem, root)
{% endfor %}
    raise RuntimeError(f"Unsupported CuTe dtype for broadcast: {dtype}")
{% endfor %}

# alltoall variations
{% for scope in alltoall_config.scope %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
@cute.jit
def alltoall{{ scope.suffix }}(team, dst, src):
    """
    Performs a {{ _scope_name }}-scoped all-to-all exchange from ``src`` to ``dst`` across all PEs in ``team``.{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function with the same arguments.{% endif %}

    In an all-to-all operation, each PE sends a distinct portion of its ``src`` array to
    every other PE, and receives a portion from every PE into its ``dst`` array.
    The ``src`` array is logically divided into ``team_n_pes(team)`` equal segments;
    segment ``i`` is sent to PE ``i``. Each PE receives one segment from every PE into
    the corresponding portion of ``dst``.

    Args:
        - ``team`` (``int``): NVSHMEM team handle identifying the participating PEs.
        - ``dst``: CuTe tensor view pointing to the symmetric destination array on this PE.
          Must be a symmetric (NVSHMEM-allocated) tensor. Must have at least
          ``size(src)`` elements (same total size as ``src``).
        - ``src``: CuTe tensor view pointing to the symmetric source array on this PE.
          Must be a symmetric (NVSHMEM-allocated) tensor. The per-PE send count is
          ``size(src) // team_n_pes(team)`` elements.

    Note:
        All PEs in ``team`` must call ``alltoall{{ scope.suffix }}`` before any PE can proceed past it.
        ``size(src)`` must be evenly divisible by ``team_n_pes(team)``.
        The per-PE element count passed to NVSHMEM is ``size(src) // team_n_pes(team)``.
        Passing an unsupported dtype raises ``RuntimeError`` at JIT compile time.
    """
    team = _resolve_team(team)
    dst_ptr = _resolve_ptr(dst)
    src_ptr = _resolve_ptr(src)
    nelem = _resolve_nelems_alltoall(src, team)
    dtype = _resolve_dtype(dst)
{% for dtype in alltoall_config.all_dtypes %}
    {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
        return {{ dtype.nvshmem_type_name }}_alltoall{{ scope.suffix }}(team, dst_ptr, src_ptr, nelem)
{% endfor %}
    raise RuntimeError(f"Unsupported CuTe dtype for alltoall: {dtype}")
{% endfor %}
