wehub-resource-sync
284 lines
11 KiB
Python
284 lines
11 KiB
Python
# 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}")
|