chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,24 @@
|
||||
# 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.
|
||||
|
||||
from .copy import *
|
||||
from .copy_async import *
|
||||
from .elementwise import *
|
||||
from .gemm import *
|
||||
from .gemm_async import *
|
||||
from .permute_layout import *
|
||||
from .reduction import *
|
||||
@@ -0,0 +1,283 @@
|
||||
# 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.
|
||||
|
||||
"""Common utilities for CUDA operator scheduling (basic helpers and copy ops)."""
|
||||
|
||||
import functools
|
||||
import operator
|
||||
import re
|
||||
from enum import Enum
|
||||
|
||||
from tvm.arith.analyzer import Analyzer
|
||||
from tvm.runtime import DataType
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import Buffer, BufferRegion, PrimFunc
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
|
||||
def next_power_of_2(x: int) -> int:
|
||||
"""Return the smallest power of 2 greater than or equal to x."""
|
||||
if x <= 1:
|
||||
return 1
|
||||
return 1 << (x - 1).bit_length()
|
||||
|
||||
|
||||
def get_st_extent(buffer_region: BufferRegion):
|
||||
"""Get the start and extent of a buffer region."""
|
||||
region = buffer_region.region
|
||||
return [r.min for r in region], [r.extent for r in region]
|
||||
|
||||
|
||||
def get_indices(nth, start, extent):
|
||||
"""Convert a fused index into multi-dimensional indices."""
|
||||
assert len(start) == len(extent)
|
||||
if len(start) == 1:
|
||||
return [start[0] + nth]
|
||||
relative = []
|
||||
for e in reversed(extent):
|
||||
relative.append(nth % e)
|
||||
nth //= e
|
||||
return [r + s for r, s in zip(reversed(relative), start)]
|
||||
|
||||
|
||||
def smem_desc_add_16B_offset(desc_val, offset):
|
||||
"""Add a 16B-aligned byte offset to the lower 32 bits of a SMEM descriptor.
|
||||
|
||||
Uses the SmemDescriptor union defined in the CUDA header (header.py).
|
||||
All callers must share a single implementation to avoid codegen conflicts.
|
||||
"""
|
||||
func_name = "tvm_builtin_smem_desc_add_16B_offset"
|
||||
source_code = f"""
|
||||
__forceinline__ __device__ uint64_t {func_name}(uint64_t desc_base, int32_t offset) {{
|
||||
SmemDescriptor desc;
|
||||
desc.desc_ = desc_base;
|
||||
desc.lo += static_cast<uint32_t>(offset);
|
||||
return desc.desc_;
|
||||
}}
|
||||
"""
|
||||
return T.cuda.func_call(
|
||||
func_name, desc_val, offset, source_code=source_code, return_type="uint64"
|
||||
)
|
||||
|
||||
|
||||
class CopyInstType(Enum):
|
||||
"""Enumeration of instruction types for memory operations."""
|
||||
|
||||
NORMAL = 0
|
||||
CP_ASYNC = 1
|
||||
|
||||
|
||||
def validate_copy_op(
|
||||
op_call: TilePrimitiveCall,
|
||||
sctx: DispatchContext, # pylint: disable=unused-argument
|
||||
) -> bool:
|
||||
"""Sanity check for copy op"""
|
||||
dst_buffer_region, src_buffer_region = op_call.args[:2]
|
||||
src: Buffer = src_buffer_region.buffer
|
||||
dst: Buffer = dst_buffer_region.buffer
|
||||
if not (src.layout and dst.layout and src.dtype == dst.dtype):
|
||||
return False
|
||||
# Extract regions and validate dimensions
|
||||
analyzer = Analyzer()
|
||||
src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
|
||||
# Extract extents and validate non-unit dimensions match
|
||||
src_extent_ = [r.extent for r in src_region if r.extent != 1]
|
||||
dst_extent_ = [r.extent for r in dst_region if r.extent != 1]
|
||||
if len(src_extent_) != len(dst_extent_) or not all(
|
||||
analyzer.can_prove_equal(s, d) for s, d in zip(src_extent_, dst_extent_)
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_vec_len(
|
||||
dst_buffer_region: BufferRegion,
|
||||
src_buffer_region: BufferRegion,
|
||||
vec_candidates: list[int],
|
||||
thread_cnt=1,
|
||||
) -> int | None:
|
||||
"""Get the vector length for the copy operation."""
|
||||
|
||||
dst: Buffer = dst_buffer_region.buffer
|
||||
src: Buffer = src_buffer_region.buffer
|
||||
# layout=None (flat local buffer) is treated as trivial for vectorization purposes
|
||||
if not (
|
||||
(dst.layout is None or dst.layout.is_trivial())
|
||||
and (src.layout is None or src.layout.is_trivial())
|
||||
):
|
||||
return None
|
||||
|
||||
# Extract regions and validate dimensions
|
||||
analyzer = Analyzer()
|
||||
src_st, src_extent = get_st_extent(src_buffer_region)
|
||||
dst_st, dst_extent = get_st_extent(dst_buffer_region)
|
||||
|
||||
# Thread and vectorization setup
|
||||
DataType(src.dtype).bits # in bits
|
||||
n_elements = functools.reduce(operator.mul, src_extent, 1)
|
||||
if n_elements % thread_cnt != 0:
|
||||
return None
|
||||
|
||||
# Find valid vector length
|
||||
for vec_len in vec_candidates:
|
||||
if vec_len > 0 and all(
|
||||
analyzer.can_prove_equal(x % vec_len, 0)
|
||||
for x in [
|
||||
src_st[-1],
|
||||
dst_st[-1],
|
||||
src.shape[-1] if len(src.shape) > 1 else 0,
|
||||
dst.shape[-1] if len(dst.shape) > 1 else 0,
|
||||
src_extent[-1],
|
||||
dst_extent[-1],
|
||||
n_elements // thread_cnt,
|
||||
]
|
||||
):
|
||||
return vec_len
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def copy_vec_load_impl(
|
||||
op_call: TilePrimitiveCall, sctx: DispatchContext, inst_type: CopyInstType
|
||||
) -> PrimFunc | None:
|
||||
"""Schedule copy operation between global and local/shared memory on CUDA across a CTA/thread.
|
||||
The implementation tries to vectorize the copy operation and parallelize over
|
||||
threads in a CTA/using a single thread.
|
||||
"""
|
||||
dst_buffer_region, src_buffer_region = op_call.args[:2]
|
||||
src: Buffer = src_buffer_region.buffer
|
||||
dst: Buffer = dst_buffer_region.buffer
|
||||
if not (
|
||||
(src.scope() == "global" and dst.scope().startswith("shared"))
|
||||
or (src.scope().startswith("shared") and dst.scope() == "global")
|
||||
or (src.scope() == "global" and dst.scope() == "local")
|
||||
or (src.scope() == "local" and dst.scope() == "global")
|
||||
or (src.scope().startswith("shared") and dst.scope() == "local")
|
||||
or (dst.scope().startswith("shared") and src.scope() == "local")
|
||||
):
|
||||
fail(f"unsupported memory scopes src={src.scope()} dst={dst.scope()}")
|
||||
|
||||
# Thread and vectorization setup
|
||||
if sctx.is_cta:
|
||||
tx = sctx.launch_params["threadIdx.x"].dom.extent
|
||||
assert "threadIdx.y" not in sctx.launch_params and "threadIdx.z" not in sctx.launch_params
|
||||
elif sctx.is_thread:
|
||||
tx = 1
|
||||
else:
|
||||
fail(f"unsupported exec_scope {sctx.scope_kind}")
|
||||
|
||||
elem_size = DataType(src.dtype).bits # in bits
|
||||
vec_len = op_call.config.get("vec_len", None)
|
||||
if vec_len is None:
|
||||
vec_len = get_vec_len(
|
||||
dst_buffer_region,
|
||||
src_buffer_region,
|
||||
[128 // elem_size, 64 // elem_size, 32 // elem_size, 1],
|
||||
thread_cnt=tx,
|
||||
)
|
||||
if vec_len is None:
|
||||
fail("no valid vector length; check alignment/extents/thread-count")
|
||||
|
||||
# cp-size (the size of data in bytes) can only be 4, 8 and 16 for cp.async
|
||||
if inst_type == CopyInstType.CP_ASYNC:
|
||||
cp_size = vec_len * elem_size // 8 # in bytes
|
||||
if cp_size not in [4, 8, 16]:
|
||||
fail("invalid cp.async cp_size; expected 4, 8 or 16 bytes")
|
||||
|
||||
src_st, src_extent = get_st_extent(src_buffer_region)
|
||||
dst_st, dst_extent = get_st_extent(dst_buffer_region)
|
||||
n_elements = functools.reduce(operator.mul, src_extent, 1)
|
||||
|
||||
if sctx.is_cta:
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def impl():
|
||||
"""Implement copy operation with vectorized loads/stores."""
|
||||
for s in T.serial(0, n_elements // (tx * vec_len)):
|
||||
for tid_x in T.thread_binding(tx, "threadIdx.x"):
|
||||
if inst_type == CopyInstType.NORMAL:
|
||||
for vec in T.vectorized(vec_len):
|
||||
fused = T.meta_var((s * tx + tid_x) * vec_len + vec)
|
||||
dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
|
||||
src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
|
||||
dst[tuple(dst_indices)] = src[tuple(src_indices)]
|
||||
elif inst_type == CopyInstType.CP_ASYNC:
|
||||
fused = T.meta_var((s * tx + tid_x) * vec_len)
|
||||
dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
|
||||
src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
|
||||
T.evaluate(T.ptx.cp_async(dst.ptr_to(dst_indices), src.ptr_to(src_indices), cp_size)) # noqa: E501
|
||||
if dst.scope().startswith("shared") and inst_type == CopyInstType.NORMAL:
|
||||
T.tvm_storage_sync("shared")
|
||||
# fmt: on
|
||||
elif sctx.is_thread:
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
for s in T.serial(0, n_elements // (vec_len)):
|
||||
if inst_type == CopyInstType.NORMAL:
|
||||
for vec in T.vectorized(vec_len):
|
||||
fused = T.meta_var(s * vec_len + vec)
|
||||
dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
|
||||
src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
|
||||
dst[tuple(dst_indices)] = src[tuple(src_indices)]
|
||||
elif inst_type == CopyInstType.CP_ASYNC:
|
||||
fused = T.meta_var(s * vec_len)
|
||||
dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
|
||||
src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
|
||||
T.evaluate(T.ptx.cp_async(dst.ptr_to(dst_indices), src.ptr_to(src_indices), cp_size)) # noqa: E501
|
||||
# fmt: on
|
||||
else:
|
||||
fail(f"unsupported exec_scope {sctx.scope_kind}")
|
||||
return impl
|
||||
|
||||
|
||||
def match_scope(scope: str | None, pattern: str) -> bool:
|
||||
"""Glob-lite scope matching: 'shared*' => prefix match; otherwise exact.
|
||||
|
||||
Returns True when scope is None (meaning "any scope is fine").
|
||||
"""
|
||||
if scope is None:
|
||||
return True
|
||||
if pattern.endswith("*"):
|
||||
return scope.startswith(pattern[:-1])
|
||||
return scope == pattern
|
||||
|
||||
|
||||
def get_thread_cnt(sctx: DispatchContext) -> int | None:
|
||||
"""Get thread count for the current execution scope."""
|
||||
scope_name = sctx.scope_kind
|
||||
if scope_name == "cta":
|
||||
return sctx.launch_params["threadIdx.x"].dom.extent
|
||||
if scope_name == "warpgroup":
|
||||
return 128
|
||||
if scope_name == "warp":
|
||||
return 32
|
||||
if scope_name == "thread":
|
||||
return 1
|
||||
return None
|
||||
|
||||
|
||||
def sm_version_ok(
|
||||
op: TilePrimitiveCall, sctx: DispatchContext, min_version: int
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Check if SM version >= min_version. Usable as a dispatch predicate."""
|
||||
target_arch = sctx.target.arch if hasattr(sctx.target, "arch") else ""
|
||||
sm_match = re.match(r"sm_(\d+)", target_arch)
|
||||
sm_version = int(sm_match.group(1)) if sm_match else 0
|
||||
ok = sm_version >= min_version
|
||||
return (ok, None if ok else f"sm_version {sm_version} < {min_version}")
|
||||
@@ -0,0 +1,27 @@
|
||||
# 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.
|
||||
|
||||
from .fallback import *
|
||||
from .gmem_smem import *
|
||||
from .ld_stmatrix import *
|
||||
from .reg import *
|
||||
from .utils import (
|
||||
_is_valid_copy,
|
||||
_is_valid_smem_tmem_copy,
|
||||
_scope_allowed,
|
||||
_single_thread_exec,
|
||||
)
|
||||
@@ -0,0 +1,556 @@
|
||||
# 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; Expr / 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
|
||||
|
||||
|
||||
def align_layouts_gs(
|
||||
g_layout,
|
||||
g_shape,
|
||||
g_region,
|
||||
s_layout,
|
||||
s_shape,
|
||||
s_region,
|
||||
elem_bits,
|
||||
thread_cnt: int,
|
||||
vec_bits_candidates: tuple | None = None,
|
||||
debug: bool = False,
|
||||
):
|
||||
"""Align G and S layouts for a synthesized G↔S copy.
|
||||
|
||||
Algorithm:
|
||||
1. Sort G iters by stride desc + canonicalize (fuses anything physically
|
||||
contig). Same for S.
|
||||
2. Group S by G's iter shape (per-iter shape list) + permute_by_groups
|
||||
with identity so S's groups line up with G's iters.
|
||||
3. Carve ``vec_len`` off G's tail (a single iter split if needed).
|
||||
4. Carve ``T = thread_cnt`` off the iter before vec (multi-iter walk,
|
||||
single split).
|
||||
5. Re-group S (already permuted) by the new finer per-iter shape. No
|
||||
further permute needed because steps 3-4 only refine the tail.
|
||||
6. Verify S's tail (vec segment) is physically contig.
|
||||
|
||||
``vec_bits_candidates`` optionally restricts the per-instruction allowed
|
||||
widths (e.g. cp.async only accepts {128, 64, 32} bits). When ``None``
|
||||
(default), uses the full {128, 64, 32, 16, 8} set plus a scalar (1)
|
||||
fallback.
|
||||
|
||||
Returns ``(g_p, s_p, vec_len)``. ``g_p.shard`` ends as
|
||||
``[outer iters..., T iter, vec iter]``; ``s_p.shard`` has the same iter
|
||||
count and matching iter-by-iter extents.
|
||||
"""
|
||||
g = g_layout.slice(list(g_shape), g_region)
|
||||
s = s_layout.slice(list(s_shape), s_region)
|
||||
# Detect a SwizzleLayout on the S side BEFORE _extract_tile strips it.
|
||||
# vec_len must fit inside one swizzle chunk (C = 2^per_element elements);
|
||||
# otherwise the vec ld/st crosses a swizzle XOR boundary and hits the
|
||||
# wrong physical bytes mid-vec.
|
||||
s_swizzle_chunk_elems = None
|
||||
if isinstance(s_layout, ComposeLayout):
|
||||
s_swizzle_chunk_elems = 1 << int(s_layout.swizzle.per_element)
|
||||
elif isinstance(s_layout, SwizzleLayout):
|
||||
s_swizzle_chunk_elems = 1 << int(s_layout.per_element)
|
||||
g = _extract_tile(g, g_region)
|
||||
s = _extract_tile(s, s_region)
|
||||
|
||||
# Only G drives the canonical form. S's iter order is derived from G's
|
||||
# pre-sort iter extents (used as the grouping shape) and then permuted
|
||||
# by G's stride-desc permutation. Independently sorting/canonicalizing
|
||||
# S would fuse iters whose strides happen to chain into one contiguous
|
||||
# range — that loses the layout's logical `(i, j) → addr` mapping for
|
||||
# layout-permuting copies (e.g. row-major GMEM → K-tiled SMEM, where
|
||||
# both layouts cover the same byte range but with different coord maps).
|
||||
g_pre_sort_extents = [int(it.extent) for it in g.shard]
|
||||
g_perm = sorted(range(len(g.shard)), key=lambda i: -int(g.shard[i].stride))
|
||||
try:
|
||||
s_grp1, seps1 = s.group(g_pre_sort_extents)
|
||||
except Exception as e:
|
||||
if debug:
|
||||
print(f" Step-2 S.group({g_pre_sort_extents}) failed: {e}")
|
||||
return _sort_by_stride_desc(g).canonicalize(), s, 1
|
||||
# S iters keep their original strides; only the *group order* is
|
||||
# rearranged to follow G's sorted iter order. Canonicalize after the
|
||||
# permute is safe — it only fuses iters whose strides genuinely chain
|
||||
# in the post-permute order, which preserves the logical structure.
|
||||
s_aligned = s_grp1.permute_by_groups(list(seps1), g_perm).canonicalize()
|
||||
g = _sort_by_stride_desc(g).canonicalize()
|
||||
if debug:
|
||||
print(f" g (sort+canon): shard={[(int(it.extent), int(it.stride)) for it in g.shard]}")
|
||||
print(
|
||||
f" s (grouped+permuted by G shape {g_pre_sort_extents}): shard="
|
||||
f"{[(int(it.extent), int(it.stride)) for it in s_aligned.shard]}"
|
||||
)
|
||||
print(f" thread_cnt: {thread_cnt}")
|
||||
|
||||
dummy_axis = g.shard[-1].axis if g.shard else None
|
||||
|
||||
for vec_len in _vec_len_candidates(elem_bits, vec_bits_candidates):
|
||||
if vec_len == 1:
|
||||
g_after_vec = [*list(g.shard), Iter(1, 1, dummy_axis)]
|
||||
else:
|
||||
# Swizzle chunk-size cap: vec must fit in one swizzle chunk.
|
||||
if s_swizzle_chunk_elems is not None and vec_len > s_swizzle_chunk_elems:
|
||||
if debug:
|
||||
print(
|
||||
f" vec_len={vec_len}: exceeds swizzle chunk "
|
||||
f"({s_swizzle_chunk_elems} elements)"
|
||||
)
|
||||
continue
|
||||
g_after_vec = _carve_tail(list(g.shard), vec_len)
|
||||
if g_after_vec is None:
|
||||
if debug:
|
||||
print(f" vec_len={vec_len}: G vec carve failed")
|
||||
continue
|
||||
outer_for_t = g_after_vec[:-1]
|
||||
outer_after_t = _carve_tail(outer_for_t, thread_cnt) if thread_cnt > 1 else outer_for_t
|
||||
if outer_after_t is None:
|
||||
if debug:
|
||||
print(f" vec_len={vec_len}: T={thread_cnt} carve failed")
|
||||
continue
|
||||
g_final_iters = [*outer_after_t, g_after_vec[-1]]
|
||||
new_shape = [int(it.extent) for it in g_final_iters]
|
||||
try:
|
||||
s_final_grp, _ = s_aligned.group(new_shape)
|
||||
except Exception as e:
|
||||
if debug:
|
||||
print(f" vec_len={vec_len}: S.group({new_shape}) failed: {e}")
|
||||
continue
|
||||
g_p = TileLayout.from_iters(g_final_iters, list(g.replica), dict(g.offset))
|
||||
s_p = s_final_grp
|
||||
|
||||
ok = _verify_s_tail_contig(s_p, vec_len)
|
||||
if debug:
|
||||
print(
|
||||
f" vec_len={vec_len}: shape={new_shape}, "
|
||||
f"g_p.shard={[(int(it.extent), int(it.stride)) for it in g_p.shard]}, "
|
||||
f"s_p.shard={[(int(it.extent), int(it.stride)) for it in s_p.shard]}, "
|
||||
f"s_tail_contig={ok}"
|
||||
)
|
||||
if not ok:
|
||||
continue
|
||||
# Alignment: per-thread starting addr = base_ptr + sizeof(elem) *
|
||||
# (region_base + tid*t_stride + outer_iter_strides). For the
|
||||
# vec_bits/8 byte vector op to be naturally aligned, every one of
|
||||
# those element-count terms must be a multiple of vec_len.
|
||||
align_terms = []
|
||||
for it in s_p.shard[:-1]:
|
||||
align_terms.append(int(it.stride))
|
||||
for it in g_p.shard[:-1]:
|
||||
align_terms.append(int(it.stride))
|
||||
align_terms.extend(s_p.offset.values())
|
||||
align_terms.extend(g_p.offset.values())
|
||||
if not _alignment_ok(vec_len, align_terms):
|
||||
if debug:
|
||||
print(f" vec_len={vec_len}: alignment check failed")
|
||||
continue
|
||||
return g_p, s_p, vec_len
|
||||
|
||||
return g, s_aligned, 1
|
||||
|
||||
|
||||
def _flat_outer_coords(outer_exts: list[int], flat_idx: int) -> list[int]:
|
||||
"""Decode a row-major flat index into per-iter coords. ``outer_exts`` is
|
||||
in stride-desc order (outermost first), so the first coord changes
|
||||
slowest."""
|
||||
coords: list[int] = []
|
||||
rem = flat_idx
|
||||
for ext in reversed(outer_exts):
|
||||
coords.append(rem % ext)
|
||||
rem //= ext
|
||||
coords.reverse()
|
||||
return coords
|
||||
|
||||
|
||||
def _outer_offsets(outer_iters_s, outer_iters_g, flat_idx):
|
||||
"""Returns ``(ds, dg)``: constant offsets on S and G sides for the
|
||||
given flat outer-loop iteration."""
|
||||
outer_exts = [int(it.extent) for it in outer_iters_s]
|
||||
coords = _flat_outer_coords(outer_exts, flat_idx)
|
||||
ds = sum(c * int(it.stride) for c, it in zip(coords, outer_iters_s))
|
||||
dg = sum(c * int(it.stride) for c, it in zip(coords, outer_iters_g))
|
||||
return ds, dg
|
||||
|
||||
|
||||
def copy_ptx_form(num_bytes: int) -> tuple[str, str]:
|
||||
"""Map copy width (bytes) to PTX ``(vec, ptx_type)`` for ``T.ptx.ld`` / ``T.ptx.st``."""
|
||||
return {
|
||||
16: ("v4", "u32"),
|
||||
8: ("v2", "u32"),
|
||||
4: ("", "u32"),
|
||||
2: ("", "u16"),
|
||||
1: ("", "u8"),
|
||||
}[num_bytes]
|
||||
|
||||
|
||||
def copy_ptx_ld_return_type(ptx_type: str) -> str:
|
||||
"""TVM dtype string for ``T.ptx.ld``'s ``return_type`` argument."""
|
||||
return {"u32": "uint32", "u16": "uint16", "u8": "uint32"}[ptx_type]
|
||||
@@ -0,0 +1,406 @@
|
||||
# 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.
|
||||
|
||||
"""Generic swizzle-aware iter pattern for CUDA copy dispatches.
|
||||
|
||||
When the per-thread outer-iter loop satisfies (C1)+(C2) below for a
|
||||
``SwizzleLayout(per_element=p, swizzle_len=sw, atom_len=at,
|
||||
swizzle_inner=True)`` on the SMEM side, the swizzled physical address
|
||||
at unrolled iter ``k`` reduces to
|
||||
|
||||
addr(k) = base_off + sum_{j : bit_j(k)=1} signed_strides[j]
|
||||
|
||||
where ``base_off`` and the ``signed_strides[j]`` are per-thread runtime
|
||||
constants set once at thread setup. Per-iter cost is then ``popcount(k)``
|
||||
register adds instead of a full ``swizzle.apply(...)`` per iter.
|
||||
|
||||
Notation. Each binary outer iter has element-stride ``2^(bj + p)`` for
|
||||
some chunk bit position ``bj >= 0`` (so ``stride / C = 2^bj`` where
|
||||
``C = 1 << p``). The chunk index ``q(M0) = M0 // C`` partitions into
|
||||
four bit ranges by where ``bj`` lands:
|
||||
|
||||
* ``[0, sw)`` — "inner" (Case 1.A in the proof)
|
||||
* ``[sw, at)`` — "mid" (Case 1.B)
|
||||
* ``[at, at + sw)`` — "outer" (Case 1.C; the bit overlaps the swizzle
|
||||
outer mask, so its addition produces a
|
||||
secondary contribution at ``bj - at``)
|
||||
* ``[at + sw, ∞)`` — "above" (Case 1.D)
|
||||
|
||||
Conditions for the linear-combination fast path:
|
||||
|
||||
(C1) bit-clear no-carry: ``bit_bj(q(M0)) = 0`` for every binary iter.
|
||||
(C2) support disjointness: no inner-outer pair ``(bj_A, bj_C)`` with
|
||||
``bj_C in [at, at+sw)`` and ``bj_A = bj_C - at`` both present.
|
||||
|
||||
(distinctness) The ``bj`` values across all binary iters must be
|
||||
distinct — two iters at the same ``bj`` collapse into bit
|
||||
``bj + 1`` whose case behavior may differ.
|
||||
|
||||
Under (C1)+(C2)+(distinctness), for each binary iter at position ``bj``:
|
||||
|
||||
T(bj) = 2^(bj + p) # element stride
|
||||
sigma_b(M0) = 1 - 2 * bit_b(q(M0)) # ∈ {+1, -1}
|
||||
|
||||
signed_strides[j] = sigma_(at + bj)(M0) * T(bj) bj in [0, sw)
|
||||
= T(bj) bj in [sw, at)
|
||||
= T(bj) + sigma_(bj - at)(M0) * T(bj - at)
|
||||
bj in [at, at + sw)
|
||||
= T(bj) bj >= at + sw
|
||||
|
||||
The ``swizzle_inner=False`` mode swaps the inner/outer roles and is not
|
||||
yet covered; ``try_recognize`` gates on this.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import tvm
|
||||
from tvm import arith
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx.expr import IntImm as _IntImm
|
||||
from tvm.tirx.layout import ComposeLayout, SwizzleLayout
|
||||
|
||||
|
||||
@dataclass
|
||||
class _BitIter:
|
||||
"""Pow2-extent outer iter, binary-split into ``n_bits`` chunk-bit flips.
|
||||
|
||||
``slot_start..slot_start + n_bits`` is this iter's range in the global
|
||||
``bit_positions`` / ``iter_strides_elems`` / ``signed_strides`` arrays.
|
||||
Slot ``slot_start + b`` corresponds to bit position ``n_bits - 1 - b``
|
||||
of this iter's per-iter coord (outermost binary bit first).
|
||||
"""
|
||||
|
||||
ext: int
|
||||
n_bits: int
|
||||
slot_start: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class _LinearIter:
|
||||
"""Outer iter contributing ``c * stride`` to the offset (no bit decomp).
|
||||
|
||||
Used when ``stride`` is a multiple of ``2^(p + at + sw)`` (pure Case 1.D
|
||||
regime: swizzle XOR has no effect on bits the iter flips). ``ext`` does
|
||||
not need to be a power of two.
|
||||
"""
|
||||
|
||||
ext: int
|
||||
stride: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class SwizzlePattern:
|
||||
"""A recognized swizzle iter pattern.
|
||||
|
||||
``bit_positions[j]`` and ``iter_strides_elems[j]`` collect the binary
|
||||
sub-iters from every BitIter in outer-iter order (outermost first).
|
||||
``outer_iters`` lists every outer iter (BitIter or LinearIter) in
|
||||
outermost-first order; ``emit_iter_offset`` walks this list to
|
||||
decompose ``mm`` per-iter. Empty lists = trivially recognized
|
||||
degenerate case (no outer iter, just base_off).
|
||||
"""
|
||||
|
||||
swizzle: SwizzleLayout
|
||||
bit_positions: list[int]
|
||||
iter_strides_elems: list[int]
|
||||
outer_iters: "list[_BitIter | _LinearIter]"
|
||||
|
||||
@property
|
||||
def n_binary_iters(self) -> int:
|
||||
return len(self.bit_positions)
|
||||
|
||||
|
||||
def get_swizzle(layout) -> SwizzleLayout | None:
|
||||
"""Return the SwizzleLayout from ``layout`` if present, else ``None``.
|
||||
|
||||
Accepts ``ComposeLayout(SwizzleLayout, TileLayout)`` (the common case
|
||||
when a TileLayout is wrapped by a swizzle), or a bare ``SwizzleLayout``.
|
||||
"""
|
||||
if isinstance(layout, ComposeLayout):
|
||||
return layout.swizzle
|
||||
if isinstance(layout, SwizzleLayout):
|
||||
return layout
|
||||
return None
|
||||
|
||||
|
||||
def _is_pow2(n: int) -> bool:
|
||||
return n > 0 and (n & (n - 1)) == 0
|
||||
|
||||
|
||||
def try_recognize(
|
||||
swizzle: SwizzleLayout,
|
||||
iter_extents: list[int],
|
||||
iter_strides: list[int],
|
||||
s_off_template,
|
||||
var_bounds: dict | None = None,
|
||||
) -> SwizzlePattern | None:
|
||||
"""Return a ``SwizzlePattern`` if (C1)+(C2)+(distinctness) hold, else ``None``.
|
||||
|
||||
``iter_extents`` / ``iter_strides``: the outer-iter list on the S side
|
||||
(excluding T iter and vec iter), in outermost-first order matching
|
||||
``s_p.shard[:-2]`` (or the atom-derived analog in ``reg.py``).
|
||||
Strides are in element units.
|
||||
|
||||
Each outer iter with ``extent=2^k`` and ``stride=s`` is conceptually
|
||||
split into ``k`` binary iters of strides ``2^(k-1)*s, ..., 2*s, s``
|
||||
(outermost first within the split — this matches ``_flat_outer_coords``
|
||||
semantics, since the highest-stride iter must change slowest in the
|
||||
flat-index decomposition).
|
||||
|
||||
``s_off_template`` is the per-thread linear base offset expression
|
||||
(with a placeholder var for the thread-id contribution). It is used
|
||||
only to check condition (C1) symbolically via ``arith.Analyzer``;
|
||||
``emit_init`` takes the resolved form separately.
|
||||
|
||||
``var_bounds`` is an optional ``{Var: tvm.ir.Range}`` map of placeholder
|
||||
bounds to ``analyzer.bind`` before the (C1) check. Without bounds,
|
||||
structurally-OK forms like ``(lane // 8) * 8 + (lane % 8) * Q`` where
|
||||
``lane < 32`` make ``(... // (C·2^bj)) % 2 == 0`` unprovable — the
|
||||
bit is in fact always 0 but the analyzer can't conclude it universally.
|
||||
Pass ``{lane_ph: Range(0, 32), warp_ph: Range(0, n_warps)}`` (or the
|
||||
scope's equivalents) to let the (C1) check fire on these templates.
|
||||
"""
|
||||
# swizzle_inner=False swaps the inner/outer xor direction — Cases 1.A
|
||||
# and 1.C roles flip. Not derived/tested yet; reject for safety.
|
||||
if not swizzle.swizzle_inner:
|
||||
return None
|
||||
|
||||
p = swizzle.per_element
|
||||
sw = swizzle.swizzle_len
|
||||
at = swizzle.atom_len
|
||||
C = 1 << p
|
||||
# Pure Case 1.D threshold: stride a multiple of this means every chunk-bit
|
||||
# the iter flips is at position >= at + sw (above the swizzle XOR region),
|
||||
# so swizzle has no effect and the contribution is purely linear in the
|
||||
# iter coord — no power-of-2 ext requirement.
|
||||
pure_1d = 1 << (p + at + sw)
|
||||
|
||||
bit_positions: list[int] = []
|
||||
iter_strides_elems: list[int] = []
|
||||
outer_iters: list = []
|
||||
|
||||
for ext, stride in zip(iter_extents, iter_strides):
|
||||
# Zero-stride iters degrade dq=0 → log2 undefined. Explicit guard.
|
||||
if stride == 0 or stride % C != 0:
|
||||
return None
|
||||
if ext <= 0:
|
||||
return None
|
||||
if ext == 1:
|
||||
# Trivial iter contributes nothing; skip without forcing pow2.
|
||||
continue
|
||||
if not _is_pow2(ext):
|
||||
# Non-pow2 ext can only be handled by the linear path. That in turn
|
||||
# requires the iter to be in pure Case 1.D (stride a multiple of
|
||||
# the swizzle period) so the swizzle does not interact with the
|
||||
# per-coord contribution.
|
||||
if stride % pure_1d != 0:
|
||||
return None
|
||||
outer_iters.append(_LinearIter(ext=ext, stride=stride))
|
||||
continue
|
||||
# pow2 ext: binary split (existing path).
|
||||
k = ext.bit_length() - 1 # log2(ext)
|
||||
slot_start = len(bit_positions)
|
||||
# Split into k binary iters; the outermost (within this split) carries
|
||||
# the largest stride so that flat-index bit decomp matches our
|
||||
# outer-iter list ordering.
|
||||
for j in range(k - 1, -1, -1):
|
||||
substride = stride * (1 << j)
|
||||
dq = substride // C
|
||||
# dq must be a single bit set (so this binary iter flips exactly
|
||||
# one bit of the chunk index). _is_pow2 also rejects dq=0.
|
||||
if not _is_pow2(dq):
|
||||
return None
|
||||
bj = dq.bit_length() - 1
|
||||
# All bj >= 0 accepted; case branching happens in emit_init.
|
||||
bit_positions.append(bj)
|
||||
iter_strides_elems.append(substride)
|
||||
outer_iters.append(_BitIter(ext=ext, n_bits=k, slot_start=slot_start))
|
||||
|
||||
# Distinctness: two binary iters at the same bj collapse to bj+1, whose
|
||||
# case behavior may differ from bj. See module docstring NB.
|
||||
if len(set(bit_positions)) != len(bit_positions):
|
||||
return None
|
||||
|
||||
bj_set = set(bit_positions)
|
||||
|
||||
# (C2) support disjointness: the only possible collision is between a
|
||||
# Case-1.A iter at bj_A and a Case-1.C iter at bj_A + at. Checking the
|
||||
# 1.C direction alone is symmetric and complete.
|
||||
for bj in bj_set:
|
||||
if at <= bj < at + sw and (bj - at) in bj_set:
|
||||
return None # inner-outer pair collision
|
||||
|
||||
# (C1) per-iter no-carry on q(M0). Must hold *symbolically over all*
|
||||
# free lane / warp placeholders in s_off_template — ``can_prove_equal``
|
||||
# returns False if the analyzer can't discharge the equality
|
||||
# universally, conservatively forcing a fallback.
|
||||
analyzer = arith.Analyzer()
|
||||
if var_bounds:
|
||||
for var, rng in var_bounds.items():
|
||||
analyzer.bind(var, rng)
|
||||
for bj in bj_set:
|
||||
divisor = C * (1 << bj)
|
||||
check = tvm.tirx.floormod(
|
||||
tvm.tirx.floordiv(s_off_template, _IntImm("int32", divisor)),
|
||||
_IntImm("int32", 2),
|
||||
)
|
||||
if not analyzer.can_prove_equal(check, _IntImm("int32", 0)):
|
||||
return None
|
||||
|
||||
return SwizzlePattern(
|
||||
swizzle=swizzle,
|
||||
bit_positions=bit_positions,
|
||||
iter_strides_elems=iter_strides_elems,
|
||||
outer_iters=outer_iters,
|
||||
)
|
||||
|
||||
|
||||
def emit_init(pattern: SwizzlePattern, s_off_resolved):
|
||||
"""Emit at thread setup (call from inside the @T.prim_func body):
|
||||
|
||||
1. ``base_off = swizzle.apply(s_off_resolved)`` — runtime, per-thread,
|
||||
computed once.
|
||||
2. ``signed_strides[j]`` for each binary iter j, written into a local
|
||||
buffer using the sigma formula above.
|
||||
|
||||
Returns ``(signed_strides_buffer_or_None, base_off_primexpr)``. The
|
||||
buffer is ``None`` when ``pattern.n_binary_iters == 0`` (no outer
|
||||
iter, no signed_strides needed).
|
||||
|
||||
``s_off_resolved`` is the per-thread offset with the real tid Var
|
||||
substituted in (not the placeholder).
|
||||
"""
|
||||
swizzle = pattern.swizzle
|
||||
p = swizzle.per_element
|
||||
sw = swizzle.swizzle_len
|
||||
at = swizzle.atom_len
|
||||
C = 1 << p
|
||||
|
||||
base_off = swizzle.apply(s_off_resolved)["m"]
|
||||
|
||||
n = pattern.n_binary_iters
|
||||
if n == 0:
|
||||
return None, base_off
|
||||
|
||||
signed_strides = T.alloc_buffer([n], "int32", scope="local")
|
||||
q = tvm.tirx.floordiv(s_off_resolved, C)
|
||||
|
||||
def _sigma_bit(bit_pos: int):
|
||||
# 1 - 2 * bit_(bit_pos)(q); ∈ {+1, -1}.
|
||||
row_bit = tvm.tirx.bitwise_and(
|
||||
tvm.tirx.shift_right(q, _IntImm("int32", bit_pos)),
|
||||
_IntImm("int32", 1),
|
||||
)
|
||||
return _IntImm("int32", 1) - row_bit * _IntImm("int32", 2)
|
||||
|
||||
for j, (bj, stride) in enumerate(zip(pattern.bit_positions, pattern.iter_strides_elems)):
|
||||
stride_pow = stride # = 2^(bj + p) elements
|
||||
if 0 <= bj < sw:
|
||||
# Case 1.A (inner): signed_stride = sigma_(at + bj) · T.
|
||||
value = _sigma_bit(at + bj) * _IntImm("int32", stride_pow)
|
||||
elif sw <= bj < at:
|
||||
# Case 1.B (mid): signed_stride = +T.
|
||||
value = _IntImm("int32", stride_pow)
|
||||
elif at <= bj < at + sw:
|
||||
# Case 1.C (outer): signed_stride = T + sigma_(bj - at) · T_sec.
|
||||
# Invariant: bj >= at, so T_sec = T >> at = 2^(bj - at + p)
|
||||
# = T(bj - at) is well-defined (no underflow).
|
||||
stride_sec = stride_pow >> at
|
||||
value = _IntImm("int32", stride_pow) + _sigma_bit(bj - at) * _IntImm(
|
||||
"int32", stride_sec
|
||||
)
|
||||
else: # bj >= at + sw, Case 1.D (above)
|
||||
# No swizzle effect at this bit; signed_stride = +T.
|
||||
value = _IntImm("int32", stride_pow)
|
||||
# NB: Buffer.__setitem__ syntax (``signed_strides[j] = value``) is
|
||||
# intercepted by the TIRx script parser but not by raw Python when
|
||||
# this function is called from outside an @T.inline body. Use the
|
||||
# low-level buffer_store builder instead.
|
||||
T.buffer_store(signed_strides, value, [_IntImm("int32", j)])
|
||||
|
||||
return signed_strides, base_off
|
||||
|
||||
|
||||
def emit_iter_offset(pattern: SwizzlePattern, signed_strides, base_off, k):
|
||||
"""Compute the per-mm physical S offset = ``base_off`` + sum of per-iter
|
||||
contributions.
|
||||
|
||||
``k`` is the flat outer iter index ∈ ``[0, prod(it.ext for it in outer_iters))``.
|
||||
Decomposed innermost-first across ``pattern.outer_iters`` into per-iter
|
||||
coords ``c_i``. Each iter contributes:
|
||||
|
||||
* ``_BitIter``: ``sum_b bit_(n_bits-1-b)(c_i) * signed_strides[slot_start + b]``,
|
||||
i.e. each binary bit of ``c_i`` selects its precomputed sigma-stride.
|
||||
The slot order (outermost-first within the iter) means the highest
|
||||
bit of ``c_i`` indexes the slot at ``slot_start``.
|
||||
* ``_LinearIter``: ``c_i * stride`` (no bit decomposition; used when
|
||||
``stride`` is a multiple of ``2^(p + at + sw)`` so swizzle has no
|
||||
XOR effect and ``ext`` need not be pow2).
|
||||
|
||||
Two paths per iter:
|
||||
* Python int ``k`` — coords and bits known at parse time; emits only
|
||||
the necessary adds, no runtime shift/mask.
|
||||
* TIRx Var ``k`` — emits floormod/floordiv + bit-and/shift; relies on
|
||||
downstream unroll + constant-fold.
|
||||
"""
|
||||
if not pattern.outer_iters:
|
||||
return base_off
|
||||
|
||||
off = base_off
|
||||
remaining = k
|
||||
is_const = isinstance(k, int)
|
||||
for it in reversed(pattern.outer_iters): # innermost first
|
||||
ext = it.ext
|
||||
if is_const:
|
||||
c = remaining % ext
|
||||
remaining = remaining // ext
|
||||
else:
|
||||
c = tvm.tirx.floormod(remaining, _IntImm("int32", ext))
|
||||
remaining = tvm.tirx.floordiv(remaining, _IntImm("int32", ext))
|
||||
if isinstance(it, _LinearIter):
|
||||
if is_const:
|
||||
if c != 0:
|
||||
off = off + c * it.stride
|
||||
else:
|
||||
off = off + c * _IntImm("int32", it.stride)
|
||||
continue
|
||||
# _BitIter
|
||||
for b in range(it.n_bits):
|
||||
bit_pos = it.n_bits - 1 - b
|
||||
slot = it.slot_start + b
|
||||
if is_const:
|
||||
if (c >> bit_pos) & 1:
|
||||
off = off + signed_strides[slot]
|
||||
else:
|
||||
bit = tvm.tirx.bitwise_and(
|
||||
tvm.tirx.shift_right(c, _IntImm("int32", bit_pos)),
|
||||
_IntImm("int32", 1),
|
||||
)
|
||||
off = off + bit * signed_strides[slot]
|
||||
return off
|
||||
|
||||
|
||||
def emit_fallback_offset(swizzle: SwizzleLayout, s_off_resolved, ds_k):
|
||||
"""Slow but always-correct path: full ``swizzle.apply(s_off + ds_k)``
|
||||
per iter. Use when ``try_recognize`` returns ``None``.
|
||||
|
||||
``ds_k`` is the outer-iter delta for unrolled iter k — typically a
|
||||
Expr (a function of the unroll var that simplifies to a constant
|
||||
after unrolling) or a Python int. ``s_off_resolved`` is the per-thread
|
||||
base linear offset with the real tid Var substituted.
|
||||
"""
|
||||
return swizzle.apply(s_off_resolved + ds_k)["m"]
|
||||
@@ -0,0 +1,116 @@
|
||||
# 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.
|
||||
|
||||
"""Scalar single-thread copy fallback (priority=0)."""
|
||||
|
||||
import warnings
|
||||
|
||||
import tvm
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import Buffer, PrimFunc
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import (
|
||||
predicate,
|
||||
register_dispatch,
|
||||
)
|
||||
from tvm.tirx.operator.tile_primitive.registry import DispatchContext
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ._common import _TID_AXIS_FOR_SCOPE
|
||||
from .reg import _axis_decl
|
||||
from .utils import _is_valid_copy
|
||||
|
||||
|
||||
def _region_st_extent(buffer_region):
|
||||
region = buffer_region.region
|
||||
return [r.min for r in region], [r.extent for r in region]
|
||||
|
||||
|
||||
def _emit_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
src: Buffer = op_call.src.buffer
|
||||
dst: Buffer = op_call.dst.buffer
|
||||
src_st, src_extent = _region_st_extent(op_call.src)
|
||||
dst_st, dst_extent = _region_st_extent(op_call.dst)
|
||||
|
||||
warnings.warn(
|
||||
f"copy/fallback (scalar single-thread) picked for {src.scope()} -> "
|
||||
f"{dst.scope()} at scope_kind={sctx.scope_kind}; all faster variants "
|
||||
f"rejected.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
def _copy_body(dst_buf, src_buf):
|
||||
dst_indices = [i for i in range(len(dst_buf.shape)) if dst_extent[i] != 1]
|
||||
src_indices = [i for i in range(len(src_buf.shape)) if src_extent[i] != 1]
|
||||
assert len(dst_indices) == len(src_indices)
|
||||
copy_extents = [dst_extent[i] for i in dst_indices]
|
||||
|
||||
def _dst_coord(lvs):
|
||||
if isinstance(lvs, tvm.tirx.Var):
|
||||
lvs = [lvs]
|
||||
coord = list(dst_st)
|
||||
for k, lv in enumerate(lvs):
|
||||
coord[dst_indices[k]] += lv
|
||||
return coord
|
||||
|
||||
def _src_coord(lvs):
|
||||
if isinstance(lvs, tvm.tirx.Var):
|
||||
lvs = [lvs]
|
||||
coord = list(src_st)
|
||||
for k, lv in enumerate(lvs):
|
||||
coord[src_indices[k]] += lv
|
||||
return coord
|
||||
|
||||
with T.grid(*copy_extents) as lvs:
|
||||
T.buffer_store(dst_buf, src_buf[tuple(_src_coord(lvs))], _dst_coord(lvs))
|
||||
|
||||
scope_kind = sctx.scope_kind
|
||||
|
||||
if scope_kind == "thread":
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
_copy_body(dst, src)
|
||||
|
||||
return impl
|
||||
|
||||
tid_axis_name = _TID_AXIS_FOR_SCOPE[scope_kind]
|
||||
# first-active tid = composition of per-axis offsets (radix-32, since a warp is 32 lanes)
|
||||
first_tid = int(sctx.intra["laneid"][1])
|
||||
if scope_kind == "warpgroup":
|
||||
first_tid += 32 * int(sctx.intra["wid_in_wg"][1])
|
||||
elif scope_kind == "cta":
|
||||
first_tid += 32 * int(sctx.intra["warpid"][1])
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
tid = _axis_decl(tid_axis_name, sctx)
|
||||
if tid == first_tid:
|
||||
_copy_body(dst, src)
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
@register_dispatch(
|
||||
"copy",
|
||||
"cuda",
|
||||
variant="fallback",
|
||||
priority=0,
|
||||
when=[predicate("validate_copy_op", _is_valid_copy)],
|
||||
)
|
||||
def copy_schedule_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
return _emit_fallback(op_call, sctx)
|
||||
@@ -0,0 +1,328 @@
|
||||
# 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 dispatch for ``global ↔ shared`` (no register side).
|
||||
|
||||
There's no per-thread register side to inherit a partition from — both sides
|
||||
are cross-thread storage. The partition is synthesized from the surrounding
|
||||
scope context (warp / warpgroup / cta / thread): ``thread_cnt`` is derived
|
||||
from ``sctx.intra`` and each thread takes ``n_elements / thread_cnt``
|
||||
consecutive fused-index slots. Layout / partition algorithm lives in
|
||||
``_common.py`` and is shared with ``ldgsts.py``.
|
||||
"""
|
||||
|
||||
import tvm
|
||||
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.operator.tile_primitive.dispatcher import (
|
||||
predicate,
|
||||
register_dispatch,
|
||||
)
|
||||
from tvm.tirx.operator.tile_primitive.registry import DispatchContext
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ._common import (
|
||||
_TID_AXIS_FOR_SCOPE,
|
||||
_thread_cnt,
|
||||
align_layouts_gs,
|
||||
copy_ptx_form,
|
||||
copy_ptx_ld_return_type,
|
||||
)
|
||||
from ._swizzle_iter import (
|
||||
emit_init,
|
||||
emit_iter_offset,
|
||||
get_swizzle,
|
||||
try_recognize,
|
||||
)
|
||||
from .reg import _all_threads_active, _axis_decl, _ptr_off
|
||||
from .utils import _is_valid_copy, _scope_allowed
|
||||
|
||||
_GMEM_SMEM_PAIRS = [
|
||||
("global", "shared*"),
|
||||
("shared*", "global"),
|
||||
]
|
||||
|
||||
|
||||
def _divides_thread_cnt(
|
||||
op_call: TilePrimitiveCall, sctx: DispatchContext
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Reject copies whose region element count does not divide ``thread_cnt``.
|
||||
|
||||
Without this guard the emit's ``[outer, T, vec]`` partition has no
|
||||
integer solution: either every thread gets fractional work, or
|
||||
``thread_cnt=0`` (degenerate scope) hits a modulo-by-zero. Both cases
|
||||
indicate a poorly-shaped copy (e.g. 1024-thread CTA writing a 64-elem
|
||||
tail) that this dispatch refuses to paper over with a slow scalar emit.
|
||||
"""
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
thread_cnt = _thread_cnt(sctx)
|
||||
if thread_cnt <= 0:
|
||||
return False, f"degenerate thread_cnt={thread_cnt} (scope has empty intra)"
|
||||
g_br = op_call.src if op_call.src.buffer.scope() == "global" else op_call.dst
|
||||
n_elements = 1
|
||||
for r in g_br.region:
|
||||
ext = r.extent
|
||||
try:
|
||||
n_elements *= int(ext)
|
||||
except (TypeError, ValueError):
|
||||
return False, f"non-constant region extent {ext}"
|
||||
if n_elements % thread_cnt != 0:
|
||||
return False, (f"region size {n_elements} not divisible by thread_cnt={thread_cnt}")
|
||||
return True, None
|
||||
|
||||
|
||||
def _is_gmem_smem(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=_GMEM_SMEM_PAIRS),
|
||||
lambda: _divides_thread_cnt(op_call, sctx),
|
||||
):
|
||||
ok, msg = check()
|
||||
if not ok:
|
||||
return False, msg
|
||||
return True, None
|
||||
|
||||
|
||||
def _emit_gmem_smem(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() == "global":
|
||||
g_buf, g_br, s_buf, s_br = src, op_call.src, dst, op_call.dst
|
||||
g_is_src = True
|
||||
else:
|
||||
g_buf, g_br, s_buf, s_br = dst, op_call.dst, src, op_call.src
|
||||
g_is_src = False
|
||||
|
||||
g_region = [(r.min, r.min + r.extent) for r in g_br.region]
|
||||
s_region = [(r.min, r.min + r.extent) for r in s_br.region]
|
||||
|
||||
elem_bits = DataType(src.dtype).bits
|
||||
thread_cnt = _thread_cnt(sctx)
|
||||
|
||||
with sctx.target:
|
||||
g_p, s_p, vec_len = align_layouts_gs(
|
||||
g_buf.layout,
|
||||
g_buf.shape,
|
||||
g_region,
|
||||
s_buf.layout,
|
||||
s_buf.shape,
|
||||
s_region,
|
||||
elem_bits,
|
||||
thread_cnt,
|
||||
)
|
||||
|
||||
# vec_len=1 is the scalar fallback — uses the same unified
|
||||
# [outer x thread x vec] coord scheme below.
|
||||
|
||||
vec_bits = vec_len * elem_bits
|
||||
num_bytes = vec_bits // 8
|
||||
vec, ptx_type = copy_ptx_form(num_bytes)
|
||||
|
||||
# Partition guarantees ``prod(s_p.shard.extents) == prod(g_p.shard.extents)
|
||||
# == n_elements`` (the total transfer count). Express the per-thread
|
||||
# per-round address as a 3D coord ``(f, tid, 0)`` against shape
|
||||
# ``[total_outer, thread_cnt, vec_len]``, and let ``layout.apply`` flatten
|
||||
# it through whatever multi-iter T / outer-iter structure ``align_layouts_gs``
|
||||
# picked. This makes the emit oblivious to how many iters the partition
|
||||
# split T or outer across.
|
||||
n_elements = 1
|
||||
for it in s_p.shard:
|
||||
n_elements *= int(it.extent)
|
||||
assert n_elements % (thread_cnt * vec_len) == 0, (
|
||||
f"partition produced {n_elements} elements but thread_cnt({thread_cnt}) * "
|
||||
f"vec_len({vec_len}) = {thread_cnt * vec_len} doesn't divide it"
|
||||
)
|
||||
total_outer = n_elements // (thread_cnt * vec_len)
|
||||
apply_shape = [
|
||||
_IntImm("int32", total_outer),
|
||||
_IntImm("int32", thread_cnt),
|
||||
_IntImm("int32", vec_len),
|
||||
]
|
||||
|
||||
s_zero = [0] * len(s_buf.shape)
|
||||
g_zero = [0] * len(g_buf.shape)
|
||||
|
||||
tid_axis_name = _TID_AXIS_FOR_SCOPE[sctx.scope_kind] if thread_cnt > 1 else None
|
||||
|
||||
# Walk shard from the vec iter backward to find the prefix that covers
|
||||
# the T region exactly (∏ext == thread_cnt). The iters consumed are T
|
||||
# iters; the leading prefix is the outer iter list — handed to
|
||||
# ``try_recognize`` so the swizzle fast path can decide whether the
|
||||
# outer iter strides match a pattern it can lower to signed_strides.
|
||||
if thread_cnt > 1:
|
||||
acc, _i = 1, len(s_p.shard) - 2
|
||||
while _i >= 0 and acc < thread_cnt:
|
||||
_ext = int(s_p.shard[_i].extent)
|
||||
if acc * _ext > thread_cnt:
|
||||
break
|
||||
acc *= _ext
|
||||
_i -= 1
|
||||
outer_iters_s = list(s_p.shard[: _i + 1]) if acc == thread_cnt else []
|
||||
else:
|
||||
outer_iters_s = list(s_p.shard[:-1])
|
||||
|
||||
# SwizzleLayout on s_buf: try the closed-form signed-strides pattern
|
||||
# (precomputed once per thread, then per-iter is a sum of register
|
||||
# adds); fall back to per-iter ``swizzle.apply`` (one full XOR +
|
||||
# decompose per iter). Closure picked at parse time so the TIRx parser
|
||||
# doesn't AST-evaluate a "dead" ternary branch.
|
||||
swizzle = get_swizzle(s_buf.layout)
|
||||
swizzle_pattern = None
|
||||
if swizzle is not None and outer_iters_s:
|
||||
if tid_axis_name is not None:
|
||||
_tid_placeholder = _TirVar(tid_axis_name, "int32")
|
||||
else:
|
||||
_tid_placeholder = _IntImm("int32", 0)
|
||||
s_off_template = s_p.apply(
|
||||
_IntImm("int32", 0),
|
||||
_tid_placeholder,
|
||||
_IntImm("int32", 0),
|
||||
shape=apply_shape,
|
||||
)["m"]
|
||||
# Bind the tid placeholder's range so the (C1) analyzer check can
|
||||
# discharge ``bit_bj(s_off // C) == 0`` for high bj's. Outer iter
|
||||
# stride here is ``thread_cnt * vec_len`` ⇒ bj ∈ [log2(thread_cnt),
|
||||
# ...]; without bounds the analyzer can't prove the lane's high bits
|
||||
# are 0 and rejects.
|
||||
var_bounds = {}
|
||||
if tid_axis_name is not None:
|
||||
var_bounds[_tid_placeholder] = tvm.ir.Range.from_min_extent(0, thread_cnt)
|
||||
swizzle_pattern = try_recognize(
|
||||
swizzle,
|
||||
[int(it.extent) for it in outer_iters_s],
|
||||
[int(it.stride) for it in outer_iters_s],
|
||||
s_off_template,
|
||||
var_bounds=var_bounds or None,
|
||||
)
|
||||
|
||||
class _SwizzleState:
|
||||
def __init__(self):
|
||||
self.signed_strides = None
|
||||
self.base_off = None
|
||||
|
||||
state = _SwizzleState()
|
||||
|
||||
def _decl_tid():
|
||||
if tid_axis_name is not None:
|
||||
return _axis_decl(tid_axis_name, sctx)
|
||||
return _IntImm("int32", 0)
|
||||
|
||||
def _setup_swizzle(tid):
|
||||
if swizzle_pattern is None:
|
||||
return
|
||||
s_off_resolved = s_p.apply(
|
||||
_IntImm("int32", 0),
|
||||
tid,
|
||||
_IntImm("int32", 0),
|
||||
shape=apply_shape,
|
||||
)["m"]
|
||||
state.signed_strides, state.base_off = emit_init(
|
||||
swizzle_pattern,
|
||||
s_off_resolved,
|
||||
)
|
||||
|
||||
if swizzle_pattern is not None:
|
||||
|
||||
def _s_off(f, s_lin):
|
||||
return emit_iter_offset(
|
||||
swizzle_pattern,
|
||||
state.signed_strides,
|
||||
state.base_off,
|
||||
f,
|
||||
)
|
||||
elif swizzle is not None:
|
||||
_sw = swizzle
|
||||
|
||||
def _s_off(f, s_lin):
|
||||
return _sw.apply(s_lin)["m"]
|
||||
else:
|
||||
|
||||
def _s_off(f, s_lin):
|
||||
return s_lin
|
||||
|
||||
v0 = _IntImm("int32", 0)
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
tid = _decl_tid()
|
||||
_setup_swizzle(tid)
|
||||
tmp = T.alloc_local((vec_len,), src.dtype)
|
||||
tmp_ptr = tmp.ptr_to([0])
|
||||
# NB: pass typed ptr_to(...) directly to _ptr_off; caching in a
|
||||
# local var turns it into void* + offset = byte arithmetic →
|
||||
# misaligned vector ops.
|
||||
#
|
||||
# Use a serial TIR loop and let ptxas unroll downstream. Mirrors
|
||||
# the reg.py rationale in commit ac7ecf70f0: explicit ``T.unroll``
|
||||
# materializes the per-iter scratch (s_lin/g_lin/s_off/s_ptr/g_ptr)
|
||||
# as N copies of each ``alignas(64)`` declaration. For large
|
||||
# ``total_outer`` (e.g. thread-scope fp32 swizzled copies of 32x256
|
||||
# at vec=4 ⇒ 2048 iters; ldgsts test4 ⇒ ~4k iters once both
|
||||
# g2s/s2g sites add up) this floods the kernel and nvcc times out.
|
||||
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_off = _s_off(f, s_lin)
|
||||
s_ptr = _ptr_off(s_buf.ptr_to(s_zero), s_off)
|
||||
g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin)
|
||||
if g_is_src:
|
||||
T.ptx.ld(
|
||||
g_ptr,
|
||||
copy_ptx_ld_return_type(ptx_type),
|
||||
ptx_type,
|
||||
dst=tmp_ptr,
|
||||
space="global",
|
||||
vec=vec,
|
||||
)
|
||||
T.ptx.st(
|
||||
s_ptr, src=tmp_ptr, space="shared", vec=vec, ptx_type=ptx_type
|
||||
)
|
||||
else:
|
||||
T.ptx.ld(
|
||||
s_ptr,
|
||||
copy_ptx_ld_return_type(ptx_type),
|
||||
ptx_type,
|
||||
dst=tmp_ptr,
|
||||
space="shared",
|
||||
vec=vec,
|
||||
)
|
||||
T.ptx.st(
|
||||
g_ptr, src=tmp_ptr, space="global", vec=vec, ptx_type=ptx_type
|
||||
)
|
||||
# fmt: on
|
||||
return impl
|
||||
|
||||
|
||||
@register_dispatch(
|
||||
"copy",
|
||||
"cuda",
|
||||
variant="gmem_smem",
|
||||
priority=10,
|
||||
when=[predicate("gmem_smem_applicable", _is_gmem_smem)],
|
||||
)
|
||||
def copy_schedule_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
return _emit_gmem_smem(op_call, sctx)
|
||||
@@ -0,0 +1,450 @@
|
||||
# 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 dispatch variant: ldmatrix / stmatrix (TBD algorithm).
|
||||
|
||||
Handles register ↔ shared copies on CUDA via PTX ``ldmatrix`` / ``stmatrix``.
|
||||
Direction (ld vs st) and exec scope (warp / warpgroup) are decided inside
|
||||
``_emit`` from the src/dst scopes and ``sctx.scope_kind``.
|
||||
"""
|
||||
|
||||
from math import prod
|
||||
|
||||
import tvm
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import PrimFunc
|
||||
from tvm.tirx import Var as _TirVar
|
||||
from tvm.tirx.expr import IntImm as _IntImm
|
||||
from tvm.tirx.layout import S, TileLayout
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import fail, predicate, register_dispatch
|
||||
from tvm.tirx.operator.tile_primitive.registry import DispatchContext
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ._common import ( # noqa: F401 (_carve_tail reserved for future variants)
|
||||
_carve_tail,
|
||||
_extract_tile,
|
||||
)
|
||||
from ._swizzle_iter import emit_init, emit_iter_offset, get_swizzle, try_recognize
|
||||
from .reg import _all_threads_active, _ptr_off
|
||||
from .utils import _is_valid_copy, _scope_allowed
|
||||
|
||||
_REG_SMEM_PAIRS = [
|
||||
("local", "shared*"),
|
||||
("shared*", "local"),
|
||||
]
|
||||
|
||||
_VALID_R_LANE_AXES = {"laneid", "tid_in_wg", "tx"}
|
||||
|
||||
|
||||
def _compute_r_perm(r):
|
||||
"""Permutation: thread iters first (stride-desc), then memory iters (stride-desc)."""
|
||||
|
||||
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 _is_ldstmatrix(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 ("warp", "warpgroup", "cta"):
|
||||
return False, f"unsupported exec_scope {sctx.scope_kind} (need warp, warpgroup, or cta)"
|
||||
for check in (
|
||||
lambda: _all_threads_active(sctx),
|
||||
lambda: _is_valid_copy(op_call, sctx),
|
||||
lambda: _scope_allowed(op_call, sctx, allowed_pairs=_REG_SMEM_PAIRS),
|
||||
):
|
||||
ok, msg = check()
|
||||
if not ok:
|
||||
return False, msg
|
||||
return True, None
|
||||
|
||||
|
||||
def _emit(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
|
||||
# Step 1: identify reg / smem sides and pull their tensor shape + layout.
|
||||
src_br = op_call.src
|
||||
dst_br = op_call.dst
|
||||
if src_br.buffer.scope() == "local":
|
||||
r_br, s_br = src_br, dst_br
|
||||
direction = "st" # reg -> smem (stmatrix)
|
||||
else:
|
||||
r_br, s_br = dst_br, src_br
|
||||
direction = "ld" # smem -> reg (ldmatrix)
|
||||
r_buf = r_br.buffer
|
||||
s_buf = s_br.buffer
|
||||
r_shape = list(r_buf.shape)
|
||||
r_layout = r_buf.layout
|
||||
s_shape = list(s_buf.shape)
|
||||
s_layout = s_buf.layout
|
||||
|
||||
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_layout.slice(r_shape, r_region)
|
||||
s_sliced = s_layout.slice(s_shape, s_region)
|
||||
r = r_sliced.canonicalize()
|
||||
s = s_sliced.canonicalize()
|
||||
|
||||
# Step 2.5: peel any S-side swizzle wrapper to expose the underlying
|
||||
# TileLayout. The ComposeLayout doesn't have ``.replica`` / ``.shard``,
|
||||
# so we must peel *before* the structural checks below. Capture the
|
||||
# swizzle separately for use at emit time.
|
||||
# NB: read swizzle from the *buffer* layout, not the post-canon ``s``.
|
||||
# When the underlying tile is trivial, ``ComposeLayoutNode::Canonicalize``
|
||||
# returns a bare ``SwizzleLayout``; isinstance(s, ComposeLayout) is then
|
||||
# False and we'd miss the swizzle here.
|
||||
s_swizzle = get_swizzle(s_buf.layout)
|
||||
if s_swizzle is not None and s_swizzle.per_element < 3:
|
||||
# ldmatrix/stmatrix .b16 reads/writes 8 fp16 = 128b per lane in one
|
||||
# contiguous chunk. The swizzle preserves the lowest ``per_element``
|
||||
# bits of the address (in-chunk offset). For the per-lane 128b unit
|
||||
# to stay contiguous post-swizzle, ``2^per_element >= 8`` ⇒ p >= 3.
|
||||
fail(
|
||||
f"swizzle per_element={s_swizzle.per_element} < 3 incompatible "
|
||||
f"with .b16 ldmatrix/stmatrix (need 8-fp16 chunk integrity)"
|
||||
)
|
||||
s = _extract_tile(s, s_region)
|
||||
|
||||
# Step 3: ldstmatrix doesn't broadcast — require zero replica on both sides.
|
||||
if len(r.replica) != 0:
|
||||
fail(f"R layout has replica {list(r.replica)}; ldstmatrix requires no replica")
|
||||
if len(s.replica) != 0:
|
||||
fail(f"S layout has replica {list(s.replica)}; ldstmatrix requires no replica")
|
||||
|
||||
# Step 4: R must have exactly one kind of lane axis from the valid set.
|
||||
r_thread_axes = {it.axis.name for it in r.shard if it.axis.is_thread()}
|
||||
if len(r_thread_axes) != 1:
|
||||
fail(f"R must have exactly one thread axis name; got {sorted(r_thread_axes)}")
|
||||
r_lane_axis = next(iter(r_thread_axes))
|
||||
if r_lane_axis not in _VALID_R_LANE_AXES:
|
||||
fail(f"R thread axis {r_lane_axis!r} not in {sorted(_VALID_R_LANE_AXES)}")
|
||||
|
||||
# Step 5: group S by R's iter extents (one S group per R iter, outer→inner).
|
||||
r_group_shape = [int(it.extent) for it in r.shard]
|
||||
s_grp, s_seps = s.group(r_group_shape)
|
||||
|
||||
# Step 6: permute R so thread iters come first (stride-desc), then memory
|
||||
# iters (stride-desc).
|
||||
r_perm = _compute_r_perm(r)
|
||||
r = r.permute_dims(r_perm)
|
||||
|
||||
# Step 7: apply R's perm to S in group units (1-to-1 with R's iters), and
|
||||
# rebuild s_seps to track group boundaries in the new order.
|
||||
s = s_grp.permute_by_groups(list(s_seps), r_perm)
|
||||
old_sizes = [s_seps[i + 1] - s_seps[i] for i in range(len(s_seps) - 1)]
|
||||
s_seps = [0]
|
||||
for pi in r_perm:
|
||||
s_seps.append(s_seps[-1] + old_sizes[pi])
|
||||
|
||||
# Step 7.5: canonicalize both R and S after permute. Fuses adjacent
|
||||
# contig iters — keeps step 8's group input clean. Push target so
|
||||
# scope-aware fusers run (laneid+wid_in_wg → tid_in_wg, etc.).
|
||||
with sctx.target:
|
||||
r = r.canonicalize()
|
||||
s = s.canonicalize()
|
||||
|
||||
t_total = prod(int(it.extent) for it in r.shard if it.axis.is_thread())
|
||||
m_total = prod(int(it.extent) for it in r.shard if not it.axis.is_thread())
|
||||
if t_total % 32 != 0:
|
||||
fail(f"R thread section total {t_total} not divisible by 32")
|
||||
|
||||
def _strs(lay, seps):
|
||||
# Atoms 8 / 4 / 2 (segs 1, 2, 5) must be single iters — their strides
|
||||
# feed downstream stride checks (lane partition + fragment 2-fp16
|
||||
# contig). The num atom (seg 4) may be MULTI-ITER: we return its iter
|
||||
# list and let layout.apply handle the decomposition at emit time.
|
||||
fixed_segs = [list(lay.shard[seps[i] : seps[i + 1]]) for i in (1, 2, 5)]
|
||||
if not all(len(g) == 1 for g in fixed_segs):
|
||||
return None
|
||||
num_iters = list(lay.shard[seps[4] : seps[5]])
|
||||
return (
|
||||
int(fixed_segs[0][0].stride), # 8 atom stride
|
||||
int(fixed_segs[1][0].stride), # 4 atom stride
|
||||
num_iters, # num atom iter list (multi-iter OK)
|
||||
int(fixed_segs[2][0].stride), # 2 atom stride
|
||||
)
|
||||
|
||||
def _try_num(r_in, s_in, num):
|
||||
"""Try grouping (r_in, s_in) with [T/32, 8, 4, M/(2num), num, 2].
|
||||
|
||||
Returns (rg, rsep, sg, ssep, trans, p, num) if structural checks pass,
|
||||
else None. ``trans`` is the ldmatrix .trans flag; ``p`` is the
|
||||
per-tile-row S stride used at emit.
|
||||
"""
|
||||
gs = [t_total // 32, 8, 4, m_total // (num * 2), num, 2]
|
||||
try:
|
||||
rg, rsep = r_in.group(gs)
|
||||
sg, ssep = s_in.group(gs)
|
||||
except Exception:
|
||||
return None
|
||||
# R seg 0 (T/32 outer): require single iter with stride 32. When
|
||||
# T/32 == 1 the segment is trivial — skip.
|
||||
if t_total > 32:
|
||||
seg0 = list(rg.shard[rsep[0] : rsep[1]])
|
||||
if len(seg0) != 1 or int(seg0[0].stride) != 32:
|
||||
return None
|
||||
rs, ss = _strs(rg, rsep), _strs(sg, ssep)
|
||||
if rs is None or ss is None:
|
||||
return None
|
||||
r8, r4, _r_num_iters, r2 = rs
|
||||
s8, s4, s_num_iters, s2 = ss
|
||||
if (r8, r4, r2) != (4, 1, 1):
|
||||
return None
|
||||
# S num atom: every iter must have stride > 0 and multiple of 8 (the
|
||||
# per-tile spacing geometry of ldmatrix m8n8; 8 fp16 = 16 bytes = one
|
||||
# tile column dimension).
|
||||
if num > 1 and not all(
|
||||
int(it.stride) > 0 and int(it.stride) % 8 == 0 for it in s_num_iters
|
||||
):
|
||||
return None
|
||||
# m_outer (seg 3) iters: each per-mm advance must keep the per-lane
|
||||
# SMEM address 16-byte aligned (ldmatrix .b16 reads 8 fp16 = 16 bytes
|
||||
# per lane), so the m_outer S-stride must also be a multiple of 8.
|
||||
# Without this, mm > 0 iterations land at unaligned addresses and
|
||||
# silently read garbage even though the layout group succeeds.
|
||||
# Skip extent-1 trivial iters — they contribute no per-mm advance,
|
||||
# so their (placeholder) stride is irrelevant.
|
||||
m_outer_iters = list(sg.shard[ssep[3] : ssep[4]])
|
||||
if not all(int(it.extent) == 1 or int(it.stride) % 8 == 0 for it in m_outer_iters):
|
||||
return None
|
||||
if (s4, s2) == (2, 1) and s8 > 0 and s8 % 8 == 0:
|
||||
return (rg, rsep, sg, ssep, False, s8, num)
|
||||
if s8 == 1 and s2 > 0 and s2 % 8 == 0 and s4 == 2 * s2:
|
||||
return (rg, rsep, sg, ssep, True, s2, num)
|
||||
return None
|
||||
|
||||
# Try the **sorted** variant: 5D-group, sub-group R's M/2 by S's M/2
|
||||
# extents, sort the sub-groups by descending S-stride, rebuild. This
|
||||
# makes the m_outer iter list carry the largest S-strides on top, which
|
||||
# maximizes the §2 swizzle fast-path applicability later. If anything
|
||||
# in the rebuild raises (e.g. M/2 can't be sub-grouped by S's extents),
|
||||
# we silently fall back to the no-sort path below.
|
||||
r_sort = s_sort = None
|
||||
try:
|
||||
gs5 = [t_total // 32, 8, 4, m_total // 2, 2]
|
||||
rg5, rsep5 = r.group(gs5)
|
||||
sg5, ssep5 = s.group(gs5)
|
||||
r_m_iters = list(rg5.shard[rsep5[3] : rsep5[4]])
|
||||
s_m_iters = list(sg5.shard[ssep5[3] : ssep5[4]])
|
||||
s_m_extents = [int(it.extent) for it in s_m_iters]
|
||||
# Sub-group R's M/2 iters by S's M/2 iter extents. This 1-to-1's
|
||||
# the R sub-groups with the S iters so we can permute them together.
|
||||
r_m_sub = TileLayout.from_iters(r_m_iters)
|
||||
r_m_grouped, r_m_seps = r_m_sub.group(s_m_extents)
|
||||
# Sort S iters by S-stride descending; permute R sub-groups in lockstep.
|
||||
perm = sorted(range(len(s_m_iters)), key=lambda i: -int(s_m_iters[i].stride))
|
||||
if perm != list(range(len(perm))):
|
||||
r_m_permuted = r_m_grouped.permute_by_groups(list(r_m_seps), perm)
|
||||
s_m_permuted = [s_m_iters[i] for i in perm]
|
||||
r_sort = TileLayout.from_iters(
|
||||
list(rg5.shard[: rsep5[3]])
|
||||
+ list(r_m_permuted.shard)
|
||||
+ list(rg5.shard[rsep5[4] :]),
|
||||
offset=dict(rg5.offset),
|
||||
)
|
||||
s_sort = TileLayout.from_iters(
|
||||
list(sg5.shard[: ssep5[3]]) + list(s_m_permuted) + list(sg5.shard[ssep5[4] :]),
|
||||
offset=dict(sg5.offset),
|
||||
)
|
||||
# If perm is identity, sorted == unsorted; no need to build duplicate layouts.
|
||||
except Exception:
|
||||
r_sort = s_sort = None
|
||||
|
||||
# Enumerate num largest-first; for each num try sorted then unsorted.
|
||||
chosen = None
|
||||
for num in (4, 2, 1):
|
||||
if m_total % (num * 2):
|
||||
continue
|
||||
if r_sort is not None:
|
||||
res = _try_num(r_sort, s_sort, num)
|
||||
if res is not None:
|
||||
chosen = res
|
||||
break
|
||||
res = _try_num(r, s, num)
|
||||
if res is not None:
|
||||
chosen = res
|
||||
break
|
||||
|
||||
if chosen is None:
|
||||
fail("ldstmatrix layout doesn't fit any num ∈ {4,2,1}")
|
||||
r, r_seps, s, s_seps, trans, p, num = chosen
|
||||
|
||||
# Step 10: emit one ldmatrix/stmatrix per mm, per warp.
|
||||
|
||||
def _get_warp_idx_in_T():
|
||||
# T.warp_id_in_wg() / T.warp_id() must be called from inside a
|
||||
# @T.prim_func body — wrap so the prim_func parser calls us at parse
|
||||
# time (Python `if` here is plain control flow, not TIR-intercepted).
|
||||
if r_lane_axis == "laneid":
|
||||
return 0
|
||||
if r_lane_axis == "tid_in_wg":
|
||||
return T.warp_id_in_wg()
|
||||
return T.warp_id() # "tx"
|
||||
|
||||
def _seg4_coord(laneid_expr):
|
||||
# num=1: seg 4 trivially extent-1, pass 0. num>1: use lane//8 (tile
|
||||
# index in ldmatrix lane convention); layout.apply decomposes through
|
||||
# the seg's iter structure (single or multi-iter).
|
||||
if num > 1:
|
||||
return laneid_expr // 8
|
||||
return 0
|
||||
|
||||
apply_shape = [t_total // 32, 8, 4, m_total // (num * 2), num, 2]
|
||||
r_mem_axis = r.shard[r_seps[5]].axis.name
|
||||
s_mem_axis = s.shard[s_seps[5]].axis.name
|
||||
m_outer = m_total // (num * 2)
|
||||
s_zero = [0] * len(s_buf.shape)
|
||||
|
||||
# Swizzle fast-path setup. When S is swizzled, the per-mm `tile_off +
|
||||
# row_off` is a logical offset; the physical SMEM address is
|
||||
# `swizzle.apply(logical)`. The slow path computes that per iter; the
|
||||
# fast path (§2.E of the swizzle-iter plan) reduces it to
|
||||
# `base_off + sum_j bit_j(mm) · signed_strides[j]` where base_off and
|
||||
# signed_strides are per-thread constants set once. We try to
|
||||
# recognize the m_outer iter list as such a pattern; if it fails (e.g.
|
||||
# the analyzer can't discharge condition C1 over the lane/warp
|
||||
# placeholders) we silently fall through to the slow path.
|
||||
swizzle_pattern = None
|
||||
s_off_template = None
|
||||
lane_ph = warp_ph = None
|
||||
if s_swizzle is not None:
|
||||
m_outer_iters = list(s.shard[s_seps[3] : s_seps[4]])
|
||||
iter_extents = [int(it.extent) for it in m_outer_iters]
|
||||
iter_strides = [int(it.stride) for it in m_outer_iters]
|
||||
# Build s_off at mm=0 with placeholder vars for lane and warp_idx.
|
||||
lane_ph = _TirVar("lane_ph", "int32")
|
||||
seg4_ph = (lane_ph // 8) if num > 1 else _IntImm("int32", 0)
|
||||
if r_lane_axis == "laneid":
|
||||
warp_ph_expr = _IntImm("int32", 0)
|
||||
else:
|
||||
warp_ph = _TirVar("warp_ph", "int32")
|
||||
warp_ph_expr = warp_ph
|
||||
s_off_template = s.apply(
|
||||
warp_ph_expr,
|
||||
_IntImm("int32", 0),
|
||||
_IntImm("int32", 0),
|
||||
_IntImm("int32", 0),
|
||||
seg4_ph,
|
||||
_IntImm("int32", 0),
|
||||
shape=apply_shape,
|
||||
)[s_mem_axis] + (lane_ph % 8) * _IntImm("int32", p)
|
||||
# Bind lane / warp placeholder bounds for the (C1) analyzer. ``lane_ph``
|
||||
# is the per-warp lane id ∈ [0, 32); ``warp_ph`` (when present) is the
|
||||
# warp index inside the scope: warpgroup ⇒ [0, 4), cta ⇒ [0, t_total/32).
|
||||
var_bounds = {lane_ph: tvm.ir.Range.from_min_extent(0, 32)}
|
||||
if warp_ph is not None:
|
||||
var_bounds[warp_ph] = tvm.ir.Range.from_min_extent(0, t_total // 32)
|
||||
swizzle_pattern = try_recognize(
|
||||
s_swizzle,
|
||||
iter_extents,
|
||||
iter_strides,
|
||||
s_off_template,
|
||||
var_bounds=var_bounds,
|
||||
)
|
||||
|
||||
class _SwizzleState:
|
||||
def __init__(self):
|
||||
self.signed_strides = None
|
||||
self.base_off = None
|
||||
|
||||
state = _SwizzleState()
|
||||
|
||||
def _resolve_s_off(laneid_var, warp_var):
|
||||
# Build the placeholder→runtime-var map and substitute. Keep this in a
|
||||
# regular Python helper — the @T.prim_func parser intercepts dict
|
||||
# literals when written directly in the body.
|
||||
vmap = {lane_ph: laneid_var}
|
||||
if warp_ph is not None:
|
||||
vmap[warp_ph] = warp_var
|
||||
return tvm.tirx.stmt_functor.substitute(s_off_template, vmap)
|
||||
|
||||
def _setup_swizzle(s_off_resolved):
|
||||
if swizzle_pattern is None:
|
||||
return
|
||||
state.signed_strides, state.base_off = emit_init(
|
||||
swizzle_pattern,
|
||||
s_off_resolved,
|
||||
)
|
||||
|
||||
def _smem_off(mm_idx, logical_off):
|
||||
# Three paths:
|
||||
# * pattern matched: physical off = base_off + Σ bit_j(mm)·ss[j].
|
||||
# * swizzle present, pattern missed: per-iter swizzle.apply(logical).
|
||||
# * no swizzle: identity.
|
||||
if swizzle_pattern is not None:
|
||||
return emit_iter_offset(
|
||||
swizzle_pattern,
|
||||
state.signed_strides,
|
||||
state.base_off,
|
||||
mm_idx,
|
||||
)
|
||||
if s_swizzle is not None:
|
||||
return s_swizzle.apply(logical_off)["m"]
|
||||
return logical_off
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
r_local = r_buf.local(m_total, layout=TileLayout(S[(m_total,)]))
|
||||
laneid = T.lane_id()
|
||||
warp_idx_in_T = _get_warp_idx_in_T()
|
||||
# Resolve s_off_template by substituting placeholders → actual
|
||||
# scope-id vars (via _resolve_s_off helper to keep the dict literal
|
||||
# out of the parser's view). Only the swizzle fast path needs this;
|
||||
# without swizzle we keep using the per-iter s.apply directly.
|
||||
if swizzle_pattern is not None:
|
||||
_setup_swizzle(_resolve_s_off(laneid, warp_idx_in_T))
|
||||
for mm in T.unroll(m_outer):
|
||||
tile_off = s.apply(
|
||||
warp_idx_in_T, 0, 0, mm, _seg4_coord(laneid), 0, shape=apply_shape,
|
||||
)[s_mem_axis]
|
||||
row_off = (laneid % 8) * p
|
||||
logical_off = tile_off + row_off
|
||||
smem_ptr = _ptr_off(s_buf.ptr_to(s_zero), _smem_off(mm, logical_off))
|
||||
handles = [
|
||||
r_local.ptr_to([
|
||||
r.apply(0, 0, 0, mm, i, 0, shape=apply_shape)[r_mem_axis]
|
||||
])
|
||||
for i in range(num)
|
||||
]
|
||||
if direction == "ld":
|
||||
T.ptx.ldmatrix(trans, num, ".b16", smem_ptr, *handles)
|
||||
else:
|
||||
T.ptx.stmatrix(
|
||||
trans, num, ".b16", smem_ptr, *handles,
|
||||
shape="m8n8", space="shared",
|
||||
)
|
||||
# fmt: on
|
||||
return impl
|
||||
|
||||
|
||||
@register_dispatch(
|
||||
"copy",
|
||||
"cuda",
|
||||
variant="ldstmatrix",
|
||||
priority=10,
|
||||
when=[predicate("ldstmatrix_applicable", _is_ldstmatrix)],
|
||||
)
|
||||
def copy_schedule_ldstmatrix(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
return _emit(op_call, sctx)
|
||||
|
||||
|
||||
__all__ = ["copy_schedule_ldstmatrix"]
|
||||
@@ -0,0 +1,596 @@
|
||||
# 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
|
||||
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, S, SwizzleLayout, TileLayout
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
|
||||
from tvm.tirx.operator.tile_primitive.registry import DispatchContext
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ._common import _alignment_ok, copy_ptx_form, copy_ptx_ld_return_type
|
||||
from ._swizzle_iter import (
|
||||
emit_fallback_offset,
|
||||
emit_init,
|
||||
emit_iter_offset,
|
||||
get_swizzle,
|
||||
try_recognize,
|
||||
)
|
||||
from .utils import _is_valid_copy, _scope_allowed
|
||||
|
||||
|
||||
def _extract_tile(layout, region):
|
||||
"""Strip swizzle off ``layout`` so we can perm/group it as a TileLayout.
|
||||
|
||||
``region`` is the per-axis ``(start, end)`` pair list — we only consume
|
||||
its extents when ``layout`` is a bare ``SwizzleLayout`` (rebuilding a
|
||||
trivial TileLayout for it). Plain ``TileLayout`` / ``ComposeLayout``
|
||||
don't need the extent, so symbolic regions are fine for them.
|
||||
"""
|
||||
if isinstance(layout, ComposeLayout):
|
||||
return layout.tile_layout
|
||||
if isinstance(layout, SwizzleLayout):
|
||||
# TODO: keep swizzle info around for later (addressing in emit).
|
||||
extents = [int(end - start) for (start, end) in region]
|
||||
return TileLayout(S[tuple(extents)])
|
||||
return layout
|
||||
|
||||
|
||||
_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_offset(r_p, s_p, placeholders: dict[str, _TirVar]):
|
||||
"""Per-thread S base offset. Coord per R position is placeholder (thread
|
||||
axis) or 0 (memory axis); apply_to_shape decomposes across s_p iters.
|
||||
Includes layout-level offsets (e.g. from slicing a non-zero S region)."""
|
||||
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)
|
||||
off = _IntImm("int32", 0)
|
||||
for c, it in zip(per_iter, s_p.shard, strict=True):
|
||||
off = off + c * it.stride
|
||||
for _axis, val in s_p.offset.items():
|
||||
off = off + val
|
||||
return off
|
||||
|
||||
|
||||
_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) -> 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 on either side are partition-virtual, not storage-physical,
|
||||
# so they don't enter the per-thread address — exclude them.
|
||||
shared_terms = list(s_p.offset.values()) + list(r_p.offset.values())
|
||||
for n in _vec_len_candidates(elem_bits):
|
||||
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 _s_thread_offset_with_vars(r_p, s_p, axis_var_map: dict):
|
||||
coord = [
|
||||
axis_var_map[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)
|
||||
off = _IntImm("int32", 0)
|
||||
for c, it in zip(per_iter, s_p.shard, strict=True):
|
||||
off = off + c * it.stride
|
||||
for _ax, val in s_p.offset.items():
|
||||
off = off + val
|
||||
return off
|
||||
|
||||
|
||||
def _substitute_axes(s_off_template, placeholders: dict[str, _TirVar], sctx: DispatchContext):
|
||||
"""Inside an impl body: declare real scope_ids and rewrite the
|
||||
placeholder-built ``s_off_template`` to use them."""
|
||||
vmap = {placeholders[name]: _axis_decl(name, sctx) for name in placeholders}
|
||||
return tvm.tirx.stmt_functor.substitute(s_off_template, vmap)
|
||||
|
||||
|
||||
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=base_ptr.ty,
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
vec_len = _choose_vec_len(elem_bits, atoms, r_p, s_p)
|
||||
vec_bits = vec_len * elem_bits
|
||||
outer = _split_atoms_for_vec(atoms, vec_len)
|
||||
per_thread_r_total = 1
|
||||
for it in r_iters:
|
||||
per_thread_r_total *= int(it.extent)
|
||||
per_thread_r_shape = [per_thread_r_total or 1]
|
||||
|
||||
# Build the per-thread S offset OUTSIDE the impl using placeholder Vars
|
||||
# (one per thread axis). Inside the impl we'll declare the real scope_ids
|
||||
# via T.lane_id/T.thread_id_in_wg/... and substitute them in.
|
||||
placeholders = _make_thread_placeholders(r_p)
|
||||
s_off_template = _s_thread_offset(r_p, s_p, placeholders)
|
||||
|
||||
# R-side base offset from slicing (e.g. ``R[i*8:i*8+8]`` ⇒ ``i*8``). The
|
||||
# canonicalize() result lives in ``r_p.offset``; sum across axes (memory
|
||||
# or thread — irrelevant once it's all 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
|
||||
vec, ptx_type = copy_ptx_form(num_bytes)
|
||||
space = "shared" if s_is_shared else "global"
|
||||
|
||||
total_outer = 1
|
||||
for a in outer:
|
||||
total_outer *= a[0]
|
||||
|
||||
# Swizzle handling: recognize the iter-pattern on S side from the atom
|
||||
# extents/strides (atom = (extent, s_stride, r_mul); a[1] is the S-side
|
||||
# stride per outer round, equivalent to outer_iter strides in gmem_smem).
|
||||
swizzle = get_swizzle(s_buf.layout)
|
||||
swizzle_pattern = None
|
||||
if swizzle is not None:
|
||||
swizzle_pattern = try_recognize(
|
||||
swizzle,
|
||||
[a[0] for a in outer],
|
||||
[a[1] for a in outer],
|
||||
s_off_template,
|
||||
)
|
||||
|
||||
class _SwizzleState:
|
||||
def __init__(self):
|
||||
self.signed_strides = None
|
||||
self.base_off = None
|
||||
|
||||
state = _SwizzleState()
|
||||
|
||||
def _setup_swizzle(s_off):
|
||||
if swizzle_pattern is None:
|
||||
return
|
||||
state.signed_strides, state.base_off = emit_init(swizzle_pattern, s_off)
|
||||
|
||||
def _s_iter_off(f, ds, s_off):
|
||||
if swizzle_pattern is not None:
|
||||
return emit_iter_offset(swizzle_pattern, state.signed_strides, state.base_off, f)
|
||||
if swizzle is not None:
|
||||
return emit_fallback_offset(swizzle, s_off, ds)
|
||||
return s_off + ds
|
||||
|
||||
# fmt: off
|
||||
s_zero_indices = [0] * len(s_buf.shape)
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
s_off = _substitute_axes(s_off_template, placeholders, sctx)
|
||||
_setup_swizzle(s_off)
|
||||
r_local = r_buf.local(*per_thread_r_shape)
|
||||
# Keep as a serial TIR loop and let ptxas unroll downstream. An
|
||||
# explicit ``T.unroll`` materializes the per-iter scratch
|
||||
# (ds/dr/s_ptr/r_ptr, swizzle ``v_<n>[]`` signed-strides) as N
|
||||
# copies of each buffer declaration; on kernels with many R↔S copy
|
||||
# sites and large ``total_outer`` (FA4 writeback) this floods the
|
||||
# function with ``alignas(64) int`` arrays and pressures registers.
|
||||
for f in range(total_outer):
|
||||
ds, dr = _outer_const_offsets(outer, f)
|
||||
s_ptr = _ptr_off(s_buf.ptr_to(s_zero_indices), _s_iter_off(f, ds, s_off))
|
||||
r_ptr = _ptr_off(r_local.ptr_to([0]), r_off_base + dr)
|
||||
if r_is_src:
|
||||
T.ptx.st(s_ptr, src=r_ptr, space=space, vec=vec, ptx_type=ptx_type)
|
||||
else:
|
||||
T.ptx.ld(
|
||||
s_ptr,
|
||||
copy_ptx_ld_return_type(ptx_type),
|
||||
ptx_type,
|
||||
dst=r_ptr,
|
||||
space=space,
|
||||
vec=vec,
|
||||
)
|
||||
# fmt: on
|
||||
import os
|
||||
|
||||
if os.environ.get("R2S_DUMP"):
|
||||
print("=== emitted impl ===")
|
||||
print(impl.script())
|
||||
return impl
|
||||
|
||||
|
||||
@register_dispatch(
|
||||
"copy",
|
||||
"cuda",
|
||||
variant="reg",
|
||||
priority=10,
|
||||
when=[predicate("reg_applicable", _is_reg_copy)],
|
||||
)
|
||||
def copy_schedule_reg(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
return _emit_reg(op_call, sctx)
|
||||
@@ -0,0 +1,101 @@
|
||||
# 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 helpers for copy operator dispatches on CUDA targets."""
|
||||
|
||||
from collections.abc import Iterable
|
||||
|
||||
from tvm.tirx import Buffer
|
||||
from tvm.tirx.operator.tile_primitive.registry import DispatchContext
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..common import match_scope, validate_copy_op
|
||||
|
||||
|
||||
def _is_valid_smem_tmem_copy(op_call: TilePrimitiveCall, sctx: DispatchContext):
|
||||
"""Validate smem->tmem copy operation.
|
||||
|
||||
The new tcgen05.cp.32x128b.warpx4 dispatch requires the destination tmem
|
||||
buffer to declare warpx4 broadcast as ``R[4 : 32@TLane]``. The legacy
|
||||
128-row dispatch path (no replica) goes through a separate code path and
|
||||
is not handled here.
|
||||
"""
|
||||
dst_region, src_region = op_call.args[:2]
|
||||
src: Buffer = src_region.buffer
|
||||
dst: Buffer = dst_region.buffer
|
||||
if not (src.scope().startswith("shared") and dst.scope() == "tmem"):
|
||||
return (False, f"expected shared->tmem, got {src.scope()}->{dst.scope()}")
|
||||
if not (src.layout and dst.layout):
|
||||
return (False, "both buffers must have layouts")
|
||||
if dst.allocated_addr is None:
|
||||
return (False, "tmem buffer must have allocated_addr")
|
||||
# Require warpx4 router on TMEM side so this dispatch only handles the
|
||||
# 32x128b.warpx4 case; other shapes (128x256b/128x128b etc.) fall back
|
||||
# to the legacy dispatch.
|
||||
rep = dst.layout.replica
|
||||
if not (
|
||||
len(rep) == 1
|
||||
and int(rep[0].extent) == 4
|
||||
and int(rep[0].stride) == 32
|
||||
and "TLane" in str(rep[0].axis)
|
||||
):
|
||||
return (False, f"requires R[4:32@TLane] on tmem, got replica={list(rep)}")
|
||||
return (True, None)
|
||||
|
||||
|
||||
def _single_thread_exec(op_call: TilePrimitiveCall, sctx: DispatchContext):
|
||||
"""Predicate: exec scope must be single-thread."""
|
||||
exec_scope = sctx.scope_kind
|
||||
ok = exec_scope == "thread"
|
||||
return (ok, None if ok else f"expected thread exec_scope, got {exec_scope}")
|
||||
|
||||
|
||||
DEFAULT_ALLOWED_PAIRS: tuple[tuple[str, str], ...] = (
|
||||
("global", "shared*"),
|
||||
("shared*", "global"),
|
||||
("global", "local"),
|
||||
("local", "global"),
|
||||
("shared*", "local"),
|
||||
("local", "shared*"),
|
||||
)
|
||||
|
||||
|
||||
def _scope_allowed(
|
||||
op_call: TilePrimitiveCall,
|
||||
sctx: DispatchContext,
|
||||
allowed_pairs: Iterable[tuple[str, str]] = DEFAULT_ALLOWED_PAIRS,
|
||||
):
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
dst_buffer_region, src_buffer_region = (op_call.dst, op_call.src)
|
||||
src_scope = src_buffer_region.buffer.scope()
|
||||
dst_scope = dst_buffer_region.buffer.scope()
|
||||
ok = any(
|
||||
(
|
||||
match_scope(src_scope, src_pat) and match_scope(dst_scope, dst_pat)
|
||||
for src_pat, dst_pat in allowed_pairs
|
||||
)
|
||||
)
|
||||
if not ok:
|
||||
allowed_str = ", ".join((f"{a}->{b}" for a, b in allowed_pairs))
|
||||
return (
|
||||
False,
|
||||
f"unsupported memory scopes src={src_scope} dst={dst_scope}; allowed: {allowed_str}",
|
||||
)
|
||||
return (True, None)
|
||||
|
||||
|
||||
def _is_valid_copy(op_call: TilePrimitiveCall, sctx: DispatchContext):
|
||||
return (validate_copy_op(op_call, sctx), "validate_copy_op failed")
|
||||
@@ -0,0 +1,29 @@
|
||||
# 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.
|
||||
|
||||
"""Implementation of copy_async operator dispatches for CUDA targets.
|
||||
|
||||
Registered op: copy_async (4 variants).
|
||||
See the @register_dispatch blocks in each submodule for detailed documentation
|
||||
with before/after IR examples.
|
||||
"""
|
||||
|
||||
from .dsmem import *
|
||||
from .ldgsts import *
|
||||
from .tcgen05_cp import *
|
||||
from .tcgen05_ldst import *
|
||||
from .tma import *
|
||||
@@ -0,0 +1,226 @@
|
||||
# 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 dispatch variant: dsmem (shared::cta -> shared::cluster)."""
|
||||
|
||||
import functools
|
||||
import operator
|
||||
|
||||
import tvm
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import Buffer, PrimFunc
|
||||
from tvm.tirx.operator.tile_primitive import (
|
||||
DispatchContext,
|
||||
fail,
|
||||
predicate,
|
||||
register_dispatch,
|
||||
)
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..common import validate_copy_op
|
||||
from ..exec_scope_utils import single_thread
|
||||
from .utils import find_contiguous_region, to_tile_layout
|
||||
|
||||
|
||||
def _is_shared_to_shared(op_call: TilePrimitiveCall) -> bool:
|
||||
"""Check if both src and dst are in shared memory."""
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
src_scope = op_call.src.buffer.scope()
|
||||
dst_scope = op_call.dst.buffer.scope()
|
||||
return src_scope.startswith("shared") and dst_scope.startswith("shared")
|
||||
|
||||
|
||||
def copy_dsmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
"""Implement shared-to-shared cross-CTA copy using cp.async.bulk.
|
||||
|
||||
Uses cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes
|
||||
to copy data from the executing CTA's shared memory to a remote CTA's shared
|
||||
memory within the same cluster.
|
||||
|
||||
The copy region is decomposed into contiguous byte chunks based on layout
|
||||
analysis of both src and dst buffers. Non-contiguous dimensions are iterated
|
||||
over, emitting one cp.async.bulk instruction per contiguous chunk.
|
||||
"""
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
|
||||
# Extract config
|
||||
remote_cta_id = op_call.config.get("remote_cta_id", None)
|
||||
if remote_cta_id is None:
|
||||
fail("remote_cta_id not set in config")
|
||||
mbar = op_call.config.get("mbar", None)
|
||||
if mbar is None:
|
||||
fail("mbar not set in config")
|
||||
|
||||
# Extract buffer regions
|
||||
dst_buffer_region = op_call.dst
|
||||
src_buffer_region = op_call.src
|
||||
src_buf: Buffer = src_buffer_region.buffer
|
||||
dst_buf: Buffer = dst_buffer_region.buffer
|
||||
|
||||
src_st = [r.min for r in src_buffer_region.region]
|
||||
src_ext = [r.extent for r in src_buffer_region.region]
|
||||
dst_st = [r.min for r in dst_buffer_region.region]
|
||||
dst_ext = [r.extent for r in dst_buffer_region.region]
|
||||
|
||||
dtype_bytes = tvm.DataType(src_buf.dtype).bits // 8
|
||||
|
||||
# Get tile layouts for both buffers
|
||||
src_tile_layout = to_tile_layout(src_buf.layout, src_buf.shape)
|
||||
dst_tile_layout = to_tile_layout(dst_buf.layout, dst_buf.shape)
|
||||
|
||||
# Slice layouts to copy region
|
||||
src_region_tuples = [(src_st[i], src_st[i] + src_ext[i]) for i in range(len(src_st))]
|
||||
sliced_src = src_tile_layout.slice([s for s in src_buf.shape], src_region_tuples)
|
||||
if sliced_src is None:
|
||||
fail("Cannot slice src layout for DSMEM copy")
|
||||
|
||||
dst_region_tuples = [(dst_st[i], dst_st[i] + dst_ext[i]) for i in range(len(dst_st))]
|
||||
sliced_dst = dst_tile_layout.slice([s for s in dst_buf.shape], dst_region_tuples)
|
||||
if sliced_dst is None:
|
||||
fail("Cannot slice dst layout for DSMEM copy")
|
||||
|
||||
# Group src layout by region extents, then group dst by src's shard extents
|
||||
# This creates 1:1 shard correspondence between the two layouts
|
||||
grouped_src, src_seps = sliced_src.canonicalize().group(src_ext)
|
||||
src_shard_extents = [s.extent for s in grouped_src.shard]
|
||||
grouped_dst, dst_seps = sliced_dst.canonicalize().group(src_shard_extents)
|
||||
|
||||
# Find contiguous regions in both layouts
|
||||
src_contig_indices, _ = find_contiguous_region(grouped_src)
|
||||
dst_contig_indices, _ = find_contiguous_region(grouped_dst)
|
||||
|
||||
# Intersect: walk from innermost outward, include only matching shard indices
|
||||
shared_contig_indices = []
|
||||
for s_idx, d_idx in zip(src_contig_indices, dst_contig_indices):
|
||||
if s_idx != d_idx:
|
||||
break
|
||||
shared_contig_indices.append(s_idx)
|
||||
|
||||
# Compute chunk size
|
||||
if shared_contig_indices:
|
||||
chunk_elements = functools.reduce(
|
||||
operator.mul, [grouped_src.shard[i].extent for i in shared_contig_indices], 1
|
||||
)
|
||||
else:
|
||||
chunk_elements = 1
|
||||
|
||||
chunk_bytes = chunk_elements * dtype_bytes
|
||||
if chunk_bytes < 16 or chunk_bytes % 16 != 0:
|
||||
fail(
|
||||
f"Layouts not compatible for bulk DSMEM copy: "
|
||||
f"chunk_bytes={chunk_bytes} (need >= 16 and multiple of 16)"
|
||||
)
|
||||
|
||||
# Build iteration space over non-contiguous (outer) shards
|
||||
shared_contig_set = set(shared_contig_indices)
|
||||
outer_shard_indices = [i for i in range(len(grouped_src.shard)) if i not in shared_contig_set]
|
||||
outer_extents = [grouped_src.shard[i].extent for i in outer_shard_indices]
|
||||
outer_src_strides = [grouped_src.shard[i].stride for i in outer_shard_indices]
|
||||
outer_dst_strides = [grouped_dst.shard[i].stride for i in outer_shard_indices]
|
||||
|
||||
# Helper to compute element offsets from loop variables (called via T.meta_var)
|
||||
def compute_offsets(loop_vars):
|
||||
if len(outer_extents) == 1:
|
||||
lvs = [loop_vars]
|
||||
else:
|
||||
lvs = list(loop_vars)
|
||||
src_off = 0
|
||||
dst_off = 0
|
||||
for j, v in enumerate(lvs):
|
||||
src_off = src_off + v * outer_src_strides[j]
|
||||
dst_off = dst_off + v * outer_dst_strides[j]
|
||||
return src_off, dst_off
|
||||
|
||||
src_tile = to_tile_layout(src_buf.layout, src_buf.shape)
|
||||
dst_tile = to_tile_layout(dst_buf.layout, dst_buf.shape)
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
# Map mbar to remote CTA (complete_tx targets the destination's mbar)
|
||||
remote_mbar = T.ptx.map_shared_rank(mbar, remote_cta_id)
|
||||
|
||||
if not outer_extents:
|
||||
# Single contiguous chunk — no iteration needed
|
||||
src_ptr = src_buf.ptr_to(src_st)
|
||||
cluster_dst = T.ptx.map_shared_rank(dst_buf.ptr_to(dst_st), remote_cta_id)
|
||||
T.ptx.cp_async.bulk.s2c(cluster_dst, src_ptr, chunk_bytes, remote_mbar)
|
||||
else:
|
||||
for loop_vars in T.grid(*outer_extents):
|
||||
src_elem_offset, dst_elem_offset = T.meta_var(compute_offsets(loop_vars))
|
||||
|
||||
src_buf_w = T.decl_buffer(
|
||||
src_buf.shape, src_buf.dtype, src_buf.data,
|
||||
elem_offset=src_buf.elem_offset + src_elem_offset,
|
||||
scope=src_buf.scope(),
|
||||
layout=src_tile,
|
||||
)
|
||||
dst_buf_w = T.decl_buffer(
|
||||
dst_buf.shape, dst_buf.dtype, dst_buf.data,
|
||||
elem_offset=dst_buf.elem_offset + dst_elem_offset,
|
||||
scope=dst_buf.scope(),
|
||||
layout=dst_tile,
|
||||
)
|
||||
|
||||
src_ptr = src_buf_w.ptr_to(src_st)
|
||||
cluster_dst = T.ptx.map_shared_rank(dst_buf_w.ptr_to(dst_st), remote_cta_id)
|
||||
T.ptx.cp_async.bulk.s2c(cluster_dst, src_ptr, chunk_bytes, remote_mbar)
|
||||
# fmt: on
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
# === Variant: copy_async/dsmem (priority=10) ===
|
||||
#
|
||||
# When: valid async copy at single-thread scope where both src and dst are in
|
||||
# shared memory. Used for intra-cluster DSMEM copies (shared::cta -> shared::cluster).
|
||||
#
|
||||
# Before (TilePrimitiveCall):
|
||||
# Tx.copy_async(
|
||||
# dst_smem[0:128, 0:64],
|
||||
# src_smem[0:128, 0:64],
|
||||
# config={"mbar": mbar, "remote_cta_id": cta_id}
|
||||
# )
|
||||
#
|
||||
# After (emits cp.async.bulk.shared::cluster.shared::cta):
|
||||
# cluster_dst = mapa(dst_smem.ptr, cta_id)
|
||||
# cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes
|
||||
# [cluster_dst], [src_smem.ptr], size, [mbar]
|
||||
@register_dispatch(
|
||||
"copy_async",
|
||||
"cuda",
|
||||
variant="dsmem",
|
||||
priority=10,
|
||||
when=[
|
||||
predicate(
|
||||
"validate_copy_op", lambda op, sctx: (validate_copy_op(op, sctx), "not a valid copy op")
|
||||
),
|
||||
predicate(
|
||||
"single_thread",
|
||||
lambda op, sctx: (
|
||||
single_thread(op, sctx),
|
||||
f"unsupported exec_scope {sctx.exec_scope}, expected single thread",
|
||||
),
|
||||
),
|
||||
predicate(
|
||||
"is_shared_to_shared",
|
||||
lambda op, sctx: (_is_shared_to_shared(op), "not a shared-to-shared copy"),
|
||||
),
|
||||
],
|
||||
)
|
||||
def copy_async_dispatch_dsmem(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
return copy_dsmem_impl(op, sctx)
|
||||
@@ -0,0 +1,275 @@
|
||||
# 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`` dispatch for ``global → shared`` via ``cp.async``
|
||||
(SASS: ``LDGSTS``).
|
||||
|
||||
Shares the partition / layout-alignment algorithm with
|
||||
``cuda/copy/gmem_smem.py`` (sync ``T.copy`` global ↔ shared); differs at
|
||||
emit time only:
|
||||
|
||||
* direction: ``cp.async`` is global → shared only (hardware restriction).
|
||||
* cp_size: PTX ``cp.async`` only accepts 4 / 8 / 16 bytes, so the vec-width
|
||||
candidate set is restricted to ``{32, 64, 128}`` bits.
|
||||
* emit: ``T.evaluate(T.ptx.cp_async(dst, src, cp_size))`` instead of the
|
||||
synchronous ``T.cuda.copy_{vec_bits}b(dst, src)``.
|
||||
|
||||
Note: ``cp.async`` does **not** sync at emit time — caller is responsible
|
||||
for ``commit_group`` / ``wait_group`` / ``cta_sync`` plumbing around the
|
||||
async pipeline.
|
||||
"""
|
||||
|
||||
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.operator.tile_primitive.dispatcher import (
|
||||
predicate,
|
||||
register_dispatch,
|
||||
)
|
||||
from tvm.tirx.operator.tile_primitive.registry import DispatchContext
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..copy._common import (
|
||||
_TID_AXIS_FOR_SCOPE,
|
||||
_thread_cnt,
|
||||
align_layouts_gs,
|
||||
)
|
||||
from ..copy._swizzle_iter import (
|
||||
emit_init,
|
||||
emit_iter_offset,
|
||||
get_swizzle,
|
||||
try_recognize,
|
||||
)
|
||||
from ..copy.reg import _all_threads_active, _axis_decl, _ptr_off
|
||||
from ..copy.utils import _is_valid_copy, _scope_allowed
|
||||
|
||||
# cp.async is unidirectional: global → shared.
|
||||
_LDGSTS_PAIRS = [("global", "shared*")]
|
||||
# cp.async cp_size ∈ {4, 8, 16} bytes ⇒ vec_bits ∈ {32, 64, 128}.
|
||||
_LDGSTS_VEC_BITS = (128, 64, 32)
|
||||
|
||||
|
||||
def _divides_thread_cnt_ldgsts(
|
||||
op_call: TilePrimitiveCall, sctx: DispatchContext
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Mirror of ``gmem_smem._divides_thread_cnt``: reject copies whose
|
||||
region element count doesn't divide ``thread_cnt`` (and reject
|
||||
``thread_cnt=0`` scopes outright). See that docstring for rationale."""
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
thread_cnt = _thread_cnt(sctx)
|
||||
if thread_cnt <= 0:
|
||||
return False, f"degenerate thread_cnt={thread_cnt} (scope has empty intra)"
|
||||
g_br = op_call.src if op_call.src.buffer.scope() == "global" else op_call.dst
|
||||
n_elements = 1
|
||||
for r in g_br.region:
|
||||
ext = r.extent
|
||||
try:
|
||||
n_elements *= int(ext)
|
||||
except (TypeError, ValueError):
|
||||
return False, f"non-constant region extent {ext}"
|
||||
if n_elements % thread_cnt != 0:
|
||||
return False, (f"region size {n_elements} not divisible by thread_cnt={thread_cnt}")
|
||||
return True, None
|
||||
|
||||
|
||||
def _is_ldgsts(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=_LDGSTS_PAIRS),
|
||||
lambda: _divides_thread_cnt_ldgsts(op_call, sctx),
|
||||
):
|
||||
ok, msg = check()
|
||||
if not ok:
|
||||
return False, msg
|
||||
return True, None
|
||||
|
||||
|
||||
def _emit_ldgsts(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
src: Buffer = op_call.src.buffer
|
||||
dst: Buffer = op_call.dst.buffer
|
||||
# Predicate above guarantees src is global, dst is shared.
|
||||
g_buf, g_br = src, op_call.src
|
||||
s_buf, s_br = dst, op_call.dst
|
||||
|
||||
g_region = [(r.min, r.min + r.extent) for r in g_br.region]
|
||||
s_region = [(r.min, r.min + r.extent) for r in s_br.region]
|
||||
|
||||
elem_bits = DataType(src.dtype).bits
|
||||
thread_cnt = _thread_cnt(sctx)
|
||||
|
||||
with sctx.target:
|
||||
g_p, s_p, vec_len = align_layouts_gs(
|
||||
g_buf.layout,
|
||||
g_buf.shape,
|
||||
g_region,
|
||||
s_buf.layout,
|
||||
s_buf.shape,
|
||||
s_region,
|
||||
elem_bits,
|
||||
thread_cnt,
|
||||
vec_bits_candidates=_LDGSTS_VEC_BITS,
|
||||
)
|
||||
|
||||
vec_bits = vec_len * elem_bits
|
||||
cp_size = vec_bits // 8 # cp.async cp_size is in bytes
|
||||
if cp_size not in (4, 8, 16):
|
||||
# align_layouts_gs already restricted candidates to _LDGSTS_VEC_BITS,
|
||||
# so reaching here means no candidate worked at all.
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import fail
|
||||
|
||||
fail(f"ldgsts: cannot find a cp.async-compatible vec_len for elem_bits={elem_bits}")
|
||||
|
||||
# Mirror gmem_smem.py: build 3D `(f, tid, 0)` against
|
||||
# `[total_outer, thread_cnt, vec_len]` and let `s_p.apply(coord, shape)`
|
||||
# flatten + resplit into whatever multi-iter T / outer-iter structure
|
||||
# `align_layouts_gs` picked. Emit is oblivious to how many shard iters
|
||||
# cover T.
|
||||
n_elements = 1
|
||||
for it in s_p.shard:
|
||||
n_elements *= int(it.extent)
|
||||
assert n_elements % (thread_cnt * vec_len) == 0, (
|
||||
f"partition produced {n_elements} elements but thread_cnt({thread_cnt}) * "
|
||||
f"vec_len({vec_len}) = {thread_cnt * vec_len} doesn't divide it"
|
||||
)
|
||||
total_outer = n_elements // (thread_cnt * vec_len)
|
||||
apply_shape = [
|
||||
_IntImm("int32", total_outer),
|
||||
_IntImm("int32", thread_cnt),
|
||||
_IntImm("int32", vec_len),
|
||||
]
|
||||
|
||||
s_zero = [0] * len(s_buf.shape)
|
||||
g_zero = [0] * len(g_buf.shape)
|
||||
|
||||
tid_axis_name = _TID_AXIS_FOR_SCOPE[sctx.scope_kind] if thread_cnt > 1 else None
|
||||
|
||||
# T-iters-walk-back to recover outer_iters_s for the fast-path
|
||||
# recognizer. Same trick as gmem_smem.py.
|
||||
if thread_cnt > 1:
|
||||
acc, _i = 1, len(s_p.shard) - 2
|
||||
while _i >= 0 and acc < thread_cnt:
|
||||
_ext = int(s_p.shard[_i].extent)
|
||||
if acc * _ext > thread_cnt:
|
||||
break
|
||||
acc *= _ext
|
||||
_i -= 1
|
||||
outer_iters_s = list(s_p.shard[: _i + 1]) if acc == thread_cnt else []
|
||||
else:
|
||||
outer_iters_s = list(s_p.shard[:-1])
|
||||
|
||||
swizzle = get_swizzle(s_buf.layout)
|
||||
swizzle_pattern = None
|
||||
if swizzle is not None and outer_iters_s:
|
||||
if tid_axis_name is not None:
|
||||
_tid_placeholder = _TirVar(tid_axis_name, "int32")
|
||||
else:
|
||||
_tid_placeholder = _IntImm("int32", 0)
|
||||
s_off_template = s_p.apply(
|
||||
_IntImm("int32", 0),
|
||||
_tid_placeholder,
|
||||
_IntImm("int32", 0),
|
||||
shape=apply_shape,
|
||||
)["m"]
|
||||
swizzle_pattern = try_recognize(
|
||||
swizzle,
|
||||
[int(it.extent) for it in outer_iters_s],
|
||||
[int(it.stride) for it in outer_iters_s],
|
||||
s_off_template,
|
||||
)
|
||||
|
||||
class _SwizzleState:
|
||||
def __init__(self):
|
||||
self.signed_strides = None
|
||||
self.base_off = None
|
||||
|
||||
state = _SwizzleState()
|
||||
|
||||
def _decl_tid():
|
||||
if tid_axis_name is not None:
|
||||
return _axis_decl(tid_axis_name, sctx)
|
||||
return _IntImm("int32", 0)
|
||||
|
||||
def _setup_swizzle(tid):
|
||||
if swizzle_pattern is None:
|
||||
return
|
||||
s_off_resolved = s_p.apply(
|
||||
_IntImm("int32", 0),
|
||||
tid,
|
||||
_IntImm("int32", 0),
|
||||
shape=apply_shape,
|
||||
)["m"]
|
||||
state.signed_strides, state.base_off = emit_init(
|
||||
swizzle_pattern,
|
||||
s_off_resolved,
|
||||
)
|
||||
|
||||
if swizzle_pattern is not None:
|
||||
|
||||
def _s_off(f, s_lin):
|
||||
return emit_iter_offset(
|
||||
swizzle_pattern,
|
||||
state.signed_strides,
|
||||
state.base_off,
|
||||
f,
|
||||
)
|
||||
elif swizzle is not None:
|
||||
_sw = swizzle
|
||||
|
||||
def _s_off(f, s_lin):
|
||||
return _sw.apply(s_lin)["m"]
|
||||
else:
|
||||
|
||||
def _s_off(f, s_lin):
|
||||
return s_lin
|
||||
|
||||
v0 = _IntImm("int32", 0)
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
tid = _decl_tid()
|
||||
_setup_swizzle(tid)
|
||||
for f in T.unroll(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_off = _s_off(f, s_lin)
|
||||
s_ptr = _ptr_off(s_buf.ptr_to(s_zero), s_off)
|
||||
g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin)
|
||||
T.evaluate(T.ptx.cp_async(s_ptr, g_ptr, cp_size))
|
||||
# cp.async is caller-synced — no cta_sync here (commit_group /
|
||||
# wait_group / cta_sync are the caller's responsibility).
|
||||
# fmt: on
|
||||
return impl
|
||||
|
||||
|
||||
@register_dispatch(
|
||||
"copy_async",
|
||||
"cuda",
|
||||
variant="ldgsts",
|
||||
priority=20,
|
||||
when=[predicate("ldgsts_applicable", _is_ldgsts)],
|
||||
)
|
||||
def copy_schedule_ldgsts(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
return _emit_ldgsts(op_call, sctx)
|
||||
@@ -0,0 +1,485 @@
|
||||
# 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.
|
||||
|
||||
"""smem->tmem dispatch via tcgen05.cp.32x128b.warpx4.
|
||||
|
||||
``tcgen05.cp`` is inherently async; this dispatch emits the cp loop only and
|
||||
leaves completion signaling (``tcgen05.commit`` against a barrier) to the
|
||||
caller. Callers who want sync semantics should issue ``tcgen05.commit``
|
||||
themselves after the copy.
|
||||
|
||||
Algorithm
|
||||
---------
|
||||
Given ``Tx.copy_async(t_region, s_region)`` where t is in tmem (with
|
||||
R[4:32@TLane] indicating warpx4 broadcast), and s is in shared memory:
|
||||
|
||||
A. Slice + canonicalize both layouts at the given regions.
|
||||
B. Verify ``t.replica == [4:32@TLane]`` (warpx4 router).
|
||||
C. Compute permutation that puts TLane first, then TCol stride-descending;
|
||||
apply to t.permute_dims and to s via group + permute_by_groups.
|
||||
D. Canonicalize again.
|
||||
E. Isolate broadcast: split-by-stride-zero on both t and s; their split
|
||||
sequences must match (same distinct prefix prods + broadcast extents).
|
||||
Drop stride-0 iters → ``t_iso`` and ``s_iso``.
|
||||
F. Group both into ``(32, middle, elem_per_128b)``. Validate:
|
||||
- t_lane = (32, 1@TLane)
|
||||
- t_col = (elem_per_128b, 1@TCol)
|
||||
- s_col = (elem_per_128b, 1)
|
||||
- s_lane refines into (4, 8) on m axis with strides (SDO_stride, atom_K_stride)
|
||||
- atom_K_byte ∈ {16, 32, 64, 128} → swizzle_mode 0..3
|
||||
- swizzle_mode matches s_buf.layout's SwizzleLayout (if any)
|
||||
G. Alignment checks:
|
||||
- t_iso TCol offset ≡ 0 (mod 32-bit)
|
||||
- s_iso m offset ≡ 0 (mod 16B for sw=0; mod atom_size for sw>0)
|
||||
- middle iter strides 16B-aligned
|
||||
H. middle 1-1 correspondence (simple-mode): t_middle and s_middle have same
|
||||
iter count and matching extents per position.
|
||||
I. Emit:
|
||||
- SmemDescriptor encoded once at SMEM base (hoisted via post_buffer_def_stmt).
|
||||
- Loop over middle iters; each cp uses ``desc.add_16B_offset(init + loop)``
|
||||
and writes to ``tmem_addr + t_col0 + Σ i_j * t_step_j``.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import operator
|
||||
|
||||
import tvm
|
||||
from tvm.arith import Analyzer
|
||||
from tvm.runtime import DataType
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import Buffer, PrimFunc
|
||||
from tvm.tirx.layout import ComposeLayout, SwizzleLayout, TCol, TileLayout, TLane
|
||||
from tvm.tirx.layout import m as m_axis
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
|
||||
from tvm.tirx.stmt import AllocBuffer, Evaluate, SeqStmt, TilePrimitiveCall
|
||||
|
||||
from ..copy import _is_valid_smem_tmem_copy, _single_thread_exec
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# -----------------------------------------------------------------------------
|
||||
def _compute_perm(t):
|
||||
def key(p):
|
||||
it = p[1]
|
||||
return (0 if it.axis == TLane else 1, -int(it.stride))
|
||||
|
||||
return [i for i, _ in sorted(enumerate(t.shard), key=key)]
|
||||
|
||||
|
||||
def _split_by_zero(lay):
|
||||
"""Split lay.shard into segments at stride==0 positions.
|
||||
Returns (split_seq, kept_iters_with_nonzero_stride)."""
|
||||
new_seq = []
|
||||
keep = []
|
||||
cur = 1
|
||||
for it in lay.shard:
|
||||
e, st = int(it.extent), int(it.stride)
|
||||
if st == 0:
|
||||
if cur > 1:
|
||||
new_seq.append(cur)
|
||||
new_seq.append(e)
|
||||
cur = 1
|
||||
else:
|
||||
cur *= e
|
||||
keep.append(it)
|
||||
if cur > 1:
|
||||
new_seq.append(cur)
|
||||
return new_seq, keep
|
||||
|
||||
|
||||
def _align_middles(t_middle, s_middle):
|
||||
"""Sub-group both middles by union-of-boundaries so they become 1-1.
|
||||
|
||||
Both inputs must be post-canonicalize iter lists with equal extent products.
|
||||
The shape is the consecutive ratios of sorted(B_t U B_s) where B_x is the
|
||||
set of cumulative extent boundaries of x_middle. Each segment then contains
|
||||
at most one iter per side (whole or sub-divided), so trivially single iter.
|
||||
|
||||
Returns (new_t_middle, new_s_middle) with len() == len() == k segments,
|
||||
each segment a single Iter on each side.
|
||||
"""
|
||||
|
||||
def cum_bounds(iters):
|
||||
b, p = [], 1
|
||||
for it in iters:
|
||||
p *= int(it.extent)
|
||||
b.append(p)
|
||||
return b
|
||||
|
||||
t_bounds = cum_bounds(t_middle)
|
||||
s_bounds = cum_bounds(s_middle)
|
||||
if not t_bounds and not s_bounds:
|
||||
return t_middle, s_middle
|
||||
N = t_bounds[-1] if t_bounds else s_bounds[-1]
|
||||
if (s_bounds and s_bounds[-1] != N) or (t_bounds and t_bounds[-1] != N):
|
||||
raise ValueError(f"middle extent mismatch: t={N} s={s_bounds[-1] if s_bounds else 0}")
|
||||
|
||||
cuts = sorted(set(t_bounds) | set(s_bounds))
|
||||
shape, prev = [], 1
|
||||
for c in cuts:
|
||||
if c % prev != 0:
|
||||
raise ValueError(
|
||||
f"middle align failed: cut {c} not divisible by prev cut {prev} "
|
||||
f"(t_bounds={t_bounds}, s_bounds={s_bounds})"
|
||||
)
|
||||
shape.append(c // prev)
|
||||
prev = c
|
||||
|
||||
def subgroup(iters):
|
||||
if len(iters) == 1 and shape == [int(iters[0].extent)]:
|
||||
return iters
|
||||
lay, _seps = TileLayout.from_iters(iters, [], {}).group(shape)
|
||||
seps = list(_seps)
|
||||
out = []
|
||||
for i in range(len(shape)):
|
||||
seg = list(lay.shard[seps[i] : seps[i + 1]])
|
||||
seg_canon = list(TileLayout.from_iters(seg, [], {}).canonicalize().shard)
|
||||
if len(seg_canon) != 1:
|
||||
raise ValueError(
|
||||
f"middle sub-group seg[{i}] not single iter after canon: {seg_canon}"
|
||||
)
|
||||
out.append(seg_canon[0])
|
||||
return out
|
||||
|
||||
return subgroup(t_middle), subgroup(s_middle)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Plan (state object)
|
||||
# -----------------------------------------------------------------------------
|
||||
def _build_plan(op_call: TilePrimitiveCall, sctx: DispatchContext):
|
||||
"""Run A..H and return a dispatch plan.
|
||||
|
||||
Plan fields:
|
||||
- s_buf, t_buf
|
||||
- dtype, dtype_bits
|
||||
- elem_per_128b, elem_per_32b
|
||||
- SmemSwizzleMode (int)
|
||||
- SDO_field, atom_K_byte
|
||||
- middle_iters: list of (extent, s_step_16B, t_step_32bcol)
|
||||
- init_off_16B (Expr)
|
||||
- t_col0 (Expr, TMEM 32-bit col offset for cp's first call)
|
||||
"""
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
dst_region, src_region = op_call.args[:2]
|
||||
s_buf: Buffer = src_region.buffer
|
||||
t_buf: Buffer = dst_region.buffer
|
||||
dtype = s_buf.dtype
|
||||
dtype_bits = DataType(dtype).bits
|
||||
elem_per_128b = 128 // dtype_bits
|
||||
elem_per_32b = 32 // dtype_bits
|
||||
|
||||
# C: slice + canonicalize.
|
||||
s_region = [(r.min, r.min + r.extent) for r in src_region.region]
|
||||
t_region = [(r.min, r.min + r.extent) for r in dst_region.region]
|
||||
s = s_buf.layout.slice(list(s_buf.shape), s_region).canonicalize()
|
||||
t = t_buf.layout.slice(list(t_buf.shape), t_region).canonicalize()
|
||||
|
||||
# If s is ComposeLayout (SwizzleLayout∘TileLayout), peel off the swizzle
|
||||
# for stride analysis; record swizzle_len for cross-check.
|
||||
s_swizzle_mode_from_layout = 0
|
||||
if isinstance(s, ComposeLayout):
|
||||
s_swizzle_mode_from_layout = int(s.swizzle.swizzle_len)
|
||||
s = s.tile_layout
|
||||
elif isinstance(s, SwizzleLayout):
|
||||
raise ValueError("s slice produced bare SwizzleLayout (unexpected)")
|
||||
|
||||
# B: warpx4 router check.
|
||||
rep = t.replica
|
||||
if not (
|
||||
len(rep) == 1
|
||||
and int(rep[0].extent) == 4
|
||||
and int(rep[0].stride) == 32
|
||||
and rep[0].axis == TLane
|
||||
):
|
||||
raise ValueError(
|
||||
f"warpx4 router fail: t.replica = "
|
||||
f"{[(int(r.extent), int(r.stride), str(r.axis)) for r in rep]}"
|
||||
)
|
||||
|
||||
# C: permute (TLane first, TCol stride desc).
|
||||
perm = _compute_perm(t)
|
||||
t_shape_for_group = [int(it.extent) for it in t.shard]
|
||||
s_grp, seps = s.group(t_shape_for_group)
|
||||
s_p = s_grp.permute_by_groups(list(seps), perm).canonicalize()
|
||||
t_p = t.permute_dims(perm).canonicalize()
|
||||
|
||||
# E: isolate broadcast.
|
||||
seq_t, keep_t = _split_by_zero(t_p)
|
||||
seq_s, keep_s = _split_by_zero(s_p)
|
||||
if seq_t != seq_s:
|
||||
raise ValueError(f"isolate split mismatch: t={seq_t} s={seq_s}")
|
||||
s_iso = TileLayout.from_iters(keep_s, list(s_p.replica), dict(s_p.offset))
|
||||
t_iso = TileLayout.from_iters(keep_t, list(t_p.replica), dict(t_p.offset))
|
||||
|
||||
# F: group into (32, middle, elem_per_128b).
|
||||
def shard_prod(lay):
|
||||
return functools.reduce(operator.mul, [int(it.extent) for it in lay.shard], 1)
|
||||
|
||||
n_lane, n_col = 32, elem_per_128b
|
||||
n_mid_t = shard_prod(t_iso) // (n_lane * n_col)
|
||||
n_mid_s = shard_prod(s_iso) // (n_lane * n_col)
|
||||
t_grp, t_seps = t_iso.group([n_lane, n_mid_t, n_col])
|
||||
s_grp2, s_seps = s_iso.group([n_lane, n_mid_s, n_col])
|
||||
t_seps = list(t_seps)
|
||||
s_seps = list(s_seps)
|
||||
|
||||
def _canon_segment(iters):
|
||||
return TileLayout.from_iters(iters, [], {}).canonicalize().shard
|
||||
|
||||
t_lane = list(_canon_segment(list(t_grp.shard[t_seps[0] : t_seps[1]])))
|
||||
t_middle = list(_canon_segment(list(t_grp.shard[t_seps[1] : t_seps[2]])))
|
||||
t_col = list(_canon_segment(list(t_grp.shard[t_seps[2] : t_seps[3]])))
|
||||
s_lane = list(s_grp2.shard[s_seps[0] : s_seps[1]])
|
||||
s_middle = list(_canon_segment(list(s_grp2.shard[s_seps[1] : s_seps[2]])))
|
||||
s_col = list(_canon_segment(list(s_grp2.shard[s_seps[2] : s_seps[3]])))
|
||||
|
||||
# F.5: align middles via union-cut sub-grouping. Both t_middle and s_middle
|
||||
# are post-canonicalize. To make their structure 1-1 we sub-group both by
|
||||
# the union of their internal cumulative-extent boundaries.
|
||||
t_middle, s_middle = _align_middles(t_middle, s_middle)
|
||||
|
||||
# F.1: lane / col validation.
|
||||
if len(t_lane) != 1:
|
||||
raise ValueError(f"t_lane must canonicalize to single iter, got {t_lane}")
|
||||
if len(t_col) != 1:
|
||||
raise ValueError(f"t_col must canonicalize to single iter, got {t_col}")
|
||||
if len(s_col) != 1:
|
||||
raise ValueError(f"s_col must canonicalize to single iter, got {s_col}")
|
||||
li = t_lane[0]
|
||||
if not (int(li.extent) == 32 and int(li.stride) == 1 and li.axis == TLane):
|
||||
raise ValueError(f"t_lane must be (32, 1@TLane), got {li}")
|
||||
ci = t_col[0]
|
||||
if not (int(ci.extent) == elem_per_128b and int(ci.stride) == 1 and ci.axis == TCol):
|
||||
raise ValueError(f"t_col must be ({elem_per_128b}, 1@TCol), got {ci}")
|
||||
sci = s_col[0]
|
||||
if not (int(sci.extent) == elem_per_128b and int(sci.stride) == 1):
|
||||
raise ValueError(f"s_col must be ({elem_per_128b}, 1, m), got {sci}")
|
||||
|
||||
# F.2: s_lane → group (4, 8) → (SDO_stride, atom_K_stride)
|
||||
s_lane_layout = TileLayout.from_iters(s_lane, [], {})
|
||||
s_lane_grp, s_lane_seps = s_lane_layout.group([4, 8])
|
||||
s_lane_seps = list(s_lane_seps)
|
||||
blk_4 = list(s_lane_grp.shard[s_lane_seps[0] : s_lane_seps[1]])
|
||||
blk_8 = list(s_lane_grp.shard[s_lane_seps[1] : s_lane_seps[2]])
|
||||
if len(blk_4) != 1 or len(blk_8) != 1:
|
||||
raise ValueError(
|
||||
f"s_lane must group into single iter per block: blk_4={blk_4}, blk_8={blk_8}"
|
||||
)
|
||||
SDO_byte = int(blk_4[0].stride) * dtype_bits // 8
|
||||
atom_K_byte = int(blk_8[0].stride) * dtype_bits // 8
|
||||
sw_candidates = {16: 0, 32: 1, 64: 2, 128: 3}
|
||||
if atom_K_byte not in sw_candidates:
|
||||
raise ValueError(f"atom_K_byte {atom_K_byte} not in {{16,32,64,128}}")
|
||||
derived_sw = sw_candidates[atom_K_byte]
|
||||
if s_swizzle_mode_from_layout != derived_sw:
|
||||
raise ValueError(
|
||||
f"swizzle mode mismatch: s_layout swizzle_len="
|
||||
f"{s_swizzle_mode_from_layout} but atom_K_byte={atom_K_byte} "
|
||||
f"implies sw={derived_sw}"
|
||||
)
|
||||
|
||||
analyzer = Analyzer()
|
||||
|
||||
# G: alignments.
|
||||
# G.1: t_iso TCol offset ≡ 0 (mod 32-bit element count).
|
||||
t_col_offset_expr = 0
|
||||
for ax, val in t_iso.offset.items():
|
||||
if ax == TCol:
|
||||
t_col_offset_expr = val
|
||||
break
|
||||
if not analyzer.can_prove_equal(t_col_offset_expr % elem_per_32b, 0):
|
||||
raise ValueError(f"t TCol offset {t_col_offset_expr} not provably 32b-aligned")
|
||||
|
||||
# G.2: s_iso m offset alignment.
|
||||
s_m_offset_expr = 0
|
||||
for ax, val in s_iso.offset.items():
|
||||
if ax == m_axis:
|
||||
s_m_offset_expr = val
|
||||
break
|
||||
elem_per_16B = 16 * 8 // dtype_bits
|
||||
if derived_sw == 0:
|
||||
align_elem = elem_per_16B
|
||||
align_label = "16B"
|
||||
else:
|
||||
atom_size_byte = 8 * atom_K_byte
|
||||
align_elem = atom_size_byte * 8 // dtype_bits
|
||||
align_label = f"atom={atom_size_byte}B"
|
||||
if not analyzer.can_prove_equal(s_m_offset_expr % align_elem, 0):
|
||||
raise ValueError(
|
||||
f"s offset {s_m_offset_expr} not provably aligned to {align_label} "
|
||||
f"({align_elem} {dtype} elements)"
|
||||
)
|
||||
|
||||
# H: middle 1-1 correspondence.
|
||||
if len(t_middle) != len(s_middle):
|
||||
raise ValueError(
|
||||
f"t_middle iter count {len(t_middle)} != s_middle {len(s_middle)} "
|
||||
"(simple-mode requires 1-1)"
|
||||
)
|
||||
middle_iters = []
|
||||
for i, (ti, si) in enumerate(zip(t_middle, s_middle)):
|
||||
if int(ti.extent) != int(si.extent):
|
||||
raise ValueError(f"middle[{i}] extent: t={int(ti.extent)} s={int(si.extent)}")
|
||||
n = int(ti.extent)
|
||||
if n == 1:
|
||||
continue
|
||||
if ti.axis != TCol:
|
||||
raise ValueError(f"middle[{i}] t axis must be TCol, got {ti.axis}")
|
||||
s_stride_byte = int(si.stride) * dtype_bits // 8
|
||||
if s_stride_byte % 16 != 0:
|
||||
raise ValueError(f"s_middle[{i}] stride {s_stride_byte}B not 16B-aligned")
|
||||
middle_iters.append((n, s_stride_byte // 16, int(ti.stride) // elem_per_32b))
|
||||
|
||||
SDO_field = SDO_byte // 16
|
||||
init_off_16B = s_m_offset_expr * dtype_bits // 8 // 16
|
||||
t_col0 = t_col_offset_expr // elem_per_32b
|
||||
|
||||
return {
|
||||
"s_buf": s_buf,
|
||||
"t_buf": t_buf,
|
||||
"dtype": dtype,
|
||||
"dtype_bits": dtype_bits,
|
||||
"elem_per_128b": elem_per_128b,
|
||||
"elem_per_32b": elem_per_32b,
|
||||
"swizzle_mode": derived_sw,
|
||||
"SDO_field": SDO_field,
|
||||
"atom_K_byte": atom_K_byte,
|
||||
"middle_iters": middle_iters,
|
||||
"init_off_16B": init_off_16B,
|
||||
"t_col0": t_col0,
|
||||
}
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Descriptor caching: one (smem_buf, ldo, sdo, swizzle) → one desc_buf,
|
||||
# encoded once at SMEM base, hoisted to right after SMEM alloc via
|
||||
# add_post_buffer_def_stmt.
|
||||
# -----------------------------------------------------------------------------
|
||||
def _get_or_create_desc(sctx, s_buf, ldo, sdo, swizzle):
|
||||
# Cache descriptor template at SMEM 0; patch addr per cp.
|
||||
cache_key = f"smem_tmem_desc:{int(ldo)}:{int(sdo)}:{int(swizzle)}"
|
||||
cached = sctx.cache_get(cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
desc_buf = tvm.tirx.decl_buffer((1,), "uint64", name="cp_desc", scope="local")
|
||||
encode_call = T.ptx.tcgen05.encode_matrix_descriptor(
|
||||
desc_buf.data, T.reinterpret("handle", T.uint64(0)), ldo, sdo, swizzle
|
||||
)
|
||||
wrap = SeqStmt([AllocBuffer(desc_buf), Evaluate(encode_call)])
|
||||
sctx.add_post_buffer_def_stmt(s_buf, wrap)
|
||||
sctx.cache_set(cache_key, desc_buf)
|
||||
return desc_buf
|
||||
|
||||
|
||||
def _desc_set_addr(desc_val, addr_ptr):
|
||||
"""Patch a SMEM matrix descriptor's 14-bit address field with cvta(addr)>>4 —
|
||||
matches the hand-rolled ``replace_smem_desc_addr`` (descriptor encoded at 0)."""
|
||||
start_addr = T.cast(
|
||||
T.bitwise_and(
|
||||
T.shift_right(T.cuda.cvta_generic_to_shared(addr_ptr), T.uint32(4)),
|
||||
T.uint32(0x3FFF),
|
||||
),
|
||||
"uint64",
|
||||
)
|
||||
return T.bitwise_or(T.bitwise_and(desc_val, T.bitwise_not(T.uint64(0x3FFF))), start_addr)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Core impl: emits the cp loop given a plan + cp config. Async only — caller
|
||||
# is responsible for issuing ``tcgen05.commit`` against a barrier if they
|
||||
# need synchronization.
|
||||
# -----------------------------------------------------------------------------
|
||||
def copy_smem_tmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
|
||||
plan = _build_plan(op_call, sctx)
|
||||
s_buf = plan["s_buf"]
|
||||
t_buf = plan["t_buf"]
|
||||
SDO_field = plan["SDO_field"]
|
||||
sw = plan["swizzle_mode"]
|
||||
middle_iters = plan["middle_iters"]
|
||||
init_off_16B = plan["init_off_16B"]
|
||||
t_col0 = plan["t_col0"]
|
||||
|
||||
# cp ignores LDO for data; non-zero LDO bloats the addr-patch codegen.
|
||||
LDO_field = 0
|
||||
|
||||
cta_group = op_call.config.get("cta_group", 1)
|
||||
|
||||
desc_buf = _get_or_create_desc(sctx, s_buf, LDO_field, SDO_field, sw)
|
||||
t_addr = t_buf.allocated_addr
|
||||
s_rank = len(s_buf.shape)
|
||||
|
||||
def _cp_desc(off_16B):
|
||||
addr = T.ptr_byte_offset(s_buf.ptr_to([0] * s_rank), off_16B * 16, s_buf.dtype)
|
||||
return _desc_set_addr(desc_buf[0], addr)
|
||||
|
||||
# Flatten the N-D middle iteration into a single T.unroll. Each iteration's
|
||||
# per-dim index is (flat // stride) % extent, summed into the t/s offsets.
|
||||
# Works uniformly for n_mid ∈ {0, 1, 2, ...}; total == 1 (no middle dims) is
|
||||
# special-cased to avoid a degenerate T.unroll(1).
|
||||
total = functools.reduce(operator.mul, [n for n, _, _ in middle_iters], 1)
|
||||
|
||||
# fmt: off
|
||||
if total == 1:
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
T.ptx.tcgen05.cp(
|
||||
t_addr[0] + t_col0,
|
||||
_cp_desc(init_off_16B),
|
||||
shape="32x128b", cta_group=cta_group, multicast="warpx4",
|
||||
)
|
||||
else:
|
||||
def compute_offsets(flat):
|
||||
t_off = 0
|
||||
s_off = 0
|
||||
div = 1
|
||||
for n, s_step, t_step in middle_iters:
|
||||
idx = (flat // div) % n
|
||||
div = div * n
|
||||
t_off = t_off + idx * t_step
|
||||
s_off = s_off + idx * s_step
|
||||
return t_off, s_off
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
for flat in T.unroll(total):
|
||||
t_off, s_off = T.meta_var(compute_offsets(flat))
|
||||
T.ptx.tcgen05.cp(
|
||||
t_addr[0] + t_col0 + t_off,
|
||||
_cp_desc(init_off_16B + s_off),
|
||||
shape="32x128b", cta_group=cta_group, multicast="warpx4",
|
||||
)
|
||||
# fmt: on
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
# === Variant: copy_async/smem->tmem (priority=10) ===
|
||||
@register_dispatch(
|
||||
"copy_async",
|
||||
"cuda",
|
||||
variant="smem->tmem",
|
||||
priority=10,
|
||||
when=[
|
||||
predicate("validate_smem_tmem_copy", _is_valid_smem_tmem_copy),
|
||||
predicate("exec_scope", _single_thread_exec),
|
||||
],
|
||||
)
|
||||
def copy_async_schedule_smem_tmem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
return copy_smem_tmem_impl(op_call, sctx)
|
||||
@@ -0,0 +1,458 @@
|
||||
# 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 dispatch: ``tcgen05.ld`` / ``tcgen05.st`` (tmem <-> local registers).
|
||||
|
||||
Both are inherently async; this dispatch emits the PTX instruction only and
|
||||
leaves completion (``tcgen05.wait.ld`` / ``tcgen05.wait.st``) to the caller.
|
||||
Callers that want sync semantics should issue the matching wait after the copy.
|
||||
"""
|
||||
|
||||
import tvm
|
||||
from tvm.arith import Analyzer
|
||||
from tvm.runtime import DataType
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import Buffer, PrimFunc
|
||||
from tvm.tirx.layout import (
|
||||
S,
|
||||
TCol,
|
||||
TileLayout,
|
||||
TLane,
|
||||
tcgen05_atom_layout,
|
||||
tid_in_wg,
|
||||
tmem_datapath_layout,
|
||||
)
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..common import get_st_extent
|
||||
from ..copy import _is_valid_copy, _scope_allowed
|
||||
from ..exec_scope_utils import exec_scope_ok
|
||||
|
||||
# Per-warp fp32-column factor for each instr_shape (mirrors
|
||||
# ``_TCGEN05_COL_FACTOR_FP32`` in ``tvm.tirx.layout``; .16x64b → 2,
|
||||
# .16x128b → 4, .16x256b → 8). Source: PTX ISA Table 49.
|
||||
_TCGEN05_COL_FACTOR_FP32 = {"16x64b": 2, "16x128b": 4, "16x256b": 8}
|
||||
|
||||
|
||||
def _match_tcgen05_atom_layout(buf):
|
||||
"""Return ``(instr_shape, rep, frag_rows)`` if ``buf.layout`` matches a
|
||||
tcgen05 ``.16x*b`` atom layout for some supported ``instr_shape``.
|
||||
|
||||
The local buffer shape ``(frag_rows, K)`` (``frag_rows`` ∈ {64, 128})
|
||||
together with the dtype determines the candidate ``rep`` for each
|
||||
``instr_shape``; we just probe the three shapes x two frag_rows and
|
||||
structurally compare. ``None`` if no atom layout matches.
|
||||
"""
|
||||
if len(buf.shape) != 2:
|
||||
return None
|
||||
rows, cols = int(buf.shape[0]), int(buf.shape[1])
|
||||
if rows not in (64, 128):
|
||||
return None
|
||||
dtype = buf.dtype
|
||||
layout_c = buf.layout.canonicalize()
|
||||
for shape in _TCGEN05_COL_FACTOR_FP32:
|
||||
try:
|
||||
cand = tcgen05_atom_layout(shape, (rows, cols), dtype).canonicalize()
|
||||
except ValueError:
|
||||
continue
|
||||
try:
|
||||
tvm.ir.assert_structural_equal(layout_c, cand)
|
||||
except (AssertionError, ValueError):
|
||||
continue
|
||||
# Recover rep from cols (same arithmetic the factory uses).
|
||||
elem_per_32b = 32 // DataType(dtype).bits
|
||||
rep = cols // (_TCGEN05_COL_FACTOR_FP32[shape] * elem_per_32b)
|
||||
return shape, rep, rows
|
||||
return None
|
||||
|
||||
|
||||
def _classify_tmem_datapath(tmem_buf):
|
||||
"""Return ``"D"`` / ``"F"`` if ``tmem_buf.layout`` matches a known tcgen05
|
||||
datapath (PTX ISA §9.7.16.10.5), else ``None``.
|
||||
|
||||
Layout D (M=128, identity row→lane) is the default returned by
|
||||
``_default_tmem_layout``. Layout F (M=64 non-``.ws``, scattered) is the
|
||||
explicit opt-in produced by ``tmem_pool.alloc(..., datapath="F")``.
|
||||
The dispatch uses this to pair each ``.16x*b`` / ``.32x32b`` atom with a
|
||||
compatible layout — see ``_check_tmem_layout_for_atom``.
|
||||
"""
|
||||
if tmem_buf.layout is None:
|
||||
return None
|
||||
buf_layout = tmem_buf.layout.canonicalize()
|
||||
rows = int(tmem_buf.shape[0])
|
||||
if rows == 128:
|
||||
cand = tmem_datapath_layout("D", 128, tmem_buf.shape[1]).canonicalize()
|
||||
try:
|
||||
tvm.ir.assert_structural_equal(buf_layout, cand)
|
||||
return "D"
|
||||
except (AssertionError, ValueError):
|
||||
return None
|
||||
if rows == 64:
|
||||
cand = tmem_datapath_layout("F", 64, tmem_buf.shape[1]).canonicalize()
|
||||
try:
|
||||
tvm.ir.assert_structural_equal(buf_layout, cand)
|
||||
return "F"
|
||||
except (AssertionError, ValueError):
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
# Compatibility matrix between the TMEM buffer's datapath layout and the
|
||||
# tcgen05 ld/st atom requested by ``T.copy_async``:
|
||||
#
|
||||
# datapath x atom | accepted? | rationale
|
||||
# ---------------------------- | --------- | --------------------------------
|
||||
# D (M=128 full) x .32x32b | yes | full 128 lanes, all 32 per warp
|
||||
# D (M=128 full) x .16x*b M=64| yes | reads first half-slab (lanes
|
||||
# | | 0..15 of each warp partition)
|
||||
# | | — the rest of acc is wasted
|
||||
# | | for this atom but valid data
|
||||
# D (M=128 full) x .16x*b M=128| yes | reads all 128 lanes via row=0
|
||||
# | | and row=16 PTX issues
|
||||
# F (M=64 scatter)x .16x*b M=64| yes | canonical pairing - F's row
|
||||
# | | indexing matches the atom's
|
||||
# | | scatter access
|
||||
# F (M=64 scatter)x .16x*b M=128| no | F only writes the low slab; the
|
||||
# | | high slab (row=16) is garbage
|
||||
# F (M=64 scatter)x .32x32b | no | F only utilizes 16 of each
|
||||
# | | warp's 32 lanes
|
||||
_TMEM_ATOM_COMPAT = {
|
||||
("D", "32x32b", 128): True,
|
||||
("D", "16x*b", 64): True,
|
||||
("D", "16x*b", 128): True,
|
||||
("F", "32x32b", 128): False,
|
||||
("F", "16x*b", 64): True,
|
||||
("F", "16x*b", 128): False,
|
||||
}
|
||||
|
||||
|
||||
def _check_tmem_layout_for_atom(tmem_buf, atom_kind, frag_rows):
|
||||
"""Raise ``ValueError`` if the TMEM buffer's datapath layout is
|
||||
incompatible with the requested ``tcgen05`` atom.
|
||||
|
||||
``atom_kind`` is ``"32x32b"`` or ``"16x*b"``; ``frag_rows`` is the
|
||||
register-side fragment row count (128 for ``.32x32b`` and ``.16x*b``
|
||||
M=128 variants, 64 for ``.16x*b`` M=64). If the buffer's layout is
|
||||
unrecognized (i.e. it isn't Layout D or Layout F), the dispatch falls
|
||||
back to the structural assertions below.
|
||||
"""
|
||||
datapath = _classify_tmem_datapath(tmem_buf)
|
||||
if datapath is None:
|
||||
return None
|
||||
allowed = _TMEM_ATOM_COMPAT.get((datapath, atom_kind, frag_rows), False)
|
||||
if not allowed:
|
||||
raise ValueError(
|
||||
f"tcgen05 dispatch: TMEM buffer with datapath={datapath!r} is "
|
||||
f"incompatible with atom={atom_kind!r} (frag_rows={frag_rows}). "
|
||||
f"See PTX ISA §9.7.16.10.5 for datapath/atom pairings; the "
|
||||
f"buffer was allocated via tmem_pool.alloc(..., "
|
||||
f"datapath={datapath!r})."
|
||||
)
|
||||
return datapath
|
||||
|
||||
|
||||
def copy_tmem_local_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
dst_buffer_region, src_buffer_region = op_call.dst, op_call.src
|
||||
dst: Buffer = dst_buffer_region.buffer
|
||||
src: Buffer = src_buffer_region.buffer
|
||||
|
||||
if src.scope() == "tmem" and dst.scope() == "local":
|
||||
direction = "tmem2local"
|
||||
tmem_region, local_region = src_buffer_region, dst_buffer_region
|
||||
elif src.scope() == "local" and dst.scope() == "tmem":
|
||||
direction = "local2tmem"
|
||||
local_region, tmem_region = src_buffer_region, dst_buffer_region
|
||||
else:
|
||||
raise ValueError(f"Unsupported src scope {src.scope()} and dst scope {dst.scope()}")
|
||||
|
||||
tmem_buf, local_buf = tmem_region.buffer, local_region.buffer
|
||||
|
||||
assert tmem_buf.layout is not None
|
||||
assert local_buf.layout is not None
|
||||
assert tmem_buf.dtype == local_buf.dtype
|
||||
assert tmem_buf.allocated_addr is not None
|
||||
|
||||
analyzer = Analyzer()
|
||||
elem_size = DataType(local_buf.dtype).bits
|
||||
elem_per_32b = 32 // elem_size
|
||||
assert len(local_buf.shape) == len(tmem_buf.shape) == 2
|
||||
|
||||
# Try the .16x* (M=64) path first by structural-matching the register-side
|
||||
# layout against ``tcgen05_atom_layout(instr_shape, (64, K), dtype)``. The
|
||||
# TMEM-side layout is the standard (128, W):(1@TLane, 1@TCol); the M=64
|
||||
# fragment lives at lanes 0..15 of each warp's accessible slab (per PTX
|
||||
# 9.7.16.8.1), so each warp issues with row_offset=0 and collectively the
|
||||
# 4 warps cover all 64 rows.
|
||||
atom_match = _match_tcgen05_atom_layout(local_buf)
|
||||
|
||||
if atom_match is not None:
|
||||
shape, num, frag_rows = atom_match
|
||||
return _emit_16xnb_path(
|
||||
shape=shape,
|
||||
num=num,
|
||||
frag_rows=frag_rows,
|
||||
direction=direction,
|
||||
tmem_buf=tmem_buf,
|
||||
local_buf=local_buf,
|
||||
tmem_region=tmem_region,
|
||||
local_region=local_region,
|
||||
elem_per_32b=elem_per_32b,
|
||||
analyzer=analyzer,
|
||||
)
|
||||
|
||||
# Fall through to the existing .32x32b (M=128) path.
|
||||
return _emit_32x32b_path(
|
||||
direction=direction,
|
||||
tmem_buf=tmem_buf,
|
||||
local_buf=local_buf,
|
||||
tmem_region=tmem_region,
|
||||
local_region=local_region,
|
||||
elem_per_32b=elem_per_32b,
|
||||
analyzer=analyzer,
|
||||
)
|
||||
|
||||
|
||||
def _emit_32x32b_path(
|
||||
*, direction, tmem_buf, local_buf, tmem_region, local_region, elem_per_32b, analyzer
|
||||
) -> PrimFunc:
|
||||
"""Original M=128 fragment path using ``tcgen05.{ld,st}.32x32b.xN``."""
|
||||
# local: 128xWIDTH <-> tmem: 128xSHAPE[1]
|
||||
# ``.32x32b`` accesses 32 lanes per warp — the full warp partition — so
|
||||
# the TMEM buffer must be Layout D (M=128 full datapath). Reject Layout F.
|
||||
_check_tmem_layout_for_atom(tmem_buf, "32x32b", 128)
|
||||
assert analyzer.can_prove_equal(local_buf.shape[0], 128)
|
||||
assert analyzer.can_prove_equal(tmem_buf.shape[0], 128)
|
||||
|
||||
# Check width is valid for 32x32b, and determine num
|
||||
width = local_region.region[1].extent
|
||||
candidates = [1, 2, 4, 8, 16, 32, 64, 128]
|
||||
|
||||
if not analyzer.can_prove_equal(tvm.tirx.floormod(width, elem_per_32b), 0):
|
||||
raise ValueError(f"Width {width} is not valid for tcgen05.ld/st with shape 32x32b")
|
||||
|
||||
num = None
|
||||
for n in candidates:
|
||||
if analyzer.can_prove_equal(tvm.tirx.floordiv(width, elem_per_32b), n):
|
||||
num = n
|
||||
break
|
||||
else:
|
||||
raise ValueError(f"Width {width} is not valid for tcgen05.ld/st with shape 32x32b")
|
||||
|
||||
tmem_st, tmem_extent = get_st_extent(tmem_region)
|
||||
local_st, local_extent = get_st_extent(local_region)
|
||||
# tmem layout (128, WIDTH):(1@TLane, 1@TCol)
|
||||
tmem_layout = TileLayout(S[(128, tmem_buf.shape[1]) : (1 @ TLane, 1 @ TCol)]).canonicalize()
|
||||
# local layout
|
||||
TileLayout(S[(128, width) : (1 @ tid_in_wg, 1)]).canonicalize()
|
||||
|
||||
tvm.ir.assert_structural_equal(tmem_buf.layout.canonicalize(), tmem_layout)
|
||||
# local: [0:128, 0:WIDTH] <-> tmem: [0:128, st:st+WIDTH]
|
||||
assert analyzer.can_prove_equal(tmem_st[0], 0)
|
||||
assert analyzer.can_prove_equal(tmem_extent[0], 128)
|
||||
|
||||
assert analyzer.can_prove_equal(local_st[0], 0)
|
||||
assert analyzer.can_prove_equal(local_extent[0], 128)
|
||||
|
||||
offset = tmem_st[1]
|
||||
assert analyzer.can_prove_equal(tvm.tirx.floormod(offset, elem_per_32b), 0)
|
||||
offset_32b = tvm.tirx.floordiv(offset, elem_per_32b)
|
||||
assert analyzer.can_prove_equal(tmem_extent[1], width), (
|
||||
f"tmem_extent[1]: {tmem_extent[1]}, width: {width}"
|
||||
)
|
||||
|
||||
# assert analyzer.can_prove_equal(local_st[1], 0)
|
||||
assert analyzer.can_prove_equal(local_extent[1], width)
|
||||
|
||||
op = T.ptx.tcgen05.ld if direction == "tmem2local" else T.ptx.tcgen05.st
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
local_storage = local_buf.view(local_buf.shape[1] * elem_per_32b, layout=TileLayout(S[num * elem_per_32b])) # noqa: E501
|
||||
local_32b = local_storage.view("uint32")
|
||||
op(tmem_buf.allocated_addr[0], *[local_32b[local_st[1] // elem_per_32b+i] for i in range(num)], shape="32x32b", num=num, row=0, col=offset_32b) # noqa: E501
|
||||
# fmt: on
|
||||
return impl
|
||||
|
||||
|
||||
def _emit_16xnb_path(
|
||||
*,
|
||||
shape,
|
||||
num,
|
||||
frag_rows,
|
||||
direction,
|
||||
tmem_buf,
|
||||
local_buf,
|
||||
tmem_region,
|
||||
local_region,
|
||||
elem_per_32b,
|
||||
analyzer,
|
||||
) -> PrimFunc:
|
||||
"""``.16x*b`` fragment path using ``tcgen05.{ld,st}.<shape>.x<num>`` (one
|
||||
of ``.16x64b``, ``.16x128b``, ``.16x256b``).
|
||||
|
||||
Each of the warpgroup's 4 warps issues the atom with ``row_offset=0`` to
|
||||
cover lanes 0..15 of its 32-lane TMEM partition (one 16-row slab); the
|
||||
four warps collectively span M=64 rows. When ``frag_rows == 128`` the
|
||||
dispatch emits a second issue with ``row_offset=16`` to also cover lanes
|
||||
16..31 of each warp's partition, doubling the fragment's row coverage to
|
||||
M=128. The two atoms share the same column footprint; the layout factory
|
||||
surfaces the combined per-thread register vector with the second slab's
|
||||
regs in the high half of the m-axis (so the dispatch can split regs
|
||||
contiguously between the two PTX calls).
|
||||
"""
|
||||
# Per-atom column footprint in fp32 columns:
|
||||
# .16x64b → 2N .16x128b → 4N .16x256b → 8N
|
||||
col_factor_fp32 = {"16x64b": 2, "16x128b": 4, "16x256b": 8}[shape]
|
||||
# Per-thread register count per 16-row slab (in 32-bit units):
|
||||
# .16x64b.xN → N .16x128b.xN → 2N .16x256b.xN → 4N
|
||||
regs_per_thread_per_slab = {"16x64b": num, "16x128b": 2 * num, "16x256b": 4 * num}[shape]
|
||||
n_slabs = frag_rows // 64 # 1 for M=64, 2 for M=128
|
||||
assert n_slabs in (1, 2)
|
||||
regs_per_thread = regs_per_thread_per_slab * n_slabs
|
||||
# Logical column width that the local buffer view exposes (in element units).
|
||||
width_elems = col_factor_fp32 * num * elem_per_32b
|
||||
# Per-thread storage in element units (same total bits as the register vector).
|
||||
per_thread_elems = regs_per_thread * elem_per_32b
|
||||
|
||||
# Local-side: shape (frag_rows, K_cols)
|
||||
assert analyzer.can_prove_equal(local_buf.shape[0], frag_rows), (
|
||||
f".16x*b path expects local_buf rows={frag_rows}, got {local_buf.shape[0]}"
|
||||
)
|
||||
assert analyzer.can_prove_equal(local_buf.shape[1], width_elems), (
|
||||
f".16x*b path expects local_buf cols={width_elems}, got {local_buf.shape[1]}"
|
||||
)
|
||||
|
||||
# TMEM-side: structurally classify the buffer's datapath (D or F) and
|
||||
# reject incompatible pairings. The PTX is identical in either case (the
|
||||
# warp partition rule and the atom's lane access pattern are baked into
|
||||
# the hardware); the layout classification just keeps the buffer's
|
||||
# logical row indexing in sync with the physical TMEM occupation.
|
||||
datapath = _check_tmem_layout_for_atom(tmem_buf, "16x*b", frag_rows)
|
||||
|
||||
if datapath == "F":
|
||||
# Layout F: buffer shape (64, W), scattered row→lane.
|
||||
assert analyzer.can_prove_equal(tmem_buf.shape[0], 64), (
|
||||
f".16x*b Layout F expects tmem_buf rows=64, got {tmem_buf.shape[0]}"
|
||||
)
|
||||
tmem_rows = 64
|
||||
else:
|
||||
# Layout D (or untagged legacy buffers): shape (128, W), identity.
|
||||
# The legacy structural check below still fires for untagged buffers
|
||||
# so we don't silently accept arbitrary layouts.
|
||||
assert analyzer.can_prove_equal(tmem_buf.shape[0], 128), (
|
||||
f".16x*b path expects tmem_buf rows=128, got {tmem_buf.shape[0]}"
|
||||
)
|
||||
if datapath is None:
|
||||
tmem_layout = TileLayout(
|
||||
S[(128, tmem_buf.shape[1]) : (1 @ TLane, 1 @ TCol)]
|
||||
).canonicalize()
|
||||
tvm.ir.assert_structural_equal(tmem_buf.layout.canonicalize(), tmem_layout)
|
||||
tmem_rows = 128
|
||||
|
||||
tmem_st, tmem_extent = get_st_extent(tmem_region)
|
||||
local_st, local_extent = get_st_extent(local_region)
|
||||
|
||||
# Rows must span the full frag. The COLUMN extent may be a sub-multiple of
|
||||
# the atom's full width ``width_elems`` — i.e. a per-chunk column slice of a
|
||||
# wider frag (e.g. an epilogue that loads one big (128, MMA_N) frag in
|
||||
# EPI_TILE-wide chunks). The atom layout maps consecutive columns to
|
||||
# consecutive registers within each slab, so a column slice occupies a
|
||||
# contiguous register window; we emit ``num_eff`` (the slice's atom rep) at
|
||||
# the slab base + the column's register offset. When the slice IS the full
|
||||
# atom (the common case), num_eff == num and reg offset == 0 (no change).
|
||||
assert analyzer.can_prove_equal(local_st[0], 0)
|
||||
assert analyzer.can_prove_equal(local_extent[0], frag_rows)
|
||||
assert analyzer.can_prove_equal(tmem_st[0], 0)
|
||||
assert analyzer.can_prove_equal(tmem_extent[0], frag_rows)
|
||||
# local and tmem column slices must match and divide the atom's full width.
|
||||
assert analyzer.can_prove_equal(local_extent[1], tmem_extent[1])
|
||||
slice_w = int(local_extent[1])
|
||||
assert width_elems % slice_w == 0, f"slice width {slice_w} must divide atom width {width_elems}"
|
||||
num_eff = num * slice_w // width_elems
|
||||
regs_eff = regs_per_thread_per_slab * slice_w // width_elems
|
||||
del tmem_rows # only used for the structural check above
|
||||
|
||||
col_off = tmem_st[1]
|
||||
assert analyzer.can_prove_equal(tvm.tirx.floormod(col_off, elem_per_32b), 0)
|
||||
col_off_32b = tvm.tirx.floordiv(col_off, elem_per_32b)
|
||||
local_col_off = local_st[1]
|
||||
assert analyzer.can_prove_equal(tvm.tirx.floormod(local_col_off, elem_per_32b), 0)
|
||||
local_col_off_elems = local_col_off
|
||||
|
||||
is_load = direction == "tmem2local"
|
||||
op = T.ptx.tcgen05.ld if is_load else T.ptx.tcgen05.st
|
||||
# We intentionally do *not* emit ``.pack::16b`` / ``.unpack::16b`` for
|
||||
# 16-bit dtypes. That qualifier would store one 16-bit element per 32-bit
|
||||
# TMEM cell (LOW half only, HIGH half wasted) — fine for some CUTLASS
|
||||
# epilogues but a 2x TMEM waste vs. the existing ``.32x32b`` convention,
|
||||
# which packs two 16-bit elements per cell. By using the plain ``.b32``
|
||||
# form we keep TMEM dense (2 elements per 32-bit cell); the per-thread
|
||||
# register file holds two packed 16-bit values per 32-bit register, and
|
||||
# the layout factory's iters describe that packing.
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
# Per-thread 1-D flat view of the local storage, then a uint32 view
|
||||
# for the register-pointer arguments of the PTX builtin.
|
||||
local_storage = local_buf.view(per_thread_elems, layout=TileLayout(S[per_thread_elems]))
|
||||
local_32b = local_storage.view("uint32")
|
||||
# Register offset of the column slice within each slab. The old
|
||||
# ``local_col_off // elem_per_32b`` is only correct when the slice IS the
|
||||
# full atom; in general consecutive columns advance registers at the rate
|
||||
# (regs_per_thread_per_slab / width_elems). For a full-atom load the
|
||||
# offset is 0 either way, so existing callers are unaffected.
|
||||
local_reg_base = local_col_off_elems * regs_per_thread_per_slab // width_elems
|
||||
for slab in range(n_slabs):
|
||||
reg_base = slab * regs_per_thread_per_slab
|
||||
op(
|
||||
tmem_buf.allocated_addr[0],
|
||||
*[local_32b[local_reg_base + reg_base + i] for i in range(regs_eff)],
|
||||
shape=shape, num=num_eff, row=slab * 16, col=col_off_32b,
|
||||
)
|
||||
# fmt: on
|
||||
return impl
|
||||
|
||||
|
||||
# === Variant: copy_async/tmem<->local (priority=10) ===
|
||||
#
|
||||
# When: one buffer is in tmem (tensor memory, Blackwell SM100+) and the other
|
||||
# is in local scope, at warpgroup exec scope.
|
||||
#
|
||||
# Emits: T.ptx.tcgen05.ld / T.ptx.tcgen05.st (async). The caller is
|
||||
# responsible for issuing the matching ``T.ptx.tcgen05.wait.ld`` /
|
||||
# ``T.ptx.tcgen05.wait.st`` when synchronization is required.
|
||||
@register_dispatch(
|
||||
"copy_async",
|
||||
"cuda",
|
||||
variant="tmem<->local",
|
||||
priority=10,
|
||||
when=[
|
||||
predicate("validate_copy_op", _is_valid_copy),
|
||||
predicate("exec_scope", exec_scope_ok, expected_scopes=["warpgroup"]),
|
||||
predicate(
|
||||
"storage_scope", _scope_allowed, allowed_pairs=[("tmem", "local"), ("local", "tmem")]
|
||||
),
|
||||
],
|
||||
)
|
||||
def copy_async_schedule_tmem_local_async(
|
||||
op_call: TilePrimitiveCall, sctx: DispatchContext
|
||||
) -> PrimFunc:
|
||||
return copy_tmem_local_impl(op_call, sctx)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,78 @@
|
||||
# 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 helpers for copy_async operator dispatch variants.
|
||||
|
||||
The TMA-specific lowering moved to ``tma.py``. What remains here are the tiny
|
||||
layout helpers other variants (e.g. ``dsmem.py``) still import.
|
||||
"""
|
||||
|
||||
from tvm.arith import Analyzer
|
||||
from tvm.tirx.layout import ComposeLayout, Layout, S, SwizzleLayout, TileLayout
|
||||
|
||||
|
||||
def find_contiguous_region(layout: TileLayout) -> tuple:
|
||||
"""Return the maximal stride-1 contiguous memory-shard chain.
|
||||
|
||||
Starts from stride==1 and repeatedly picks the shard whose stride equals
|
||||
the running product of extents, stopping when no shard matches. Returns
|
||||
the maximal chain; callers that need a shorter prefix should take one
|
||||
themselves (e.g. to satisfy TMA's rank<=5 or a per-path reduction step).
|
||||
Stride/extent comparisons go through an ``Analyzer`` so symbolic strides
|
||||
work.
|
||||
"""
|
||||
|
||||
analyzer = Analyzer()
|
||||
memory_shards = [
|
||||
(i, s)
|
||||
for i, s in enumerate(layout.shard)
|
||||
if s.axis.is_memory() and not analyzer.can_prove_equal(s.extent, 1)
|
||||
]
|
||||
if not memory_shards:
|
||||
return [], 1
|
||||
|
||||
contiguous_indices: list[int] = []
|
||||
contiguous_extent = 1
|
||||
expected_stride = 1
|
||||
consumed: set[int] = set()
|
||||
|
||||
while True:
|
||||
for idx, shard in memory_shards:
|
||||
if idx in consumed:
|
||||
continue
|
||||
if analyzer.can_prove_equal(shard.stride, expected_stride):
|
||||
consumed.add(idx)
|
||||
contiguous_indices.append(idx)
|
||||
contiguous_extent *= shard.extent
|
||||
expected_stride = contiguous_extent
|
||||
break
|
||||
else:
|
||||
break
|
||||
|
||||
if not contiguous_indices:
|
||||
return [], 0
|
||||
return contiguous_indices, contiguous_extent
|
||||
|
||||
|
||||
def to_tile_layout(layout: Layout, shape: list[int]) -> TileLayout:
|
||||
"""Normalize any layout kind to a TileLayout for pointer arithmetic."""
|
||||
|
||||
if isinstance(layout, ComposeLayout):
|
||||
return layout.tile_layout
|
||||
if isinstance(layout, SwizzleLayout):
|
||||
return TileLayout(S[tuple(shape)])
|
||||
return layout
|
||||
@@ -0,0 +1,38 @@
|
||||
# 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.
|
||||
|
||||
"""CUDA elementwise dispatch.
|
||||
|
||||
Split by storage scope to mirror ``cuda/copy/``:
|
||||
|
||||
reg.py — operands all in ``local`` (registers) → induced partition
|
||||
smem.py — operands all in ``shared*`` → synthesized partition
|
||||
|
||||
Each op in ``ops.ALL_OPS`` is registered under both variants. Per-op packed
|
||||
PTX/CUDA intrinsics live in ``vec_emit/`` (``binary_f32x2`` / ``cast_vec2``
|
||||
/ ``fma_f32x2``) and are attached to the relevant ``OpSpec.vec_impls``.
|
||||
"""
|
||||
|
||||
from .register import *
|
||||
|
||||
# Suppress submodule-attribute leakage. Without an explicit ``__all__`` here,
|
||||
# ``from tvm.backend.cuda.operator.tile_primitive.elementwise import *`` (run by
|
||||
# tile_primitive/__init__.py) re-exports the implicit submodule attributes
|
||||
# (``ops``, ``reg``, ``smem``, ``vec_emit``) — and ``ops`` in particular
|
||||
# shadows the top-level ``tile_primitive/ops.py`` (BinaryReduce / UnaryReduce
|
||||
# / ...) when downstream code does ``from tile_primitive import ops``.
|
||||
__all__: list[str] = []
|
||||
@@ -0,0 +1,412 @@
|
||||
# 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 layout / vec-selection / emit helpers for ``reg.py`` and ``smem.py``.
|
||||
|
||||
Borrows directly from ``cuda/copy/reg.py`` (induced partition) and
|
||||
``cuda/copy/_common.py`` (synthesized partition), extended to N operands.
|
||||
|
||||
The dispatch split mirrors copy:
|
||||
reg.py — all operands in ``local`` → partition induced by anchor's layout
|
||||
smem.py — all operands in ``shared*`` → partition synthesized from ``sctx.intra``
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import operator
|
||||
|
||||
from tvm.arith.analyzer import Analyzer
|
||||
from tvm.runtime import DataType
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import BufferRegion
|
||||
from tvm.tirx.layout import Axis, Iter, TileLayout
|
||||
|
||||
from ..common import get_indices, get_st_extent
|
||||
|
||||
# Re-use copy's primitives (PR-640) — same algorithm, same scope_id machinery.
|
||||
from ..copy._common import _TID_AXIS_FOR_SCOPE, _extract_tile, _thread_cnt
|
||||
from ..copy.reg import _all_threads_active, _axis_decl, _compute_perm_r
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Plan helpers
|
||||
# -----------------------------------------------------------------------------
|
||||
def buffer_regions(plan) -> list[BufferRegion]:
|
||||
"""All BufferRegion args (dst + buffer-region srcs), in plan order."""
|
||||
out: list[BufferRegion] = [plan.dst]
|
||||
for s in plan.srcs:
|
||||
if s.buf_region is not None:
|
||||
out.append(s.buf_region)
|
||||
return out
|
||||
|
||||
|
||||
def dtype_name(dtype) -> str:
|
||||
dtype_obj = getattr(dtype, "dtype", None)
|
||||
if dtype_obj is not None:
|
||||
return str(dtype_obj)
|
||||
return str(dtype)
|
||||
|
||||
|
||||
def dtype_bits(dtype) -> int:
|
||||
return DataType(dtype_name(dtype)).bits
|
||||
|
||||
|
||||
def compute_dtype_of(plan) -> str:
|
||||
"""Widest dtype in bits across dst + buffer/scalar srcs (dst breaks ties)."""
|
||||
candidates = [dtype_name(plan.dst.buffer.dtype)]
|
||||
for s in plan.srcs:
|
||||
if s.buf_region is not None:
|
||||
candidates.append(dtype_name(s.buf_region.buffer.dtype))
|
||||
elif s.scalar is not None:
|
||||
candidates.append(scalar_dtype(s.scalar))
|
||||
widest = candidates[0]
|
||||
widest_bits = DataType(widest).bits
|
||||
for d in candidates[1:]:
|
||||
b = DataType(d).bits
|
||||
if b > widest_bits:
|
||||
widest, widest_bits = d, b
|
||||
return widest
|
||||
|
||||
|
||||
def scalar_dtype(scalar) -> str:
|
||||
dtype = getattr(scalar, "dtype", None)
|
||||
if dtype is not None:
|
||||
return str(dtype)
|
||||
ty = getattr(scalar, "ty", None)
|
||||
if ty is None and hasattr(scalar, "expr_ty"):
|
||||
ty = scalar.expr_ty()
|
||||
dtype = getattr(ty, "dtype", None)
|
||||
if dtype is None:
|
||||
raise AttributeError(f"{type(scalar).__name__} has no dtype-bearing PrimType")
|
||||
return str(dtype)
|
||||
|
||||
|
||||
def n_elements(buf_region: BufferRegion) -> int:
|
||||
_, ext = get_st_extent(buf_region)
|
||||
return functools.reduce(operator.mul, ext, 1)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Anchor selection (reg.py)
|
||||
# -----------------------------------------------------------------------------
|
||||
def pick_anchor(plan) -> BufferRegion:
|
||||
"""Anchor is always ``plan.dst`` — every operand must have a layout
|
||||
(enforced by predicate); dst's layout drives iteration. No choice to make.
|
||||
"""
|
||||
return plan.dst
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Broadcast support (NumPy-style right-aligned, anchor = result shape)
|
||||
# -----------------------------------------------------------------------------
|
||||
def _tensor_shape_of(region) -> tuple[int, ...]:
|
||||
"""Per-dim region extent (post-slice tensor shape, NOT layout shape).
|
||||
|
||||
Accepts either ``[(start, end), ...]`` pairs (as built locally from a
|
||||
``BufferRegion``) or the ``BufferRegion.region`` sequence of ``Range``
|
||||
objects directly. ``Range.extent`` is already simplified by the
|
||||
front-end, so we avoid computing ``end - start`` on raw Expr (which
|
||||
yields an un-simplified ``Sub`` and breaks ``int(...)``).
|
||||
"""
|
||||
out = []
|
||||
a = Analyzer()
|
||||
for r in region:
|
||||
if hasattr(r, "extent"):
|
||||
ext = r.extent
|
||||
else:
|
||||
start, end = r
|
||||
ext = a.simplify(end - start)
|
||||
out.append(int(ext))
|
||||
return tuple(out)
|
||||
|
||||
|
||||
def shape_broadcast_compat(op_shape, anchor_shape) -> tuple[bool, str | None]:
|
||||
"""NumPy-style: right-align op against anchor; per-dim extent must equal
|
||||
anchor's or be 1. anchor is the result shape; op broadcasts TO anchor.
|
||||
"""
|
||||
pad = len(anchor_shape) - len(op_shape)
|
||||
if pad < 0:
|
||||
return False, f"op rank {len(op_shape)} > anchor rank {len(anchor_shape)}"
|
||||
for d in range(len(op_shape)):
|
||||
e_op = int(op_shape[d])
|
||||
e_a = int(anchor_shape[pad + d])
|
||||
if e_op != e_a and e_op != 1:
|
||||
return False, f"dim {d}: op extent {e_op} vs anchor {e_a} (need equal or 1)"
|
||||
return True, None
|
||||
|
||||
|
||||
def _broadcast_lift(op_layout, op_tensor_shape, anchor_tensor_shape):
|
||||
"""Lift ``op_layout`` to ``anchor_tensor_shape`` by inserting stride-0
|
||||
iters for padded leading dims and replacing extent-1 buckets with a
|
||||
single stride-0 iter of anchor's extent. Offset and replica list are
|
||||
preserved untouched.
|
||||
|
||||
The lift preserves the physical-address function: new iters have
|
||||
stride 0 so they contribute ``coord * 0 = 0`` to the address
|
||||
regardless of which virtual index is supplied, and dropped extent-1
|
||||
iters contributed ``0 * stride = 0`` already.
|
||||
"""
|
||||
pad = len(anchor_tensor_shape) - len(op_tensor_shape)
|
||||
assert pad >= 0, "shape_broadcast_compat should have rejected this"
|
||||
|
||||
try:
|
||||
grouped, seps = op_layout.group(list(op_tensor_shape))
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
raise ValueError(
|
||||
f"op layout {op_layout} not groupable by tensor shape {op_tensor_shape}: {e}"
|
||||
) from e
|
||||
|
||||
new_shard: list = []
|
||||
# (1) Padded leading dims — one stride-0 iter each.
|
||||
for d in range(pad):
|
||||
new_shard.append(Iter(int(anchor_tensor_shape[d]), 0, Axis.get("m")))
|
||||
# (2) Aligned dims.
|
||||
for d_op in range(len(op_tensor_shape)):
|
||||
e_op = int(op_tensor_shape[d_op])
|
||||
e_a = int(anchor_tensor_shape[pad + d_op])
|
||||
bucket = list(grouped.shard[seps[d_op] : seps[d_op + 1]])
|
||||
if e_op == e_a:
|
||||
new_shard.extend(bucket)
|
||||
elif e_op == 1:
|
||||
new_shard.append(Iter(e_a, 0, Axis.get("m")))
|
||||
else:
|
||||
raise ValueError(
|
||||
f"dim {d_op}: op extent {e_op} vs anchor {e_a}"
|
||||
" (shape_broadcast_compat should have rejected)"
|
||||
)
|
||||
|
||||
return TileLayout.from_iters(new_shard, grouped.replica, grouped.offset)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Shared preprocess: slice each operand by its region, broadcast-lift to
|
||||
# anchor's tensor shape. Output: every operand has a layout whose logical
|
||||
# shape equals ``anchor_tensor_shape``. Reg / smem diverge from here.
|
||||
# -----------------------------------------------------------------------------
|
||||
def preprocess_operand(op_br, anchor_tshape):
|
||||
"""Slice ``op_br``'s buffer layout by region (region offset absorbed into
|
||||
layout.offset), then broadcast-lift to ``anchor_tshape`` if shapes differ.
|
||||
|
||||
Raises ``ValueError`` if the lift is not broadcast-compatible (caller
|
||||
should have verified via ``shape_broadcast_compat`` in the predicate).
|
||||
"""
|
||||
op_layout = op_br.buffer.layout
|
||||
op_shape = op_br.buffer.shape
|
||||
op_region = [(r.min, r.min + r.extent) for r in op_br.region]
|
||||
sliced = op_layout.slice(list(op_shape), op_region).canonicalize()
|
||||
sliced = _extract_tile(sliced, op_region)
|
||||
op_tshape = _tensor_shape_of(op_br.region)
|
||||
if op_tshape != tuple(anchor_tshape):
|
||||
sliced = _broadcast_lift(sliced, op_tshape, anchor_tshape)
|
||||
return sliced
|
||||
|
||||
|
||||
def preprocess_operands(plan):
|
||||
"""Shared entry for reg.py and smem.py: returns
|
||||
``(anchor_tensor_shape, {op_br: sliced_lifted_layout})``.
|
||||
|
||||
Every output layout has logical shape == ``anchor_tensor_shape``.
|
||||
Broadcast iters carry stride 0 with the default mem axis. Reg's induced
|
||||
partition (`align_operands_to_anchor`) and smem's synthesized partition
|
||||
both build on this output.
|
||||
"""
|
||||
anchor_tshape = _tensor_shape_of(plan.dst.region)
|
||||
out: dict = {}
|
||||
for br in buffer_regions(plan):
|
||||
out[br] = preprocess_operand(br, anchor_tshape)
|
||||
return anchor_tshape, out
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Multi-operand layout alignment for reg.py (induced)
|
||||
# -----------------------------------------------------------------------------
|
||||
def _align_layouts_no_post_canon(r_layout, r_shape, r_region, s_layout, s_shape, s_region):
|
||||
"""Variant of copy ``reg.py:align_layouts_raw`` that omits the final
|
||||
``canonicalize()`` on ``r_p``.
|
||||
|
||||
Copy's version returns ``r_p = r.permute_dims(perm).canonicalize()`` —
|
||||
that post-permute canonicalize can fuse adjacent iters (e.g. wgmma
|
||||
layout's 5 iters collapse to 2), but ``s_seps`` is built from
|
||||
``perm`` of length ``len(r.shard) pre-canon``. The two lengths then
|
||||
disagree and ``s_p.shard[s_seps[k]:s_seps[k+1]]`` indexes into the
|
||||
wrong sub-range.
|
||||
|
||||
Dropping the final canonicalize keeps ``r_p.shard`` and ``s_seps`` in
|
||||
1-to-1 correspondence. Copy's tests don't hit this because R is
|
||||
typically 1D and doesn't fuse further after permute.
|
||||
"""
|
||||
r = r_layout.slice(list(r_shape), r_region).canonicalize()
|
||||
s = s_layout.slice(list(s_shape), s_region).canonicalize()
|
||||
s = _extract_tile(s, s_region)
|
||||
# Broadcast lift: when op's post-slice tensor shape != anchor's, expand
|
||||
# s via stride-0 iters so group() below can partition along anchor's
|
||||
# iter structure. Legality must be enforced upstream by the predicate.
|
||||
r_tshape = _tensor_shape_of(r_region)
|
||||
s_tshape = _tensor_shape_of(s_region)
|
||||
if s_tshape != r_tshape:
|
||||
s = _broadcast_lift(s, s_tshape, r_tshape)
|
||||
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_p = r.permute_dims(perm) # NO post-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
|
||||
|
||||
|
||||
def align_operands_to_anchor(anchor_br, layout_others_br):
|
||||
"""Align every layout-bearing non-anchor operand to ``anchor_br``.
|
||||
|
||||
Returns ``(anchor_p, per_op_aligned)`` where ``per_op_aligned[op_br] =
|
||||
(op_p, op_seps)``. Trivial-layout operands are NOT included here —
|
||||
caller indexes them directly via their region. Scalar srcs likewise
|
||||
live outside this map.
|
||||
|
||||
Caller must enter ``with sctx.target:`` so ``canonicalize()`` runs the
|
||||
scope-aware fusers (e.g. laneid+wid_in_wg → tid_in_wg).
|
||||
|
||||
Uses ``_align_layouts_no_post_canon`` (not copy's ``align_layouts_raw``
|
||||
directly) so ``anchor_p.shard`` length matches ``op_seps`` groupings.
|
||||
"""
|
||||
anchor_layout = anchor_br.buffer.layout
|
||||
anchor_shape = anchor_br.buffer.shape
|
||||
anchor_region = [(r.min, r.min + r.extent) for r in anchor_br.region]
|
||||
|
||||
per_op_aligned: dict = {}
|
||||
anchor_p = None
|
||||
if not layout_others_br:
|
||||
# Just slice + permute anchor alone (no post-canon — keep iters
|
||||
# 1-to-1 with how they'd appear with srcs).
|
||||
r = anchor_layout.slice(list(anchor_shape), anchor_region).canonicalize()
|
||||
perm = _compute_perm_r(r)
|
||||
anchor_p = r.permute_dims(perm)
|
||||
return anchor_p, per_op_aligned
|
||||
|
||||
for op_br in layout_others_br:
|
||||
op_layout = op_br.buffer.layout
|
||||
op_shape = op_br.buffer.shape
|
||||
op_region = [(r.min, r.min + r.extent) for r in op_br.region]
|
||||
r_p, op_p, op_seps = _align_layouts_no_post_canon(
|
||||
anchor_layout,
|
||||
anchor_shape,
|
||||
anchor_region,
|
||||
op_layout,
|
||||
op_shape,
|
||||
op_region,
|
||||
)
|
||||
if anchor_p is None:
|
||||
anchor_p = r_p
|
||||
per_op_aligned[op_br] = (op_p, op_seps)
|
||||
return anchor_p, per_op_aligned
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# vec_chunk selection
|
||||
# -----------------------------------------------------------------------------
|
||||
def pick_vec_chunk(spec, op_call, sctx, plan, max_layout_vec_len: int):
|
||||
"""Pick widest ``(vec_chunk, vec_impl)`` such that:
|
||||
- ``vec_impl.vec_len`` divides ``max_layout_vec_len`` AND ``vec_impl.applies(...)``
|
||||
- Or no vec_impl matches → scalar fallback ``(max_layout_vec_len, None)``
|
||||
|
||||
``spec.vec_impls`` is assumed pre-sorted widest-first.
|
||||
"""
|
||||
if max_layout_vec_len <= 0:
|
||||
return 1, None
|
||||
for impl in getattr(spec, "vec_impls", []):
|
||||
if impl.vec_len > max_layout_vec_len:
|
||||
continue
|
||||
if max_layout_vec_len % impl.vec_len != 0:
|
||||
continue
|
||||
ok, _ = impl.applies(op_call, sctx, plan)
|
||||
if ok:
|
||||
return impl.vec_len, impl
|
||||
return max_layout_vec_len, None
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Emit-time helpers
|
||||
# -----------------------------------------------------------------------------
|
||||
def _broadcast_indices(dst_indices, dst_start, dst_extent, op_start, op_ext):
|
||||
"""NumPy-style right-aligned broadcast: derive op's per-dim indices from
|
||||
dst's. For matching extents copies dst's coord (rebased); for op_ext[d]
|
||||
== 1 returns the constant start (the only valid index for that dim).
|
||||
"""
|
||||
pad = len(dst_extent) - len(op_ext)
|
||||
return [
|
||||
(dst_indices[i + pad] - dst_start[i + pad]) + op_start[i]
|
||||
if int(op_ext[i]) != 1
|
||||
else op_start[i]
|
||||
for i in range(len(op_ext))
|
||||
]
|
||||
|
||||
|
||||
def fetch_src_value(src, fused, dst_indices, dst_start, dst_extent):
|
||||
"""Per-element load Expr for one src. Handles buffer / scalar / broadcast srcs."""
|
||||
if src.is_scalar:
|
||||
return src.scalar
|
||||
region = src.buf_region
|
||||
src_st, src_ext = get_st_extent(region)
|
||||
if src.index_fn is not None:
|
||||
idx = src.index_fn(dst_indices, dst_start, dst_extent, src_st, src_ext)
|
||||
elif tuple(int(e) for e in src_ext) != tuple(int(e) for e in dst_extent):
|
||||
# Broadcast — derive src indices from dst's via right-aligned compat.
|
||||
idx = _broadcast_indices(dst_indices, dst_start, dst_extent, src_st, src_ext)
|
||||
else:
|
||||
idx = get_indices(fused, src_st, src_ext)
|
||||
return region.buffer[tuple(idx)]
|
||||
|
||||
|
||||
def emit_scope_sync(scope_kind: str):
|
||||
"""Returns an ``@T.inline`` sync helper matched to the exec scope."""
|
||||
|
||||
@T.inline
|
||||
def sync():
|
||||
if scope_kind == "cta":
|
||||
T.cuda.cta_sync()
|
||||
elif scope_kind == "warpgroup":
|
||||
T.cuda.warpgroup_sync(8)
|
||||
elif scope_kind == "warp":
|
||||
T.cuda.warp_sync()
|
||||
|
||||
return sync
|
||||
|
||||
|
||||
__all__ = [
|
||||
"_TID_AXIS_FOR_SCOPE",
|
||||
"_all_threads_active",
|
||||
"_axis_decl",
|
||||
"_broadcast_indices",
|
||||
"_tensor_shape_of",
|
||||
"_thread_cnt",
|
||||
"align_operands_to_anchor",
|
||||
"buffer_regions",
|
||||
"compute_dtype_of",
|
||||
"dtype_bits",
|
||||
"dtype_name",
|
||||
"emit_scope_sync",
|
||||
"fetch_src_value",
|
||||
"n_elements",
|
||||
"pick_anchor",
|
||||
"pick_vec_chunk",
|
||||
"preprocess_operand",
|
||||
"preprocess_operands",
|
||||
"shape_broadcast_compat",
|
||||
]
|
||||
@@ -0,0 +1,121 @@
|
||||
# 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.
|
||||
|
||||
"""Per-op data model + ALL_OPS registry.
|
||||
|
||||
``OpSpec`` describes one elementwise op. ``VecImpl`` describes one packed-PTX
|
||||
or CUDA-intrinsic emit available for that op (e.g. ``add_f32x2``); a list of
|
||||
these (widest-first) lets ``reg.py``/``smem.py`` pick the widest matching
|
||||
both the layout and the op's available intrinsics, like copy picks
|
||||
``copy_{128,64,32,16,8}b`` based on bit-width and tail contiguity.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from tvm.ir.expr import Expr
|
||||
from tvm.tirx import BufferRegion, TilePrimitiveCall
|
||||
|
||||
|
||||
@dataclass
|
||||
class SrcSpec:
|
||||
"""One operand of an elementwise op.
|
||||
|
||||
Either a ``BufferRegion`` (per-element load) or a scalar ``Expr``.
|
||||
``index_fn``, if given, derives per-element indices for broadcasting srcs:
|
||||
``index_fn(dst_indices, dst_start, dst_extent, src_start, src_extent) -> list[Expr]``
|
||||
Default is the standard ``get_indices`` over the src's own region.
|
||||
"""
|
||||
|
||||
buf_region: BufferRegion | None = None
|
||||
scalar: Expr | None = None
|
||||
index_fn: Callable | None = None
|
||||
|
||||
@property
|
||||
def is_scalar(self) -> bool:
|
||||
return self.scalar is not None
|
||||
|
||||
@property
|
||||
def buffer(self):
|
||||
return self.buf_region.buffer if self.buf_region is not None else None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Plan:
|
||||
"""Parsed elementwise op ready for a schedule to consume."""
|
||||
|
||||
dst: BufferRegion
|
||||
srcs: list[SrcSpec]
|
||||
extras: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VecImpl:
|
||||
"""One packed-vector implementation registered for an op.
|
||||
|
||||
Mirrors the ``copy_{Nb}`` menu in copy: each entry says "I can process
|
||||
``vec_len`` consecutive elements per call". The schedule picks the widest
|
||||
one whose ``vec_len`` divides the layout's contig tail AND whose
|
||||
``applies()`` returns ``True``.
|
||||
"""
|
||||
|
||||
vec_len: int # elements per packed call
|
||||
applies: Callable[[TilePrimitiveCall, Any, Plan], tuple[bool, str | None]]
|
||||
# emit(dst_ptr, src_ptrs, extras) -> Stmt
|
||||
# dst_ptr: typed ptr to ``vec_len`` consecutive dst elements
|
||||
# src_ptrs[i]: typed ptr to ``vec_len`` consecutive src[i] elements,
|
||||
# OR a scalar Expr if src[i].is_scalar.
|
||||
# Runs in Python at @T.prim_func build time -- branching on src kind is a
|
||||
# normal Python ``if``, not a TVMScript shape limitation. This is what
|
||||
# collapses the old 4x2 shape-explosion in schema.py's factories.
|
||||
emit: Callable
|
||||
|
||||
|
||||
@dataclass
|
||||
class OpSpec:
|
||||
"""Metadata for an elementwise op."""
|
||||
|
||||
name: str
|
||||
# parse(op_call) -> (Plan, msg|None); msg explains why parse failed.
|
||||
parse: Callable[[TilePrimitiveCall], tuple[Plan | None, str | None]]
|
||||
# Scalar compute used by the fallback emit path (wrapped in Tx.vectorized).
|
||||
# compute_scalar(src_vals_at_one_idx, extras, dst_dtype) -> Expr
|
||||
compute_scalar: Callable[[list, dict, str], Any]
|
||||
# Optional dtype check on plan.extras (e.g. unary bias/scale dtype agreement).
|
||||
check_extras: Callable | None = None
|
||||
# Widest-first vec impls. Schedule picks first matching layout+applies.
|
||||
vec_impls: list[VecImpl] = field(default_factory=list)
|
||||
|
||||
|
||||
def _build_all_ops() -> dict[str, OpSpec]:
|
||||
"""Aggregate per-family op specs. Deferred imports avoid cycles
|
||||
(vec_emit/* imports VecImpl from this module)."""
|
||||
from .binary import BINARY_OPS
|
||||
from .cast import CAST_OPS
|
||||
from .fma import FMA_OPS
|
||||
from .unary import UNARY_OPS
|
||||
|
||||
return {**UNARY_OPS, **BINARY_OPS, **CAST_OPS, **FMA_OPS}
|
||||
|
||||
|
||||
ALL_OPS: dict[str, OpSpec] = _build_all_ops()
|
||||
|
||||
|
||||
__all__ = ["ALL_OPS", "OpSpec", "Plan", "SrcSpec", "VecImpl"]
|
||||
@@ -0,0 +1,139 @@
|
||||
# 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.
|
||||
|
||||
"""Binary elementwise ops: add / sub / mul / fdiv / maximum.
|
||||
|
||||
Includes constant-lhs commute logic. Broadcasting (extent=1 dims) is
|
||||
handled at the layout level in dispatch's ``_broadcast_lift``, not here —
|
||||
parser just records each src as-is.
|
||||
|
||||
``add``/``sub``/``mul`` attach a ``VecImpl`` for sm_100+ packed f32x2;
|
||||
``fdiv``/``maximum`` have no packed PTX (scalar fallback only — ``max`` lowers
|
||||
to a single ``FMNMX``/``max.f32``, which is exact, so there is no rounding/ftz
|
||||
variant to pack).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import operator
|
||||
from typing import Any
|
||||
|
||||
from tvm.script import tirx as Tx
|
||||
from tvm.tirx import BufferRegion, TilePrimitiveCall
|
||||
|
||||
from ..vec_emit.binary_f32x2 import BINARY_F32X2_IMPLS
|
||||
from . import OpSpec, Plan, SrcSpec
|
||||
|
||||
_COMMUTATIVE = frozenset({"add", "mul", "maximum"})
|
||||
|
||||
|
||||
def _parse_binary_for(op_name: str):
|
||||
"""Build a ``parse(op_call) -> (Plan, msg)`` for a specific binary op."""
|
||||
|
||||
def parse(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
|
||||
_dst: BufferRegion = op.args[0]
|
||||
_src1 = op.args[1]
|
||||
_src2 = op.args[2]
|
||||
|
||||
s1_scalar = not isinstance(_src1, BufferRegion)
|
||||
s2_scalar = not isinstance(_src2, BufferRegion)
|
||||
if s1_scalar and s2_scalar:
|
||||
return None, "both inputs are constants"
|
||||
|
||||
# Move constant to rhs (commute if allowed; else reject).
|
||||
if s1_scalar:
|
||||
if op_name not in _COMMUTATIVE:
|
||||
return None, f"non-commutative op {op_name} cannot have constant lhs"
|
||||
_src1, _src2 = _src2, _src1
|
||||
s2_scalar = True
|
||||
|
||||
# If rhs is a smaller buffer (broadcast), swap if commutative so the
|
||||
# bigger one is in src1 — keeps src1 == dst convention.
|
||||
if not s2_scalar:
|
||||
s1_n = functools.reduce(operator.mul, [r.extent for r in _src1.region], 1)
|
||||
s2_n = functools.reduce(operator.mul, [r.extent for r in _src2.region], 1)
|
||||
if s1_n < s2_n:
|
||||
if op_name not in _COMMUTATIVE:
|
||||
return None, f"non-commutative op {op_name} cannot swap to broadcast"
|
||||
_src1, _src2 = _src2, _src1
|
||||
|
||||
srcs: list[SrcSpec] = [SrcSpec(buf_region=_src1)]
|
||||
if s2_scalar:
|
||||
srcs.append(SrcSpec(scalar=_src2))
|
||||
else:
|
||||
srcs.append(SrcSpec(buf_region=_src2))
|
||||
|
||||
extras: dict[str, Any] = {}
|
||||
rm = op.config.get("rounding_mode", None)
|
||||
if rm is not None:
|
||||
extras["rounding_mode"] = rm
|
||||
return Plan(dst=_dst, srcs=srcs, extras=extras), None
|
||||
|
||||
return parse
|
||||
|
||||
|
||||
def _compute_add(src_vals, extras, dt):
|
||||
return src_vals[0] + src_vals[1]
|
||||
|
||||
|
||||
def _compute_sub(src_vals, extras, dt):
|
||||
return src_vals[0] - src_vals[1]
|
||||
|
||||
|
||||
def _compute_mul(src_vals, extras, dt):
|
||||
return src_vals[0] * src_vals[1]
|
||||
|
||||
|
||||
def _compute_fdiv(src_vals, extras, dt):
|
||||
return src_vals[0] / src_vals[1]
|
||||
|
||||
|
||||
def _compute_maximum(src_vals, extras, dt):
|
||||
return Tx.max(src_vals[0], src_vals[1])
|
||||
|
||||
|
||||
BINARY_OPS: dict[str, OpSpec] = {
|
||||
"add": OpSpec(
|
||||
"add",
|
||||
_parse_binary_for("add"),
|
||||
_compute_add,
|
||||
vec_impls=[BINARY_F32X2_IMPLS["add"]],
|
||||
),
|
||||
"sub": OpSpec(
|
||||
"sub",
|
||||
_parse_binary_for("sub"),
|
||||
_compute_sub,
|
||||
vec_impls=[BINARY_F32X2_IMPLS["sub"]],
|
||||
),
|
||||
"mul": OpSpec(
|
||||
"mul",
|
||||
_parse_binary_for("mul"),
|
||||
_compute_mul,
|
||||
vec_impls=[BINARY_F32X2_IMPLS["mul"]],
|
||||
),
|
||||
"fdiv": OpSpec(
|
||||
"fdiv",
|
||||
_parse_binary_for("fdiv"),
|
||||
_compute_fdiv,
|
||||
),
|
||||
"maximum": OpSpec(
|
||||
"maximum",
|
||||
_parse_binary_for("maximum"),
|
||||
_compute_maximum,
|
||||
),
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
# 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.
|
||||
|
||||
"""Cast op: ``Tx.cast(dst, src)``. Outer ``Tx.cast(..., dst.dtype)`` in the
|
||||
schedule handles the scalar conversion; the vec-impl packs pairs via
|
||||
CUDA intrinsics like ``__float22half2_rn``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from tvm.tirx import BufferRegion, TilePrimitiveCall
|
||||
|
||||
from ..vec_emit.cast_vec2 import CAST_VEC2_IMPL
|
||||
from . import OpSpec, Plan, SrcSpec
|
||||
|
||||
|
||||
def _parse_cast(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
|
||||
_dst: BufferRegion = op.args[0]
|
||||
_src = op.args[1]
|
||||
if not isinstance(_src, BufferRegion):
|
||||
return None, "cast src must be a buffer region"
|
||||
return Plan(dst=_dst, srcs=[SrcSpec(buf_region=_src)], extras={}), None
|
||||
|
||||
|
||||
def _compute_cast(src_vals, extras, dt):
|
||||
# Schedule wraps with Tx.cast(..., dst.dtype) — just pass through.
|
||||
return src_vals[0]
|
||||
|
||||
|
||||
CAST_OPS: dict[str, OpSpec] = {
|
||||
"cast": OpSpec("cast", _parse_cast, _compute_cast, vec_impls=[CAST_VEC2_IMPL]),
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
# 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.
|
||||
|
||||
"""FMA op: ``Tx.fma(dst, a, b, c)`` → ``dst = a*b + c``.
|
||||
|
||||
Attaches ``fma_f32x2`` VecImpl for sm_100+ f32; falls back to scalar
|
||||
``a*b + c`` otherwise.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from tvm.tirx import BufferRegion, TilePrimitiveCall
|
||||
|
||||
from ..vec_emit.fma_f32x2 import FMA_F32X2_IMPL
|
||||
from . import OpSpec, Plan, SrcSpec
|
||||
|
||||
|
||||
def _parse_fma(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
|
||||
_dst: BufferRegion = op.args[0]
|
||||
args = op.args[1:4]
|
||||
srcs: list[SrcSpec] = []
|
||||
for a in args:
|
||||
if isinstance(a, BufferRegion):
|
||||
srcs.append(SrcSpec(buf_region=a))
|
||||
else:
|
||||
srcs.append(SrcSpec(scalar=a))
|
||||
return Plan(dst=_dst, srcs=srcs, extras={}), None
|
||||
|
||||
|
||||
def _compute_fma(src_vals, extras, dt):
|
||||
return src_vals[0] * src_vals[1] + src_vals[2]
|
||||
|
||||
|
||||
FMA_OPS: dict[str, OpSpec] = {
|
||||
"fma": OpSpec("fma", _parse_fma, _compute_fma, vec_impls=[FMA_F32X2_IMPL]),
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
# 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.
|
||||
|
||||
"""Unary elementwise ops: zero / fill / reciprocal / sqrt / exp / exp2 / silu.
|
||||
|
||||
All carry the same ``T.<unary>(dst, src[, bias, scale])`` shape (bias / scale
|
||||
optional; ``silu`` ignores bias/scale to preserve legacy behavior).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from tvm.ir import is_prim_expr
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import BufferRegion, TilePrimitiveCall
|
||||
from tvm.tirx.expr import FloatImm
|
||||
|
||||
from .._common import scalar_dtype
|
||||
from . import OpSpec, Plan, SrcSpec
|
||||
|
||||
|
||||
def _parse_unary(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
|
||||
"""T.<unary>(dst, src[, bias, scale]) → Plan."""
|
||||
_dst: BufferRegion = op.args[0]
|
||||
_src = op.args[1]
|
||||
_bias = op.args[2] if len(op.args) > 2 else None
|
||||
_scale = op.args[3] if len(op.args) > 2 else None
|
||||
|
||||
srcs: list[SrcSpec] = []
|
||||
if isinstance(_src, BufferRegion):
|
||||
srcs.append(SrcSpec(buf_region=_src))
|
||||
elif is_prim_expr(_src):
|
||||
srcs.append(SrcSpec(scalar=_src))
|
||||
else:
|
||||
return None, f"unsupported src type {type(_src).__name__}"
|
||||
|
||||
extras: dict[str, Any] = {
|
||||
"scale": _scale,
|
||||
"bias_const": _bias if isinstance(_bias, FloatImm) else None,
|
||||
}
|
||||
if isinstance(_bias, BufferRegion):
|
||||
srcs.append(SrcSpec(buf_region=_bias))
|
||||
extras["has_bias_buf"] = True
|
||||
else:
|
||||
extras["has_bias_buf"] = False
|
||||
return Plan(dst=_dst, srcs=srcs, extras=extras), None
|
||||
|
||||
|
||||
def _check_unary_extras(extras: dict, compute_dtype: str) -> tuple[bool, str | None]:
|
||||
scale = extras.get("scale")
|
||||
if scale is not None and scalar_dtype(scale) != compute_dtype:
|
||||
return False, f"scale dtype {scalar_dtype(scale)} != compute dtype {compute_dtype}"
|
||||
bias_const = extras.get("bias_const")
|
||||
if bias_const is not None and scalar_dtype(bias_const) != compute_dtype:
|
||||
return (
|
||||
False,
|
||||
f"bias_const dtype {scalar_dtype(bias_const)} != compute dtype {compute_dtype}",
|
||||
)
|
||||
return True, None
|
||||
|
||||
|
||||
def _with_bias_scale(raw_op):
|
||||
"""Wrap ``raw_op`` (e.g. ``T.exp``) into a compute that applies bias/scale first."""
|
||||
|
||||
def compute(src_vals, extras, dt):
|
||||
x = src_vals[0]
|
||||
scale = extras.get("scale")
|
||||
if scale is not None:
|
||||
x = x * scale
|
||||
if extras.get("has_bias_buf"):
|
||||
x = x + src_vals[1]
|
||||
elif extras.get("bias_const") is not None:
|
||||
x = x + extras["bias_const"]
|
||||
return raw_op(x)
|
||||
|
||||
return compute
|
||||
|
||||
|
||||
def _compute_zero(src_vals, extras, dt):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _compute_fill(src_vals, extras, dt):
|
||||
return src_vals[0]
|
||||
|
||||
|
||||
def _compute_reciprocal(src_vals, extras, dt):
|
||||
x = src_vals[0]
|
||||
return T.FloatImm(x.ty, 1.0) / x
|
||||
|
||||
|
||||
def _compute_silu(src_vals, extras, dt):
|
||||
# Legacy: silu doesn't apply bias/scale.
|
||||
x = src_vals[0]
|
||||
return x / (T.FloatImm(x.ty, 1.0) + T.exp(T.FloatImm(x.ty, 0.0) - x))
|
||||
|
||||
|
||||
UNARY_OPS: dict[str, OpSpec] = {
|
||||
"zero": OpSpec("zero", _parse_unary, _compute_zero, _check_unary_extras),
|
||||
"fill": OpSpec("fill", _parse_unary, _compute_fill, _check_unary_extras),
|
||||
"reciprocal": OpSpec("reciprocal", _parse_unary, _compute_reciprocal, _check_unary_extras),
|
||||
"sqrt": OpSpec("sqrt", _parse_unary, _with_bias_scale(T.sqrt), _check_unary_extras),
|
||||
"exp": OpSpec("exp", _parse_unary, _with_bias_scale(T.exp), _check_unary_extras),
|
||||
"exp2": OpSpec("exp2", _parse_unary, _with_bias_scale(T.exp2), _check_unary_extras),
|
||||
"silu": OpSpec("silu", _parse_unary, _compute_silu, _check_unary_extras),
|
||||
}
|
||||
@@ -0,0 +1,429 @@
|
||||
# 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.
|
||||
|
||||
"""Elementwise dispatch when all operands live in ``local`` (registers).
|
||||
|
||||
Mirrors ``cuda/copy/reg.py``: the partition is *induced* by the layout that
|
||||
carries thread-axis info (the "anchor" operand). The region slice is absorbed
|
||||
into the sliced layout up front via ``align_operands_to_anchor`` — emit
|
||||
operates on a flat 1D per-thread view and indexes it with a scalar offset, so
|
||||
codegen never sees multi-dim ``get_indices`` inside ``T.vectorized``.
|
||||
|
||||
Two paths inside emit:
|
||||
* induced (anchor exists) — atom-based, exactly mirrors copy reg.py
|
||||
* trivial (no anchor) — flat full region, every thread runs the full
|
||||
loop on its private storage
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import operator
|
||||
|
||||
from tvm.arith import Analyzer
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import PrimFunc, TilePrimitiveCall
|
||||
from tvm.tirx.layout import TileLayout
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import fail
|
||||
|
||||
from ..common import get_st_extent
|
||||
from ..copy._common import _carve_tail, _verify_s_tail_contig
|
||||
from ..layout_utils import get_sublayout_from_region, layout_signature
|
||||
from ._common import (
|
||||
_TID_AXIS_FOR_SCOPE,
|
||||
_all_threads_active,
|
||||
_tensor_shape_of,
|
||||
_thread_cnt,
|
||||
align_operands_to_anchor,
|
||||
buffer_regions,
|
||||
compute_dtype_of,
|
||||
pick_anchor,
|
||||
shape_broadcast_compat,
|
||||
)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Predicate
|
||||
# -----------------------------------------------------------------------------
|
||||
def _validate_anchor_layout(anchor_br) -> tuple[bool, str | None]:
|
||||
layout = anchor_br.buffer.layout
|
||||
if layout.is_swizzle():
|
||||
return False, "anchor layout is swizzle"
|
||||
if not isinstance(layout, TileLayout):
|
||||
return False, f"anchor layout is {type(layout).__name__}, not TileLayout"
|
||||
return True, None
|
||||
|
||||
|
||||
def _validate_scope_level_anchor(anchor_br, sctx: DispatchContext) -> tuple[bool, str | None]:
|
||||
"""For warp/warpgroup/cta scope, require dst to be scope-level: after
|
||||
canonicalizing with the target its thread axes are the scope's intra-thread
|
||||
axis (laneid/tid_in_wg/tx) and, sorted by stride, tile a complete ``T:1``
|
||||
chain over all ``T`` threads of the scope. Rejects thread-local ``.local()``
|
||||
views; thread scope is exempt.
|
||||
"""
|
||||
scope = sctx.scope_kind
|
||||
if scope == "thread":
|
||||
return True, None
|
||||
expected_axis = _TID_AXIS_FOR_SCOPE.get(scope)
|
||||
if expected_axis is None:
|
||||
return True, None
|
||||
expected_cnt = _thread_cnt(sctx)
|
||||
|
||||
# Canonicalize the sliced anchor with the target so warp/lane axes fuse.
|
||||
st, ext = get_st_extent(anchor_br)
|
||||
sliced = get_sublayout_from_region(anchor_br.buffer.layout, anchor_br.buffer.shape, st, ext)
|
||||
with sctx.target:
|
||||
canon = sliced.canonicalize()
|
||||
shard = getattr(canon, "shard", None)
|
||||
if shard is None:
|
||||
return False, f"{scope}-scope op operand layout is not a TileLayout after slicing"
|
||||
|
||||
thread_iters = [it for it in shard if it.axis.is_thread()]
|
||||
if not thread_iters:
|
||||
return (
|
||||
False,
|
||||
f"{scope}-scope op needs a {scope}-level operand whose layout carries "
|
||||
f"thread axes ({expected_axis} composing to {expected_cnt}:1); got a "
|
||||
f"thread-local view with no thread axes — pass the {scope}-level tensor, "
|
||||
f"not its `.local()` (per-thread) view",
|
||||
)
|
||||
bad = sorted({it.axis.name for it in thread_iters if it.axis.name != expected_axis})
|
||||
if bad:
|
||||
return (
|
||||
False,
|
||||
f"{scope}-scope op operand carries thread axes {bad}; after "
|
||||
f"canonicalization a {scope}-level layout must use only {expected_axis!r}",
|
||||
)
|
||||
# Sorted by stride the thread iters must tile a complete chain 1, e0,
|
||||
# e0*e1, ... up to the scope thread count — i.e. cover all T threads with
|
||||
# no gap or overlap (extents alone would miss gaps/overlaps).
|
||||
running = 1
|
||||
for it in sorted(thread_iters, key=lambda i: int(i.stride)):
|
||||
stride, extent = int(it.stride), int(it.extent)
|
||||
if stride != running:
|
||||
return (
|
||||
False,
|
||||
f"{scope}-scope op operand thread axes do not tile a complete "
|
||||
f"{expected_cnt}:1 (sorted by stride: expected {running}, got {stride})",
|
||||
)
|
||||
running *= extent
|
||||
if running != expected_cnt:
|
||||
return (
|
||||
False,
|
||||
f"{scope}-scope op operand thread axes span {running} threads, not the "
|
||||
f"full {expected_cnt} of the {scope}",
|
||||
)
|
||||
return True, None
|
||||
|
||||
|
||||
def _check_layout_operands_agree(plan, sctx) -> tuple[bool, str | None]:
|
||||
"""Replica sigs must match across non-trivial-layout operands.
|
||||
|
||||
``align_operands_to_anchor`` normalizes thread + local parts via
|
||||
permute/group, but the replica part isn't touched by alignment — if
|
||||
operands disagree there, alignment can't fix it and emit will be wrong.
|
||||
Thread / local mismatches that alignment can't resolve will raise
|
||||
cleanly at align time, so we don't pre-check them.
|
||||
"""
|
||||
# All operands have a layout (predicate already enforced this); just
|
||||
# iterate them all.
|
||||
layout_brs = list(buffer_regions(plan))
|
||||
if len(layout_brs) < 2:
|
||||
return True, None
|
||||
analyzer = Analyzer()
|
||||
replica_sigs = []
|
||||
for br in layout_brs:
|
||||
st, ext = get_st_extent(br)
|
||||
with sctx.target:
|
||||
sliced = get_sublayout_from_region(br.buffer.layout, br.buffer.shape, st, ext)
|
||||
canon = sliced.canonicalize()
|
||||
sig = layout_signature(canon)
|
||||
if sig is None:
|
||||
return False, "layout has no signature (not a TileLayout?)"
|
||||
# layout_signature returns (thread_sig, local_sig, replica_sig)
|
||||
replica_sigs.append(sig[2])
|
||||
for s in replica_sigs[1:]:
|
||||
# Compare replica entries (axis_key, extent, stride) element-wise.
|
||||
if len(s) != len(replica_sigs[0]):
|
||||
return False, "replica sig mismatch (different number of replica iters)"
|
||||
for (k_a, e_a, st_a), (k_b, e_b, st_b) in zip(replica_sigs[0], s):
|
||||
if k_a != k_b:
|
||||
return False, "replica sig mismatch (axis key)"
|
||||
if not analyzer.can_prove_equal(e_a, e_b):
|
||||
return False, "replica sig mismatch (extent)"
|
||||
if not analyzer.can_prove_equal(st_a, st_b):
|
||||
return False, "replica sig mismatch (stride)"
|
||||
return True, None
|
||||
|
||||
|
||||
def is_reg_ewise(spec):
|
||||
"""Predicate factory: dispatch accepted iff all operands in ``local`` scope."""
|
||||
|
||||
def check(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 scope {sctx.scope_kind}"
|
||||
ok, reason = _all_threads_active(sctx)
|
||||
if not ok:
|
||||
return False, reason
|
||||
plan, msg = spec.parse(op_call)
|
||||
if msg is not None or plan is None:
|
||||
return False, msg
|
||||
for br in buffer_regions(plan):
|
||||
if br.buffer.scope() != "local":
|
||||
return False, f"operand scope {br.buffer.scope()} != local"
|
||||
if br.buffer.layout is None:
|
||||
return False, f"operand {br} has no layout"
|
||||
if spec.check_extras is not None:
|
||||
ok2, reason2 = spec.check_extras(plan.extras, compute_dtype_of(plan))
|
||||
if not ok2:
|
||||
return False, reason2
|
||||
anchor = pick_anchor(plan)
|
||||
ok3, reason3 = _validate_anchor_layout(anchor)
|
||||
if not ok3:
|
||||
return False, reason3
|
||||
ok_scope, reason_scope = _validate_scope_level_anchor(anchor, sctx)
|
||||
if not ok_scope:
|
||||
return False, reason_scope
|
||||
# Shape compat (NumPy-style broadcast): anchor's tensor shape is the
|
||||
# result shape; every operand must broadcast TO anchor.
|
||||
anchor_tshape = _tensor_shape_of(anchor.region)
|
||||
for br in buffer_regions(plan):
|
||||
if br is anchor:
|
||||
continue
|
||||
op_tshape = _tensor_shape_of(br.region)
|
||||
ok_b, reason_b = shape_broadcast_compat(op_tshape, anchor_tshape)
|
||||
if not ok_b:
|
||||
return False, f"shape incompat: {reason_b}"
|
||||
ok4, reason4 = _check_layout_operands_agree(plan, sctx)
|
||||
if not ok4:
|
||||
return False, reason4
|
||||
return True, None
|
||||
|
||||
return check
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Shared helpers
|
||||
# -----------------------------------------------------------------------------
|
||||
def _prod(it) -> int:
|
||||
return functools.reduce(operator.mul, it, 1)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Main entry
|
||||
# -----------------------------------------------------------------------------
|
||||
def emit_reg(op_call: TilePrimitiveCall, spec, sctx: DispatchContext) -> PrimFunc:
|
||||
plan, msg = spec.parse(op_call)
|
||||
if msg is not None or plan is None:
|
||||
fail(msg or "parse failed")
|
||||
return _emit_induced(plan, spec, sctx, op_call, pick_anchor(plan))
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Induced path — anchor with non-trivial layout drives partition
|
||||
# -----------------------------------------------------------------------------
|
||||
def _strip_thread(layout):
|
||||
"""Return a new TileLayout with thread iters removed from shard."""
|
||||
mem_iters = [it for it in layout.shard if not it.axis.is_thread()]
|
||||
return TileLayout.from_iters(mem_iters, list(layout.replica), dict(layout.offset))
|
||||
|
||||
|
||||
def _pick_vec_and_carve(spec, op_call, sctx, plan, per_op_mem_layouts):
|
||||
"""Pick ``(vec_len, vec_impl, carved_layouts)``.
|
||||
|
||||
Enumerate ``spec.vec_impls`` widest-first. For each candidate ``vec_len``:
|
||||
1. Try ``_carve_tail`` on each operand's mem-only layout (may split a
|
||||
boundary iter so the tail product equals ``vec_len``).
|
||||
2. Verify the carved tail is physically contiguous (stride-1 chain whose
|
||||
product equals ``vec_len``).
|
||||
3. Call ``impl.applies(op_call, sctx, plan)``.
|
||||
First candidate that passes all three on EVERY operand wins. Otherwise
|
||||
fall back to scalar: ``vec_len=1``, ``vec_impl=None``, original (uncarved)
|
||||
layouts.
|
||||
"""
|
||||
impls = sorted(getattr(spec, "vec_impls", []), key=lambda i: -i.vec_len)
|
||||
for impl in impls:
|
||||
cand = impl.vec_len
|
||||
carved_try = {}
|
||||
all_ok = True
|
||||
for op_br, layout in per_op_mem_layouts.items():
|
||||
new_iters = _carve_tail(list(layout.shard), cand)
|
||||
if new_iters is None:
|
||||
all_ok = False
|
||||
break
|
||||
new_layout = TileLayout.from_iters(new_iters, list(layout.replica), dict(layout.offset))
|
||||
if not _verify_s_tail_contig(new_layout, cand):
|
||||
all_ok = False
|
||||
break
|
||||
carved_try[op_br] = new_layout
|
||||
if not all_ok:
|
||||
continue
|
||||
ok, _ = impl.applies(op_call, sctx, plan)
|
||||
if not ok:
|
||||
continue
|
||||
return cand, impl, carved_try
|
||||
# Scalar fallback — use uncarved mem layouts as-is.
|
||||
return 1, None, dict(per_op_mem_layouts)
|
||||
|
||||
|
||||
def _emit_induced(plan, spec, sctx, op_call, anchor_br) -> PrimFunc:
|
||||
# Every buffer-region operand has a layout (enforced by predicate);
|
||||
# trivial / identity layouts are fine — the algorithm is robust to
|
||||
# layouts with no thread axes (strip is no-op, placeholders empty).
|
||||
layout_others = [br for br in buffer_regions(plan) if br is not anchor_br]
|
||||
|
||||
# Step 1: slice + permute (region offset absorbed into op_p.offset; per-iter
|
||||
# strides reflect post-slice physical addressing). No (st, ext) leaks out.
|
||||
with sctx.target:
|
||||
anchor_p, per_op_aligned = align_operands_to_anchor(anchor_br, layout_others)
|
||||
|
||||
# Step 2: post-align thread-equality check. ``align`` is supposed to
|
||||
# normalize the thread part; we verify it did. (Replica was pre-checked
|
||||
# in the predicate.)
|
||||
def _thread_iters(layout):
|
||||
c = layout.canonicalize()
|
||||
return [(it.axis, int(it.extent), int(it.stride)) for it in c.shard if it.axis.is_thread()]
|
||||
|
||||
anchor_thread = _thread_iters(anchor_p)
|
||||
for op_br, (op_p, _) in per_op_aligned.items():
|
||||
if _thread_iters(op_p) != anchor_thread:
|
||||
fail("thread part mismatch between anchor and operand after alignment")
|
||||
|
||||
# Step 3: drop thread iters; from here on operands have mem-only layouts.
|
||||
per_op_mem = {anchor_br: _strip_thread(anchor_p)}
|
||||
for op_br, (op_p, _) in per_op_aligned.items():
|
||||
per_op_mem[op_br] = _strip_thread(op_p)
|
||||
|
||||
# Step 4: enumerate spec.vec_impls widest-first; try carve tail for each
|
||||
# operand. First candidate that all operands can carve + impl.applies wins.
|
||||
# Otherwise scalar fallback (vec=1, no inner loop).
|
||||
vec_len, vec_impl, per_op_carved = _pick_vec_and_carve(spec, op_call, sctx, plan, per_op_mem)
|
||||
|
||||
# Step 5: totals + emit. per_thread_total = ∏ extents of (carved) mem
|
||||
# layout. All operands have the same per_thread_total (alignment invariant).
|
||||
per_thread_total = _prod(int(it.extent) for it in per_op_carved[anchor_br].shard)
|
||||
outer_total = per_thread_total // vec_len if vec_len > 0 else per_thread_total
|
||||
|
||||
if vec_impl is not None:
|
||||
result = _emit_induced_packed(
|
||||
plan,
|
||||
vec_impl,
|
||||
vec_len,
|
||||
outer_total,
|
||||
per_thread_total,
|
||||
per_op_carved,
|
||||
anchor_br,
|
||||
)
|
||||
else:
|
||||
result = _emit_induced_scalar(
|
||||
plan,
|
||||
spec,
|
||||
outer_total,
|
||||
per_thread_total,
|
||||
per_op_carved,
|
||||
anchor_br,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _make_views_meta(per_op_carved, per_thread_total):
|
||||
"""Build the per-operand 1D buffer view dict.
|
||||
|
||||
Each view aliases the operand's physical storage as a 1D shape of
|
||||
``per_thread_total`` elements, with layout = the operand's carved mem-only
|
||||
TileLayout. Scalar indexing into the view goes through this layout's
|
||||
iter strides at codegen time.
|
||||
"""
|
||||
return {
|
||||
op_br: T.decl_buffer(
|
||||
(per_thread_total,),
|
||||
op_br.buffer.dtype,
|
||||
op_br.buffer.data,
|
||||
scope="local",
|
||||
layout=per_op_carved[op_br],
|
||||
)
|
||||
for op_br in per_op_carved
|
||||
}
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Emit — packed (one PTX/CUDA call per outer chunk; no T.vectorized inside)
|
||||
# -----------------------------------------------------------------------------
|
||||
def _emit_induced_packed(
|
||||
plan, vec_impl, vec_len, outer_total, per_thread_total, per_op_carved, anchor_br
|
||||
) -> PrimFunc:
|
||||
extras = plan.extras
|
||||
srcs = plan.srcs
|
||||
dst_br = plan.dst
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
views = T.meta_var(_make_views_meta(per_op_carved, per_thread_total))
|
||||
# Serial loop (not T.unroll): T.unroll materializes each per-iter
|
||||
# ``dst_lane_indices`` / ``src_args`` buffer as a fresh int[1]
|
||||
# declaration, multiplying by outer_total. ptxas unrolls the
|
||||
# static-bound loop without that scratch explosion.
|
||||
for f in range(outer_total):
|
||||
# Pass logical 1D coord; each buffer's own layout maps it to
|
||||
# physical at access time (handles wgmma, broadcast, etc.).
|
||||
dst_lane_indices = [[f * vec_len + k] for k in range(vec_len)]
|
||||
src_args = T.meta_var(
|
||||
[
|
||||
src.scalar
|
||||
if src.is_scalar
|
||||
else (
|
||||
views[src.buf_region],
|
||||
[[f * vec_len + k] for k in range(vec_len)],
|
||||
)
|
||||
for src in srcs
|
||||
]
|
||||
)
|
||||
T.evaluate(vec_impl.emit(views[dst_br], dst_lane_indices, src_args, extras))
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Emit — scalar fallback (vec_len = 1; one element per outer iter; no
|
||||
# T.vectorized inside, so no codegen vec-packing of multi-dim indices).
|
||||
# -----------------------------------------------------------------------------
|
||||
def _emit_induced_scalar(
|
||||
plan, spec, outer_total, per_thread_total, per_op_carved, anchor_br
|
||||
) -> PrimFunc:
|
||||
extras = plan.extras
|
||||
srcs = plan.srcs
|
||||
dst_br = plan.dst
|
||||
dst_dtype = dst_br.buffer.dtype
|
||||
compute = spec.compute_scalar
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
views = T.meta_var(_make_views_meta(per_op_carved, per_thread_total))
|
||||
# Serial loop (not T.unroll) — see _emit_induced_packed for why.
|
||||
for f in range(outer_total):
|
||||
# Logical 1D coord = f (vec_len = 1 in scalar path); each
|
||||
# buffer's layout maps to physical at access time.
|
||||
src_vals = T.meta_var(
|
||||
[src.scalar if src.is_scalar else views[src.buf_region][f] for src in srcs]
|
||||
)
|
||||
views[dst_br][f] = T.cast(compute(src_vals, extras, dst_dtype), dst_dtype)
|
||||
|
||||
return impl
|
||||
@@ -0,0 +1,61 @@
|
||||
# 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.
|
||||
|
||||
"""Register each op in ``ALL_OPS`` for both dispatch variants (``reg``, ``smem``).
|
||||
|
||||
Mirrors copy PR-640's two-variant model: scope-pair drives the dispatch
|
||||
selection, the underlying algorithm (induced vs synthesized) follows.
|
||||
"""
|
||||
|
||||
from tvm.tirx import PrimFunc, TilePrimitiveCall
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
|
||||
|
||||
from .ops import ALL_OPS
|
||||
from .reg import emit_reg, is_reg_ewise
|
||||
from .smem import emit_smem, is_smem_ewise
|
||||
|
||||
|
||||
def _register_reg(spec) -> None:
|
||||
@register_dispatch(
|
||||
spec.name,
|
||||
"cuda",
|
||||
variant="reg",
|
||||
priority=10,
|
||||
when=[predicate(f"{spec.name}_reg", is_reg_ewise(spec))],
|
||||
)
|
||||
def _dispatch(op: TilePrimitiveCall, sctx: DispatchContext, _spec=spec) -> PrimFunc:
|
||||
return emit_reg(op, _spec, sctx)
|
||||
|
||||
|
||||
def _register_smem(spec) -> None:
|
||||
@register_dispatch(
|
||||
spec.name,
|
||||
"cuda",
|
||||
variant="smem",
|
||||
priority=10,
|
||||
when=[predicate(f"{spec.name}_smem", is_smem_ewise(spec))],
|
||||
)
|
||||
def _dispatch(op: TilePrimitiveCall, sctx: DispatchContext, _spec=spec) -> PrimFunc:
|
||||
return emit_smem(op, _spec, sctx)
|
||||
|
||||
|
||||
for _spec in ALL_OPS.values():
|
||||
_register_reg(_spec)
|
||||
_register_smem(_spec)
|
||||
|
||||
|
||||
__all__: list[str] = []
|
||||
@@ -0,0 +1,264 @@
|
||||
# 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.
|
||||
|
||||
"""Elementwise dispatch when all operands live in ``shared*``.
|
||||
|
||||
Mirrors ``cuda/copy/gmem_smem.py``: no operand carries a per-thread partition,
|
||||
so partition is *synthesized* from ``sctx.intra`` (``thread_cnt = ∏ intra``).
|
||||
Each thread takes ``ceildiv(total, vec_chunk * thread_cnt)`` strided chunks.
|
||||
|
||||
The shared buffers are indexed multi-dim via ``get_indices(fused, dst_st,
|
||||
dst_ext)`` and the buffer's own layout resolves to physical addresses at
|
||||
codegen time. Packed-vec emit requires the innermost dim to have stride 1
|
||||
(non-swizzle slice) so lanes are physically contiguous; checked in
|
||||
``_max_layout_vec``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import PrimFunc, TilePrimitiveCall
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import fail
|
||||
|
||||
from ..common import get_indices, get_st_extent, get_thread_cnt
|
||||
from ._common import (
|
||||
_TID_AXIS_FOR_SCOPE,
|
||||
_all_threads_active,
|
||||
_axis_decl,
|
||||
_broadcast_indices,
|
||||
_tensor_shape_of,
|
||||
buffer_regions,
|
||||
compute_dtype_of,
|
||||
dtype_bits,
|
||||
emit_scope_sync,
|
||||
fetch_src_value,
|
||||
n_elements,
|
||||
pick_vec_chunk,
|
||||
shape_broadcast_compat,
|
||||
)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Predicate
|
||||
# -----------------------------------------------------------------------------
|
||||
def is_smem_ewise(spec):
|
||||
"""Predicate factory: dispatch accepted iff all operands in ``shared*``."""
|
||||
|
||||
def check(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 scope {sctx.scope_kind}"
|
||||
ok, reason = _all_threads_active(sctx)
|
||||
if not ok:
|
||||
return False, reason
|
||||
plan, msg = spec.parse(op_call)
|
||||
if msg is not None or plan is None:
|
||||
return False, msg
|
||||
for br in buffer_regions(plan):
|
||||
if not br.buffer.scope().startswith("shared"):
|
||||
return False, f"operand scope {br.buffer.scope()} != shared*"
|
||||
if br.buffer.layout is None:
|
||||
return False, "shared operand has no layout"
|
||||
if spec.check_extras is not None:
|
||||
ok2, reason2 = spec.check_extras(plan.extras, compute_dtype_of(plan))
|
||||
if not ok2:
|
||||
return False, reason2
|
||||
# NumPy-style right-aligned broadcast: anchor = plan.dst; every src
|
||||
# must be shape-compatible with anchor (extent matches or is 1).
|
||||
anchor_tshape = _tensor_shape_of(plan.dst.region)
|
||||
for s in plan.srcs:
|
||||
if s.buf_region is None:
|
||||
continue
|
||||
src_tshape = _tensor_shape_of(s.buf_region.region)
|
||||
ok_b, reason_b = shape_broadcast_compat(src_tshape, anchor_tshape)
|
||||
if not ok_b:
|
||||
return False, f"shape incompat: {reason_b}"
|
||||
return True, None
|
||||
|
||||
return check
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# vec_chunk selection
|
||||
# -----------------------------------------------------------------------------
|
||||
def _max_layout_vec(plan, total: int, thread_cnt: int) -> int:
|
||||
"""Widest vec_chunk dividing all operands' innermost extents AND
|
||||
``total / thread_cnt``, within dtype-bit candidates ``{128,64,32,16,8}``."""
|
||||
max_bits = dtype_bits(plan.dst.buffer.dtype)
|
||||
for s in plan.srcs:
|
||||
if s.buf_region is not None:
|
||||
max_bits = max(max_bits, dtype_bits(s.buf_region.buffer.dtype))
|
||||
per_thread = total // thread_cnt if thread_cnt > 0 else total
|
||||
if total % thread_cnt != 0:
|
||||
return 1
|
||||
|
||||
inners = [int(plan.dst.region[-1].extent)]
|
||||
for s in plan.srcs:
|
||||
if s.buf_region is None or s.index_fn is not None:
|
||||
continue
|
||||
inners.append(int(s.buf_region.region[-1].extent))
|
||||
|
||||
for cand_bits in (128, 64, 32, 16, 8):
|
||||
n = cand_bits // max_bits
|
||||
if n <= 0:
|
||||
continue
|
||||
if per_thread % n != 0:
|
||||
continue
|
||||
if all(i % n == 0 for i in inners):
|
||||
return n
|
||||
return 1
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Main entry
|
||||
# -----------------------------------------------------------------------------
|
||||
def emit_smem(op_call: TilePrimitiveCall, spec, sctx: DispatchContext) -> PrimFunc:
|
||||
plan, msg = spec.parse(op_call)
|
||||
if msg is not None or plan is None:
|
||||
fail(msg or "parse failed")
|
||||
|
||||
# Use cuda/common.py:get_thread_cnt rather than copy/_common.py:_thread_cnt
|
||||
# — the latter computes ``∏ sctx.intra`` which silently returns 0 for
|
||||
# sub-warp counts at cta scope (warpid extent rounds down to 0). The
|
||||
# former reads launch_params["threadIdx.x"].dom.extent and is correct
|
||||
# for all scopes.
|
||||
thread_cnt = get_thread_cnt(sctx)
|
||||
if thread_cnt is None:
|
||||
fail(f"unsupported scope {sctx.scope_kind} for smem emit")
|
||||
thread_cnt = int(thread_cnt)
|
||||
if thread_cnt <= 0:
|
||||
fail(f"non-positive thread_cnt {thread_cnt}")
|
||||
assert "threadIdx.y" not in sctx.launch_params and "threadIdx.z" not in sctx.launch_params, (
|
||||
"smem emit currently assumes 1D threadIdx"
|
||||
)
|
||||
|
||||
total = n_elements(plan.dst)
|
||||
vec_max = _max_layout_vec(plan, total, thread_cnt)
|
||||
vec_chunk, vec_impl = pick_vec_chunk(spec, op_call, sctx, plan, vec_max)
|
||||
|
||||
if vec_impl is not None:
|
||||
return _emit_packed(plan, vec_impl, vec_chunk, total, thread_cnt, sctx)
|
||||
return _emit_scalar(plan, spec, vec_chunk, total, thread_cnt, sctx)
|
||||
|
||||
|
||||
def _tid_expr(sctx: DispatchContext):
|
||||
"""Per-scope tid expr. ``thread`` scope returns 0; collective scopes use
|
||||
``_axis_decl`` (T.lane_id / T.thread_id_in_wg / threadIdx.x)."""
|
||||
if sctx.scope_kind == "thread":
|
||||
return 0
|
||||
axis_name = _TID_AXIS_FOR_SCOPE[sctx.scope_kind]
|
||||
return _axis_decl(axis_name, sctx)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Per-lane src index helper (handles broadcast via right-aligned compat)
|
||||
# -----------------------------------------------------------------------------
|
||||
def _src_lane_indices(src_br, dst_lane_indices, dst_st, dst_ext, vec_chunk, fused0):
|
||||
"""Return the per-lane multi-dim index list for ``src_br``.
|
||||
|
||||
If src's region shape matches dst's, fall through to ``get_indices``
|
||||
(same as the legacy non-broadcast path). Otherwise derive each lane's
|
||||
src index from the corresponding dst lane index via right-aligned
|
||||
broadcast compat.
|
||||
"""
|
||||
src_st, src_ext = get_st_extent(src_br)
|
||||
if tuple(int(e) for e in src_ext) == tuple(int(e) for e in dst_ext):
|
||||
return [get_indices(fused0 + k, src_st, src_ext) for k in range(vec_chunk)]
|
||||
return [
|
||||
_broadcast_indices(dst_lane_indices[k], dst_st, dst_ext, src_st, src_ext)
|
||||
for k in range(vec_chunk)
|
||||
]
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Emit — packed
|
||||
# -----------------------------------------------------------------------------
|
||||
def _emit_packed(plan, vec_impl, vec_chunk, total, thread_cnt, sctx) -> PrimFunc:
|
||||
extras = plan.extras
|
||||
srcs = plan.srcs
|
||||
dst_buf = plan.dst.buffer
|
||||
dst_st, dst_ext = get_st_extent(plan.dst)
|
||||
sync = emit_scope_sync(sctx.scope_kind)
|
||||
n_outer = (total + vec_chunk * thread_cnt - 1) // (vec_chunk * thread_cnt)
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
tid = _tid_expr(sctx)
|
||||
for s in T.serial(0, n_outer):
|
||||
# First lane's fused index for this thread, this chunk.
|
||||
fused0 = T.meta_var(s * vec_chunk * thread_cnt + tid * vec_chunk)
|
||||
# Predicate the call (skip the trailing partial chunk).
|
||||
if fused0 + vec_chunk <= total:
|
||||
dst_lane_indices = T.meta_var(
|
||||
[get_indices(fused0 + k, dst_st, dst_ext) for k in range(vec_chunk)]
|
||||
)
|
||||
src_args = T.meta_var(
|
||||
[
|
||||
srcs[i].scalar
|
||||
if srcs[i].is_scalar
|
||||
else (
|
||||
srcs[i].buf_region.buffer,
|
||||
_src_lane_indices(
|
||||
srcs[i].buf_region,
|
||||
dst_lane_indices,
|
||||
dst_st,
|
||||
dst_ext,
|
||||
vec_chunk,
|
||||
fused0,
|
||||
),
|
||||
)
|
||||
for i in range(len(srcs))
|
||||
]
|
||||
)
|
||||
T.evaluate(vec_impl.emit(dst_buf, dst_lane_indices, src_args, extras))
|
||||
sync()
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Emit — scalar fallback
|
||||
# -----------------------------------------------------------------------------
|
||||
def _emit_scalar(plan, spec, vec_chunk, total, thread_cnt, sctx) -> PrimFunc:
|
||||
extras = plan.extras
|
||||
srcs = plan.srcs
|
||||
dst_buf = plan.dst.buffer
|
||||
dst_st, dst_ext = get_st_extent(plan.dst)
|
||||
dst_dtype = dst_buf.dtype
|
||||
compute = spec.compute_scalar
|
||||
sync = emit_scope_sync(sctx.scope_kind)
|
||||
n_outer = (total + vec_chunk * thread_cnt - 1) // (vec_chunk * thread_cnt)
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
tid = _tid_expr(sctx)
|
||||
for s in T.serial(0, n_outer):
|
||||
for vec in T.vectorized(vec_chunk):
|
||||
fused = T.meta_var(s * vec_chunk * thread_cnt + tid * vec_chunk + vec)
|
||||
if fused < total:
|
||||
dst_idx = T.meta_var(get_indices(fused, dst_st, dst_ext))
|
||||
src_vals = T.meta_var(
|
||||
[fetch_src_value(src, fused, dst_idx, dst_st, dst_ext) for src in srcs]
|
||||
)
|
||||
dst_buf[tuple(dst_idx)] = T.cast(
|
||||
compute(src_vals, extras, dst_dtype), dst_dtype
|
||||
)
|
||||
sync()
|
||||
|
||||
return impl
|
||||
@@ -0,0 +1,40 @@
|
||||
# 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.
|
||||
|
||||
"""Packed-vector emit functions for elementwise ops.
|
||||
|
||||
Each module here exposes one or more ``VecImpl`` instances that an ``OpSpec``
|
||||
(in ``ops/``) attaches to its ``vec_impls`` list. ``reg.py``/``smem.py`` then
|
||||
pick the widest matching one at dispatch time, mirroring how copy picks
|
||||
``copy_{Nb}`` from a menu.
|
||||
|
||||
VecImpl emit contract:
|
||||
emit(dst_buf, dst_lane_indices, src_args, extras) -> Expr
|
||||
|
||||
* dst_buf: Buffer
|
||||
* dst_lane_indices: list[list[Expr]] of length ``vec_len``; each entry is the
|
||||
multi-dim indices for one lane (precomputed by schedule).
|
||||
* src_args[i]: one of
|
||||
- Expr (scalar src — broadcast across all lanes)
|
||||
- tuple (Buffer, list[list[Expr]] of length ``vec_len``) — buffer src
|
||||
with per-lane indices
|
||||
* extras: dict (rounding_mode, etc.)
|
||||
|
||||
Returns the PTX/CUDA call result; the schedule wraps in ``T.evaluate`` at
|
||||
the call site. All Python-side shape branching (scalar vs buffer src) happens
|
||||
in this emit function -- collapses the old 4x2 schema.py factory explosion.
|
||||
"""
|
||||
@@ -0,0 +1,97 @@
|
||||
# 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.
|
||||
|
||||
"""Packed f32x2 VecImpls for binary add/sub/mul on sm_100+.
|
||||
|
||||
PTX op family: ``{add,sub,mul}.<rm>.ftz.f32x2``. Each call processes 2 f32s
|
||||
per operand. The old ``_make_binary_packed_f32x2_factory`` (240+ lines, 8
|
||||
``@T.prim_func`` shape combos per op) collapses to one ``emit`` per op
|
||||
because operand-shape branching is now Python-level (outside any
|
||||
``@T.prim_func``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from tvm.ir.expr import Expr
|
||||
from tvm.script import tirx as T
|
||||
|
||||
from .._common import dtype_name, scalar_dtype
|
||||
from ..ops import VecImpl
|
||||
|
||||
|
||||
def _lane(arg, k):
|
||||
"""Read lane ``k`` of one operand argument.
|
||||
|
||||
arg is either a scalar Expr (broadcast) or ``(Buffer, lane_indices)``.
|
||||
"""
|
||||
if isinstance(arg, tuple):
|
||||
buf, lane_indices = arg
|
||||
return buf[tuple(lane_indices[k])]
|
||||
return arg
|
||||
|
||||
|
||||
def _f32x2_applies(op_name):
|
||||
"""Predicate: f32 dst+srcs, sm_100+, no broadcasting srcs, two srcs."""
|
||||
|
||||
def applies(op_call, sctx, plan):
|
||||
from ...common import sm_version_ok
|
||||
|
||||
if dtype_name(plan.dst.buffer.dtype) != "float32":
|
||||
return False, "dst dtype not f32"
|
||||
if not sm_version_ok(op_call, sctx, min_version=100)[0]:
|
||||
return False, "sm version < 100"
|
||||
if len(plan.srcs) != 2:
|
||||
return False, "binary requires 2 srcs"
|
||||
for s in plan.srcs:
|
||||
if s.is_scalar:
|
||||
if scalar_dtype(s.scalar) != "float32":
|
||||
return False, "scalar src dtype not f32"
|
||||
else:
|
||||
if dtype_name(s.buf_region.buffer.dtype) != "float32":
|
||||
return False, "buffer src dtype not f32"
|
||||
if s.index_fn is not None:
|
||||
return False, "broadcasting src not supported by f32x2 packed"
|
||||
return True, None
|
||||
|
||||
return applies
|
||||
|
||||
|
||||
def _emit_binary_f32x2_for(op_name):
|
||||
op_func = getattr(T.ptx, f"{op_name}_f32x2")
|
||||
|
||||
def emit(dst_buf, dst_lane_indices, src_args, extras) -> Expr:
|
||||
a_arg, b_arg = src_args
|
||||
rm = extras.get("rounding_mode", "rz")
|
||||
return op_func(
|
||||
T.address_of(dst_buf[tuple(dst_lane_indices[0])]),
|
||||
T.cuda.make_float2(_lane(a_arg, 0), _lane(a_arg, 1)),
|
||||
T.cuda.make_float2(_lane(b_arg, 0), _lane(b_arg, 1)),
|
||||
rounding=rm,
|
||||
ftz=True,
|
||||
)
|
||||
|
||||
return emit
|
||||
|
||||
|
||||
BINARY_F32X2_IMPLS: dict[str, VecImpl] = {
|
||||
name: VecImpl(
|
||||
vec_len=2,
|
||||
applies=_f32x2_applies(name),
|
||||
emit=_emit_binary_f32x2_for(name),
|
||||
)
|
||||
for name in ("add", "sub", "mul")
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
# 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.
|
||||
|
||||
"""Packed vec_len=2 cast via CUDA pair intrinsics.
|
||||
|
||||
Each supported (src_dtype, dst_dtype) pair has a CUDA builtin that converts a
|
||||
packed-2 source to a packed-2 destination in one instruction
|
||||
(e.g. ``__float22half2_rn``). The intrinsic takes pointers to the first
|
||||
element of each packed pair on either side.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from tvm.ir.expr import Expr
|
||||
from tvm.script import tirx as T
|
||||
|
||||
from .._common import dtype_name
|
||||
from ..ops import VecImpl
|
||||
|
||||
_VEC2_CAST_INTRINSICS = {
|
||||
("float32", "float16"): "__float22half2_rn",
|
||||
("float16", "float32"): "__half22float2",
|
||||
("bfloat16", "float32"): "__bfloat1622float2",
|
||||
("float32", "bfloat16"): "__float22bfloat162_rn",
|
||||
}
|
||||
_DTYPE_X2_NAME = {"float32": "float2", "float16": "half2", "bfloat16": "nv_bfloat162"}
|
||||
|
||||
|
||||
def _intrinsic_name(src_dtype, dst_dtype):
|
||||
return f"tvm_builtin_cast_{src_dtype}x2_{dst_dtype}x2"
|
||||
|
||||
|
||||
def _intrinsic_source(src_dtype, dst_dtype):
|
||||
intrinsic = _VEC2_CAST_INTRINSICS[(src_dtype, dst_dtype)]
|
||||
return (
|
||||
f"\n__forceinline__ __device__ void {_intrinsic_name(src_dtype, dst_dtype)}"
|
||||
f"(void* dst, void* src) {{\n"
|
||||
f" (({_DTYPE_X2_NAME[dst_dtype]}*)dst)[0] = "
|
||||
f"{intrinsic}((({_DTYPE_X2_NAME[src_dtype]}*)src)[0]);\n"
|
||||
"}\n"
|
||||
)
|
||||
|
||||
|
||||
def _cast_vec2_applies(op_call, sctx, plan):
|
||||
if len(plan.srcs) != 1 or plan.srcs[0].is_scalar:
|
||||
return False, "cast requires 1 buffer src"
|
||||
src = plan.srcs[0]
|
||||
if src.index_fn is not None:
|
||||
return False, "broadcasting src not supported by cast vec2"
|
||||
src_dtype = dtype_name(src.buf_region.buffer.dtype)
|
||||
dst_dtype = dtype_name(plan.dst.buffer.dtype)
|
||||
if (src_dtype, dst_dtype) not in _VEC2_CAST_INTRINSICS:
|
||||
return False, f"no vec2 intrinsic for {src_dtype}->{dst_dtype}"
|
||||
return True, None
|
||||
|
||||
|
||||
def _emit_cast_vec2(dst_buf, dst_lane_indices, src_args, extras) -> Expr:
|
||||
src_arg = src_args[0]
|
||||
# cast_vec2 requires buffer src (guarded by applies()).
|
||||
assert isinstance(src_arg, tuple), "cast vec2 src must be a buffer"
|
||||
src_buf, src_lane_indices = src_arg
|
||||
src_dtype = dtype_name(src_buf.dtype)
|
||||
dst_dtype = dtype_name(dst_buf.dtype)
|
||||
func_name = _intrinsic_name(src_dtype, dst_dtype)
|
||||
source_code = _intrinsic_source(src_dtype, dst_dtype)
|
||||
return T.cuda.func_call(
|
||||
func_name,
|
||||
T.address_of(dst_buf[tuple(dst_lane_indices[0])]),
|
||||
T.address_of(src_buf[tuple(src_lane_indices[0])]),
|
||||
source_code=source_code,
|
||||
)
|
||||
|
||||
|
||||
CAST_VEC2_IMPL = VecImpl(
|
||||
vec_len=2,
|
||||
applies=_cast_vec2_applies,
|
||||
emit=_emit_cast_vec2,
|
||||
)
|
||||
@@ -0,0 +1,79 @@
|
||||
# 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.
|
||||
|
||||
"""Packed f32x2 VecImpl for FMA on sm_100+.
|
||||
|
||||
PTX: ``fma.<rm>.ftz.f32x2 d, a, b, c`` — 2 f32 FMAs per call. Same Python-side
|
||||
shape collapse as binary_f32x2.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from tvm.ir.expr import Expr
|
||||
from tvm.script import tirx as T
|
||||
|
||||
from .._common import dtype_name, scalar_dtype
|
||||
from ..ops import VecImpl
|
||||
from .binary_f32x2 import _lane
|
||||
|
||||
|
||||
def _fma_f32x2_applies(op_call, sctx, plan):
|
||||
from ...common import sm_version_ok
|
||||
|
||||
if dtype_name(plan.dst.buffer.dtype) != "float32":
|
||||
return False, "dst dtype not f32"
|
||||
if not sm_version_ok(op_call, sctx, min_version=100)[0]:
|
||||
return False, "sm version < 100"
|
||||
if len(plan.srcs) != 3:
|
||||
return False, "fma requires 3 srcs"
|
||||
a, b, c = plan.srcs
|
||||
if a.is_scalar:
|
||||
return False, "fma 'a' must be a buffer (no scalar-a packed FMA)"
|
||||
if dtype_name(a.buf_region.buffer.dtype) != "float32":
|
||||
return False, "src a dtype not f32"
|
||||
if a.index_fn is not None:
|
||||
return False, "broadcasting src a not supported"
|
||||
for s in (b, c):
|
||||
if s.is_scalar:
|
||||
if scalar_dtype(s.scalar) != "float32":
|
||||
return False, "scalar b/c dtype not f32"
|
||||
else:
|
||||
if dtype_name(s.buf_region.buffer.dtype) != "float32":
|
||||
return False, "buffer b/c dtype not f32"
|
||||
if s.index_fn is not None:
|
||||
return False, "broadcasting src b/c not supported"
|
||||
return True, None
|
||||
|
||||
|
||||
def _emit_fma_f32x2(dst_buf, dst_lane_indices, src_args, extras) -> Expr:
|
||||
a_arg, b_arg, c_arg = src_args
|
||||
rm = extras.get("rounding_mode", "rz")
|
||||
return T.ptx.fma_f32x2(
|
||||
T.address_of(dst_buf[tuple(dst_lane_indices[0])]),
|
||||
T.cuda.make_float2(_lane(a_arg, 0), _lane(a_arg, 1)),
|
||||
T.cuda.make_float2(_lane(b_arg, 0), _lane(b_arg, 1)),
|
||||
T.cuda.make_float2(_lane(c_arg, 0), _lane(c_arg, 1)),
|
||||
rounding=rm,
|
||||
ftz=True,
|
||||
)
|
||||
|
||||
|
||||
FMA_F32X2_IMPL = VecImpl(
|
||||
vec_len=2,
|
||||
applies=_fma_f32x2_applies,
|
||||
emit=_emit_fma_f32x2,
|
||||
)
|
||||
@@ -0,0 +1,104 @@
|
||||
# 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.
|
||||
"""Execution scope utilities for CUDA op dispatches."""
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import PrimFunc
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
|
||||
def macro_or_prim_func(macro: Callable, need_macro: bool = False) -> Callable:
|
||||
"""Wrap a macro in a ``prim_func`` unless the caller explicitly wants the macro."""
|
||||
if need_macro:
|
||||
return macro
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def func():
|
||||
macro()
|
||||
|
||||
return func
|
||||
|
||||
|
||||
def thread_selector(sctx: DispatchContext, inner_impl, macro: bool = False) -> Callable:
|
||||
"""Narrow execution to a single, deterministic thread within ``sctx.exec_scope``.
|
||||
|
||||
The elected thread is stable across invocations so that synchronization
|
||||
primitives (for example PTX ``elect_sync``) behave correctly.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
sctx : DispatchContext
|
||||
The dispatch context. Only ``sctx.scope_kind`` is consulted; the
|
||||
caller is responsible for having narrowed into the desired scope via an
|
||||
``if`` guard with a canonical thread-filter predicate before reaching here.
|
||||
inner_impl : T.inline
|
||||
The body to execute inside the selected thread.
|
||||
macro : bool
|
||||
If True, return the macro directly; otherwise wrap it in a ``prim_func``.
|
||||
"""
|
||||
assert not isinstance(inner_impl, PrimFunc), "inner_impl must be a macro, not a PrimFunc"
|
||||
name = sctx.scope_kind
|
||||
if name == "thread":
|
||||
return macro_or_prim_func(inner_impl, need_macro=macro)
|
||||
if name == "cta":
|
||||
|
||||
@T.inline()
|
||||
def impl():
|
||||
T.lane_id([32])
|
||||
if T.ptx.elect_sync():
|
||||
inner_impl()
|
||||
|
||||
return macro_or_prim_func(impl, need_macro=macro)
|
||||
if name == "warp":
|
||||
|
||||
@T.inline()
|
||||
def impl():
|
||||
T.lane_id([32])
|
||||
if T.ptx.elect_sync():
|
||||
inner_impl()
|
||||
|
||||
return macro_or_prim_func(impl, need_macro=macro)
|
||||
if name == "warpgroup":
|
||||
|
||||
@T.inline()
|
||||
def impl():
|
||||
warp_id = T.warp_id_in_wg([4])
|
||||
T.lane_id([32])
|
||||
if warp_id == 0:
|
||||
if T.ptx.elect_sync():
|
||||
inner_impl()
|
||||
|
||||
return macro_or_prim_func(impl, need_macro=macro)
|
||||
raise ValueError(f"thread_selector: unsupported exec_scope {name!r}")
|
||||
|
||||
|
||||
def single_thread(op_call: TilePrimitiveCall, sctx: DispatchContext) -> bool:
|
||||
"""Predicate for dispatchers that require a single-thread execution scope."""
|
||||
del op_call
|
||||
return sctx.is_thread
|
||||
|
||||
|
||||
def exec_scope_ok(
|
||||
op_call: TilePrimitiveCall, sctx: DispatchContext, expected_scopes: list[str]
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Predicate helper: check that ``sctx.scope_kind`` is in *expected_scopes*."""
|
||||
del op_call
|
||||
ok = sctx.scope_kind in expected_scopes
|
||||
return ok, None if ok else f"unsupported exec_scope {sctx.scope_kind}"
|
||||
@@ -0,0 +1,25 @@
|
||||
# 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.
|
||||
|
||||
"""CUDA synchronous ``gemm`` lowerings (warp-level ``mma.sync`` tensor core).
|
||||
|
||||
Importing this package registers every synchronous CUDA ``gemm`` dispatch
|
||||
candidate as a side effect (each submodule calls ``register_dispatch`` at
|
||||
import time). It is the synchronous counterpart to ``gemm_async``.
|
||||
"""
|
||||
|
||||
from .mma_m16n8k_ import *
|
||||
@@ -0,0 +1,595 @@
|
||||
# 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.
|
||||
|
||||
"""Warp-level ``mma.sync`` GEMM lowering for the synchronous ``gemm`` op on CUDA."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from tvm.arith.analyzer import Analyzer
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import PrimFunc
|
||||
from tvm.tirx.layout import TileLayout
|
||||
from tvm.tirx.operator.tile_primitive import (
|
||||
DispatchContext,
|
||||
fail,
|
||||
predicate,
|
||||
register_dispatch,
|
||||
)
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaInst:
|
||||
"""One concrete mma.sync instruction we can emit. A given shape appears once
|
||||
per dtype signature. Adding an instruction is just adding an entry to
|
||||
MMA_INSTRUCTIONS; the feasibility checks below stay generic. The lane->coord
|
||||
mapping is the fixed m16n8 family structure (hardcoded in the match); the
|
||||
only per-instruction fragment detail is ``k_pack`` (it is not always
|
||||
32/dtype_bits, so it is stored explicitly rather than derived)."""
|
||||
|
||||
name: str # distinguishing tag, e.g. "m16n8k16.bf16"
|
||||
m: int
|
||||
n: int
|
||||
k: int
|
||||
dtype: tuple[str, str, str, str] # the single (A, B, C, D) dtype signature
|
||||
k_pack: int # contiguous-along-K elements packed into one register for A and B
|
||||
|
||||
|
||||
# One entry per concrete instruction. The m16n8 family fixes M/N at 16/8; K is
|
||||
# 16 or 8. f16/bf16 inputs with f32 accumulation, packing 2 bf16/f16 per b32
|
||||
# along K (f16 accumulation, int8, fp8, etc. are added as more entries).
|
||||
_BF16_F32 = ("bfloat16", "bfloat16", "float32", "float32")
|
||||
_F16_F32 = ("float16", "float16", "float32", "float32")
|
||||
MMA_INSTRUCTIONS = (
|
||||
MmaInst("m16n8k16.bf16", 16, 8, 16, _BF16_F32, k_pack=2),
|
||||
MmaInst("m16n8k16.f16", 16, 8, 16, _F16_F32, k_pack=2),
|
||||
MmaInst("m16n8k8.bf16", 16, 8, 8, _BF16_F32, k_pack=2),
|
||||
MmaInst("m16n8k8.f16", 16, 8, 8, _F16_F32, k_pack=2),
|
||||
)
|
||||
|
||||
|
||||
def _split(grouped, seps):
|
||||
"""Split a grouped layout into one shard-only sub-layout per group, plus the
|
||||
layout's offset (returned separately rather than distributed into the subs)."""
|
||||
subs = [
|
||||
TileLayout.from_iters(list(grouped.shard[seps[g] : seps[g + 1]]), [], {})
|
||||
for g in range(len(seps) - 1)
|
||||
]
|
||||
return subs, dict(grouped.offset)
|
||||
|
||||
|
||||
def _combine(layouts, offset):
|
||||
"""Concatenate sub-layouts' shards and apply the separately tracked offset."""
|
||||
shard = [it for lay in layouts for it in lay.shard]
|
||||
return TileLayout.from_iters(shard, [], offset)
|
||||
|
||||
|
||||
def _canon_perm(iters):
|
||||
"""Permutation putting thread iters left, memory iters right; stride-desc within each."""
|
||||
thr = sorted(
|
||||
(i for i, it in enumerate(iters) if it.axis.is_thread()),
|
||||
key=lambda i: -int(iters[i].stride),
|
||||
)
|
||||
mem = sorted(
|
||||
(i for i, it in enumerate(iters) if not it.axis.is_thread()),
|
||||
key=lambda i: -int(iters[i].stride),
|
||||
)
|
||||
return thr + mem
|
||||
|
||||
|
||||
def _canon(layout):
|
||||
"""Reorder an anchor sub-layout into canonical (thread-left/memory-right) order."""
|
||||
return layout.permute_dims(_canon_perm(list(layout.shard)))
|
||||
|
||||
|
||||
def _align(anchor, follower, name):
|
||||
"""Regroup follower by anchor's iter extents, then permute its groups to follow
|
||||
the anchor's canonical order (follower keeps its own strides)."""
|
||||
extents = [int(it.extent) for it in anchor.shard]
|
||||
perm = _canon_perm(list(anchor.shard))
|
||||
try:
|
||||
grp, seps = follower.group(extents)
|
||||
except Exception as e: # follower dim layout incompatible with anchor decomposition
|
||||
fail(f"gemm mma: {name} not alignable to anchor extents {extents}: {e}")
|
||||
return grp.permute_by_groups(seps, perm)
|
||||
|
||||
|
||||
def _region_totals(layout):
|
||||
"""(product of thread-axis iter extents, product of memory-axis iter extents).
|
||||
|
||||
Computed once on each anchor to fix that dim's thread/memory region lengths;
|
||||
followers reuse the anchor's split rather than re-deriving it from their own
|
||||
iters (which can misclassify -- e.g. B's N, whose register slot is actually a
|
||||
lane iter, would otherwise report memory length 1)."""
|
||||
thr, mem = 1, 1
|
||||
for it in layout.shard:
|
||||
if it.axis.is_thread():
|
||||
thr *= int(it.extent)
|
||||
else:
|
||||
mem *= int(it.extent)
|
||||
return thr, mem
|
||||
|
||||
|
||||
def _frag_group(layout, lane, mem, thread_total, mem_total):
|
||||
"""Group one logical-dim sub-layout into the fragment shape, then optionally
|
||||
verify it.
|
||||
|
||||
``lane`` and ``mem`` are lists of ``(extent, stride, want_thread)``;
|
||||
``thread_total`` and ``mem_total`` are this dim's thread/memory region lengths
|
||||
(from its anchor via _region_totals). The group shape is
|
||||
``[thread_total // prod(lane), *lane, mem_total // prod(mem), *mem]``: the lane
|
||||
extents are carved off the thread region and the mem extents off the memory
|
||||
region (innermost last, e.g. ``[(reg, ...)]`` for an accumulator dim or
|
||||
``[(kHi, ...), (k_pack, ...)]`` for A/B's K). The input must already be in
|
||||
canonical (thread-left/memory-right) order.
|
||||
|
||||
Every carved group is verified: it must be a single iter, its axis must be a
|
||||
thread axis when ``want_thread`` else a memory axis (only is_thread is checked,
|
||||
since scope varies the exact thread axis; e.g. B's N register slot is actually
|
||||
a lane, so want_thread=True there), and a non-None ``stride`` pins that iter's
|
||||
stride. Raises on a tiling or verification failure, so a non-matching caller
|
||||
layout is declined via the caller's try/except.
|
||||
"""
|
||||
lane_ext = [e for e, _, _ in lane]
|
||||
mem_ext = [e for e, _, _ in mem]
|
||||
lane_prod, mem_prod = 1, 1
|
||||
for e in lane_ext:
|
||||
lane_prod *= e
|
||||
for e in mem_ext:
|
||||
mem_prod *= e
|
||||
grouped, seps = layout.group(
|
||||
[thread_total // lane_prod, *lane_ext, mem_total // mem_prod, *mem_ext]
|
||||
)
|
||||
# group order: [thread_rest, *lane (from idx 1), mem_rest, *mem (after)].
|
||||
specs = [(1 + j, s, t) for j, (_, s, t) in enumerate(lane)]
|
||||
specs += [(2 + len(lane) + j, s, t) for j, (_, s, t) in enumerate(mem)]
|
||||
for idx, stride, want_thread in specs:
|
||||
grp = grouped.shard[seps[idx] : seps[idx + 1]]
|
||||
if len(grp) != 1:
|
||||
raise ValueError(f"frag group {idx} is not a single iter")
|
||||
if int(grp[0].extent) == 1:
|
||||
# An extent-1 group iterates nothing, so its axis/stride is
|
||||
# meaningless and gets dropped downstream (cf. _same_iters /
|
||||
# _reg_layout). This is the kHi == 1 case of m16n8k8: there is a
|
||||
# single high-K register group, which .group() may materialize as a
|
||||
# degenerate split of the (thread) lane axis.
|
||||
continue
|
||||
if grp[0].axis.is_thread() != want_thread:
|
||||
raise ValueError(f"frag group {idx} thread/memory axis mismatch")
|
||||
if stride is not None and int(grp[0].stride) != stride:
|
||||
raise ValueError(f"frag group {idx} stride {int(grp[0].stride)} != {stride}")
|
||||
return grouped, seps
|
||||
|
||||
|
||||
def _grp(grouped, seps, i):
|
||||
"""The iters of group ``i`` of a grouped layout: ``shard[seps[i]:seps[i+1]]``."""
|
||||
return grouped.shard[seps[i] : seps[i + 1]]
|
||||
|
||||
|
||||
def _ext(iters):
|
||||
"""Product of the extents of an iter list."""
|
||||
p = 1
|
||||
for it in iters:
|
||||
p *= int(it.extent)
|
||||
return p
|
||||
|
||||
|
||||
def _reg_layout(groups, offset):
|
||||
"""Per-thread register layout from per-logical-dim iter groups (dropping
|
||||
thread-axis offset terms), plus the matching local-view shape (each dim is
|
||||
the product of that group's extents).
|
||||
|
||||
Extent-1 iters are dropped from the layout: they iterate nothing (offset
|
||||
always 0, so harmless to the mapping) but would otherwise pin a degenerate
|
||||
axis -- e.g. B's N has no real register, so its "register" slot is a single
|
||||
extent-1 lane iter that must not make the register buffer thread-axis."""
|
||||
iters = [it for g in groups for it in g if int(it.extent) != 1]
|
||||
layout = TileLayout.from_iters(
|
||||
iters, [], {ax: v for ax, v in offset.items() if not ax.is_thread()}
|
||||
)
|
||||
return layout, [_ext(g) for g in groups]
|
||||
|
||||
|
||||
def _same_iters(a, b):
|
||||
"""True iff iter lists ``a`` and ``b`` match elementwise on (extent, stride,
|
||||
axis), ignoring extent-1 iters (they iterate nothing, so their stride/axis is
|
||||
meaningless -- e.g. the degenerate thread-rest left by a group shape's '1')."""
|
||||
a = [it for it in a if int(it.extent) != 1]
|
||||
b = [it for it in b if int(it.extent) != 1]
|
||||
if len(a) != len(b):
|
||||
return False
|
||||
return all(
|
||||
int(x.extent) == int(y.extent)
|
||||
and int(x.stride) == int(y.stride)
|
||||
and x.axis.name == y.axis.name
|
||||
for x, y in zip(a, b, strict=True)
|
||||
)
|
||||
|
||||
|
||||
def _full_active_lanes(op: TilePrimitiveCall, sctx: DispatchContext):
|
||||
"""The active thread set (sctx.intra) must be complete and un-narrowed.
|
||||
|
||||
mma.sync.aligned is collective over every active thread; an enclosing if
|
||||
that narrows any intra axis makes the .aligned instruction undefined. So
|
||||
each intra axis must be at offset 0 with its full extent: laneid=32,
|
||||
wid_in_wg=4 (warpgroup), and warpid=warps-per-CTA (cta) from the launch
|
||||
config. Any other axis (e.g. cta_id at cluster scope) is not supported.
|
||||
"""
|
||||
full = {"laneid": 32, "wid_in_wg": 4}
|
||||
if "warpid" in sctx.intra:
|
||||
tx = sctx.launch_params.get("threadIdx.x")
|
||||
if tx is None:
|
||||
return False, "cta scope needs threadIdx.x in launch_params"
|
||||
try:
|
||||
full["warpid"] = int(tx.dom.extent) // 32
|
||||
except (TypeError, ValueError):
|
||||
return False, f"non-static threadIdx.x extent {tx.dom.extent}"
|
||||
for axis, rng in sctx.intra.items():
|
||||
if axis not in full:
|
||||
return False, f"unsupported active-set axis {axis!r}"
|
||||
extent, offset = int(rng[0]), int(rng[1])
|
||||
if extent != full[axis] or offset != 0:
|
||||
return False, (
|
||||
f"active {axis} is [{offset}, {offset + extent}), need full [0, {full[axis]})"
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def _no_replica(op: TilePrimitiveCall, sctx: DispatchContext):
|
||||
"""All operand layouts must have no replica (no broadcast/duplicated axes)."""
|
||||
for region, name in zip(op.args[:4], ("D", "A", "B", "C")):
|
||||
if region.buffer.layout.replica:
|
||||
return False, f"{name} layout has replica {region.buffer.layout.replica}"
|
||||
return True
|
||||
|
||||
|
||||
@register_dispatch(
|
||||
"gemm",
|
||||
"cuda",
|
||||
variant="mma.m16n8k*",
|
||||
priority=10,
|
||||
when=[
|
||||
predicate("full_active_lanes", _full_active_lanes),
|
||||
predicate("no_replica", _no_replica),
|
||||
],
|
||||
)
|
||||
def gemm_cuda_mma_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
"""``gemm`` -> warp-level ``mma.sync`` of the m16n8k* family.
|
||||
|
||||
This is the ``"mma.m16n8k*"`` variant. It targets the m16n8k* tensor-core
|
||||
instructions -- currently m16n8k16 / m16n8k8 with bf16/f16 inputs and f32
|
||||
accumulation (see MMA_INSTRUCTIONS); other K (and other shapes/dtypes) are
|
||||
added as more entries. Pure-register path: A/B fragments and C/D
|
||||
accumulators all live in registers.
|
||||
"""
|
||||
# gemm op args: D = alpha * A @ B + beta * C
|
||||
# D (args[0]) is the output; C (args[3]) is the beta-accumulator input.
|
||||
D_region, A_region, B_region, C_region, transpose_A, transpose_B, alpha, beta = op.args
|
||||
D, A, B, C = D_region.buffer, A_region.buffer, B_region.buffer, C_region.buffer
|
||||
|
||||
# Pure-register mma path: A/B fragments and C/D accumulators all live in
|
||||
# registers ("local"). The caller is responsible for staging A/B into
|
||||
# registers (e.g. via ldmatrix) beforehand.
|
||||
for buf, name in ((D, "D"), (A, "A"), (B, "B"), (C, "C")):
|
||||
if buf.scope() != "local":
|
||||
fail(f"gemm mma requires {name} in register (local) scope, got {buf.scope()}")
|
||||
|
||||
# transpose_A/transpose_B only describe the input's logical orientation; we
|
||||
# normalize to the standard form A=[M,K], B=[K,N] (D/C are always [M,N]).
|
||||
# transpose_A: False -> buffer is [M,K], True -> [K,M]
|
||||
# transpose_B: False -> buffer is [K,N], True -> [N,K]
|
||||
# The .row.col K-major requirement is not enforced here -- it is checked
|
||||
# later by the per-instruction fragment match against the real layout.
|
||||
analyzer = Analyzer()
|
||||
|
||||
# mma.sync computes D = A·B + C natively (no scalar scaling), so we support
|
||||
# only alpha=1 and beta in {0, 1}; beta selects whether C is the accumulator
|
||||
# (1 -> c_ptr=C, 0 -> c_ptr=0). General alpha/beta is declined.
|
||||
def _const_scalar(expr):
|
||||
s = analyzer.simplify(expr)
|
||||
try:
|
||||
return float(s.value)
|
||||
except (AttributeError, TypeError, ValueError):
|
||||
return None
|
||||
|
||||
if _const_scalar(alpha) != 1.0:
|
||||
fail(f"gemm mma supports only alpha=1, got alpha={alpha}")
|
||||
if _const_scalar(beta) not in (0.0, 1.0):
|
||||
fail(f"gemm mma supports only beta in {{0, 1}}, got beta={beta}")
|
||||
|
||||
def _mat_extents(region, name):
|
||||
ext = [r.extent for r in region.region if not analyzer.can_prove_equal(r.extent, 1)]
|
||||
if len(ext) != 2:
|
||||
fail(f"gemm mma expects 2D {name}, got non-unit extents {ext}")
|
||||
return ext
|
||||
|
||||
A_ext = _mat_extents(A_region, "A")
|
||||
B_ext = _mat_extents(B_region, "B")
|
||||
D_M, D_N = _mat_extents(D_region, "D")
|
||||
C_M, C_N = _mat_extents(C_region, "C")
|
||||
M, K = (A_ext[1], A_ext[0]) if transpose_A else (A_ext[0], A_ext[1])
|
||||
B_K, N = (B_ext[1], B_ext[0]) if transpose_B else (B_ext[0], B_ext[1])
|
||||
assert analyzer.can_prove_equal(B_K, K), f"gemm mma: A K={K} != B K={B_K}"
|
||||
assert analyzer.can_prove_equal(D_M, M) and analyzer.can_prove_equal(D_N, N), (
|
||||
f"gemm mma: D dims ({D_M}, {D_N}) != (M={M}, N={N})"
|
||||
)
|
||||
assert analyzer.can_prove_equal(C_M, M) and analyzer.can_prove_equal(C_N, N), (
|
||||
f"gemm mma: C dims ({C_M}, {C_N}) != (M={M}, N={N})"
|
||||
)
|
||||
|
||||
# Tiling into instructions needs static extents.
|
||||
def _const(expr, name):
|
||||
try:
|
||||
return int(analyzer.simplify(expr))
|
||||
except (TypeError, ValueError):
|
||||
fail(f"gemm mma needs static {name} extent, got {expr}")
|
||||
|
||||
M, N, K = _const(M, "M"), _const(N, "N"), _const(K, "K")
|
||||
|
||||
# Slice each operand's layout to its region, then group it into its 2D
|
||||
# buffer-order shape; the split below maps those groups to standard
|
||||
# (M,K)/(K,N). group() raises if the layout can't be tiled that way, so a
|
||||
# caller layout that doesn't match the operand shape is declined cleanly.
|
||||
def _slice_group(buf, region, shape2d, name):
|
||||
# slice() itself groups internally, so both slice and group can raise the
|
||||
# ICHECK when the layout can't be tiled as shape2d -- guard both.
|
||||
canon = None
|
||||
try:
|
||||
sliced = buf.layout.slice(buf.shape, region.region)
|
||||
if sliced is not None:
|
||||
canon = sliced.canonicalize()
|
||||
except Exception as e: # ICHECK failure -> layout not tileable as shape2d
|
||||
fail(f"gemm mma: {name} layout not tileable as {tuple(shape2d)}: {e}")
|
||||
if canon is None:
|
||||
fail(f"gemm mma: cannot slice {name} layout to its region")
|
||||
# All thread iters must share one thread axis (e.g. all laneid). _frag_group
|
||||
# only checks is_thread (not the exact axis, since scope varies), so a layout
|
||||
# mixing two thread axes (e.g. laneid + wid_in_wg) would carve ambiguously --
|
||||
# decline it here, like ldmatrix validating its thread structure.
|
||||
thread_axes = {it.axis.name for it in canon.shard if it.axis.is_thread()}
|
||||
if len(thread_axes) > 1:
|
||||
fail(f"gemm mma: {name} has >1 thread axis {sorted(thread_axes)}, only one supported")
|
||||
try:
|
||||
return canon.group(list(shape2d))
|
||||
except Exception as e: # ICHECK failure -> layout not tileable as shape2d
|
||||
fail(f"gemm mma: {name} layout not tileable as {tuple(shape2d)}: {e}")
|
||||
|
||||
A_grouped, A_seps = _slice_group(A, A_region, (K, M) if transpose_A else (M, K), "A")
|
||||
B_grouped, B_seps = _slice_group(B, B_region, (N, K) if transpose_B else (K, N), "B")
|
||||
C_grouped, C_seps = _slice_group(C, C_region, (M, N), "C")
|
||||
D_grouped, D_seps = _slice_group(D, D_region, (M, N), "D")
|
||||
|
||||
# Split each operand into per-logical-dim sub-layouts (+ its offset), mapping
|
||||
# the buffer-order subs to standard (M,K)/(K,N) per the transpose flags.
|
||||
(DM, DN), D_off = _split(D_grouped, D_seps)
|
||||
(CM, CN), C_off = _split(C_grouped, C_seps)
|
||||
A_subs, A_off = _split(A_grouped, A_seps)
|
||||
B_subs, B_off = _split(B_grouped, B_seps)
|
||||
AM, AK = (A_subs[1], A_subs[0]) if transpose_A else (A_subs[0], A_subs[1])
|
||||
BK, BN = (B_subs[1], B_subs[0]) if transpose_B else (B_subs[0], B_subs[1])
|
||||
|
||||
# Anchor-align so every operand decomposes each shared logical dim the same
|
||||
# way: M anchor = DM -> AM, CM ; N anchor = DN -> BN, CN ; K anchor = AK -> BK.
|
||||
# _align uses each anchor's raw (pre-canon) per-iter extents to group the
|
||||
# follower and reorders the follower's groups into the anchor's canonical
|
||||
# order, so the followers come out canonical. Canon the anchors themselves
|
||||
# afterwards. Every sub-layout is then in canonical (thread-left/memory-right)
|
||||
# order before the loop, so _frag_group groups directly without re-canon.
|
||||
AM = _align(DM, AM, "A.M")
|
||||
CM = _align(DM, CM, "C.M")
|
||||
BN = _align(DN, BN, "B.N")
|
||||
CN = _align(DN, CN, "C.N")
|
||||
BK = _align(AK, BK, "B.K")
|
||||
DM, DN, AK = _canon(DM), _canon(DN), _canon(AK)
|
||||
|
||||
# Each dim's thread/memory region lengths, fixed once from its anchor (3
|
||||
# anchors x 2 parts = 6 lengths). Every operand of that dim reuses them in
|
||||
# _frag_group, so a follower whose register slot is actually a lane (B's N)
|
||||
# still gets the anchor's memory length instead of its own (mis)classified one.
|
||||
m_thr, m_mem = _region_totals(DM)
|
||||
n_thr, n_mem = _region_totals(DN)
|
||||
k_thr, k_mem = _region_totals(AK)
|
||||
|
||||
# Per-instruction selection: try each candidate in order and use the first
|
||||
# whose shape / dtype / (later) fragment layout all fit. A failing check just
|
||||
# moves on to the next instruction; if none fit, decline.
|
||||
sig = (str(A.dtype), str(B.dtype), str(C.dtype), str(D.dtype))
|
||||
for inst in MMA_INSTRUCTIONS:
|
||||
assert inst.m % 8 == 0 and inst.n % 8 == 0 and inst.k % 8 == 0, (
|
||||
f"mma instruction {inst.name} m/n/k must be multiples of 8"
|
||||
)
|
||||
if M % inst.m or N % inst.n or K % inst.k:
|
||||
continue
|
||||
if sig != inst.dtype:
|
||||
continue
|
||||
# Group every operand into this instruction's fragment shape. The m16n8
|
||||
# lane split is g (8 lanes) along M and t (4 lanes) along N/K:
|
||||
# C/D accumulator: M = g + 8*rM (inst.m//8 regs), N = 2*t + rN (inst.n//4 regs)
|
||||
# A multiplicand: M as C/D's M, K = 2*t + p + 8*kHi
|
||||
# B multiplicand: K as A's K, N = g (lane 8, no reg: B has no M so N
|
||||
# reuses the 8-lane g group)
|
||||
# so K's memory tail is [kHi, k_pack] with k_pack the innermost (stride-1)
|
||||
# contiguous pack and kHi = inst.k // (4 * k_pack) high-K register groups.
|
||||
# _frag_group raises if the caller layout can't be tiled or (when any
|
||||
# stride is given) fails the fragment checks -> move on to the next
|
||||
# instruction. lane/mem are [(extent, stride), ...]: the lane stride pins
|
||||
# the laneid stride (g=4, t=1), a mem stride pins a register iter (rN /
|
||||
# k_pack = 1; rM / kHi free = None). C shares D's accumulator fragment and
|
||||
# B's K shares A's K. B's N is the pure 8-lane g group, but aligned to the
|
||||
# accumulator's N (lane 4 + reg 2) it splits into lane g_hi (4, laneid
|
||||
# stride 8) and a "register" g_lo (2, laneid stride 4) -- a lane iter in the
|
||||
# register slot (B's N has no real register). Region lengths come from the
|
||||
# per-dim anchor (m/n/k _thr,_mem), so this split still tiles correctly.
|
||||
# _frag_group(layout, lane, mem, thread_total, mem_total); each carve is
|
||||
# (extent, stride, want_thread): lanes are thread, registers memory, except
|
||||
# B.N's register slot (g_lo) which is itself a lane (want_thread=True).
|
||||
kHi = inst.k // (4 * inst.k_pack)
|
||||
try:
|
||||
DM_g, DM_seps = _frag_group(
|
||||
DM, [(8, 4, True)], [(inst.m // 8, None, False)], m_thr, m_mem
|
||||
)
|
||||
DN_g, DN_seps = _frag_group(DN, [(4, 1, True)], [(inst.n // 4, 1, False)], n_thr, n_mem)
|
||||
CM_g, CM_seps = _frag_group(
|
||||
CM, [(8, 4, True)], [(inst.m // 8, None, False)], m_thr, m_mem
|
||||
)
|
||||
CN_g, CN_seps = _frag_group(CN, [(4, 1, True)], [(inst.n // 4, 1, False)], n_thr, n_mem)
|
||||
AM_g, AM_seps = _frag_group(
|
||||
AM, [(8, 4, True)], [(inst.m // 8, None, False)], m_thr, m_mem
|
||||
)
|
||||
AK_g, AK_seps = _frag_group(
|
||||
AK, [(4, 1, True)], [(kHi, None, False), (inst.k_pack, 1, False)], k_thr, k_mem
|
||||
)
|
||||
BK_g, BK_seps = _frag_group(
|
||||
BK, [(4, 1, True)], [(kHi, None, False), (inst.k_pack, 1, False)], k_thr, k_mem
|
||||
)
|
||||
BN_g, BN_seps = _frag_group(BN, [(4, 8, True)], [(2, 4, True)], n_thr, n_mem)
|
||||
except Exception:
|
||||
continue
|
||||
# M.to (M's warp tiling, group 0) must match across D, A, C so the same
|
||||
# logical M-block lands on the same warp in all three operands.
|
||||
m_to = _grp(DM_g, DM_seps, 0)
|
||||
if not (
|
||||
_same_iters(m_to, _grp(AM_g, AM_seps, 0)) and _same_iters(m_to, _grp(CM_g, CM_seps, 0))
|
||||
):
|
||||
continue
|
||||
# N.to (N's warp tiling, group 0) must match across D, B, C so the same
|
||||
# logical N-block lands on the same warp in all three operands.
|
||||
n_to = _grp(DN_g, DN_seps, 0)
|
||||
if not (
|
||||
_same_iters(n_to, _grp(BN_g, BN_seps, 0)) and _same_iters(n_to, _grp(CN_g, CN_seps, 0))
|
||||
):
|
||||
continue
|
||||
# K.to (K's warp tiling, group 0) must match across A, B so the same
|
||||
# logical K-block lands on the same warp in both operands.
|
||||
if not _same_iters(_grp(AK_g, AK_seps, 0), _grp(BK_g, BK_seps, 0)):
|
||||
continue
|
||||
break
|
||||
else:
|
||||
fail(f"no mma instruction fits M={M}, N={N}, K={K}, dtypes={sig}")
|
||||
|
||||
# Per-operand register layout + matching local-view shape, grouped per logical
|
||||
# dim (offset drops thread-axis terms). Iter order = shape dim order:
|
||||
# D/C -> [M.mo, N.mo, rM, rN] A -> [M.mo, K.mo, rM, kHi, k_pack]
|
||||
# B -> [K.mo, N.mo, kHi, k_pack] (k_pack innermost / contiguous)
|
||||
D_reg, d_shape = _reg_layout(
|
||||
[
|
||||
_grp(DM_g, DM_seps, 2),
|
||||
_grp(DN_g, DN_seps, 2),
|
||||
_grp(DM_g, DM_seps, 3),
|
||||
_grp(DN_g, DN_seps, 3),
|
||||
],
|
||||
D_off,
|
||||
)
|
||||
C_reg, c_shape = _reg_layout(
|
||||
[
|
||||
_grp(CM_g, CM_seps, 2),
|
||||
_grp(CN_g, CN_seps, 2),
|
||||
_grp(CM_g, CM_seps, 3),
|
||||
_grp(CN_g, CN_seps, 3),
|
||||
],
|
||||
C_off,
|
||||
)
|
||||
A_reg, a_shape = _reg_layout(
|
||||
[
|
||||
_grp(AM_g, AM_seps, 2),
|
||||
_grp(AK_g, AK_seps, 2),
|
||||
_grp(AM_g, AM_seps, 3),
|
||||
_grp(AK_g, AK_seps, 3),
|
||||
_grp(AK_g, AK_seps, 4),
|
||||
],
|
||||
A_off,
|
||||
)
|
||||
B_reg, b_shape = _reg_layout(
|
||||
[
|
||||
_grp(BK_g, BK_seps, 2),
|
||||
_grp(BN_g, BN_seps, 2),
|
||||
_grp(BK_g, BK_seps, 3),
|
||||
_grp(BK_g, BK_seps, 4),
|
||||
],
|
||||
B_off,
|
||||
)
|
||||
|
||||
# Emit one mma per (m, n) output tile, accumulating over K. The tile / init /
|
||||
# K loops use T.unroll: the UnrollLoop pass fully expands them in TIR (their
|
||||
# bounds are compile-time constants), so the local-buffer indices resolve to
|
||||
# static register slots -- mma register operands must be constant.
|
||||
#
|
||||
# mma is d = a·b + c. D's accumulator is initialized once per output tile --
|
||||
# copying C when beta==1, clearing to 0 when beta==0 -- then every K step
|
||||
# accumulates in place with c = d, giving a single uniform mma form.
|
||||
M_tiles, N_tiles, K_tiles = d_shape[0], d_shape[1], a_shape[1]
|
||||
shape_str = f"m{inst.m}n{inst.n}k{inst.k}"
|
||||
a_type, b_type, c_type, d_type = inst.dtype
|
||||
use_c = _const_scalar(beta) == 1.0
|
||||
|
||||
# Per-register counts in the fixed PTX enumeration order (derived from the
|
||||
# instruction, NOT hardcoded, so m16n8k8 with kHi==1 also works):
|
||||
# D/C accumulator: rM = inst.m // 8 regs along M, rN = inst.n // 4 along N
|
||||
# c_id = 2 * rM + rN (4 f32 for m16n8k16, also 4 for k8)
|
||||
# A multiplicand: rM = inst.m // 8, kHi = inst.k // (4 * inst.k_pack)
|
||||
# b32 = rM + 2 * kHi (4 b32 for k16, 2 b32 for k8)
|
||||
# B multiplicand: kHi = inst.k // (4 * inst.k_pack)
|
||||
# b32 = kHi (2 b32 for k16, 1 b32 for k8)
|
||||
n_rM = inst.m // 8
|
||||
n_rN = inst.n // 4
|
||||
n_kHi = inst.k // (4 * inst.k_pack)
|
||||
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
d_local = D.local(*d_shape, layout=D_reg)
|
||||
c_local = C.local(*c_shape, layout=C_reg)
|
||||
a_local = A.local(*a_shape, layout=A_reg)
|
||||
b_local = B.local(*b_shape, layout=B_reg)
|
||||
for m in T.unroll(M_tiles):
|
||||
for n in T.unroll(N_tiles):
|
||||
# Initialize D[m, n]: copy C (beta==1) or clear to 0 (beta==0).
|
||||
for rM in T.unroll(n_rM):
|
||||
for rN in T.unroll(n_rN):
|
||||
if use_c:
|
||||
d_local[m, n, rM, rN] = c_local[m, n, rM, rN]
|
||||
else:
|
||||
d_local[m, n, rM, rN] = T.float32(0)
|
||||
# Accumulate over K in place: d = a·b + d.
|
||||
for k in T.unroll(K_tiles):
|
||||
# D: 4 f32 in PTX order c_id = 2*rM + rN.
|
||||
d_ptrs = [
|
||||
d_local.ptr_to([m, n, rM, rN]) for rM in range(n_rM) for rN in range(n_rN)
|
||||
]
|
||||
# A: b32 regs in PTX order b32 = rM + 2*kHi (kHi outer, rM inner).
|
||||
a_ptrs = [
|
||||
a_local.ptr_to([m, k, rM, kHi, 0])
|
||||
for kHi in range(n_kHi)
|
||||
for rM in range(n_rM)
|
||||
]
|
||||
# B: b32 regs in PTX order b32 = kHi.
|
||||
b_ptrs = [b_local.ptr_to([k, n, kHi, 0]) for kHi in range(n_kHi)]
|
||||
# Accumulate in place into D's own regs: c = d.
|
||||
T.ptx.mma(
|
||||
shape_str,
|
||||
"row",
|
||||
"col",
|
||||
d_type,
|
||||
a_type,
|
||||
b_type,
|
||||
c_type,
|
||||
d_ptrs,
|
||||
a_ptrs,
|
||||
b_ptrs,
|
||||
d_ptrs,
|
||||
)
|
||||
|
||||
return impl
|
||||
@@ -0,0 +1,18 @@
|
||||
# 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.
|
||||
|
||||
from .tcgen05 import *
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,62 @@
|
||||
# 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.
|
||||
|
||||
"""GEMM-related utilities for CUDA op dispatches."""
|
||||
|
||||
from tvm.arith.analyzer import Analyzer
|
||||
from tvm.tirx import Buffer
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
|
||||
def validate_gemm_op(op_call: TilePrimitiveCall, sctx: DispatchContext) -> bool:
|
||||
"""Sanity check for gemm op"""
|
||||
C_buffer_region, A_buffer_region, B_buffer_region = op_call.args[:3]
|
||||
C: Buffer = C_buffer_region.buffer
|
||||
A: Buffer = A_buffer_region.buffer
|
||||
B: Buffer = B_buffer_region.buffer
|
||||
if not (C.layout and A.layout and B.layout and A.dtype == B.dtype):
|
||||
return False
|
||||
# Extract regions and validate dimensions
|
||||
analyzer = Analyzer()
|
||||
C_region, A_region, B_region = (
|
||||
C_buffer_region.region,
|
||||
A_buffer_region.region,
|
||||
B_buffer_region.region,
|
||||
)
|
||||
# Extract extents and validate non-unit dimensions match
|
||||
transA, transB = op_call.args[3:5]
|
||||
C_extent_ = [r.extent for r in C_region if r.extent != 1]
|
||||
A_extent_ = [r.extent for r in A_region if r.extent != 1]
|
||||
B_extent_ = [r.extent for r in B_region if r.extent != 1]
|
||||
assert len(C_extent_) == len(A_extent_) == len(B_extent_) == 2, (
|
||||
"Only 2D C, A, B are supported for gemm"
|
||||
)
|
||||
if transA:
|
||||
A_extent_ = [A_extent_[1], A_extent_[0]]
|
||||
if transB:
|
||||
B_extent_ = [B_extent_[1], B_extent_[0]]
|
||||
# C: MxN, A: MxK, B: NxK
|
||||
if not all(
|
||||
[
|
||||
analyzer.can_prove_equal(C_extent_[0], A_extent_[0]),
|
||||
analyzer.can_prove_equal(C_extent_[1], B_extent_[0]),
|
||||
analyzer.can_prove_equal(A_extent_[1], B_extent_[1]),
|
||||
]
|
||||
):
|
||||
return False
|
||||
return True
|
||||
@@ -0,0 +1,326 @@
|
||||
# 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.
|
||||
|
||||
"""Layout analysis utilities for local-memory op dispatches.
|
||||
|
||||
Provides functions for analyzing TileLayout thread/local partitions,
|
||||
computing local region info, layout signature comparison, and thread
|
||||
variable resolution. Used by cast.py, unary.py, and binary.py.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import operator
|
||||
from collections import defaultdict
|
||||
|
||||
from tvm.arith import Analyzer
|
||||
from tvm.tirx.layout import TileLayout
|
||||
|
||||
|
||||
def get_sublayout_from_region(layout, buffer_shape, region_st, region_extent):
|
||||
"""Get sublayout by slicing the layout with the buffer region.
|
||||
|
||||
Args:
|
||||
layout: The buffer's TileLayout.
|
||||
buffer_shape: The buffer's shape.
|
||||
region_st: Region start indices.
|
||||
region_extent: Region extents.
|
||||
|
||||
Returns:
|
||||
Sublayout if slicing succeeds, otherwise the original layout.
|
||||
"""
|
||||
if not layout:
|
||||
return layout
|
||||
region = [(region_st[i], region_st[i] + region_extent[i]) for i in range(len(region_st))]
|
||||
sliced = layout.slice(list(buffer_shape), region)
|
||||
return sliced if sliced is not None else layout
|
||||
|
||||
|
||||
def get_layout_thread_local_partition(layout):
|
||||
"""Extract thread and local dimension info from layout.
|
||||
|
||||
Returns:
|
||||
tuple | None: On success, (thread_groups, local_dim_indices, local_extents).
|
||||
- thread_groups: dict {axis: (dim_indices, extents)} for each thread axis
|
||||
- local_dim_indices: list of dimension indices for local (memory) axes
|
||||
- local_extents: list of extents for local dimensions
|
||||
Returns None if layout is not supported.
|
||||
|
||||
Validates:
|
||||
- No stride==0 on thread dims (broadcast/overlap = cross-thread semantics)
|
||||
- Local dims may have arbitrary strides (alignment uses actual layout strides)
|
||||
- No thread axes in replica
|
||||
|
||||
Example:
|
||||
Layout (2, 8, 4, 2):(2@warpid, 4@laneid, 1@laneid, 1@m) returns:
|
||||
- thread_groups = {warpid: ([0], [2]), laneid: ([1, 2], [8, 4])}
|
||||
- local_dim_indices = [3], local_extents = [2]
|
||||
"""
|
||||
if not isinstance(layout, TileLayout):
|
||||
return None
|
||||
|
||||
shard = getattr(layout, "shard", None)
|
||||
if not shard:
|
||||
return None
|
||||
|
||||
# Partition dimensions into thread and local (memory) axes
|
||||
thread_dim_indices = [i for i, it in enumerate(shard) if it.axis.is_thread()]
|
||||
local_dim_indices = [i for i, it in enumerate(shard) if not it.axis.is_thread()]
|
||||
|
||||
if not thread_dim_indices or not local_dim_indices:
|
||||
return None
|
||||
|
||||
analyzer = Analyzer()
|
||||
for idx in thread_dim_indices:
|
||||
if analyzer.can_prove_equal(shard[idx].stride, 0):
|
||||
return None
|
||||
|
||||
# Replica must not contain thread axes
|
||||
replica = getattr(layout, "replica", None)
|
||||
if replica and any(it.axis.is_thread() for it in replica):
|
||||
return None
|
||||
|
||||
# Group thread dimensions by axis
|
||||
thread_groups_dict = defaultdict(list)
|
||||
for idx in thread_dim_indices:
|
||||
thread_groups_dict[shard[idx].axis].append(idx)
|
||||
|
||||
thread_groups = {}
|
||||
|
||||
for axis, dim_indices in thread_groups_dict.items():
|
||||
dim_indices = sorted(dim_indices)
|
||||
extents = [shard[i].extent for i in dim_indices]
|
||||
thread_groups[axis] = (dim_indices, extents)
|
||||
|
||||
local_extents = [shard[i].extent for i in local_dim_indices]
|
||||
return (thread_groups, local_dim_indices, local_extents)
|
||||
|
||||
|
||||
def cast_layout_supported_for_local(layout) -> bool:
|
||||
"""Check that layout is valid for local cast (warp/warpgroup/cta/cluster):
|
||||
filter out cross-thread semantics."""
|
||||
return get_layout_thread_local_partition(layout) is not None
|
||||
|
||||
|
||||
def get_local_region(orig_layout: TileLayout, buffer_shape, region_st, region_extent):
|
||||
"""Compute local storage shape, iteration starts, and extents with validation of region.
|
||||
|
||||
Args:
|
||||
orig_layout: The original (unsliced) TileLayout.
|
||||
buffer_shape: The buffer shape.
|
||||
region_st: Region start in shape space.
|
||||
region_extent: Region extent in shape space.
|
||||
|
||||
Returns:
|
||||
(local_shape, local_st, local_ext), or ([1], [0], [1]) if no local dims.
|
||||
Returns None if the region is invalid (non-contiguous slicing).
|
||||
- local_shape: full storage extents per local dim.
|
||||
- local_st: region start per local dim.
|
||||
- local_ext: region extent per local dim.
|
||||
|
||||
Example:
|
||||
Layout (2, 8, 4, 2):(8@m, 2@laneid, 2@m, 1@m), Shape [16, 8], Region [8:16, :] returns:
|
||||
- local_shape = [2, 8], local_st = [1, 0], local_ext = [1, 8]
|
||||
"""
|
||||
grouped, seps = orig_layout.group(list(buffer_shape))
|
||||
|
||||
local_shape = []
|
||||
local_st = []
|
||||
local_ext = []
|
||||
analyzer = Analyzer()
|
||||
|
||||
for d in range(len(buffer_shape)):
|
||||
shard_range = list(range(seps[d], seps[d + 1]))
|
||||
has_local = any(not grouped.shard[s].axis.is_thread() for s in shard_range)
|
||||
if not has_local:
|
||||
continue
|
||||
|
||||
has_thread = any(grouped.shard[s].axis.is_thread() for s in shard_range)
|
||||
|
||||
if not has_thread:
|
||||
# Pure local shape dim: use shape-level values directly.
|
||||
local_shape.append(buffer_shape[d])
|
||||
local_st.append(region_st[d])
|
||||
local_ext.append(region_extent[d])
|
||||
else:
|
||||
# Decompose start element
|
||||
remaining_st = region_st[d]
|
||||
st_coords = []
|
||||
for i, s_idx in enumerate(shard_range):
|
||||
sub_prod = 1
|
||||
for j in range(i + 1, len(shard_range)):
|
||||
sub_prod = sub_prod * grouped.shard[shard_range[j]].extent
|
||||
st_coords.append(remaining_st // sub_prod)
|
||||
remaining_st = remaining_st % sub_prod
|
||||
|
||||
# Decompose end element
|
||||
remaining_end = region_st[d] + region_extent[d] - 1
|
||||
end_coords = []
|
||||
for i, s_idx in enumerate(shard_range):
|
||||
sub_prod = 1
|
||||
for j in range(i + 1, len(shard_range)):
|
||||
sub_prod = sub_prod * grouped.shard[shard_range[j]].extent
|
||||
end_coords.append(remaining_end // sub_prod)
|
||||
remaining_end = remaining_end % sub_prod
|
||||
|
||||
# check the rectangularity and contiguity of the sliced region
|
||||
cur_local_shape, cur_local_st, cur_local_end = 1, 0, 0
|
||||
for k in reversed(range(len(st_coords))):
|
||||
if grouped.shard[seps[d] + k].axis.is_thread():
|
||||
# for thread dims, region must be contiguous and span full extent
|
||||
if not (
|
||||
analyzer.can_prove_equal(st_coords[k], 0)
|
||||
and analyzer.can_prove_equal(
|
||||
end_coords[k], grouped.shard[seps[d] + k].extent - 1
|
||||
)
|
||||
):
|
||||
return None
|
||||
else:
|
||||
if not analyzer.can_prove_equal(end_coords[k] - st_coords[k], 1) and not (
|
||||
analyzer.can_prove_equal(st_coords[k], 0)
|
||||
and analyzer.can_prove_equal(
|
||||
end_coords[k], grouped.shard[seps[d] + k].extent - 1
|
||||
)
|
||||
):
|
||||
# to ensure contiguity, if the region spans multiple values
|
||||
# in this dim, it must span the full extent
|
||||
return None
|
||||
cur_local_shape *= grouped.shard[seps[d] + k].extent
|
||||
cur_local_st = cur_local_st * grouped.shard[seps[d] + k].extent + st_coords[k]
|
||||
cur_local_end = (
|
||||
cur_local_end * grouped.shard[seps[d] + k].extent + end_coords[k]
|
||||
)
|
||||
|
||||
# double check the validity of the sliced region
|
||||
assert region_extent[d] == functools.reduce(
|
||||
operator.mul, [end - st + 1 for st, end in zip(st_coords, end_coords)], 1
|
||||
)
|
||||
|
||||
# append the local info without thread dims
|
||||
local_shape.append(cur_local_shape)
|
||||
local_st.append(cur_local_st)
|
||||
local_ext.append(cur_local_end - cur_local_st + 1)
|
||||
|
||||
if not local_shape:
|
||||
return [1], [0], [1] # treat no local dim case as 1D local shape with 1 element
|
||||
return local_shape, local_st, local_ext
|
||||
|
||||
|
||||
def compute_linear_offset(region_st, local_dims, layout):
|
||||
"""Compute linear offset using layout's actual strides.
|
||||
|
||||
Physical offset = sum(region_st[dim] * layout.shard[dim].stride) for all local dims.
|
||||
"""
|
||||
offset = 0
|
||||
for dim_idx in local_dims:
|
||||
offset = offset + region_st[dim_idx] * layout.shard[dim_idx].stride
|
||||
return offset
|
||||
|
||||
|
||||
def _axis_key(axis):
|
||||
if hasattr(axis, "name") and axis.name:
|
||||
return str(axis.name)
|
||||
return str(axis)
|
||||
|
||||
|
||||
def layout_signature(layout):
|
||||
"""Return semantic signature from canonicalized TileLayout.
|
||||
|
||||
Returns (thread_sig, local_sig, replica_sig).
|
||||
Each sig is a list of (axis_key, extent, stride) in shard/replica order.
|
||||
"""
|
||||
if not isinstance(layout, TileLayout):
|
||||
return None
|
||||
shard = getattr(layout, "shard", None)
|
||||
if not shard:
|
||||
return None
|
||||
|
||||
thread_sig = []
|
||||
local_sig = []
|
||||
for it in shard:
|
||||
item = (_axis_key(it.axis), it.extent, it.stride)
|
||||
if it.axis.is_thread():
|
||||
thread_sig.append(item)
|
||||
else:
|
||||
local_sig.append(item)
|
||||
|
||||
replica_sig = []
|
||||
replica = getattr(layout, "replica", None) or []
|
||||
for it in replica:
|
||||
replica_sig.append((_axis_key(it.axis), it.extent, it.stride))
|
||||
return (thread_sig, local_sig, replica_sig)
|
||||
|
||||
|
||||
def sig_equal(analyzer: Analyzer, src_sig, dst_sig) -> bool:
|
||||
"""Compare two layout signatures with semantic equality (Analyzer).
|
||||
|
||||
Signatures come from layout_signature(layout) and are:
|
||||
(thread_sig, local_sig, replica_sig)
|
||||
Each sig element is (axis_key, extent, stride).
|
||||
"""
|
||||
if src_sig is None or dst_sig is None:
|
||||
return False
|
||||
|
||||
src_thread_sig, src_local_sig, src_replica_sig = src_sig
|
||||
dst_thread_sig, dst_local_sig, dst_replica_sig = dst_sig
|
||||
|
||||
if len(src_thread_sig) != len(dst_thread_sig):
|
||||
return False
|
||||
if len(src_local_sig) != len(dst_local_sig):
|
||||
return False
|
||||
if len(src_replica_sig) != len(dst_replica_sig):
|
||||
return False
|
||||
|
||||
def _list_equal(a_list, b_list) -> bool:
|
||||
for (a_key, a_ext, a_str), (b_key, b_ext, b_str) in zip(a_list, b_list):
|
||||
if a_key != b_key:
|
||||
return False
|
||||
if not analyzer.can_prove_equal(a_ext, b_ext):
|
||||
return False
|
||||
if not analyzer.can_prove_equal(a_str, b_str):
|
||||
return False
|
||||
return True
|
||||
|
||||
return (
|
||||
_list_equal(src_thread_sig, dst_thread_sig)
|
||||
and _list_equal(src_local_sig, dst_local_sig)
|
||||
and _list_equal(src_replica_sig, dst_replica_sig)
|
||||
)
|
||||
|
||||
|
||||
def resolve_thread_var(axis, sctx):
|
||||
"""Map the axis to the corresponding thread variable."""
|
||||
axis_name = getattr(axis, "name", None)
|
||||
if not axis_name:
|
||||
try:
|
||||
axis_name = str(axis)
|
||||
except Exception:
|
||||
axis_name = ""
|
||||
|
||||
for key, itervar in sctx.launch_params.items():
|
||||
if getattr(itervar.var, "name", "") == axis_name:
|
||||
return itervar.var
|
||||
|
||||
if axis_name:
|
||||
axis_name_lower = axis_name.lower()
|
||||
for key in sctx.launch_params:
|
||||
if axis_name_lower in key.lower() or (axis_name == "tx" and "threadIdx.x" in key):
|
||||
return sctx.launch_params[key].var
|
||||
|
||||
if "threadIdx.x" in sctx.launch_params:
|
||||
return sctx.launch_params["threadIdx.x"].var
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,18 @@
|
||||
# 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.
|
||||
|
||||
from .warp_xor_swizzle import *
|
||||
@@ -0,0 +1,430 @@
|
||||
# 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.
|
||||
|
||||
"""CUDA permute_layout dispatch: warp register-staged in-place transpose with
|
||||
optional per-lane XOR-swizzle to avoid SMEM bank conflicts on the write phase.
|
||||
|
||||
The dispatcher reasons about the **layout's shard**, not the buffer's
|
||||
declared shape (the two can differ — a buffer with ``shape=(PIPE, M, K)``
|
||||
may carry a layout whose shard has more dims internally, with grouping
|
||||
mapping shard segments onto buffer dims). Concretely:
|
||||
|
||||
src_sliced = src.layout.slice(src.shape, region).canonicalize()
|
||||
dst_sliced = dst.layout.slice(dst.shape, region).canonicalize()
|
||||
# If the two sliced shards have different structures (which is common —
|
||||
# a linear layout collapses to 1D under canon while a transposed one
|
||||
# keeps its multi-dim structure), regroup src to dst's shape.
|
||||
if src_sliced.shard != dst_sliced.shard:
|
||||
src_sliced, _ = src_sliced.group(dst.shard.extents)
|
||||
extent = [int(it.extent) for it in dst_sliced.shard] # iteration shape
|
||||
src_str = [int(it.stride) for it in src_sliced.shard]
|
||||
dst_str = [int(it.stride) for it in dst_sliced.shard]
|
||||
|
||||
The algorithm:
|
||||
|
||||
regs[P]
|
||||
for r in 0..P:
|
||||
j = r XOR ((lane >> SHIFT) & MASK)
|
||||
i = lane + j * 32 # flat logical index
|
||||
idx = decompose(i, extent) # iter multi-dim index
|
||||
regs[r] = src[project(idx, src.shape, slice_starts)]
|
||||
warp_sync()
|
||||
for r in 0..P:
|
||||
j = r XOR ((lane >> SHIFT) & MASK)
|
||||
i = lane + j * 32
|
||||
idx = decompose(i, extent)
|
||||
dst[project(idx, dst.shape, slice_starts)] = regs[r]
|
||||
warp_sync()
|
||||
|
||||
where ``project`` mixed-radix-folds the iter shard dims back onto the
|
||||
buffer's iterated slice dims (so the emit's index matches buf.shape rank,
|
||||
which TIR's BufferLoad/Store requires).
|
||||
|
||||
SHIFT and MASK are chosen by simulating the bank pattern at the **shard
|
||||
granularity** (where strides are affine), trying k = 0, 1, …, log2(P)
|
||||
and picking the smallest k that makes both phases bank-conflict-free.
|
||||
|
||||
Correctness rests on:
|
||||
|
||||
* For each lane, ``r ↦ r XOR const`` is a bijection on ``[0, P)``.
|
||||
* Therefore (lane, r) ↔ flat over [0, V).
|
||||
* Both layouts are verified bijections on the slice (every logical
|
||||
position has a unique byte offset under that layout).
|
||||
* The mixed-radix projection from iter shard idx to buf coord is exactly
|
||||
what TIR's BufferLoad does internally when buf.shape rank < shard rank
|
||||
— so iter shard's strides and the buffer-indexed byte offset agree.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
from tvm.runtime import DataType
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import Buffer, BufferRegion, IntImm, PrimFunc
|
||||
from tvm.tirx.layout import TileLayout, _flatten_coord
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext, fail, register_dispatch
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..common import get_indices, get_st_extent
|
||||
|
||||
# ---------- helpers ----------------------------------------------------------
|
||||
|
||||
|
||||
def _as_buffer_and_region(arg):
|
||||
"""Normalize a Buffer or BufferRegion to (buffer, start_list, extent_list)."""
|
||||
if isinstance(arg, Buffer):
|
||||
buf = arg
|
||||
extent = list(buf.shape)
|
||||
st = [0] * len(extent)
|
||||
elif isinstance(arg, BufferRegion):
|
||||
buf = arg.buffer
|
||||
st, extent = get_st_extent(arg)
|
||||
else:
|
||||
raise TypeError(f"unexpected permute_layout arg type: {type(arg)}")
|
||||
return buf, list(st), list(extent)
|
||||
|
||||
|
||||
def _as_int(x):
|
||||
"""Return int(x) if x is int-like, else None."""
|
||||
if isinstance(x, int):
|
||||
return x
|
||||
if isinstance(x, IntImm):
|
||||
return int(x.value)
|
||||
if hasattr(x, "value") and isinstance(x.value, int):
|
||||
return int(x.value)
|
||||
try:
|
||||
return int(x)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _layout_shard_int(layout):
|
||||
"""Return (extents, strides) as int lists from a TileLayout's shard, or (None, None)."""
|
||||
if not isinstance(layout, TileLayout):
|
||||
return None, None
|
||||
extents, strides = [], []
|
||||
for it in layout.shard:
|
||||
e = _as_int(it.extent)
|
||||
s = _as_int(it.stride)
|
||||
if e is None or s is None:
|
||||
return None, None
|
||||
extents.append(e)
|
||||
strides.append(s)
|
||||
return extents, strides
|
||||
|
||||
|
||||
def _decompose_row_major(i, extent):
|
||||
out, rem = [], i
|
||||
for e in reversed(extent):
|
||||
out.append(rem % e)
|
||||
rem //= e
|
||||
return list(reversed(out))
|
||||
|
||||
|
||||
def _eval_offset(idx, strides):
|
||||
return sum(i * s for i, s in zip(idx, strides))
|
||||
|
||||
|
||||
def _check_bijection(extent, strides):
|
||||
"""Iteration extents + strides define a bijection on [0, V)?"""
|
||||
V = math.prod(extent)
|
||||
seen = set()
|
||||
for i in range(V):
|
||||
off = _eval_offset(_decompose_row_major(i, extent), strides)
|
||||
if off in seen:
|
||||
return False
|
||||
seen.add(off)
|
||||
return len(seen) == V
|
||||
|
||||
|
||||
def _bank_free(extent, strides, dtype_bytes, P, k):
|
||||
"""For every register slot r ∈ [0, P), do the 32 lanes hit 32 distinct banks?"""
|
||||
T, BANKS, BANK_W = 32, 32, 4
|
||||
shift = 5 - k
|
||||
mask = (1 << k) - 1
|
||||
for r in range(P):
|
||||
seen = set()
|
||||
for lane in range(T):
|
||||
j = r ^ ((lane >> shift) & mask)
|
||||
flat = lane + j * T
|
||||
idx = _decompose_row_major(flat, extent)
|
||||
off_bytes = _eval_offset(idx, strides) * dtype_bytes
|
||||
bank = (off_bytes // BANK_W) % BANKS
|
||||
if bank in seen:
|
||||
return False
|
||||
seen.add(bank)
|
||||
return True
|
||||
|
||||
|
||||
def _choose_xor_k(extent, src_strides, dst_strides, dtype_bytes, P):
|
||||
max_k = int(math.log2(P)) if P > 0 else 0
|
||||
for k in range(max_k + 1):
|
||||
if _bank_free(extent, src_strides, dtype_bytes, P, k) and _bank_free(
|
||||
extent, dst_strides, dtype_bytes, P, k
|
||||
):
|
||||
return k
|
||||
return None
|
||||
|
||||
|
||||
# ---------- validator + dispatch impl ---------------------------------------
|
||||
|
||||
|
||||
def _gather(op_call):
|
||||
op_call = TilePrimitiveCall.downcast(op_call)
|
||||
dst_arg, src_arg = op_call.args[0], op_call.args[1]
|
||||
src_buf, src_st, src_ext = _as_buffer_and_region(src_arg)
|
||||
dst_buf, dst_st, dst_ext = _as_buffer_and_region(dst_arg)
|
||||
return src_buf, src_st, src_ext, dst_buf, dst_st, dst_ext
|
||||
|
||||
|
||||
def _why_reject(op_call, sctx):
|
||||
if not sctx.is_warp:
|
||||
return f"scope {sctx.scope_kind!r} is not 'warp'"
|
||||
if "threadIdx.y" in sctx.launch_params or "threadIdx.z" in sctx.launch_params:
|
||||
return "multi-dim threadIdx is not supported"
|
||||
|
||||
src_buf, src_st, src_ext, dst_buf, dst_st, dst_ext = _gather(op_call)
|
||||
|
||||
if src_buf.dtype != dst_buf.dtype:
|
||||
return f"dtype mismatch: dst={dst_buf.dtype} vs src={src_buf.dtype}"
|
||||
|
||||
src_ext_i = [_as_int(e) for e in src_ext]
|
||||
dst_ext_i = [_as_int(e) for e in dst_ext]
|
||||
if None in src_ext_i or None in dst_ext_i:
|
||||
return "extents must be compile-time integers"
|
||||
if src_ext_i != dst_ext_i:
|
||||
return f"slice shape mismatch: src={src_ext_i} vs dst={dst_ext_i}"
|
||||
|
||||
dtype_bytes = DataType(src_buf.dtype).bits // 8
|
||||
if dtype_bytes not in (1, 2, 4, 8, 16):
|
||||
return f"unsupported dtype byte width: {dtype_bytes}"
|
||||
|
||||
if not isinstance(src_buf.layout, TileLayout):
|
||||
return "src buffer's layout is not a plain TileLayout"
|
||||
if not isinstance(dst_buf.layout, TileLayout):
|
||||
return "dst buffer's layout is not a plain TileLayout"
|
||||
|
||||
# Slice + canonicalize both layouts. The result's shard describes the
|
||||
# iteration domain; runtime starts (like ``ks``) are folded into the
|
||||
# layout's offset, separate from the shard's affine part.
|
||||
src_region = [(s, s + e) for s, e in zip(src_st, src_ext)]
|
||||
dst_region = [(s, s + e) for s, e in zip(dst_st, dst_ext)]
|
||||
src_sliced = src_buf.layout.slice(list(src_buf.shape), src_region)
|
||||
dst_sliced = dst_buf.layout.slice(list(dst_buf.shape), dst_region)
|
||||
if src_sliced is None or dst_sliced is None:
|
||||
return "layout.slice failed"
|
||||
src_sliced = src_sliced.canonicalize()
|
||||
dst_sliced = dst_sliced.canonicalize()
|
||||
|
||||
# Iteration shape: regroup dst onto the iterated buf dims; the result's
|
||||
# shard may stay finer than iter_buf_extents (one buf dim ↔ several shard
|
||||
# dims via seps), which is fine. Then regroup src to match dst's shard
|
||||
# extents exactly so both phases share the same iteration index space.
|
||||
iter_buf_extents = [e for e in src_ext_i if e != 1]
|
||||
try:
|
||||
dst_grouped, dst_seps = dst_sliced.group(iter_buf_extents)
|
||||
src_grouped, _ = src_sliced.group([int(it.extent) for it in dst_grouped.shard])
|
||||
except Exception as e:
|
||||
return f"layout.group failed: {e}"
|
||||
|
||||
dst_ext_, dst_str_ = _layout_shard_int(dst_grouped)
|
||||
src_ext_, src_str_ = _layout_shard_int(src_grouped)
|
||||
if dst_ext_ is None or src_ext_ is None:
|
||||
return "regrouped layout shard contains non-integer extent/stride"
|
||||
if src_ext_ != dst_ext_:
|
||||
return f"src shard {src_ext_} doesn't match dst shard {dst_ext_} after regrouping"
|
||||
|
||||
extent = dst_ext_
|
||||
V = math.prod(extent)
|
||||
T = 32
|
||||
if V == 0 or V % T != 0:
|
||||
return f"volume {V} not divisible by warp size {T}"
|
||||
P = V // T
|
||||
if P == 0 or (P & (P - 1)) != 0 or P > T:
|
||||
return f"per-thread count {P} must be power of 2 in [1, {T}]"
|
||||
|
||||
if not _check_bijection(extent, src_str_):
|
||||
return "src layout (regrouped) is not a bijection on the slice"
|
||||
if not _check_bijection(extent, dst_str_):
|
||||
return "dst layout is not a bijection on the slice"
|
||||
return None
|
||||
|
||||
|
||||
def _impl(op_call, sctx):
|
||||
src_buf, src_st, src_ext, dst_buf, dst_st, dst_ext = _gather(op_call)
|
||||
src_ext_i = [_as_int(e) for e in src_ext]
|
||||
|
||||
src_region = [(s, s + e) for s, e in zip(src_st, src_ext)]
|
||||
dst_region = [(s, s + e) for s, e in zip(dst_st, dst_ext)]
|
||||
src_sliced = src_buf.layout.slice(list(src_buf.shape), src_region).canonicalize()
|
||||
dst_sliced = dst_buf.layout.slice(list(dst_buf.shape), dst_region).canonicalize()
|
||||
|
||||
iter_buf_extents = [e for e in src_ext_i if e != 1]
|
||||
dst_grouped, dst_seps = dst_sliced.group(iter_buf_extents)
|
||||
src_grouped, _ = src_sliced.group([int(it.extent) for it in dst_grouped.shard])
|
||||
|
||||
extent, dst_str_ = _layout_shard_int(dst_grouped)
|
||||
_, src_str_ = _layout_shard_int(src_grouped)
|
||||
V = math.prod(extent)
|
||||
P = V // 32
|
||||
dtype_bytes = DataType(src_buf.dtype).bits // 8
|
||||
|
||||
k_opt = _choose_xor_k(extent, src_str_, dst_str_, dtype_bytes, P)
|
||||
if k_opt is None:
|
||||
fail(f"no XOR-bits k ∈ [0, log2(P)={int(math.log2(P))}] makes both phases bank-free")
|
||||
|
||||
shift = 5 - k_opt
|
||||
mask = (1 << k_opt) - 1
|
||||
|
||||
iter_buf_dims = [i for i, e in enumerate(src_ext_i) if e != 1]
|
||||
seps = list(dst_seps)
|
||||
|
||||
def _project(iter_idx, st_list):
|
||||
buf_idx = list(st_list)
|
||||
for bi in range(len(seps) - 1):
|
||||
lo, hi = seps[bi], seps[bi + 1]
|
||||
flat = _flatten_coord(iter_idx[lo:hi], extent[lo:hi])
|
||||
buf_idx[iter_buf_dims[bi]] = st_list[iter_buf_dims[bi]] + flat
|
||||
return tuple(buf_idx)
|
||||
|
||||
tid_x = sctx.launch_params["threadIdx.x"]
|
||||
dtype = src_buf.dtype
|
||||
|
||||
# Shared 32/64b: base ptr + stride offset avoids buf[] flatten IMAD path.
|
||||
direct = (
|
||||
dtype_bytes in (4, 8)
|
||||
and "shared" in str(src_buf.scope())
|
||||
and "shared" in str(dst_buf.scope())
|
||||
)
|
||||
ptx_t = f"b{dtype_bytes * 8}"
|
||||
# ld/st only move bits, so use an unsigned container of the matching width:
|
||||
# ``ptx.ld(..., bN)`` rejects a float return dtype, and the permute is a pure
|
||||
# byte shuffle, so a float32/float64 tile loads/stores correctly as uint.
|
||||
bits_dtype = f"uint{dtype_bytes * 8}"
|
||||
|
||||
def _iter_off(iter_idx, strides):
|
||||
return sum(iter_idx[d] * strides[d] for d in range(len(strides)))
|
||||
|
||||
# fmt: off
|
||||
if direct:
|
||||
@T.prim_func
|
||||
def impl():
|
||||
warp_size = T.meta_var(32)
|
||||
lane_id = T.meta_var(tid_x % warp_size)
|
||||
regs = T.alloc_buffer((P,), bits_dtype, scope="local")
|
||||
base_src = T.meta_var(src_buf.ptr_to(list(src_st)))
|
||||
base_dst = T.meta_var(dst_buf.ptr_to(list(dst_st)))
|
||||
# Phase 1: read via L_src
|
||||
for r in T.unroll(0, P):
|
||||
j = T.meta_var(r ^ ((lane_id >> shift) & mask))
|
||||
flat = T.meta_var(lane_id + j * warp_size)
|
||||
iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
|
||||
off = T.meta_var(_iter_off(iter_idx, src_str_))
|
||||
ptr = T.meta_var(T.ptr_byte_offset(base_src, off * dtype_bytes, dtype))
|
||||
regs[r] = T.ptx.ld(ptr, bits_dtype, ptx_t, space="shared")
|
||||
T.cuda.warp_sync()
|
||||
# Phase 2: write via L_dst
|
||||
for r in T.unroll(0, P):
|
||||
j = T.meta_var(r ^ ((lane_id >> shift) & mask))
|
||||
flat = T.meta_var(lane_id + j * warp_size)
|
||||
iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
|
||||
off = T.meta_var(_iter_off(iter_idx, dst_str_))
|
||||
ptr = T.meta_var(T.ptr_byte_offset(base_dst, off * dtype_bytes, dtype))
|
||||
T.evaluate(T.ptx.st(ptr, regs[r], space="shared", ptx_type=ptx_t))
|
||||
T.cuda.warp_sync()
|
||||
else:
|
||||
@T.prim_func
|
||||
def impl():
|
||||
warp_size = T.meta_var(32)
|
||||
lane_id = T.meta_var(tid_x % warp_size)
|
||||
regs = T.alloc_buffer((P,), dtype, scope="local")
|
||||
# Phase 1: read via L_src
|
||||
for r in T.unroll(0, P):
|
||||
j = T.meta_var(r ^ ((lane_id >> shift) & mask))
|
||||
flat = T.meta_var(lane_id + j * warp_size)
|
||||
iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
|
||||
src_idx = T.meta_var(_project(iter_idx, src_st))
|
||||
regs[r] = src_buf[tuple(src_idx)]
|
||||
T.cuda.warp_sync()
|
||||
# Phase 2: write via L_dst
|
||||
for r in T.unroll(0, P):
|
||||
j = T.meta_var(r ^ ((lane_id >> shift) & mask))
|
||||
flat = T.meta_var(lane_id + j * warp_size)
|
||||
iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
|
||||
dst_idx = T.meta_var(_project(iter_idx, dst_st))
|
||||
dst_buf[tuple(dst_idx)] = regs[r]
|
||||
T.cuda.warp_sync()
|
||||
# fmt: on
|
||||
return impl
|
||||
|
||||
|
||||
# === Variant: permute_layout/warp_xor_swizzle (priority=20) ============
|
||||
#
|
||||
# When: warp scope; matching dst/src dtype + slice shape; both buffers carry
|
||||
# a plain TileLayout; after slice + canonicalize (and regrouping src to dst's
|
||||
# structure if needed), the iteration extents form a power-of-2 ≤32 elements
|
||||
# per lane; both layouts are bijections on the slice; and there exists an
|
||||
# XOR-bits ``k`` that makes both phases bank-conflict-free.
|
||||
#
|
||||
# Buffer ``shape`` rank does NOT need to equal layout ``shard`` rank — the
|
||||
# dispatcher uses the layout shard for iteration (after slice+canon) and
|
||||
# projects back onto ``buf.shape`` via mixed-radix grouping for the emit.
|
||||
#
|
||||
# Before (TilePrimitiveCall):
|
||||
# with T.warp():
|
||||
# # SFA_smem: u32 (PIPE, BLK_SFA//32, 32), layout shard 4D
|
||||
# # (PIPE, BLK_SFA//128, 4, 32) strides (BLK_SFA, 128, 32, 1)
|
||||
# # SFA_post: same shape; layout shard 4D, strides (BLK_SFA, 128, 1, 4)
|
||||
# Tx.permute_layout(SFA_post[ks, :, :], SFA_smem[ks, :, :])
|
||||
#
|
||||
# After (BLK_SFA=128, P=4, k=2, shift=3):
|
||||
# lane_id = threadIdx.x % 32
|
||||
# regs = T.alloc_buffer((4,), "uint32", scope="local")
|
||||
# for r in T.unroll(4):
|
||||
# j = r ^ ((lane_id >> 3) & 0x3)
|
||||
# flat = lane_id + j * 32
|
||||
# (g, l) = decompose(flat, extent=[4, 32])
|
||||
# regs[r] = src[ks, g, l]
|
||||
# T.cuda.warp_sync()
|
||||
# for r in T.unroll(4):
|
||||
# j = r ^ ((lane_id >> 3) & 0x3)
|
||||
# flat = lane_id + j * 32
|
||||
# (g, l) = decompose(flat, extent=[4, 32])
|
||||
# dst[ks, g, l] = regs[r]
|
||||
# T.cuda.warp_sync()
|
||||
@register_dispatch(
|
||||
"permute_layout",
|
||||
"cuda",
|
||||
variant="warp_xor_swizzle",
|
||||
priority=20,
|
||||
)
|
||||
def permute_layout_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
|
||||
reason = _why_reject(op, sctx)
|
||||
if reason is not None:
|
||||
fail(reason)
|
||||
return _impl(op, sctx)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"_bank_free",
|
||||
"_check_bijection",
|
||||
"_choose_xor_k",
|
||||
"_decompose_row_major",
|
||||
"_eval_offset",
|
||||
"permute_layout_dispatch",
|
||||
]
|
||||
@@ -0,0 +1,20 @@
|
||||
# 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.
|
||||
|
||||
from .local import *
|
||||
from .shared import *
|
||||
from .sm100_packed import *
|
||||
@@ -0,0 +1,484 @@
|
||||
# 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.
|
||||
|
||||
"""CUDA reduction operator dispatch: local-memory variant.
|
||||
|
||||
Registered ops: sum, max, min.
|
||||
|
||||
When: dst and src are both local-scope buffers with matching dtype, on CUDA.
|
||||
|
||||
(A) Thread scope -- sequential per-element reduction
|
||||
(_emit_reduction_local_thread_wise):
|
||||
|
||||
Before:
|
||||
Tx.sum(B_local[0:2, 0:3], A_local[0:2, 0:3, 0:4], [-1], False)
|
||||
|
||||
After (scheduled PrimFunc, spatial_len=6, reduction_len=4):
|
||||
for spa in range(6):
|
||||
B_local[spa] = T.float32(0.0) # init (skipped if accum)
|
||||
for red in range(4):
|
||||
B_local[spa] = B_local[spa] + A_local[spa * 4 + red]
|
||||
|
||||
(B) Warp/Warpgroup scope -- layout-driven reduction
|
||||
(_emit_reduction_local_view):
|
||||
Requires TileLayout with valid thread-partition. Decomposes layout to
|
||||
identify thread-local elements, then optionally shuffles partial sums.
|
||||
|
||||
thread_reduce=False: local-only, no shuffle (warp and warpgroup).
|
||||
thread_reduce=True: local reduction + cross-thread shfl_xor steps (warp only).
|
||||
accum=True + shuffle: saves old dst before reduce+shuffle, combines after (warp only).
|
||||
|
||||
Before:
|
||||
Tx.warp.sum(red_view[0:16, 0:4], acc_view[0:16, 0:128], [-1], False,
|
||||
thread_reduce=True)
|
||||
|
||||
After (scheduled PrimFunc, local_total=2, local_red=32, 2 shuffle steps):
|
||||
src_local = acc_view.view(64)
|
||||
dst_local = red_view.view(2)
|
||||
for spa in range(2):
|
||||
dst_local[spa] = T.float32(0.0)
|
||||
for red in range(32):
|
||||
dst_local[spa] = dst_local[spa] + src_local[...]
|
||||
dst_local[spa] = dst_local[spa] + shfl_xor(..., 1, 32, 32)
|
||||
dst_local[spa] = dst_local[spa] + shfl_xor(..., 2, 32, 32)
|
||||
"""
|
||||
|
||||
import functools
|
||||
import operator
|
||||
from typing import Any
|
||||
|
||||
from tvm.arith.analyzer import Analyzer
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import BufferRegion, PrimFunc
|
||||
from tvm.tirx.layout import TileLayout, laneid
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
|
||||
from tvm.tirx.operator.tile_primitive.common import ReduceOpType
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..common import get_indices, get_st_extent
|
||||
from ..layout_utils import get_local_region, get_sublayout_from_region
|
||||
from .utils import (
|
||||
_REDUCE_OP_TO_STR,
|
||||
_analyze_axes,
|
||||
_analyze_layout_dims,
|
||||
_build_local_dim_map,
|
||||
_compute_shuffle_masks,
|
||||
_match_reduction_storage_scope,
|
||||
_reduction_args,
|
||||
_validate_reduction_layout,
|
||||
reduce_default_value_table,
|
||||
reduce_op_table,
|
||||
)
|
||||
|
||||
|
||||
def _analyze_shuffle_reduce(src_layout, dst_layout):
|
||||
"""Analyze src/dst layouts for laneid shard->replica reduce pattern.
|
||||
|
||||
Returns (reduce_width, local_elems) if the pattern matches, or None.
|
||||
- reduce_width: number of lanes participating in each group's reduction
|
||||
- local_elems: per-thread element count (product of non-laneid shard extents)
|
||||
"""
|
||||
if src_layout.is_swizzle() or dst_layout.is_swizzle():
|
||||
return None
|
||||
|
||||
src_canon = src_layout.canonicalize()
|
||||
dst_canon = dst_layout.canonicalize()
|
||||
|
||||
# Extract laneid iters from shard and replica
|
||||
src_laneid_shard = [it for it in src_canon.shard if it.axis == laneid]
|
||||
dst_laneid_replica = [it for it in dst_canon.replica if it.axis == laneid]
|
||||
|
||||
# src shard must contain laneid (data distributed across lanes)
|
||||
if not src_laneid_shard:
|
||||
return None
|
||||
# dst replica must contain laneid (result broadcast to lanes)
|
||||
if not dst_laneid_replica:
|
||||
return None
|
||||
|
||||
# laneid span must be 32 (full warp)
|
||||
src_laneid_span = 1 + sum(abs(int(it.stride)) * (int(it.extent) - 1) for it in src_laneid_shard)
|
||||
if src_laneid_span != 32:
|
||||
return None
|
||||
|
||||
reduce_width = functools.reduce(operator.mul, [int(it.extent) for it in dst_laneid_replica], 1)
|
||||
if reduce_width <= 0 or reduce_width > 32 or (reduce_width & (reduce_width - 1)) != 0:
|
||||
return None # must be power of 2
|
||||
|
||||
# local_elems = product of non-laneid shard extents in src
|
||||
src_non_laneid = [it for it in src_canon.shard if it.axis != laneid]
|
||||
local_elems = functools.reduce(operator.mul, [int(it.extent) for it in src_non_laneid], 1)
|
||||
|
||||
return reduce_width, local_elems
|
||||
|
||||
|
||||
def _gen_warp_shuffle_reduce(src, dst, reduce_width, local_elems, accum, op_type, init_value):
|
||||
"""Generate warp shuffle reduce codegen for laneid shard->replica pattern.
|
||||
|
||||
Unified for both full warp (reduce_width=32) and partial warp (e.g. reduce_width=8).
|
||||
"""
|
||||
is_same_buffer = src.same_as(dst)
|
||||
op_str = _REDUCE_OP_TO_STR[op_type]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
src_local = src.local(local_elems)
|
||||
dst_local = dst.local(local_elems)
|
||||
for k in T.serial(local_elems):
|
||||
if not is_same_buffer:
|
||||
dst_local[k] = src_local[k]
|
||||
dst_local[k] = T.cuda.warp_reduce(dst_local[k], op_str, reduce_width)
|
||||
# fmt: on
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
def validate_reduction_local(
|
||||
op: TilePrimitiveCall, sctx: DispatchContext
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Validate reduction in local memory."""
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
dst_br, src_br = op.output, op.input
|
||||
dst, src = dst_br.buffer, src_br.buffer
|
||||
|
||||
if not (src.scope() == "local" and dst.scope() == "local" and sctx.is_target("cuda")):
|
||||
return False, "expected local scope and CUDA target"
|
||||
if src.dtype != dst.dtype:
|
||||
return False, f"dtype mismatch: src={src.dtype} dst={dst.dtype}"
|
||||
|
||||
if sctx.is_thread:
|
||||
return True, None # thread-wise reduction
|
||||
elif sctx.scope_kind in ["warp", "warpgroup"]:
|
||||
if not sctx.is_warp and op.config.get("thread_reduce", False):
|
||||
return (
|
||||
False,
|
||||
"thread_reduce=True is only supported in warp scope; "
|
||||
"warpgroup local reduction is thread-local only",
|
||||
)
|
||||
# VIEW: need layouts and layout analysis
|
||||
if not (src.layout and dst.layout):
|
||||
return False, "layouts required for view-based local reduction"
|
||||
if not (isinstance(src.layout, TileLayout) and isinstance(dst.layout, TileLayout)):
|
||||
return False, "TileLayout required for view-based local reduction"
|
||||
if src.layout.is_swizzle() or dst.layout.is_swizzle():
|
||||
return False, "swizzle layout unsupported for local reduction"
|
||||
|
||||
analyzer = Analyzer()
|
||||
|
||||
# Validate get_local_region succeeds for both
|
||||
src_st, src_extent = get_st_extent(src_br)
|
||||
dst_st, dst_extent = get_st_extent(dst_br)
|
||||
|
||||
if sctx.is_warp:
|
||||
# Check for laneid shard->replica shuffle reduce pattern first.
|
||||
# This pattern has laneid in dst replica (broadcast), which the
|
||||
# general validation below would reject.
|
||||
shuffle_info = _analyze_shuffle_reduce(src.layout, dst.layout)
|
||||
if shuffle_info is not None:
|
||||
return True, None
|
||||
|
||||
for layout, buf, st, ext, name in [
|
||||
(src.layout, src, src_st, src_extent, "src"),
|
||||
(dst.layout, dst, dst_st, dst_extent, "dst"),
|
||||
]:
|
||||
for it in layout.shard:
|
||||
if it.axis.is_thread() and analyzer.can_prove_equal(it.stride, 0):
|
||||
return False, f"thread dim with zero stride in {name}"
|
||||
replica = getattr(layout, "replica", None) or []
|
||||
if any(it.axis.is_thread() for it in replica):
|
||||
return False, f"thread axis in replica for {name}"
|
||||
if get_local_region(layout, list(buf.shape), st, ext) is None:
|
||||
return False, f"get_local_region failed for {name}"
|
||||
|
||||
# Validate layout compatibility
|
||||
# Spatial dims match, reduce dims in dst have local_extent==1
|
||||
reduce_axes = tuple(int(a) for a in op.reduce_axes)
|
||||
src_ndim = len(src_br.region)
|
||||
try:
|
||||
reduce_dims, _ = _analyze_axes(src_ndim, reduce_axes)
|
||||
except AssertionError as e:
|
||||
return False, str(e)
|
||||
src_sliced = get_sublayout_from_region(src.layout, src.shape, src_st, src_extent)
|
||||
dst_sliced = get_sublayout_from_region(dst.layout, dst.shape, dst_st, dst_extent)
|
||||
ok, msg = _validate_reduction_layout(
|
||||
src_sliced, dst_sliced, list(src_extent), list(dst_extent), reduce_dims
|
||||
)
|
||||
return ok, msg
|
||||
else:
|
||||
return False, f"unsupported exec_scope {sctx.scope_kind} for local reduction"
|
||||
|
||||
|
||||
def _emit_reduction_local_thread_wise(
|
||||
dst_br: BufferRegion,
|
||||
src_br: BufferRegion,
|
||||
accum: bool,
|
||||
reduce_op: ReduceOpType,
|
||||
reduce_dims: list[int],
|
||||
spatial_dims: list[int],
|
||||
) -> PrimFunc:
|
||||
dst, src = dst_br.buffer, src_br.buffer
|
||||
dtype = src.dtype
|
||||
src_st, src_extent = get_st_extent(src_br)
|
||||
dst_st, dst_extent = get_st_extent(dst_br)
|
||||
src_ndim = len(src_extent)
|
||||
spa_extents = [src_extent[d] for d in spatial_dims]
|
||||
red_extents = [src_extent[d] for d in reduce_dims]
|
||||
spatial_len = functools.reduce(operator.mul, spa_extents, 1)
|
||||
reduction_len = functools.reduce(operator.mul, red_extents, 1)
|
||||
|
||||
op_func = reduce_op_table.get(reduce_op)
|
||||
assert op_func is not None
|
||||
init_value = reduce_default_value_table(dtype).get(reduce_op)
|
||||
|
||||
def get_src_indices(spa_fused, red_fused):
|
||||
spa_indices = []
|
||||
rem = spa_fused
|
||||
for e in reversed(spa_extents):
|
||||
spa_indices.append(rem % e)
|
||||
rem //= e
|
||||
spa_indices.reverse()
|
||||
|
||||
red_indices = []
|
||||
rem = red_fused
|
||||
for e in reversed(red_extents):
|
||||
red_indices.append(rem % e)
|
||||
rem //= e
|
||||
red_indices.reverse()
|
||||
|
||||
full = [None] * src_ndim
|
||||
for i, d in enumerate(spatial_dims):
|
||||
full[d] = spa_indices[i] + src_st[d]
|
||||
for i, d in enumerate(reduce_dims):
|
||||
full[d] = red_indices[i] + src_st[d]
|
||||
return full
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
for spa in T.serial(spatial_len):
|
||||
dst_idx = T.meta_var(get_indices(spa, dst_st, dst_extent))
|
||||
if not accum:
|
||||
dst[tuple(dst_idx)] = init_value
|
||||
for red in T.serial(reduction_len):
|
||||
src_idx = T.meta_var(get_src_indices(spa, red))
|
||||
dst[tuple(dst_idx)] = op_func(dst[tuple(dst_idx)], src[tuple(src_idx)])
|
||||
# fmt: on
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
def _emit_reduction_local_view(
|
||||
dst_br: BufferRegion,
|
||||
src_br: BufferRegion,
|
||||
accum: bool,
|
||||
reduce_op: ReduceOpType,
|
||||
config: dict[str, Any],
|
||||
reduce_dims: set[int],
|
||||
spatial_dims: list[int],
|
||||
src_local_info,
|
||||
dst_local_info,
|
||||
shuffle_masks: list[int],
|
||||
) -> PrimFunc:
|
||||
dst, src = dst_br.buffer, src_br.buffer
|
||||
dtype = src.dtype
|
||||
|
||||
op_func = reduce_op_table.get(reduce_op)
|
||||
assert op_func is not None
|
||||
init_value = reduce_default_value_table(dtype).get(reduce_op)
|
||||
|
||||
src_local_shape, src_local_st, src_local_ext = src_local_info
|
||||
dst_local_shape, dst_local_st, dst_local_ext = dst_local_info
|
||||
|
||||
# Build maps from original dim index to position in get_local_region output
|
||||
src_dim_map = _build_local_dim_map(src.layout, list(src.shape))
|
||||
dst_dim_map = _build_local_dim_map(dst.layout, list(dst.shape))
|
||||
|
||||
# Only include reduction dims that have local parts in src
|
||||
src_ndim = len(src_br.region)
|
||||
reduce_local_dims = [d for d in reduce_dims if src_dim_map[d] is not None]
|
||||
reduction_local_ext = [src_local_ext[src_dim_map[d]] for d in reduce_local_dims]
|
||||
reduction_local_st = [src_local_st[src_dim_map[d]] for d in reduce_local_dims]
|
||||
|
||||
reduction_local_total = functools.reduce(operator.mul, reduction_local_ext, 1)
|
||||
dst_local_total = functools.reduce(operator.mul, dst_local_ext, 1)
|
||||
|
||||
def _get_src_local_index(dst_fused, red_fused):
|
||||
"""Compute src local multi-dim index from dst fused index and reduction fused index."""
|
||||
dst_indices = get_indices(dst_fused, dst_local_st, dst_local_ext)
|
||||
red_indices = get_indices(red_fused, reduction_local_st, reduction_local_ext)
|
||||
|
||||
# Interleave into src local indices (skipping pure-thread dims)
|
||||
src_local = []
|
||||
ri = 0
|
||||
for d in range(src_ndim):
|
||||
if src_dim_map[d] is None:
|
||||
continue # pure-thread in src, not in src.local()
|
||||
if d in reduce_dims:
|
||||
src_local.append(red_indices[ri])
|
||||
ri += 1
|
||||
else:
|
||||
# Spatial dim: use corresponding dst local position
|
||||
src_local.append(dst_indices[dst_dim_map[d]])
|
||||
|
||||
return src_local
|
||||
|
||||
# is_same_buffer = src.same_as(dst)
|
||||
shuffle = bool(config.get("thread_reduce", False))
|
||||
in_place = dst.same_as(src)
|
||||
|
||||
def shuffle_data(mask, dst_local, dst_idx):
|
||||
@T.inline
|
||||
def inner_shuffle(v, shuffle_mask):
|
||||
dst_local[tuple(dst_idx)] = op_func(
|
||||
v, T.tvm_warp_shuffle_xor(mask, v, shuffle_mask, 32, 32)
|
||||
)
|
||||
|
||||
for i in range(len(shuffle_masks)):
|
||||
inner_shuffle(dst_local[tuple(dst_idx)], shuffle_masks[i])
|
||||
|
||||
need_save_accum = accum and shuffle
|
||||
|
||||
# fmt: off
|
||||
if need_save_accum:
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
src_local = src.local(*src_local_shape)
|
||||
dst_local = dst.local(*dst_local_shape)
|
||||
old_val = T.alloc_buffer([1], dtype, scope="local")
|
||||
|
||||
for spa in T.serial(dst_local_total):
|
||||
dst_idx = T.meta_var(get_indices(spa, dst_local_st, dst_local_ext))
|
||||
old_val[0] = dst_local[tuple(dst_idx)]
|
||||
if not in_place:
|
||||
dst_local[tuple(dst_idx)] = init_value
|
||||
for red in T.serial(reduction_local_total):
|
||||
src_idx = T.meta_var(_get_src_local_index(spa, red))
|
||||
dst_local[tuple(dst_idx)] = op_func(dst_local[tuple(dst_idx)], src_local[tuple(src_idx)]) # noqa: E501
|
||||
if shuffle:
|
||||
mask = T.tvm_warp_activemask()
|
||||
shuffle_data(mask, dst_local, dst_idx)
|
||||
dst_local[tuple(dst_idx)] = op_func(dst_local[tuple(dst_idx)], old_val[0])
|
||||
else:
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
src_local = src.local(*src_local_shape)
|
||||
dst_local = dst.local(*dst_local_shape)
|
||||
|
||||
for spa in T.serial(dst_local_total):
|
||||
dst_idx = T.meta_var(get_indices(spa, dst_local_st, dst_local_ext))
|
||||
if not in_place:
|
||||
if not accum:
|
||||
dst_local[tuple(dst_idx)] = init_value
|
||||
for red in T.serial(reduction_local_total):
|
||||
src_idx = T.meta_var(_get_src_local_index(spa, red))
|
||||
dst_local[tuple(dst_idx)] = op_func(dst_local[tuple(dst_idx)], src_local[tuple(src_idx)]) # noqa: E501
|
||||
if shuffle:
|
||||
mask = T.tvm_warp_activemask()
|
||||
shuffle_data(mask, dst_local, dst_idx)
|
||||
# fmt: on
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
def reduction_local_impl(
|
||||
op: TilePrimitiveCall, op_type: ReduceOpType, sctx: DispatchContext
|
||||
) -> PrimFunc | None:
|
||||
dst_br, src_br, reduce_axes, accum, config = _reduction_args(op)
|
||||
src_ndim = len(src_br.region)
|
||||
reduce_dims, spatial_dims = _analyze_axes(src_ndim, reduce_axes)
|
||||
|
||||
if sctx.is_thread:
|
||||
return _emit_reduction_local_thread_wise(
|
||||
dst_br, src_br, accum, op_type, reduce_dims, spatial_dims
|
||||
)
|
||||
elif sctx.scope_kind in ["warp", "warpgroup"]:
|
||||
src = src_br.buffer
|
||||
dst = dst_br.buffer
|
||||
|
||||
if sctx.is_warp:
|
||||
# --- Try laneid shard->replica shuffle reduce ---
|
||||
shuffle_info = _analyze_shuffle_reduce(src.layout, dst.layout)
|
||||
if shuffle_info is not None:
|
||||
reduce_width, local_elems = shuffle_info
|
||||
if op_type not in _REDUCE_OP_TO_STR:
|
||||
fail(f"unsupported reduce op: {op_type}")
|
||||
dtype = src.dtype
|
||||
init_value = reduce_default_value_table(dtype).get(op_type)
|
||||
return _gen_warp_shuffle_reduce(
|
||||
src, dst, reduce_width, local_elems, accum, op_type, init_value
|
||||
)
|
||||
elif config.get("thread_reduce", False):
|
||||
fail(
|
||||
"thread_reduce=True is only supported in warp scope; "
|
||||
"warpgroup local reduction is thread-local only"
|
||||
)
|
||||
|
||||
# --- Existing WGMMA layout path below ---
|
||||
src_st, src_extent = get_st_extent(src_br)
|
||||
dst_st, dst_extent = get_st_extent(dst_br)
|
||||
|
||||
src_local_info = get_local_region(src.layout, list(src.shape), src_st, src_extent)
|
||||
dst_local_info = get_local_region(dst.layout, list(dst.shape), dst_st, dst_extent)
|
||||
assert src_local_info is not None and dst_local_info is not None
|
||||
|
||||
src_dim_info = _analyze_layout_dims(src.layout, list(src.shape))
|
||||
shuffle_masks = (
|
||||
_compute_shuffle_masks(src_dim_info, reduce_dims)
|
||||
if config.get("thread_reduce", False)
|
||||
else []
|
||||
)
|
||||
|
||||
return _emit_reduction_local_view(
|
||||
dst_br,
|
||||
src_br,
|
||||
accum,
|
||||
op_type,
|
||||
config,
|
||||
reduce_dims,
|
||||
spatial_dims,
|
||||
src_local_info,
|
||||
dst_local_info,
|
||||
shuffle_masks,
|
||||
)
|
||||
else:
|
||||
fail(f"unsupported exec_scope {sctx.scope_kind} for reduction_local_impl")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registration: local memory reduction (priority=10)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
for op_name, op_type in [
|
||||
("sum", ReduceOpType.SUM),
|
||||
("max", ReduceOpType.MAX),
|
||||
("min", ReduceOpType.MIN),
|
||||
]:
|
||||
|
||||
@register_dispatch(
|
||||
op_name,
|
||||
"cuda",
|
||||
variant="local",
|
||||
priority=10,
|
||||
when=[
|
||||
predicate("storage_scope", _match_reduction_storage_scope, expected_scope=["local"]),
|
||||
predicate("local_valid", validate_reduction_local),
|
||||
],
|
||||
)
|
||||
def _local_dispatch(op: TilePrimitiveCall, sctx: DispatchContext, _op_type=op_type) -> PrimFunc:
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
return reduction_local_impl(op, _op_type, sctx)
|
||||
@@ -0,0 +1,298 @@
|
||||
# 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.
|
||||
|
||||
"""CUDA reduction operator dispatch: shared-memory variant.
|
||||
|
||||
Registered ops: sum, max, min.
|
||||
|
||||
When: dst and src are both shared-memory buffers, exec scope is one of
|
||||
{cta, warpgroup, warp, thread}, threadIdx.x bound, reduce axes valid.
|
||||
|
||||
(A) CTA/warpgroup/warp scope -- adaptive-group shuffle tree
|
||||
(_emit_reduction_shared_cta):
|
||||
group_size = min(next_power_of_2(reduction_len), 32).
|
||||
Each group of threads reduces one spatial position via shfl_xor.
|
||||
|
||||
Before:
|
||||
Tx.cta.sum(B_smem[0:4], A_smem[0:4, 0:8], [-1], False)
|
||||
|
||||
After (scheduled PrimFunc, group_size=8, spatial_par=4):
|
||||
thread_data[0] = T.float32(0.0)
|
||||
thread_data[0] = thread_data[0] + A_smem[tid_in_scope] # gather
|
||||
# log2(8) = 3 shuffle-xor steps with width=8
|
||||
thread_data[0] = thread_data[0] + shfl_xor(thread_data[0], 1, 8, 32)
|
||||
thread_data[0] = thread_data[0] + shfl_xor(thread_data[0], 2, 8, 32)
|
||||
thread_data[0] = thread_data[0] + shfl_xor(thread_data[0], 4, 8, 32)
|
||||
if tid_in_scope % 8 == 0:
|
||||
B_smem[tid_in_scope // 8] = thread_data[0]
|
||||
|
||||
(B) Thread scope -- sequential loop (_emit_reduction_shared_thread):
|
||||
|
||||
Before:
|
||||
if tid == 65:
|
||||
Tx.sum(B_smem[0:4], A_smem[0:4, 0:8], [-1], False)
|
||||
|
||||
After (scheduled PrimFunc):
|
||||
for spa in range(4):
|
||||
B_smem[spa] = T.float32(0.0) # init (skipped if accum)
|
||||
for red in range(8):
|
||||
B_smem[spa] = B_smem[spa] + A_smem[spa * 8 + red]
|
||||
"""
|
||||
|
||||
import functools
|
||||
import math
|
||||
import operator
|
||||
|
||||
from tvm.arith.analyzer import Analyzer
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import BufferRegion, PrimFunc
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
|
||||
from tvm.tirx.operator.tile_primitive.common import ReduceOpType
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..common import get_indices, get_st_extent, next_power_of_2
|
||||
from .utils import (
|
||||
_analyze_axes,
|
||||
_match_reduction_storage_scope,
|
||||
_reduction_args,
|
||||
build_src_indices,
|
||||
reduce_default_value_table,
|
||||
reduce_op_table,
|
||||
)
|
||||
|
||||
|
||||
def validate_reduction_shared(
|
||||
op: TilePrimitiveCall, sctx: DispatchContext
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Validate reduction in shared memory."""
|
||||
if sctx.scope_kind not in ["cta", "warpgroup", "warp", "thread"]:
|
||||
return False, f"unsupported exec_scope {sctx.scope_kind} for shared reduction"
|
||||
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
dst, src = op.output.buffer, op.input.buffer
|
||||
if not (src.scope().startswith("shared") and dst.scope().startswith("shared")):
|
||||
return False, "expected shared scope for both src and dst"
|
||||
if src.dtype != dst.dtype:
|
||||
return False, f"dtype mismatch: src={src.dtype} dst={dst.dtype}"
|
||||
|
||||
if "threadIdx.x" not in sctx.launch_params:
|
||||
return False, "threadIdx.x not in launch_params"
|
||||
if "threadIdx.y" in sctx.launch_params or "threadIdx.z" in sctx.launch_params:
|
||||
return False, "multi-dimensional thread binding not supported for shared reduction"
|
||||
|
||||
reduce_axes = tuple(int(a) for a in op.reduce_axes)
|
||||
src_region = op.input.region
|
||||
dst_region = op.output.region
|
||||
src_ndim = len(src_region)
|
||||
try:
|
||||
reduce_dims, spatial_dims = _analyze_axes(src_ndim, reduce_axes)
|
||||
except AssertionError as e:
|
||||
return False, str(e)
|
||||
|
||||
# Validate dst shape matches spatial dims of src
|
||||
src_extent = [r.extent for r in src_region]
|
||||
dst_extent = [r.extent for r in dst_region]
|
||||
expected_dst_len = functools.reduce(operator.mul, [src_extent[d] for d in spatial_dims], 1)
|
||||
actual_dst_len = functools.reduce(operator.mul, dst_extent, 1)
|
||||
analyzer = Analyzer()
|
||||
if not analyzer.can_prove_equal(expected_dst_len, actual_dst_len):
|
||||
return (False, f"dst size {actual_dst_len} != expected spatial size {expected_dst_len}")
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
def _emit_reduction_shared_cta(
|
||||
dst_br: BufferRegion,
|
||||
src_br: BufferRegion,
|
||||
accum: bool,
|
||||
reduce_op: ReduceOpType,
|
||||
sctx: DispatchContext,
|
||||
reduce_dims: list[int],
|
||||
spatial_dims: list[int],
|
||||
) -> PrimFunc:
|
||||
exec_scope_name = sctx.scope_kind
|
||||
|
||||
def get_thread_cnt():
|
||||
if exec_scope_name == "cta":
|
||||
return sctx.launch_params["threadIdx.x"].dom.extent
|
||||
elif exec_scope_name == "warpgroup":
|
||||
return 128
|
||||
elif exec_scope_name == "warp":
|
||||
return 32
|
||||
elif exec_scope_name == "thread":
|
||||
return 1
|
||||
|
||||
thread_cnt = get_thread_cnt()
|
||||
dst, src = dst_br.buffer, src_br.buffer
|
||||
src_st, src_extent = get_st_extent(src_br)
|
||||
dst_st, dst_extent = get_st_extent(dst_br)
|
||||
dtype = src.dtype
|
||||
|
||||
# Compute spatial/reduction from the explicit axes
|
||||
spatial_len = functools.reduce(operator.mul, [src_extent[d] for d in spatial_dims], 1)
|
||||
reduction_len = functools.reduce(operator.mul, [src_extent[d] for d in reduce_dims], 1)
|
||||
|
||||
op_func = reduce_op_table.get(reduce_op)
|
||||
assert op_func is not None
|
||||
init_value = reduce_default_value_table(dtype).get(reduce_op)
|
||||
|
||||
# Adaptive group size: nearest power-of-2 for reduction length, capped at warp size and thread count. # noqa: E501
|
||||
group_size = min(next_power_of_2(int(reduction_len)), 32, int(thread_cnt))
|
||||
group_size = max(group_size, 1) # ensure at least 1
|
||||
n_shuffles = int(math.log2(group_size)) if group_size > 1 else 0
|
||||
spatial_par = int(thread_cnt) // group_size
|
||||
|
||||
def get_tid_in_scope():
|
||||
tx_var = sctx.launch_params["threadIdx.x"].var
|
||||
if exec_scope_name == "cta":
|
||||
return tx_var
|
||||
elif exec_scope_name in ("warp", "warpgroup"):
|
||||
return tx_var % thread_cnt
|
||||
elif exec_scope_name == "thread":
|
||||
return 0
|
||||
|
||||
def shuffle_data(thread_data):
|
||||
@T.inline
|
||||
def inner_shuffle(mask, v, shuffle_mask):
|
||||
v[0] = op_func(v[0], T.tvm_warp_shuffle_xor(mask, v[0], shuffle_mask, group_size, 32))
|
||||
|
||||
if n_shuffles > 0:
|
||||
mask = T.tvm_warp_activemask()
|
||||
for i in range(n_shuffles):
|
||||
inner_shuffle(mask, thread_data, 1 << i)
|
||||
|
||||
@T.inline
|
||||
def sync():
|
||||
if exec_scope_name == "cta":
|
||||
T.cuda.cta_sync()
|
||||
elif exec_scope_name == "warpgroup":
|
||||
T.cuda.warpgroup_sync(8) # TODO: fix this hardcoded value
|
||||
elif exec_scope_name == "warp":
|
||||
T.cuda.warp_sync()
|
||||
elif exec_scope_name == "thread":
|
||||
pass
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func
|
||||
def impl():
|
||||
tid_in_scope = get_tid_in_scope()
|
||||
thread_data = T.alloc_buffer([1], dtype=dtype, scope="local")
|
||||
group_id = T.meta_var(T.floordiv(tid_in_scope, group_size))
|
||||
lane_in_grp = T.meta_var(tid_in_scope % group_size)
|
||||
for step in T.serial(T.ceildiv(spatial_len, spatial_par)):
|
||||
spa_fused = T.meta_var(step * spatial_par + group_id)
|
||||
if spa_fused < spatial_len:
|
||||
thread_data[0] = init_value
|
||||
for t in T.serial(T.ceildiv(reduction_len, group_size)):
|
||||
red_fused = T.meta_var(t * group_size + lane_in_grp)
|
||||
if red_fused < reduction_len:
|
||||
src_indices = T.meta_var(build_src_indices(spa_fused, red_fused, spatial_dims, reduce_dims, src_extent, src_st)) # noqa: E501
|
||||
thread_data[0] = op_func(thread_data[0], src[tuple(src_indices)])
|
||||
shuffle_data(thread_data)
|
||||
if lane_in_grp == 0:
|
||||
dst_indices = T.meta_var(get_indices(spa_fused, dst_st, dst_extent))
|
||||
dst[tuple(dst_indices)] = T.if_then_else(T.bool(accum), op_func(dst[tuple(dst_indices)], thread_data[0]), thread_data[0]) # noqa: E501
|
||||
|
||||
sync()
|
||||
# fmt: on
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
def _emit_reduction_shared_thread(
|
||||
dst_br: BufferRegion,
|
||||
src_br: BufferRegion,
|
||||
accum: bool,
|
||||
reduce_op: ReduceOpType,
|
||||
sctx: DispatchContext,
|
||||
reduce_dims: list[int],
|
||||
spatial_dims: list[int],
|
||||
) -> PrimFunc:
|
||||
dst, src = dst_br.buffer, src_br.buffer
|
||||
src_st, src_extent = get_st_extent(src_br)
|
||||
dst_st, dst_extent = get_st_extent(dst_br)
|
||||
dtype = src.dtype
|
||||
|
||||
# Compute spatial/reduction from the explicit axes
|
||||
spatial_len = functools.reduce(operator.mul, [src_extent[d] for d in spatial_dims], 1)
|
||||
reduction_len = functools.reduce(operator.mul, [src_extent[d] for d in reduce_dims], 1)
|
||||
|
||||
op_func = reduce_op_table.get(reduce_op)
|
||||
assert op_func is not None
|
||||
init_value = reduce_default_value_table(dtype).get(reduce_op)
|
||||
|
||||
@T.prim_func
|
||||
def impl():
|
||||
for spa_fused in T.serial(spatial_len):
|
||||
dst_indices = T.meta_var(get_indices(spa_fused, dst_st, dst_extent))
|
||||
if not accum:
|
||||
dst[tuple(dst_indices)] = init_value
|
||||
for red_fused in T.serial(reduction_len):
|
||||
src_indices = T.meta_var(
|
||||
build_src_indices(
|
||||
spa_fused, red_fused, spatial_dims, reduce_dims, src_extent, src_st
|
||||
)
|
||||
)
|
||||
dst[tuple(dst_indices)] = op_func(dst[tuple(dst_indices)], src[tuple(src_indices)])
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
def reduction_shared_impl(
|
||||
op: TilePrimitiveCall, op_type: ReduceOpType, sctx: DispatchContext
|
||||
) -> PrimFunc | None:
|
||||
dst_br, src_br, reduce_axes, accum, config = _reduction_args(op)
|
||||
src_ndim = len(src_br.region)
|
||||
reduce_dims, spatial_dims = _analyze_axes(src_ndim, reduce_axes)
|
||||
if sctx.scope_kind in ["cta", "warpgroup", "warp"]:
|
||||
return _emit_reduction_shared_cta(
|
||||
dst_br, src_br, accum, op_type, sctx, reduce_dims, spatial_dims
|
||||
)
|
||||
elif sctx.is_thread:
|
||||
return _emit_reduction_shared_thread(
|
||||
dst_br, src_br, accum, op_type, sctx, reduce_dims, spatial_dims
|
||||
)
|
||||
else:
|
||||
fail(f"unsupported exec_scope {sctx.scope_kind} for reduction_shared_impl")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registration: shared memory reduction (priority=10)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
for op_name, op_type in [
|
||||
("sum", ReduceOpType.SUM),
|
||||
("max", ReduceOpType.MAX),
|
||||
("min", ReduceOpType.MIN),
|
||||
]:
|
||||
|
||||
@register_dispatch(
|
||||
op_name,
|
||||
"cuda",
|
||||
variant="shared",
|
||||
priority=10,
|
||||
when=[
|
||||
predicate("storage_scope", _match_reduction_storage_scope, expected_scope=["shared*"]),
|
||||
predicate("shared_valid", validate_reduction_shared),
|
||||
],
|
||||
)
|
||||
def _shared_dispatch(
|
||||
op: TilePrimitiveCall, sctx: DispatchContext, _op_type=op_type
|
||||
) -> PrimFunc:
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
return reduction_shared_impl(op, _op_type, sctx)
|
||||
@@ -0,0 +1,249 @@
|
||||
# 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.
|
||||
|
||||
"""CUDA reduction operator dispatch: SM100+ packed optimized variant.
|
||||
|
||||
Registered ops: sum, max, min.
|
||||
|
||||
When: thread scope, all local buffers, float32, 1D src with len >= 8,
|
||||
SM100+ (uses packed PTX instructions not available on older GPUs).
|
||||
|
||||
Before (TilePrimitiveCall -- sum example):
|
||||
Tx.sum(dst_local[0:1], src_local[0:32]) # float32, reduce 32 -> 1 (thread scope)
|
||||
|
||||
After -- packed_add_sum (uses add.f32x2 to reduce pairs):
|
||||
# Iteratively reduce: 32 -> 16 -> 8 -> 4 -> 2 -> 1
|
||||
# Each step: add.f32x2 combines adjacent pairs
|
||||
for i in T.serial(16):
|
||||
T.cuda.func_call("add_f32x2", &buf[i*2], &buf[i*2], &buf[i*2+2])
|
||||
# ... repeat halving until scalar result
|
||||
dst_local[0] = buf[0]
|
||||
|
||||
After -- 3input_maxmin (uses 3-input PTX max/min):
|
||||
# Tree reduction with 3-input instructions:
|
||||
# max(a, b, c) in one PTX instruction
|
||||
for i in T.serial(n // 3):
|
||||
T.cuda.func_call("max3_f32", &buf[i*3], &buf[i*3+1], &buf[i*3+2])
|
||||
|
||||
With accum=True: accumulator folded into first element/pair of the reduction.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import operator
|
||||
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import BufferRegion, PrimFunc
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext
|
||||
from tvm.tirx.operator.tile_primitive.common import ReduceOpType
|
||||
from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..common import sm_version_ok
|
||||
from ..exec_scope_utils import exec_scope_ok
|
||||
from .utils import (
|
||||
_dst_len_ok,
|
||||
_dtype_ok,
|
||||
_local_scope_match,
|
||||
_reduction_len_ok,
|
||||
_src_ndim_ok,
|
||||
reduce_op_table,
|
||||
)
|
||||
|
||||
|
||||
def _emit_reduction_local_thread_packed_add_sum(
|
||||
dst_buffer_region: BufferRegion,
|
||||
src_buffer_region: BufferRegion,
|
||||
accum: bool,
|
||||
reduce_op: ReduceOpType,
|
||||
sctx: DispatchContext,
|
||||
) -> PrimFunc:
|
||||
dst, src = dst_buffer_region.buffer, src_buffer_region.buffer
|
||||
src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
|
||||
dtype = src.dtype
|
||||
|
||||
src_extent = [r.extent for r in src_region]
|
||||
[r.extent for r in dst_region]
|
||||
src_st = [r.min for r in src_region]
|
||||
dst_st = [r.min for r in dst_region]
|
||||
|
||||
reduction_len = functools.reduce(operator.mul, src_extent, 1)
|
||||
|
||||
src_base = src_st[0]
|
||||
num_full_chunks = reduction_len // 8
|
||||
remainder = reduction_len % 8
|
||||
remainder_base = num_full_chunks * 8
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
local_sum = T.alloc_buffer([8], dtype, scope="local")
|
||||
# First pass: copy first 8 elements (with optional accumulator)
|
||||
for i in T.unroll(8):
|
||||
if accum and i == 0:
|
||||
local_sum[i] = src[src_base + i] + dst[tuple(dst_st)]
|
||||
else:
|
||||
local_sum[i] = src[src_base + i]
|
||||
|
||||
# Process remaining full chunks of 8
|
||||
for outer in T.serial(num_full_chunks - 1):
|
||||
for j in T.unroll(4):
|
||||
T.ptx.add_f32x2(
|
||||
T.address_of(local_sum[2 * j]),
|
||||
T.cuda.make_float2(local_sum[2 * j], local_sum[2 * j + 1]),
|
||||
T.cuda.make_float2(
|
||||
src[src_base + 8 * (outer + 1) + 2 * j],
|
||||
src[src_base + 8 * (outer + 1) + 2 * j + 1],
|
||||
),
|
||||
ftz=True,
|
||||
)
|
||||
|
||||
# Handle remainder elements (0 to 7)
|
||||
for i in T.serial(remainder):
|
||||
local_sum[0] = local_sum[0] + src[src_base + remainder_base + i]
|
||||
|
||||
# Final packed add sum: 8 -> 4 -> 2 -> 1
|
||||
T.ptx.add_f32x2(
|
||||
T.address_of(local_sum[0]),
|
||||
T.cuda.make_float2(local_sum[0], local_sum[1]),
|
||||
T.cuda.make_float2(local_sum[2], local_sum[3]),
|
||||
ftz=True,
|
||||
)
|
||||
T.ptx.add_f32x2(
|
||||
T.address_of(local_sum[4]),
|
||||
T.cuda.make_float2(local_sum[4], local_sum[5]),
|
||||
T.cuda.make_float2(local_sum[6], local_sum[7]),
|
||||
ftz=True,
|
||||
)
|
||||
T.ptx.add_f32x2(
|
||||
T.address_of(local_sum[0]),
|
||||
T.cuda.make_float2(local_sum[0], local_sum[1]),
|
||||
T.cuda.make_float2(local_sum[4], local_sum[5]),
|
||||
ftz=True,
|
||||
)
|
||||
dst[tuple(dst_st)] = local_sum[0] + local_sum[1]
|
||||
# fmt: on
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
def _emit_reduction_local_thread_3input_maxmin(
|
||||
dst_buffer_region: BufferRegion,
|
||||
src_buffer_region: BufferRegion,
|
||||
accum: bool,
|
||||
reduce_op: ReduceOpType,
|
||||
sctx: DispatchContext,
|
||||
) -> PrimFunc:
|
||||
dst, src = dst_buffer_region.buffer, src_buffer_region.buffer
|
||||
src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
|
||||
dtype = src.dtype
|
||||
|
||||
src_extent = [r.extent for r in src_region]
|
||||
src_st = [r.min for r in src_region]
|
||||
dst_st = [r.min for r in dst_region]
|
||||
|
||||
reduction_len = functools.reduce(operator.mul, src_extent, 1)
|
||||
|
||||
op_func = reduce_op_table[reduce_op]
|
||||
reduce3_func = T.ptx.reduce3_max_f32 if reduce_op == ReduceOpType.MAX else T.ptx.reduce3_min_f32
|
||||
|
||||
src_base = src_st[0]
|
||||
num_full_chunks = reduction_len // 8
|
||||
remainder = reduction_len % 8
|
||||
remainder_base = num_full_chunks * 8
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(check_well_formed=False)
|
||||
def impl():
|
||||
temp = T.alloc_buffer([4], dtype, scope="local")
|
||||
# First pass: process first 8 elements into 4 temps
|
||||
for i in T.unroll(4):
|
||||
if accum and i == 0:
|
||||
temp[i] = reduce3_func(src[src_base + 2 * i], src[src_base + 2 * i + 1], dst[tuple(dst_st)]) # noqa: E501
|
||||
else:
|
||||
temp[i] = op_func(src[src_base + 2 * i], src[src_base + 2 * i + 1])
|
||||
|
||||
# Process remaining full chunks of 8
|
||||
for outer in T.serial(num_full_chunks - 1):
|
||||
for i in T.unroll(4):
|
||||
temp[i] = reduce3_func(
|
||||
temp[i],
|
||||
src[src_base + 8 * (outer + 1) + 2 * i],
|
||||
src[src_base + 8 * (outer + 1) + 2 * i + 1],
|
||||
)
|
||||
|
||||
# Process remainder elements (0 to 7 elements)
|
||||
for i in T.serial(remainder):
|
||||
temp[0] = op_func(temp[0], src[src_base + remainder_base + i])
|
||||
|
||||
# Final merge: combine 4 temps into result
|
||||
dst[tuple(dst_st)] = op_func(temp[0], temp[1])
|
||||
dst[tuple(dst_st)] = reduce3_func(dst[tuple(dst_st)], temp[2], temp[3])
|
||||
# fmt: on
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
def _sm100_packed_add_sum_impl(op: TilePrimitiveCall, op_type: ReduceOpType, sctx: DispatchContext):
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
return _emit_reduction_local_thread_packed_add_sum(op.output, op.input, op.accum, op_type, sctx)
|
||||
|
||||
|
||||
def _sm100_3input_maxmin_impl(op: TilePrimitiveCall, op_type: ReduceOpType, sctx: DispatchContext):
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
return _emit_reduction_local_thread_3input_maxmin(op.output, op.input, op.accum, op_type, sctx)
|
||||
|
||||
|
||||
_optimized_local_reduction_predicates = [
|
||||
predicate("exec_scope", exec_scope_ok, expected_scopes=["thread"]),
|
||||
predicate("local_scope", _local_scope_match),
|
||||
predicate("dst_len", _dst_len_ok, expected_len=1),
|
||||
predicate("src_ndim", _src_ndim_ok, expected_ndim=1),
|
||||
predicate("dtype", _dtype_ok, expected_dtype="float32"),
|
||||
predicate("sm_version", sm_version_ok, min_version=100),
|
||||
predicate("reduction_len", _reduction_len_ok, min_len=8),
|
||||
]
|
||||
|
||||
_optimized_impl_table = {
|
||||
ReduceOpType.SUM: ("packed_add_sum", _sm100_packed_add_sum_impl),
|
||||
ReduceOpType.MAX: ("3input_maxmin", _sm100_3input_maxmin_impl),
|
||||
ReduceOpType.MIN: ("3input_maxmin", _sm100_3input_maxmin_impl),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registration: SM100+ optimized local reduction (priority=20)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
for op_name, op_type in [
|
||||
("sum", ReduceOpType.SUM),
|
||||
("max", ReduceOpType.MAX),
|
||||
("min", ReduceOpType.MIN),
|
||||
]:
|
||||
variant_name, optimized_impl = _optimized_impl_table[op_type]
|
||||
|
||||
@register_dispatch(
|
||||
op_name,
|
||||
"cuda",
|
||||
variant=variant_name,
|
||||
priority=20,
|
||||
when=_optimized_local_reduction_predicates,
|
||||
)
|
||||
def _optimized_dispatch(
|
||||
op: TilePrimitiveCall, sctx: DispatchContext, _impl=optimized_impl, _op_type=op_type
|
||||
) -> PrimFunc:
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
return _impl(op, _op_type, sctx)
|
||||
@@ -0,0 +1,262 @@
|
||||
# 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 helpers for reduction operator dispatches on CUDA targets."""
|
||||
|
||||
import functools
|
||||
import math
|
||||
import operator
|
||||
|
||||
from tvm.arith.analyzer import Analyzer
|
||||
from tvm.script import tirx as T
|
||||
from tvm.tirx import BufferRegion
|
||||
from tvm.tirx.operator.tile_primitive import DispatchContext
|
||||
from tvm.tirx.operator.tile_primitive.common import ReduceOpType
|
||||
from tvm.tirx.stmt import TilePrimitiveCall
|
||||
|
||||
from ..common import match_scope
|
||||
|
||||
reduce_op_table = {
|
||||
ReduceOpType.SUM: lambda a, b: a + b,
|
||||
ReduceOpType.MAX: T.max,
|
||||
ReduceOpType.MIN: T.min,
|
||||
}
|
||||
|
||||
|
||||
def reduce_default_value_table(dtype):
|
||||
return {
|
||||
ReduceOpType.SUM: 0.0,
|
||||
ReduceOpType.MAX: T.min_value(dtype),
|
||||
ReduceOpType.MIN: T.max_value(dtype),
|
||||
}
|
||||
|
||||
|
||||
def _reduction_args(
|
||||
op: TilePrimitiveCall,
|
||||
) -> tuple[BufferRegion, BufferRegion, tuple[int, ...], bool, dict]:
|
||||
"""Parse ReduceOp -> (dst, src, reduce_axes, accum, config)."""
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
dst = op.output
|
||||
src = op.input
|
||||
reduce_axes = tuple(int(a) for a in op.reduce_axes)
|
||||
accum = op.accum
|
||||
config = op.config
|
||||
return dst, src, reduce_axes, accum, config
|
||||
|
||||
|
||||
def _match_reduction_storage_scope(
|
||||
op: TilePrimitiveCall, sctx: DispatchContext, expected_scope: list[str]
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Check that dst and src scopes match one of the expected patterns."""
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
dst_scope = op.output.buffer.scope()
|
||||
src_scope = op.input.buffer.scope()
|
||||
|
||||
ok = any(match_scope(dst_scope, p) and match_scope(src_scope, p) for p in expected_scope)
|
||||
msg = f"storage scope mismatch: dst {dst_scope}, src {src_scope}; expected {expected_scope}"
|
||||
return (ok, None if ok else msg)
|
||||
|
||||
|
||||
def _analyze_axes(src_ndim: int, reduce_axes: tuple[int, ...]) -> tuple[list[int], list[int]]:
|
||||
"""Normalize negative axes -> (reduce_dim_set, spatial_dim_list)."""
|
||||
reduce_dims = set()
|
||||
for ax in reduce_axes:
|
||||
a = ax if ax >= 0 else ax + src_ndim
|
||||
assert 0 <= a < src_ndim, f"reduce axis {ax} out of range for ndim={src_ndim}"
|
||||
reduce_dims.add(a)
|
||||
spatial_dims = [d for d in range(src_ndim) if d not in reduce_dims]
|
||||
return sorted(reduce_dims), spatial_dims
|
||||
|
||||
|
||||
def _analyze_layout_dims(layout, shape):
|
||||
"""layout.group(shape) -> decompose each dim into thread/local iters.
|
||||
|
||||
Returns list of per-dim (thread_extent, local_extent, thread_strides):
|
||||
thread_extent = product of thread iter extents in this dim
|
||||
local_extent = product of local iter extents in this dim
|
||||
thread_strides = list of (stride, extent) for thread iters in this dim
|
||||
"""
|
||||
grouped, seps = layout.group(list(shape))
|
||||
result = []
|
||||
for d in range(len(shape)):
|
||||
shard_range = list(range(seps[d], seps[d + 1]))
|
||||
thread_extent = 1
|
||||
local_extent = 1
|
||||
thread_strides = []
|
||||
for s_idx in shard_range:
|
||||
it = grouped.shard[s_idx]
|
||||
if it.axis.is_thread():
|
||||
thread_extent *= it.extent
|
||||
thread_strides.append((it.stride, it.extent))
|
||||
else:
|
||||
local_extent *= it.extent
|
||||
result.append((thread_extent, local_extent, thread_strides))
|
||||
return result
|
||||
|
||||
|
||||
def _compute_shuffle_masks(dim_info, reduce_dims: set[int]) -> list[int]:
|
||||
"""From reduction dims' thread iter (stride, extent) pairs, compute XOR masks.
|
||||
|
||||
For each thread iter in a reduction dim:
|
||||
masks += [stride * 2^i for i in range(log2(extent))]
|
||||
Sorted ascending.
|
||||
"""
|
||||
masks = []
|
||||
for d in reduce_dims:
|
||||
_, _, thread_strides = dim_info[d]
|
||||
for stride, extent in thread_strides:
|
||||
ext_int = int(extent) if hasattr(extent, "__int__") else extent
|
||||
n_bits = int(math.log2(ext_int))
|
||||
for i in range(n_bits):
|
||||
stride_int = int(stride) if hasattr(stride, "__int__") else stride
|
||||
masks.append(stride_int * (1 << i))
|
||||
masks.sort()
|
||||
return masks
|
||||
|
||||
|
||||
def _build_local_dim_map(layout, buffer_shape):
|
||||
"""Map original dim index to position in get_local_region output (None if pure-thread)."""
|
||||
grouped, seps = layout.group(list(buffer_shape))
|
||||
dim_map = {}
|
||||
local_pos = 0
|
||||
for d in range(len(buffer_shape)):
|
||||
shard_range = list(range(seps[d], seps[d + 1]))
|
||||
has_local = any(not grouped.shard[s].axis.is_thread() for s in shard_range)
|
||||
if has_local:
|
||||
dim_map[d] = local_pos
|
||||
local_pos += 1
|
||||
else:
|
||||
dim_map[d] = None
|
||||
return dim_map
|
||||
|
||||
|
||||
def _validate_reduction_layout(
|
||||
src_layout, dst_layout, src_shape, dst_shape, reduce_dims: list[int]
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Validate that spatial dims of src/dst have matching thread+local structure,
|
||||
and that reduction dims in dst have local_extent == 1.
|
||||
"""
|
||||
src_dim_info = _analyze_layout_dims(src_layout, src_shape)
|
||||
dst_dim_info = _analyze_layout_dims(dst_layout, dst_shape)
|
||||
analyzer = Analyzer()
|
||||
|
||||
# Spatial dims: src/dst must match in both thread and local extents.
|
||||
# Reduce dims: src/dst thread extent must match, and dst local extent must be 1.
|
||||
|
||||
# get expected simplified dst layout
|
||||
expected_dst_dim = []
|
||||
for src_idx in range(len(src_shape)):
|
||||
if analyzer.can_prove_equal(src_dim_info[src_idx][0], 1) and analyzer.can_prove_equal(
|
||||
src_dim_info[src_idx][1], 1
|
||||
):
|
||||
continue # skip if extent=1
|
||||
if src_idx in reduce_dims: # reduce dims
|
||||
if not analyzer.can_prove_equal(src_dim_info[src_idx][0], 1):
|
||||
expected_dst_dim.append((src_dim_info[src_idx][0], 1))
|
||||
else: # spatial dims
|
||||
expected_dst_dim.append((src_dim_info[src_idx][0], src_dim_info[src_idx][1]))
|
||||
|
||||
# check dst layout
|
||||
check_idx = 0
|
||||
for dst_idx in range(len(dst_shape)):
|
||||
if analyzer.can_prove_equal(dst_dim_info[dst_idx][0], 1) and analyzer.can_prove_equal(
|
||||
dst_dim_info[dst_idx][1], 1
|
||||
):
|
||||
continue
|
||||
if not (
|
||||
analyzer.can_prove_equal(dst_dim_info[dst_idx][0], expected_dst_dim[check_idx][0])
|
||||
and analyzer.can_prove_equal(dst_dim_info[dst_idx][1], expected_dst_dim[check_idx][1])
|
||||
):
|
||||
return False, "mismatch dst/src layout for reduction"
|
||||
check_idx += 1
|
||||
if check_idx != len(expected_dst_dim):
|
||||
return False, "mismatch dst/src layout for reduction"
|
||||
return True, None
|
||||
|
||||
|
||||
def build_src_indices(spa_fused, red_fused, spatial_dims, reduce_dims, src_extent, src_st):
|
||||
"""Combine spatial and reduction indices into full src index tuple."""
|
||||
|
||||
# Build index helpers that work with the explicit axis split
|
||||
def get_spatial_or_reduction_src_indices(spa_or_red_fused, is_spatial):
|
||||
dims = spatial_dims if is_spatial else reduce_dims
|
||||
spa_extents = [src_extent[d] for d in dims]
|
||||
indices = []
|
||||
rem = spa_or_red_fused
|
||||
for e in reversed(spa_extents):
|
||||
indices.append(rem % e)
|
||||
rem //= e
|
||||
indices.reverse()
|
||||
return [idx + src_st[d] for idx, d in zip(indices, dims)]
|
||||
|
||||
spa_vals = get_spatial_or_reduction_src_indices(spa_fused, is_spatial=True)
|
||||
red_vals = get_spatial_or_reduction_src_indices(red_fused, is_spatial=False)
|
||||
full = [None] * len(src_extent)
|
||||
for i, d in enumerate(spatial_dims):
|
||||
full[d] = spa_vals[i]
|
||||
for i, d in enumerate(reduce_dims):
|
||||
full[d] = red_vals[i]
|
||||
return full
|
||||
|
||||
|
||||
_REDUCE_OP_TO_STR = {ReduceOpType.SUM: "sum", ReduceOpType.MAX: "max", ReduceOpType.MIN: "min"}
|
||||
|
||||
|
||||
def _dtype_ok(op: TilePrimitiveCall, sctx: DispatchContext, expected_dtype: str):
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
dtype = op.input.buffer.dtype
|
||||
ok = dtype == expected_dtype
|
||||
return (ok, None if ok else f"dtype {dtype} != {expected_dtype}")
|
||||
|
||||
|
||||
def _reduction_len_ok(op: TilePrimitiveCall, sctx: DispatchContext, min_len: int):
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
src_extent = [r.extent for r in op.input.region]
|
||||
reduction_len = functools.reduce(operator.mul, src_extent, 1)
|
||||
ok = reduction_len >= min_len
|
||||
return (ok, None if ok else f"reduction_len {reduction_len} < {min_len}")
|
||||
|
||||
|
||||
def _dst_len_ok(op: TilePrimitiveCall, sctx: DispatchContext, expected_len: int):
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
dst_extent = [r.extent for r in op.output.region]
|
||||
dst_len = functools.reduce(operator.mul, dst_extent, 1)
|
||||
ok = dst_len == expected_len
|
||||
return (ok, None if ok else f"dst_len {dst_len} != {expected_len}")
|
||||
|
||||
|
||||
def _src_ndim_ok(op: TilePrimitiveCall, sctx: DispatchContext, expected_ndim: int):
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
src_extent = [r.extent for r in op.input.region]
|
||||
ok = len(src_extent) == expected_ndim
|
||||
return (ok, None if ok else f"src ndim {len(src_extent)} != {expected_ndim}")
|
||||
|
||||
|
||||
def _local_scope_match(op: TilePrimitiveCall, sctx: DispatchContext):
|
||||
op = TilePrimitiveCall.downcast(op)
|
||||
src, dst = op.input.buffer, op.output.buffer
|
||||
ok = all(
|
||||
[
|
||||
src.scope() == "local",
|
||||
dst.scope() == "local",
|
||||
src.dtype == dst.dtype,
|
||||
sctx.is_target("cuda"),
|
||||
]
|
||||
)
|
||||
if not ok:
|
||||
return (False, "src/dst must be local scope with matching dtype on CUDA")
|
||||
return (True, None)
|
||||
@@ -0,0 +1,117 @@
|
||||
# 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.
|
||||
|
||||
"""TMA (Tensor Memory Accelerator) utilities for CUDA op dispatches."""
|
||||
|
||||
import copy
|
||||
from enum import Enum
|
||||
|
||||
import tvm
|
||||
from tvm.arith.analyzer import Analyzer
|
||||
from tvm.tirx.layout import ComposeLayout, Layout, S, SwizzleLayout, TileLayout
|
||||
|
||||
|
||||
class SwizzleMode(Enum):
|
||||
"""The swizzle mode of the TMA."""
|
||||
|
||||
SWIZZLE_NONE = 0
|
||||
SWIZZLE_32B_ATOM = 1
|
||||
SWIZZLE_64B_ATOM = 2
|
||||
SWIZZLE_128B_ATOM = 3
|
||||
|
||||
|
||||
def mma_atom_layout(dtype: str, swizzle_mode: SwizzleMode | int) -> SwizzleLayout:
|
||||
"""Generate the MMA-compatible shared-memory atom layout."""
|
||||
bits = tvm.DataType(dtype).bits
|
||||
if isinstance(swizzle_mode, int):
|
||||
swizzle_mode = SwizzleMode(swizzle_mode)
|
||||
return SwizzleLayout(
|
||||
per_element=(128 // bits).bit_length() - 1, swizzle_len=swizzle_mode.value, atom_len=3
|
||||
)
|
||||
|
||||
|
||||
def mma_atom_shape(dtype: str, swizzle_mode: SwizzleMode | int, shape: list[int] | None = None):
|
||||
"""Generate the MMA-compatible shared-memory atom shape."""
|
||||
bits = tvm.DataType(dtype).bits
|
||||
if isinstance(swizzle_mode, int):
|
||||
swizzle_mode = SwizzleMode(swizzle_mode)
|
||||
atom_shape = {
|
||||
SwizzleMode.SWIZZLE_32B_ATOM: [8, 256],
|
||||
SwizzleMode.SWIZZLE_64B_ATOM: [8, 512],
|
||||
SwizzleMode.SWIZZLE_128B_ATOM: [8, 1024],
|
||||
}[swizzle_mode]
|
||||
atom_shape[-1] //= bits
|
||||
if shape is None:
|
||||
return atom_shape
|
||||
atom_shape = [1] * (len(shape) - len(atom_shape)) + atom_shape
|
||||
return atom_shape
|
||||
|
||||
|
||||
def mma_shared_layout(dtype: str, swizzle_mode: SwizzleMode | int, shape) -> Layout:
|
||||
"""Generate the MMA-compatible shared-memory layout for shape and dtype.
|
||||
|
||||
It uses a default tiling strategy to tile the TMA atom layout into the shared memory.
|
||||
"""
|
||||
if isinstance(swizzle_mode, int):
|
||||
swizzle_mode = SwizzleMode(swizzle_mode)
|
||||
if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
|
||||
return TileLayout(S[tuple(shape)]).canonicalize()
|
||||
atom_shape = mma_atom_shape(dtype, swizzle_mode, shape)
|
||||
layout = mma_atom_layout(dtype, swizzle_mode)
|
||||
tile_to_shape = copy.copy(atom_shape)
|
||||
tile_to_shape[-2] = shape[-2]
|
||||
return layout.tile_to(tile_to_shape, atom_shape).tile_to(shape, tile_to_shape).canonicalize()
|
||||
|
||||
|
||||
# Backward-compatible aliases kept during the alloc_mma migration.
|
||||
tma_atom_layout = mma_atom_layout
|
||||
tma_atom_shape = mma_atom_shape
|
||||
tma_shared_layout = mma_shared_layout
|
||||
|
||||
|
||||
def tma_atom_compatible(dst_shape, dst_st, dst_extent, atom_shape):
|
||||
"""Check if the copy region in dst is compatible with the TMA atom shape."""
|
||||
analyzer = Analyzer()
|
||||
for i, _ in enumerate(dst_st):
|
||||
if any(
|
||||
not analyzer.can_prove_equal(x % atom_shape[i], 0)
|
||||
for x in [dst_shape[i], dst_st[i], dst_extent[i]]
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_swizzle_mode_from_layout(layout: Layout) -> SwizzleMode | None:
|
||||
"""Extract swizzle mode from a shared memory layout."""
|
||||
if isinstance(layout, ComposeLayout):
|
||||
swizzle = layout.swizzle # SwizzleLayout is named 'swizzle' in ComposeLayout
|
||||
swizzle_len = swizzle.swizzle_len
|
||||
elif isinstance(layout, SwizzleLayout):
|
||||
swizzle_len = layout.swizzle_len
|
||||
elif isinstance(layout, TileLayout):
|
||||
# TileLayout without SwizzleLayout means no swizzle (mode 0)
|
||||
return SwizzleMode.SWIZZLE_NONE
|
||||
else:
|
||||
return None
|
||||
|
||||
# Map swizzle_len to SwizzleMode
|
||||
return {
|
||||
0: SwizzleMode.SWIZZLE_NONE,
|
||||
1: SwizzleMode.SWIZZLE_32B_ATOM,
|
||||
2: SwizzleMode.SWIZZLE_64B_ATOM,
|
||||
3: SwizzleMode.SWIZZLE_128B_ATOM,
|
||||
}.get(swizzle_len)
|
||||
Reference in New Issue
Block a user