# 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.

"""Non-ldmatrix copy dispatch for register ↔ memory.

This file owns every copy where one side is per-thread local (``R`` =
register). That R side carries the partition: its ``TileLayout`` ``shard``
has thread-axis iters telling us which thread owns which logical coordinate.
The other side (``S``) can be ``shared*`` or ``global`` — the algorithm is
identical either way.

Slice/canonicalize both sides, align via perm+group, then emit a per-thread
vectorized copy loop. Direction-symmetric: covers R2S / S2R / R2G / G2R.
"""

import tvm
from tvm.arith import Analyzer, ConstIntBound
from tvm.runtime import DataType
from tvm.script import tirx as T
from tvm.tirx import Buffer, PrimFunc
from tvm.tirx import Var as _TirVar
from tvm.tirx.expr import IntImm as _IntImm
from tvm.tirx.layout import ComposeLayout, Iter, TileLayout
from tvm.tirx.operator.tile_primitive.registry import DispatchContext
from tvm.tirx.tile_primitive import TilePrimitiveCall

from ..layout_utils import recompose_swizzle, strip_swizzle_to_tile
from ._common import _alignment_ok, copy_ptx_form, copy_ptx_ld_chain
from .utils import _is_valid_copy, _scope_allowed
from .vec_forced import _ld_cache_config


def _extract_tile(layout, region):
    """Strip swizzle off ``layout`` so it can be perm/grouped as a TileLayout."""
    return strip_swizzle_to_tile(layout, lambda: [int(end - start) for (start, end) in region])


_REG_PAIRS = [
    ("local", "shared*"),
    ("shared*", "local"),
    ("local", "global"),
    ("global", "local"),
]
_SCOPE_RANK = {"thread": 0, "warp": 1, "warpgroup": 2, "cta": 3}
_VALID_R_SUBSCOPES = {"thread", "warp", "warpgroup"}


def _all_threads_active(sctx: DispatchContext) -> tuple[bool, str | None]:
    if sctx.scope_kind == "thread":
        return True, None
    required: dict[str, int] = {}
    if sctx.scope_kind in ("warp", "warpgroup", "cta"):
        required["laneid"] = 32
    if sctx.scope_kind == "warpgroup":
        required["wid_in_wg"] = 4
    if sctx.scope_kind == "cta":
        tx_iv = sctx.launch_params.get("threadIdx.x")
        if tx_iv is None:
            return False, "cta scope missing threadIdx.x launch_params"
        try:
            required["warpid"] = int(tx_iv.dom.extent) // 32
        except (TypeError, ValueError):
            return False, f"non-static threadIdx.x extent: {tx_iv.dom.extent}"
    for axis_name, expected in required.items():
        if axis_name not in sctx.intra:
            return False, f"sctx.intra missing {axis_name!r}"
        ext_raw, off_raw = sctx.intra[axis_name]
        try:
            ext, off = int(ext_raw), int(off_raw)
        except (TypeError, ValueError):
            return False, f"non-static range for {axis_name}: ({ext_raw}, {off_raw})"
        if ext != expected or off != 0:
            return False, f"{axis_name} narrowed to [{off}, {off + ext}) vs full [0, {expected})"
    return True, None


def _r_side_layout_valid(
    op_call: TilePrimitiveCall, sctx: DispatchContext
) -> tuple[bool, str | None]:
    op_call = TilePrimitiveCall.downcast(op_call)
    src: Buffer = op_call.src.buffer
    dst: Buffer = op_call.dst.buffer
    r_buf = src if src.scope() == "local" else dst
    layout = r_buf.layout
    if layout is None:
        return False, "R has no layout"
    if layout.is_swizzle():
        return False, "R layout is swizzle"
    if not isinstance(layout, TileLayout):
        return False, f"R layout is {type(layout).__name__}, not TileLayout"

    scope_rank = _SCOPE_RANK[sctx.scope_kind]
    for it in layout.shard:
        ax = it.axis
        if not ax.is_thread():
            continue
        ax_scope = ax.get_scope()
        ax_sub = ax.get_subscope()
        if ax_scope is None or ax_sub is None:
            return False, f"R thread axis {ax.name!r} missing scope/subscope"
        if ax_sub.name not in _VALID_R_SUBSCOPES:
            return False, f"R thread axis {ax.name!r} subscope={ax_sub.name!r} (not register-level)"
        if ax_scope.name not in _SCOPE_RANK or _SCOPE_RANK[ax_scope.name] > scope_rank:
            return (
                False,
                f"R thread axis {ax.name!r} scope={ax_scope.name!r} > exec {sctx.scope_kind!r}",
            )
    r_br = op_call.src if src.scope() == "local" else op_call.dst
    region = [(r.min, r.min + r.extent) for r in r_br.region]
    sliced = layout.slice(list(r_buf.shape), region)
    if sliced is None:
        return False, "R layout slice failed"
    analyzer = Analyzer()
    for axis, off in sliced.offset.items():
        if axis.is_thread() and not analyzer.can_prove_equal(off, 0):
            return False, f"R sliced offset on thread axis {axis.name!r} = {off}"
    return True, None


