feat(rocke): split-C and grouped LDS-halo direct conv fwd kernels

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py

@@ -2161,6 +2161,7 @@ def _cache_dispatch(args, arch: str, cases: list) -> int:
             "direct_fwd",
             "direct_dgrad",
             "direct_dgrad_helper",
+            "direct_fwd_helper",
             "direct_wgrad",
         ),
     )

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py

@@ -34,10 +34,13 @@
 
 from __future__ import annotations
 
+import collections
+import hashlib
 import json
+import struct
 import time
 from concurrent.futures import ProcessPoolExecutor, as_completed
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from pathlib import Path
 from typing import Dict, List, Optional, Sequence, Tuple
 
@@ -77,6 +80,7 @@
     _caps_jobs,
     _caps_of,
     _helper_identity,
+    _split_reduce_identity,
     count_direct_jobs,
     dgrad_rtc_k_atoms,
     dgrad_rtc_spills,
@@ -91,7 +95,9 @@
     DIRECT_DGRAD_REORGANIZE_RTC as _REORGANIZE_RTC,
     DIRECT_DGRAD_TRANSPOSE as _TRANSPOSE,
     DIRECT_DGRAD_TRANSPOSE_RTC as _TRANSPOSE_RTC,
+    DIRECT_HALO_VARIANTS as _HALO_VARIANTS,
     DIRECT_MFMA_DGRAD as _MFMA_DGRAD,
+    DIRECT_SPLIT_REDUCE as _SPLIT_REDUCE,
     direct_dgrad_weight_transforms,
     direct_dgrad_workspaces,
     direct_spec_for_problem,
@@ -222,16 +228,17 @@ def compile_all_direct(
         f"AOT compile-all: enumerating direct variants for {arch}/{dtype} "
         f"({', '.join(directions)})"
     )
