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

from nvshmem.bindings.device.cute import *

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__ = ["p", "g", "put", "get", "put_nbi", "get_nbi", "put_block", "get_block", "put_nbi_block", "get_nbi_block", "put_warp", "get_warp", "put_nbi_warp", "get_nbi_warp", "put_signal_block", "put_signal", "put_signal_nbi", "put_signal_warp", "put_signal_nbi_block", "put_signal_nbi_warp"]


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


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


@cute.jit
def _resolve_nelems(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_dtype(dst, src=None):
    return dst.element_type


# put variations
{% for scope in put_config.scope %}
{% for blocking_semantics in put_config.blocking_semantics %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
{% set _is_nbi = blocking_semantics.name == "_nbi" %}

@cute.jit
def put{{ blocking_semantics.name }}{{ scope.suffix }}(dst, src, pe):
    """
    {{ "Non-blockingly copies" if _is_nbi else "Copies" }} data from local ``src`` to symmetric ``dst`` on PE ``pe``.
    This is a {{ _scope_name }}-level operation.{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function with the same arguments.{% endif %}

    Args:
        - ``dst``: CuTe tensor view pointing to the symmetric destination on PE ``pe``.
          Must be a symmetric (NVSHMEM-allocated) tensor accessible by all PEs.
        - ``src``: CuTe tensor view pointing to the local source data on this PE.
        - ``pe`` (``int``): Target PE to copy to.

    Note:
        The number of elements transferred is ``min(size(dst), size(src))``.
        ``dst`` and ``src`` must have the same element dtype.
        {% if _is_nbi %}This is a non-blocking operation: the transfer may not be complete when the call returns.
        Use a fence or synchronization primitive to ensure completion.{% else %}This is a blocking operation: local data movement completes before the call returns.{% endif %}
    """
    dst_ptr = _resolve_ptr(dst)
    src_ptr = _resolve_ptr(src)
    nelems = _resolve_nelems(dst, src)
    dtype = _resolve_dtype(dst, src)
{% for dtype in put_config.src_dst %}
    {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
        return {{ dtype.nvshmem_type_name }}_put{{ blocking_semantics.name }}{{ scope.suffix }}(dst_ptr, src_ptr, nelems, pe)
{% endfor %}
    raise RuntimeError(f"Unsupported CuTe dtype for RMA dispatch: {dtype}")
{% endfor %}
{% endfor %}

# get variations
{% for scope in get_config.scope %}
{% for blocking_semantics in get_config.blocking_semantics %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
{% set _is_nbi = blocking_semantics.name == "_nbi" %}

@cute.jit
def get{{ blocking_semantics.name }}{{ scope.suffix }}(dst, src, pe):
    """
    {{ "Non-blockingly copies" if _is_nbi else "Copies" }} data from symmetric ``src`` on PE ``pe`` to local ``dst``.
    This is a {{ _scope_name }}-level operation.{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function with the same arguments.{% endif %}

    Args:
        - ``dst``: CuTe tensor view pointing to the local destination on this PE.
        - ``src``: CuTe tensor view pointing to the symmetric source on PE ``pe``.
          Must be a symmetric (NVSHMEM-allocated) tensor accessible by all PEs.
        - ``pe`` (``int``): Source PE to copy from.

    Note:
        The number of elements transferred is ``min(size(dst), size(src))``.
        ``dst`` and ``src`` must have the same element dtype.
        {% if _is_nbi %}This is a non-blocking operation: the transfer may not be complete when the call returns.
        Use a fence or synchronization primitive to ensure completion.{% else %}This is a blocking operation: data is available in ``dst`` before the call returns.{% endif %}
    """
    dst_ptr = _resolve_ptr(dst)
    src_ptr = _resolve_ptr(src)
    nelems = _resolve_nelems(dst, src)
    dtype = _resolve_dtype(dst, src)
{% for dtype in get_config.src_dst %}
    {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
        return {{ dtype.nvshmem_type_name }}_get{{ blocking_semantics.name }}{{ scope.suffix }}(dst_ptr, src_ptr, nelems, pe)
{% endfor %}
    raise RuntimeError(f"Unsupported CuTe dtype for RMA dispatch: {dtype}")
{% endfor %}
{% endfor %}

# put_signal variations
{% for scope in put_signal_config.scope %}
{% for blocking_semantics in put_signal_config.blocking_semantics %}
{% set _scope_name = "CTA" if scope.suffix == "_block" else ("warp" if scope.suffix == "_warp" else "thread") %}
{% set _is_nbi = blocking_semantics.name == "_nbi" %}

@cute.jit
def put_signal{{ blocking_semantics.name }}{{ scope.suffix }}(dst, src, signal_var, signal_val, signal_op, pe):
    """
    {{ "Non-blockingly puts" if _is_nbi else "Puts" }} data from ``src`` to symmetric ``dst`` on PE ``pe``, then signals ``signal_var``.
    This is a {{ _scope_name }}-level operation.{% if scope.suffix != "" %} All threads in the {{ _scope_name }} must call this function with the same arguments.{% endif %}

    The signal operation atomically updates ``signal_var`` on the remote PE after the data transfer,
    allowing the remote PE to detect transfer completion via a signal variable.

    Args:
        - ``dst``: CuTe tensor view pointing to the symmetric destination on PE ``pe``.
          Must be a symmetric (NVSHMEM-allocated) tensor.
        - ``src``: CuTe tensor view pointing to the local source data on this PE.
        - ``signal_var``: CuTe tensor view pointing to a symmetric signal variable (dtype ``uint64``)
          on PE ``pe``. Must be a 1-element NVSHMEM-allocated tensor.
        - ``signal_val`` (``int``): Value used to update the signal variable. Cast to ``uint64``.
        - ``signal_op``: Signal operation type. Supported values are ``NVSHMEM_SIGNAL_SET``
          (set the signal to ``signal_val``) and ``NVSHMEM_SIGNAL_ADD``
          (atomically add ``signal_val`` to the signal variable).
        - ``pe`` (``int``): Target PE.

    Note:
        The number of data elements transferred is ``min(size(dst), size(src))``.
        ``dst`` and ``src`` must have the same element dtype.
        {% if _is_nbi %}This is a non-blocking operation: neither the data transfer nor the signal update
        are guaranteed to be visible to the remote PE when the call returns.{% else %}This is a blocking operation: the data transfer and signal update complete before the call returns.{% endif %}
    """
    dst_ptr = _resolve_ptr(dst)
    src_ptr = _resolve_ptr(src)
    signal_var_ptr = _resolve_ptr(signal_var)
    nelems = _resolve_nelems(dst, src)
    signal_val = cute_cast(signal_val, cutlass.Uint64)
    dtype = _resolve_dtype(dst, src)
{% for dtype in put_signal_config.src_dst %}
    {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
        return {{ dtype.nvshmem_type_name }}_put_signal{{ blocking_semantics.name }}{{ scope.suffix }}(dst_ptr, src_ptr, nelems, signal_var_ptr, signal_val, signal_op, pe)
{% endfor %}
    raise RuntimeError(f"Unsupported CuTe dtype for RMA dispatch: {dtype}")
{% endfor %}
{% endfor %}

# p variations

@cute.jit
def p(dst, src, pe):
    """
    Writes a single scalar value ``src`` to the symmetric location ``dst`` on PE ``pe``.
    This is a thread-level point operation (scalar put).

    Unlike ``put``, which transfers an array of elements, ``p`` transfers exactly one scalar
    value. ``dst`` must point to a single-element symmetric location.

    Args:
        - ``dst``: CuTe tensor view pointing to a single-element symmetric destination on PE ``pe``.
          Must be a symmetric (NVSHMEM-allocated) tensor with exactly one element.
        - ``src``: Scalar value to write. The value is cast to the dtype of ``dst``.
        - ``pe`` (``int``): Target PE.

    Note:
        This is a blocking, thread-level operation.
        ``src`` is cast to the element dtype of ``dst`` before the transfer.
    """
    dst_ptr = _resolve_ptr(dst)
    dtype = _resolve_dtype(dst)
{% for dtype in p_config.src_dst %}
    {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
        return {{ dtype.nvshmem_type_name }}_p(dst_ptr, cute_cast(src, cutlass.{{ dtype.cute_type_name }}), pe)
{% endfor %}
    raise RuntimeError(f"Unsupported CuTe dtype for RMA dispatch: {dtype}")

# g variations

@cute.jit
def g(src, pe):
    """
    Reads and returns a single scalar value from the symmetric location ``src`` on PE ``pe``.
    This is a thread-level get operation (scalar get).

    Unlike ``get``, which transfers an array of elements, ``g`` retrieves exactly one scalar
    value. ``src`` must point to a single-element symmetric location.

    Args:
        - ``src``: CuTe tensor view pointing to a single-element symmetric source on PE ``pe``.
          Must be a symmetric (NVSHMEM-allocated) tensor with exactly one element.
        - ``pe`` (``int``): Source PE to read from.

    Returns:
        The scalar value stored at ``src`` on PE ``pe``, with the same dtype as ``src``.

    Note:
        This is a blocking, thread-level operation. The returned value is immediately available.
    """
    src_ptr = _resolve_ptr(src)
    pe = cute_cast(pe, cutlass.Int32)
    dtype = _resolve_dtype(src)
{% for dtype in g_config.src_dst %}
    {{ "if" if loop.first else "elif" }} const_expr(dtype == cutlass.{{ dtype.cute_type_name }}):
        return {{ dtype.nvshmem_type_name }}_g(src_ptr, pe)
{% endfor %}
    raise RuntimeError(f"Unsupported CuTe dtype for RMA dispatch: {dtype}")