def _s_side_slice_ok(op_call: TilePrimitiveCall) -> tuple[bool, str | None]:
    """S is the non-local side (shared* or global). Slice must succeed."""
    op_call = TilePrimitiveCall.downcast(op_call)
    src_br = op_call.src
    dst_br = op_call.dst
    s_br = dst_br if src_br.buffer.scope() == "local" else src_br
    s_buf: Buffer = s_br.buffer
    layout = s_buf.layout
    if layout is None:
        return False, "S has no layout"
    region = [(r.min, r.min + r.extent) for r in s_br.region]
    if layout.slice(list(s_buf.shape), region) is None:
        return False, "S layout slice failed"
    return True, None


def _is_reg_copy(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
    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=_REG_PAIRS),
        lambda: _r_side_layout_valid(op_call, sctx),
        lambda: _s_side_slice_ok(op_call),
    ):
        ok, msg = check()
        if not ok:
            return False, msg
    return True, None


def _compute_perm_r(r):
    # thread axes first, then by stride descending
    def key(p):
        it = p[1]
        return (0 if it.axis.is_thread() else 1, -int(it.stride))

    return [i for i, _ in sorted(enumerate(r.shard), key=key)]


def align_layouts_raw(r_sliced, s_sliced, s_region):
    """Returns (r_p, s_p, s_seps, r_perm).

    ``r_p`` for ``_s_thread_offset``; ``r_perm`` for ``_split_thread_loop`` pairing.
    """
    r = r_sliced.canonicalize()
    s = s_sliced.canonicalize()
    s = _extract_tile(s, s_region)
    perm = _compute_perm_r(r)
    r_shape_for_group = [int(it.extent) for it in r.shard]
    s_grp, seps = s.group(r_shape_for_group)
    s_p = s_grp.permute_by_groups(list(seps), perm)
    r_perm = r.permute_dims(perm)
    r_p = r_perm.canonicalize()
    sizes = [seps[i + 1] - seps[i] for i in range(len(seps) - 1)]
    s_seps = [0]
    for p in perm:
        s_seps.append(s_seps[-1] + sizes[p])
    return r_p, s_p, s_seps, r_perm


def _split_thread_loop(r_perm, s_p, s_seps):
    """Drop R's thread-axis positions and return per-R-position bundles:
    (r_iters, s_groups) — same length lists; s_groups[k] is the list of S
    iters belonging to the k-th kept R position.

    ``r_perm`` is the permuted-but-uncanonicalized R layout: its iter list lines
    up one-to-one with ``s_seps`` (which is built from the same ``perm``), so a
    multi-group memory axis keeps each of its groups paired with the matching S
    group instead of collapsing into one (see ``align_layouts_raw``)."""
    r_iters = []
    s_groups = []
    for k, r_it in enumerate(r_perm.shard):
        if r_it.axis.is_thread():
            continue
        r_iters.append(r_it)
        s_groups.append(list(s_p.shard[s_seps[k] : s_seps[k + 1]]))
    return r_iters, s_groups


