feat(rocke): shape-generic (AOT) Winograd forward transforms for gfx1250

dnn-providers/hip-kernel-provider/rocke/library/dispatch/conv_winograd_plan.py

@@ -13,13 +13,17 @@
   input transform, K3 batched GEMM (:mod:`kernels.common.conv_winograd_gemm_wmma`)
   and K4 output transform; a batch split runs K2 / K3 / K4 once per chunk of
   ``N / chunks`` images on byte offsets of x and y, sharing U;
+* ``unfused2d_rt`` (3x3, :mod:`kernels.common.conv_winograd_fwd_rt`): the same
+  pipeline (without the merged launch) on shape-generic transform binaries,
+  every problem extent a kernarg, so each binary is compiled once ahead of
+  time and named by its kernel name;
 * ``unfused_gen`` (1-D, 7x1, 5x5, segmented long filters,
   :mod:`kernels.common.conv_winograd_gen`): K1, K2, K3 with ``taps`` segments,
   K4;
 * ``fused1d`` (:mod:`kernels.common.conv_winograd_fused_1d_wmma`): the
   generalized K1 writes U, then one fused kernel goes from x to y.
 
-The steps are :class:`ConvLaunchStep` records with no AOT identity: each
+The steps are :class:`ConvLaunchStep` records with no conv-kernel identity: each
 carries the kernel to build (``kernel.build(arch)``) and its kernargs by name,
 resolved by :meth:`ConvLaunchPlan.values` against the buffers ``x`` / ``w`` /
 ``y`` and the plan's workspaces (``U``, ``V``, ``M``, ``e``). Transform knobs,
@@ -232,6 +236,17 @@ def _build_with_arch(fn, spec, arch):
 # ---------------------------------------------------------------------------
 
 
+def _plan_chunks(problem: ConvProblem, plan: WinogradPlan) -> int:
+    """The plan's batch split: a positive integer dividing N (refused, never
+    normalized, otherwise)."""
+    chunks = plan.chunks
+    if not isinstance(chunks, int) or chunks < 1:
+        raise ValueError(f"batch split must be a positive integer (got {chunks!r})")
+    if problem.N % chunks:
+        raise ValueError(f"batch split {chunks} does not divide N={problem.N}")
+    return chunks
+
+
 def _unfused2d(problem: ConvProblem, dtype: str, plan: WinogradPlan):
     from kernels.common import conv_winograd_fwd as W
     from kernels.common.conv_winograd_gemm_wmma import (
@@ -240,9 +255,7 @@ def _unfused2d(problem: ConvProblem, dtype: str, plan: WinogradPlan):
     )
 
     mh, mw = plan.tile
-    chunks = max(1, int(plan.chunks))
-    if problem.N % chunks:
-        raise ValueError(f"batch split {chunks} does not divide N={problem.N}")
+    chunks = _plan_chunks(problem, plan)
     prescale = dtype == "bf16" and max(mh, mw) >= 4
     operand = "bf16" if ((mh, mw) == (2, 2) and dtype == "bf16") else "fp16"
     xf = W.tuned_transform_options((mh, mw), prescale, problem)
@@ -309,6 +322,74 @@ def stage_step(name, chunk):
     return tuple(steps), workspaces
 
 