-    return compile_jobs(
+    all_jobs = enumerate_direct_jobs(
+        arch=arch,
+        target=target,
+        directions=directions,
+        dtype=dtype,
+        jobs=jobs,
+        log=log,
+    )
+    rc = compile_jobs(
         cache=cache,
-        all_jobs=enumerate_direct_jobs(
-            arch=arch,
-            target=target,
-            directions=directions,
-            dtype=dtype,
-            jobs=jobs,
-            log=log,
-        ),
+        all_jobs=all_jobs,
         build=direct_kernel_for_job,
         arch=arch,
         dtype=dtype,
@@ -240,6 +247,118 @@ def compile_all_direct(
         limit=limit,
         log=log,
     )
+    if "fwd" in directions:
+        prune_halo_entries(cache, all_jobs, log=log)
+    return rc
+
+
+def _elf_code_facts(hsaco: bytes) -> Optional[Tuple[bytes, int]]:
+    """``(code digest, scratch bytes)`` of an ELF64 code object, or None for
+    anything else.
+
+    The code is the ``.text`` section followed by every kernel descriptor
+    (``<kernel>.kd``): what the GPU executes and how it is launched, so two
+    binaries with equal digests run identically whatever their symbol names.
+    Only its SHA-256 is kept, so a sweep's worth of facts stays small.
+    ``scratch`` is the largest descriptor's private segment size -- non-zero
+    when the kernel spills.
+    """
+    if hsaco[:5] != b"\x7fELF\x02" or len(hsaco) < 64:
+        return None
+    (shoff,) = struct.unpack_from("<Q", hsaco, 0x28)
+    shentsize, shnum, shstrndx = struct.unpack_from("<HHH", hsaco, 0x3A)
+    secs = [
+        struct.unpack_from("<IIQQQQII", hsaco, shoff + i * shentsize)
+        for i in range(shnum)
+    ]
+    # (name, type, flags, addr, offset, size, link, info)
+    strtab_off = secs[shstrndx][4]
+
+    def cstr(off: int) -> str:
+        return hsaco[off : hsaco.index(b"\0", off)].decode()
+
+    text = b""
+    kds = []
+    for name_off, sh_type, _, _, off, size, link, _ in secs:
+        if cstr(strtab_off + name_off) == ".text":
+            text = hsaco[off : off + size]
+        if sh_type != 2:  # SHT_SYMTAB
+            continue
+        names = secs[link][4]
+        for sym in range(off, off + size, 24):
+            st_name, _, _, shndx, value, st_size = struct.unpack_from(
+                "<IBBHQQ", hsaco, sym
+            )
+            if not cstr(names + st_name).endswith(".kd") or shndx >= shnum:
+                continue
+            _, _, _, addr, sec_off, _, _, _ = secs[shndx]
+            at = sec_off + value - addr
+            kd = hsaco[at : at + st_size]
+            # Bytes 16..24 locate the entry point relative to the descriptor:
+            # file layout, not behaviour, so they are left out.
+            kds.append(kd[:16] + kd[24:])
+    if not kds:
+        return None
+    scratch = max(struct.unpack_from("<I", kd, 4)[0] for kd in kds)
+    return hashlib.sha256(text + b"".join(sorted(kds))).digest(), scratch
+
+
+def prune_halo_entries(cache: KernelCache, jobs: Sequence[BuildJob], log=print) -> int:
+    """Drop the cached grouped-halo and split-C forward entries no sweep should
+    be offered (the dense unsplit grid is left as it is).
+
+    Two kinds, both only visible in the compiled code: binaries that spill to
+    scratch (never the best of their shape), and ``waves_per_eu=3`` binaries
+    whose code equals their ``waves_per_eu=None`` twin's -- the hint changed
+    nothing, so the twin already is that kernel. The binaries stay in the blob
+    store: a later ``--compile-all`` re-links those entries without compiling
+    and this drops them again. Returns the number of entries dropped.
+    """
+    wanted = {
+        j.identity
+        for j in jobs
+        if j.identity.direction == "direct_fwd"
+        and (
+            j.identity.algorithm == "direct_grouped_halo"
+            or (
+                j.identity.algorithm == "direct_nongrouped"
+                and json.loads(j.identity.knobs).get("split_c", 1) > 1
+            )
+        )
+    }
+    present = {
+        ident: path for ident, path in cache.list_all("direct_fwd") if ident in wanted
+    }
+    facts = {
+        ident: _elf_code_facts(path.read_bytes()) for ident, path in present.items()
+    }
+    drop: Dict[KernelIdentity, str] = {}
+    twins: Dict[KernelIdentity, Dict[object, KernelIdentity]] = {}
+    for ident, fact in facts.items():
+        if fact is not None and fact[1] > 0:
+            drop[ident] = "spills"
+        knobs = json.loads(ident.knobs)
+        we = knobs.pop("waves_per_eu", None)
+        key = replace(ident, knobs=json.dumps(knobs, sort_keys=True))
+        twins.setdefault(key, {})[we] = ident
+    for pair in twins.values():
+        hinted, plain = pair.get(3), pair.get(None)
+        if hinted is None or plain is None or hinted in drop:
+            continue
+        a, b = facts[hinted], facts[plain]
+        if a is not None and b is not None and a[0] == b[0]:
+            drop[hinted] = "same code as waves_per_eu=None"
+    for ident in drop:
+        cache.discard(ident)
+    if drop:
+        reasons = collections.Counter(drop.values())
+        log(
+            f"AOT compile-all: dropped {len(drop)} of {len(present)} grouped-halo "
+            "and split-C forward entries ("
+            + ", ".join(f"{n} {why}" for why, n in sorted(reasons.items()))
+            + ")"
+        )
+    return len(drop)
 
 
 # ---------------------------------------------------------------------------
@@ -265,6 +384,10 @@ class DirectStep:
     buffers: Tuple[str, Optional[str], str]
     conv_args: object = None
     scalars: Tuple[Tuple[str, int], ...] = ()
+    # Per A/B/D slot, ``(byte offset, byte size)`` of the window of its
+    # buffer the step binds, or None for the whole buffer: a batch-split
+    # launch sees only its images (see _batch_chunks).
+    views: Tuple[Optional[Tuple[int, int]], ...] = ()
 
 
 @dataclass(frozen=True)
@@ -306,6 +429,28 @@ def _conv_step(identity, hsaco_path, meta, spec, problem, direction, buffers):
     )
 
 
+def _batch_chunks(problem) -> Optional[List[int]]:
+    """Images per launch for a forward problem the LDS-halo kernels can only
+    run in pieces, or None when one launch serves it.
+
+    Those kernels address every operand through one 32-bit buffer resource and
+    redirect masked accesses past ``_OOB_BASE``, so no operand may reach it
+    (a 768x768x256 bf16 batch of 8 is 2.4 GB). The images are independent, so
+    the batch is split into launches whose activation and output windows each
+    fit; the chunks are balanced rather than greedy.
+    """
+    from kernels.common.conv_direct_nongrouped import _OOB_BASE
+
+    in_img = 2 * problem.H * problem.W * problem.total_c
+    out_img = 2 * problem.Ho * problem.Wo * problem.total_k
+    per_launch = _OOB_BASE // max(in_img, out_img)
+    if problem.N <= per_launch or per_launch < 1:
+        return None
+    n_steps = -(-problem.N // per_launch)
+    base, extra = divmod(problem.N, n_steps)
+    return [base + (i < extra) for i in range(n_steps)]
+
+
 def plan_for(
     cache: KernelCache,
     identity: KernelIdentity,
@@ -328,10 +473,46 @@ def plan_for(
         return None, "cache entry has no kernel_name (rebuild the cache)"
     variant = identity.algorithm
     knobs = json.loads(identity.knobs) if identity.knobs else {}
+    if variant in _HALO_VARIANTS and identity.direction == "direct_fwd":
+        chunks = _batch_chunks(problem)
+        if chunks is not None:
+            if knobs.get("split_c", 1) > 1:
+                return None, "a batch-split launch does not take split-C kernels"
+            return _batch_split_plan(
+                cache, identity, hsaco_path, meta, problem, arch, knobs, chunks
+            )
     spec, why = direct_spec_for_problem(variant, problem, knobs, arch)
     if spec is None:
         return None, why
 
+    if identity.direction == "direct_fwd" and getattr(spec, "split_c", 1) > 1:
+        # Split-C: the conv pass writes fp32 partial sums into a workspace,
+        # and the reduction folds them into y. Both take the forward ABI.
+        from kernels.common.conv_args import ConvArgs
+
+        reduce_id = _split_reduce_identity(
+            identity.arch, identity.wave_size, spec.split_c, identity.dtype_a
+        )
+        entry = cache.get(reduce_id)
+        if entry is None:
+            return None, f"missing cached {_SPLIT_REDUCE} helper (rebuild the cache)"
+        reduce_spec = make_spec(_SPLIT_REDUCE, problem, {"split_c": spec.split_c})
+        conv = _conv_step(
+            identity, hsaco_path, meta, spec, problem, "fwd", ("x", "w", "ws_split")
+        )
+        fold = DirectStep(
+            kernel_name=entry[1]["kernel_name"],
+            hsaco_path=cache.hsaco_path(reduce_id),
+            signature=conv.signature,
+            grid=reduce_spec.grid(problem),
+            block=(reduce_spec.threads, 1, 1),
+            buffers=("ws_split", "w", "y"),
+            conv_args=ConvArgs.from_problem(problem, direction="fwd"),
+        )
+        return (
+            DirectPlan(identity, (conv, fold), {"ws_split": spec.workspace_bytes}),
+            "ok",
+        )
     if identity.direction == "direct_fwd":
         step = _conv_step(
             identity, hsaco_path, meta, spec, problem, "fwd", ("x", "w", "y")
@@ -402,6 +583,34 @@ def plan_for(
     return DirectPlan(identity, tuple(steps), workspaces), "ok"
 
 
+def _batch_split_plan(
+    cache, identity, hsaco_path, meta, problem, arch, knobs, chunks
+) -> Tuple[Optional[DirectPlan], str]:
+    """One forward launch per batch chunk (see :func:`_batch_chunks`)."""
+    from kernels.common.conv_direct_nongrouped import tile_w_candidates
+
+    widths = tile_w_candidates(
+        problem.Wo, make_spec(identity.algorithm, problem, knobs).atom_tile
+    )
+    if knobs["tile_w"] not in widths:
+        return None, f"tile_w={knobs['tile_w']} does not fit Wo={problem.Wo}"
+    in_img = 2 * problem.H * problem.W * problem.total_c
+    out_img = 2 * problem.Ho * problem.Wo * problem.total_k
+    steps = []
+    n0 = 0
+    for n in chunks:
+        sub = replace(problem, N=n)
+        spec = make_spec(identity.algorithm, sub, knobs)
+        ok, why = validate_spec(identity.algorithm, spec, arch)
+        if not ok:
+            return None, why
+        step = _conv_step(identity, hsaco_path, meta, spec, sub, "fwd", ("x", "w", "y"))
+        views = ((n0 * in_img, n * in_img), None, (n0 * out_img, n * out_img))
+        steps.append(replace(step, views=views))
+        n0 += n
+    return DirectPlan(identity, tuple(steps)), "ok"
+
+
 def direct_plans(
     cache: KernelCache, problem, direction: str, arch: str
 ) -> Tuple[List[DirectPlan], List[Tuple[KernelIdentity, str]]]:
@@ -423,6 +632,10 @@ def direct_plans(
         # grouped problem would bury the grouped kernels' reasons in noise.
         if identity.algorithm == "direct_nongrouped" and problem.groups != 1:
             continue
+        # ... and the grouped LDS-halo kernel only for more than one; on a
+        # dense problem it is the non-grouped kernel with extra divides.
+        if identity.algorithm == "direct_grouped_halo" and problem.groups == 1:
+            continue
         plan, why = plan_for(cache, identity, hsaco_path, meta, problem, arch)
         if plan is None:
             rejected.append((identity, why))
@@ -436,14 +649,23 @@ def launch_values(
 ) -> dict:
     """Kernarg values of ``step`` with its buffers bound to ``ptrs``/``sizes``."""
     a, b, d = step.buffers
+
+    def bind(slot: int, name: str) -> Tuple[int, int]:
+        view = step.views[slot] if step.views else None
+        if view is None:
+            return ptrs[name], sizes[name]
+        return ptrs[name] + view[0], view[1]
+
+    (a_ptr, a_bytes), (d_ptr, d_bytes) = bind(0, a), bind(2, d)
     if step.conv_args is None:
         return {
-            "A": ptrs[a],
-            "D": ptrs[d],
-            "A_bytes": sizes[a],
-            "D_bytes": sizes[d],
+            "A": a_ptr,
+            "D": d_ptr,
+            "A_bytes": a_bytes,
+            "D_bytes": d_bytes,
             **dict(step.scalars),
         }
+    b_ptr, b_bytes = bind(1, b)
     return step.conv_args.to_launch_values(
-        ptrs[a], ptrs[b], ptrs[d], sizes[a], sizes[b], sizes[d]
+        a_ptr, b_ptr, d_ptr, a_bytes, b_bytes, d_bytes
     )

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py

@@ -133,6 +133,9 @@ def comgr_input_key(comgr_input) -> str:
     # looked up by exact identity from the pipeline's main entry and never
     # offered as candidates on their own.
     "direct_dgrad_helper": "conv_direct_dgrad_helpers",
+    # The split-C reduction of the non-grouped forward kernel, found the same
+    # way from a split entry.
+    "direct_fwd_helper": "conv_direct_fwd_helpers",
 }
 
 
@@ -233,6 +236,11 @@ def link(self, identity: KernelIdentity, key: str, meta: dict) -> None:
         )
         self._loaded.clear()
 
+    def discard(self, identity: KernelIdentity) -> None:
+        """Remove ``identity``'s entry. Its binary stays: others may share it."""
+        self._meta_path(identity).unlink(missing_ok=True)
+        self._loaded.clear()
+
     def put(
         self,
         identity: KernelIdentity,