chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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}")
|
||||
Reference in New Issue
Block a user