gemini-code-assist[bot] commented on code in PR #19657: URL: https://github.com/apache/tvm/pull/19657#discussion_r3343976472
########## python/tvm/tirx/operator/tile_primitive/cuda/copy/_common.py: ########## @@ -0,0 +1,540 @@ +# 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. + +"""Shared partition / layout algorithm for synthesized-partition copy +dispatches (currently ``gmem_smem`` and ``ldgsts``). + +``gmem_smem`` (sync ``Tx.copy`` global ↔ shared) and ``ldgsts`` (async +``Tx.copy_async`` global → shared via cp.async / SASS LDGSTS) share the +same algorithm to pick a vec-isolating + thread-distributing layout for +``G ↔ S`` copies. Only emit-time details differ (which copy instruction +to call, allowed vec widths). All the layout/partition logic lives here. +""" + +from tvm import arith +from tvm.tirx.layout import ComposeLayout, Iter, S, SwizzleLayout, TileLayout +from tvm.tirx.operator.tile_primitive.registry import DispatchContext + + +def _alignment_ok(vec_len: int, terms) -> bool: + """Every term must be a multiple of ``vec_len``. Constants checked + directly; PrimExpr / symbolic terms checked via ``arith.Analyzer``. + + ``vec_len=1`` always passes (the scalar fallback). When a symbolic + term can't be proved divisible, returns ``False`` conservatively — + the candidate loop will then try a smaller ``vec_len``. + """ + if vec_len <= 1: + return True + analyzer = arith.Analyzer() + for t in terms: + if isinstance(t, int): + if t % vec_len != 0: + return False + else: + if not analyzer.can_prove_equal(t % vec_len, 0): + return False + return True + + +# scope_kind → name of the scope_id that decomposes the scope into per-thread. +_TID_AXIS_FOR_SCOPE = { + "warp": "laneid", + "warpgroup": "tid_in_wg", + "cta": "tx", +} + + +def _thread_cnt(sctx: DispatchContext) -> int: + """Total threads active in the current scope = ∏ intra-axis extents. + + For thread scope ``sctx.intra`` is empty → returns 1. + """ + n = 1 + for ext, _off in sctx.intra.values(): + n *= int(ext) + return n + + +# ----------------------------------------------------------------------------- +# Layout primitives +# ----------------------------------------------------------------------------- + + +def _contig_group(iters: list) -> list[int]: + """Indices (in iters) of the maximal physical-contiguous chain starting + at the stride=1 iter, ordered stride-ascending. + + Returns [] if no stride=1 iter exists. + """ + one_idx = next( + (i for i, it in enumerate(iters) if int(it.stride) == 1), + None, + ) + if one_idx is None: + return [] + chain = [one_idx] + acc = int(iters[one_idx].extent) + used = {one_idx} + while True: + nxt = next( + (i for i, it in enumerate(iters) if i not in used and int(it.stride) == acc), + None, + ) + if nxt is None: + break + chain.append(nxt) + acc *= int(iters[nxt].extent) + used.add(nxt) + return chain + + +def _try_split_vec(iters: list, vec_len: int): + """Try to walk ``vec_len`` consecutive elements along the contig chain. + + Returns ``(new_iters, selected_positions)`` on success, ``None`` on + failure. ``new_iters`` may contain a freshly-split iter (replacing one + entry with its "outer" half, with the "inner" half appended at the end); + ``selected_positions`` are positions in ``new_iters`` that together + cover the ``vec_len`` contig elements. + """ + chain = _contig_group(iters) + if not chain: + return None + rem = vec_len + new_iters = list(iters) + selected: list[int] = [] + for orig_idx in chain: + if rem == 0: + break + it = new_iters[orig_idx] + ext = int(it.extent) + if ext <= rem: + if rem % ext != 0: + return None + selected.append(orig_idx) + rem //= ext + else: + if ext % rem != 0: + return None + stride = int(it.stride) + outer = Iter(ext // rem, stride * rem, it.axis) + inner = Iter(rem, stride, it.axis) + new_iters[orig_idx] = outer + new_iters.append(inner) + selected.append(len(new_iters) - 1) + rem = 0 + break + if rem != 0: + return None + return new_iters, selected + + +def _isolated_shape(iters: list, selected: list[int]) -> tuple[list[int], list[tuple[int, int]]]: + """Build the isolated shape: each selected iter is its own segment; + adjacent unselected iters are merged into a single segment. + + Returns ``(shape, segments)`` where ``segments[i] = (start, end)`` is + the half-open range in ``iters`` covered by shape entry ``i``. + """ + sel_set = set(selected) + shape: list[int] = [] + segments: list[tuple[int, int]] = [] + cur_start = None + cur_ext = 1 + for i, it in enumerate(iters): + if i in sel_set: + if cur_start is not None: + shape.append(cur_ext) + segments.append((cur_start, i)) + cur_start = None + cur_ext = 1 + shape.append(int(it.extent)) + segments.append((i, i + 1)) + else: + if cur_start is None: + cur_start = i + cur_ext *= int(it.extent) + if cur_start is not None: + shape.append(cur_ext) + segments.append((cur_start, len(iters))) + return shape, segments + + +def _vec_perm(iters: list, selected: list[int]) -> list[int]: + """Reorder ``iters`` into ``[outer, vec]``, both ordered by stride + descending so the stride=1 iter ends up at the very last position.""" + sel_set = set(selected) + unsel_sorted = sorted( + (i for i in range(len(iters)) if i not in sel_set), + key=lambda i: -int(iters[i].stride), + ) + sel_sorted = sorted(selected, key=lambda i: -int(iters[i].stride)) + return list(unsel_sorted) + sel_sorted + + +def _try_split_thread(iters: list, vec_selected: list[int], thread_cnt: int): + """After ``_try_split_vec``, carve ``thread_cnt`` from the OUTER tail + (smallest-stride outer iter, then towards bigger stride if needed). + + Unlike vec split, this doesn't require physical contiguity — T + consecutive fused indices map to per-thread offsets via the layout's + stride (which may be > 1). + + Returns ``(new_iters, thread_selected_positions)`` on success, ``None`` + on failure (outer doesn't divide T cleanly, or no outer iters left). + """ + if thread_cnt == 1: + return list(iters), [] + vec_set = set(vec_selected) + outer = [i for i in range(len(iters)) if i not in vec_set] + if not outer: + return None + outer_by_stride_desc = sorted(outer, key=lambda i: -int(iters[i].stride)) + rem = thread_cnt + new_iters = list(iters) + thread_selected: list[int] = [] + for orig_idx in reversed(outer_by_stride_desc): + if rem == 0: + break + it = new_iters[orig_idx] + ext = int(it.extent) + if ext <= rem: + if rem % ext != 0: + return None + thread_selected.append(orig_idx) + rem //= ext + else: + if ext % rem != 0: + return None + stride = int(it.stride) + new_iters[orig_idx] = Iter(ext // rem, stride * rem, it.axis) + new_iters.append(Iter(rem, stride, it.axis)) + thread_selected.append(len(new_iters) - 1) + rem = 0 + break + if rem != 0: + return None + return new_iters, thread_selected + + +def _three_segment_perm(iters: list, t_selected: list[int], vec_selected: list[int]) -> list[int]: + """Reorder ``iters`` into ``[outer, T, vec]`` segments. Within each + segment, stride descending so stride=1 sits at the very end.""" + t_set = set(t_selected) + vec_set = set(vec_selected) + outer = sorted( + (i for i in range(len(iters)) if i not in t_set and i not in vec_set), + key=lambda i: -int(iters[i].stride), + ) + t_sorted = sorted(t_selected, key=lambda i: -int(iters[i].stride)) + vec_sorted = sorted(vec_selected, key=lambda i: -int(iters[i].stride)) + return list(outer) + t_sorted + vec_sorted + + +def _shape_perm_for_isolated( + shape_segments: list[tuple[int, int]], iter_perm: list[int] +) -> list[int]: + """Given segments (one per shape entry, each = (start, end) in + pre-perm iter positions) and the iter permutation, compute the + corresponding shape permutation.""" + seg_of = [0] * sum(end - start for start, end in shape_segments) + for seg_idx, (start, end) in enumerate(shape_segments): + for k in range(start, end): + seg_of[k] = seg_idx + seen: set[int] = set() + perm: list[int] = [] + for orig_idx in iter_perm: + seg_idx = seg_of[orig_idx] + if seg_idx not in seen: + seen.add(seg_idx) + perm.append(seg_idx) + return perm + + +def _verify_s_tail_contig(s_p: TileLayout, vec_len: int) -> bool: + """Check the last iters of ``s_p`` form a stride=1 contig chain whose + extent product equals ``vec_len``.""" + iters = list(s_p.shard) + if not iters: + return vec_len == 1 + last = iters[-1] + if int(last.stride) != 1: + return False + acc = int(last.extent) + if acc == vec_len: + return True + for k in range(len(iters) - 2, -1, -1): + it = iters[k] + if int(it.stride) != acc: + break + acc *= int(it.extent) + if acc == vec_len: + return True + if acc > vec_len: + return False + return acc >= vec_len and acc % vec_len == 0 + + +_VEC_BITS_CANDIDATES = (128, 64, 32, 16, 8) + + +def _vec_len_candidates(elem_bits: int, allowed_bits: tuple | None = None) -> list[int]: + """Vec-length candidates (in elements) for the given element width. + + ``allowed_bits`` optionally filters the per-instruction allowed widths + (e.g. cp.async only accepts {128, 64, 32} bits = 16/8/4 bytes). + Defaults to ``_VEC_BITS_CANDIDATES`` if not specified. + """ + bits_tuple = allowed_bits if allowed_bits is not None else _VEC_BITS_CANDIDATES + out: list[int] = [] + for vb in bits_tuple: + 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 and allowed_bits is None: + # Scalar fallback is only added for the unrestricted candidate set; + # an instruction-specific list (cp.async etc.) keeps its strictness. + out.append(1) + return out + + +def _extract_tile(layout, region): + """Strip swizzle so we can perm/group as a TileLayout.""" + if isinstance(layout, ComposeLayout): + return layout.tile_layout + if isinstance(layout, SwizzleLayout): + extents = [int(end - start) for (start, end) in region] + return TileLayout(S[tuple(extents)]) + return layout + + +def _sort_by_stride_desc(layout: TileLayout) -> TileLayout: + """Reorder shard so list order = traversal order (outer first, stride=1 + last). Required before canonicalize() can fuse non-adjacent-but-contig + iters.""" + iters = list(layout.shard) + perm = sorted(range(len(iters)), key=lambda i: -int(iters[i].stride)) + if perm == list(range(len(iters))): + return layout + return layout.permute_dims(perm) + + +def _carve_tail(iters: list, chunk: int): + """Carve ``chunk`` elements off the tail. Walk back across multiple + iters as needed; at most one iter is split. + + Per iter (from last to first), let ``ext`` = iter extent and ``rem`` + = remaining chunk to fill: + + * ``ext == rem``: eat this iter whole, done. + * ``ext < rem``: must divide ``rem``; eat whole, ``rem //= ext``. + * ``ext > rem``: must divide ``ext``; split into + ``(ext/rem, stride*rem) + (rem, stride)``, take the inner. Done. + + Returns the new iter list on success, ``None`` on failure. + """ + if not iters or chunk <= 0: + return None + rem = chunk + work = list(iters) + for idx in range(len(work) - 1, -1, -1): + it = work[idx] + ext = int(it.extent) + if ext == rem: + return work + if ext < rem: + if rem % ext != 0: + return None + rem //= ext + continue + if ext % rem != 0: + return None + stride = int(it.stride) + work[idx] = Iter(ext // rem, stride * rem, it.axis) + work.insert(idx + 1, Iter(rem, stride, it.axis)) + return work + return None Review Comment:  ### Critical Correctness Bug in `_carve_tail` In the current implementation of `_carve_tail`, when the suffix of `work` matches the requested `chunk` size (either exactly when `ext == rem` or after splitting when `ext > rem`), the function returns the modified `work` list containing all original elements. However, the caller of `_carve_tail` (e.g., `align_layouts_gs`) expects the carved tail to be represented by a **single** element at the very end of the returned list (i.e., `g_after_vec[-1]`), and treats all preceding elements (`g_after_vec[:-1]`) as the remaining outer/thread loops. Because `_carve_tail` does not fuse the consumed suffix into a single element, the preceding elements of the carved tail are incorrectly left in `outer_for_t`, leading to redundant/incorrect outer loop iterations and misaligned vector memory accesses. #### Suggested Fix Slice and replace the consumed suffix with a single `Iter(chunk, 1, work[-1].axis)` representing the carved tail. ```python def _carve_tail(iters: list, chunk: int): if not iters or chunk <= 0: return None rem = chunk work = list(iters) for idx in range(len(work) - 1, -1, -1): it = work[idx] ext = int(it.extent) if ext == rem: return work[:idx] + [Iter(chunk, 1, work[-1].axis)] if ext < rem: if rem % ext != 0: return None rem //= ext continue if ext % rem != 0: return None stride = int(it.stride) work[idx] = Iter(ext // rem, stride * rem, it.axis) return work[:idx + 1] + [Iter(chunk, 1, work[-1].axis)] return None ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
