diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py index a65c93921dd9..41557824ec0e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py @@ -108,6 +108,10 @@ # conv split-K resolver reads no environment; this benchmark maps the variable # onto its explicit ``override`` argument. _ENV_CONV_SPLIT_K = "ROCKE_CONV_SPLIT_K" +# A forward split only pays while the unsplit grid underfills the device: a +# degree whose grid would exceed this many workgroups is not offered (it also +# keeps large layers from timing hundreds of split candidates that cannot win). +_FWD_SPLIT_K_MAX_CTAS = 16384 def _conv_split_k_override() -> int | None: @@ -1246,7 +1250,8 @@ def main() -> int: dest="split_k", metavar="N", help=( - "wgrad split-K degree: " + "wgrad split-K degree (and the forward one, for the cached " + "split-K forward binaries): " "0 = sweep the degrees of --split-k-degrees plus the one the " "split-K resolver picks for each kernel, " "1 = disabled, " @@ -4786,7 +4791,7 @@ def _u8(t): # A split-K wgrad binary is one kernel for every degree > 1 # (the degree is a kernarg), so --split-k decides whether it # is timed at one degree or at each degree it can use. - if direction == "wgrad" and ident.split_k > 1: + if direction in ("wgrad", "fwd") and ident.split_k > 1: degrees = _runtime_split_ks( problem, ident, args.split_k, args.split_k_degrees ) @@ -4795,10 +4800,20 @@ def _u8(t): f" [skip] {ident.label()}: " + ( "no split-K degree > 1 fits this problem" + + ( + " (K tiles, i32 workspace, workgroup cap)" + if direction == "fwd" + else "" + ) if args.split_k == 0 - else f"--split-k {args.split_k} does not fit " - f"this problem (degree 1 is a separate " - f"non-split-K binary)" + else ( + "the forward split-K policy does not " + "split this problem (the plain binary serves it)" + if direction == "fwd" and args.split_k == -1 + else f"--split-k {args.split_k} does not fit " + f"this problem (degree 1 is a separate " + f"non-split-K binary)" + ) ), flush=True, ) @@ -4971,9 +4986,16 @@ def _launch_problem(direction, problem): (:func:`kernels.common.conv_implicit_gemm_dgrad.fold_width`) -- the same bytes in all three tensors, in fewer and wider sub-GEMMs -- exactly as ``dispatch.grouped_convolution`` launches it, so what is benchmarked here - is what dispatch runs. Every other problem launches as it is. + is what dispatch runs. + + A forward folds the same way: the identity behind ``fold_width`` does not + depend on the direction, so a patch embedding (stride = filter, no + padding) or a 2x2 stride-2 downsample runs as a ``Y x 1`` convolution + over ``X * C`` contiguous channels -- vector-wide A loads instead of + ``C``-element gathers -- on exactly the case's tensors. Every other + problem launches as it is. """ - if direction != "dgrad": + if direction not in ("dgrad", "fwd"): return problem from kernels.common.conv_implicit_gemm_dgrad import fold_width @@ -5103,7 +5125,7 @@ def _signature_for_identity(ident, dtype): ) if ident.direction == "dgrad": return conv_args_signature(dtype, direction="dgrad", is_3d=ident.is_3d) - return conv_args_signature(dtype, is_3d=ident.is_3d) + return conv_args_signature(dtype, is_3d=ident.is_3d, two_stage=ident.two_stage) def _ws_replicas_of(ident) -> int: @@ -5349,6 +5371,7 @@ def _split_k_degrees( degrees: tuple = _SPLIT_K_AUTO, ws_replicas: int = 0, include_resolved: bool = True, + direction: str = "wgrad", ) -> tuple: """Split-K degrees to launch a wgrad split-K kernel with, largest first. @@ -5365,21 +5388,29 @@ def _split_k_degrees( are a property of the problem, not of how the degree was chosen -- so a pinned degree the shape cannot take drops out here rather than failing at launch. + + ``direction="fwd"`` serves the split-K forward binary: the resolver's + degree is the forward policy's + (:func:`rocke.helpers.split_k.select_split_k_fwd`), added only where it + splits -- where it does not, the plain forward binary is the one to time -- + and a degree is also dropped when its grid exceeds + ``_FWD_SPLIT_K_MAX_CTAS`` workgroups or its fp32 workspace passes the i32 + byte range. """ from kernels.common.conv_args import ConvArgs from kernels.common.conv_split_k import resolve_conv_split_k - wg_k = ConvArgs.from_problem( + args_ = ConvArgs.from_problem( problem, - direction="wgrad", + direction=direction, tile_m=tile_m, tile_n=tile_n, tile_k=tile_k, - ).gemm_k - k_tiles = -(-wg_k // tile_k) + ) + k_tiles = -(-args_.gemm_k // tile_k) grid_groups = max(1, problem.groups // max(1, group_merge)) candidates = set(degrees) - if include_resolved: + if include_resolved and direction == "wgrad": try: candidates.add( resolve_conv_split_k( @@ -5400,35 +5431,132 @@ def _split_k_degrees( ) except ValueError: pass # this binary's reduction cannot serve the problem; nothing to add + elif include_resolved and direction == "fwd": + try: + candidates.add( + resolve_conv_split_k( + "fwd", + problem, + arch=arch, + dtype=dtype, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + block_size=block_size, + override=_conv_split_k_override(), + allow_unsplit=False, + ).split_k + ) + except ValueError: + pass # the forward split does not pay here; nothing to add + + def _fits(sk: int) -> bool: + if direction != "fwd": + return True + if args_.num_pid_m * args_.num_pid_n * grid_groups * sk > _FWD_SPLIT_K_MAX_CTAS: + return False + # The forward workspace is sk fp32 slabs of the output; it is + # addressed through an i32 byte count like every operand. + from kernels.common.conv_fwd_split_k_reduce import ( + fwd_split_k_workspace_nbytes, + ) + + return fwd_split_k_workspace_nbytes(args_.gemm_m, problem.K, sk) <= _INT32_MAX + return tuple( sk for sk in sorted(candidates, reverse=True) - if 1 < sk <= k_tiles and grid_groups * sk <= _MAX_GRID_DIM_Z + if 1 < sk <= k_tiles and grid_groups * sk <= _MAX_GRID_DIM_Z and _fits(sk) + ) + + +# Cluster split-K forward gate (see _runtime_split_ks): fewer than this many +# output tiles (groups included), and at least this deep a reduction. +_FWD_SPLIT_K_CLUSTER_MAX_TILES = 256 +_FWD_SPLIT_K_CLUSTER_MIN_GEMM_K = 1500 + + +def _fwd_cluster_split_pays(problem, ident) -> bool: + """Is ``problem`` a small-output / deep-reduction layer the cluster + split-K forward binary ``ident`` is for?""" + from kernels.common.conv_args import ConvArgs + from kernels.common.conv_implicit_gemm import fwd_split_k_cluster_problem_reason + + args_ = ConvArgs.from_problem( + problem, tile_m=ident.tile_m, tile_n=ident.tile_n, tile_k=ident.tile_k + ) + tiles = args_.num_pid_m * args_.num_pid_n * max(1, problem.groups) + return ( + # The fold walks 4-channel quads of the packed [M][K] output. + fwd_split_k_cluster_problem_reason(problem) is None + and tiles < _FWD_SPLIT_K_CLUSTER_MAX_TILES + and args_.gemm_k >= _FWD_SPLIT_K_CLUSTER_MIN_GEMM_K ) def _runtime_split_ks( problem, ident, requested: int = 0, degrees: tuple = _SPLIT_K_AUTO ) -> tuple: - """Degrees to launch a cached split-K wgrad binary at. + """Degrees to launch a cached split-K wgrad or forward binary at. ``requested`` is ``--split-k``, with the same meaning it has on the JIT sweep: 0 sweeps ``degrees`` (``--split-k-degrees``) plus the resolver's pick, -1 launches at the resolver's pick only, and a positive value pins - that degree. The cache records ``split_k=2`` only as a marker that the - binary is the split-K one -- the degree itself is a kernarg -- so nothing - here is baked in and every mode is reachable at launch. + that degree -- for the forward split as for wgrad. Where the forward + policy does not split the problem, -1 yields no degree, so the plain + forward binary alone is timed. The cache records ``split_k=2`` only as a + marker that the binary is the split-K one -- the degree itself is a + kernarg -- so nothing here is baked in and every mode is reachable at + launch. Degree 1 is not among them: that is the separate non-atomic kernel, a different binary, so ``--split-k 1`` yields no degree and the caller skips the entry rather than silently timing it at some other degree. """ + if ident.direction == "fwd" and ident.split_k_cluster: + # The cluster binary bakes its degree. It is only offered where the + # single launch pays: few output tiles (the grid alone cannot fill + # the device) and a deep reduction; elsewhere -- e.g. many small + # grouped tiles -- the cluster barrier costs more than it saves. + if requested not in (0, -1, ident.split_k_cluster): + return () + if not _fwd_cluster_split_pays(problem, ident): + return () + return _split_k_degrees( + problem, + arch=ident.arch, + dtype=ident.dtype_d, + tile_m=ident.tile_m, + tile_n=ident.tile_n, + tile_k=ident.tile_k, + block_size=ident.warp_m * ident.warp_n * ident.wave_size, + group_merge=1, + two_stage=True, + degrees=(ident.split_k_cluster,), + include_resolved=False, + direction="fwd", + ) if requested == 0: swept, include_resolved = degrees, True elif requested == -1: swept, include_resolved = (), True else: swept, include_resolved = (requested,), False + if ident.direction == "fwd": + return _split_k_degrees( + problem, + arch=ident.arch, + dtype=ident.dtype_d, + tile_m=ident.tile_m, + tile_n=ident.tile_n, + tile_k=ident.tile_k, + block_size=ident.warp_m * ident.warp_n * ident.wave_size, + group_merge=1, + two_stage=True, + degrees=swept, + include_resolved=include_resolved, + direction="fwd", + ) return _split_k_degrees( problem, arch=ident.arch, @@ -5563,6 +5691,20 @@ def _cached_bind( return None live_extras.append(extras) values = _launch_values_for(direction, problem, ident, ptrs, sizes, extras) + cluster = None + if direction == "fwd" and ident.split_k_cluster: + # One (1, 1, S) cluster per output tile: the K slices of a tile are + # consecutive in z, so z must hold whole clusters and the launch + # degree must be the baked one. + S = ident.split_k_cluster + if split_k != S or grid[2] % S or any(g < 1 for g in grid): + print( + f" [skip] {ident.label()}: cluster split-K launched at degree " + f"{split_k} with grid {grid}", + flush=True, + ) + return None + cluster = (1, 1, S) cfg = LaunchConfig( grid=grid, block=( @@ -5571,6 +5713,7 @@ def _cached_bind( 1, ), stream=0, + cluster=cluster, ) run = launcher.bind(values, config=cfg) if ( @@ -5589,7 +5732,7 @@ def run(_launch=_launch, dw_ptr=ptrs[2], dw_bytes=sizes[2]): rt.memset(dw_ptr, 0, dw_bytes) _launch() - elif direction == "wgrad" and ident.two_stage: + elif direction == "wgrad" and ident.needs_stage2: # Two-stage is three steps and all of them are the algorithm's cost: # zero the scratch (Stage 1 accumulates), Stage 1, then the fold/cast # into dW. @@ -5621,12 +5764,87 @@ def run( launcher(values, config=cfg) s2_launcher(s2_values, config=s2_cfg) + if direction == "fwd" and ident.needs_stage2: + # Split-K forward: Stage 1 stores every slab element (no clearing), + # then Stage 2 folds the slabs and casts into D. Both are its cost. + # The cluster form folds in-kernel and has no Stage 2. + s2_launcher, s2_values, s2_cfg = _fwd_stage2( + problem, + case_dtype, + arch, + stage2_kernels, + KernelLauncher, + extras, + d_ptr=ptrs[2], + d_bytes=sizes[2], + ) + + # Both stages pre-packed (bind), like the single-kernel path: a timed + # loop of launcher(values, ...) repacks the arguments on every call, + # which costs more host time than these small kernels run, so the + # loop timed the host instead of the GPU. + stage1 = run + stage2 = s2_launcher.bind(s2_values, config=s2_cfg) + + def run(stage1=stage1, stage2=stage2): + stage1() + stage2() + label = ident.label() if split_k is not None: label += f" @split_k={split_k}" return run, label +def _fwd_stage2( + problem, dtype_d, arch, compiled, KernelLauncher, extras, *, d_ptr, d_bytes +): + """``(launcher, values, config)`` of the split-K forward Stage 2 fold. + + The fold is shape-generic (the element count is a kernarg) and specialised + for the degree, which unrolls its slab loop so every load issues at once -- + a runtime loop over the slabs costs ~2.5x on the small layers split-K + serves. It is compiled once per (dtype, vector width, degree) in a run, + like wgrad's Stage 2. + """ + from rocke import compile_kernel + from rocke.runtime.launcher import LaunchConfig + from kernels.common.conv_fwd_split_k_reduce import ( + FwdSplitKReduceSpec, + build_conv_fwd_split_k_reduce, + fwd_split_k_reduce_grid, + fwd_split_k_reduce_signature, + fwd_split_k_reduce_vec, + ) + + elems = extras["elems"] + spec = FwdSplitKReduceSpec( + dtype_d=dtype_d, vec=fwd_split_k_reduce_vec(elems), degree=extras["split_k"] + ) + key = ("fwd_split_k", dtype_d, spec.vec, spec.degree) + if key not in compiled: + art = compile_kernel(build_conv_fwd_split_k_reduce(spec, arch=arch), arch=arch) + compiled[key] = KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=fwd_split_k_reduce_signature(spec), + ) + values = { + "ws_ptr": extras["ws_ptr"], + "d_ptr": d_ptr, + "elems": elems, + "ws_bytes": extras["ws_bytes"], + "d_bytes": d_bytes, + "ks_count": extras["split_k"], + } + cfg = LaunchConfig( + grid=fwd_split_k_reduce_grid(spec, elems), + block=(spec.block_size, 1, 1), + stream=0, + ) + return compiled[key], values, cfg + + def _extras_and_grid(direction, problem, ident, rt, dtype, split_k=None): """Trailing kernargs and the launch grid for one cached kernel. @@ -5694,6 +5912,31 @@ def _extras_and_grid(direction, problem, ident, rt, dtype, split_k=None): extras["ws_bytes"] = ws_bytes return extras, _grid_for(direction, problem, ident, split_k) + if ident.two_stage: + # Split-K forward: one fp32 slab of the output per K slice. Every slab + # element is stored by Stage 1, so it is not cleared. + from kernels.common.conv_fwd_split_k_reduce import ( + fwd_split_k_workspace_nbytes, + ) + + if split_k is None or split_k < 1: + raise ValueError(f"split-K forward launched without a degree ({split_k})") + elems = ( + ConvArgs.from_problem( + problem, tile_m=ident.tile_m, tile_n=ident.tile_n + ).gemm_m + * problem.K + ) + ws_bytes = fwd_split_k_workspace_nbytes(elems, 1, split_k) + if ws_bytes > _INT32_MAX: + raise ValueError(f"split-K workspace of {ws_bytes} bytes exceeds i32") + extras["split_k"] = split_k + extras["elems"] = elems + extras["tile_k"] = ident.tile_k + extras["ws_ptr"] = rt.alloc(ws_bytes) + extras["ws_bytes"] = ws_bytes + return extras, _grid_for(direction, problem, ident, split_k) + return extras, _grid_for(direction, problem, ident) @@ -5715,6 +5958,9 @@ def _grid_for(direction, problem, ident, split_k=None): problem, direction="dgrad", tile_m=tm, tile_n=tn ).grid() gx, gy, gz = ConvArgs.from_problem(problem, tile_m=tm, tile_n=tn).grid() + if getattr(ident, "two_stage", False): + # Split-K forward: the group and the K slice share z. + return gx, gy, gz * (split_k or 1) # A group-merged fwd kernel runs one workgroup per Gm groups. Its merged # GEMM-N (Gm) fits one tile, as the true one (kpg = 1) does, so gx agrees. return gx, gy, gz // max(1, ident.group_merge) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py index 92f8379f4ab6..052cf19c75ae 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py @@ -250,6 +250,7 @@ def enumerate_jobs( split_ks=split_ks, streamk_modes=streamk_modes, streamk_reductions=streamk_reductions, + llvm_flavor=flavor, ) log(f" {total} candidates to check with {jobs} process(es)") @@ -540,7 +541,15 @@ def compile_all( ("grouped", lambda i: "yes" if i.grouped else "no"), ( "split-K", - lambda i: ("two-stage" if i.two_stage else "atomic") if i.split_k > 1 else "no", + lambda i: ( + ( + "cluster" + if i.split_k_cluster + else "two-stage" if i.needs_stage2 else "atomic" + ) + if i.split_k > 1 + else "no" + ), ), ("group_merge", lambda i: i.group_merge), ( @@ -844,6 +853,18 @@ def _launch_values_for(direction, problem, identity, ptrs, sizes, extras): tm, tn, tk = identity.tile_m, identity.tile_n, identity.tile_k if direction == "fwd": + if "ws_ptr" in extras: + # Split-K forward: the degree and the workspace, with the slice + # width rounded to the binary's tile_k. + return ConvArgs.from_problem( + problem, tile_m=tm, tile_n=tn, tile_k=tk + ).to_launch_values( + *ptrs, + *sizes, + split_k=extras["split_k"], + ws_ptr=extras["ws_ptr"], + ws_bytes=extras["ws_bytes"], + ) return ConvArgs.from_problem(problem, tile_m=tm, tile_n=tn).to_launch_values( *ptrs, *sizes ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_abi.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_abi.py index 8b2a0cef4ed8..c200785ca959 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_abi.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_abi.py @@ -109,8 +109,13 @@ def _grid() -> List[ArgSpec]: return [("p_num_pid_m", "i32"), ("p_num_pid_n", "i32")] -def _fwd_arg_names(*, is_3d: bool = False) -> List[ArgSpec]: - """Forward implicit-GEMM conv: ``A`` (N[D]HWC), ``B`` (K[Z]YXC), ``D`` (N[D]HWK).""" +def _fwd_arg_names(*, is_3d: bool = False, split_k: bool = False) -> List[ArgSpec]: + """Forward implicit-GEMM conv: ``A`` (N[D]HWC), ``B`` (K[Z]YXC), ``D`` (N[D]HWK). + + ``split_k`` is the split-K binary's ABI: the fp32 workspace its slices + store into, then the slice width and degree (wgrad's ``ks`` / + ``ks_count``). + """ out: List[ArgSpec] = [ ("A", "a"), ("B", "b"), @@ -140,6 +145,9 @@ def _fwd_arg_names(*, is_3d: bool = False) -> List[ArgSpec]: # k -> ([z,] y, x, c): divisors cpg, X[, Y]. out += _magic("k_", (["Y"] if is_3d else []) + ["X", "cpg"]) out += _grid() + if split_k: + out += [("ws_ptr", "f32*"), ("ws_bytes", "i32")] + out += [("ks", "i32"), ("ks_count", "i32")] return out @@ -313,17 +321,18 @@ def conv_arg_names( callers that already dispatch on it (the manifest emitter, the grouped dispatcher, the AOT cache) stop branching. - ``two_stage`` belongs to wgrad's deterministic epilogue and is rejected - elsewhere rather than ignored: silently accepting it would hand back a - signature that does not match any kernel the flag can build. + ``two_stage`` is wgrad's deterministic epilogue, and on fwd the split-K + binary (fp32 workspace slabs + a Stage 2 reduce). dgrad has neither, so it + is rejected there rather than ignored: silently accepting it would hand + back a signature that does not match any kernel the flag can build. ``streamk`` (wgrad only) is the stream-K reduction, or ``None`` for a kernel without stream-K. Stream-K partitions K itself, so it excludes the two-stage split-K variant. """ - if two_stage and direction != "wgrad": + if two_stage and direction not in ("wgrad", "fwd"): raise ValueError( - f"two_stage is a wgrad-only variant (got direction={direction!r})" + f"two_stage is a wgrad / fwd split-K variant (got direction={direction!r})" ) if streamk is not None: if direction != "wgrad": @@ -338,7 +347,7 @@ def conv_arg_names( if two_stage: raise ValueError("streamk and two_stage are exclusive: both partition K_wg") if direction == "fwd": - return _fwd_arg_names(is_3d=is_3d) + return _fwd_arg_names(is_3d=is_3d, split_k=two_stage) if direction == "wgrad": return _wgrad_arg_names(is_3d=is_3d, two_stage=two_stage, streamk=streamk) if direction == "dgrad": diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_args.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_args.py index c1b2e180b971..2abe10661704 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_args.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_args.py @@ -422,7 +422,8 @@ class ConvArgs: # mismatched by accident. tile_m: int = 0 tile_n: int = 0 - # wgrad only: the granularity the split-K slice width is rounded to. + # wgrad, and the fwd split-K binary: the granularity the split-K slice + # width is rounded to. tile_k: int = 0 def __post_init__(self) -> None: @@ -622,11 +623,11 @@ def _named_values(self) -> Dict[str, int]: def _variant_values(self, split_k: int) -> Dict[str, int]: """Scalars whose meaning depends on the variant, not just the name.""" - if self.algorithm == "direct" or self.direction != "wgrad": + if self.algorithm == "direct" or self.direction == "dgrad": if split_k != 1: raise ValueError( - f"split_k={split_k}: only implicit-GEMM wgrad splits the " - f"reduction ({self.algorithm} {self.direction} has no " + f"split_k={split_k}: only implicit-GEMM wgrad and fwd split " + f"the reduction ({self.algorithm} {self.direction} has no " f"split-K arguments)" ) if self.algorithm == "direct": @@ -638,10 +639,21 @@ def _variant_values(self, split_k: int) -> Dict[str, int]: # The tail past wg_K reads zero through the descriptor's bounds, # which is what lets all ``split_k`` slices share one uniform # bound computation. - values["ks"] = ceil_div(self.gemm_k, self.tile_k * split_k) * self.tile_k - values["ks_count"] = split_k + values.update(self._split_values(split_k)) return values + def _split_values(self, split_k: int) -> Dict[str, int]: + """``ks`` / ``ks_count``: the split-K slice width and degree.""" + if self.tile_k <= 0: + raise ValueError( + "a split-K launch needs the kernel's tile_k: the slice width is " + "rounded to it" + ) + return { + "ks": ceil_div(self.gemm_k, self.tile_k * split_k) * self.tile_k, + "ks_count": split_k, + } + def _divisor(self, name: str) -> int: g = self.geom return { @@ -758,6 +770,14 @@ def to_launch_values( supplied = {k: v for k, v in supplied.items() if v is not None} named = self._named_values() named.update(self._variant_values(split_k)) + if self.direction == "fwd" and "ks" in names: + # The split-K forward binary (the ABI has its workspace pair). + named.update(self._split_values(split_k)) + elif split_k != 1 and self.direction == "fwd": + raise ValueError( + f"split_k={split_k}: a forward split-K launch needs the " + f"workspace (ws_ptr / ws_bytes)" + ) if streamk is not None: named.update(self._streamk_values(streamk)) for name in names[6:]: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_fwd_split_k_reduce.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_fwd_split_k_reduce.py new file mode 100644 index 000000000000..d0b69af79e93 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_fwd_split_k_reduce.py @@ -0,0 +1,158 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Stage 2 of the split-K forward convolution: fold the slabs and cast. + +Stage 1 (``ImplicitGemmConvSpec`` with ``split_k > 1``) stores each K slice's +fp32 partial output into its own slab of a ``[ks_count][M][K]`` workspace -- +plain stores, no atomics, every element written exactly once. This kernel sums +the ``ks_count`` slabs element-wise in a fixed order and converts the result to +``dtype_d`` into the output tensor ``D`` (N[D]HWK, the same packed layout as one +slab). The result is deterministic. + +The degree is a kernarg, so one binary per (``dtype_d``, ``vec``) serves every +split. Each thread owns ``vec`` consecutive output elements: one vector load per +slab, ``vec`` adds, one vector store. ``vec`` must divide the element count +(``M * K``), which a K divisible by ``vec`` guarantees. + +Kernel signature:: + + ws_ptr : f32 global ptr, readonly -- workspace [ks_count * elems] + d_ptr : dtype_d global ptr -- output [elems] + elems : i32 -- elements per slab (M * K) + ws_bytes : i32 -- workspace byte size (ABI boundary) + d_bytes : i32 -- output byte size (ABI boundary) + ks_count : i32 -- number of slabs (the split-K degree) + +Grid: ``(ceil(elems / (vec * block_size)), 1, 1)``; block ``(block_size, 1, 1)``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Tuple + +from rocke.core.ir import F32, I32, IRBuilder, KernelDef, PtrType +from rocke.helpers.io import io_ir_type, pack_f32_to, store_scalar_from_f32 +from rocke.helpers.spec import SignatureBuilder, ceil_div, kernel_name_join + + +@dataclass(frozen=True) +class FwdSplitKReduceSpec: + """One Stage 2 binary: the output dtype and elements per thread. + + ``degree`` > 0 specialises the fold for exactly that many slabs: the slab + loop is unrolled, so all of a thread's loads issue back to back instead of + one per dependent loop trip (``ks_count`` must then equal it). 0 keeps the + runtime loop over ``ks_count``. + """ + + dtype_d: str = "bf16" + vec: int = 4 + block_size: int = 256 + degree: int = 0 + name: str = "conv_fwd_splitk_reduce" + + def kernel_name(self) -> str: + parts = [f"v{self.vec}", f"b{self.block_size}"] + if self.degree: + parts.append(f"s{self.degree}") + return kernel_name_join(self.name, *parts, self.dtype_d) + + def validate(self) -> None: + if self.vec not in (1, 2, 4): + raise ValueError(f"vec must be 1, 2 or 4 (got {self.vec})") + if self.dtype_d not in ("fp16", "bf16"): + raise ValueError(f"dtype_d must be fp16 or bf16 (got {self.dtype_d})") + if self.degree < 0: + raise ValueError(f"degree must be >= 0 (got {self.degree})") + + +def build_conv_fwd_split_k_reduce( + spec: FwdSplitKReduceSpec, *, arch: str = "gfx950" # noqa: ARG001 +) -> KernelDef: + spec.validate() + vec = spec.vec + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = spec.block_size + + ws_ptr = b.param( + "ws_ptr", PtrType(F32, "global"), noalias=True, readonly=True, align=16 + ) + d_ptr = b.param( + "d_ptr", + PtrType(io_ir_type(spec.dtype_d), "global"), + noalias=True, + writeonly=True, + align=16, + ) + elems = b.param("elems", I32) + _ws_bytes = b.param("ws_bytes", I32) # noqa: F841 -- ABI boundary + _d_bytes = b.param("d_bytes", I32) # noqa: F841 -- ABI boundary + ks_count = b.param("ks_count", I32) + + gid = b.add(b.mul(b.block_id_x(), b.const_i32(spec.block_size)), b.thread_id_x()) + idx = b.mul(gid, b.const_i32(vec)) + + def _load(off): + if vec == 1: + return [b.global_load_f32(ws_ptr, off)] + v = b.global_load_vN(ws_ptr, off, F32, vec) + return [b.vec_extract(v, i) for i in range(vec)] + + with b.scf_if(b.cmp_lt(idx, elems)): + if spec.degree: + # Unrolled: every slab's load is issued before the first add. + slabs = [ + _load(idx if s == 0 else b.add(b.mul(b.const_i32(s), elems), idx)) + for s in range(spec.degree) + ] + total = slabs[0] + for part in slabs[1:]: + total = [b.fadd(a, p) for a, p in zip(total, part)] + else: + zero = b.const_f32(0.0) + init = [(f"acc{i}", zero) for i in range(vec)] + loop = b.scf_for_iter( + b.const_i32(0), ks_count, b.const_i32(1), init, iv_name="slab" + ) + with loop as (s, accs): + parts = _load(b.add(b.mul(s, elems), idx)) + b.scf_yield(*[b.fadd(a, p) for a, p in zip(accs, parts)]) + total = list(loop.results) + if vec == 1: + store_scalar_from_f32(b, d_ptr, idx, total[0], dtype=spec.dtype_d) + else: + b.global_store_vN( + d_ptr, idx, pack_f32_to(b, total, dtype=spec.dtype_d), vec + ) + return b.kernel + + +def fwd_split_k_reduce_grid( + spec: FwdSplitKReduceSpec, elems: int +) -> Tuple[int, int, int]: + return (max(1, ceil_div(elems, spec.vec * spec.block_size)), 1, 1) + + +def fwd_split_k_reduce_signature(spec: FwdSplitKReduceSpec) -> list: + return ( + SignatureBuilder() + .ptr("ws_ptr", "fp32") + .ptr("d_ptr", spec.dtype_d) + .scalar("elems", "i32") + .scalar("ws_bytes", "i32") + .scalar("d_bytes", "i32") + .scalar("ks_count", "i32") + .build() + ) + + +def fwd_split_k_reduce_vec(elems: int) -> int: + """The widest ``vec`` that divides ``elems``.""" + return next(v for v in (4, 2, 1) if elems % v == 0) + + +def fwd_split_k_workspace_nbytes(M: int, K: int, ks_count: int) -> int: + """Bytes of the ``[ks_count][M][K]`` fp32 workspace.""" + return 4 * ks_count * M * K diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py index 993759459bca..e2638456b369 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py @@ -328,6 +328,28 @@ class ImplicitGemmConvSpec: # cshuffle store width derived from kpg is a separate Gm=1 change.) group_merge: int = 1 + # Split-K (WMMA only). ``split_k > 1`` builds the split-K binary, not a + # fixed degree: the degree and the slice width arrive as the ``ks_count`` / + # ``ks`` kernargs (as in wgrad), so one binary serves every degree and the + # AOT cache records it as 2. ``-1`` (auto) is resolved by the builder + # through the conv split-K resolver (:func:`resolve_fwd_split_k_spec`) + # before validation, as for wgrad and dgrad. Every CTA reduces one slice + # of the [Z*]Y*X*cpg reduction and stores its fp32 partial tile into slab + # ``slice`` of the ``ws_ptr`` workspace, laid out ``[ks_count][M][K]``; + # Stage 2 (:mod:`kernels.common.conv_fwd_split_k_reduce`) sums the slabs and + # casts into D. Every slab element is written by exactly one CTA (a slice + # past the reduction stores zeros), so the workspace needs no clearing and + # the result is deterministic. 1 (the default) emits none of this. + split_k: int = 1 + # Single-launch split-K through a gfx1250 workgroup cluster (0 = off). + # The ``split_k_cluster`` slices of one output tile form one (1, 1, S) + # cluster: each stores its fp32 partial tile into its workspace slab as + # above, then the cluster barrier (cluster-scope release / acquire, no L2 + # writeback) makes the slabs visible and each slice folds 1/S of the + # tile's rows across the S slabs and stores them cast into D. No Stage 2 + # launch. Bakes the degree: the launch passes ks_count == S. + split_k_cluster: int = 0 + @property def block_size(self) -> int: return self.warp_m * self.warp_n * self.wave_size @@ -424,6 +446,10 @@ def kernel_name(self) -> str: # without it two degrees collide on one name and the compile # cache hands the second the first's binary. f"gm{self.group_merge}": self.group_merge > 1, + # The split-K binary has its own ABI and epilogue; the degree is + # a kernarg, so only the variant reaches the name. + "splitk": self.split_k > 1, + f"skc{self.split_k_cluster}": self.split_k_cluster > 0, }, ) @@ -606,6 +632,123 @@ def fwd_group_merge_available( # --------------------------------------------------------------------- +def fwd_split_k_cluster_problem_reason(problem: ConvProblem) -> Optional[str]: + """Why the cluster split-K forward cannot run ``problem``, or ``None``. + + The cluster fold walks 4-channel quads of the packed ``[M][K]`` output and + workspace, bounded only by the quad's first channel, so every quad has to + lie inside one group's channels: ``K`` and the per-group ``kpg`` must be + multiples of 4. The binary is shape-generic (K and kpg are kernargs), so + this is checked both on the build problem and on every launched problem. + """ + kpg = problem.K // max(1, problem.groups) + if problem.K % 4 or kpg % 4: + return ( + f"split_k_cluster folds 4-channel quads: K ({problem.K}) and the " + f"per-group channels ({kpg}) must be multiples of 4" + ) + return None + + +def fwd_split_k_reason( + spec: ImplicitGemmConvSpec, family: str, target=None +) -> Optional[str]: + """Why ``spec``'s split-K forward cannot be built, or ``None``. + + The split-K body is the plain WMMA K loop (``mem`` / ``wavelet``) run over + one slice, with an fp32 store into the workspace slab. The paths that drive + the K loop differently (async DMA, ``unroll_k``), remap the grid + (``chiplet_swizzle``), merge groups, or take the pointwise fast path are not + wired for it. Stage 1 stores raw partial sums and the reduction (Stage 2 + or the cluster fold) only adds and casts them, so the accumulator + epilogue (bias / scale / ReLU / clamp) must be the identity. ``target`` + (an ``ArchTarget``) gates the cluster form on workgroup-cluster launch. + """ + if spec.split_k < 1: + return f"split_k must be >= 1 (got {spec.split_k})" + if family != "wmma": + return "split-K forward is WMMA-only (built and tested on wave32)" + if not spec.acc_epilogue.is_identity(): + return ( + "split-K forward supports only the identity acc_epilogue: the " + "partial sums are reduced after the per-slice accumulation" + ) + for flag, label in ( + (spec.async_dma, "async_dma"), + (spec.unroll_k, "unroll_k"), + (spec.chiplet_swizzle, "chiplet_swizzle"), + (spec.group_merge > 1, "group_merge"), + (spec.problem.is_pointwise, "the pointwise fast path"), + (spec.problem.is_3d, "3-D convolution"), + ): + if flag: + return f"split-K forward does not support {label}" + if spec.epilogue == "default" and (spec.vector_size_c or 1) > 1: + return "split-K forward with the default epilogue stores scalars (vector_size_c must be 1)" + if spec.split_k_cluster: + if spec.split_k < 2: + return "split_k_cluster needs the split-K binary (split_k > 1)" + if not 2 <= spec.split_k_cluster <= 15: + return f"split_k_cluster must be 2..15 (got {spec.split_k_cluster})" + if spec.pipeline != "mem": + return "split_k_cluster is wired for the mem pipeline only" + if spec.tile_n % 4: + return "split_k_cluster folds 4-channel quads (tile_n % 4 == 0)" + if target is not None and not target.memory.has_cluster_launch: + return f"split_k_cluster needs workgroup-cluster launch, which {target.gfx} lacks" + _why = fwd_split_k_cluster_problem_reason(spec.problem) + if _why is not None: + return _why + return None + + +def resolve_fwd_split_k_spec( + spec: ImplicitGemmConvSpec, arch: str +) -> ImplicitGemmConvSpec: + """``spec`` with an auto (``-1``) ``split_k`` resolved; any other spec is + returned unchanged. + + A spec the split form cannot take (:func:`fwd_split_k_reason`: a fused + accumulator epilogue, async DMA, ``unroll_k`` and the like) resolves to 1, + as the problems that form cannot take do, so an auto request still builds + the plain kernel. The cluster settings are the caller's and stay out of + that check: their own errors still come from validation. + + The degree comes from the shared conv split-K resolver + (:func:`kernels.common.conv_split_k.resolve_conv_split_k`), so a builder + and the benchmark agree on it. It may resolve to 1 (the plain kernel): + where the split does not pay, or where this spec's arch cannot build it. + A resolved degree > 1 only selects the split-K binary -- the launch passes + the degree as ``ks_count`` -- so the caller launches with + ``resolve_conv_split_k`` 's answer, which this spec now records. The C++ + engine, like its wgrad port, takes explicit degrees only. + """ + if spec.split_k != -1: + return spec + from kernels.common.conv_split_k import resolve_conv_split_k + + resolved = resolve_conv_split_k( + "fwd", + spec.problem, + arch=arch, + dtype=spec.data.dtype_d, + tile_m=spec.tile_m, + tile_n=spec.tile_n, + tile_k=spec.tile_k, + block_size=spec.block_size, + ) + degree = resolved.split_k + if degree > 1: + from rocke.core.arch import ArchTarget + + target = ArchTarget.from_gfx(arch) + family = "wmma" if target.wave_size == 32 else "mma" + probe = dc_replace(spec, split_k=degree, split_k_cluster=0) + if fwd_split_k_reason(probe, family, target) is not None: + degree = 1 + return dc_replace(spec, split_k=degree) + + def is_valid_spec_for_problem( spec: ImplicitGemmConvSpec, problem: ConvProblem, arch: str = "gfx950" ) -> Tuple[bool, str]: @@ -618,6 +761,12 @@ def is_valid_spec_for_problem( ok, reason = is_valid_spec(spec, arch) if not ok: return False, reason + if spec.split_k_cluster: + # The cluster binary takes K and kpg as kernargs: the launched problem + # has to satisfy the fold's quad rule, not only the build problem. + _why = fwd_split_k_cluster_problem_reason(problem) + if _why is not None: + return False, _why # caps gridDim.y (and .x/.z) at 65535. The M-tile axis maps # to gridDim.y so reject specs whose grid would silently truncate and @@ -663,7 +812,8 @@ def implicit_gemm_lds_bytes(spec: ImplicitGemmConvSpec) -> Tuple[int, int, int, _double = spec.async_dma or spec.unroll_k _ab_lds = _ab_bytes * (2 if _double else 1) # cshuffle stages tile_m×tile_n elements at dtype_d (fp16/bf16 = 2B, fp32 = 4B). - _c_dtype_bytes = 4 if spec.data.dtype_d == "fp32" else 2 + # The split-K binary stages its fp32 partial tile, whatever dtype_d is. + _c_dtype_bytes = 4 if spec.data.dtype_d == "fp32" or spec.split_k > 1 else 2 _c_lds = ( spec.tile_m * spec.tile_n * _c_dtype_bytes if spec.epilogue == "cshuffle" else 0 ) @@ -854,6 +1004,11 @@ def is_valid_spec(spec: ImplicitGemmConvSpec, arch: str = "gfx950") -> Tuple[boo if flag: return False, f"WMMA conv does not support {label} on {arch}" + if spec.split_k != 1 or spec.split_k_cluster: + _why = fwd_split_k_reason(spec, family, target) + if _why is not None: + return False, _why + # The tile loaders. An explicit vector width (or the async chunk width the # tile admits) has to split the tile evenly over the block's threads; the # builder would otherwise only find out halfway through a build. The @@ -1104,7 +1259,10 @@ def _build_implicit_gemm_conv_impl( only exists on gfx950 (e.g. the default wide ``16x16x32`` f16 atom) for ``gfx942`` fails with a clean structured error instead of an ``LLVM ERROR`` from comgr. + + ``spec.split_k == -1`` is resolved first (:func:`resolve_fwd_split_k_spec`). """ + spec = resolve_fwd_split_k_spec(spec, arch) spec.validate() ok, why = is_valid_spec(spec, arch=arch) if not ok: @@ -1203,6 +1361,16 @@ def _build_implicit_gemm_conv_impl( # Magic-number pairs, one per unmerge divisor (host-precomputed). p_num_pid_m = params["p_num_pid_m"] p_num_pid_n = params["p_num_pid_n"] + # Split-K (see ImplicitGemmConvSpec.split_k): the fp32 workspace, then the + # slice width and degree -- the same trailing pair wgrad takes. + split = spec.split_k > 1 + if split: + p_ws = b.param("ws_ptr", PtrType(F32, "global"), noalias=True, align=16) + p_ws_bytes = b.param("ws_bytes", I32) + p_ks = b.param("ks", I32) + p_ks_count = b.param("ks_count", I32) + if spec.split_k_cluster: + b.set_cluster_dims(1, 1, spec.split_k_cluster) def _magic(prefix: str, name: str): """The ``(mult, shift)`` SSA pair for one unmerge divisor.""" @@ -1322,8 +1490,22 @@ def _magic(prefix: str, name: str): # every group count Gm divides, so the build-time problem's count cannot # decide whether block_id_z is needed. grouped = p.groups > 1 or merged + if split: + # The group and the K slice share z: z = group * ks_count + slice, as + # in wgrad. Slice ``s`` reduces [s*ks, (s+1)*ks); the host rounds ks + # to whole K tiles, and the tail past K_gemm reads zero through the + # descriptor bounds, so every slice runs the same bound computation. + z_id = b.block_id_z() + if grouped: + group_z = b.to_sgpr_u32(b.div(z_id, p_ks_count)) + slice_id = b.sub(z_id, b.mul(group_z, p_ks_count)) + else: + group_z = None + slice_id = z_id + split_k_lo = b.to_sgpr_u32(b.mul(slice_id, p_ks)) + split_k_hi = b.to_sgpr_u32(b.add(split_k_lo, p_ks)) if grouped: - group_idx = b.block_id_z() + group_idx = group_z if split else b.block_id_z() # Merged: the base advances by the merged kpg, which is Gm exactly # (depthwise only), so it is the build-time constant, not p_kpg. k_out_group_base = b.mul(group_idx, c_gm if merged else p_kpg) @@ -1988,13 +2170,31 @@ def d_addr(b_: IRBuilder, m_val: Value, n_val: Value): # Merged: g is the merged group and the merged kpg is Gm, so the same # expression yields k_out = merged_group*Gm + g_n. _k_out_grp_base = ( - b.mul(b.block_id_z(), c_gm if merged else p_kpg) if grouped else None + b.mul(group_idx if split else b.block_id_z(), c_gm if merged else p_kpg) + if grouped + else None ) def d_addr(b_: IRBuilder, m_val: Value, n_val: Value): k_out = n_val if _k_out_grp_base is None else b_.add(_k_out_grp_base, n_val) return D_desc.offset(b_, m=m_val, k_out=k_out) + if split: + # Stage 1 stores into slab ``slice`` of the [ks_count][M][K] fp32 + # workspace. The slab is a packed N[D]HWK tensor like D, so D's + # descriptor addresses it; only the slab base is added. + _slab_base = b.mul(slice_id, b.mul(p_M, p_K)) + _d_addr_out = d_addr + + def d_addr(b_: IRBuilder, m_val: Value, n_val: Value): + off, valid = _d_addr_out(b_, m_val, n_val) + return b_.add(_slab_base, off), valid + + d_rsrc = make_buffer_resource(b, p_ws, num_bytes=p_ws_bytes).rsrc + _epi_spec = dc_replace(spec, data=dc_replace(spec.data, dtype_d="fp32")) + else: + _epi_spec = spec + # The GEMM-N extent the tile covers: kpg, or kpg*Gm == Gm under merge. The # true kpg (1 on depthwise) would clamp away Gm-1 of every Gm real outputs # -- silently, as a wrong answer rather than a crash. @@ -2015,7 +2215,7 @@ def _emit_fwd_epilogue(final_accs_in): """Store the accumulators through ``d_addr`` with runtime bounds.""" if spec.epilogue == "cshuffle": cshuffle_kwargs = { - "out_dtype": spec.data.dtype_d, + "out_dtype": _epi_spec.data.dtype_d, "max_store_vec": _max_store_vec, "no_alias": spec.cshuffle_no_alias, } @@ -2041,7 +2241,7 @@ def _emit_fwd_epilogue(final_accs_in): # predates the helper-based path. emit_direct_epilogue_wmma( b, - spec, + _epi_spec, op, final_accs_in, warp_m_idx, @@ -2161,8 +2361,8 @@ def _fwd_epilogue_with_acc_transform(accs_in): tid=tid, n_math_warps=n_math_warps, math_block_size=spec.block_size, - k_lo=c0, - k_hi=c_K_gemm, + k_lo=split_k_lo if split else c0, + k_hi=split_k_hi if split else c_K_gemm, block_k=block_k, A_smem=A_smem, B_smem=B_smem, @@ -2234,7 +2434,11 @@ def _fwd_epilogue_with_acc_transform(accs_in): final_accs = for_op.results elif not spec.async_dma: # mem / compv3 / compv4 / basic: all use scf_for_iter — only bound changes to runtime. - for_op = b.scf_for_iter(c0, c_K_gemm, c_block_k, accs, iv_name="k0") + for_op = ( + b.scf_for_iter(split_k_lo, split_k_hi, c_block_k, accs, iv_name="k0") + if split + else b.scf_for_iter(c0, c_K_gemm, c_block_k, accs, iv_name="k0") + ) with for_op as (k0, iter_vars): emit_load_phase(k0, A_smem, B_smem) b.sync() @@ -2283,9 +2487,85 @@ def compute_dyn(k_offset_val, buf_pair, state): epilogue_override(b, spec, final_accs, grid, d_rsrc, extra_context) else: _emit_fwd_epilogue(final_accs) + if split and spec.split_k_cluster: + _emit_cluster_fold( + b, + spec, + tid=tid, + slice_id=slice_id, + m0=block_m_off_v, + n0=block_n_off_v, + block_m=block_m, + block_n=block_n, + p_M=p_M, + p_K=p_K, + n_bound=c_N_bound, + k_base=k_out_group_base, + ws=p_ws, + D=D, + ) return b.kernel +def _emit_cluster_fold( + b: IRBuilder, + spec: "ImplicitGemmConvSpec", + *, + tid, + slice_id, + m0, + n0, + block_m: int, + block_n: int, + p_M, + p_K, + n_bound, + k_base, + ws, + D, +) -> None: + """Fold the cluster's S slabs of this tile into D (``split_k_cluster``). + + Slice ``s`` owns rows ``[s*R, (s+1)*R)`` of the tile, R = ceil(TM / S); + each thread walks 4-channel quads of those rows (D and the slabs are + packed ``[M][K]``; :func:`fwd_split_k_cluster_problem_reason` requires + K % 4 == 0 and kpg % 4 == 0 at build and launch admission, so a quad + whose first channel is in bounds lies wholly in its group). + """ + from rocke.helpers.io import pack_f32_to + + S = spec.split_k_cluster + rows = -(-block_m // S) + quads_per_row = block_n // 4 + n_quads = rows * quads_per_row + threads = spec.block_size + b.cluster_barrier() + row0 = b.add(m0, b.mul(slice_id, b.const_i32(rows))) + slab = b.mul(p_M, p_K) + for t in range(-(-n_quads // threads)): + q = b.add(tid, b.const_i32(t * threads)) + r = b.div(q, b.const_i32(quads_per_row)) + c = b.mul(b.sub(q, b.mul(r, b.const_i32(quads_per_row))), b.const_i32(4)) + m = b.add(row0, r) + n = b.add(n0, c) + ok = b.land( + b.land(b.cmp_lt(q, b.const_i32(n_quads)), b.cmp_lt(m, p_M)), + b.land(b.cmp_lt(m, b.add(m0, b.const_i32(block_m))), b.cmp_lt(n, n_bound)), + ) + with b.scf_if(ok): + k_out = n if k_base is None else b.add(k_base, n) + off = b.add(b.mul(m, p_K), k_out) + parts = [] + for s in range(S): + o = off if s == 0 else b.add(b.mul(b.const_i32(s), slab), off) + v = b.global_load_vN(ws, o, F32, 4) + parts.append([b.vec_extract(v, i) for i in range(4)]) + tot = parts[0] + for pp in parts[1:]: + tot = [b.fadd(x, y) for x, y in zip(tot, pp)] + b.global_store_vN(D, off, pack_f32_to(b, tot, dtype=spec.data.dtype_d), 4) + + def async_tile_loaders( spec: "ImplicitGemmConvSpec", arch: str = "gfx950" ) -> Tuple[AsyncTileLoader, AsyncTileLoader]: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py index a4d7c09b20a0..d1ce5a500759 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py @@ -26,7 +26,14 @@ the workspace only; the workspace must fit the i32 ``ws_bytes`` kernarg; WMMA (wave32) arches have no workspace path, and their packed atomic needs neither cshuffle nor an even store vector; -* the gridDim.z budget (``z = groups * split_k``). +* the gridDim.z budget (``z = groups * split_k``); +* for the forward split (the two-stage WMMA kernel, built on + :data:`FWD_SPLIT_K_ARCHS` only): its GEMM ``N*[Do*]Ho*Wo x kpg x + [Z*]Y*X*cpg``, the problems it cannot take (pointwise, 3-D) and the i32 + workspace of ``split_k`` fp32 slabs of the output. Its auto degree comes from + :func:`rocke.helpers.split_k.select_split_k_fwd` -- no arch it is built for + has a fitted cost model. Its partials meet in fp32 slabs, never in 16-bit + atomics, so no accuracy cap applies. An explicit degree, an override, and an auto degree on an arch without a cost model (gfx942 included: it has not been fitted) all go through the @@ -53,7 +60,7 @@ from dataclasses import dataclass from rocke.core.arch.target import ArchTarget -from rocke.helpers.split_k import select_split_k_wgrad +from rocke.helpers.split_k import select_split_k_fwd, select_split_k_wgrad from rocke.helpers.split_k_planner import ( MAX_GRID_DIM_Z, SplitKKernel, @@ -64,6 +71,7 @@ ) from kernels.common._conv_implicit_gemm_common import ConvProblem +from kernels.common.conv_fwd_split_k_reduce import fwd_split_k_workspace_nbytes from kernels.common.conv_implicit_gemm_dgrad import ( _dg_K, _dg_M, @@ -83,6 +91,7 @@ __all__ = [ "ATOMIC_16BIT_MAX_SPLIT_K", + "FWD_SPLIT_K_ARCHS", "MAX_WGRAD_WS_BYTES", "ConvSplitK", "parse_split_k_override", @@ -99,6 +108,13 @@ # ws_bytes is an i32 kernarg, so a two-stage workspace above this overflows it. MAX_WGRAD_WS_BYTES = (1 << 31) - 4 +# Targets the split-K forward kernel is built and validated for (the AOT cache +# builds it for these only). Elsewhere an auto forward request resolves to 1. +FWD_SPLIT_K_ARCHS = ("gfx1250",) +# The forward workspace is addressed through an i32 byte count like every +# operand. +MAX_FWD_WS_BYTES = (1 << 31) - 1 + _DIRECTIONS = ("fwd", "wgrad", "dgrad") _REDUCTIONS = ("auto", "atomic", "workspace") @@ -109,7 +125,8 @@ class ConvSplitK: Attributes: split_k: the degree to build and launch (1 = no split). - two_stage: reduce through the f32 workspace and Stage 2 (wgrad only). + two_stage: reduce through the f32 workspace and Stage 2 (wgrad; the + forward split always does, so it is set whenever it splits). requested: the degree before the launch limits applied -- the explicit value, the override, or the auto choice. reason: why, for logs and diagnostics. @@ -331,6 +348,88 @@ def _rule(requested: int, source: str, *, atomic_cap: bool) -> ConvSplitK: return ConvSplitK(plan.split_k, two_stage, plan.split_k, plan.reason) +def fwd_split_k_problem_reason(p: ConvProblem, arch: str) -> str | None: + """Why no forward split-K kernel can serve ``p`` on ``arch``, or ``None``. + + Only what the problem and the target decide; the spec-level gates (WMMA, + pipeline, epilogue) are the builder's + (:func:`kernels.common.conv_implicit_gemm.fwd_split_k_reason`). + """ + if arch.split(":", 1)[0] not in FWD_SPLIT_K_ARCHS: + return f"split-K forward is built for {', '.join(FWD_SPLIT_K_ARCHS)} only" + if p.is_pointwise: + return "split-K forward does not support the pointwise fast path" + if p.is_3d: + return "split-K forward does not support 3-D convolution" + return None + + +def _fwd_degree_limit(p: ConvProblem, degree: int) -> int: + """``degree`` held to what one forward split-K launch can encode: the + gridDim.z budget (``z = group * split_k + slice``) and the i32 workspace. + The workspace halves with the degree, so the largest degree that fits is + found by halving (an explicit degree is not held; the launch rejects it).""" + degree = min(degree, MAX_GRID_DIM_Z // max(1, p.groups)) + while degree > 1 and fwd_split_k_workspace_nbytes(p.M, p.K, degree) > ( + MAX_FWD_WS_BYTES + ): + degree //= 2 + return max(1, degree) + + +def _resolve_fwd( + p: ConvProblem, + *, + arch: str, + tile_m: int, + tile_n: int, + tile_k: int, + split_k: int, + allow_unsplit: bool, + override: int | None, + num_cus: int | None, +) -> ConvSplitK: + if split_k != -1: + return ConvSplitK(split_k, split_k > 1, split_k, f"explicit split_k={split_k}") + why = fwd_split_k_problem_reason(p, arch) + if override is not None: + # Held to what the kernel can build, as for dgrad: an unbuildable + # split becomes 1 rather than a spec the builder then rejects. + degree = 1 if why is not None else _fwd_degree_limit(p, override) + if why is not None and override != 1: + note = f" (no split: {why})" + elif degree < override: + note = f" (held to {degree})" + else: + note = "" + return ConvSplitK( + degree, degree > 1, override, f"override split_k={override}{note}" + ) + if why is not None: + decision_degree, reason = 1, f"no split: {why}" + else: + d = select_split_k_fwd( + M=p.M, + N=p.kpg, + K=p.K_gemm, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + groups=p.groups, + arch=arch, + num_cus=num_cus, + ) + decision_degree = _fwd_degree_limit(p, d.split_k) + reason = d.reason + if decision_degree < d.split_k: + reason += f" (held to {decision_degree} by gridDim.z / the workspace)" + if decision_degree == 1 and not allow_unsplit: + # The caller holds a split-K binary: degree 1 is the separate plain + # kernel, and the policy says splitting this problem does not pay. + raise ValueError(f"fwd split-K does not serve this problem: {reason}") + return ConvSplitK(decision_degree, decision_degree > 1, decision_degree, reason) + + def _resolve_dgrad( p: ConvProblem, *, @@ -459,7 +558,11 @@ def resolve_conv_split_k( on, for a part built differently from its arch's default; ``None`` keeps the arch default. - Forward convolution has no split-K path and always resolves to 1. + Forward convolution splits only on :data:`FWD_SPLIT_K_ARCHS` (the + two-stage WMMA kernel); elsewhere, and for the problems that kernel cannot + take, an auto or override request resolves to 1. ``block_size``, + ``group_merge``, ``ws_replicas`` and ``num_xcds`` do not affect the + forward answer. """ if direction not in _DIRECTIONS: raise ValueError(f"direction must be one of {_DIRECTIONS}, got {direction!r}") @@ -471,8 +574,6 @@ def resolve_conv_split_k( epilogue = None if override is not None and override < 1: raise ValueError(f"override must be None or >= 1, got {override}") - if direction == "fwd": - return ConvSplitK(1, False, 1, "fwd implicit GEMM has no split-K path") return _resolve_cached( direction, problem, @@ -531,6 +632,18 @@ def _resolve_cached( "num_cus": num_cus, "num_xcds": num_xcds, } + if direction == "fwd": + return _resolve_fwd( + problem, + arch=arch, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + split_k=split_k, + allow_unsplit=allow_unsplit, + override=override, + num_cus=num_cus, + ) if direction == "dgrad": return _resolve_dgrad(problem, **common) return _resolve_wgrad( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py index fe7b2c4058e9..b7785bfaeda0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py @@ -41,7 +41,12 @@ group-merged depthwise kernels (``group_merge`` in 4..64, cshuffle, plain and ``unroll_k`` loops, load widths derived by the builder, on the tile subset the ``CACHE_FWD_GROUP_MERGE_*`` constants allow), served only to - depthwise problems whose merged GEMM-N fits the tile. + depthwise problems whose merged GEMM-N fits the tile. Plus, on the + ``CACHE_FWD_SPLIT_K_ARCHS`` targets, the split-K kernel (``split_k`` > 1, + recorded as 2 with ``two_stage``; fp32 workspace slabs + the Stage 2 fold, + the degree a kernarg swept at launch) on the small tiles, pipelines and + widths of ``CACHE_FWD_SPLIT_K_*``, and its cluster form where the target + launches workgroup clusters and the flavor lowers them. * **wgrad** -- tile_m/tile_n also take 144 and 192, and warp_m/warp_n also take 3 and 6 (``CACHE_WGRAD_EXTRA_*``): wgrad's GEMM N is Z*Y*X*cpg, which carries a factor of 9 for a 3x3 filter, and the base ladder steps by 2x, so @@ -116,6 +121,7 @@ from dataclasses import dataclass from typing import Dict, Iterator, List, Optional, Sequence, Tuple +from kernels.common.conv_split_k import FWD_SPLIT_K_ARCHS from kernels.conv.kernel_identity import KernelIdentity, current_llvm_flavor @@ -310,6 +316,39 @@ # sweep does: one workgroup per merged group cannot fill the device. CACHE_WGRAD_SPLIT_KS: Tuple[int, ...] = (1, 2) +# Split-K forward (WMMA only; ImplicitGemmConvSpec.split_k): built for the +# small-M / deep-K layers it exists for, so only small output tiles, the two +# WMMA pipelines and the cshuffle epilogue. Recorded as split_k = 2 with +# two_stage set: the degree is a kernarg, swept at launch. +CACHE_FWD_SPLIT_K_TILE_M: Tuple[int, ...] = (16, 32, 64) +CACHE_FWD_SPLIT_K_TILE_N: Tuple[int, ...] = (16, 32, 64, 128) +CACHE_FWD_SPLIT_K_TILE_K: Tuple[int, ...] = (32, 64) +CACHE_FWD_SPLIT_K_PIPELINES: Tuple[str, ...] = ("mem", "wavelet") +CACHE_FWD_SPLIT_K_VECS: Tuple[int, ...] = (4, 8) +# Single-launch cluster split-K forward (ImplicitGemmConvSpec.split_k_cluster): +# the S slices of one output tile form a (1, 1, S) workgroup cluster and fold +# their slabs after a cluster barrier, so there is no Stage 2 launch. The +# degree is baked (the cluster shape), so one binary per S, on the tiles +# (tile_m, tile_n, warp_m, warp_n) that win the small-output / deep-K layers +# it is for; mem pipeline, cshuffle, tile_k 64, 16-byte vectors. +CACHE_FWD_SPLIT_K_CLUSTER_DEGREES: Tuple[int, ...] = (4, 8, 12, 15) +CACHE_FWD_SPLIT_K_CLUSTER_TILES: Tuple[Tuple[int, int, int, int], ...] = ( + (16, 32, 1, 1), + (32, 32, 2, 1), + (64, 32, 2, 1), + (64, 64, 2, 2), + (128, 128, 2, 2), +) +CACHE_FWD_SPLIT_K_CLUSTER_TILE_K = 64 +CACHE_FWD_SPLIT_K_CLUSTER_VEC = 8 +# Targets the split-K forward kernels are cached for: the ones they were +# validated on, which are also the ones the conv split-K resolver splits a +# forward on (one tuple, so the cache and auto-deduction cannot drift). The +# cluster form also needs workgroup-cluster launch on the target and the LLVM +# flavor that lowers the cluster attribute and barrier. +CACHE_FWD_SPLIT_K_ARCHS: Tuple[str, ...] = FWD_SPLIT_K_ARCHS +CACHE_FWD_SPLIT_K_CLUSTER_FLAVORS: Tuple[str, ...] = ("llvm23",) + # ---- wgrad stream-K axis (opt-in) ------------------------------------------- # Stream-K (WgradConvSpec.streamk) bakes its launch mode and its partial-sum # reduction into the kernel; the CTA pool -- and with it the whole iteration @@ -682,6 +721,150 @@ def _fwd_jobs( wave_size, (tile_m, tile_n, tile_k, warp_m, warp_n, wt, pipeline, epilogue, atom), ) + if _fwd_split_target(arch): + yield from _fwd_split_jobs( + arch, + da, + db, + dd, + flavor, + wave_size, + (tile_m, tile_n, tile_k, warp_m, warp_n, wt, pipeline, epilogue, atom), + cluster=_fwd_split_cluster_target(target, flavor), + ) + + +def _fwd_split_target(arch: str) -> bool: + """Whether split-K forward kernels are cached for ``arch`` at all.""" + return arch in CACHE_FWD_SPLIT_K_ARCHS + + +def _fwd_split_cluster_target(target, flavor: str) -> bool: + """Whether cluster split-K forward kernels are cached for ``target`` and + ``flavor``: workgroup-cluster launch and an LLVM that lowers it.""" + return ( + target.memory.has_cluster_launch and flavor in CACHE_FWD_SPLIT_K_CLUSTER_FLAVORS + ) + + +def _fwd_split_geometry(geometry: tuple) -> bool: + """Whether split-K forward kernels are built for ``geometry`` (WMMA only).""" + tile_m, tile_n, tile_k, _wm, _wn, _wt, pipeline, epilogue, atom = geometry + return ( + atom.family == "wmma" + and epilogue == "cshuffle" + and pipeline in CACHE_FWD_SPLIT_K_PIPELINES + and tile_m in CACHE_FWD_SPLIT_K_TILE_M + and tile_n in CACHE_FWD_SPLIT_K_TILE_N + and tile_k in CACHE_FWD_SPLIT_K_TILE_K + ) + + +def _fwd_split_cluster_geometry(geometry: tuple) -> bool: + """Whether cluster split-K forward kernels are built for ``geometry``.""" + tile_m, tile_n, tile_k, warp_m, warp_n, _wt, pipeline, epilogue, atom = geometry + return ( + atom.family == "wmma" + and epilogue == "cshuffle" + and pipeline == "mem" + and tile_k == CACHE_FWD_SPLIT_K_CLUSTER_TILE_K + and (tile_m, tile_n, warp_m, warp_n) in CACHE_FWD_SPLIT_K_CLUSTER_TILES + ) + + +def _fwd_split_jobs( + arch: str, + da: str, + db: str, + dd: str, + flavor: str, + wave_size: int, + geometry: tuple, + *, + cluster: bool, +) -> Iterator[BuildJob]: + """Split-K forward kernels for one geometry (see ``CACHE_FWD_SPLIT_K_*``). + + ``cluster`` adds the cluster binaries (see :func:`_fwd_split_cluster_target`). + """ + tile_m, tile_n, tile_k, warp_m, warp_n, wt, pipeline, epilogue, atom = geometry + if cluster and _fwd_split_cluster_geometry(geometry): + vec = CACHE_FWD_SPLIT_K_CLUSTER_VEC + for degree in CACHE_FWD_SPLIT_K_CLUSTER_DEGREES: + cfg = dict( + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + warp_m=warp_m, + warp_n=warp_n, + warp_tile_m=wt, + warp_tile_n=wt, + warp_tile_k=atom.k, + pipeline=pipeline, + epilogue=epilogue, + wave_size=wave_size, + vector_size_a=vec, + vector_size_b=vec, + vector_size_c=vec, + split_k=2, + split_k_cluster=degree, + ) + yield BuildJob( + identity=KernelIdentity( + arch=arch, + direction="fwd", + algorithm="implicit_gemm", + llvm_flavor=flavor, + dtype_a=da, + dtype_b=db, + dtype_d=dd, + grouped=True, + two_stage=True, + **cfg, + ), + direction="fwd", + spec_kwargs=cfg, + caps=dict(grouped=True), + ) + if not _fwd_split_geometry(geometry): + return + for vec_ab, vec_c in itertools.product( + CACHE_FWD_SPLIT_K_VECS, CACHE_FWD_SPLIT_K_VECS + ): + cfg = dict( + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + warp_m=warp_m, + warp_n=warp_n, + warp_tile_m=wt, + warp_tile_n=wt, + warp_tile_k=atom.k, + pipeline=pipeline, + epilogue=epilogue, + wave_size=wave_size, + vector_size_a=vec_ab, + vector_size_b=vec_ab, + vector_size_c=vec_c, + split_k=2, + ) + yield BuildJob( + identity=KernelIdentity( + arch=arch, + direction="fwd", + algorithm="implicit_gemm", + llvm_flavor=flavor, + dtype_a=da, + dtype_b=db, + dtype_d=dd, + grouped=True, + two_stage=True, + **cfg, + ), + direction="fwd", + spec_kwargs=cfg, + caps=dict(grouped=True), + ) def _fwd_merged_k_loops() -> Tuple[bool, ...]: @@ -1502,10 +1685,14 @@ def count_jobs( split_ks: Sequence[int] = CACHE_WGRAD_SPLIT_KS, streamk_modes: Sequence[str] = (), streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, + llvm_flavor: Optional[str] = None, ) -> int: """Raw candidate count ``kernel_sweep.enumerate_jobs`` walks, before dedup/validation. + ``llvm_flavor`` is the flavor the enumeration records (see + :func:`_job_flavor`); it gates the cluster split-K forward jobs. + Cheap: it runs only the geometry filter and multiplies out the inner axes, without building a single identity -- which is where enumeration spends its time. Used as the denominator of the enumeration progress. @@ -1519,6 +1706,10 @@ def count_jobs( layouts = len(_layout_caps()) merged = sum(1 for sk in split_ks if sk > 1) * len(_WGRAD_GROUP_MERGES) wave_size = target.wave_size + fwd_split = _fwd_split_target(target.gfx) + fwd_cluster = fwd_split and _fwd_split_cluster_target( + target, _job_flavor(llvm_flavor) + ) total = 0 for direction in directions: if direction not in _AOT_ALIAS_PIPELINES: @@ -1531,6 +1722,10 @@ def count_jobs( total += vecs * len(_fwd_k_loops(pipeline)) * layouts if _fwd_merged_geometry(geo): total += len(_FWD_GROUP_MERGES) * len(_fwd_merged_k_loops()) + if fwd_split and _fwd_split_geometry(geo): + total += len(CACHE_FWD_SPLIT_K_VECS) ** 2 + if fwd_cluster and _fwd_split_cluster_geometry(geo): + total += len(CACHE_FWD_SPLIT_K_CLUSTER_DEGREES) elif direction == "wgrad": total += ( vecs * sum(len(_wgrad_two_stages(sk)) for sk in split_ks) * layouts diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_identity.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_identity.py index b6b9432d99c8..c0d2e848289c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_identity.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_identity.py @@ -111,6 +111,12 @@ class KernelIdentity: waves_per_eu: int | None = None acc_epilogue: str = "none" split_k: int = 1 + # Single-launch cluster split-K forward (ImplicitGemmConvSpec. + # split_k_cluster): the degree is baked as the (1, 1, S) cluster shape and + # there is no Stage 2. Left out of the hash while 0, like ws_replicas. + split_k_cluster: int = 0 + # The binary takes the split-K workspace ABI. Whether a Stage 2 launch + # follows is ``needs_stage2`` (False for the cluster forward). two_stage: bool = False # Two-stage scratch slabs per group (WgradConvSpec.ws_replicas). The host # sizes the scratch and builds Stage 2 from it, so it has to come from the @@ -209,6 +215,19 @@ def is_direct(self) -> bool: """Direct-conv kernels: every capability field is baked, even a 0.""" return self.algorithm.startswith("direct") + @property + def needs_stage2(self) -> bool: + """Whether a launch of this binary must be followed by the Stage 2 + reduce kernel before the output is final. + + ``two_stage`` alone means the binary takes the split-K workspace ABI + (``ws_ptr``, ``ws_bytes``, ``ks``, ``ks_count``). The cluster split-K + forward takes that ABI too but folds its slabs in-kernel, so it has no + Stage 2. Anything deciding whether to launch Stage 2 asks this, never + ``two_stage``. + """ + return self.two_stage and not self.split_k_cluster + def stable_hash(self) -> str: """Deterministic SHA-1 hex digest (40 chars) of the identity. @@ -231,6 +250,8 @@ def stable_hash(self) -> str: for name in ("dy_halo", "dy_halo_setprio", "dy_halo_kouter_pad", "halo_max_wo"): if not fields[name]: del fields[name] + if not fields["split_k_cluster"]: + del fields["split_k_cluster"] # Same for the stream-K pair at "no stream-K". if fields["streamk"] == "off" and not fields["streamk_reduction"]: del fields["streamk"], fields["streamk_reduction"] @@ -502,6 +523,18 @@ def identity_supports_problem( if not identity.is_direct and problem_grouped and not identity.grouped: return False, "kernel is ungrouped, problem is grouped" + # The cluster split-K forward folds 4-channel quads bounded by their first + # channel; the binary takes K and kpg as kernargs, so the launched problem + # is checked with the same predicate the spec validator uses. + if identity.direction == "fwd" and identity.split_k_cluster: + from kernels.common.conv_implicit_gemm import ( + fwd_split_k_cluster_problem_reason, + ) + + why = fwd_split_k_cluster_problem_reason(problem) + if why is not None: + return False, why + # Split-K wgrad without two-stage atomic-adds straight into dW. For a # 16-bit dW that is a packed <2 x dtype> atomic, which needs an even # dW row (wg_N = [Z*]Y*X*cpg) and store width; on an odd row it diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_label.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_label.py index 9cabfaa8d662..521a54dea0ec 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_label.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_label.py @@ -69,6 +69,7 @@ waves_per_eu=None, acc_epilogue="none", split_k=1, + split_k_cluster=0, two_stage=False, ws_replicas=0, streamk="off", @@ -266,6 +267,7 @@ def _num(fields: Sequence[str], prefix: str, sep: Sequence[str] = ()) -> _Tag: _num(("num_load_waves",), "nlw"), _Tag(("cshuffle_no_alias",), "noalias"), _num(("split_k",), "sk"), + _num(("split_k_cluster",), "skc"), _Tag(("two_stage",), "2stage"), _num(("ws_replicas",), "wsr"), _num(("group_merge",), "gm"), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_abi_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_abi_emit.c index f83e5891d652..87d4cd418257 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_abi_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_abi_emit.c @@ -13,8 +13,9 @@ * 8 direct fwd 9 direct dgrad * 10 fwd problem block 11 fwd problem block 3-D * 12 deep_fused_conv_pool signature (2-D) - * 13 fwd two_stage (rejected) 14 direction "bwd" (rejected) + * 13 fwd two_stage (split-K) 14 direction "bwd" (rejected) * 15 direct wgrad 16 direct "bwd" (rejected) + * 17 fwd 3-D two_stage (split-K) 18 dgrad two_stage (rejected) */ #include #include @@ -122,6 +123,12 @@ int main(int argc, char** argv) case 16: ok = rocke_conv_direct_arg_names("bwd", &list); break; + case 17: + ok = rocke_conv_arg_names("fwd", true, true, &list); + break; + case 18: + ok = rocke_conv_arg_names("dgrad", false, true, &list); + break; default: fprintf(stderr, "unknown config index %d\n", idx); return 2; diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_abi_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_abi_emit.py index 05695c2d82c4..9d1e8c64af47 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_abi_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_abi_emit.py @@ -24,10 +24,12 @@ # 10 -- conv_fwd_problem_block() # 11 -- conv_fwd_problem_block(is_3d) # 12 -- deep_fused_conv_pool_signature(2-D spec) -# 13 -- conv_arg_names("fwd", two_stage) (rejected) +# 13 -- conv_arg_names("fwd", two_stage) (the split-K forward) # 14 -- conv_arg_names("bwd") (rejected) # 15 -- conv_direct_arg_names("wgrad") # 16 -- conv_direct_arg_names("bwd") (rejected) +# 17 -- conv_arg_names("fwd", is_3d, two_stage) +# 18 -- conv_arg_names("dgrad", two_stage) (rejected) import sys from kernels.common.conv_abi import ( @@ -67,6 +69,8 @@ def _deep_signature(): 14: lambda: conv_arg_names(direction="bwd"), 15: lambda: conv_direct_arg_names(direction="wgrad"), 16: lambda: conv_direct_arg_names(direction="bwd"), + 17: lambda: conv_arg_names(direction="fwd", is_3d=True, two_stage=True), + 18: lambda: conv_arg_names(direction="dgrad", two_stage=True), } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_fwd_split_k_reduce_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_fwd_split_k_reduce_emit.c new file mode 100644 index 000000000000..c10eea5a00ed --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_fwd_split_k_reduce_emit.c @@ -0,0 +1,151 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT + * + * tests/parity/conv_fwd_split_k_reduce_emit.c -- C-side emitter for the + * split-K forward Stage 2 (conv_fwd_split_k_reduce) parity harness. Selects one + * of N sampled spec configs by argv[1] (the config index), builds the + * rocke_fwd_split_k_reduce_spec_t identically to the Python emitter + * conv_fwd_split_k_reduce_emit.py, builds the kernel via + * rocke_build_conv_fwd_split_k_reduce_new and lowers via + * rocke_lower_kernel_to_llvm (per-config arch, flavor AUTO), printing the .ll + * to stdout so the two outputs can be byte-compared. + * + * Config index table (must stay in sync with the Python emitter): + * 0 bf16, vec 4, runtime slab loop, gfx1250 + * 1 fp16, vec 4, runtime slab loop, gfx1250 + * 2 bf16, vec 2, runtime slab loop, gfx1250 + * 3 fp16, vec 1 (scalar store), runtime slab loop, gfx1250 + * 4 bf16, vec 4, unrolled fold, degree 4, gfx1250 + * 5 fp16, vec 2, unrolled fold, degree 3, gfx1250 + * 6 bf16, vec 1, unrolled fold, degree 16, gfx1250 + * 7 fp16, vec 4, block 128, unrolled fold, degree 8, gfx1250 + * 8 bf16, vec 4, runtime slab loop, gfx950 + * 9 vec 3 (rejected by both engines) + * 10 dtype_d NULL (rejected by both engines, not defaulted) + */ +#include +#include +#include + +#include "rocke/instance_conv_fwd_split_k_reduce.h" +#include "rocke/ir.h" +#include "rocke/ir_serialize.h" +#include "rocke/lower_llvm.h" +#include "rocke/verify.h" + +/* Fill the spec for config index `idx`. Returns 0 on success, -1 if unknown. */ +static int make_cfg(int idx, rocke_fwd_split_k_reduce_spec_t* spec, const char** arch) +{ + static const struct + { + const char* dtype_d; + int vec; + int block_size; + int degree; + const char* arch; + } kCfgs[] = { + {"bf16", 4, 256, 0, "gfx1250"}, + {"fp16", 4, 256, 0, "gfx1250"}, + {"bf16", 2, 256, 0, "gfx1250"}, + {"fp16", 1, 256, 0, "gfx1250"}, + {"bf16", 4, 256, 4, "gfx1250"}, + {"fp16", 2, 256, 3, "gfx1250"}, + {"bf16", 1, 256, 16, "gfx1250"}, + {"fp16", 4, 128, 8, "gfx1250"}, + {"bf16", 4, 256, 0, "gfx950"}, + {"bf16", 3, 256, 0, "gfx1250"}, + {NULL, 4, 256, 0, "gfx1250"}, + }; + if(idx < 0 || idx >= (int)(sizeof(kCfgs) / sizeof(kCfgs[0]))) + return -1; + *spec = rocke_fwd_split_k_reduce_spec_default(); + spec->dtype_d = kCfgs[idx].dtype_d; + spec->vec = kCfgs[idx].vec; + spec->block_size = kCfgs[idx].block_size; + spec->degree = kCfgs[idx].degree; + *arch = kCfgs[idx].arch; + return 0; +} + +int main(int argc, char** argv) +{ + if(argc < 2) + { + fprintf(stderr, "usage: %s [ll|ir|verify]\n", argv[0]); + return 2; + } + int idx = atoi(argv[1]); + const char* mode = (argc > 2) ? argv[2] : "ll"; + + rocke_fwd_split_k_reduce_spec_t spec; + const char* arch = "gfx950"; + if(make_cfg(idx, &spec, &arch) != 0) + { + fprintf(stderr, "unknown config index %d\n", idx); + return 2; + } + + rocke_ir_builder_t b; + rocke_kernel_def_t* kernel = rocke_build_conv_fwd_split_k_reduce_new(&b, &spec, arch); + if(kernel == NULL) + { + const char* m = rocke_ir_builder_error(&b); + fprintf(stderr, "build failed: %s\n", m ? m : "(no message)"); + rocke_ir_builder_free(&b); + return 1; + } + + int ret = 0; + if(strcmp(mode, "ll") == 0) + { + char* llvm_text = NULL; + rocke_status_t st + = rocke_lower_kernel_to_llvm(kernel, ROCKE_LLVM_FLAVOR_AUTO, arch, &llvm_text); + if(st != ROCKE_OK || !llvm_text) + { + fprintf(stderr, "lower failed: status=%d\n", (int)st); + rocke_ir_builder_free(&b); + return 1; + } + fputs(llvm_text, stdout); + free(llvm_text); + } + else if(strcmp(mode, "ir") == 0) + { + char* ir_text = NULL; + rocke_status_t st = rocke_ir_serialize(kernel, &ir_text); + if(st != ROCKE_OK || !ir_text) + { + fprintf(stderr, "serialize failed: status=%d\n", (int)st); + rocke_ir_builder_free(&b); + return 1; + } + fputs(ir_text, stdout); + free(ir_text); + } + else if(strcmp(mode, "verify") == 0) + { + rocke_diag_t* diags = NULL; + size_t ndiag = 0; + rocke_verify(kernel, &diags, &ndiag); + for(size_t i = 0; i < ndiag; i++) + { + char* s = rocke_diag_to_string(&diags[i]); + if(s) + { + fputs(s, stdout); + fputc('\n', stdout); + free(s); + } + } + rocke_diags_free(diags, ndiag); + } + else + { + fprintf(stderr, "unknown mode '%s' (expected ll|ir|verify)\n", mode); + ret = 2; + } + + rocke_ir_builder_free(&b); + return ret; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_fwd_split_k_reduce_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_fwd_split_k_reduce_emit.py new file mode 100644 index 000000000000..883fc7507f99 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_fwd_split_k_reduce_emit.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +# +# tests/parity/conv_fwd_split_k_reduce_emit.py -- Python reference emitter for +# the split-K forward Stage 2 (conv_fwd_split_k_reduce) parity harness. Selects +# one of N sampled spec configs by argv[1], builds the FwdSplitKReduceSpec, +# builds the kernel via build_conv_fwd_split_k_reduce(spec, arch=) +# and prints the LLVM IR to stdout so it can be byte-compared with the C +# emitter conv_fwd_split_k_reduce_emit.c. +# +# Config index map: +# 0 -- bf16, vec 4, runtime slab loop, gfx1250 +# 1 -- fp16, vec 4, runtime slab loop, gfx1250 +# 2 -- bf16, vec 2, runtime slab loop, gfx1250 +# 3 -- fp16, vec 1 (scalar store), runtime slab loop, gfx1250 +# 4 -- bf16, vec 4, unrolled fold, degree 4, gfx1250 +# 5 -- fp16, vec 2, unrolled fold, degree 3, gfx1250 +# 6 -- bf16, vec 1, unrolled fold, degree 16, gfx1250 +# 7 -- fp16, vec 4, block 128, unrolled fold, degree 8, gfx1250 +# 8 -- bf16, vec 4, runtime slab loop, gfx950 +# 9 -- vec 3 (rejected by both engines) +# 10 -- dtype_d None / NULL (rejected by both engines, not defaulted) +from kernels.common.conv_fwd_split_k_reduce import ( + FwdSplitKReduceSpec, + build_conv_fwd_split_k_reduce, +) +from _emit_common import run_emit + +_CONFIGS = { + 0: (dict(dtype_d="bf16", vec=4), "gfx1250"), + 1: (dict(dtype_d="fp16", vec=4), "gfx1250"), + 2: (dict(dtype_d="bf16", vec=2), "gfx1250"), + 3: (dict(dtype_d="fp16", vec=1), "gfx1250"), + 4: (dict(dtype_d="bf16", vec=4, degree=4), "gfx1250"), + 5: (dict(dtype_d="fp16", vec=2, degree=3), "gfx1250"), + 6: (dict(dtype_d="bf16", vec=1, degree=16), "gfx1250"), + 7: (dict(dtype_d="fp16", vec=4, block_size=128, degree=8), "gfx1250"), + 8: (dict(dtype_d="bf16", vec=4), "gfx950"), + 9: (dict(dtype_d="bf16", vec=3), "gfx1250"), + 10: (dict(dtype_d=None, vec=4), "gfx1250"), +} + + +def _spec(idx: int): + if idx not in _CONFIGS: + raise SystemExit(f"unknown config index {idx}") + kw, arch = _CONFIGS[idx] + return FwdSplitKReduceSpec(**kw), arch + + +def main() -> int: + return run_emit( + _spec, + build_conv_fwd_split_k_reduce, + usage="usage: conv_fwd_split_k_reduce_emit.py \n", + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_emit.c index 04f609fd97df..0df43026bc32 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_emit.c @@ -379,6 +379,145 @@ static int make_cfg_raw(int idx, rocke_implicit_gemm_conv_spec_t* spec, const ch } } +/* Split-K forward (split_k > 1) and its single-launch cluster form + * (split_k_cluster) on gfx1250 WMMA, plus one MFMA rejection. Keep in lockstep + * with _split_k_spec() in conv_implicit_gemm_emit.py. + * 40: wavelet/cshuffle, 16x16x64 tile (small M, deep K) + * 41: mem/cshuffle, 64x64 tile, 2x2 warps + * 42: mem/default (vector_size_c=1), 32x32x32 tile + * 43: grouped (groups=2) wavelet/cshuffle + * 44: cluster S=4, mem/cshuffle, 16x16x64 tile + * 45: cluster S=8, mem/cshuffle, 64x64 tile + * 46: cluster S=15, mem/cshuffle, 16x16x64 tile + * 47: cluster S=4, grouped (groups=2), 32x32 tile, 1x2 warps + * 48: cluster S=12, mem/default, 32x32x32 tile + * 49: split_k=2 on gfx950 (MFMA): rejected by both engines + * 50: split_k=2 with a bias acc_epilogue: rejected by both engines + * 51: split_k=0: rejected by both engines + * 52: cluster S=4 with K=kpg=10 (not a multiple of 4): rejected by both + * 53: cluster S=4 on gfx1201 (no cluster launch): rejected by both engines */ +static int make_cfg_split_k(int idx, rocke_implicit_gemm_conv_spec_t* spec, const char** arch) +{ + const rocke_conv_problem_t small_m + = rocke_conv_problem_make(1, 7, 7, 256, 64, 3, 3, 1, 1, 1, 1, 1, 1); + const rocke_conv_problem_t csh_64 + = rocke_conv_problem_make(2, 12, 12, 128, 96, 3, 3, 1, 1, 1, 1, 1, 1); + const rocke_conv_problem_t dflt + = rocke_conv_problem_make(1, 9, 9, 64, 40, 3, 3, 1, 1, 1, 1, 1, 1); + rocke_conv_problem_t grouped + = rocke_conv_problem_make(1, 10, 10, 128, 64, 3, 3, 1, 1, 1, 1, 1, 1); + const rocke_conv_problem_t odd_k + = rocke_conv_problem_make(1, 7, 7, 256, 10, 3, 3, 1, 1, 1, 1, 1, 1); + grouped.groups = 2; + + *spec = rocke_implicit_gemm_conv_spec_default(); + spec->tile_m = 16; + spec->tile_n = 16; + spec->tile_k = 64; + spec->warp_m = 1; + spec->warp_n = 1; + spec->warp_tile_m = 16; + spec->warp_tile_n = 16; + spec->warp_tile_k = 32; + spec->wave_size = 32; + spec->has_vector_size_a = true; + spec->vector_size_a = 8; + spec->has_vector_size_b = true; + spec->vector_size_b = 8; + spec->has_vector_size_c = true; + spec->vector_size_c = 8; + spec->pipeline = "wavelet"; + spec->epilogue = "cshuffle"; + spec->num_load_waves = 1; + spec->split_k = 2; + *arch = "gfx1250"; + + switch(idx) + { + case 40: + spec->problem = small_m; + return 0; + case 41: + case 45: + spec->problem = csh_64; + spec->tile_m = 64; + spec->tile_n = 64; + spec->warp_m = 2; + spec->warp_n = 2; + spec->pipeline = "mem"; + spec->split_k_cluster = (idx == 45) ? 8 : 0; + return 0; + case 42: + case 48: + spec->problem = dflt; + spec->tile_m = 32; + spec->tile_n = 32; + spec->tile_k = 32; + spec->warp_m = 2; + spec->warp_n = 1; + spec->epilogue = "default"; + spec->vector_size_c = 1; + spec->pipeline = "mem"; + spec->split_k_cluster = (idx == 48) ? 12 : 0; + return 0; + case 43: + spec->problem = grouped; + return 0; + case 44: + case 46: + spec->problem = small_m; + spec->pipeline = "mem"; + spec->split_k_cluster = (idx == 44) ? 4 : 15; + return 0; + case 47: + spec->problem = grouped; + spec->tile_m = 32; + spec->tile_n = 32; + spec->warp_m = 1; + spec->warp_n = 2; + spec->pipeline = "mem"; + spec->split_k_cluster = 4; + return 0; + case 49: + spec->problem = csh_64; + spec->tile_m = 64; + spec->tile_n = 64; + spec->warp_m = 2; + spec->warp_n = 2; + spec->warp_tile_m = 32; + spec->warp_tile_n = 32; + spec->warp_tile_k = 16; + spec->wave_size = 64; + spec->pipeline = "mem"; + *arch = "gfx950"; + return 0; + case 50: + spec->problem = small_m; + spec->pipeline = "mem"; + spec->acc_epilogue.bias = 1.0; + return 0; + case 51: + spec->problem = small_m; + spec->split_k = 0; + return 0; + case 52: + spec->problem = odd_k; + spec->pipeline = "mem"; + spec->split_k_cluster = 4; + spec->vector_size_c = 2; + return 0; + case 53: + spec->problem = small_m; + spec->pipeline = "mem"; + spec->split_k_cluster = 4; + spec->warp_tile_k = 16; + *arch = "gfx1201"; + return 0; + default: + return -1; + } +} + /* The default epilogue stores one element per lane, so the validator rejects * it with the vec_c the spec would auto-derive from kpg (8 here). Configs that * leave vector_size_c unset get vector_size_c=1 -- otherwise they do not build @@ -386,7 +525,12 @@ static int make_cfg_raw(int idx, rocke_implicit_gemm_conv_spec_t* spec, const ch * conv_implicit_gemm_emit.py. */ static int make_cfg(int idx, rocke_implicit_gemm_conv_spec_t* spec, const char** arch) { - if(make_cfg_raw(idx, spec, arch) != 0) + if(idx >= 40 && idx <= 53) + { + if(make_cfg_split_k(idx, spec, arch) != 0) + return -1; + } + else if(make_cfg_raw(idx, spec, arch) != 0) return -1; if(strcmp(spec->epilogue, "default") == 0 && !spec->has_vector_size_c) { diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_emit.py index ebf26a6f7d97..0ce49528d5fe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_emit.py @@ -9,6 +9,7 @@ # lower_kernel_to_llvm(arch=) to stdout so it can be byte-compared # with the C emitter conv_implicit_gemm_emit.c. from kernels.common.conv_implicit_gemm import ( + ConvAccumulatorEpilogue, ConvProblem, ImplicitGemmConvSpec, build_implicit_gemm_conv, @@ -598,9 +599,113 @@ def _spec_raw(idx: int): ), "gfx1250", ) + if 40 <= idx <= 53: + return _split_k_spec(idx) raise SystemExit(f"unknown config index {idx}") +# Split-K forward (split_k > 1) and its single-launch cluster form +# (split_k_cluster) on gfx1250 WMMA, plus one MFMA rejection. Keep in lockstep +# with make_cfg_split_k() in conv_implicit_gemm_emit.c. +# 40: wavelet/cshuffle, 16x16x64 tile (small M, deep K) +# 41: mem/cshuffle, 64x64 tile, 2x2 warps +# 42: mem/default (vector_size_c=1), 32x32x32 tile +# 43: grouped (groups=2) wavelet/cshuffle +# 44: cluster S=4, mem/cshuffle, 16x16x64 tile +# 45: cluster S=8, mem/cshuffle, 64x64 tile +# 46: cluster S=15, mem/cshuffle, 16x16x64 tile +# 47: cluster S=4, grouped (groups=2), 32x32 tile, 1x2 warps +# 48: cluster S=12, mem/default, 32x32x32 tile +# 49: split_k=2 on gfx950 (MFMA): rejected by both engines +# 50: split_k=2 with a bias acc_epilogue: rejected by both engines +# 51: split_k=0: rejected by both engines +# 52: cluster S=4 with K=kpg=10 (not a multiple of 4): rejected by both +# 53: cluster S=4 on gfx1201 (no cluster launch): rejected by both engines +_SMALL_M = dict(N=1, Hi=7, Wi=7, C=256, K=64, fy=3, fx=3, pH=1, pW=1) +_CSH_64 = dict(N=2, Hi=12, Wi=12, C=128, K=96, fy=3, fx=3, pH=1, pW=1) +_DEFAULT = dict(N=1, Hi=9, Wi=9, C=64, K=40, fy=3, fx=3, pH=1, pW=1) +_GROUPED = dict(N=1, Hi=10, Wi=10, C=128, K=64, fy=3, fx=3, pH=1, pW=1, groups=2) +_ODD_K = dict(N=1, Hi=7, Wi=7, C=256, K=10, fy=3, fx=3, pH=1, pW=1) + + +def _split_k_spec(idx: int): + base = dict( + tile_m=16, + tile_n=16, + tile_k=64, + warp_m=1, + warp_n=1, + warp_tile_m=16, + warp_tile_n=16, + warp_tile_k=32, + wave_size=32, + vector_size_a=8, + vector_size_b=8, + vector_size_c=8, + pipeline="wavelet", + epilogue="cshuffle", + num_load_waves=1, + split_k=2, + ) + tile_64 = dict(tile_m=64, tile_n=64, warp_m=2, warp_n=2) + tile_def = dict( + tile_m=32, + tile_n=32, + tile_k=32, + warp_m=2, + warp_n=1, + epilogue="default", + vector_size_c=1, + ) + prob, kw = { + 40: (_SMALL_M, {}), + 41: (_CSH_64, dict(tile_64, pipeline="mem")), + 42: (_DEFAULT, dict(tile_def, pipeline="mem")), + 43: (_GROUPED, {}), + 44: (_SMALL_M, dict(pipeline="mem", split_k_cluster=4)), + 45: (_CSH_64, dict(tile_64, pipeline="mem", split_k_cluster=8)), + 46: (_SMALL_M, dict(pipeline="mem", split_k_cluster=15)), + 47: ( + _GROUPED, + dict( + tile_m=32, + tile_n=32, + warp_m=1, + warp_n=2, + pipeline="mem", + split_k_cluster=4, + ), + ), + 48: (_DEFAULT, dict(tile_def, pipeline="mem", split_k_cluster=12)), + 49: ( + _CSH_64, + dict( + tile_m=64, + tile_n=64, + warp_m=2, + warp_n=2, + warp_tile_m=32, + warp_tile_n=32, + warp_tile_k=16, + wave_size=64, + pipeline="mem", + ), + ), + 50: ( + _SMALL_M, + dict(pipeline="mem", acc_epilogue=ConvAccumulatorEpilogue(bias=1.0)), + ), + 51: (_SMALL_M, dict(split_k=0)), + 52: (_ODD_K, dict(pipeline="mem", split_k_cluster=4, vector_size_c=2)), + 53: (_SMALL_M, dict(pipeline="mem", split_k_cluster=4, warp_tile_k=16)), + }[idx] + cfg = dict(base, **kw) + return ( + ImplicitGemmConvSpec(problem=_cp(**prob), **cfg), + {49: "gfx950", 53: "gfx1201"}.get(idx, "gfx1250"), + ) + + def main() -> int: return run_emit( _spec, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_builder_signature_contract.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_builder_signature_contract.py index 0e1da17ff0f0..b8539da5f34f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_builder_signature_contract.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_builder_signature_contract.py @@ -312,6 +312,7 @@ "kernels.common.conv_implicit_gemm.ImplicitGemmConvSpec": ("problem",), "kernels.common.conv_implicit_gemm_wgrad.WgradConvSpec": ("problem",), "kernels.common.conv_wgrad_workspace_reduce.WgradReduceSpec": ("problem",), + "kernels.common.conv_fwd_split_k_reduce.FwdSplitKReduceSpec": (), "kernels.common.conv_implicit_gemm_dgrad.DgradConvSpec": ("problem",), "kernels.common.conv_direct_grouped.DirectConvProblem": ( "H", diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_dgrad_patch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_dgrad_patch.py index 822b02fd9bd4..9732c1e62bae 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_dgrad_patch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_dgrad_patch.py @@ -178,6 +178,86 @@ def dgrad(p, dy, w): torch.testing.assert_close(got.reshape(want.shape), want) +class TestForwardWidthFold(unittest.TestCase): + """The fold is direction-free: a forward runs folded on the same tensors.""" + + @unittest.skipUnless(_HAS_TORCH, "torch is optional") + def test_folded_forward_is_the_same_numbers(self): + import torch + + from kernels.common.conv_implicit_gemm_dgrad import fold_width + + def fwd(p, x, w): + # NHWC X / KYXC W -> NHWC Y, through torch's NCHW reference. + y = torch.nn.functional.conv2d( + x.permute(0, 3, 1, 2), + w.permute(0, 3, 1, 2), + stride=(p.sH, p.sW), + padding=(p.pH, p.pW), + ) + return y.permute(0, 2, 3, 1).contiguous() + + torch.manual_seed(0) + for kw in ( + # ViT patch embedding, ConvNeXt stem, a 2x2 stride-2 downsample, + # and a padded strided case the fold still covers. + { + "N": 2, + "Hi": 32, + "Wi": 32, + "C": 3, + "K": 8, + "Y": 16, + "X": 16, + "sH": 16, + "sW": 16, + }, + {"N": 2, "Hi": 12, "Wi": 12, "C": 3, "K": 8}, + { + "N": 1, + "Hi": 8, + "Wi": 8, + "C": 16, + "K": 8, + "Y": 2, + "X": 2, + "sH": 2, + "sW": 2, + }, + { + "N": 2, + "Hi": 11, + "Wi": 16, + "C": 3, + "K": 4, + "Y": 3, + "X": 4, + "sH": 2, + "sW": 2, + "pH": 1, + "pW": 2, + }, + ): + p = _problem(**kw) + q = fold_width(p) + self.assertNotEqual(q, p) + x = torch.randn(p.N, p.Hi, p.Wi, p.C, dtype=torch.float64) + w = torch.randn(p.K, p.Y, p.X, p.C, dtype=torch.float64) + want = fwd(p, x, w) + got = fwd(q, x.reshape(q.N, q.Hi, q.Wi, q.C), w.reshape(q.K, q.Y, q.X, q.C)) + torch.testing.assert_close(got, want) + + def test_benchmark_launches_forward_folded(self): + from benchmarks.common.benchmark_implicit_gemm_conv import _launch_problem + from kernels.common.conv_implicit_gemm_dgrad import fold_width + + patch = _problem(Wi=224, Hi=224, X=16, Y=16, sW=16, sH=16) + self.assertEqual(_launch_problem("fwd", patch), fold_width(patch)) + plain = _problem(X=3, Y=3, sW=1, sH=1, pW=1, pH=1) + self.assertIs(_launch_problem("fwd", plain), plain) + self.assertIs(_launch_problem("wgrad", patch), patch) + + class TestPatchRecord(unittest.TestCase): def test_is_patch_problem(self): from kernels.common.conv_implicit_gemm_dgrad import is_patch_problem diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_dgrad_spec_policy.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_dgrad_spec_policy.py index 5a88c15b114c..d72c725d849c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_dgrad_spec_policy.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_dgrad_spec_policy.py @@ -351,6 +351,7 @@ def test_off_keeps_the_hash_on_is_distinct(self): "dy_halo_setprio", "dy_halo_kouter_pad", "halo_max_wo", + "split_k_cluster", ): del fields[name] del fields["streamk"], fields["streamk_reduction"] diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_fwd_split_k.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_fwd_split_k.py new file mode 100644 index 000000000000..9712f1dfb51c --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_fwd_split_k.py @@ -0,0 +1,1120 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Split-K forward implicit-GEMM conv (``ImplicitGemmConvSpec.split_k``). + +CPU: the validator gate, the ABI, that ``split_k = 1`` builds the same +kernel as before, and the auto (``-1``) degree policy and its resolver. GPU +(gfx1250): Stage 1 + the Stage 2 reduce against the +reference across degrees, tiles, pipelines, epilogues, groups, strides, a +partial last K tile and degrees past the K-tile count (empty slices). +""" + +from __future__ import annotations + +import ctypes +import importlib.util +import os +import unittest +from unittest import mock + +from kernels.common._conv_implicit_gemm_common import ConvDataSpec, ConvProblem +from kernels.common.conv_abi import conv_arg_names +from kernels.common.conv_implicit_gemm import ( + ImplicitGemmConvSpec, + build_implicit_gemm_conv, + is_valid_spec, +) + +_HAS_TORCH = importlib.util.find_spec("torch") is not None + +_WMMA = dict( + warp_tile_m=16, + warp_tile_n=16, + warp_tile_k=32, + wave_size=32, + vector_size_a=8, + vector_size_b=8, +) +_PROBE = ConvProblem(N=1, Hi=16, Wi=16, C=128, K=128, Y=3, X=3, pH=1, pW=1, groups=2) +_BF16 = ConvDataSpec("bf16", "bf16", "bf16") + + +def _spec(problem=_PROBE, data=_BF16, **kw): + cfg = dict( + _WMMA, + tile_m=16, + tile_n=16, + tile_k=64, + warp_m=1, + warp_n=1, + pipeline="wavelet", + epilogue="cshuffle", + vector_size_c=8, + num_load_waves=1, + ) + cfg.update(kw) + return ImplicitGemmConvSpec(problem=problem, data=data, **cfg) + + +class TestSplitKGate(unittest.TestCase): + def test_wmma_split_k_is_valid(self): + ok, why = is_valid_spec(_spec(split_k=2), "gfx1250") + self.assertTrue(ok, why) + + def test_mfma_is_rejected(self): + spec = _spec( + split_k=2, + wave_size=64, + pipeline="mem", + warp_tile_m=32, + warp_tile_n=32, + warp_tile_k=16, + tile_m=64, + tile_n=64, + ) + ok, why = is_valid_spec(spec, "gfx950") + self.assertFalse(ok) + self.assertIn("WMMA-only", why) + + def test_group_merge_is_rejected(self): + dw = ConvProblem(N=1, Hi=16, Wi=16, C=64, K=64, Y=3, X=3, pH=1, pW=1, groups=64) + ok, why = is_valid_spec( + _spec( + problem=dw, + split_k=2, + group_merge=4, + vector_size_a=None, + vector_size_b=None, + vector_size_c=None, + ), + "gfx1250", + ) + self.assertFalse(ok) + + def test_non_identity_acc_epilogue_is_rejected(self): + # Stage 1 stores raw partials and the reduction only adds them: a bias + # would be added once per slice, ReLU / clamp applied before the sum. + from kernels.common.conv_implicit_gemm import ConvAccumulatorEpilogue + + for epi in ( + ConvAccumulatorEpilogue(bias=1.0), + ConvAccumulatorEpilogue(scale=2.0), + ConvAccumulatorEpilogue(relu=True), + ConvAccumulatorEpilogue(clamp_min=-1.0, clamp_max=1.0), + ): + for kw in ( + dict(split_k=2), + dict(split_k=2, split_k_cluster=4, pipeline="mem"), + ): + with self.subTest(epi=epi, **kw): + ok, why = is_valid_spec(_spec(acc_epilogue=epi, **kw), "gfx1250") + self.assertFalse(ok) + self.assertIn("identity acc_epilogue", why) + # The unsplit kernel keeps its fused epilogue. + ok, why = is_valid_spec( + _spec(acc_epilogue=ConvAccumulatorEpilogue(relu=True)), "gfx1250" + ) + self.assertTrue(ok, why) + + def test_nonpositive_split_k_is_rejected(self): + for sk in (0, -1, -4): + with self.subTest(split_k=sk): + ok, why = is_valid_spec(_spec(split_k=sk), "gfx1250") + self.assertFalse(ok) + self.assertIn("split_k must be >= 1", why) + + def test_name_tags_the_variant(self): + self.assertTrue(_spec(split_k=2).kernel_name().endswith("_splitk")) + self.assertNotIn("splitk", _spec().kernel_name()) + + def test_abi_appends_workspace_and_slice(self): + base = [n for n, _ in conv_arg_names(direction="fwd")] + split = [n for n, _ in conv_arg_names(direction="fwd", two_stage=True)] + self.assertEqual(split[: len(base)], base) + self.assertEqual(split[len(base) :], ["ws_ptr", "ws_bytes", "ks", "ks_count"]) + + def test_kernel_params_follow_the_abi(self): + kernel = build_implicit_gemm_conv(_spec(split_k=2), arch="gfx1250") + names = [p.name for p in kernel.params] + abi = [n for n, _ in conv_arg_names(direction="fwd", two_stage=True)] + self.assertEqual(names, abi) + + +class TestClusterSplitKGate(unittest.TestCase): + """``split_k_cluster``: the single-launch fold through a workgroup cluster.""" + + def test_mem_cluster_is_valid_and_sets_the_cluster(self): + spec = _spec(split_k=2, split_k_cluster=4, pipeline="mem") + ok, why = is_valid_spec(spec, "gfx1250") + self.assertTrue(ok, why) + kernel = build_implicit_gemm_conv(spec, arch="gfx1250") + self.assertEqual(kernel.cluster_dims, (1, 1, 4)) + self.assertTrue(spec.kernel_name().endswith("_splitk_skc4")) + # The ABI is the split-K one: the degree still arrives as ks_count. + abi = [n for n, _ in conv_arg_names(direction="fwd", two_stage=True)] + self.assertEqual([p.name for p in kernel.params], abi) + + def test_rejections(self): + for kw, why in ( + (dict(split_k=1, split_k_cluster=4, pipeline="mem"), "split_k > 1"), + (dict(split_k=2, split_k_cluster=4), "mem pipeline"), + (dict(split_k=2, split_k_cluster=16, pipeline="mem"), "2..15"), + ): + ok, reason = is_valid_spec(_spec(**kw), "gfx1250") + self.assertFalse(ok, kw) + self.assertIn(why, reason) + + def test_off_by_default(self): + kernel = build_implicit_gemm_conv(_spec(split_k=2), arch="gfx1250") + self.assertIsNone(kernel.cluster_dims) + + def test_target_without_cluster_launch_is_rejected(self): + spec = _spec(split_k=2, split_k_cluster=4, pipeline="mem", warp_tile_k=16) + ok, why = is_valid_spec(spec, "gfx1201") + self.assertFalse(ok) + self.assertIn("workgroup-cluster launch", why) + # The two-stage split-K form needs no cluster. + ok, why = is_valid_spec(_spec(split_k=2, warp_tile_k=16), "gfx1201") + self.assertTrue(ok, why) + + def test_channels_not_a_multiple_of_4_are_rejected(self): + # The fold reads 4-channel quads bounded by their first channel: with + # kpg = 10 the quad at channel 8 would run past the row. + for problem in ( + ConvProblem(N=1, Hi=7, Wi=7, C=256, K=10, Y=3, X=3, pH=1, pW=1), + ConvProblem(N=1, Hi=7, Wi=7, C=128, K=20, Y=3, X=3, pH=1, pW=1, groups=2), + ): + with self.subTest(K=problem.K, groups=problem.groups): + spec = _spec( + problem=problem, + split_k=2, + split_k_cluster=4, + pipeline="mem", + vector_size_c=2, + ) + ok, why = is_valid_spec(spec, "gfx1250") + self.assertFalse(ok) + self.assertIn("multiples of 4", why) + + def test_launch_admission_checks_the_launched_problem(self): + # The binary is shape-generic: a spec built on a good problem must + # still refuse a launched one whose channels break the quad rule. + from kernels.common.conv_implicit_gemm import is_valid_spec_for_problem + from kernels.conv.kernel_identity import ( + KernelIdentity, + identity_supports_problem, + ) + + spec = _spec(split_k=2, split_k_cluster=4, pipeline="mem", vector_size_c=2) + good = ConvProblem(N=1, Hi=7, Wi=7, C=256, K=12, Y=3, X=3, pH=1, pW=1) + odd = ConvProblem(N=1, Hi=7, Wi=7, C=256, K=10, Y=3, X=3, pH=1, pW=1) + grouped_odd = ConvProblem( + N=1, Hi=7, Wi=7, C=256, K=12, Y=3, X=3, pH=1, pW=1, groups=2 + ) + self.assertEqual(is_valid_spec_for_problem(spec, good, "gfx1250"), (True, "ok")) + for problem in (odd, grouped_odd): + ok, why = is_valid_spec_for_problem(spec, problem, "gfx1250") + self.assertFalse(ok) + self.assertIn("multiples of 4", why) + + ident = KernelIdentity.from_label( + "fwd-implicit_gemm-bf16-2d-t16x16x64-w1x1-a16x16x32-v8x8x2-mem-" + "cshuffle-grp-sk2-skc4-2stage-wave32-gfx1250-llvm23" + ) + self.assertTrue(identity_supports_problem(ident, good, "llvm23")[0]) + for problem in (odd, grouped_odd): + ok, why = identity_supports_problem(ident, problem, "llvm23") + self.assertFalse(ok) + self.assertIn("multiples of 4", why) + + +class TestClusterSplitKCache(unittest.TestCase): + """The AOT cache carries one cluster binary per degree, named apart.""" + + def test_identity_keeps_its_hash_at_zero(self): + from dataclasses import asdict + + from kernels.conv.kernel_identity import KernelIdentity + + ident = KernelIdentity.from_label( + "fwd-implicit_gemm-bf16-2d-t32x32x64-w2x1-a16x16x32-v8x8x8-mem-" + "cshuffle-grp-sk2-2stage-wave32-gfx1250-llvm23" + ) + self.assertEqual(ident.split_k_cluster, 0) + # An identity written before the field existed hashes the same. + old = asdict(ident) + del old["split_k_cluster"] + self.assertEqual( + KernelIdentity.from_dict(old).stable_hash(), ident.stable_hash() + ) + + def test_only_the_two_stage_form_needs_stage2(self): + # Both split-K forward forms take the workspace ABI (two_stage); only + # the two-stage one is followed by the Stage 2 reduce launch. + from kernels.conv.kernel_identity import KernelIdentity + + prefix = "fwd-implicit_gemm-bf16-2d-t32x32x64-w2x1-a16x16x32-v8x8x8-mem-" + two_stage = KernelIdentity.from_label( + prefix + "cshuffle-grp-sk2-2stage-wave32-gfx1250-llvm23" + ) + cluster = KernelIdentity.from_label( + prefix + "cshuffle-grp-sk2-skc4-2stage-wave32-gfx1250-llvm23" + ) + unsplit = KernelIdentity.from_label( + prefix + "cshuffle-grp-wave32-gfx1250-llvm23" + ) + self.assertTrue(two_stage.two_stage and two_stage.needs_stage2) + self.assertTrue(cluster.two_stage) + self.assertFalse(cluster.needs_stage2) + self.assertFalse(unsplit.needs_stage2) + + def test_grid_has_one_binary_per_degree(self): + from rocke.core.arch import ArchTarget + + from benchmarks.common.kernel_sweep import enumerate_jobs + from kernels.conv.aot_jobs import ( + CACHE_FWD_SPLIT_K_CLUSTER_DEGREES, + CACHE_FWD_SPLIT_K_CLUSTER_TILES, + ) + from kernels.conv.kernel_identity import KernelIdentity + + # Cluster jobs exist only at the flavor that lowers clusters, so the + # grid is enumerated there, whatever the host's own flavor is. + with mock.patch.dict(os.environ, {"ROCKE_LLVM_FLAVOR": "llvm23"}): + jobs = [ + j + for j in enumerate_jobs( + arch="gfx1250", + dtype="bf16", + target=ArchTarget.from_gfx("gfx1250"), + directions=["fwd"], + ) + if j.identity.split_k_cluster + ] + self.assertEqual( + len(jobs), + len(CACHE_FWD_SPLIT_K_CLUSTER_DEGREES) + * len(CACHE_FWD_SPLIT_K_CLUSTER_TILES), + ) + self.assertEqual(len({j.identity.stable_hash() for j in jobs}), len(jobs)) + for j in jobs: + label = j.identity.label() + self.assertIn(f"skc{j.identity.split_k_cluster}", label) + self.assertEqual(KernelIdentity.from_label(label), j.identity) + + def test_launch_gate(self): + from benchmarks.common.benchmark_implicit_gemm_conv import _runtime_split_ks + from kernels.conv.kernel_identity import KernelIdentity + + ident = KernelIdentity.from_label( + "fwd-implicit_gemm-bf16-2d-t64x32x64-w2x1-a16x16x32-v8x8x8-mem-" + "cshuffle-grp-sk2-skc4-2stage-wave32-gfx1250-llvm23" + ) + small = ConvProblem(N=1, Hi=7, Wi=7, C=512, K=512, Y=3, X=3, pH=1, pW=1) + many_tiles = ConvProblem( + N=64, Hi=7, Wi=7, C=1024, K=1024, Y=3, X=3, pH=1, pW=1, groups=32 + ) + self.assertEqual(_runtime_split_ks(small, ident), (4,)) + self.assertEqual(_runtime_split_ks(small, ident, 8), ()) + self.assertEqual(_runtime_split_ks(many_tiles, ident), ()) + # The fold's quad rule holds at launch too. + odd = ConvProblem(N=1, Hi=7, Wi=7, C=512, K=510, Y=3, X=3, pH=1, pW=1) + self.assertEqual(_runtime_split_ks(odd, ident), ()) + + +class TestSplitKDegrees(unittest.TestCase): + """``--split-k`` / ``--split-k-degrees`` on a cached split-K forward binary.""" + + _IDENT = ( + "fwd-implicit_gemm-bf16-2d-t16x16x64-w1x1-a16x16x32-v8x8x8-wavelet-" + "cshuffle-grp-sk2-2stage-wave32-gfx1250-llvm23" + ) + # 9*512 = 4608 = 72 K tiles of 64: every degree below fits. 4 x 4 output + # tiles, so the policy splits as far as it goes (32). + _SMALL = ConvProblem(N=1, Hi=7, Wi=7, C=512, K=64, Y=3, X=3, pH=1, pW=1) + # 8*28*28/16 * 64/16 = 1568 output tiles of 16 x 16 fill the device unsplit. + _FULL = ConvProblem(N=8, Hi=28, Wi=28, C=64, K=64, Y=3, X=3, pH=1, pW=1) + + def test_explicit_sweep_uses_the_given_degrees(self): + from benchmarks.common.benchmark_implicit_gemm_conv import _runtime_split_ks + from kernels.conv.kernel_identity import KernelIdentity + + ident = KernelIdentity.from_label(self._IDENT) + small = self._SMALL + # 0 sweeps the list plus the resolver's pick, as for wgrad. + self.assertEqual(_runtime_split_ks(small, ident, 0, (3,)), (32, 3)) + self.assertEqual(_runtime_split_ks(small, ident, 0, (5, 3)), (32, 5, 3)) + self.assertEqual(_runtime_split_ks(small, ident, 6, (3,)), (6,)) + + def test_auto_launches_the_resolved_degree(self): + from benchmarks.common.benchmark_implicit_gemm_conv import _runtime_split_ks + from kernels.common.conv_split_k import resolve_conv_split_k + from kernels.conv.kernel_identity import KernelIdentity + + ident = KernelIdentity.from_label(self._IDENT) + resolved = resolve_conv_split_k( + "fwd", + self._SMALL, + arch="gfx1250", + dtype="bf16", + tile_m=16, + tile_n=16, + tile_k=64, + block_size=32, + ) + self.assertEqual(resolved.split_k, 32) + # Whatever the list says: -1 is the resolver's pick only. + self.assertEqual(_runtime_split_ks(self._SMALL, ident, -1, (3,)), (32,)) + + def test_auto_skips_the_split_binary_where_splitting_does_not_pay(self): + from benchmarks.common.benchmark_implicit_gemm_conv import _runtime_split_ks + from kernels.conv.kernel_identity import KernelIdentity + + ident = KernelIdentity.from_label(self._IDENT) + big = TestSplitKDegrees._FULL + self.assertEqual(_runtime_split_ks(big, ident, -1), ()) + # An explicit degree is still the caller's. + self.assertEqual(_runtime_split_ks(big, ident, 2), (2,)) + + def test_the_override_variable_forces_the_auto_degree(self): + from benchmarks.common.benchmark_implicit_gemm_conv import _runtime_split_ks + from kernels.conv.kernel_identity import KernelIdentity + + ident = KernelIdentity.from_label(self._IDENT) + with mock.patch.dict(os.environ, {"ROCKE_CONV_SPLIT_K": "4"}): + self.assertEqual(_runtime_split_ks(self._SMALL, ident, -1), (4,)) + with mock.patch.dict(os.environ, {"ROCKE_CONV_SPLIT_K": "off"}): + self.assertEqual(_runtime_split_ks(self._SMALL, ident, -1), ()) + + +class TestSplitKPolicy(unittest.TestCase): + """``select_split_k_fwd``: the auto forward degree from the GEMM, the tile + and the arch.""" + + def _pick(self, M, N, K, tile=(16, 16, 64), groups=1, arch="gfx1250"): + from rocke.helpers.split_k import select_split_k_fwd + + tm, tn, tk = tile + return select_split_k_fwd( + M=M, N=N, K=K, tile_m=tm, tile_n=tn, tile_k=tk, groups=groups, arch=arch + ).split_k + + def test_fills_toward_the_target_grid(self): + from rocke.helpers.split_k import FWD_SPLIT_K_TARGET_CTAS_PER_CU + + target = 256 * FWD_SPLIT_K_TARGET_CTAS_PER_CU + # Deep enough that only the grid limits the degree. + for tiles in (1, 3, 64, 100, 255, 256, target // 2 - 1): + with self.subTest(tiles=tiles): + got = self._pick(16 * tiles, 16, 1 << 20) + want = min(32, 1 << ((target // tiles).bit_length() - 1)) + self.assertEqual(got, want) + self.assertLessEqual(got * tiles, target) + + def test_a_grid_that_fills_the_device_is_not_split(self): + from rocke.helpers.split_k import FWD_SPLIT_K_TARGET_CTAS_PER_CU + + target = 256 * FWD_SPLIT_K_TARGET_CTAS_PER_CU + self.assertEqual(self._pick(16 * (target // 2), 16, 1 << 20), 1) + # Groups multiply the grid. + self.assertEqual(self._pick(16 * 8, 16, 1 << 20, groups=target // 16), 1) + self.assertGreater(self._pick(16 * 8, 16, 1 << 20, groups=2), 1) + + def test_shallow_reductions_are_not_split(self): + from rocke.helpers.split_k import FWD_SPLIT_K_MIN_GEMM_K + + self.assertEqual(self._pick(16, 16, FWD_SPLIT_K_MIN_GEMM_K - 1), 1) + self.assertGreater(self._pick(16, 16, FWD_SPLIT_K_MIN_GEMM_K), 1) + + def test_every_slice_keeps_two_k_tiles(self): + for k_tiles in (8, 9, 15, 16, 33, 64, 65): + with self.subTest(k_tiles=k_tiles): + got = self._pick(16, 16, 64 * k_tiles) + self.assertGreaterEqual(k_tiles // got, 2) + self.assertEqual(got & (got - 1), 0) + # The same reduction in narrower K tiles splits further. + self.assertGreater( + self._pick(16, 16, 640, (16, 16, 32)), self._pick(16, 16, 640) + ) + + def test_a_smaller_part_splits_less(self): + from rocke.helpers.split_k import select_split_k_fwd + + kw = dict(M=16 * 64, N=16, K=1 << 20, tile_m=16, tile_n=16, tile_k=64) + full = select_split_k_fwd(**kw).split_k + self.assertEqual(select_split_k_fwd(**kw, num_cus=128).split_k, full // 2) + + def test_degrees_stay_on_the_ladder_and_under_the_cap(self): + from rocke.helpers.split_k import FWD_SPLIT_K_MAX + + for M in (1, 49, 196, 784, 3136): + for K in (512, 2304, 4608, 17280): + with self.subTest(M=M, K=K): + got = self._pick(M, 256, K, (32, 32, 64)) + self.assertEqual(got & (got - 1), 0) + self.assertLessEqual(got, FWD_SPLIT_K_MAX) + + +class TestSplitKResolver(unittest.TestCase): + """``resolve_conv_split_k("fwd", ...)``: where the forward splits, the + explicit and override paths, and the launch limits.""" + + def _resolve(self, problem, arch="gfx1250", tile=(16, 16, 64), **kw): + from kernels.common.conv_split_k import resolve_conv_split_k + + tm, tn, tk = tile + return resolve_conv_split_k( + "fwd", + problem, + arch=arch, + dtype="bf16", + tile_m=tm, + tile_n=tn, + tile_k=tk, + block_size=32, + **kw, + ) + + def test_auto_follows_the_policy_on_gfx1250(self): + from rocke.helpers.split_k import select_split_k_fwd + + p = TestSplitKDegrees._SMALL + got = self._resolve(p) + want = select_split_k_fwd( + M=p.M, N=p.kpg, K=p.K_gemm, tile_m=16, tile_n=16, tile_k=64 + ).split_k + self.assertEqual((got.split_k, got.two_stage), (want, True)) + self.assertGreater(want, 1) + + def test_no_split_where_the_kernel_is_not_built(self): + p = TestSplitKDegrees._SMALL + for arch in ("gfx950", "gfx942", "gfx1201"): + with self.subTest(arch=arch): + r = self._resolve(p, arch=arch) + self.assertEqual((r.split_k, r.two_stage), (1, False)) + # Even when forced: the kernel cannot be built there. + self.assertEqual(self._resolve(p, arch=arch, override=8).split_k, 1) + pointwise = ConvProblem(N=1, Hi=7, Wi=7, C=4096, K=64, Y=1, X=1) + self.assertEqual(self._resolve(pointwise).split_k, 1) + three_d = ConvProblem( + N=1, + Hi=4, + Wi=4, + C=512, + K=64, + Y=3, + X=3, + pH=1, + pW=1, + Di=4, + Z=3, + sD=1, + pD=1, + dD=1, + ) + self.assertEqual(self._resolve(three_d).split_k, 1) + + def test_explicit_degrees_are_kept(self): + p = TestSplitKDegrees._SMALL + for sk in (1, 2, 3, 64): + with self.subTest(split_k=sk): + r = self._resolve(p, split_k=sk) + self.assertEqual((r.split_k, r.two_stage), (sk, sk > 1)) + + def test_override_is_held_to_the_launch_limits(self): + p = TestSplitKDegrees._SMALL + self.assertEqual(self._resolve(p, override=1).split_k, 1) + self.assertEqual(self._resolve(p, override=6).split_k, 6) + # gridDim.z carries group * split_k + slice. + grouped = ConvProblem( + N=1, Hi=7, Wi=7, C=40000, K=40000, Y=3, X=3, pH=1, pW=1, groups=10000 + ) + self.assertEqual(self._resolve(grouped, override=8).split_k, 6) + # The fp32 slabs are addressed through an i32 byte count. + from kernels.common.conv_fwd_split_k_reduce import fwd_split_k_workspace_nbytes + from kernels.common.conv_split_k import MAX_FWD_WS_BYTES + + wide = ConvProblem(N=16, Hi=128, Wi=128, C=64, K=256, Y=3, X=3, pH=1, pW=1) + got = self._resolve(wide, override=64).split_k + self.assertLess(got, 64) + self.assertLessEqual( + fwd_split_k_workspace_nbytes(wide.M, wide.K, got), MAX_FWD_WS_BYTES + ) + self.assertGreater( + fwd_split_k_workspace_nbytes(wide.M, wide.K, 2 * got), MAX_FWD_WS_BYTES + ) + + def test_a_split_k_binary_holder_gets_an_error_not_degree_1(self): + big = TestSplitKDegrees._FULL + self.assertEqual(self._resolve(big).split_k, 1) + with self.assertRaisesRegex(ValueError, "does not serve this problem"): + self._resolve(big, allow_unsplit=False) + + def test_the_cache_builds_where_the_resolver_splits(self): + from kernels.common.conv_split_k import FWD_SPLIT_K_ARCHS + from kernels.conv.aot_jobs import CACHE_FWD_SPLIT_K_ARCHS + + self.assertEqual(CACHE_FWD_SPLIT_K_ARCHS, FWD_SPLIT_K_ARCHS) + + +class TestSplitKAutoBuild(unittest.TestCase): + """A spec with ``split_k=-1`` builds the binary its resolved degree + selects, in both engines' lowering entry.""" + + def test_auto_builds_the_split_binary_where_it_pays(self): + # 128 output tiles over 9 K tiles: the policy splits. + auto = build_implicit_gemm_conv(_spec(split_k=-1), arch="gfx1250") + split = build_implicit_gemm_conv(_spec(split_k=2), arch="gfx1250") + self.assertEqual(auto.name, split.name) + self.assertIn("ks_count", [p.name for p in auto.params]) + + def test_auto_builds_the_plain_kernel_where_it_does_not(self): + big = TestSplitKDegrees._FULL + auto = build_implicit_gemm_conv(_spec(problem=big, split_k=-1), arch="gfx1250") + plain = build_implicit_gemm_conv(_spec(problem=big), arch="gfx1250") + self.assertEqual(auto.name, plain.name) + self.assertNotIn("ks_count", [p.name for p in auto.params]) + + def test_auto_with_a_spec_the_split_cannot_take_builds_the_plain_kernel(self): + # The problem would split (see above), but a fused accumulator + # epilogue keeps the spec out of the split form: auto resolves to 1 + # and builds the plain kernel with its epilogue, not an error. + from kernels.common.conv_implicit_gemm import ( + ConvAccumulatorEpilogue, + resolve_fwd_split_k_spec, + ) + + for kw in ( + dict(acc_epilogue=ConvAccumulatorEpilogue(relu=True)), + dict(acc_epilogue=ConvAccumulatorEpilogue(bias=1.0)), + dict(acc_epilogue=ConvAccumulatorEpilogue(clamp_min=-1.0, clamp_max=1.0)), + dict(acc_epilogue=ConvAccumulatorEpilogue(scale=2.0)), + ): + with self.subTest(**kw): + spec = _spec(split_k=-1, **kw) + self.assertEqual(resolve_fwd_split_k_spec(spec, "gfx1250").split_k, 1) + epi = dict(acc_epilogue=ConvAccumulatorEpilogue(relu=True)) + auto = build_implicit_gemm_conv(_spec(split_k=-1, **epi), arch="gfx1250") + plain = build_implicit_gemm_conv(_spec(**epi), arch="gfx1250") + self.assertEqual(auto.name, plain.name) + self.assertNotIn("ks_count", [p.name for p in auto.params]) + # The identity epilogue still splits. + self.assertEqual( + resolve_fwd_split_k_spec(_spec(split_k=-1), "gfx1250").split_k, 4 + ) + + def test_the_backend_resolves_auto_for_the_cpp_engine(self): + from rocke.core import backend + + seen = [] + + def _fake_lower_family(family, spec, *a, **kw): + seen.append(spec.split_k) + + with mock.patch.object( + backend, "_lower_family", _fake_lower_family + ), mock.patch.object(backend, "_import_engine", lambda: None): + backend.lower_conv_implicit_gemm(_spec(split_k=-1), arch="gfx1250") + backend.lower_conv_implicit_gemm(_spec(split_k=-1), arch="gfx950") + self.assertEqual(seen, [4, 1]) + + +class TestSplitKAotScope(unittest.TestCase): + """The AOT grid builds split-K forward jobs only where they are supported: + gfx1250, and the cluster ones only at the flavor that lowers clusters. + ``count_jobs`` agrees with the walk in every case.""" + + def _walk(self, arch, flavor): + from rocke.core.arch import ArchTarget + + from kernels.conv import aot_jobs as aj + + target = ArchTarget.from_gfx(arch) + jobs = list( + aj._fwd_jobs( + arch, "bf16", target.wave_size, "wmma", target, llvm_flavor=flavor + ) + ) + counted = aj.count_jobs( + dtype="bf16", target=target, directions=["fwd"], llvm_flavor=flavor + ) + split = [j for j in jobs if j.identity.split_k > 1] + cluster = [j for j in split if j.identity.split_k_cluster] + return len(jobs), counted, len(split), len(cluster) + + def test_gfx1250_llvm23_has_both(self): + walked, counted, split, cluster = self._walk("gfx1250", "llvm23") + self.assertEqual(walked, counted) + self.assertGreater(cluster, 0) + self.assertGreater(split, cluster) + + def test_gfx1250_llvm22_has_no_cluster_jobs(self): + walked, counted, split, cluster = self._walk("gfx1250", "llvm22") + self.assertEqual(walked, counted) + self.assertGreater(split, 0) + self.assertEqual(cluster, 0) + + def test_other_wmma_targets_have_no_split_jobs(self): + for arch in ("gfx1201", "gfx1151"): + with self.subTest(arch=arch): + walked, counted, split, _ = self._walk(arch, "llvm23") + self.assertEqual(walked, counted) + self.assertEqual(split, 0) + + +class TestSplitKSpecDict(unittest.TestCase): + """The Python-to-C++ spec dict carries the split-K fields as they are.""" + + def test_fields_are_forwarded(self): + from rocke.core.backend import conv_implicit_gemm_spec_to_dict + + d = conv_implicit_gemm_spec_to_dict( + _spec(split_k=2, split_k_cluster=4, pipeline="mem", num_load_waves=3) + ) + self.assertEqual( + (d["split_k"], d["split_k_cluster"], d["num_load_waves"]), (2, 4, 3) + ) + + def test_a_spec_without_them_is_rejected(self): + # A spec object missing a field is malformed: it must not be read as + # an unsplit, cluster-free kernel with four load waves. + from rocke.core.backend import conv_implicit_gemm_spec_to_dict + + spec = _spec(split_k=2) + for missing in ("split_k", "split_k_cluster", "num_load_waves"): + + class _Without: + def __getattr__(self, name, _missing=missing): + if name == _missing: + raise AttributeError(name) + return getattr(spec, name) + + with self.subTest(missing=missing): + with self.assertRaises(AttributeError): + conv_implicit_gemm_spec_to_dict(_Without()) + + +def _engine_or_skip(case): + try: + import rocke_engine # type: ignore + except ImportError: + case.skipTest("rocke_engine (the C++ engine extension) is not built") + return rocke_engine + + +class TestEngineAgreement(unittest.TestCase): + """The Python and C++ engines, reached through the backend, agree on the + accumulator epilogue and on split-K admission.""" + + def test_split_k_with_a_fused_epilogue_is_rejected_by_both(self): + _engine_or_skip(self) + from kernels.common.conv_implicit_gemm import ConvAccumulatorEpilogue + from rocke.core.backend import lower_conv_implicit_gemm + + for epi in ( + ConvAccumulatorEpilogue(bias=1.0), + ConvAccumulatorEpilogue(relu=True), + ConvAccumulatorEpilogue(clamp_min=-1.0, clamp_max=1.0), + ): + spec = _spec(split_k=2, acc_epilogue=epi) + for backend in ("python", "cpp"): + with self.subTest(epi=epi, backend=backend): + with self.assertRaisesRegex(Exception, "identity acc_epilogue"): + lower_conv_implicit_gemm(spec, arch="gfx1250", backend=backend) + + def test_unsplit_fused_epilogue_lowers_identically(self): + _engine_or_skip(self) + from kernels.common.conv_implicit_gemm import ConvAccumulatorEpilogue + from rocke.core.backend import lower_conv_implicit_gemm + + # fp16: the C++ conv implicit-GEMM builder stages A/B in LDS as f16 + # whatever the dtype (a separate, pre-existing mirror gap), so bf16 + # cannot be compared through this path yet. + fp16 = ConvDataSpec("fp16", "fp16", "fp16") + spec = _spec( + data=fp16, + acc_epilogue=ConvAccumulatorEpilogue( + bias=0.5, scale=2.0, relu=True, clamp_max=6.0 + ), + ) + py = lower_conv_implicit_gemm(spec, arch="gfx1250", backend="python") + cpp = lower_conv_implicit_gemm(spec, arch="gfx1250", backend="cpp") + self.assertEqual(py.llvm_text, cpp.llvm_text) + plain = lower_conv_implicit_gemm( + _spec(data=fp16), arch="gfx1250", backend="cpp" + ) + self.assertNotEqual(cpp.llvm_text, plain.llvm_text) + + def test_reduce_rejects_an_explicit_none_dtype_in_both(self): + engine = _engine_or_skip(self) + from kernels.common.conv_fwd_split_k_reduce import ( + FwdSplitKReduceSpec, + build_conv_fwd_split_k_reduce, + ) + + with self.assertRaisesRegex(ValueError, r"dtype_d must be fp16 or bf16"): + build_conv_fwd_split_k_reduce( + FwdSplitKReduceSpec(dtype_d=None), arch="gfx1250" + ) + with self.assertRaisesRegex( + Exception, r"dtype_d must be fp16 or bf16 \(got None\)" + ): + engine.conv_fwd_split_k_reduce_lower_llvm( + dict(dtype_d=None, vec=4), arch="gfx1250" + ) + # An absent key keeps the engine default (the dataclass default, bf16). + ll = engine.conv_fwd_split_k_reduce_lower_llvm(dict(vec=4), arch="gfx1250") + self.assertIn("bf16", ll) + + +class TestSplitKReduceEngine(unittest.TestCase): + """The C++ Stage 2 builder fails cleanly on a name it cannot hold.""" + + def test_overlong_name_is_an_error_not_a_stale_builder(self): + try: + import rocke_engine # type: ignore + except ImportError: + self.skipTest("rocke_engine (the C++ engine extension) is not built") + spec = dict(dtype_d="bf16", vec=4, name="x" * 300) + with self.assertRaisesRegex(Exception, "kernel name does not fit"): + rocke_engine.conv_fwd_split_k_reduce_lower_llvm(spec, arch="gfx1250") + with self.assertRaisesRegex(Exception, "kernel name does not fit"): + rocke_engine.conv_fwd_split_k_reduce_serialize_ir(spec, arch="gfx1250") + # A name that fits still builds. + ll = rocke_engine.conv_fwd_split_k_reduce_lower_llvm( + dict(dtype_d="bf16", vec=4), arch="gfx1250" + ) + self.assertIn("conv_fwd_splitk_reduce", ll) + + +# --------------------------------------------------------------------------- +# GPU numerics +# --------------------------------------------------------------------------- + + +def _gpu_skip() -> str: + if not _HAS_TORCH: + return "torch not importable" + try: + from rocke.runtime.hip_module import get_device_arch + + arch = get_device_arch() + except Exception as e: # pragma: no cover - no device + return f"no GPU ({e})" + return ( + "" if arch == "gfx1250" else f"split-K fwd is WMMA (gfx1250); device is {arch}" + ) + + +def _u8(t): + return (ctypes.c_uint8 * t.nbytes).from_address(t.data_ptr()) + + +def _run(problem, spec_kw, degrees, dtype="bf16", unrolled=True, cluster=False): + """Max relative error per degree of one split-K binary on ``problem``.""" + import torch + + from rocke import compile_kernel + from builders.common.conv_reference import conv_reference_gfx1250 + from kernels.common.conv_abi import conv_args_signature + from kernels.common.conv_args import ConvArgs + from kernels.common.conv_fwd_split_k_reduce import ( + FwdSplitKReduceSpec, + build_conv_fwd_split_k_reduce, + fwd_split_k_reduce_grid, + fwd_split_k_reduce_signature, + fwd_split_k_reduce_vec, + fwd_split_k_workspace_nbytes, + ) + from rocke.runtime.hip_module import Runtime + from rocke.runtime.launcher import ( + KernelLauncher, + LaunchConfig, + synchronize_and_release, + ) + + data = ConvDataSpec(dtype, dtype, dtype) + + def _stage1(degree): + # The cluster binary bakes its degree (one per degree); the plain one + # serves every degree. + spec = _spec( + problem=problem, + data=data, + split_k=2, + split_k_cluster=degree if cluster else 0, + **spec_kw, + ) + ok, why = is_valid_spec(spec, "gfx1250") + assert ok, why + art = compile_kernel( + build_implicit_gemm_conv(spec, arch="gfx1250"), arch="gfx1250" + ) + return spec, KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=conv_args_signature(dtype, two_stage=True), + cluster_dims=(1, 1, degree) if cluster else None, + ) + + if not cluster: + spec, s1 = _stage1(2) + elems = problem.M * problem.K + + def _stage2(degree): + # The degree-specialised fold (what the benchmark launches), or the + # runtime-loop one. + r_spec = FwdSplitKReduceSpec( + dtype_d=dtype, + vec=fwd_split_k_reduce_vec(elems), + degree=degree if unrolled else 0, + ) + r_art = compile_kernel(build_conv_fwd_split_k_reduce(r_spec), arch="gfx1250") + return r_spec, KernelLauncher( + hsaco=r_art.hsaco, + kernel_name=r_art.kernel_name, + signature=fwd_split_k_reduce_signature(r_spec), + ) + + tdt = {"fp16": torch.float16, "bf16": torch.bfloat16}[dtype] + torch.manual_seed(0) + A = ( + torch.empty(problem.N, problem.Hi, problem.Wi, problem.C) + .uniform_(-1, 1) + .to(tdt) + ) + B = ( + torch.empty(problem.K, problem.Y, problem.X, problem.cpg) + .uniform_(-1, 1) + .to(tdt) + ) + ref = conv_reference_gfx1250(A, B, problem, out_dtype=tdt).float().cpu() + D = torch.empty(problem.N, problem.Ho, problem.Wo, problem.K, dtype=tdt) + # Host <-> device copies go through pinned buffers only: a pageable copy + # registers the host range with the driver, which must not overlap a + # kernel in flight. + A, B = A.pin_memory(), B.pin_memory() + + rt = Runtime() + bufs = [rt.alloc(A.nbytes), rt.alloc(B.nbytes), rt.alloc(D.nbytes)] + rt.memcpy_h2d(bufs[0], _u8(A), A.nbytes) + rt.memcpy_h2d(bufs[1], _u8(B), B.nbytes) + errs = {} + try: + for degree in degrees: + if cluster: + spec, s1 = _stage1(degree) + r_spec, s2 = _stage2(degree) + ws_bytes = fwd_split_k_workspace_nbytes(problem.M, problem.K, degree) + ws = rt.alloc(ws_bytes) + # Poison the workspace and the output: Stage 1 must write every slab + # element, Stage 2 every output element. + rt.memset(ws, 0x7F, ws_bytes) + rt.memset(bufs[2], 0x7F, D.nbytes) + args = ConvArgs.from_problem( + problem, tile_m=spec.tile_m, tile_n=spec.tile_n, tile_k=spec.tile_k + ) + values = args.to_launch_values( + *[int(x) for x in bufs], + A.nbytes, + B.nbytes, + D.nbytes, + split_k=degree, + ws_ptr=int(ws), + ws_bytes=ws_bytes, + ) + s1_cfg = LaunchConfig( + grid=args.grid(degree), + block=(spec.launch_block_size, 1, 1), + fence=True, + ) + if cluster: + # One launch: the cluster folds the slabs into D itself. Run + # it through bind() too, the path the benchmark times. + s1(values, config=s1_cfg) + out = torch.empty_like(D).pin_memory() + rt.memcpy_d2h(_u8(out), bufs[2], D.nbytes) + errs[degree] = float( + (out.float() - ref).abs().max() / ref.abs().max().clamp(min=1.0) + ) + rt.memset(bufs[2], 0x7F, D.nbytes) + s1.bind(values, config=s1_cfg)() + synchronize_and_release() + rt.memcpy_d2h(_u8(out), bufs[2], D.nbytes) + rt.free(ws) + diff = (out.float() - ref).abs().max() + errs[degree] = max( + errs[degree], float(diff / ref.abs().max().clamp(min=1.0)) + ) + continue + s1(values, config=s1_cfg) + s2( + dict( + ws_ptr=int(ws), + d_ptr=int(bufs[2]), + elems=elems, + ws_bytes=ws_bytes, + d_bytes=D.nbytes, + ks_count=degree, + ), + config=LaunchConfig( + grid=fwd_split_k_reduce_grid(r_spec, elems), + block=(r_spec.block_size, 1, 1), + fence=True, + ), + ) + out = torch.empty_like(D).pin_memory() + rt.memcpy_d2h(_u8(out), bufs[2], D.nbytes) + rt.free(ws) + diff = (out.float() - ref).abs().max() + errs[degree] = float(diff / ref.abs().max().clamp(min=1.0)) + finally: + for x in bufs: + rt.free(x) + return errs + + +_GPU_SKIP = _gpu_skip() + +# (problem, spec knobs, degrees) +_CASES = [ + # small M, deep K; 9*256 = 2304 = 36 K tiles of 64 + ( + "small_m_wavelet", + ConvProblem(N=1, Hi=7, Wi=7, C=256, K=64, Y=3, X=3, pH=1, pW=1), + {}, + (1, 2, 4, 16, 64), + ), + ( + "mem_cshuffle_64", + ConvProblem(N=2, Hi=12, Wi=12, C=128, K=96, Y=3, X=3, pH=1, pW=1), + dict(tile_m=64, tile_n=64, warp_m=2, warp_n=2, pipeline="mem"), + (1, 3, 8), + ), + ( + "mem_default", + ConvProblem(N=1, Hi=9, Wi=9, C=64, K=40, Y=3, X=3, pH=1, pW=1), + dict( + tile_m=32, + tile_n=32, + tile_k=32, + warp_m=2, + warp_n=1, + pipeline="mem", + epilogue="default", + vector_size_c=1, + ), + (2, 5), + ), + # partial last K tile: 9*24 = 216 is not a multiple of 64 + ( + "k_tail_stride2", + ConvProblem(N=2, Hi=15, Wi=13, C=24, K=32, Y=3, X=3, sH=2, sW=2, pH=1, pW=1), + {}, + (2, 3, 4), + ), + ( + "grouped", + ConvProblem(N=1, Hi=10, Wi=10, C=128, K=64, Y=3, X=3, pH=1, pW=1, groups=2), + {}, + (1, 4), + ), + ( + "pointwise_shape_2x2s2", + ConvProblem(N=1, Hi=14, Wi=14, C=256, K=128, Y=2, X=2, sH=2, sW=2), + {}, + (2, 8), + ), +] + + +@unittest.skipIf(_GPU_SKIP, _GPU_SKIP) +class TestSplitKNumerics(unittest.TestCase): + def test_cases(self): + for dtype in ("bf16", "fp16"): + tol = 2e-2 + for name, problem, kw, degrees in _CASES: + with self.subTest(case=name, dtype=dtype): + for unrolled in (True, False): + errs = _run(problem, kw, degrees, dtype, unrolled) + for degree, err in errs.items(): + self.assertLess( + err, + tol, + f"{name} {dtype} degree {degree} " + f"unrolled={unrolled}: {err}", + ) + + +@unittest.skipIf(_GPU_SKIP, _GPU_SKIP) +class TestSplitKAutoNumerics(unittest.TestCase): + """The binary an auto (-1) spec builds, launched at the resolved degree.""" + + def test_small_m_at_the_resolved_degree(self): + from kernels.common.conv_split_k import resolve_conv_split_k + + problem = _CASES[0][1] + spec = _spec(problem=problem, split_k=-1) + degree = resolve_conv_split_k( + "fwd", + problem, + arch="gfx1250", + dtype="bf16", + tile_m=spec.tile_m, + tile_n=spec.tile_n, + tile_k=spec.tile_k, + block_size=spec.block_size, + ).split_k + self.assertGreater(degree, 1) + auto = build_implicit_gemm_conv(spec, arch="gfx1250") + split = build_implicit_gemm_conv( + _spec(problem=problem, split_k=2), arch="gfx1250" + ) + self.assertEqual(auto.name, split.name) + err = _run(problem, {}, (degree,))[degree] + self.assertLess(err, 2e-2, f"degree {degree}: {err}") + + +# (problem, spec knobs, degrees) for the single-launch cluster fold (mem only). +_CLUSTER_CASES = [ + ( + "small_m", + ConvProblem(N=1, Hi=7, Wi=7, C=256, K=64, Y=3, X=3, pH=1, pW=1), + dict(pipeline="mem"), + (2, 4, 15), + ), + ( + "mem_cshuffle_64", + ConvProblem(N=2, Hi=12, Wi=12, C=128, K=96, Y=3, X=3, pH=1, pW=1), + dict(tile_m=64, tile_n=64, warp_m=2, warp_n=2, pipeline="mem"), + (3, 8), + ), + ( + "k_tail_stride2", + ConvProblem(N=2, Hi=15, Wi=13, C=24, K=32, Y=3, X=3, sH=2, sW=2, pH=1, pW=1), + dict(tile_m=32, tile_n=32, warp_m=2, warp_n=1, pipeline="mem"), + (2, 4), + ), + ( + "grouped", + ConvProblem(N=1, Hi=10, Wi=10, C=128, K=64, Y=3, X=3, pH=1, pW=1, groups=2), + dict(tile_m=32, tile_n=32, warp_m=1, warp_n=2, pipeline="mem"), + (2, 4), + ), +] + + +@unittest.skipIf(_GPU_SKIP, _GPU_SKIP) +class TestClusterSplitKNumerics(unittest.TestCase): + def test_cases(self): + for dtype in ("bf16", "fp16"): + for name, problem, kw, degrees in _CLUSTER_CASES: + with self.subTest(case=name, dtype=dtype): + errs = _run(problem, kw, degrees, dtype, cluster=True) + for degree, err in errs.items(): + self.assertLess(err, 2e-2, f"{name} {dtype} S={degree}: {err}") + + +if __name__ == "__main__": + unittest.main() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_prefetch_addr.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_prefetch_addr.py index 9407c9263b09..ebe86f8e1743 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_prefetch_addr.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_prefetch_addr.py @@ -323,6 +323,7 @@ def test_off_keeps_the_historical_hash(self): del fields["tap_outer_k"] # dgrad's tap-outer K loop, off del fields["dy_halo"], fields["halo_max_wo"] # dgrad dY halo, off del fields["dy_halo_setprio"], fields["dy_halo_kouter_pad"] + del fields["split_k_cluster"] # cluster split-K forward, off del fields["streamk"], fields["streamk_reduction"] # stream-K, off legacy = hashlib.sha1( json.dumps(fields, sort_keys=True, separators=(",", ":")).encode() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index ebcd48c4b8dd..51f8a7711273 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -606,6 +606,7 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_fwd_small_wmma.py" "tests/library/tests/test_direct_dgrad_wmma.py" "tests/library/tests/test_direct_fwd_wmma.py" + "tests/library/tests/test_conv_fwd_split_k.py" "tests/library/tests/test_conv_winograd_fwd_cpu.py" "tests/library/tests/test_conv_winograd_gemm_wmma.py" "tests/library/tests/test_conv_winograd_gen_cpu.py" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp index c0523ca03f5f..4b1d679c0ae8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp @@ -51,6 +51,7 @@ extern "C" { #include "rocke/instance_conv_direct_grouped.h" #include "rocke/instance_conv_direct_nongrouped.h" #include "rocke/instance_conv_fwd_small_wmma.h" +#include "rocke/instance_conv_fwd_split_k_reduce.h" #include "rocke/instance_conv_implicit_gemm.h" #include "rocke/instance_conv_implicit_gemm_wgrad.h" #include "rocke/instance_conv_wgrad_workspace_reduce.h" @@ -1592,6 +1593,27 @@ void fill_conv_problem(rocke_conv_problem_t* p, const py::dict& d) /* ========================= conv_implicit_gemm ======================= */ +/* ConvAccumulatorEpilogue dict {bias, scale, relu, clamp_min, clamp_max} + * (clamp_* None => unset). Absent keys keep the identity default. */ +void fill_conv_acc_epilogue(rocke_conv_acc_epilogue_t* e, const py::dict& d) +{ + if(d.contains("bias") && !d["bias"].is_none()) + e->bias = d["bias"].cast(); + if(d.contains("scale") && !d["scale"].is_none()) + e->scale = d["scale"].cast(); + e->relu = dict_bool(d, "relu", e->relu); + if(d.contains("clamp_min") && !d["clamp_min"].is_none()) + { + e->has_clamp_min = true; + e->clamp_min = d["clamp_min"].cast(); + } + if(d.contains("clamp_max") && !d["clamp_max"].is_none()) + { + e->has_clamp_max = true; + e->clamp_max = d["clamp_max"].cast(); + } +} + rocke_implicit_gemm_conv_spec_t conv_igemm_build_spec(const py::dict& d, std::deque& store) { @@ -1648,6 +1670,17 @@ rocke_implicit_gemm_conv_spec_t conv_igemm_build_spec(const py::dict& d, s.k0_k1_split = dict_bool(d, "k0_k1_split", s.k0_k1_split); s.groups = dict_int(d, "groups", s.groups); s.group_merge = dict_int(d, "group_merge", s.group_merge); + /* Split-K forward: a different ABI and epilogue (and the cluster form a + * cluster launch shape); dropped, the engine would build the plain kernel. */ + s.split_k = dict_int(d, "split_k", s.split_k); + s.split_k_cluster = dict_int(d, "split_k_cluster", s.split_k_cluster); + /* The wavelet load-wave count sizes the launch block and the load-wave + * loaders; dropped, a non-default count builds a different kernel. */ + s.num_load_waves = dict_int(d, "num_load_waves", s.num_load_waves); + /* The fused accumulator epilogue: dropped, the engine would build the + * identity epilogue (and admit split-K specs Python rejects). */ + if(d.contains("acc_epilogue") && py::isinstance(d["acc_epilogue"])) + fill_conv_acc_epilogue(&s.acc_epilogue, d["acc_epilogue"].cast()); if(d.contains("vector_size_a") && !d["vector_size_a"].is_none()) { s.has_vector_size_a = true; @@ -3278,6 +3311,67 @@ std::vector conv_wgrad_reduce_verify(const py::dict& d, const std:: rocke_build_wgrad_workspace_reduce_new(&b, &s, arch.c_str())); } +/* ================= conv_fwd_split_k_reduce (fwd split-K Stage 2) =============== */ + +rocke_fwd_split_k_reduce_spec_t conv_fwd_split_k_reduce_build_spec(const py::dict& d, + std::deque& store) +{ + auto keep = [&](const std::string& s) -> const char* { + store.push_back(s); + return store.back().c_str(); + }; + rocke_fwd_split_k_reduce_spec_t s = rocke_fwd_split_k_reduce_spec_default(); + /* An explicit dtype_d=None is a malformed spec (Python validate() rejects + * it), not a request for the bf16 default: pass NULL so the builder + * rejects it. Only an absent key keeps the default. */ + if(d.contains("dtype_d") && d["dtype_d"].is_none()) + s.dtype_d = NULL; + s.vec = dict_int(d, "vec", s.vec); + s.block_size = dict_int(d, "block_size", s.block_size); + /* The unrolled fold bakes the slab count: it must match the launch's + * ks_count, so it is parsed rather than left at the runtime-loop default. */ + s.degree = dict_int(d, "degree", s.degree); + { + std::string v; + if(dict_str(d, "dtype_d", v)) + s.dtype_d = keep(v); + if(dict_str(d, "name", v)) + s.name = keep(v); + } + return s; +} + +std::string conv_fwd_split_k_reduce_lower_llvm(const py::dict& d, const std::string& arch) +{ + require_arch(arch); + std::deque store; + rocke_fwd_split_k_reduce_spec_t s = conv_fwd_split_k_reduce_build_spec(d, store); + char* ll = nullptr; + char err[ROCKE_ERR_MSG_CAP]; + err[0] = '\0'; + rocke_status_t st = rocke_fwd_split_k_reduce_lower_to_llvm( + &s, arch.c_str(), resolve_python_llvm_flavor(), &ll, err, sizeof err); + return take_lowered(st, ll, err, "rocke_engine.conv_fwd_split_k_reduce_lower_llvm"); +} + +std::string conv_fwd_split_k_reduce_serialize_ir(const py::dict& d, const std::string& arch) +{ + require_arch(arch); + ROCKE_FAMILY_SERIALIZE_BODY("rocke_engine.conv_fwd_split_k_reduce_serialize_ir", + rocke_fwd_split_k_reduce_spec_t, + conv_fwd_split_k_reduce_build_spec, + rocke_build_conv_fwd_split_k_reduce_new(&b, &s, arch.c_str())); +} + +std::vector conv_fwd_split_k_reduce_verify(const py::dict& d, const std::string& arch) +{ + require_arch(arch); + ROCKE_FAMILY_VERIFY_BODY("rocke_engine.conv_fwd_split_k_reduce_verify", + rocke_fwd_split_k_reduce_spec_t, + conv_fwd_split_k_reduce_build_spec, + rocke_build_conv_fwd_split_k_reduce_new(&b, &s, arch.c_str())); +} + /* ---- pooling (nested problem, generic lower, build takes arch) ---- */ rocke_pooling2d_spec_t pool_build_spec(const py::dict& d, std::deque& store) { @@ -5101,6 +5195,11 @@ PYBIND11_MODULE(rocke_engine, m) conv_wgrad_reduce_lower_llvm, conv_wgrad_reduce_serialize_ir, conv_wgrad_reduce_verify); + /* ---- fwd split-K Stage 2 (folds the [ks_count][M][K] fp32 slabs) ---- */ + reg3("conv_fwd_split_k_reduce", + conv_fwd_split_k_reduce_lower_llvm, + conv_fwd_split_k_reduce_serialize_ir, + conv_fwd_split_k_reduce_verify); m.def( "conv_wgrad_workspace_bytes", [](const py::dict& d) -> size_t { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/epilogues.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/epilogues.cpp index 4fe0a6a71f69..27f1a0b31c22 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/epilogues.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/epilogues.cpp @@ -697,15 +697,18 @@ void rocke_cshuffle_epilogue_store(rocke_ir_builder_t* b, if(_fp32_out) { + /* Python: b.smem_load_vN(c_smem, row, col, dtype=F32, n=...) -- the + * generic vector load, not smem_load_vN_f32 (whose result name + * differs), so the f32 tile reads mirror it. */ if(sv == 1) { - rocke_value_t* vf32 = rocke_b_smem_load_vN_f32(b, c_smem, idx, 2, 1); + rocke_value_t* vf32 = rocke_b_smem_load_vN(b, c_smem, idx, 2, rocke_f32(), 1); rocke_b_buffer_store_f32( b, d_rsrc, safe, rocke_b_const_i32(b, 0), rocke_b_vec_extract(b, vf32, 0)); } else { - rocke_value_t* vf32 = rocke_b_smem_load_vN_f32(b, c_smem, idx, 2, sv); + rocke_value_t* vf32 = rocke_b_smem_load_vN(b, c_smem, idx, 2, rocke_f32(), sv); rocke_b_buffer_store_vN_f32(b, d_rsrc, safe, rocke_b_const_i32(b, 0), vf32, sv); } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_abi.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_abi.h index 21e7b3fd2747..4456e1bed990 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_abi.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_abi.h @@ -60,9 +60,11 @@ typedef struct /* The Python kind string ("a", "b", "d", "f32*", "i32*", "i32"). */ const char* rocke_conv_arg_kind_str(rocke_conv_arg_kind_t kind); -/* Implicit-GEMM ABI for one direction ("fwd", "wgrad", "dgrad"). Returns - * false -- leaving *out empty -- for an unknown direction or for two_stage - * outside wgrad, the cases conv_arg_names() raises on. */ +/* Implicit-GEMM ABI for one direction ("fwd", "wgrad", "dgrad"). two_stage + * is wgrad's two-stage epilogue, and on fwd the split-K binary (a trailing + * ws_ptr, ws_bytes, ks, ks_count). Returns false -- leaving *out empty -- for + * an unknown direction or for two_stage on dgrad, the cases conv_arg_names() + * raises on. */ bool rocke_conv_arg_names(const char* direction, bool is_3d, bool two_stage, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_fwd_split_k_reduce.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_fwd_split_k_reduce.h new file mode 100644 index 000000000000..112fe0c9f8d5 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_fwd_split_k_reduce.h @@ -0,0 +1,113 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT + * + * rocke/instance_conv_fwd_split_k_reduce.h + * + * C99 port of: + * library/kernels/common/conv_fwd_split_k_reduce.py + * (FwdSplitKReduceSpec + build_conv_fwd_split_k_reduce) + * + * Stage 2 of the split-K forward convolution. Stage 1 (the implicit-GEMM + * forward with split_k > 1) stores each K slice's fp32 partial output into its + * own slab of a [ks_count][M][K] workspace -- plain stores, every element + * written exactly once. This kernel sums the ks_count slabs element-wise in a + * fixed order and converts the result to dtype_d into D (N[D]HWK, the same + * packed layout as one slab). The result is deterministic. + * + * Kernel signature: + * ws_ptr : f32 global ptr (readonly) -- workspace [ks_count * elems] + * d_ptr : dtype_d global ptr -- output [elems] + * elems : i32 -- elements per slab (M * K) + * ws_bytes : i32 -- workspace byte size (ABI boundary) + * d_bytes : i32 -- output byte size (ABI boundary) + * ks_count : i32 -- number of slabs (the split-K degree) + * + * Grid: (ceil(elems / (vec * block_size)), 1, 1); block (block_size, 1, 1). + */ +#ifndef ROCKE_INSTANCE_CONV_FWD_SPLIT_K_REDUCE_H +#define ROCKE_INSTANCE_CONV_FWD_SPLIT_K_REDUCE_H + +#include +#include + +#include "rocke/helper_rocke.helpers.spec.h" +#include "rocke/ir.h" +#include "rocke/lower_llvm.h" /* for rocke_llvm_flavor_t */ + +#ifdef __cplusplus +extern "C" { +#endif + +/* -------------------------------------------------- FwdSplitKReduceSpec + * + * Mirror of Python FwdSplitKReduceSpec (frozen dataclass): + * + * dtype_d: str = "bf16" + * vec: int = 4 -- output elements per thread (1, 2 or 4) + * block_size: int = 256 + * degree: int = 0 -- > 0 unrolls the fold for exactly that many + * slabs (ks_count must then equal it); 0 keeps + * the runtime loop over ks_count + * name: str = "conv_fwd_splitk_reduce" + */ +typedef struct rocke_fwd_split_k_reduce_spec +{ + const char* dtype_d; /* default "bf16" (spec_default); NULL is rejected */ + int vec; /* default 4 */ + int block_size; /* default 256 */ + int degree; /* default 0 */ + const char* name; /* default "conv_fwd_splitk_reduce" */ +} rocke_fwd_split_k_reduce_spec_t; + +/* Default-constructed spec (every field == the Python dataclass default). */ +rocke_fwd_split_k_reduce_spec_t rocke_fwd_split_k_reduce_spec_default(void); + +/* spec.kernel_name() -> NUL-terminated into out (capacity out_cap). + * ROCKE_ERR_VALUE on a NULL dtype_d or a name that does not fit. */ +rocke_status_t rocke_fwd_split_k_reduce_kernel_name(const rocke_fwd_split_k_reduce_spec_t* spec, + char* out, + size_t out_cap); + +/* spec.validate(). Returns true if valid; false + reason string otherwise. */ +bool rocke_fwd_split_k_reduce_is_valid_spec(const rocke_fwd_split_k_reduce_spec_t* spec, + char* reason, + size_t reason_cap); + +/* Build the reduction kernel IR into the (initialised) builder b. + * Returns b->kernel on success, NULL on error (check b's sticky error). */ +rocke_kernel_def_t* rocke_build_conv_fwd_split_k_reduce(rocke_ir_builder_t* b, + const rocke_fwd_split_k_reduce_spec_t* spec, + const char* arch); + +/* Validate spec, init b with spec.kernel_name(), then build. b is zeroed + * first, so on any failure (including a name that does not fit) it holds the + * error and can be queried and rocke_ir_builder_free()d. */ +rocke_kernel_def_t* rocke_build_conv_fwd_split_k_reduce_new( + rocke_ir_builder_t* b, const rocke_fwd_split_k_reduce_spec_t* spec, const char* arch); + +/* Launch grid: (max(1, ceil(elems / (vec * block_size))), 1, 1). */ +void rocke_fwd_split_k_reduce_grid(const rocke_fwd_split_k_reduce_spec_t* spec, + long long elems, + int* out_gx, + int* out_gy, + int* out_gz); + +/* The kernel signature (6 entries: ws_ptr, d_ptr, elems, ws_bytes, d_bytes, + * ks_count). arena must not be NULL. */ +rocke_status_t rocke_fwd_split_k_reduce_signature(rocke_arena_t* arena, + const rocke_fwd_split_k_reduce_spec_t* spec, + const rocke_sig_entry_t** out_items, + size_t* out_count); + +/* Convenience: build + lower to LLVM IR in one shot. */ +rocke_status_t rocke_fwd_split_k_reduce_lower_to_llvm(const rocke_fwd_split_k_reduce_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap); + +#ifdef __cplusplus +} +#endif +#endif /* ROCKE_INSTANCE_CONV_FWD_SPLIT_K_REDUCE_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm.h index a1692ab6be14..1d8347d134a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm.h @@ -168,6 +168,21 @@ typedef struct rocke_implicit_gemm_conv_spec * default path stays byte-identical. */ int group_merge; /* default 1 */ + /* Split-K forward (WMMA only). split_k > 1 builds the split-K binary, not a + * fixed degree: the degree and the slice width arrive as the ks_count / ks + * kernargs (as in wgrad), after a ws_ptr / ws_bytes fp32 workspace pair. + * Every CTA reduces one slice of the [Z*]Y*X*cpg reduction and stores its + * fp32 partial tile into slab `slice` of the [ks_count][M][K] workspace; + * Stage 2 (instance_conv_fwd_split_k_reduce.h) folds the slabs into D. + * Default 1 emits none of this. Mirrors Python ImplicitGemmConvSpec.split_k. */ + int split_k; /* default 1 */ + /* Single-launch split-K through a gfx1250 workgroup cluster (0 = off). The + * split_k_cluster slices of one output tile form one (1, 1, S) cluster; + * after a cluster barrier each slice folds 1/S of the tile's rows across + * the S slabs into D. Requires split_k > 1. Mirrors + * ImplicitGemmConvSpec.split_k_cluster. */ + int split_k_cluster; /* default 0 */ + /* #8624 vector-sizes-as-args: per-operand load/store vector widths * (elements). has_* false => Python None => auto-select (choose_load_vec / * CShuffleEpilogue.from_grid default). When set they override the per-A/B @@ -252,6 +267,25 @@ bool rocke_conv_fwd_group_merge_available(const rocke_implicit_gemm_conv_spec_t* char* reason, size_t reason_cap); +/* Admission gate for spec.split_k != 1 / spec.split_k_cluster, called from + * rocke_implicit_gemm_conv_is_valid_spec after the WMMA gates. `family` is the + * MMA family ("wmma" / "mma"); `arch` (NULL = unchecked, Python target=None) + * gates the cluster form on workgroup-cluster launch. Returns true when the + * split-K forward can be built; on false writes the reason. Mirrors Python + * fwd_split_k_reason. */ +bool rocke_conv_fwd_split_k_available(const rocke_implicit_gemm_conv_spec_t* s, + const char* family, + const char* arch, + char* reason, + size_t reason_cap); + +/* The cluster split-K forward's problem rule (K and kpg multiples of 4), on + * the build problem and on every launched one. Mirrors Python + * fwd_split_k_cluster_problem_reason; returns true when `p` can run. */ +bool rocke_conv_fwd_split_k_cluster_problem_ok(const rocke_conv_problem_t* p, + char* reason, + size_t reason_cap); + /* spec.kernel_name() -> NUL-terminated into out (capacity out_cap). Returns * ROCKE_OK or ROCKE_ERR_VALUE (buffer too small). */ rocke_status_t rocke_implicit_gemm_conv_spec_kernel_name(const rocke_implicit_gemm_conv_spec_t* s, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h index 59e02df03d6c..a576b518eccf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h @@ -359,6 +359,18 @@ typedef struct rocke_conv_build_ctx rocke_value_t* group_idx; rocke_value_t* k_out_group_base; + /* ---- forward split-K state (spec.split_k > 1; all NULL / false off it) ---- + * The trailing ws_ptr / ws_bytes / ks / ks_count kernargs, the CTA's K + * slice (z = group * ks_count + slice), and the slab base the epilogue adds + * to every D offset. Mirrors the Python builder's `split` locals. */ + bool fwd_split; + rocke_value_t* fwd_p_ws; + rocke_value_t* fwd_p_ws_bytes; + rocke_value_t* fwd_p_ks; + rocke_value_t* fwd_p_ks_count; + rocke_value_t* fwd_slice_id; + rocke_value_t* fwd_slab_base; + /* ---- wgrad stream-K state (NULL off the stream-K path) ---- * * wgrad_c_kpg = b.const_i32(kpg) (grouped dY k_out fold, Python c_kpg) * wgrad_replica_sel = linear CTA index hashed into the workspace replica in @@ -793,6 +805,7 @@ void rocke_conv_emit_direct_epilogue_wmma(rocke_ir_builder_t* b, rocke_value_t* ir_d_pitch_pw, const rocke_tensor_descriptor_t* D_desc, rocke_value_t* k_out_group_base, + rocke_value_t* slab_base, rocke_value_t* bound_m, rocke_value_t* bound_n); @@ -809,6 +822,8 @@ void rocke_conv_emit_cshuffle_epilogue(rocke_ir_builder_t* b, const rocke_mmaop_t* op, const rocke_tensor_descriptor_t* D_desc, rocke_value_t* k_out_group_base, + rocke_value_t* slab_base, + const char* out_dtype, rocke_value_t* bound_m, rocke_value_t* bound_n); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_abi.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_abi.cpp index 76ae1370edcb..a7ef0171303b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_abi.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_abi.cpp @@ -122,7 +122,9 @@ void leading(rocke_conv_arg_list_t* out, const char* a, const char* b, const cha const int kLeadingArgs = 6; -void fwd_arg_names(rocke_conv_arg_list_t* out, bool is_3d) +/* split_k: the split-K binary's ABI -- the fp32 workspace its slices store + * into, then the slice width and degree (wgrad's ks / ks_count). */ +void fwd_arg_names(rocke_conv_arg_list_t* out, bool is_3d, bool split_k) { leading(out, "A", "B", "D"); dims(out, is_3d); @@ -150,6 +152,13 @@ void fwd_arg_names(rocke_conv_arg_list_t* out, bool is_3d) /* k -> ([z,] y, x, c): divisors cpg, X[, Y]. */ magic_2d3d(out, "k_", is_3d, "Y", "X", "cpg"); grid(out); + if(split_k) + { + push(out, "ws_ptr", ROCKE_CONV_ARG_F32_PTR); + push_i32(out, "ws_bytes"); + push_i32(out, "ks"); + push_i32(out, "ks_count"); + } } /* The stream-K launch block: the host-computed iteration partition. The @@ -292,11 +301,12 @@ bool rocke_conv_arg_names_ex(const char* direction, out->count = 0; if(direction == NULL) return false; - /* two_stage belongs to wgrad's deterministic epilogue; accepting it - * elsewhere would describe a kernel the flag cannot build. The same goes - * for streamk, which additionally partitions K itself and so cannot be - * combined with the two-stage split. */ - if(two_stage && strcmp(direction, "wgrad") != 0) + /* two_stage is wgrad's deterministic epilogue, and on fwd the split-K + * binary (fp32 workspace slabs + a Stage 2 reduce). dgrad has neither, so + * accepting it there would describe a kernel the flag cannot build. The + * same goes for streamk (wgrad only), which additionally partitions K + * itself and so cannot be combined with the two-stage split. */ + if(two_stage && strcmp(direction, "wgrad") != 0 && strcmp(direction, "fwd") != 0) return false; if(streamk != NULL) { @@ -304,7 +314,7 @@ bool rocke_conv_arg_names_ex(const char* direction, return false; } if(strcmp(direction, "fwd") == 0) - fwd_arg_names(out, is_3d); + fwd_arg_names(out, is_3d, two_stage); else if(strcmp(direction, "wgrad") == 0) wgrad_arg_names(out, is_3d, two_stage, streamk); else if(strcmp(direction, "dgrad") == 0) @@ -367,7 +377,7 @@ void rocke_conv_fwd_problem_block(bool is_3d, rocke_conv_arg_list_t* out) { rocke_conv_arg_list_t full; full.count = 0; - fwd_arg_names(&full, is_3d); + fwd_arg_names(&full, is_3d, false); out->count = 0; for(int i = kLeadingArgs; i < full.count; ++i) push(out, full.items[i].name, full.items[i].kind); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_fwd_split_k_reduce.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_fwd_split_k_reduce.cpp new file mode 100644 index 000000000000..f86752845885 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_fwd_split_k_reduce.cpp @@ -0,0 +1,382 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT + * + * C99 port of library/kernels/common/conv_fwd_split_k_reduce.py + * + * Stage 2 of the split-K forward convolution: sum the ks_count fp32 slabs + * Stage 1 stored and cast into D. Each thread owns `vec` consecutive output + * elements: one vector load per slab, `vec` adds, one vector store. Every SSA + * value is bound in the Python builder's evaluation order. + */ + +#include +#include + +#include "rocke/instance_conv_fwd_split_k_reduce.h" + +#include "rocke/arena.h" /* rocke_arena_alloc */ +#include "rocke/error_boundary.hpp" /* ckc::guard_builder */ +#include "rocke/helper_rocke.helpers.io.h" /* io_ir_type, pack_f32_to, store_scalar_from_f32 */ +#include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join, SignatureBuilder */ +#include "rocke/ir.h" +#include "rocke/ir_internal.h" /* rocke_i_set_err */ +#include "rocke/lower_llvm.h" + +/* The widest per-thread vector: vec <= 4. */ +#define ROCKE_FWD_SPLITK_REDUCE_MAX_VEC 4 + +/* ---- spec helpers -------------------------------------------------------- */ + +rocke_fwd_split_k_reduce_spec_t rocke_fwd_split_k_reduce_spec_default(void) +{ + rocke_fwd_split_k_reduce_spec_t s; + memset(&s, 0, sizeof(s)); + s.dtype_d = "bf16"; + s.vec = 4; + s.block_size = 256; + s.degree = 0; + s.name = "conv_fwd_splitk_reduce"; + return s; +} + +rocke_status_t rocke_fwd_split_k_reduce_kernel_name(const rocke_fwd_split_k_reduce_spec_t* spec, + char* out, + size_t out_cap) +{ + /* Python: kernel_name_join(name, f"v{vec}", f"b{block_size}", + * [f"s{degree}"] if degree, dtype_d) */ + char v_buf[24]; + char b_buf[24]; + char s_buf[24]; + const char* parts[4]; + int n_parts = 0; + if(spec == NULL || out == NULL || spec->dtype_d == NULL) + return ROCKE_ERR_VALUE; + snprintf(v_buf, sizeof(v_buf), "v%d", spec->vec); + snprintf(b_buf, sizeof(b_buf), "b%d", spec->block_size); + parts[n_parts++] = v_buf; + parts[n_parts++] = b_buf; + if(spec->degree) + { + snprintf(s_buf, sizeof(s_buf), "s%d", spec->degree); + parts[n_parts++] = s_buf; + } + parts[n_parts++] = spec->dtype_d; + return rocke_kernel_name_join(spec->name ? spec->name : "conv_fwd_splitk_reduce", + parts, + (size_t)n_parts, + NULL, + NULL, + 0, + out, + out_cap, + NULL); +} + +bool rocke_fwd_split_k_reduce_is_valid_spec(const rocke_fwd_split_k_reduce_spec_t* spec, + char* reason, + size_t reason_cap) +{ + const char* dtype_d; + if(spec == NULL) + { + if(reason && reason_cap > 0) + snprintf(reason, reason_cap, "spec is NULL"); + return false; + } + if(spec->vec != 1 && spec->vec != 2 && spec->vec != 4) + { + if(reason && reason_cap > 0) + snprintf(reason, reason_cap, "vec must be 1, 2 or 4 (got %d)", spec->vec); + return false; + } + /* Python validate() rejects dtype_d=None; NULL is not a request for the + * default (rocke_fwd_split_k_reduce_spec_default sets that). */ + dtype_d = spec->dtype_d; + if(dtype_d == NULL) + { + if(reason && reason_cap > 0) + snprintf(reason, reason_cap, "dtype_d must be fp16 or bf16 (got None)"); + return false; + } + if(strcmp(dtype_d, "fp16") != 0 && strcmp(dtype_d, "bf16") != 0) + { + if(reason && reason_cap > 0) + snprintf(reason, reason_cap, "dtype_d must be fp16 or bf16 (got %s)", dtype_d); + return false; + } + if(spec->degree < 0) + { + if(reason && reason_cap > 0) + snprintf(reason, reason_cap, "degree must be >= 0 (got %d)", spec->degree); + return false; + } + return true; +} + +/* ---- main builder -------------------------------------------------------- */ + +/* Python _load(off): one f32 load (vec 1) or one vector load + extracts. */ +static void fwd_reduce_load( + rocke_ir_builder_t* b, rocke_value_t* ws_ptr, rocke_value_t* off, int vec, rocke_value_t** out) +{ + if(vec == 1) + { + out[0] = rocke_b_global_load_f32(b, ws_ptr, off, 4); + return; + } + rocke_value_t* v = rocke_b_global_load_vN(b, ws_ptr, off, rocke_f32(), vec, 0); + for(int i = 0; i < vec; ++i) + out[i] = rocke_b_vec_extract(b, v, i); +} + +rocke_kernel_def_t* rocke_build_conv_fwd_split_k_reduce(rocke_ir_builder_t* b, + const rocke_fwd_split_k_reduce_spec_t* spec, + const char* arch) +{ + char reason[256]; + (void)arch; + if(b == NULL || spec == NULL) + return NULL; + if(!rocke_fwd_split_k_reduce_is_valid_spec(spec, reason, sizeof(reason))) + { + rocke_i_set_err(b, ROCKE_ERR_VALUE, "fwd_split_k_reduce: %s", reason); + return NULL; + } + + const int vec = spec->vec; + const char* dtype_d = spec->dtype_d; /* non-NULL: validated above */ + rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", spec->block_size); + + /* ---- params ---------------------------------------------------------- */ + rocke_param_opts_t ro; + memset(&ro, 0, sizeof(ro)); + ro.noalias = ro.noalias_set = true; + ro.readonly = ro.readonly_set = true; + ro.align = 16; + ro.align_set = true; + + rocke_param_opts_t wo; + memset(&wo, 0, sizeof(wo)); + wo.noalias = wo.noalias_set = true; + wo.writeonly = wo.writeonly_set = true; + wo.align = 16; + wo.align_set = true; + + rocke_value_t* ws_ptr + = rocke_b_param(b, "ws_ptr", rocke_ptr_type(b, rocke_f32(), "global"), &ro); + rocke_value_t* d_ptr = rocke_b_param( + b, "d_ptr", rocke_ptr_type(b, rocke_b_io_ir_type(b, dtype_d), "global"), &wo); + rocke_value_t* elems = rocke_b_param(b, "elems", rocke_i32(), NULL); + rocke_b_param(b, "ws_bytes", rocke_i32(), NULL); /* ABI boundary */ + rocke_b_param(b, "d_bytes", rocke_i32(), NULL); /* ABI boundary */ + rocke_value_t* ks_count = rocke_b_param(b, "ks_count", rocke_i32(), NULL); + + /* gid = block_id_x * block_size + thread_id_x; idx = gid * vec */ + rocke_value_t* bid = rocke_b_block_id_x(b); + rocke_value_t* c_bs = rocke_b_const_i32(b, spec->block_size); + rocke_value_t* base = rocke_b_mul(b, bid, c_bs); + rocke_value_t* tid = rocke_b_thread_id_x(b); + rocke_value_t* gid = rocke_b_add(b, base, tid); + rocke_value_t* idx = rocke_b_mul(b, gid, rocke_b_const_i32(b, vec)); + + rocke_value_t* total[ROCKE_FWD_SPLITK_REDUCE_MAX_VEC]; + rocke_if_t guard = rocke_b_scf_if(b, rocke_b_cmp_lt(b, idx, elems)); + rocke_b_region_enter(b, guard.then_region); + if(spec->degree) + { + /* Unrolled: every slab's load is issued before the first add. */ + rocke_value_t** slabs = (rocke_value_t**)rocke_arena_alloc( + &b->arena, sizeof(rocke_value_t*) * (size_t)spec->degree * (size_t)vec); + if(slabs == NULL) + { + rocke_i_set_err(b, ROCKE_ERR_OOM, "fwd_split_k_reduce: slab alloc failed"); + rocke_b_region_leave(b); + return NULL; + } + for(int s = 0; s < spec->degree; ++s) + { + rocke_value_t* off = idx; + if(s != 0) + { + rocke_value_t* c_s = rocke_b_const_i32(b, s); + rocke_value_t* s_base = rocke_b_mul(b, c_s, elems); + off = rocke_b_add(b, s_base, idx); + } + fwd_reduce_load(b, ws_ptr, off, vec, &slabs[(size_t)s * (size_t)vec]); + } + for(int i = 0; i < vec; ++i) + total[i] = slabs[i]; + for(int s = 1; s < spec->degree; ++s) + { + for(int i = 0; i < vec; ++i) + total[i] = rocke_b_fadd(b, total[i], slabs[(size_t)s * (size_t)vec + (size_t)i]); + } + } + else + { + rocke_value_t* zero = rocke_b_const_f32(b, 0.0); + static const char* const kAccNames[ROCKE_FWD_SPLITK_REDUCE_MAX_VEC] + = {"acc0", "acc1", "acc2", "acc3"}; + rocke_iter_arg_t init[ROCKE_FWD_SPLITK_REDUCE_MAX_VEC]; + for(int i = 0; i < vec; ++i) + { + init[i].name = kAccNames[i]; + init[i].init = zero; + } + rocke_value_t* c_lo = rocke_b_const_i32(b, 0); + rocke_value_t* c_one = rocke_b_const_i32(b, 1); + rocke_for_t loop = rocke_b_scf_for_iter(b, + c_lo, + ks_count, + c_one, + init, + vec, + "slab", + /*unroll=*/false, + /*elide_trailing_barrier=*/true); + rocke_b_region_enter(b, loop.body); + { + rocke_value_t* s_base = rocke_b_mul(b, loop.iv, elems); + rocke_value_t* off = rocke_b_add(b, s_base, idx); + rocke_value_t* part[ROCKE_FWD_SPLITK_REDUCE_MAX_VEC]; + rocke_value_t* sums[ROCKE_FWD_SPLITK_REDUCE_MAX_VEC]; + fwd_reduce_load(b, ws_ptr, off, vec, part); + for(int i = 0; i < vec; ++i) + sums[i] = rocke_b_fadd(b, loop.iter_vars[i], part[i]); + rocke_b_scf_yield(b, sums, vec); + } + rocke_b_region_leave(b); + for(int i = 0; i < vec; ++i) + total[i] = loop.op->results[i]; + } + if(vec == 1) + { + rocke_b_store_scalar_from_f32(b, d_ptr, idx, total[0], dtype_d); + } + else + { + rocke_value_t* packed = rocke_b_pack_f32_to(b, total, vec, dtype_d); + rocke_b_global_store_vN(b, d_ptr, idx, packed, vec, 0); + } + rocke_b_region_leave(b); + + if(!rocke_ir_builder_ok(b)) + return NULL; + return b->kernel; +} + +rocke_kernel_def_t* rocke_build_conv_fwd_split_k_reduce_new( + rocke_ir_builder_t* b, const rocke_fwd_split_k_reduce_spec_t* spec, const char* arch) +{ + if(b != NULL) + { + /* Zero first, so every early return below leaves a builder the caller + * can safely query and rocke_ir_builder_free(). */ + memset(b, 0, sizeof(*b)); + } + return ckc::guard_builder(b, [&]() -> rocke_kernel_def_t* { + char name[256]; + char reason[256]; + if(b == NULL) + return NULL; + if(spec == NULL) + { + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, "build_conv_fwd_split_k_reduce: null spec"); + return NULL; + } + /* Python: spec.validate() runs before IRBuilder(spec.kernel_name()). */ + if(!rocke_fwd_split_k_reduce_is_valid_spec(spec, reason, sizeof(reason))) + { + rocke_i_set_err(b, ROCKE_ERR_VALUE, "fwd_split_k_reduce: %s", reason); + return NULL; + } + if(rocke_fwd_split_k_reduce_kernel_name(spec, name, sizeof(name)) != ROCKE_OK) + { + rocke_i_set_err_msg( + b, ROCKE_ERR_VALUE, "build_conv_fwd_split_k_reduce: kernel name does not fit"); + return NULL; + } + if(rocke_ir_builder_init(b, name) != ROCKE_OK) + return NULL; + return rocke_build_conv_fwd_split_k_reduce(b, spec, arch); + }); +} + +void rocke_fwd_split_k_reduce_grid(const rocke_fwd_split_k_reduce_spec_t* spec, + long long elems, + int* out_gx, + int* out_gy, + int* out_gz) +{ + const long long per_block = (long long)spec->vec * (long long)spec->block_size; + long long gx = per_block > 0 ? (elems + per_block - 1) / per_block : 1; + *out_gx = (int)(gx < 1 ? 1 : gx); + *out_gy = 1; + *out_gz = 1; +} + +rocke_status_t rocke_fwd_split_k_reduce_signature(rocke_arena_t* arena, + const rocke_fwd_split_k_reduce_spec_t* spec, + const rocke_sig_entry_t** out_items, + size_t* out_count) +{ + rocke_signature_builder_t sb; + rocke_status_t st; + if(arena == NULL || spec == NULL || spec->dtype_d == NULL || out_items == NULL + || out_count == NULL) + return ROCKE_ERR_VALUE; + st = rocke_signature_builder_init(&sb, arena); + if(st != ROCKE_OK) + return st; + rocke_signature_builder_ptr(&sb, "ws_ptr", "fp32", NULL); + rocke_signature_builder_ptr(&sb, "d_ptr", spec->dtype_d, NULL); + rocke_signature_builder_scalar(&sb, "elems", "i32"); + rocke_signature_builder_scalar(&sb, "ws_bytes", "i32"); + rocke_signature_builder_scalar(&sb, "d_bytes", "i32"); + rocke_signature_builder_scalar(&sb, "ks_count", "i32"); + return rocke_signature_builder_build(&sb, out_items, out_count); +} + +rocke_status_t rocke_fwd_split_k_reduce_lower_to_llvm(const rocke_fwd_split_k_reduce_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap) +{ + rocke_ir_builder_t b; + rocke_kernel_def_t* k; + rocke_status_t st; + + if(out_ll != NULL) + *out_ll = NULL; + if(spec == NULL || out_ll == NULL) + return ROCKE_ERR_VALUE; + if(arch == NULL) + arch = "gfx950"; + + k = rocke_build_conv_fwd_split_k_reduce_new(&b, spec, arch); + if(k == NULL) + { + /* Capture status BEFORE free (free memsets struct to zero). */ + st = rocke_ir_builder_status(&b); + if(err != NULL && err_cap > 0) + { + const char* m = rocke_ir_builder_error(&b); + size_t n; + if(m == NULL) + m = "build_conv_fwd_split_k_reduce failed"; + n = strlen(m); + if(n >= err_cap) + n = err_cap - 1; + memcpy(err, m, n); + err[n] = '\0'; + } + rocke_ir_builder_free(&b); + return (st == ROCKE_OK) ? ROCKE_ERR_VALUE : st; + } + st = rocke_lower_kernel_to_llvm_ex(k, flavor, arch, out_ll, err, err_cap); + rocke_ir_builder_free(&b); + return st; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_build_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_build_glue.cpp index 98c61d793ef3..cb9ed6566825 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_build_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_build_glue.cpp @@ -326,6 +326,31 @@ bool rocke_conv_build_ctx_init(rocke_conv_build_ctx_t* ctx, ctx->p_pW_neg = NULL; } + /* ---- split-K kernargs (spec.split_k > 1) ---- + * Python: the fp32 workspace, then the slice width and degree -- the same + * trailing pair wgrad takes -- emitted right after the problem block. The + * cluster form also declares its (1, 1, S) launch shape. */ + ctx->fwd_split = spec->split_k > 1; + if(ctx->fwd_split) + { + rocke_param_opts_t ws_opts; + memset(&ws_opts, 0, sizeof(ws_opts)); + ws_opts.noalias = true; + ws_opts.noalias_set = true; + ws_opts.align = 16; + ws_opts.align_set = true; + ctx->fwd_p_ws + = rocke_b_param(b, "ws_ptr", rocke_ptr_type(b, rocke_f32(), "global"), &ws_opts); + ctx->fwd_p_ws_bytes = rocke_b_param(b, "ws_bytes", rocke_i32(), NULL); + ctx->fwd_p_ks = rocke_b_param(b, "ks", rocke_i32(), NULL); + ctx->fwd_p_ks_count = rocke_b_param(b, "ks_count", rocke_i32(), NULL); + if(spec->split_k_cluster && b->kernel != NULL) + { + const int64_t dims[3] = {1, 1, spec->split_k_cluster}; + rocke_attr_set_int_list(b, &b->kernel->attrs, "cluster_dims", dims, 3); + } + } + /* ---- resolve op + atom + frag widths ---- (811-815) */ ctx->op = rocke_conv_resolve_op(b, spec, ctx->arch); if(ctx->op == NULL) @@ -502,11 +527,35 @@ bool rocke_conv_build_ctx_init(rocke_conv_build_ctx_t* ctx, * * Merged: "group" is the merged group (gridDim.z runs 0..groups/Gm) and * the path stays engaged even at Gm == groups -- see ctx->grouped. */ + rocke_value_t* split_group_z = NULL; + if(ctx->fwd_split) + { + /* Split-K: the group and the K slice share z: z = group * ks_count + + * slice, as in wgrad. Slice s reduces [s*ks, (s+1)*ks); the host rounds + * ks to whole K tiles, and the tail past K_gemm reads zero through the + * descriptor bounds, so every slice runs the same bound computation. */ + rocke_value_t* z_id = rocke_b_block_id_z(b); + if(ctx->grouped) + { + split_group_z = rocke_b_to_sgpr_u32(b, rocke_b_div(b, z_id, ctx->fwd_p_ks_count)); + rocke_value_t* g_ks = rocke_b_mul(b, split_group_z, ctx->fwd_p_ks_count); + ctx->fwd_slice_id = rocke_b_sub(b, z_id, g_ks); + } + else + { + ctx->fwd_slice_id = z_id; + } + rocke_value_t* lo_mul = rocke_b_mul(b, ctx->fwd_slice_id, ctx->fwd_p_ks); + ctx->kloop_k_lo = rocke_b_to_sgpr_u32(b, lo_mul); + rocke_value_t* hi_add = rocke_b_add(b, ctx->kloop_k_lo, ctx->fwd_p_ks); + ctx->kloop_k_hi = rocke_b_to_sgpr_u32(b, hi_add); + } if(ctx->grouped) { /* Python: group_idx = b.block_id_z(); k_out_group_base = b.mul(group_idx, b.const_i32(kpg)) - * Bind subexpressions in Python's left-to-right order to pin SSA ids. */ - ctx->group_idx = rocke_b_block_id_z(b); + * Bind subexpressions in Python's left-to-right order to pin SSA ids. + * Split-K: the group is the z quotient decoded above. */ + ctx->group_idx = ctx->fwd_split ? split_group_z : rocke_b_block_id_z(b); /* AOT: the runtime p_kpg param (was const_i32(p.kpg)). Merged: the * merged kpg, which is Gm exactly (depthwise only), so c_gm. */ ctx->k_out_group_base @@ -895,8 +944,9 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv(rocke_ir_builder_t* b, if(ctx.grouped) { /* Merged: g is the merged group and the merged kpg is Gm, so the - * same expression yields k_out = merged_group*Gm + g_n. */ - rocke_value_t* bid_z = rocke_b_block_id_z(b); + * same expression yields k_out = merged_group*Gm + g_n. Split-K + * reuses the decoded group (z also carries the slice). */ + rocke_value_t* bid_z = ctx.fwd_split ? ctx.group_idx : rocke_b_block_id_z(b); ctx.d_k_out_group_base = rocke_b_mul(b, bid_z, ctx.merged ? ctx.c_gm : ctx.p_kpg); } else @@ -907,6 +957,19 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv(rocke_ir_builder_t* b, ctx.ir_d_pitch_pw = NULL; } + if(ctx.fwd_split) + { + /* Stage 1 stores into slab `slice` of the [ks_count][M][K] fp32 + * workspace. The slab is a packed N[D]HWK tensor like D, so D's + * descriptor addresses it; the epilogue adds only the slab base, and + * every store goes through the workspace resource. */ + rocke_value_t* slab = rocke_b_mul(b, ctx.p_M, ctx.p_K); + ctx.fwd_slab_base = rocke_b_mul(b, ctx.fwd_slice_id, slab); + rocke_conv_buffer_resource_t ws_rsrc; + rocke_conv_make_buffer_resource(b, ctx.fwd_p_ws, ctx.fwd_p_ws_bytes, &ws_rsrc); + ctx.d_rsrc = ws_rsrc.rsrc; + } + /* ---- K-loop driver selection (1276-1347) ---- * * The driver is chosen by K-loop *structure*, not by pipeline name. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_epilogue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_epilogue.cpp index 2cf74f06ef57..384c11a15168 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_epilogue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_epilogue.cpp @@ -18,6 +18,7 @@ */ #include +#include "rocke/helper_rocke.helpers.io.h" /* rocke_b_pack_f32_to */ #include "rocke/instance_conv_implicit_gemm_internal.h" /* ===================================================================== * @@ -103,6 +104,7 @@ typedef struct rocke_conv_d_addr_ctx { const rocke_tensor_descriptor_t* D_desc; rocke_value_t* k_out_group_base; /* NULL => groups==1 */ + rocke_value_t* slab_base; /* split-K workspace slab base; NULL => store into D */ } rocke_conv_d_addr_ctx_t; static rocke_value_t* rocke_conv_d_addr(rocke_ir_builder_t* b, @@ -128,6 +130,10 @@ static rocke_value_t* rocke_conv_d_addr(rocke_ir_builder_t* b, values[1] = k_out; rocke_transforms_descriptor_offset(b, dctx->D_desc, names, values, 2, &off, &valid); + /* Split-K: the slab is a packed N[D]HWK tensor like D; only its base is + * added (Python wraps d_addr: b.add(_slab_base, off)). */ + if(dctx->slab_base != NULL) + off = rocke_b_add(b, dctx->slab_base, off); if(out_valid != NULL) *out_valid = valid; return off; @@ -234,6 +240,7 @@ void rocke_conv_emit_direct_epilogue(rocke_ir_builder_t* b, rocke_conv_d_addr_ctx_t dctx; dctx.D_desc = D_desc; dctx.k_out_group_base = k_out_group_base; + dctx.slab_base = NULL; /* split-K is WMMA-only */ rocke_direct_epilogue_store(b, &epi, accs, @@ -271,6 +278,7 @@ void rocke_conv_emit_direct_epilogue_wmma(rocke_ir_builder_t* b, rocke_value_t* ir_d_pitch_pw, const rocke_tensor_descriptor_t* D_desc, rocke_value_t* k_out_group_base, + rocke_value_t* slab_base, rocke_value_t* bound_m, rocke_value_t* bound_n) { @@ -388,6 +396,9 @@ void rocke_conv_emit_direct_epilogue_wmma(rocke_ir_builder_t* b, values[1] = k_out; rocke_transforms_descriptor_offset( b, D_desc, names, values, 2, &d_off_elems, &valid); + /* Split-K: offset into this slice's workspace slab. */ + if(slab_base != NULL) + d_off_elems = rocke_b_add(b, slab_base, d_off_elems); } d_off_bytes = rocke_b_mul(b, d_off_elems, rocke_b_const_i32(b, elem_bytes)); safe_off = rocke_b_select( @@ -428,6 +439,8 @@ void rocke_conv_emit_cshuffle_epilogue(rocke_ir_builder_t* b, const rocke_mmaop_t* op, const rocke_tensor_descriptor_t* D_desc, rocke_value_t* k_out_group_base, + rocke_value_t* slab_base, + const char* out_dtype, rocke_value_t* bound_m, rocke_value_t* bound_n) { @@ -462,6 +475,14 @@ void rocke_conv_emit_cshuffle_epilogue(rocke_ir_builder_t* b, else max_store_vec = (kpg % 8 == 0) ? 8 : (kpg % 4 == 0) ? 4 : (kpg % 2 == 0) ? 2 : 1; } + /* The stored element type: dtype_d, except the split-K binary, which + * stores fp32 partials. Python's CShuffleEpilogue.from_grid[_op] caps an + * fp32 store at 4 elements (16 bytes); the C helpers learn out_dtype only + * after construction, so the cap is applied here. A no-op for a dtype_d + * fp32 spec, whose store width already derives as <= 4. */ + const char* store_dtype = (out_dtype != NULL) ? out_dtype : spec->dtype_d; + if(store_dtype != NULL && strcmp(store_dtype, "fp32") == 0 && max_store_vec > 4) + max_store_vec = 4; int _war_barriers = (spec->pipeline && strcmp(spec->pipeline, "wavelet") == 0) ? 2 : 1; rocke_cshuffle_epilogue_t epi; if(op != NULL && op->family != NULL && strcmp(op->family, "wmma") == 0) @@ -474,7 +495,10 @@ void rocke_conv_emit_cshuffle_epilogue(rocke_ir_builder_t* b, = rocke_mfma_atom("f16", spec->warp_tile_m, spec->warp_tile_n, spec->warp_tile_k); epi = rocke_cshuffle_epilogue_from_grid(atom, grid, max_store_vec); } - epi.out_dtype = spec->dtype_d; + /* out_dtype: the stored element type. dtype_d, except the split-K binary, + * which stores fp32 partials (the store width above still derives from + * dtype_d, as in Python). */ + epi.out_dtype = store_dtype; epi.no_alias = spec->cshuffle_no_alias; epi.war_barriers = _war_barriers; @@ -503,11 +527,103 @@ void rocke_conv_emit_cshuffle_epilogue(rocke_ir_builder_t* b, rocke_conv_d_addr_ctx_t dctx; dctx.D_desc = D_desc; dctx.k_out_group_base = k_out_group_base; + dctx.slab_base = slab_base; rocke_cshuffle_epilogue_store( b, &epi, accs, num_accs, rocke_conv_d_addr, (void*)&dctx, d_rsrc, bound_m, bound_n); } } +/* ===================================================================== * + * rocke_conv_emit_cluster_fold (Python _emit_cluster_fold) + * + * Fold the cluster's S slabs of this tile into D (spec.split_k_cluster). + * Slice s owns rows [s*R, (s+1)*R) of the tile, R = ceil(TM / S); each + * thread walks 4-channel quads of those rows (D and the slabs are packed + * [M][K]; rocke_conv_fwd_split_k_cluster_problem_ok requires K % 4 == 0 and + * kpg % 4 == 0 at build and launch admission). Every SSA value is + * bound in Python's left-to-right evaluation order. + * ===================================================================== */ +static void rocke_conv_emit_cluster_fold(rocke_conv_build_ctx_t* ctx) +{ + rocke_ir_builder_t* b = ctx->b; + const rocke_implicit_gemm_conv_spec_t* spec = ctx->spec; + const int S = spec->split_k_cluster; + const int block_m = ctx->block_m; + const int block_n = ctx->block_n; + const int rows = (block_m + S - 1) / S; + const int quads_per_row = block_n / 4; + const int n_quads = rows * quads_per_row; + const int threads = rocke_implicit_gemm_conv_spec_block_size(spec); + rocke_value_t* m0 = ctx->block_m_off_v; + rocke_value_t* n0 = ctx->block_n_off_v; + rocke_value_t* n_bound = ctx->merged ? ctx->c_gm : ctx->p_kpg; + rocke_value_t* k_base = ctx->k_out_group_base; + + rocke_b_cluster_barrier(b); + rocke_value_t* c_rows = rocke_b_const_i32(b, rows); + rocke_value_t* row_step = rocke_b_mul(b, ctx->fwd_slice_id, c_rows); + rocke_value_t* row0 = rocke_b_add(b, m0, row_step); + rocke_value_t* slab = rocke_b_mul(b, ctx->p_M, ctx->p_K); + for(int t = 0; t < (n_quads + threads - 1) / threads; ++t) + { + rocke_value_t* c_t = rocke_b_const_i32(b, t * threads); + rocke_value_t* q = rocke_b_add(b, ctx->tid, c_t); + rocke_value_t* c_qpr = rocke_b_const_i32(b, quads_per_row); + rocke_value_t* r = rocke_b_div(b, q, c_qpr); + rocke_value_t* c_qpr2 = rocke_b_const_i32(b, quads_per_row); + rocke_value_t* r_q = rocke_b_mul(b, r, c_qpr2); + rocke_value_t* q_in_row = rocke_b_sub(b, q, r_q); + rocke_value_t* c_four = rocke_b_const_i32(b, 4); + rocke_value_t* c = rocke_b_mul(b, q_in_row, c_four); + rocke_value_t* m = rocke_b_add(b, row0, r); + rocke_value_t* n = rocke_b_add(b, n0, c); + rocke_value_t* c_nq = rocke_b_const_i32(b, n_quads); + rocke_value_t* q_ok = rocke_b_cmp_lt(b, q, c_nq); + rocke_value_t* m_ok = rocke_b_cmp_lt(b, m, ctx->p_M); + rocke_value_t* ok_a = rocke_b_land(b, q_ok, m_ok); + rocke_value_t* c_bm = rocke_b_const_i32(b, block_m); + rocke_value_t* m_end = rocke_b_add(b, m0, c_bm); + rocke_value_t* m_in_tile = rocke_b_cmp_lt(b, m, m_end); + rocke_value_t* n_ok = rocke_b_cmp_lt(b, n, n_bound); + rocke_value_t* ok_b = rocke_b_land(b, m_in_tile, n_ok); + rocke_value_t* ok = rocke_b_land(b, ok_a, ok_b); + + rocke_if_t guard = rocke_b_scf_if(b, ok); + rocke_b_region_enter(b, guard.then_region); + { + rocke_value_t* k_out = (k_base == NULL) ? n : rocke_b_add(b, k_base, n); + rocke_value_t* m_row = rocke_b_mul(b, m, ctx->p_K); + rocke_value_t* off = rocke_b_add(b, m_row, k_out); + /* S <= 15 (rocke_conv_fwd_split_k_available). */ + rocke_value_t* parts[15][4] = {}; + for(int s_i = 0; s_i < S; ++s_i) + { + rocke_value_t* o = off; + if(s_i != 0) + { + rocke_value_t* c_s = rocke_b_const_i32(b, s_i); + rocke_value_t* s_off = rocke_b_mul(b, c_s, slab); + o = rocke_b_add(b, s_off, off); + } + rocke_value_t* v = rocke_b_global_load_vN(b, ctx->fwd_p_ws, o, rocke_f32(), 4, 0); + for(int i = 0; i < 4; ++i) + parts[s_i][i] = rocke_b_vec_extract(b, v, i); + } + rocke_value_t* tot[4]; + for(int i = 0; i < 4; ++i) + tot[i] = parts[0][i]; + for(int s_i = 1; s_i < S; ++s_i) + { + for(int i = 0; i < 4; ++i) + tot[i] = rocke_b_fadd(b, tot[i], parts[s_i][i]); + } + rocke_value_t* packed = rocke_b_pack_f32_to(b, tot, 4, spec->dtype_d); + rocke_b_global_store_vN(b, ctx->D, off, packed, 4, 0); + } + rocke_b_region_leave(b); + } +} + /* ===================================================================== * * rocke_conv_emit_epilogue (Python lines 1349-1377) * @@ -560,14 +676,21 @@ void rocke_conv_emit_epilogue(rocke_conv_build_ctx_t* ctx) ctx->op, ctx->D_desc, ctx->d_k_out_group_base, + ctx->fwd_split ? ctx->fwd_slab_base : NULL, + ctx->fwd_split ? "fp32" : NULL, ctx->p_M, ctx->merged ? ctx->c_gm : ctx->p_kpg); } /* elif op.family == "wmma": */ else if(ctx->op != NULL && ctx->op->family != NULL && strcmp(ctx->op->family, "wmma") == 0) { + /* Split-K stores fp32 partials: Python hands this epilogue the spec + * with dtype_d replaced by fp32 (_epi_spec). */ + rocke_implicit_gemm_conv_spec_t epi_spec = *spec; + if(ctx->fwd_split) + epi_spec.dtype_d = "fp32"; rocke_conv_emit_direct_epilogue_wmma(b, - spec, + &epi_spec, ctx->op, final_accs, num_accs, @@ -581,6 +704,7 @@ void rocke_conv_emit_epilogue(rocke_conv_build_ctx_t* ctx) ctx->ir_d_pitch_pw, ctx->D_desc, ctx->d_k_out_group_base, + ctx->fwd_split ? ctx->fwd_slab_base : NULL, ctx->p_M, ctx->merged ? ctx->c_gm : ctx->p_kpg); } @@ -600,4 +724,11 @@ void rocke_conv_emit_epilogue(rocke_conv_build_ctx_t* ctx) ctx->p_M, ctx->merged ? ctx->c_gm : ctx->p_kpg); } + /* Cluster split-K: after the slab store, fold the cluster's slabs into D + * (Python: only on the stock epilogue path, never under an override). */ + if(!(ctx->ov != NULL && ctx->ov->epilogue_override != NULL) && ctx->fwd_split + && spec->split_k_cluster) + { + rocke_conv_emit_cluster_fold(ctx); + } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp index c957f3a77e1b..eac50da69cfb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp @@ -168,10 +168,11 @@ void rocke_conv_emit_kloop_simple(rocke_conv_build_ctx_t* ctx) /* Start at the slice base, not at a bare zero: wgrad hands the drivers a * runtime k_lo and only the forward conv leaves it NULL. Reading ctx->c0 * here worked while wgrad aliased the two, and silently restarted every - * split-K slice from 0 once they came apart. */ + * split-K slice from 0 once they came apart. The forward split-K binary + * also ends at its slice: kloop_k_hi (wgrad sets it to c_K_gemm). */ for_op = rocke_b_scf_for_iter(b, (ctx->kloop_k_lo != NULL) ? ctx->kloop_k_lo : ctx->c0, - ctx->c_K_gemm, + (ctx->kloop_k_hi != NULL) ? ctx->kloop_k_hi : ctx->c_K_gemm, ctx->c_block_k, iter_args, num_accs, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_spec_descriptors.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_spec_descriptors.cpp index c50bca5f71b2..dfa777220486 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_spec_descriptors.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_spec_descriptors.cpp @@ -300,6 +300,8 @@ rocke_implicit_gemm_conv_spec_t rocke_implicit_gemm_conv_spec_default(void) s.k0_k1_split = false; s.groups = 1; s.group_merge = 1; /* inert: every merged-aware site reduces to the unmerged form */ + s.split_k = 1; /* inert: the plain (unsplit) forward binary */ + s.split_k_cluster = 0; /* #8624 vector-sizes-as-args: default None => auto-select. */ s.has_vector_size_a = false; @@ -383,7 +385,9 @@ rocke_status_t rocke_implicit_gemm_conv_spec_kernel_name(const rocke_implicit_ge * self.acc_epilogue.tag(), * flags={"async": self.async_dma, "noalc": self.cshuffle_no_alias, * "unroll": self.unroll_k, - * f"gm{group_merge}": group_merge > 1}) */ + * f"gm{group_merge}": group_merge > 1, + * "splitk": split_k > 1, + * f"skc{split_k_cluster}": split_k_cluster > 0}) */ char short_buf[128]; char t_buf[48]; char w_buf[32]; @@ -391,9 +395,10 @@ rocke_status_t rocke_implicit_gemm_conv_spec_kernel_name(const rocke_implicit_ge char pe_buf[64]; char tag_buf[256]; char gm_buf[24]; + char skc_buf[24]; const char* parts[6]; - const char* flag_names[4]; - int flag_on[4]; + const char* flag_names[6]; + int flag_on[6]; rocke_status_t st; if(s == NULL || out == NULL) @@ -441,8 +446,15 @@ rocke_status_t rocke_implicit_gemm_conv_spec_kernel_name(const rocke_implicit_ge snprintf(gm_buf, sizeof(gm_buf), "gm%d", s->group_merge); flag_names[3] = gm_buf; flag_on[3] = s->group_merge > 1 ? 1 : 0; - - return rocke_kernel_name_join(s->name, parts, 6, flag_names, flag_on, 4, out, out_cap, NULL); + /* The split-K binary has its own ABI and epilogue; the degree is a kernarg, + * so only the variant reaches the name. The cluster form bakes its degree. */ + flag_names[4] = "splitk"; + flag_on[4] = s->split_k > 1 ? 1 : 0; + snprintf(skc_buf, sizeof(skc_buf), "skc%d", s->split_k_cluster); + flag_names[5] = skc_buf; + flag_on[5] = s->split_k_cluster > 0 ? 1 : 0; + + return rocke_kernel_name_join(s->name, parts, 6, flag_names, flag_on, 6, out, out_cap, NULL); } /* ---- group merging (Gm) -------------------------------------------------- * @@ -506,6 +518,129 @@ int rocke_implicit_gemm_conv_spec_grid_groups(const rocke_implicit_gemm_conv_spe return s->problem.groups / gm; } +/* Mirrors Python fwd_split_k_cluster_problem_reason: the cluster fold walks + * 4-channel quads bounded by their first channel, so K and kpg must be + * multiples of 4. Checked on the build problem and on every launched one. */ +bool rocke_conv_fwd_split_k_cluster_problem_ok(const rocke_conv_problem_t* p, + char* reason, + size_t reason_cap) +{ + const int groups = p->groups > 1 ? p->groups : 1; + const int kpg = p->K / groups; + if(p->K % 4 != 0 || kpg % 4 != 0) + { + if(reason != NULL && reason_cap > 0) + { + snprintf(reason, + reason_cap, + "split_k_cluster folds 4-channel quads: K (%d) and the " + "per-group channels (%d) must be multiples of 4", + p->K, + kpg); + } + return false; + } + return true; +} + +/* ---- split-K forward admission ------------------------------------------ * + * + * Mirrors Python fwd_split_k_reason: the split-K body is the plain WMMA K loop + * (mem / wavelet) run over one slice, with an fp32 store into the workspace + * slab. The paths that drive the K loop differently, remap the grid, merge + * groups, or take the pointwise fast path are not wired for it. */ +bool rocke_conv_fwd_split_k_available(const rocke_implicit_gemm_conv_spec_t* s, + const char* family, + const char* arch, + char* reason, + size_t reason_cap) +{ +#define ROCKE_SPLITK_REJECT(...) \ + do \ + { \ + if(reason != NULL && reason_cap > 0) \ + { \ + snprintf(reason, reason_cap, __VA_ARGS__); \ + } \ + return false; \ + } while(0) + + if(s->split_k < 1) + { + ROCKE_SPLITK_REJECT("split_k must be >= 1 (got %d)", s->split_k); + } + if(family == NULL || strcmp(family, "wmma") != 0) + { + ROCKE_SPLITK_REJECT("split-K forward is WMMA-only (built and tested on wave32)"); + } + if(!rocke_conv_acc_epilogue_is_identity(&s->acc_epilogue)) + { + ROCKE_SPLITK_REJECT("split-K forward supports only the identity acc_epilogue: the " + "partial sums are reduced after the per-slice accumulation"); + } + { + const struct + { + bool on; + const char* label; + } flags[] = { + {s->async_dma, "async_dma"}, + {s->unroll_k, "unroll_k"}, + {s->chiplet_swizzle, "chiplet_swizzle"}, + {s->group_merge > 1, "group_merge"}, + {rocke_conv_problem_is_pointwise(&s->problem), "the pointwise fast path"}, + {s->problem.is_3d, "3-D convolution"}, + }; + for(size_t i = 0; i < sizeof(flags) / sizeof(flags[0]); ++i) + { + if(flags[i].on) + { + ROCKE_SPLITK_REJECT("split-K forward does not support %s", flags[i].label); + } + } + } + if(s->epilogue != NULL && strcmp(s->epilogue, "default") == 0 && s->has_vector_size_c + && s->vector_size_c > 1) + { + ROCKE_SPLITK_REJECT("split-K forward with the default epilogue stores scalars " + "(vector_size_c must be 1)"); + } + if(s->split_k_cluster) + { + if(s->split_k < 2) + { + ROCKE_SPLITK_REJECT("split_k_cluster needs the split-K binary (split_k > 1)"); + } + if(s->split_k_cluster < 2 || s->split_k_cluster > 15) + { + ROCKE_SPLITK_REJECT("split_k_cluster must be 2..15 (got %d)", s->split_k_cluster); + } + if(s->pipeline == NULL || strcmp(s->pipeline, "mem") != 0) + { + ROCKE_SPLITK_REJECT("split_k_cluster is wired for the mem pipeline only"); + } + if(s->tile_n % 4) + { + ROCKE_SPLITK_REJECT("split_k_cluster folds 4-channel quads (tile_n %% 4 == 0)"); + } + if(arch != NULL) + { + const rocke_archtarget_t* target = rocke_archtarget_from_gfx(arch); + if(target != NULL && !target->memory.has_cluster_launch) + { + ROCKE_SPLITK_REJECT( + "split_k_cluster needs workgroup-cluster launch, which %s lacks", target->gfx); + } + } + if(!rocke_conv_fwd_split_k_cluster_problem_ok(&s->problem, reason, reason_cap)) + { + return false; + } + } + return true; +#undef ROCKE_SPLITK_REJECT +} + bool rocke_conv_fwd_group_merge_available(const rocke_implicit_gemm_conv_spec_t* s, const char* arch, char* reason, @@ -1054,7 +1189,9 @@ bool rocke_implicit_gemm_conv_is_valid_spec(const rocke_implicit_gemm_conv_spec_ is_cshuffle = (s->epilogue && strcmp(s->epilogue, "cshuffle") == 0); { - int c_dtype_bytes = (s->dtype_d && strcmp(s->dtype_d, "fp32") == 0) ? 4 : 2; + /* The split-K binary stages its fp32 partial tile, whatever dtype_d is. */ + int c_dtype_bytes + = ((s->dtype_d && strcmp(s->dtype_d, "fp32") == 0) || s->split_k > 1) ? 4 : 2; c_lds = is_cshuffle ? (s->tile_m * s->tile_n * c_dtype_bytes) : 0; } { @@ -1133,6 +1270,14 @@ bool rocke_implicit_gemm_conv_is_valid_spec(const rocke_implicit_gemm_conv_spec_ * group-aware descriptor/epilogue are family-neutral (Python PR #10064). */ } + if(s->split_k != 1 || s->split_k_cluster) + { + if(!rocke_conv_fwd_split_k_available(s, family, arch, reason, reason_cap)) + { + return false; + } + } + /* The tile loaders (Python: the async_tile_loaders / _sync_load_vecs block * at the end of is_valid_spec). An explicit vector width, or the async * chunk width the tile admits, has to split the tile evenly over the diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md index 6b0a3d2a20c9..bf4f47dccd89 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md @@ -285,6 +285,25 @@ for each thread's coalesced output chunk: buffer_store_vN_f16(d_rsrc, d_off*2, 0, v, dwords=4) ``` +### Split-K Forward (gfx1250) + +`ImplicitGemmConvSpec.split_k > 1` builds the split-K forward binary (WMMA only; `mem` or `wavelet` pipeline, 2-D, not pointwise, no `async_dma` / `unroll_k` / `chiplet_swizzle` / `group_merge`, identity `acc_epilogue` since the reduction only adds the raw partials; checked by `fwd_split_k_reason`). It is one binary for every degree: the ABI (`conv_arg_names("fwd", two_stage=True)`) appends `ws_ptr`, `ws_bytes`, `ks` and `ks_count`, and the kernel name gains a `splitk` tag. + +- The group and the K slice share `blockIdx.z = group * ks_count + slice`. Slice `s` reduces `[s*ks, (s+1)*ks)` of the `[Z*]Y*X*cpg` reduction; the host rounds `ks` to whole K tiles (`ConvArgs`), and the tail past `K_gemm` reads zero through the descriptor bounds. +- The epilogue stores the fp32 partial tile into slab `slice` of a `[ks_count][M][K]` fp32 workspace, addressed by the D descriptor plus the slab base. Every slab element is written by exactly one workgroup, so the workspace needs no clearing. A cshuffle C tile is staged in fp32. +- Stage 2, `conv_fwd_split_k_reduce` (`FwdSplitKReduceSpec(dtype_d, vec, block_size, degree)`), sums the slabs element-wise in a fixed order and casts into D, so the result is deterministic. `degree > 0` unrolls the fold for exactly that many slabs; `0` loops over `ks_count`. +- `split_k_cluster = S` (2..15, `mem` pipeline, `tile_n % 4 == 0`, K and kpg multiples of 4, a target with workgroup-cluster launch, llvm23) is the single-launch form: the S slices of one output tile form a `(1, 1, S)` workgroup cluster (`cluster_dims` attribute; the launch passes the same shape), store their slabs, meet at a cluster barrier, and each slice folds `ceil(tile_m / S)` rows of the tile across the S slabs into D in 4-channel quads. The degree is baked (`skc` tag, `ks_count == S` at launch) and there is no Stage 2. + +`split_k = -1` (auto) is resolved before validation by the conv split-K resolver (`resolve_conv_split_k("fwd", ...)`), as for wgrad and dgrad. Its degree comes from `rocke.helpers.split_k.select_split_k_fwd`: +- It splits only when the unsplit grid is under half of 8 workgroups per CU and the reduction is at least 512 deep. +- The degree is a power of two toward that target grid, at most 32, and keeps at least 2 K tiles per slice. +- It is then held to the gridDim.z budget and the i32 workspace. +- It resolves to 1 (the plain kernel) elsewhere, including on targets other than gfx1250 and on pointwise or 3-D problems. + +The partials meet in fp32 slabs, so the degree has no accuracy limit. The C++ engine takes explicit degrees only; `rocke.core.backend.lower_conv_implicit_gemm` resolves `-1` first. + +Both forms are in the gfx1250 AOT kernel cache (`kernels/conv/aot_jobs.py`; the cluster form only at llvm23) and in `benchmark_implicit_gemm_conv --run-from-cache`, where `--split-k -1` (the default) launches a cached split-K binary at the resolved degree and `0` adds `--split-k-degrees`. The grouped-convolution dispatcher does not select them. `split_k = 1` emits exactly the plain kernel. + ### Step-By-Step Trace (one workgroup) ```text diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md index dd2695b11e5e..7dc23c7e1b8b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md @@ -40,7 +40,8 @@ implementation. | File | Spec | Doc | |-----------------------------------|-------------------------------------------------------------------|------------------------------| -| `conv_implicit_gemm.py` | `ConvProblem`, `ImplicitGemmConvSpec` | `instances/convolution.md` | +| `conv_implicit_gemm.py` | `ConvProblem`, `ImplicitGemmConvSpec` (incl. split-K forward: `split_k`, `split_k_cluster`) | `instances/convolution.md` | +| `conv_fwd_split_k_reduce.py` | `FwdSplitKReduceSpec` (split-K forward Stage 2) | `instances/convolution.md` | | `conv_implicit_gemm_dgrad.py` | `DgradConvSpec` (backward data; tilde and stride-1 binaries, `tap_outer_k`, `dy_halo`) | `instances/convolution.md`, [`README_conv_implicit_gemm_dgrad.md`](../../../library/kernels/common/README_conv_implicit_gemm_dgrad.md) | | `conv_direct_grouped.py` | `DirectConvProblem`, `DirectConv16cSpec`, `DirectConv4cSpec` (incl. 4c dgrad), `DirectConvSpec` (incl. fused-weight dgrad), `DirectDepthwiseDgradWindowedSpec` | `instances/convolution.md`, [`README_conv_direct_grouped.md`](../../../library/kernels/common/README_conv_direct_grouped.md) | | `conv_direct_nongrouped.py` | `DirectNongroupedConvSpec` (`groups == 1`) | `instances/convolution.md` | @@ -58,7 +59,7 @@ implementation. | `img2col.py` | `Img2ColSpec` | `instances/convolution.md` | | `pooling.py` | `PoolingProblem`, `Pooling2DSpec`, `PoolOp` | `instances/convolution.md` | -ABI: conv kernels are ahead-of-time compiled and shape-generic. Every implicit-GEMM / direct grouped conv kernel takes `(A, B, D, A_bytes, B_bytes, D_bytes)` followed by the problem block -- extents, strides, padding, dilation, groups and the host-computed magic-division constants -- whose names and order come from `kernels.common.conv_abi` (`conv_arg_names(direction)` / `conv_direct_arg_names(direction)`; C++ twin `instance_conv_abi.h`). Hosts fill them with `kernels.common.conv_args.ConvArgs.from_problem(...).to_launch_values(...)`; wgrad appends the split-K `ks` / `ks_count` pair (and the two-stage `ws_ptr` / `ws_bytes`), dgrad its sub-GEMM table. See `instances/convolution.md`. Img2col writes `Y`. Pooling reads `X` and writes `Y`. +ABI: conv kernels are ahead-of-time compiled and shape-generic. Every implicit-GEMM / direct grouped conv kernel takes `(A, B, D, A_bytes, B_bytes, D_bytes)` followed by the problem block -- extents, strides, padding, dilation, groups and the host-computed magic-division constants -- whose names and order come from `kernels.common.conv_abi` (`conv_arg_names(direction)` / `conv_direct_arg_names(direction)`; C++ twin `instance_conv_abi.h`). Hosts fill them with `kernels.common.conv_args.ConvArgs.from_problem(...).to_launch_values(...)`; wgrad appends the split-K `ks` / `ks_count` pair (and the two-stage `ws_ptr` / `ws_bytes`), the split-K forward (`conv_arg_names("fwd", two_stage=True)`) appends `ws_ptr` / `ws_bytes` / `ks` / `ks_count`, dgrad its sub-GEMM table. See `instances/convolution.md`. Img2col writes `Y`. Pooling reads `X` and writes `Y`. Layouts: NHWC input, KYXC weight, NHWK output for conv. Grouping via `cpg`/`kpg`. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index 809e95062012..f0dc92ca5f06 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -2017,6 +2017,8 @@ Per-arch facts belong in the arch references — see §21. | `blocks_per_cu` | `StreamKGemmSpec` | 1 | Persistent dispatch density | | `reduction` | `StreamKGemmSpec` | `Atomic` | `StreamKReductionStrategy.{Atomic, Reduction, AtomicWithFixup}` | | split-K | GEMM (not yet a spec field; available primitives) | — | `b.global_atomic_add_f32` for atomic split-K; `helpers/streamk.py` for the Stream-K macro tile decode | +| `split_k` | `ImplicitGemmConvSpec` (fwd, gfx1250 WMMA) | 1 | `> 1` builds the one split-K forward binary (`splitk` tag); the degree is the `ks_count` kernarg. Slices store fp32 partial tiles into a `[ks_count][M][K]` workspace and `conv_fwd_split_k_reduce` (`degree` unrolls its fold) sums them deterministically. For small-M layers whose plain grid cannot fill the device; costs the workspace traffic and a second launch. `-1` resolves the degree (`select_split_k_fwd`; 1 where splitting does not pay); `--split-k 0 --split-k-degrees` sweeps others. Not with `async_dma`, `unroll_k`, `chiplet_swizzle`, `group_merge`, pointwise or 3-D | +| `split_k_cluster` | `ImplicitGemmConvSpec` (fwd, gfx1250, llvm23) | 0 | `S` in 2..15 folds the S slices of a tile through a `(1, 1, S)` workgroup cluster and a cluster barrier instead of a Stage 2 launch (`skc` tag, one binary per S). Needs `split_k > 1`, `mem`, `tile_n % 4 == 0`; launch with the same cluster shape | | Grouped GEMM persistent | `instances/grouped_gemm.py` | per-group launches | v2 persistent variant is a documented follow-up | | `streamk` | `WgradConvSpec` | `"off"` | `"dp_sk"` launches one CTA per data-parallel tile plus `streamk_ctas` stream-K CTAs; `"persistent"` launches only the pool and round-robins the data-parallel tiles. Both balance MAC iterations across CTAs instead of tiles, so a shape with few tiles and a long `wg_K` still fills the device. Mutually exclusive with `split_k` | | `streamk_reduction` | `WgradConvSpec` | `"linear"` | How CTAs sharing a tile combine partials. `"linear"`: the owner waits on each later contributor in turn and stores once. `"tree"`: pairwise fan-in, `ceil(log2)` rounds. Both are deterministic and need no dW pre-zeroing. `"atomic"`: f32 atomics into dW (fp32 only). `"workspace"`: f32 scratch atomics plus the Stage 2 reduce kernel | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/backend.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/backend.py index a7a7399fdfbd..56c605e971ff 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/backend.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/backend.py @@ -955,6 +955,17 @@ def _conv_problem_to_dict(p: Any) -> Dict[str, Any]: return d +def _conv_acc_epilogue_to_dict(epi: Any) -> Dict[str, Any]: + """:class:`ConvAccumulatorEpilogue` -> dict (clamp bounds None when unset).""" + return dict( + bias=float(epi.bias), + scale=float(epi.scale), + relu=bool(epi.relu), + clamp_min=None if epi.clamp_min is None else float(epi.clamp_min), + clamp_max=None if epi.clamp_max is None else float(epi.clamp_max), + ) + + def conv_implicit_gemm_spec_to_dict(spec: Any) -> Dict[str, Any]: """:class:`ImplicitGemmConvSpec` -> flat dict (problem nested). Optional fields (lds_k_pad/waves_per_eu/vector_size_*) are forwarded when set; the @@ -984,6 +995,16 @@ def conv_implicit_gemm_spec_to_dict(spec: Any) -> Dict[str, Any]: # Changes the emitted kernel and its launch grid; dropped, the engine # would build an unmerged kernel under a merged spec's name and grid. group_merge=getattr(spec, "group_merge", 1), + # The split-K forward binary (its own ABI and epilogue) and its cluster + # form; dropped, the engine would build the plain kernel. Read + # directly: a spec without them is malformed, not unsplit. + split_k=spec.split_k, + split_k_cluster=spec.split_k_cluster, + # The wavelet load-wave count changes the launch block and the loaders. + num_load_waves=spec.num_load_waves, + # The fused accumulator epilogue changes the emitted epilogue (and + # split-K admission); dropped, the engine would build the identity. + acc_epilogue=_conv_acc_epilogue_to_dict(spec.acc_epilogue), ) # dtype_* and the vector-size/optional knobs are merge-target additions # (#8624); forward them only when this branch's spec exposes them so the @@ -1840,6 +1861,12 @@ def py_fn(wi: bool) -> Tuple[str, str]: ir = serialize(k) return ll, ir + if spec.split_k == -1: + # Auto split-K is resolved on the Python side for both engines, so the + # C++ engine (explicit degrees only) builds the same kernel. + from kernels.common.conv_implicit_gemm import resolve_fwd_split_k_spec + + spec = resolve_fwd_split_k_spec(spec, arch) eng = _import_engine() sd = conv_implicit_gemm_spec_to_dict(spec) return _lower_family( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py index 5c9c8a2d77a8..9eda8cdc95f7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py @@ -180,6 +180,25 @@ def atomic_rounding_degree_cap(fraction_bits: int, tol: float = WGRAD_TOLERANCE) _ENV_FLAG = "ROCKE_GEMM_SPLIT_K" +# Forward split-K (the two-stage WMMA implicit-GEMM conv on gfx1250, see +# select_split_k_fwd). Fitted on gfx1250 measurements of every power-of-two +# degree of the cached split-K forward tiles over small-output / deep-reduction +# layers: +# +# * the best degree puts about this many workgroups on each CU -- the short +# K slices of a small tile are latency-bound, so several resident slices per +# CU still pay; +FWD_SPLIT_K_TARGET_CTAS_PER_CU = 8 +# * a slice of fewer K tiles than this is mostly prologue and epilogue (and +# one more fp32 slab for Stage 2 to read); +FWD_SPLIT_K_MIN_K_TILES_PER_SLICE = 2 +# * past this degree the slabs, not the K loop, dominate on every measured +# layer; +FWD_SPLIT_K_MAX = 32 +# * a reduction shallower than TARGET_SLICE_K elements is latency-bound as a +# whole: the Stage 2 launch costs more than splitting it saves. +FWD_SPLIT_K_MIN_GEMM_K = TARGET_SLICE_K + @dataclass(frozen=True) class SplitKDecision: @@ -491,3 +510,82 @@ def select_split_k_wgrad( f"snapped to a multiple of {step} = {split_k} " f"(num_cus={num_cus} waves_per_block={waves_per_block} xcds={xcds})", ) + + +def _pow2_floor(x: int) -> int: + return 1 << (max(1, int(x)).bit_length() - 1) + + +def select_split_k_fwd( + *, + M: int, + N: int, + K: int, + tile_m: int, + tile_n: int, + tile_k: int, + groups: int = 1, + arch: str = "gfx1250", + num_cus: int | None = None, +) -> SplitKDecision: + """Pick a split-K degree for a forward implicit-GEMM convolution. + + ``M`` / ``N`` / ``K`` are one group's GEMM (``N*Ho*Wo`` x ``kpg`` x + ``Y*X*cpg``); ``groups`` multiplies the grid. The forward split reduces + through fp32 workspace slabs (one writer per element, folded by Stage 2 or + the cluster), so unlike the 16-bit atomic reductions its accuracy does not + depend on the degree: the degree is a performance choice only. + + ``num_cus`` is the device's CU count; ``None`` takes the arch's from + ``_ARCH_NUM_CUS``, else the 256-CU default (gfx1250's own count, the arch + the policy was fitted on). :: + + base_grid = ceil(M / tile_m) * ceil(N / tile_n) * groups + target = num_cus * FWD_SPLIT_K_TARGET_CTAS_PER_CU + # engage gate: only a grid that leaves the device mostly idle, over a + # reduction deep enough to pay for the Stage 2 launch + if base_grid >= target / 2 or K < FWD_SPLIT_K_MIN_GEMM_K: split_k = 1 + split_k = pow2_floor(target / base_grid), held to + FWD_SPLIT_K_MAX and to pow2_floor(k_tiles / + FWD_SPLIT_K_MIN_K_TILES_PER_SLICE) + + The degree stays on the power-of-two ladder the policy was fitted on. It is + not snapped to a divisor of the K tiles: the slice width is rounded up to + ``tile_k`` and the tail reads zeros, so every degree is legal. Launch + limits that depend on the binary (gridDim.z, the i32 workspace) are the + caller's (:func:`kernels.common.conv_split_k.resolve_conv_split_k`). + """ + if num_cus is None: + num_cus = _ARCH_NUM_CUS.get(arch.split(":", 1)[0], _DEFAULT_NUM_CUS) + target = num_cus * FWD_SPLIT_K_TARGET_CTAS_PER_CU + base_grid = ( + ceil_div(int(M), int(tile_m)) * ceil_div(int(N), int(tile_n)) * max(1, groups) + ) + if base_grid <= 0: + return SplitKDecision(1, base_grid, target, "empty grid") + if 2 * base_grid >= target: + return SplitKDecision( + 1, base_grid, target, f"grid {base_grid} already fills the device" + ) + if K < FWD_SPLIT_K_MIN_GEMM_K: + return SplitKDecision( + 1, + base_grid, + target, + f"reduction {K} < {FWD_SPLIT_K_MIN_GEMM_K} is too shallow to split", + ) + k_tiles = ceil_div(int(K), int(tile_k)) + depth_cap = _pow2_floor(k_tiles // FWD_SPLIT_K_MIN_K_TILES_PER_SLICE) + split_k = min(_pow2_floor(target // base_grid), FWD_SPLIT_K_MAX, depth_cap) + if split_k < 2: + return SplitKDecision( + 1, base_grid, target, f"{k_tiles} K tiles leave no room to split" + ) + return SplitKDecision( + split_k, + base_grid, + target, + f"fwd fill: grid {base_grid} toward {target} workgroups " + f"({FWD_SPLIT_K_TARGET_CTAS_PER_CU}/CU on {num_cus} CUs), " + f"{ceil_div(k_tiles, split_k)} K tiles per slice -> split_k {split_k}", + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md index ba15ba741ca8..f8713936809c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md @@ -109,6 +109,9 @@ be smaller than 2 GiB. | `conv_implicit_gemm_dgrad` stride-1 knobs | ✅ | ✅ | ❌ | folded sub-GEMM record (stride capability), `tap_outer_k` (wave64 MFMA, `kpg % tile_k == 0`), `dy_halo` / `dy_halo_setprio` / `dy_halo_kouter_pad` / `halo_max_wo` (on the tap-outer path, `Ho == Hi`, `Wo == Wi`; the pad needs `lds_k_outer`, gfx950); `xcd_tile_order` remaps only on gfx950. AOT grid: stride-1 tap-outer binaries next to the tilde ones, halo sub-grid on gfx950; not yet wired into dispatch | | `conv_implicit_gemm_wgrad` | ✅ | ✅ | ✅ | gfx1151/gfx1201/gfx1250: WMMA, `mem`+`default`; split-K (WMMA packed atomic, no two-stage) on gfx1250 only, automatic for bf16/fp32 and explicit-only for fp16, `split_k=1` on gfx1151/gfx1201 | | `conv_implicit_gemm_wgrad` stream-K | ✅ | ✅ | ❌ | CK iteration-balanced partition; `linear`/`tree` deterministic fixups need the cross-workgroup flag ops | +| `conv_implicit_gemm` split-K forward (`split_k > 1`) | ❌ | ❌ | ❌ | gfx1250 only: WMMA, `mem` or `wavelet` pipeline, 2-D, not pointwise, identity `acc_epilogue`; no `async_dma`, `unroll_k`, `chiplet_swizzle` or `group_merge`. Each slice stores an fp32 partial tile into its slab of a `[ks_count][M][K]` workspace and `conv_fwd_split_k_reduce` folds the slabs into D; the degree is the `ks_count` kernarg, and `split_k = -1` resolves it (`resolve_conv_split_k`; 1 where splitting does not pay). Built into the AOT kernel cache; the dispatcher does not select it | +| `conv_implicit_gemm` cluster split-K (`split_k_cluster`) | ❌ | ❌ | ❌ | gfx1250 only (llvm23, workgroup-cluster launch): `mem` pipeline, `split_k_cluster` in 2..15, `tile_n % 4 == 0`, K and kpg multiples of 4 (checked at build and on every launched problem). The slices of one output tile form a (1, 1, S) workgroup cluster and fold through a cluster barrier, so there is no Stage 2 launch; one binary per S. Built into the AOT kernel cache; the dispatcher does not select it | +| `conv_fwd_split_k_reduce` | ❌ | ❌ | ❌ | Stage 2 of the gfx1250 split-K forward: sums the fp32 slabs in a fixed order and casts to fp16/bf16 (deterministic); `degree > 0` unrolls the fold for that many slabs. Plain global loads/stores, no arch-specific ops | | `conv_direct_nongrouped` | ✅ | ✅ | ❌ | `groups=1`; gfx942 uses 32x32x8 / 16x16x16 atoms, gfx950 also 32x32x16 / 16x16x32; MFMA-only; the only direct kernel with dilation. Built from `library/kernels`; not yet wired into dispatch | | `direct_wgrad_tiled` | ❌ | ✅ | ❌ | dense and grouped wgrad, stride 1; needs `ds_read_b64_tr_b16` and `mfma_f32_16x16x32`. Built from `library/kernels`; AOT sweep only, not yet wired into dispatch | | `dw_wgrad_lds` / `dw_wgrad_cols` / `direct_wgrad_gather` | ❌ | ✅ | ❌ | depthwise (`C % 8 == 0` / `C <= 8`) and dense small-C wgrad, partial slots + `dw_wgrad_fold`; validated on gfx950 only (the LDS kernels DMA 16 bytes per lane). AOT sweep only, not yet wired into dispatch | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py index 5580c573e67d..abcde0f9e99f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py @@ -687,6 +687,43 @@ def _enqueue( ) return () + cfg, attr = Runtime._cluster_launch_config( + grid, block, shared_bytes, stream, cluster, where + ) + Runtime._launch_cluster(ctypes.byref(cfg), fn, params, extra, where) + return (cfg, attr) + + @staticmethod + def _launch_cluster( + cfg_ref: Any, fn: _HipFunctionHandle, params: Any, extra: Any, where: str + ) -> None: + """Issue a cluster launch built by :meth:`_cluster_launch_config`. + + The one ``hipDrvLaunchKernelEx`` call site: :meth:`_enqueue` and the + callables of :meth:`prepare_launch` both go through it, so the two + paths cannot drift apart. + """ + _check( + _hipDrvLaunchKernelEx(cfg_ref, fn, params, extra), + f"hipDrvLaunchKernelEx({where})", + ) + + @staticmethod + def _cluster_launch_config( + grid: Tuple[int, int, int], + block: Tuple[int, int, int], + shared_bytes: int, + stream: int, + cluster: Tuple[int, int, int], + where: str, + ) -> Tuple[Any, Any]: + """Validate ``cluster`` and build the ``hipDrvLaunchKernelEx`` config. + + Returns ``(cfg, attr)``: the launch config and the cluster-dimension + attribute it points at, which must stay alive together. Raises before + HIP sees anything for a bad shape (see :func:`check_launch_cluster`) + or a device that reports no cluster launch support. + """ shape = check_launch_cluster(grid, cluster) device = current_device() if device_supports_cluster_launch(device) is False: @@ -709,11 +746,7 @@ def _enqueue( attrs=ctypes.pointer(attr), numAttrs=1, ) - _check( - _hipDrvLaunchKernelEx(ctypes.byref(cfg), fn, params, extra), - f"hipDrvLaunchKernelEx({where})", - ) - return (cfg, attr) + return cfg, attr def launch( self, @@ -806,10 +839,14 @@ def prepare_launch( *, shared_bytes: int = 0, stream: int = 0, + cluster: Optional[Tuple[int, int, int]] = None, ) -> "Callable[[], None]": """Build every ctypes argument of one launch once; return a zero-argument callable that only enqueues it. + ``cluster`` launches workgroup clusters of that shape, as in + :meth:`launch` (``hipDrvLaunchKernelEx`` with a cluster attribute). + For loops that launch the same kernel with the same arguments over and over -- the timed loop of a benchmark. :meth:`launch` rebuilds the args buffer, the ``extra`` array and the scalar @@ -845,6 +882,20 @@ def prepare_launch( ) keep = (args_buf, size_buf, extra) + if cluster is not None: + # Validated, device-checked and built once, like the plain path's + # ctypes arguments: the callable only enqueues, through the same + # launch helper as :meth:`_enqueue`. + cfg, attr = self._cluster_launch_config( + grid, block, shared_bytes, int(stream), cluster, "prepare_launch" + ) + cfg_ref = ctypes.byref(cfg) + + def enqueue_cluster(_keep=keep + (cfg, attr)) -> None: + Runtime._launch_cluster(cfg_ref, fn, None, extra, "prepare_launch") + + return enqueue_cluster + def enqueue(_keep=keep) -> None: _check(_hipModuleLaunchKernel(*call_args), "hipModuleLaunchKernel") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py index 0fcc8a7bf8f8..748aa22175af 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py @@ -557,6 +557,7 @@ def bind( self._packer(values), shared_bytes=config.shared_bytes, stream=resolve_stream(config.stream), + cluster=self._resolve_cluster(config), ) keep = tuple(values.values()) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index 4944ae06f04f..940343db32c7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -773,6 +773,22 @@ "sha256": "960dde16f771e8e06fe7eebf29f1a9e479847318e4af5c0837c7e9aba97d2baa", "status": "ok" }, + "conv/gfx1250/n1h7c256k64r3/splitk_cshuffle": { + "arch": "gfx1250", + "bytes": 42631, + "family": "conv", + "kernel_name": "irhash_conv_1250_splitk_cshuffle_N1H7W7C256_K64Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_splitk", + "sha256": "f82e7e69e5b2facf7c90fa2b537d9c6edfdcf26d7e0b7776404067695ea077db", + "status": "ok" + }, + "conv/gfx1250/n1h7c256k64r3/splitk_default": { + "arch": "gfx1250", + "bytes": 53537, + "family": "conv", + "kernel_name": "irhash_conv_1250_splitk_default_N1H7W7C256_K64Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_splitk", + "sha256": "d6ecd55e89f1baa7180002c9a3bd5c2d0f9488828e8d63ad58f240b658433856", + "status": "ok" + }, "conv/gfx1250/n2h14c64k64r3g4/g4_cshuffle": { "arch": "gfx1250", "bytes": 38016, @@ -1541,6 +1557,22 @@ "sha256": "9dacd91ff5ea1ba3156c9c97eebedbb5b0a764ddd616e7326bd9185f9cfdaaab", "status": "ok" }, + "conv_fwd_split_k_reduce/gfx1250/bf16_v4": { + "arch": "gfx1250", + "bytes": 3339, + "family": "conv_fwd_split_k_reduce", + "kernel_name": "conv_fwd_splitk_reduce_v4_b256_bf16", + "sha256": "9f5d2e0e9448f4b577f054434b9f1ef1d23bc2f3ea573b6e034515f862199af2", + "status": "ok" + }, + "conv_fwd_split_k_reduce/gfx1250/fp16_v2_s4": { + "arch": "gfx1250", + "bytes": 2802, + "family": "conv_fwd_split_k_reduce", + "kernel_name": "conv_fwd_splitk_reduce_v2_b256_s4_fp16", + "sha256": "e58fd98cfb623117fc877b644243ba385014a8deb4af1f1b56df2d5f5fe1bfa9", + "status": "ok" + }, "conv_wgrad/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 53578, @@ -2479,6 +2511,10 @@ } }, "expected_failures": { + "conv/gfx1250/n1h7c256k64r3/splitk_cluster4": { + "message": "cluster_barrier requires LLVM flavor llvm23, got llvm20", + "type": "ValueError" + }, "conv_dgrad_patch_wmma/gfx1250/bf16_p16_rb8": { "message": "tensor_load_to_lds requires LLVM flavor llvm23, got llvm20", "type": "ValueError" @@ -2598,24 +2634,25 @@ "gfx1151": 10, "gfx12-generic": 1, "gfx1201": 8, - "gfx1250": 29, + "gfx1250": 33, "gfx90a": 5, "gfx942": 60, "gfx950": 195 }, - "expected_failures": 28, + "expected_failures": 29, "families": { "attention": 39, "attention_d128_swa": 1, "attention_d256": 2, "attention_dense": 51, - "conv": 20, + "conv": 22, "conv_dgrad": 13, "conv_direct": 27, "conv_direct_depthwise_tiled": 5, "conv_direct_dgrad": 15, "conv_direct_nongrouped": 7, "conv_direct_wgrad": 12, + "conv_fwd_split_k_reduce": 2, "conv_wgrad": 22, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, @@ -2633,7 +2670,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 309 + "ok_cases": 313 } }, "llvm22": { @@ -3406,6 +3443,22 @@ "sha256": "e958db8b566b938a9874b69ebdb6bbfd38c084cbeb82fc5841e37de81bb7d362", "status": "ok" }, + "conv/gfx1250/n1h7c256k64r3/splitk_cshuffle": { + "arch": "gfx1250", + "bytes": 42813, + "family": "conv", + "kernel_name": "irhash_conv_1250_splitk_cshuffle_N1H7W7C256_K64Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_splitk", + "sha256": "95c9afa6d4c4921ddcb8a46d22567acba3ade401e690f437cd2f6ff3cc878f3c", + "status": "ok" + }, + "conv/gfx1250/n1h7c256k64r3/splitk_default": { + "arch": "gfx1250", + "bytes": 53719, + "family": "conv", + "kernel_name": "irhash_conv_1250_splitk_default_N1H7W7C256_K64Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_splitk", + "sha256": "d58d5823bb64160c3f0aa43536030b8f47a7bd58f9fd7bbf214ee65e2ea7ddab", + "status": "ok" + }, "conv/gfx1250/n2h14c64k64r3g4/g4_cshuffle": { "arch": "gfx1250", "bytes": 38155, @@ -4174,6 +4227,22 @@ "sha256": "1fdeb533d3a975eeaa7431efb89457bcce602ad7e865ac20677b5aa35eed6bb1", "status": "ok" }, + "conv_fwd_split_k_reduce/gfx1250/bf16_v4": { + "arch": "gfx1250", + "bytes": 3346, + "family": "conv_fwd_split_k_reduce", + "kernel_name": "conv_fwd_splitk_reduce_v4_b256_bf16", + "sha256": "1224e46c4cc351f7d63cdf21cbbfb446eaec8e64fe5460148da7bb543ac2e689", + "status": "ok" + }, + "conv_fwd_split_k_reduce/gfx1250/fp16_v2_s4": { + "arch": "gfx1250", + "bytes": 2809, + "family": "conv_fwd_split_k_reduce", + "kernel_name": "conv_fwd_splitk_reduce_v2_b256_s4_fp16", + "sha256": "d3f6b4e3314c207f3cc441e6e1015d340e5a77e0fc48fc2bbc31d53edbf3be75", + "status": "ok" + }, "conv_wgrad/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 53717, @@ -5112,6 +5181,10 @@ } }, "expected_failures": { + "conv/gfx1250/n1h7c256k64r3/splitk_cluster4": { + "message": "cluster_barrier requires LLVM flavor llvm23, got llvm22", + "type": "ValueError" + }, "conv_dgrad_patch_wmma/gfx1250/bf16_p16_rb8": { "message": "tensor_load_to_lds requires LLVM flavor llvm23, got llvm22", "type": "ValueError" @@ -5231,24 +5304,25 @@ "gfx1151": 10, "gfx12-generic": 1, "gfx1201": 8, - "gfx1250": 29, + "gfx1250": 33, "gfx90a": 5, "gfx942": 60, "gfx950": 195 }, - "expected_failures": 28, + "expected_failures": 29, "families": { "attention": 39, "attention_d128_swa": 1, "attention_d256": 2, "attention_dense": 51, - "conv": 20, + "conv": 22, "conv_dgrad": 13, "conv_direct": 27, "conv_direct_depthwise_tiled": 5, "conv_direct_dgrad": 15, "conv_direct_nongrouped": 7, "conv_direct_wgrad": 12, + "conv_fwd_split_k_reduce": 2, "conv_wgrad": 22, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, @@ -5266,7 +5340,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 309 + "ok_cases": 313 } }, "llvm23": { @@ -6039,6 +6113,30 @@ "sha256": "d94f658d1c5dc893acc5695ec58281b4d3d148ede0597402ede5c6a10e230e16", "status": "ok" }, + "conv/gfx1250/n1h7c256k64r3/splitk_cluster4": { + "arch": "gfx1250", + "bytes": 46739, + "family": "conv", + "kernel_name": "irhash_conv_1250_splitk_cluster4_N1H7W7C256_K64Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_splitk_skc4", + "sha256": "a8b939d52480eb423bfba2f749d82ae619c776765a543acee44dcf6b7a6ee0bc", + "status": "ok" + }, + "conv/gfx1250/n1h7c256k64r3/splitk_cshuffle": { + "arch": "gfx1250", + "bytes": 42877, + "family": "conv", + "kernel_name": "irhash_conv_1250_splitk_cshuffle_N1H7W7C256_K64Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_splitk", + "sha256": "a88d9f70b35d5e6057f507b5ae995f4e8c311e465433d57a2aa73121e4319273", + "status": "ok" + }, + "conv/gfx1250/n1h7c256k64r3/splitk_default": { + "arch": "gfx1250", + "bytes": 53783, + "family": "conv", + "kernel_name": "irhash_conv_1250_splitk_default_N1H7W7C256_K64Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_splitk", + "sha256": "a7ddd10ef979b41d095e627a0b0c849d086149712913c75300a5b005281237c0", + "status": "ok" + }, "conv/gfx1250/n2h14c64k64r3g4/g4_cshuffle": { "arch": "gfx1250", "bytes": 38219, @@ -6983,6 +7081,22 @@ "sha256": "f926c3ff4d613a6e38cbddb59be8aad725317d4d13c0e4e4789422cac7a6e950", "status": "ok" }, + "conv_fwd_split_k_reduce/gfx1250/bf16_v4": { + "arch": "gfx1250", + "bytes": 3410, + "family": "conv_fwd_split_k_reduce", + "kernel_name": "conv_fwd_splitk_reduce_v4_b256_bf16", + "sha256": "e6b13e9985e7fcd3aa641ba3fc078d40c6e38f3b92edc8f4c64ac86067dff3dc", + "status": "ok" + }, + "conv_fwd_split_k_reduce/gfx1250/fp16_v2_s4": { + "arch": "gfx1250", + "bytes": 2873, + "family": "conv_fwd_split_k_reduce", + "kernel_name": "conv_fwd_splitk_reduce_v2_b256_s4_fp16", + "sha256": "333a82f4dd3088e68945d7206f6346ee2c15aace62611f2ba1e74354be4cd0e4", + "status": "ok" + }, "conv_wgrad/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 53781, @@ -7975,7 +8089,7 @@ "gfx1151": 10, "gfx12-generic": 1, "gfx1201": 8, - "gfx1250": 57, + "gfx1250": 62, "gfx90a": 5, "gfx942": 60, "gfx950": 195 @@ -7986,7 +8100,7 @@ "attention_d128_swa": 1, "attention_d256": 2, "attention_dense": 51, - "conv": 20, + "conv": 23, "conv_dgrad": 13, "conv_dgrad_patch_wmma": 2, "conv_dgrad_smallc_wmma": 2, @@ -7999,6 +8113,7 @@ "conv_direct_fwd_wmma": 5, "conv_direct_nongrouped": 7, "conv_direct_wgrad": 12, + "conv_fwd_split_k_reduce": 2, "conv_wgrad": 22, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, @@ -8016,7 +8131,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 337 + "ok_cases": 342 } } }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 518d49548ff0..730b4e4d3fa4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -395,6 +395,17 @@ "98": "c0cc60f03600060e18bea53e3309dabb928ebe707f07d2d2bb577c52a97e3459", "99": "88356443b084b2986c967ca6ee6eb1180c239747e5f855c0cfb7576b50501099" }, + "conv_fwd_split_k_reduce": { + "0": "1bd4b9b473fdcbe3ff72e8ebcda3186edf877d158e72ed7de52616a6e0987134", + "1": "ef3557ed1a0b20cde9e68ab2bcc7403dcf9e5fe4340c802c7575f1820c6f12bf", + "2": "dc54a5ac9a5cd473173049b9552c2e2a62dde38b88d12254d9dcee0b956f1811", + "3": "f1c5f9f7345a6b9f9f3a0720d191ebfc66cbf7558aecfe9220d85f42f179feae", + "4": "fd034da734dffccd3abcdac1de2df130885d0cf9269bc229a1cc767faaf2035f", + "5": "03fd1243e3e297425ae8c034a227f2963f68eb2c7004db778862c75ba1a35162", + "6": "f053e415623685ef26af4edfcc7d7f179bc76771338605fcf2a76b817ef3dc3b", + "7": "4644f8f86f4fccac98dbd7e540a43d22e037d5f3306418d98281c11c67a85893", + "8": "1bd4b9b473fdcbe3ff72e8ebcda3186edf877d158e72ed7de52616a6e0987134" + }, "conv_implicit_gemm": { "0": "5a524247339bccd3daa72b440ffd50bb475324be9f280cdcd4aa1977e0c2a35f", "1": "38d45e4e6fd036098c347a6046562932c7ddb563f98fb0c4719b76572814388f", @@ -425,7 +436,20 @@ "33": "abc30b49ff5ddce35c96802f015088f900f20216e0f60b62fc5e4e95a7eb9bf8", "34": "803159fbf47cd5131cd81c4ff09244cb85463a48cfdf8ec77e87c64dcd53ded9", "35": "4136318019804d7c4e6c5f0d18e61af9ca4851fbd4de47c051d1b41c78804bae", + "36": "0e3681a1dd0fb09c09c8c6221d145f35deeb740e0c2ac7375c01bb805e956a32", + "37": "724c13caa410f36a622a365adad4f0735010a8b9c7844aaa65b4489028cb9a01", + "38": "86a1af86fbd26785acf364818a64bbe1618a913fe3450310a76168c59a42416a", + "39": "438bd913d8e7c12ada70f4fca880c13ea0256d2dc0543f1974dae7648a246aba", "4": "88f306f723739ac86d3e4324c161253ea8b3b3615c2ab38c9b4163aed2ff65fa", + "40": "7a18f929538d9b2085e5c933a31138f1e5a66fddd3f58f83441cf93f11f77cbc", + "41": "7e25638d2a17c6a67cae8b43650bb84b736dc647667d31358224e62b40e1b4e8", + "42": "9d136489ae9f0218b70e215b889e5ece74d1910d967a3c66514c78e56900a286", + "43": "9debc6ef449b9084cb2b4440ab0686d864286b467f41f29d6776a0275f9b010c", + "44": "4bfa52d255ea802ad3f37a3be2f35e6a436be40b09ef04149c776b3064db0027", + "45": "fd5b5aa3c2cfecb62393fd45ac658988a7b2dc1ee4da82537b7732e438face3e", + "46": "d6294ed883c8eabc76247455fb155ec6ca1a8a781d88faba54a43489ffc1c14a", + "47": "8264ce35448a8d462dfbc4c793a2a626ddaaa8f7b9a5208f35ad1991a6f4d318", + "48": "cb6a7a7e41b149b17a2a4e661f1f01e1ef7783f6e07e84961dd2dc62db57f57b", "5": "04b8950ec6d95e66a1965cf34bd22bf9d7a8fb818b626b43646e7d038f99ee2e", "6": "094dadc42ce9dc18e46c2e74fc3d7c7694f32fd48a920e2d4ceafafaf2662d71", "7": "225364dd9f02f18796d20570aca51616a33843e95339b5e935fe713ecba2dff3", @@ -1947,10 +1971,12 @@ "10": "b73cc46ba9d20e43672db7d5d57d6c027f0b83b6c10bc570b0f67a99701d3a8d", "11": "ce971fe801deb956878fdff0a7d6297e5bcb3d77a0eb51c30fe091bd0c497369", "12": "e618b9bb0882aa33e02fc7bb27f021b0179fbc80a2b2e9289b67242ad380fa1e", - "13": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", + "13": "1cf8be636ab08b8242029720d0a632e066d6c00ad52918d5309c15df051703fb", "14": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", "15": "18758d929385cd9e3c89e45473f607f1495e1a5f90dcd8a21513e6374473bdda", "16": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", + "17": "587f84e7509b9887236ec5e15602319c8bfeb92d87a780b43442ae1c4b9dbf5e", + "18": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", "2": "6adba17dcd82a6ee790c56b4a02dc4becee840a1b7d108ea9378f0f24428f60f", "3": "89e67f3528cfa660478d41228af6c8a9967c2a199a808bdf7d2a1cc5afe8c74b", "4": "458002ecae4794df97126822203e0a6c7c5e8561a547bdda3d1f4c4c66d2abc7", @@ -2141,6 +2167,17 @@ "98": "7e9c58f7a784fa3b86caea305286da5cafd86fb4c50b3a6b1967ca6214fe588a", "99": "7f0092afee5f20eb0ee4ed3d2f94f1d2457474e4872011c103a97d688793e9fa" }, + "conv_fwd_split_k_reduce": { + "0": "9f5d2e0e9448f4b577f054434b9f1ef1d23bc2f3ea573b6e034515f862199af2", + "1": "b2d978f92898f5f836bc030b694c64e7dbb63773eb50433749d7aebe2d0fc7c0", + "2": "73ca635658fa31bf1b869e2164a66504b7bd99f342fa854c03c48b6580c73048", + "3": "5b28b43fefc002302eafa73d745d7ce31a161f2fbe1f21325eb2e93361e3f3d0", + "4": "8360e2b8275f72be9f8a07b90df6c2270ba2e6169b311384be2becb9b0221abe", + "5": "37038d41160f27c6572e00f7c15bfb653fb2ed2dc6b37c00f5ae1f95e9f68555", + "6": "e139491b12318e10c6afac8a905ee0bec094e8a986ffcb9edd27acad1b430297", + "7": "a2c4f4aae0977638c35f94e7306132d0eaeaf918931d7b49d2f9f1a5858e9d82", + "8": "c3fa61fd73e4adaefadbac4836815788cbb15eb3d298965d9ed25a7ff07f1d81" + }, "conv_implicit_gemm": { "0": "772647833bbab21b0c102b363fad3e4f6a7682399b46ede89ec55817d01aad1c", "1": "88298d7944fc0a1a85c80c315540165f92d8f3123f29da2399a43643376ea9ff", @@ -2171,7 +2208,15 @@ "33": "6c2887ec97cbb268f565d9131c2fc0c7f9eeb9bfa349a6710c2d9ecff03afb79", "34": "41645f3357afe819e84433f2d49ce37278beedc1b1d72be0f75805031962a48b", "35": "eb7e6c286ed6ef4186a498ca359bf854689c8abed329dbac4eb4dbf16609c072", + "36": "e786475831dd79c30d72c2dbfd001b2e0ffdf2d9725c5d8ddd8b2ae8bb80e8bb", + "37": "01ba91cde2ebbc67f5ea7997384072b2d5e4d5c2ee7b3cc83515c3b63c90cb1b", + "38": "cd3cef7ea725a44cbc18b33c206114884b6c11b18a1d50fc1759a99b285a2939", + "39": "1e787d0c77c06ff45147e746f1960a1811d7226686195e150683560d87a6e8e0", "4": "77efbd65b4c5b32052e7b6cb04f2929c18a94f2212cb6320aa7889465dad953e", + "40": "829e4ee47d4fdd7e1949a51f736203e5b314ab4e9acdf3756cc760174f6259a3", + "41": "c6261f70e0d64f3275ea1f5bdcb3ebc5111941c18c536333c818c54b04108f13", + "42": "3101a83a29d0d08bf3ae9c0789046da8a643b38b0ba8ea787ded19a180c8ac6a", + "43": "218e99f47dc23541073fbca8d9c3c28058b2009dde18043a8787e5fddae7b5fc", "5": "611cf41a2bd6434d1de40966b24ef77c42c8f826b07f961b765a679b965e840d", "6": "9da3b4883cdda7b7cfbf8391cb3f6705b5ca52b0e4ba173ad50a1044804268d0", "7": "1151fa3728bc2700afbfedbff3334cc1c2a0eff2e2df1727cef1f8b6aa3d1cef", @@ -3595,10 +3640,12 @@ "10": "b73cc46ba9d20e43672db7d5d57d6c027f0b83b6c10bc570b0f67a99701d3a8d", "11": "ce971fe801deb956878fdff0a7d6297e5bcb3d77a0eb51c30fe091bd0c497369", "12": "e618b9bb0882aa33e02fc7bb27f021b0179fbc80a2b2e9289b67242ad380fa1e", - "13": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", + "13": "1cf8be636ab08b8242029720d0a632e066d6c00ad52918d5309c15df051703fb", "14": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", "15": "18758d929385cd9e3c89e45473f607f1495e1a5f90dcd8a21513e6374473bdda", "16": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", + "17": "587f84e7509b9887236ec5e15602319c8bfeb92d87a780b43442ae1c4b9dbf5e", + "18": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", "2": "6adba17dcd82a6ee790c56b4a02dc4becee840a1b7d108ea9378f0f24428f60f", "3": "89e67f3528cfa660478d41228af6c8a9967c2a199a808bdf7d2a1cc5afe8c74b", "4": "458002ecae4794df97126822203e0a6c7c5e8561a547bdda3d1f4c4c66d2abc7", @@ -3789,6 +3836,17 @@ "98": "570cd04e144051c6c77f9e8f7dc7f26397ecf01727a5487464ce3a6ef5def3c2", "99": "e6dbba784cf0c799a26bdedddb5b5579fbdbc70b8f4f7878e39920875d08cacc" }, + "conv_fwd_split_k_reduce": { + "0": "1224e46c4cc351f7d63cdf21cbbfb446eaec8e64fe5460148da7bb543ac2e689", + "1": "1b23c405e95a0a53931aaf48ac00042532e951a4ed9871b2301bdba9f7268e47", + "2": "e84858f537058d2ae976b1733d0a5e79c5b5d587780578aff1afbd5bb2ee6e3a", + "3": "3a6e11e9b6828208331c647fe74f4f6ec69a1c831e36b972e06198cba5c5ba8b", + "4": "551836622936e8574c76f7e380330d258ab593dbab18bc2e8926d8ccce549d34", + "5": "9b0c789843b2126d6102e70346facc16769341e4980f6004789fa319c7111510", + "6": "7a7952045cdc507b2a6931a7103a2541c1d0dae818ffa9ba979ea1be8eb4ecd9", + "7": "714a8b14b10ffbfea6b77bb32e209f6cde637f42c1e0e00467a50c41a00df8d2", + "8": "8aa7a9ef9fa03b2000807a864f86c5408e3d0dc9365a4d72cdf38c850757f9d7" + }, "conv_implicit_gemm": { "0": "537d9c117027edd71266496ef618c913d87f5b3c786a1aa515aa6c3ce588412a", "1": "875a128de1cbc234fef52702022e719d396ad475274cf5e9f9c96bf7d88ecb08", @@ -3819,7 +3877,15 @@ "33": "0a6df89035ced3975135105bb7866342f5a1ae4a80a469dd722bebca053e900e", "34": "05ef782f611c79e0ce45d760932b94acec45bb5fbc1fbb26a74b79d3fbd44c1e", "35": "fe59372c68f0e230e1bed50d43b46e21b2e135b8db0b6034d2b8c8a8a60f0bd8", + "36": "56cb257b09ce0512ab08e1dd09a791d59df4428d6dcbb2af5910a3b38c7354d5", + "37": "3125f74ebeaea3e183e427a926cec5aeca7ce5bae64395503556709a90558961", + "38": "0f890d6b70ed756bf80d56c9afccc6e7cbe9c599e49e6151c906840baece3432", + "39": "f9b59b05ccfb16c99f12e999b2709ca5dedacb1d3034ea3f4d37b49c56917533", "4": "fc875236d2e41709772087dc529b907989270e96cc6ed54c504195b4b5525605", + "40": "51e25c0f016d7a424867ba5ec0578f12fc5806dfbce0bf9222ac5b8a42fe902f", + "41": "78ac63a2766ca167fd7d6ca409f8fa9d436988d986d3ed6a7e6b58e374d4aafb", + "42": "a15f65af3310f808a1f526136b6634c7ba211d3b6d1a9eeaca963a8637149c80", + "43": "ab5b80a1f39ad5eddeeac36ebef2f509ae8077a447eba4626fbe7287e8f0d276", "5": "d1d4a8f8e6a59d6a06e3d8f7e488dcec8c2ad324ee23ce37ccc1909d68490a72", "6": "2b7ca9235abc29cda005fb6a6ab7be1709b3133513d5c8571a79da42b0ee2d1e", "7": "334f88c8f024cdc16d2a1b418cb38211c0af31f55640357ece9e79878b4a7917", @@ -5243,10 +5309,12 @@ "10": "b73cc46ba9d20e43672db7d5d57d6c027f0b83b6c10bc570b0f67a99701d3a8d", "11": "ce971fe801deb956878fdff0a7d6297e5bcb3d77a0eb51c30fe091bd0c497369", "12": "e618b9bb0882aa33e02fc7bb27f021b0179fbc80a2b2e9289b67242ad380fa1e", - "13": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", + "13": "1cf8be636ab08b8242029720d0a632e066d6c00ad52918d5309c15df051703fb", "14": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", "15": "18758d929385cd9e3c89e45473f607f1495e1a5f90dcd8a21513e6374473bdda", "16": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", + "17": "587f84e7509b9887236ec5e15602319c8bfeb92d87a780b43442ae1c4b9dbf5e", + "18": "08970ac3e4ce2f247d9596b812bf8433646bab5435a8d0a97b1932f90a3cbdee", "2": "6adba17dcd82a6ee790c56b4a02dc4becee840a1b7d108ea9378f0f24428f60f", "3": "89e67f3528cfa660478d41228af6c8a9967c2a199a808bdf7d2a1cc5afe8c74b", "4": "458002ecae4794df97126822203e0a6c7c5e8561a547bdda3d1f4c4c66d2abc7", @@ -5462,6 +5530,17 @@ "98": "1da4f9f7ac2232b89d5559d569f84d48fdcb7cb925d4b9f167b72450f6470e8c", "99": "cc107ea999a363cd359e4feb61ecde612189f9366a6cd33684ab37f2053b474b" }, + "conv_fwd_split_k_reduce": { + "0": "e6b13e9985e7fcd3aa641ba3fc078d40c6e38f3b92edc8f4c64ac86067dff3dc", + "1": "42f5a194c72d93d02253339943a2bb28aaf72bb7a1d3b0f11e88067c5483ba47", + "2": "98afed06bcddf92ab46fc78daea2c08785ec4a6380a26799096d57e7e81056f2", + "3": "c89645b963e65ca19adc2e100a73ed5ad8ca8c4d7ccf749cd77297dd163217e7", + "4": "cbaf6fbac633e525c78fa4885a25c87dfadd4c7c0f431682cbb0a6cceb7248de", + "5": "ec5825453140a5116961893994975885506e638c55b75bcd466995f8f179c1a4", + "6": "ffe6fad42a762aace764003fcfb25682c54577b101ef3475e06d76617dc058f3", + "7": "982dbe7395d2a5380831aaefab1811f6c3dc449c8fcf909db40a76ea38f58d84", + "8": "d1c390d9ed69870adf73b5cc3a8fd2888e7c887ac6d00d80972d052e3c542792" + }, "conv_implicit_gemm": { "0": "aedb929fc77d47ad3818072587768da2e837daff1783b62b077b4ab0a411a01c", "1": "9aa61312d75bd2b6238289c0a068e1315e8ed5ff5599e3d8522775fa80476b30", @@ -5492,7 +5571,20 @@ "33": "f2cd9b6ed4caec88601dc4704f4b2233b3e6a62a48fdb0f91a6449aa65f40b9b", "34": "552b2e8622c107c4232b97c0f7b7539d833b0f0af64828587b9d4a87b1a5f202", "35": "7a31056543c03e1e4fce5cd8b3c34ac0afbd26568f062af8e8b5d713d9e1ba5d", + "36": "0df12057e9c9695958ced869503461d86a39b6d72bc49def46f27569992eaa2b", + "37": "dd71534a9c710c9fa658611ce26094056b90dad98c006f8e18943480b59195b2", + "38": "7995b6c3f39cee64ad027629e4e36e04591ec6d0988ed4bf908a21f15592f6df", + "39": "9a4db02cbd078648423576044d049c8a77344c1c244349c77ac839779b06c0ba", "4": "294495c53dd2816f58fed394e5e64556abbf4f986560fa6ed067b3c1be3f89b4", + "40": "4dff69b03d34b61dd7a82ac1530e9a81fbe5f2ee46251c80106aa4b37bfd85e3", + "41": "c0eb0fc92a3bac4613e2b92189b23890fb0657387902fdb1a8d1efc5ed7e13e1", + "42": "3fdb6a79902d3e5c4c592e5fe6f4a0dc9faede07190ac6f511335086a18be458", + "43": "8da08138603c31a3f912b82bf8d09e2b392c96c2814fef470477eb631eb0ad3d", + "44": "e7092a963b489261fa1778e28fc84df3e5f1fd74b165cbc093b391209773e146", + "45": "b1ebd837f12df9cbc5b6903ca50c7d6f672886e51bcc712bfae5a1486f7a0450", + "46": "d2eca84b320a9007e56cde5d803029e5808e0f57d069115e1d455dc24f9d08b4", + "47": "22c93037c79bb4f9f014ca342ffc06e258b2ebbdd07d4baaf988afdfc30e6666", + "48": "5015c05bbc955939935ccd7e3a8465ae0d5c8c839eb0295e9822d16fc25022fb", "5": "8c1357be310fd2071943bf3a8bdddc41633797041d23298ee7c984808a7407fe", "6": "254c5850fb4caa8bbf2cadbcb9c4d66909ebac2ad147efd3e0dee6931cdfce80", "7": "062c1d80190ce13030dc11c49b20ba5c6666775fff1dded93803dffe33cf3c2b", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index fa9531af4855..a9d2effeb2ca 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -140,6 +140,8 @@ def build_conv( groups=1, vector_size_c=None, group_merge=1, + split_k=1, + split_k_cluster=0, ): def _build(): from kernels.common.conv_implicit_gemm import ( @@ -166,6 +168,8 @@ def _build(): groups=groups, vector_size_c=vector_size_c, group_merge=group_merge, + split_k=split_k, + split_k_cluster=split_k_cluster, ) return build_implicit_gemm_conv(spec, arch=arch) @@ -311,6 +315,20 @@ def _build(): return _build +def build_conv_fwd_split_k_reduce(arch, **spec_kw): + """Build the split-K forward Stage 2 (folds the fp32 slabs into D).""" + + def _build(): + from kernels.common.conv_fwd_split_k_reduce import ( + FwdSplitKReduceSpec, + build_conv_fwd_split_k_reduce, + ) + + return build_conv_fwd_split_k_reduce(FwdSplitKReduceSpec(**spec_kw), arch=arch) + + return _build + + def build_moe_sort(phase, tokens, topk, experts, block, arch): def _build(): from rocke.instances.common.moe_sorting import ( @@ -2260,6 +2278,49 @@ def add(family, case_id, arch, build): ), ) + # gfx1250 split-K forward (WMMA): the two-stage binary (fp32 workspace + # slabs, degree as a kernarg) on mem/cshuffle and mem/default, and the + # single-launch cluster form (S=4) that folds the slabs through a cluster + # barrier. Small M, deep K: 9*256 reduction. + conv_small_m = (1, 7, 7, 256, 64, 3, 3, 1, 1, 1, 1, 1, 1) + for tag, epi, vec_c, skc in ( + ("splitk_cshuffle", "cshuffle", None, 0), + ("splitk_default", "default", 1, 0), + ("splitk_cluster4", "cshuffle", None, 4), + ): + add( + "conv", + f"conv/gfx1250/n1h7c256k64r3/{tag}", + "gfx1250", + build_conv( + "irhash_conv_1250_" + tag, + "gfx1250", + conv_small_m, + wave_size=32, + wtm=16, + wtn=16, + wtk=32, + tile_m=32, + tile_n=32, + tile_k=32, + epilogue=epi, + vector_size_c=vec_c, + split_k=2, + split_k_cluster=skc, + ), + ) + # Its Stage 2: the runtime slab loop and the degree-unrolled fold. + for tag, kw in ( + ("bf16_v4", dict(dtype_d="bf16", vec=4)), + ("fp16_v2_s4", dict(dtype_d="fp16", vec=2, degree=4)), + ): + add( + "conv_fwd_split_k_reduce", + f"conv_fwd_split_k_reduce/gfx1250/{tag}", + "gfx1250", + build_conv_fwd_split_k_reduce("gfx1250", **kw), + ) + # Conv wgrad: implicit-GEMM backward-weight direction (NHWK x NHWC -> KYXC). # wgrad1: small 3x3 conv (N8 H8 W8 C16 K32 Y3 X3, stride/pad=defaults). # wgrad2: 1x1 conv (N2 H16 W16 C32 K32, no filter spatial). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_cluster_launch.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_cluster_launch.py index f6097fefc884..b173173fe3a2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_cluster_launch.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_cluster_launch.py @@ -219,6 +219,80 @@ def test_support_query_is_cached_per_device(): assert hip_module.device_supports_cluster_launch(2) is False +def test_prepared_launch_without_cluster_is_unchanged(hip): + enqueue = Runtime().prepare_launch( + object(), (4, 1, 1), (64, 1, 1), b"\0" * 8, stream=_STREAM + ) + enqueue() + enqueue() + assert not hip.ex + assert [c["grid"] for c in hip.plain] == [(4, 1, 1), (4, 1, 1)] + + +def test_prepared_launch_with_cluster_goes_through_launch_kernel_ex(hip): + enqueue = Runtime().prepare_launch( + object(), + (8, 2, 1), + (64, 1, 1), + b"\0" * 8, + shared_bytes=256, + stream=_STREAM, + cluster=(4, 2, 1), + ) + assert not hip.ex # nothing is enqueued until the callable runs + enqueue() + enqueue() + assert not hip.plain + assert [c["cluster"] for c in hip.ex] == [(4, 2, 1), (4, 2, 1)] + call = hip.ex[0] + assert call["grid"] == (8, 2, 1) and call["shared"] == 256 + assert call["stream"] == _STREAM and call["extra"] is not None + + +def test_launch_and_prepared_launch_share_the_cluster_launch(hip): + # Both paths issue the cluster launch through one helper, so they cannot + # drift apart (error handling, call-site name). + with mock.patch.object( + Runtime, "_launch_cluster", wraps=Runtime._launch_cluster + ) as shared: + Runtime().launch( + object(), + (4, 1, 1), + (64, 1, 1), + b"\0" * 8, + stream=_STREAM, + cluster=(4, 1, 1), + ) + Runtime().prepare_launch( + object(), + (4, 1, 1), + (64, 1, 1), + b"\0" * 8, + stream=_STREAM, + cluster=(4, 1, 1), + )() + assert [c.args[-1] for c in shared.call_args_list] == [ + "hipModuleLaunchKernel", + "prepare_launch", + ] + assert [c["cluster"] for c in hip.ex] == [(4, 1, 1), (4, 1, 1)] + + +def test_prepared_launch_validates_the_cluster(hip): + # A grid that is not a whole number of clusters is refused when the launch + # is prepared, not later inside a timed loop. + with pytest.raises(ValueError): + Runtime().prepare_launch( + object(), + (6, 1, 1), + (64, 1, 1), + b"\0" * 8, + stream=_STREAM, + cluster=(4, 1, 1), + ) + assert not hip.ex and not hip.plain + + class _FakeRuntime: def __init__(self): self.calls = [] @@ -235,6 +309,10 @@ def launch(self, fn, grid, block, args, **kw): def retain_for_stream(self, *a): pass + def prepare_launch(self, fn, grid, block, args, **kw): + self.calls.append(("prepared", kw.get("cluster"))) + return lambda: None + def _launcher(rt, cluster_dims): with mock.patch.object(launcher, "_runtime", return_value=rt): @@ -274,6 +352,19 @@ def test_launcher_without_cluster_dims(): assert rt.calls == [("blocking", None), ("blocking", (4, 1, 1))] +def test_bind_keeps_the_compiled_cluster(): + # bind() pre-packs a launch for timed loops; it must launch the same + # clusters as a plain call, or a cluster kernel silently runs as 1x1x1 + # clusters (its cluster barrier then synchronises nothing). + rt = _FakeRuntime() + kl = _launcher(rt, (1, 1, 4)) + with mock.patch.object(launcher, "_runtime", return_value=rt): + kl.bind({}, config=launcher.LaunchConfig(stream=_STREAM, grid=(2, 2, 4))) + plain = _launcher(rt, None) + plain.bind({}, config=launcher.LaunchConfig(stream=_STREAM, grid=(2, 2, 4))) + assert rt.calls == [("prepared", (1, 1, 4)), ("prepared", None)] + + def test_launcher_validates_cluster_dims_at_construction(): with pytest.raises(ValueError, match="limit of 16"): _launcher(_FakeRuntime(), (4, 4, 4)) @@ -302,3 +393,39 @@ def test_manifest_timing_passes_the_cluster_without_torch(): assert rt.launch.call_count == 5 for call in rt.launch.call_args_list: assert call.kwargs == {"cluster": (2, 1, 1)} + + +def test_prepared_cluster_launch_sets_up_once(hip): + # Validation, the device query and the ctypes launch config are built when + # the launch is prepared; each call only enqueues the same config. + with mock.patch.object( + hip_module, "check_launch_cluster", wraps=hip_module.check_launch_cluster + ) as check, mock.patch.object( + hip_module, "device_supports_cluster_launch", return_value=True + ) as supports: + enqueue = Runtime().prepare_launch( + object(), + (8, 1, 1), + (64, 1, 1), + b"\0" * 8, + stream=_STREAM, + cluster=(4, 1, 1), + ) + for _ in range(3): + enqueue() + assert check.call_count == 1 + assert supports.call_count == 1 + assert len(hip.ex) == 3 + assert all(c["cfg"] is hip.ex[0]["cfg"] for c in hip.ex) + assert [c["cluster"] for c in hip.ex] == [(4, 1, 1)] * 3 + + +def test_prepared_cluster_launch_refuses_an_unsupported_device_up_front(hip): + with mock.patch.object( + hip_module, "device_supports_cluster_launch", return_value=False + ): + with pytest.raises(HipError, match="does not support cluster launch"): + Runtime().prepare_launch( + object(), (2, 1, 1), (64, 1, 1), b"", stream=_STREAM, cluster=(2, 1, 1) + ) + assert not hip.ex diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py index 6c5070923131..9a1838baddb8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py @@ -6975,13 +6975,17 @@ def packer(values): return b"\x01\x02" class FakeRuntime: - def prepare_launch(self, fn, grid, block, args, *, shared_bytes, stream): + def prepare_launch( + self, fn, grid, block, args, *, shared_bytes, stream, cluster=None + ): self.prepared = (fn, grid, block, args, shared_bytes, stream) + self.cluster = cluster return lambda: enqueued.append(1) launcher = L.KernelLauncher.__new__(L.KernelLauncher) launcher._fn = "fn" launcher._packer = packer + launcher._cluster_dims = None rt = FakeRuntime() cfg = L.LaunchConfig(grid=(2, 1, 1), block=(64, 1, 1), stream=7) with mock.patch.object(L, "_runtime", return_value=rt): @@ -6992,6 +6996,13 @@ def prepare_launch(self, fn, grid, block, args, *, shared_bytes, stream): self.assertEqual(packed, [{"a": 1}]) self.assertEqual(len(enqueued), 3) self.assertEqual(rt.prepared, ("fn", (2, 1, 1), (64, 1, 1), b"\x01\x02", 0, 7)) + self.assertIsNone(rt.cluster) + + # A clustered kernel keeps its cluster shape on the prepared launch. + launcher._cluster_dims = (1, 1, 2) + with mock.patch.object(L, "_runtime", return_value=rt): + launcher.bind({"a": 1}, config=cfg) + self.assertEqual(rt.cluster, (1, 1, 2)) class TestLauncherFenceContract(unittest.TestCase):