def _build_atoms(r_iters, s_groups):
    """One atom per (R position, intra-group S iter): (extent, s_stride, r_mul).
    r_mul = R_stride_at_position * (product of S group extents to the right
    of this intra-position index) — i.e. how much R address advances per unit
    of this iter's loop input."""
    atoms = []
    for r_it, s_group in zip(r_iters, s_groups, strict=True):
        rs = int(r_it.stride)
        extents = [int(it.extent) for it in s_group]
        for j, s_it in enumerate(s_group):
            inner_prod = 1
            for e in extents[j + 1 :]:
                inner_prod *= e
            atoms.append((int(s_it.extent), int(s_it.stride), rs * inner_prod))
    return atoms


def _atoms_contiguous_tail_extent(atoms) -> int:
    """Like _contiguous_tail_extent but on atoms (uses s_stride for chaining)."""
    if not atoms or atoms[-1][1] != 1:
        return 0
    acc = atoms[-1][0]
    for k in range(len(atoms) - 2, -1, -1):
        if atoms[k][1] == acc:
            acc *= atoms[k][0]
        else:
            break
    return acc


def _split_atoms_for_vec(atoms, vec_len):
    """Returns outer atoms (the inner vec_len-element tail is consumed by one
    vec ld/st and dropped). Splits the boundary atom if needed."""
    outer = list(atoms)
    acc = 1
    while outer:
        ext, ss, rm = outer[-1]
        new_acc = acc * ext
        if new_acc == vec_len:
            outer.pop()
            return outer
        if new_acc > vec_len:
            inner_factor = vec_len // acc
            outer[-1] = (ext // inner_factor, ss * inner_factor, rm * inner_factor)
            return outer
        acc = new_acc
        outer.pop()
    raise ValueError(f"tail too short for vec_len {vec_len}")


def _align_layouts(op_call: TilePrimitiveCall, sctx: DispatchContext):
    op_call = TilePrimitiveCall.downcast(op_call)
    src_br = op_call.src
    dst_br = op_call.dst
    if src_br.buffer.scope() == "local":
        r_br, s_br = src_br, dst_br
    else:
        r_br, s_br = dst_br, src_br
    r_buf = r_br.buffer
    s_buf = s_br.buffer
    r_region = [(r.min, r.min + r.extent) for r in r_br.region]
    s_region = [(r.min, r.min + r.extent) for r in s_br.region]
    with sctx.target:
        r_sliced = r_buf.layout.slice(list(r_buf.shape), r_region)
        s_sliced = s_buf.layout.slice(list(s_buf.shape), s_region)
        return align_layouts_raw(r_sliced, s_sliced, s_region)


def _make_thread_placeholders(r_p) -> dict[str, _TirVar]:
    placeholders: dict[str, _TirVar] = {}
    for it in r_p.shard:
        name = it.axis.name
        if it.axis.is_thread() and name not in placeholders:
            placeholders[name] = _TirVar(name, "int32")
    return placeholders


def _s_thread_contributions(r_p, s_p, placeholders: dict[str, _TirVar]):
    """Return the per-S-iter coordinates contributing to a thread's address."""
    coord = [
        placeholders[it.axis.name] if it.axis.is_thread() else _IntImm("int32", 0)
        for it in r_p.shard
    ]
    input_shape = [int(it.extent) for it in r_p.shard]
    per_iter = s_p.apply_to_shape(coord, input_shape)
    return list(zip(per_iter, (it.stride for it in s_p.shard), strict=True))


_VEC_BITS_CANDIDATES = (128, 64, 32, 16, 8)


def _vec_len_candidates(elem_bits: int) -> list[int]:
    """Widest-first element counts to try; always ends with scalar (1)."""
    out: list[int] = []
    for vb in _VEC_BITS_CANDIDATES:
        if vb < elem_bits or vb % elem_bits != 0:
            continue
        n = vb // elem_bits
        if n not in out:
            out.append(n)
    if 1 not in out:
        out.append(1)
    return out


def _choose_vec_len(elem_bits: int, atoms, r_p, s_p, max_vec_len: int | None = None) -> int:
    """Widest candidate that:
      1. divides the atom contiguous-tail extent (so vec_len consecutive
         R-side regs map to vec_len contiguous S-side elements), AND
      2. keeps every per-thread / per-round address-offset term a
         multiple of vec_len, so the resulting vec ld/st pointer is
         naturally aligned to vec_bits/8 bytes.

    Only **mem-axis** strides contribute to physical address. Thread-axis
    iter strides live in partition-coord space (which thread owns which
    logical position), not in the buffer's storage space — they're
    redistributed through ``apply_to_shape`` into the mem iters and don't
    appear directly in the per-thread address. So neither r-side nor
    s-side thread-axis strides belong in the alignment check.

    The contig-tail atoms (whose extents the vec ld/st consumes) have
    stride 1 by definition; they live entirely inside the vec and
    contribute nothing to the per-round address delta. Only the
    **post-vec-split** outer atom strides matter for the per-round delta.
    """
    t = _atoms_contiguous_tail_extent(atoms)
    # Region-base offsets are real address contributions; thread-iter
    # strides are partition-virtual, not storage-physical — exclude them.
    shared_terms = list(s_p.offset.values()) + list(r_p.offset.values())
    for n in _vec_len_candidates(elem_bits):
        if max_vec_len is not None and n > max_vec_len:
            continue
        if n == 1:
            return n
        if t % n != 0 or t < n:
            continue
        # Post-vec-split outer atoms: these are the strides that contribute
        # to per-round address deltas after the vec consumes the inner tail.
        outer = _split_atoms_for_vec(atoms, n)
        outer_atom_terms = [a[1] for a in outer] + [a[2] for a in outer]
        if not _alignment_ok(n, outer_atom_terms + shared_terms):
            continue
        return n
    return 1


def _axis_decl(axis_name: str, sctx: DispatchContext):
    """Declare the runtime Var for one thread axis (called inside impl body).

    Each scope_id declarator emits a ``ScopeIdDef`` stmt at the current
    builder frame. ``TilePrimitiveDispatch`` re-gathers + resolves all
    ScopeIdDefs after dispatch (see ``ResolveAllScopeBinds`` in
    ``tile_primitive_dispatch.cc``), so dispatch-introduced vars are bound
    alongside kernel-declared ones.

    Extents are deferred: the kernel header is expected to declare the full
    scope-id chain (``cta_id`` / ``warpgroup_id`` / ``warp_id_in_wg`` /
    ``lane_id`` / ``thread_id`` / ``thread_id_in_wg``) — the verifier then
    fills our deferred defs from those siblings.
    """
    if axis_name == "tx":
        return sctx.launch_params["threadIdx.x"].var
    if axis_name == "laneid":
        return T.lane_id()
    if axis_name == "wid_in_wg":
        return T.warp_id_in_wg()
    if axis_name == "tid_in_wg":
        return T.thread_id_in_wg()
    if axis_name == "warpid":
        return T.warp_id()
    if axis_name == "wgid":
        return T.warpgroup_id()
    raise ValueError(f"unsupported thread axis {axis_name}")


def _thread_axis_extent(axis_name: str, sctx: DispatchContext) -> int | None:
    """Return the full hardware range of a scope-id axis when it is static."""
    if axis_name in sctx.intra:
        active = sctx.intra[axis_name]
        try:
            extent, offset = int(active[0]), int(active[1])
        except (TypeError, ValueError):
            return None
        if offset != 0:
            return None
        return extent

    fixed = {"laneid": 32, "wid_in_wg": 4, "tid_in_wg": 128}
    if axis_name in fixed:
        return fixed[axis_name]

    tx = sctx.launch_params.get("threadIdx.x")
    if tx is None:
        return None
    try:
        tx_extent = int(tx.dom.extent)
    except (TypeError, ValueError):
        return None
    divisor = {"tx": 1, "warpid": 32, "wgid": 128}.get(axis_name)
    if divisor is None or tx_extent % divisor != 0:
        return None
    return tx_extent // divisor


def _build_s_apply_layout(original_layout, r_p, s_p, outer, placeholders, sctx):
    """Build a structured tile whose ``m`` is exactly old ``s_off + ds``.

    Each bounded per-thread coordinate becomes its own Iter with a proven
    hardware-derived extent. Outer atoms retain their exact extent/stride and
    are addressed by one flattened loop coordinate through shape-aware apply.
    Slice offsets and any thread term without a finite non-negative bound stay
    in the tile's ``m`` offset, which makes ComposeLayout conservatively use
    its full-swizzle fallback when needed.
    """
    analyzer = Analyzer()
    for name, placeholder in placeholders.items():
        extent = _thread_axis_extent(name, sctx)
        if extent is not None:
            analyzer.bind(placeholder, tvm.ir.Range.from_min_extent(0, extent))

    shard = []
    thread_coords = []
    apply_shape = []
    offset = _IntImm("int32", 0)
    for coord, stride in _s_thread_contributions(r_p, s_p, placeholders):
        coord = analyzer.simplify(coord)
        if analyzer.can_prove_equal(coord, 0):
            continue
        bound = analyzer.const_int_bound(coord)
        if (
            bound.min_value >= 0
            and bound.max_value != ConstIntBound.POS_INF
            and bound.max_value < (1 << 31) - 1
        ):
            extent = int(bound.max_value) + 1
            extent_expr = _IntImm("int32", extent)
            shard.append(Iter(extent_expr, stride, "m"))
            thread_coords.append(coord)
            apply_shape.append(extent_expr)
        else:
            offset = analyzer.simplify(offset + coord * stride)

    for _axis, value in s_p.offset.items():
        offset = analyzer.simplify(offset + value)

    total_outer = 1
    for extent, stride, _r_mul in outer:
        shard.append(Iter(extent, stride, "m"))
        total_outer *= extent
    apply_shape.append(_IntImm("int32", total_outer))

    tile = TileLayout.from_iters(shard, [], {"m": offset})
    return recompose_swizzle(original_layout, tile), thread_coords, apply_shape


def _axis_substitution(placeholders: dict[str, _TirVar], sctx: DispatchContext):
    """Declare the real scope ids once and map placeholder Vars to them."""
    return {placeholders[name]: _axis_decl(name, sctx) for name in placeholders}


def _apply_s_layout(layout, thread_coords, flat_outer, shape, placeholders, sctx):
    coord = [*thread_coords, flat_outer]
    mapped = layout.apply(*coord, shape=shape)["m"]
    return tvm.tirx.stmt_functor.substitute(mapped, _axis_substitution(placeholders, sctx))


def _flat_coords(outer_atoms, flat_idx: int) -> list[int]:
    coords = []
    rem = flat_idx
    for a in reversed(outer_atoms):
        coords.append(rem % a[0])
        rem //= a[0]
    coords.reverse()
    return coords


_POINTER_OFFSET_SRC = (
    "\ntemplate <typename T>\n"
    "__forceinline__ __device__ T* tvm_builtin_pointer_offset(T* ptr, int offset) {\n"
    "    return ptr + offset;\n"
    "}\n"
)


def _ptr_off(base_ptr, off):
    return T.cuda.func_call(
        "tvm_builtin_pointer_offset",
        base_ptr,
        off,
        source_code=_POINTER_OFFSET_SRC,
        return_type="handle",
    )


def _outer_const_offsets(outer_atoms, flat_idx: int) -> tuple[int, int]:
    """Returns (s_offset_const, r_offset_const) for one outer-loop flat index."""
    coords = _flat_coords(outer_atoms, flat_idx)
    ds = sum(c * a[1] for c, a in zip(coords, outer_atoms))
    dr = sum(c * a[2] for c, a in zip(coords, outer_atoms))
    return ds, dr


def _emit_reg(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
    op_call = TilePrimitiveCall.downcast(op_call)
    src: Buffer = op_call.src.buffer
    dst: Buffer = op_call.dst.buffer
    if src.scope() == "local":
        r_buf, s_buf, r_is_src = src, dst, True
    else:
        r_buf, s_buf, r_is_src = dst, src, False

    with sctx.target:
        r_p, s_p, s_seps, r_perm = _align_layouts(op_call, sctx)
    r_iters, s_groups = _split_thread_loop(r_perm, s_p, s_seps)
    atoms = _build_atoms(r_iters, s_groups)
    elem_bits = DataType(src.dtype).bits
    max_vec_len = None
    if isinstance(s_buf.layout, ComposeLayout) and int(s_buf.layout.swizzle_len) > 0:
        max_vec_len = 1 << int(s_buf.layout.per_element)
    vec_len = _choose_vec_len(elem_bits, atoms, r_p, s_p, max_vec_len)
    vec_bits = vec_len * elem_bits
    outer = _split_atoms_for_vec(atoms, vec_len)

    placeholders = _make_thread_placeholders(r_p)
    s_apply_layout, thread_coords, s_apply_shape = _build_s_apply_layout(
        s_buf.layout, r_p, s_p, outer, placeholders, sctx
    )
    # A plain TileLayout is additive, so keep its loop-invariant thread base
    # outside the serial copy loop. A real swizzle must still apply to the
    # complete base + outer coordinate inside each iteration.
    has_swizzle = isinstance(s_apply_layout, ComposeLayout) and int(s_apply_layout.swizzle_len) > 0

    # R-side base offset from slicing (e.g. ``R[i*8:i*8+8]`` ⇒ ``i*8``), from
    # ``r_p.offset``; sum across all axes on R's local stride-1 storage.
    r_off_base = _IntImm("int32", 0)
    for _ax, val in r_p.offset.items():
        r_off_base = r_off_base + val

    s_is_shared = s_buf.scope().startswith("shared")
    num_bytes = vec_bits // 8
    tail, lanes, reg_dtype = copy_ptx_form(num_bytes)
    space = "shared" if s_is_shared else "global"

    # Honor the load cache hint (cache="nc") only on global -> register loads;
    # ``nc`` is a global qualifier, so stores and shared copies stay plain.
    _cache, _cache_hints = _ld_cache_config(op_call)
    _use_nc = _cache == "nc" and space == "global" and not r_is_src
    _l1_evict = _cache_hints.get("l1_evict", "")
    _l2_evict = _cache_hints.get("l2_evict", "")
    _prefetch = _cache_hints.get("prefetch_size", "")
    # Element offsets index the element buffer; the instruction moves whole
    # containers, so they are scaled to container index. The vector alignment
    # rule guarantees the division is exact.
    container_bits = num_bytes * 8 // lanes
    elems_per_word = container_bits // elem_bits if container_bits >= elem_bits else 1
    words_per_elem = elem_bits // container_bits if elem_bits > container_bits else 1
    # Chains are built here: a Python string bound inside the traced body is
    # not something the parser can carry. `nc` folds into the chain, so the
    # three-way branch at the call site collapses to load-or-store.
    st_chain = f"st.{space}.{tail}"
    ld_chain = copy_ptx_ld_chain(
        space,
        tail,
        nc=_use_nc,
        l1_evict=_l1_evict,
        l2_evict=_l2_evict,
        prefetch_size=_prefetch,
    )

    total_outer = 1
    for a in outer:
        total_outer *= a[0]

    # fmt: off
    s_zero_indices = [0] * len(s_buf.shape)

    @T.prim_func(check_well_formed=False)
    def impl():
        if not has_swizzle:
            s_base = _apply_s_layout(
                s_apply_layout, thread_coords, 0, s_apply_shape, placeholders, sctx
            )
        # ``dr`` and ``r_off_base`` below are physical storage offsets, so use
        # the raw-span local view rather than storage-iterator coordinates.
        r_local = r_buf.local()
        # The registers the instruction moves are containers, not elements.
        r_words = r_local.view(reg_dtype)
        # Keep a serial TIR loop and let ptxas unroll; explicit ``T.unroll``
        # replicates per-iter scratch arrays and pressures registers.
        for f in range(total_outer):
            ds, dr = _outer_const_offsets(outer, f)
            if has_swizzle:
                s_off = _apply_s_layout(
                    s_apply_layout, thread_coords, f, s_apply_shape, placeholders, sctx
                )
                s_ptr = _ptr_off(s_buf.ptr_to(s_zero_indices), s_off)
            else:
                s_ptr = _ptr_off(s_buf.ptr_to(s_zero_indices), s_base + ds)
            # The registers move through the reg_dtype word view, not a raw
            # element pointer: the container index is the element offset scaled
            # between element and word granularity.
            r_w = (r_off_base + dr) * words_per_elem // elems_per_word
            if r_is_src:
                T.ptx[st_chain](s_ptr, *[r_words[r_w + i] for i in range(lanes)])
            else:
                T.ptx[ld_chain](*[r_words[r_w + i] for i in range(lanes)], s_ptr)
    # fmt: on
    import os

    if os.environ.get("R2S_DUMP"):
        print("=== emitted impl ===")
        print(impl.script())
    return impl