+def _rt_kernargs(signature, vals: dict, sizes: dict, offsets: dict) -> tuple:
+    """Kernargs of a shape-generic transform step: its scalar values from
+    ``stage_launch``, ``*_bytes`` the chunk's buffer sizes, pointers at the
+    chunk's byte offsets."""
+    out = []
+    for a in signature:
+        nm = a["name"]
+        if nm in vals:
+            out.append(("i32", int(vals[nm])))
+        elif nm.endswith("_bytes"):
+            out.append(("i32", int(sizes[_BUF_OF[nm[: -len("_bytes")]]])))
+        else:
+            buf = _BUF_OF[nm]
+            out.append(("buf", buf, int(offsets.get(buf, 0))))
+    return tuple(out)
+
+
+def _unfused2d_rt(problem: ConvProblem, dtype: str, plan: WinogradPlan):
+    """The 3x3 pipeline on the shape-generic (AOT) transform binaries of
+    :mod:`kernels.common.conv_winograd_fwd_rt`: each step is named by its
+    binary's kernel name (its AOT identity) and its problem extents are
+    kernargs. The binaries are :func:`default_stage_specs`; K3 is the plan's
+    WMMA GEMM, as in :func:`_unfused2d`."""
+    from kernels.common import conv_winograd_fwd_rt as R
+    from kernels.common.conv_winograd_gemm_wmma import (
+        WinogradGemmProblem,
+        WinogradGemmWmmaSpec,
+    )
+
+    mh, mw = plan.tile
+    chunks = _plan_chunks(problem, plan)
+    sub = dataclasses.replace(problem, N=problem.N // chunks)
+    geo = R.RtGeometry(sub, mh, mw)
+    operand, prescale = R.pipeline_dtypes((mh, mw), dtype)
+    why = geo.check(prescale)
+    if why:
+        raise ValueError(why)
+    specs = R.default_stage_specs(sub, (mh, mw), dtype)
+    sizes = dict(geo.buffer_bytes(prescale))
+    x_chunk, y_chunk = sizes["x"], sizes["y"]
+
+    def stage_step(stage, chunk):
+        spec = specs[stage]
+        grid, block, vals = R.stage_launch(spec, geo)
+        sig = R.stage_signature(spec)
+        offs = {"x": chunk * x_chunk, "y": chunk * y_chunk}
+        return _step(
+            spec.kernel_name(),
+            functools.partial(_build_with_arch, R.build_stage, spec),
+            grid,
+            block,
+            sig,
+            _rt_kernargs(sig, vals, sizes, offs),
+        )
+
+    steps = [stage_step("K1", 0)]
+    k3 = WinogradGemmWmmaSpec(dtype=operand, **_k3_knobs(plan.k3))
+    prob = WinogradGemmProblem.from_conv(sub, (mh, mw))
+    for i in range(chunks):
+        for k in ("K2a", "K2"):
+            if k in specs:
+                steps.append(stage_step(k, i))
+        steps.append(_k3_step(k3, prob))
+        steps.append(stage_step("K4", i))
+    workspaces = {k: int(max(sizes[k], 4)) for k in ("U", "V", "M", "e")}
+    return tuple(steps), workspaces
+
+
 def _unfused_gen(problem: ConvProblem, dtype: str, plan: WinogradPlan):
     from kernels.common import conv_winograd_gen as WG
     from kernels.common.conv_winograd_gemm_wmma import WinogradGemmWmmaSpec
@@ -407,7 +488,12 @@ def _fused1d(problem: ConvProblem, dtype: str, plan: WinogradPlan):
     return tuple(steps), {"U": int(max(gsizes["U"], 4))}
 
 
-_KINDS = {"unfused2d": _unfused2d, "unfused_gen": _unfused_gen, "fused1d": _fused1d}
+_KINDS = {
+    "unfused2d": _unfused2d,
+    "unfused2d_rt": _unfused2d_rt,
+    "unfused_gen": _unfused_gen,
+    "fused1d": _fused1d,
+}
 
 
 @functools.lru_cache(maxsize=256)

dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py

@@ -23,6 +23,10 @@
 built per problem (no AOT binary), its eligibility and plan come from
 :mod:`kernels.common.conv_winograd_select`, and its :class:`ConvLaunchPlan`
 runs the transform, GEMM and output kernels in order over named workspaces.
+Its AOT twin, ``conv_winograd_fwd_aot`` (``algorithm="winograd_aot"``, also
+opt-in), runs the same 3x3 plans on the shape-generic transform binaries of
+:mod:`kernels.common.conv_winograd_fwd_rt` (plan kind ``unfused2d_rt``):
+every kernel of its plans is compiled once and named by its kernel name.
 
 Every
 record there spells the :class:`KernelIdentity` of one AOT kernel (format:
@@ -1195,7 +1199,9 @@ def _winograd_model(arch: str) -> WinogradCostModel:
     return WinogradCostModel(num_cus=_ARCH_NUM_CUS.get(arch, _DEFAULT_NUM_CUS))
 
 
-def _winograd_spec(req: ConvGroupedRequest, plan) -> ConvWinogradSpec:
+def _winograd_spec(
+    req: ConvGroupedRequest, plan, name: str = _WINOGRAD_NAME
+) -> ConvWinogradSpec:
     """The spec of ``plan`` on ``req``; raises ``ValueError`` when a kernel of
     the plan cannot run the problem or does not lower at the request's LLVM
     flavor (the K3 GEMM lowers at llvm23 only)."""
@@ -1204,9 +1210,9 @@ def _winograd_spec(req: ConvGroupedRequest, plan) -> ConvWinogradSpec:
     steps, _ = winograd_steps(problem, dtype, plan)
     why = winograd_flavor_reason(problem, dtype, plan, req.llvm_flavor)
     if why is not None:
-        raise ValueError(f"{_WINOGRAD_NAME}:{plan.label}: {why}")
+        raise ValueError(f"{name}:{plan.label}: {why}")
     return ConvWinogradSpec(
-        name=f"{_WINOGRAD_NAME}:{plan.label}",
+        name=f"{name}:{plan.label}",
         plan=plan,
         problem=problem,
         dtype=dtype,
@@ -1266,12 +1272,81 @@ def _winograd_selection(
         return None, str(e)
 
 
+# The same 3x3 plans on the shape-generic transform binaries
+# (:mod:`kernels.common.conv_winograd_fwd_rt`): every kernel of the pipeline is
+# compiled once, ahead of time, and named by its kernel name.
+_WINOGRAD_AOT_NAME = "conv_winograd_fwd_aot"
+
+
+def _as_aot(plan):
+    """The ``unfused2d_rt`` twin of a 3x3 ``unfused2d`` plan, else ``None``
+    (the 1-D and generalized plans have no shape-generic kernels)."""
+    if plan.kind != "unfused2d":
+        return None
+    return dataclasses.replace(plan, kind="unfused2d_rt", merge="")
+
+
+@functools.lru_cache(maxsize=256)
+def _winograd_aot_sweep(req: ConvGroupedRequest) -> tuple[ConvWinogradSpec, ...]:
+    """Every safe 3x3 Winograd plan of ``req`` on the AOT binaries, in model
+    cost order."""
+    if _winograd_problem_reason(req) is not None:
+        return ()
+    dtype = req.dtype.lower()
+    out = []
+    for plan in winograd_plans(
+        _conv_problem(req), dtype, model=_winograd_model(req.arch)
+    ):
+        aot = _as_aot(plan)
+        if aot is None:
+            continue
+        try:
+            out.append(_winograd_spec(req, aot, _WINOGRAD_AOT_NAME))
+        except ValueError:
+            continue
+    return tuple(out)
+
+
+@functools.lru_cache(maxsize=256)
+def _winograd_aot_selection(
+    req: ConvGroupedRequest,
+) -> tuple[ConvWinogradSpec | None, str]:
+    """The selector's plan for ``req`` on the AOT binaries (the cheapest safe
+    3x3 plan when the selector picks a 1-D / generalized one, or when a
+    binary of its plan cannot run the problem), or ``(None, why)``."""
+    why = _winograd_problem_reason(req)
+    if why is not None:
+        return None, why
+    sel = select_winograd(
+        _conv_problem(req), req.dtype.lower(), model=_winograd_model(req.arch)
+    )
+    if sel.kind == "none":
+        return None, sel.reason
+    aot = _as_aot(sel)
+    if aot is not None:
+        try:
+            return _winograd_spec(req, aot, _WINOGRAD_AOT_NAME), "ok"
+        except ValueError as e:
+            why = str(e)
+    else:
+        why = f"the selected {sel.label} plan has no AOT kernels"
+    rest = _winograd_aot_sweep(req)
+    if rest:
+        return rest[0], "ok"
+    return None, why
+
+
 def _winograd_last_step(spec: ConvWinogradSpec):
     steps, _ = winograd_steps(spec.problem, spec.dtype, spec.plan)
     return steps[-1]
 
 
-def _make_winograd_candidate() -> KernelCandidate:
+def _make_winograd_candidate(
+    name: str = _WINOGRAD_NAME,
+    algorithm: str = "winograd",
+    selection=_winograd_selection,
+    sweep=_winograd_sweep,
+) -> KernelCandidate:
     def support(req: OperatorRequest) -> tuple[bool, str]:
         errors = _request_errors(req)
         if errors:
@@ -1280,27 +1355,27 @@ def support(req: OperatorRequest) -> tuple[bool, str]:
         ok, why = selector_matches(req, candidate)
         if not ok:
             return False, why
-        spec, why = _winograd_selection(req)
+        spec, why = selection(req)
         return spec is not None, why
 
     def select(req: OperatorRequest) -> ConvWinogradSpec:
         ok, why = candidate.admits(req)
         if not ok:
-            raise ValueError(f"{_WINOGRAD_NAME} does not support request: {why}")
+            raise ValueError(f"{name} does not support request: {why}")
         assert isinstance(req, ConvGroupedRequest)
-        spec, why = _winograd_selection(req)
+        spec, why = selection(req)
         if spec is None:
-            raise ValueError(f"{_WINOGRAD_NAME} does not support request: {why}")
+            raise ValueError(f"{name} does not support request: {why}")
         return spec
 
     def build(spec: ConvWinogradSpec, arch: str):
         return _winograd_last_step(spec).kernel.build(arch)
 
     candidate = KernelCandidate(
-        name=_WINOGRAD_NAME,
+        name=name,
         family=_FAMILIES["fwd"],
-        algorithm="winograd",
-        spec_id=_WINOGRAD_NAME,
+        algorithm=algorithm,
+        spec_id=name,
         abi_version=CONV_WINOGRAD_ABI_VERSION,
         # After every listed instance; opt-in, so never picked under "auto".
         priority=1_000_000,
@@ -1314,9 +1389,7 @@ def build(spec: ConvWinogradSpec, arch: str):
         signature=lambda spec: _winograd_last_step(spec).signature,
         grid=lambda spec, req: _winograd_last_step(spec).grid,
         block=lambda spec: _winograd_last_step(spec).block,
-        sweep_space=lambda req: (
-            _winograd_sweep(req) if candidate.admits(req)[0] else ()
-        ),
+        sweep_space=lambda req: (sweep(req) if candidate.admits(req)[0] else ()),
         build=build,
         opt_in=True,
     )
@@ -1326,6 +1399,14 @@ def build(spec: ConvWinogradSpec, arch: str):
 _WINOGRAD_CANDIDATE = _make_winograd_candidate()
 _REGISTRIES["fwd"].register(_WINOGRAD_CANDIDATE)
 _EXECUTION_REGISTRIES["fwd"].register(_WINOGRAD_CANDIDATE)
+_WINOGRAD_AOT_CANDIDATE = _make_winograd_candidate(
+    _WINOGRAD_AOT_NAME,
+    "winograd_aot",
+    _winograd_aot_selection,
+    _winograd_aot_sweep,
+)
+_REGISTRIES["fwd"].register(_WINOGRAD_AOT_CANDIDATE)
+_EXECUTION_REGISTRIES["fwd"].register(_WINOGRAD_AOT_CANDIDATE)
 
 
 # ---------------------------------------------------------------------------