..  Licensed to the Apache Software Foundation (ASF) under one
    or more contributor license agreements.  See the NOTICE file
    distributed with this work for additional information
    regarding copyright ownership.  The ASF licenses this file
    to you under the Apache License, Version 2.0 (the
    "License"); you may not use this file except in compliance
    with the License.  You may obtain a copy of the License at

..    http://www.apache.org/licenses/LICENSE-2.0

..  Unless required by applicable law or agreed to in writing,
    software distributed under the License is distributed on an
    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    KIND, either express or implied.  See the License for the
    specific language governing permissions and limitations
    under the License.

copy_async → ldgsts
===================

The ``ldgsts`` variant lowers ``copy_async`` for a **global → shared** transfer to
the PTX ``cp.async`` (LDGSTS) instruction: each thread issues an *asynchronous*
vectorized copy that the hardware completes in the background, so the warp can keep
computing while the load is in flight. It reuses the exact ``[outer, threads, vec]``
partition of the synchronous :doc:`../copy/gmem_smem` variant; the differences are
all in *what* is emitted and *when* it completes. Source:
``python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py``.

What it accepts
---------------

.. code-block:: python

    _LDGSTS_PAIRS = [("global", "shared*")]      # cp.async is unidirectional
    _LDGSTS_VEC_BITS = (128, 64, 32)             # cp_size ∈ {16, 8, 4} bytes

    def _is_ldgsts(op_call, sctx):
        if not sctx.is_target("cuda"):
            return False, "non-cuda target"
        if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
            return False, f"unsupported exec_scope {sctx.scope_kind}"
        for check in (
            lambda: _all_threads_active(sctx),
            lambda: _is_valid_copy(op_call, sctx),
            lambda: _scope_allowed(op_call, sctx, allowed_pairs=_LDGSTS_PAIRS),
            lambda: _divides_thread_cnt_ldgsts(op_call, sctx),
        ):
            ok, msg = check()
            if not ok:
                return False, msg
        return True, None

.. list-table::
   :header-rows: 1
   :widths: 22 78

   * - Property
     - Requirement
   * - target / scope
     - ``cuda``; ``thread`` / ``warp`` / ``warpgroup`` / ``cta``, all active
   * - direction
     - **global → shared only** (``_LDGSTS_PAIRS``) — ``cp.async`` is unidirectional
   * - dtype / shape
     - ``_is_valid_copy``; region element count divides the thread count
   * - vector width
     - ``cp.async`` only accepts cp_size ∈ {4, 8, 16} bytes, so the candidate set is
       restricted to ``_LDGSTS_VEC_BITS`` = ``{128, 64, 32}`` bits
   * - priority
     - ``20`` — selected for ``copy_async`` global→shared (vs the bulk/TMA variants)

Demonstration program
----------------------

A CTA (128 threads) asynchronously loads a ``128×32`` ``float16`` tile global →
shared, then commits and waits before reading it back (from ``test_ldgsts.py``):

.. code-block:: python

    shape, dtype = (128, 32), "float16"
    s_layout = TileLayout(S[shape]); full = (slice(0, 128), slice(0, 32))

    @T.prim_func
    def copy_async(A_ptr: T.handle, B_ptr: T.handle):
        A = T.match_buffer(A_ptr, shape, dtype)
        B = T.match_buffer(B_ptr, shape, dtype)
        T.device_entry(); T.cta_id([1]); T.warp_id([4]); T.lane_id([32]); tid = T.thread_id([128])
        A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
        Tx.cta.copy_async(A_smem[full], A[full], dispatch="ldgsts")   # async global -> shared
        T.ptx.cp.async_.commit_group()                               # caller commits ...
        T.ptx.cp.async_.wait_group(0)                                # ... and waits
        T.cuda.cta_sync()
        Tx.cta.copy(B[full], A_smem[full])

Algorithm
---------

**1. Same partition as** :doc:`../copy/gmem_smem`. ``align_layouts_gs`` builds the
``[outer, threads, vec]`` split with the global side driving the canonical order —
but the vector candidates are clamped to ``{128, 64, 32}`` bits so the byte size is
a legal ``cp.async`` cp_size. For ``128×32 = 4096`` ``float16`` over 128 threads the
widest legal width is ``vec = 8`` (``8 × 2 B = 16 B``), giving ``outer = 4``.

**2. Emit** ``cp.async`` **instead of a synchronous copy**, and do **not** sync:

.. code-block:: python

    for f in range(total_outer):
        s_lin = s_p.apply(f, tid, v0, shape=apply_shape)["m"]
        g_lin = g_p.apply(f, tid, v0, shape=apply_shape)["m"]
        s_ptr = _ptr_off(s_buf.ptr_to(s_zero), _s_off(f, s_lin))
        g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin)
        T.evaluate(T.s_tir.cp_async_raw(s_ptr, g_ptr, cp_size))   # async; cp_size = vec_bits // 8
    # NO cta_sync — commit_group / wait_group / cta_sync are the caller's job

Completion is the caller's responsibility (``cp_async.commit_group()`` then
``cp_async.wait_group()``); the dispatch only issues the in-flight loads.

Generated TIRx IR
-----------------

.. code-block:: python

    for f in range(4):                                       # outer = 4
        s_ptr = pointer_offset(A_smem, ...)
        g_ptr = pointer_offset(A_1, ...)
        T.s_tir.cp_async_raw(s_ptr, g_ptr, 16, T.uint64(0), 0, -1, -1, "")   # cp_size = 16 B

Generated CUDA
--------------

.. code-block:: c++

    // cp.async.cg copies 16 bytes shared <- global, asynchronously
    tvm_builtin_ptx_cp_async_cg_async_cg_shared_global_16(s_ptr, g_ptr);   // x4 (outer = 4)
    // ...
    tvm_builtin_ptx_cp_async_commit_group_async_commit_group();  // asm: cp.async.commit_group;
    tvm_builtin_ptx_cp_async_wait_group_async_wait_group_0();    // asm: cp.async.wait_group 0;

where the helper is ``asm volatile("cp.async.cg.shared.global [%0], [%1], 16;")``.
Each thread issues 4 asynchronous 16-byte copies; nothing blocks until the caller's
``wait_group``.

How inputs change the algorithm
-------------------------------

The dtype/alignment set ``vec`` (hence ``cp_size`` and ``outer``), but unlike the
synchronous variant the width is capped at 16 B (cp.async maximum):

.. list-table::
   :header-rows: 1
   :widths: 24 18 20 38

   * - case
     - ``vec``
     - ``cp_size``
     - ``outer = 4096 / (128 · vec)``
   * - ``float16``, aligned
     - 8
     - 16 B
     - 4
   * - ``float32``, aligned
     - 4
     - 16 B
     - 8
   * - 8-B-aligned only
     - (clamped)
     - 8 B
     - (doubles)
   * - 4-B-aligned only
     - (clamped)
     - 4 B
     - (×4)

If the region can't satisfy even a 4-byte (32-bit) cp_size, ``align_layouts_gs``
finds no candidate and the variant declines. The **direction is fixed**: a
shared → global ``copy_async`` is never ``ldgsts`` (hardware has no store form).
