diff --git a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml index c443c9da269d..1f705a00892a 100644 --- a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml +++ b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml @@ -29,6 +29,9 @@ test_categories: - "rocke_arch_target_identity" - "rocke_ir_serialize_roundtrip" - "rocke_mma_frag_ssot" + - "rocke_streamk_iter_partition" + - "rocke_wgrad_streamk_host" + - "rocke_future_intrinsic_lowering" - "rocke_tiled_attention_2d_reentrancy" - "rocke_python_import_smoke" - "rocke_golden_static" @@ -52,6 +55,9 @@ test_categories: - "rocke_arch_target_identity" - "rocke_ir_serialize_roundtrip" - "rocke_mma_frag_ssot" + - "rocke_streamk_iter_partition" + - "rocke_wgrad_streamk_host" + - "rocke_future_intrinsic_lowering" - "rocke_tiled_attention_2d_reentrancy" - "rocke_python_import_smoke" - "rocke_golden_static" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py index ca8a82705a8a..ddd6cabc8349 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py @@ -230,6 +230,20 @@ def _build_implicit_only_args(args) -> list[str]: a += ["--split-k", str(args.split_k)] if args.split_k_prune is not None: a += ["--split-k-prune", str(args.split_k_prune)] + if args.streamk == "off" and ( + args.streamk_reductions is not None or args.streamk_ctas != -1 + ): + print( + "warning: --streamk-reductions / --streamk-ctas are ignored " + "without --streamk", + file=sys.stderr, + ) + if args.streamk != "off": + a += ["--streamk", args.streamk] + if args.streamk_reductions is not None: + a += ["--streamk-reductions", args.streamk_reductions] + if args.streamk_ctas != -1: + a += ["--streamk-ctas", str(args.streamk_ctas)] return a @@ -426,6 +440,35 @@ def main() -> int: "Only effective with --split-k 0. Forwarded to implicit-GEMM only." ), ) + implicit_grp.add_argument( + "--streamk", + default="off", + choices=["off", "dp_sk", "persistent", "all"], + help=( + "wgrad stream-K leg (off, dp_sk, persistent or all launch modes). " + "Forwarded to implicit-GEMM only." + ), + ) + implicit_grp.add_argument( + "--streamk-reductions", + default=None, + dest="streamk_reductions", + metavar="LIST", + help=( + "comma-separated stream-K reductions (linear,tree,workspace,atomic) " + "for --streamk. Forwarded to implicit-GEMM only." + ), + ) + implicit_grp.add_argument( + "--streamk-ctas", + type=int, + default=-1, + dest="streamk_ctas", + metavar="N", + help=( + "stream-K CTA pool (-1 = arch default). Forwarded to implicit-GEMM only." + ), + ) miopen_grp = parser.add_argument_group( "MIOpen input", 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 77688e97211d..e2938cab6290 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 @@ -77,6 +77,7 @@ EarlyStop, add_early_stop_arg, ) +from benchmarks.common.kernel_sweep import _AOT_ALIAS_PIPELINES # noqa: E402 from benchmarks.common.kernel_sweep import ( # noqa: E402 CACHE_EPILOGUES as _EPILOGUES, CACHE_PIPELINES as _PIPELINES, @@ -84,6 +85,8 @@ CACHE_TILE_MN as _TILE_MN, CACHE_WARP_MN as _WARP_MN, CACHE_WARP_TILE_MN as _WARP_TILE_MN, + CACHE_WGRAD_STREAMK_MODES as _STREAMK_MODES, + CACHE_WGRAD_STREAMK_REDUCTIONS as _STREAMK_REDUCTIONS, ) # gfx1250-only: wider CUs allow larger tiles and more warp concurrency. @@ -148,6 +151,10 @@ class WgradCombo(NamedTuple): # 16-bit atomics and allocates no scratch), which is why the axis is only # expanded on the two-stage work list. ws_replicas: int = 0 + # Stream-K launch mode ("dp_sk" / "persistent") and reduction. "off" on + # every other leg; the stream-K leg keeps split_k = 1 (stream-K owns K). + streamk: str = "off" + streamk_reduction: str = "" @dataclass @@ -182,6 +189,10 @@ class Result: # argument as group_merge -- an --ws-replicas sweep prints one row per R # and they must be tellable apart. ws_replicas: int = 0 + # Stream-K leg only: launch mode, reduction and the resolved CTA pool. + streamk: str = "off" + streamk_reduction: str = "" + streamk_ctas: int = 0 # --------------------------------------------------------------------------- @@ -927,6 +938,43 @@ def main() -> int: ).replace("%(sweep)s", str(list(_WS_REPLICAS_SWEEP))), ) + parser.add_argument( + "--streamk", + default="off", + choices=["off", "dp_sk", "persistent", "all"], + dest="streamk", + help=( + "wgrad stream-K kernels (gfx942/gfx950): off (default) = none; " + "dp_sk / persistent = that launch mode; all = both. The JIT sweep " + "adds a stream-K leg (default epilogue, split_k=1) over the " + "--streamk-reductions; --compile-all adds the stream-K kernels to " + "the cache; --run-from-cache runs the cached ones that match." + ), + ) + parser.add_argument( + "--streamk-reductions", + default=",".join(_STREAMK_REDUCTIONS), + dest="streamk_reductions", + metavar="LIST", + help=( + "comma-separated stream-K reductions to sweep with --streamk " + "(default: %(default)s). 'atomic' accumulates into dW with f32 " + "atomics, so it only applies when dW is fp32." + ), + ) + parser.add_argument( + "--streamk-ctas", + type=int, + default=-1, + dest="streamk_ctas", + metavar="N", + help=( + "stream-K CTA pool, a launch parameter (default: -1 = the arch " + "default: one CTA per CU for linear/tree, several per CU for " + "workspace/atomic)." + ), + ) + parser.add_argument( "--csv-top", type=int, @@ -1143,6 +1191,37 @@ def main() -> int: print(f"--csv-top must be >= 1 (got {args.csv_top})", file=sys.stderr) return 2 + if args.streamk == "off" and ( + args.streamk_reductions != ",".join(_STREAMK_REDUCTIONS) + or args.streamk_ctas != -1 + ): + print( + "warning: --streamk-reductions / --streamk-ctas are ignored " + "without --streamk", + file=sys.stderr, + ) + args.streamk_modes = { + "off": (), + "all": _STREAMK_MODES, + }.get(args.streamk, (args.streamk,)) + args.streamk_reductions = tuple( + r.strip() for r in args.streamk_reductions.split(",") if r.strip() + ) + for r in args.streamk_reductions: + if r not in _STREAMK_REDUCTIONS: + print( + f"--streamk-reductions: unknown reduction {r!r} " + f"(expected some of {', '.join(_STREAMK_REDUCTIONS)})", + file=sys.stderr, + ) + return 2 + if args.streamk_ctas == 0 or args.streamk_ctas < -1: + print( + f"--streamk-ctas must be -1 (auto) or > 0 (got {args.streamk_ctas})", + file=sys.stderr, + ) + return 2 + if args.miopen_cmd is None and args.miopen_file is None and args.json_file is None: if args.Di is not None and args.Z is None: print("--Z (filter depth) is required when --Di is set", file=sys.stderr) @@ -1638,6 +1717,16 @@ def _build_wgrad_one(args_tuple): group_merge=group_merge, lds_k_outer=lds_k_outer, async_dma=async_dma, + # Stream-K leg: the launch mode and reduction pick the kernel; the CTA + # pool stays at -1 here (it is a kernarg, resolved at launch). + **( + {} + if combo.streamk == "off" + else { + "streamk": combo.streamk, + "streamk_reduction": combo.streamk_reduction, + } + ), ) ok, _ = is_valid_wgrad_spec(spec, arch) if not ok: @@ -2716,6 +2805,53 @@ def _degrees(pipeline: str, async_dma: bool) -> tuple: else: pending_2s = [] + # Stream-K leg (--streamk): one kernel per geometry x pipeline x launch + # mode x reduction. Stream-K runs the default epilogue on a plain K loop + # and owns the K partition, so there is no cshuffle, async or split-K + # axis here; the pipelines that alias another's code are skipped as in + # the AOT grid. The CTA pool is a kernarg (--streamk-ctas). + _sk_reductions = tuple( + r for r in args.streamk_reductions if r != "atomic" or dtype == "fp32" + ) + pending_sk = [] + if args.streamk_modes and arch not in ("gfx942", "gfx950"): + print(f" --streamk: stream-K wgrad is gfx942/gfx950 only; skipped on {arch}.") + elif args.streamk_modes and _sk_reductions: + sk_combos = [ + WgradCombo( + *_geom, + _pipeline, + "default", + False, + 1, + streamk=_mode, + streamk_reduction=_red, + ) + for _geom in itertools.product( + _TILE_MN, _TILE_MN, _TILE_K, _WARP_MN, _WARP_MN, _WARP_TILE_MN + ) + if _geom[3] * _geom[5] <= _geom[0] and _geom[4] * _geom[5] <= _geom[1] + for _pipeline in _PIPELINES + if _pipeline not in _AOT_ALIAS_PIPELINES["wgrad"] + for _mode in args.streamk_modes + for _red in _sk_reductions + ] + if args.sample is not None: + sk_combos = _sample_combos(sk_combos, args.sample, args.seed + case_idx) + print( + f" stream-K: building {len(sk_combos)} combos " + f"(modes {list(args.streamk_modes)} x reductions " + f"{list(_sk_reductions)}) ...", + flush=True, + ) + pending_sk = _build_ir_parallel( + [(c, problem, dtype, arch) for c in sk_combos], + _build_and_compile_wgrad_one, + jobs, + ) + n_skipped += len(sk_combos) - len(pending_sk) + n_built += len(pending_sk) + # --------------------------------------------------------------------------- # Phase 3 – GPU run: load modules and time each kernel serially. # --------------------------------------------------------------------------- @@ -2989,7 +3125,6 @@ def _launch_spk(_v=_vals_snap, _c=_cfg_snap): ) from kernels.common.conv_wgrad_workspace_reduce import ( WgradReduceSpec, - wgrad_reduce_grid, wgrad_reduce_signature, ) @@ -3034,9 +3169,6 @@ def _launch_spk(_v=_vals_snap, _c=_cfg_snap): # Must match Stage 1 or the fold covers the wrong slab count. ws_replicas=spec.ws_replicas, ) - s2_grid = wgrad_reduce_grid(s2_spec) - s2_block = (s2_spec.tile_m * s2_spec.tile_n, 1, 1) - try: s1_launcher = KernelLauncher( hsaco=s1_art.hsaco, @@ -3081,19 +3213,16 @@ def _launch_spk(_v=_vals_snap, _c=_cfg_snap): ws_ptr=int(ws_dev), ws_bytes=ws_nbytes, ) - s2_values = { - "ws_ptr": ws_dev, - "dw_ptr": dW_dev2, - "wg_M": spec.wg_M, - "wg_N": spec.wg_N, - "ws_bytes": ws_nbytes, - "dw_bytes": dW_t.nbytes, - "groups": spec.problem.groups, - } + s2_values, s2_cfg = _stage2_launch( + s2_spec, + ws_ptr=ws_dev, + ws_bytes=ws_nbytes, + dw_ptr=dW_dev2, + dw_bytes=dW_t.nbytes, + ) s1_cfg = LaunchConfig( grid=s1_grid, block=(spec.block_size, 1, 1), stream=0 ) - s2_cfg = LaunchConfig(grid=s2_grid, block=s2_block, stream=0) def _launch_two_stage( _s1=s1_launcher, @@ -3198,6 +3327,133 @@ def _launch_two_stage( rt2.free(X_dev2) rt2.free(dW_dev2) + # --------------------------------------------------------------------------- + # Phase 3c – stream-K GPU run. Each launch zeroes what its reduction + # accumulates into (dW for atomic, the scratch for workspace, the flag + # table for linear/tree) and the workspace reduction runs its Stage 2. + # --------------------------------------------------------------------------- + if pending_sk: + from dataclasses import replace as _dc_replace + + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_signature + + print(f"\nSweeping {len(pending_sk)} stream-K wgrad configs ...", flush=True) + rt3 = Runtime() + dY_dev3 = rt3.alloc(dY_t.nbytes) + X_dev3 = rt3.alloc(X_t.nbytes) + dW_dev3 = rt3.alloc(dW_t.nbytes) + rt3.memcpy_h2d(dY_dev3, _u8(dY_t), dY_t.nbytes) + rt3.memcpy_h2d(X_dev3, _u8(X_t), X_t.nbytes) + _sk_stage2: dict = {} + _sk_stopped = False + try: + for combo, spec, _, artifact in pending_sk: + spec = _dc_replace(spec, streamk_ctas=args.streamk_ctas) + _mode = _streamk_mode_label(spec.streamk, spec.streamk_reduction) + _cfg_str = ( + f"tile={spec.tile_m}x{spec.tile_n}x{spec.tile_k} " + f"warp={spec.warp_m}x{spec.warp_n} " + f"atom={spec.warp_tile_m}x{spec.warp_tile_n}x{spec.warp_tile_k} " + f"{spec.pipeline}/{spec.epilogue:9s}" + ) + try: + launcher = KernelLauncher( + hsaco=artifact.hsaco, + kernel_name=artifact.kernel_name, + signature=wgrad_streamk_signature(spec), + ) + run, free_ws, pool = _streamk_launch( + rt=rt3, + launcher=launcher, + spec=spec, + problem=p, + arch=arch, + ptrs=(int(dY_dev3), int(X_dev3), int(dW_dev3)), + sizes=(dY_t.nbytes, X_t.nbytes, dW_t.nbytes), + stage2=_sk_stage2, + KernelLauncher=KernelLauncher, + LaunchConfig=LaunchConfig, + ) + except (HipError, ValueError) as e: + n_skipped += 1 + print(f"[skip] stream-K {_cfg_str} {_mode}: {e}", flush=True) + continue + try: + kernel_passed = None + if args.verify or args.dump_fail: + _sk_stopped, kernel_passed = _verify_kernel( + rt=rt3, + launch_fn=run, + out_dev=dW_dev3, + out_t=dW_t, + zero_init_out=True, + ref_out=ref_out, + kernel_name=artifact.kernel_name, + dump_fail=args.dump_fail, + extra_tensors={"dY": dY_t, "X": X_t}, + u8=_u8, + arch=arch, + compute_dtype=dtype, + ) + if _sk_stopped: + break + ms = _stop.measure( + run, + warmup=args.warmup, + iters=args.iters, + stream=0, + passed=kernel_passed, + ) + if ms is None: + _stop.report(f"{artifact.kernel_name} {_mode}") + continue + synchronize_and_release(0) + finally: + free_ws() + + cur_tflops = (flop / ms) * 1e-9 + cur_gbps = (bytes_xfer / ms) * 1e-6 + n_run += 1 + _va, _vb, _vc = WgradConvSpec.default_vector_sizes( + p.C, p.K, dtype, split_k=1 + ) + results.append( + Result( + kernel_name=artifact.kernel_name, + tile_m=spec.tile_m, + tile_n=spec.tile_n, + tile_k=spec.tile_k, + warp_m=spec.warp_m, + warp_n=spec.warp_n, + warp_tile_mn=spec.warp_tile_m, + warp_tile_k=spec.warp_tile_k, + pipeline=spec.pipeline, + epilogue=spec.epilogue, + split_k=1, + ms=ms, + tflops=cur_tflops, + gbps=cur_gbps, + vec_a=_va, + vec_b=_vb, + vec_c=_vc, + passed=kernel_passed, + streamk=spec.streamk, + streamk_reduction=spec.streamk_reduction, + streamk_ctas=pool, + ) + ) + print( + f"[{n_run:4d}] {_cfg_str} {_mode:<10s} ctas={pool:<5d} " + f"{cur_tflops:6.1f} TFLOPS {ms:.3f} ms", + flush=True, + ) + finally: + rt3.free(dY_dev3) + rt3.free(X_dev3) + rt3.free(dW_dev3) + if _sk_stopped: + return 1, [] + print(f"\nWgrad sweep done: {n_built} compiled, {n_skipped} skipped.", flush=True) if not results: @@ -3217,6 +3473,8 @@ def _launch_two_stage( print("-" * 92) for rank, r in enumerate(results[:top_n], 1): mode = f"spk{r.split_k}2s" if r.two_stage else f"spk{r.split_k}" + if r.streamk != "off": + mode = _streamk_mode_label(r.streamk, r.streamk_reduction) if r.group_merge > 1: mode += f" gm{r.group_merge}" if r.ws_replicas > 0: @@ -3619,6 +3877,8 @@ def _cache_dispatch(args, arch, target, cases) -> int: directions=directions or args.compile_directions, jobs=max(1, int(args.jobs or 1)), limit=args.limit, + streamk_modes=args.streamk_modes, + streamk_reductions=args.streamk_reductions, ), ) return rc @@ -3650,7 +3910,6 @@ def _run_from_cache(args, arch, target, cases, cache, directions) -> int: from rocke.runtime import synchronize_and_release from rocke.runtime.hip_module import HipError, Runtime from rocke.runtime.launcher import KernelLauncher, LaunchConfig - from benchmarks.common.kernel_sweep import _launch_values_for _TORCH_DT = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} @@ -3708,6 +3967,27 @@ def _u8(t): flush=True, ) continue + if direction == "wgrad": + # Stream-K kernels run only when asked for (--streamk), and + # only the requested modes/reductions, as in the JIT sweep. + n_before = len(candidates) + candidates = [ + c + for c in candidates + if c[0].streamk == "off" + or ( + c[0].streamk in args.streamk_modes + and c[0].streamk_reduction in args.streamk_reductions + ) + ] + if not candidates: + print( + f"Case {case_idx} {problem.short()} {case_dtype} wgrad: " + f"no compatible kernel ({n_before} stream-K kernels " + f"filtered out; see --streamk / --streamk-reductions)", + flush=True, + ) + continue n_compatible = len(candidates) n_vec_pruned = 0 if not args.all_vector_sizes: @@ -3798,6 +4078,31 @@ def _u8(t): print(f" [skip] {ident.label()}: {e}", flush=True) continue + # A stream-K binary partitions K itself over a CTA pool (a + # kernarg) and has its own ABI, workspace and zeroing + # contract per reduction, so it gets its own launch setup and + # one launch; its workspace is released after the loop. + sk_bound = sk_free = None + if direction == "wgrad" and ident.streamk != "off": + try: + sk_run, sk_free, sk_label = _cached_streamk_run( + args, + problem, + ident, + rt, + arch, + launcher, + (A_dev, B_dev, D_dev), + (A_t.nbytes, B_t.nbytes, D_t.nbytes), + stage2_kernels, + KernelLauncher, + LaunchConfig, + ) + except ValueError as e: + print(f" [skip] {ident.label()}: {e}", flush=True) + continue + sk_bound = (sk_run, sk_label) + # A split-K wgrad binary is one kernel for every degree > 1 # (the degree is a kernarg), so it is timed at each degree it # can use. @@ -3814,78 +4119,25 @@ def _u8(t): degrees = (None,) for split_k in degrees: - try: - extras, grid = _extras_and_grid( - direction, problem, ident, rt, case_dtype, split_k=split_k - ) - except ValueError as e: - print(f" [skip] {ident.label()}: {e}", flush=True) - continue - values = _launch_values_for( + bound = sk_bound or _cached_bind( direction, problem, ident, + meta, + rt, + launcher, + case_dtype, + split_k, (A_dev, B_dev, D_dev), (A_t.nbytes, B_t.nbytes, D_t.nbytes), - extras, - ) - cfg = LaunchConfig( - grid=grid, - block=( - meta.get("block_size") - or (ident.warp_m * ident.warp_n * ident.wave_size), - 1, - 1, - ), - stream=0, + arch, + stage2_kernels, + KernelLauncher, + LaunchConfig, ) - run = launcher.bind(values, config=cfg) - if direction == "wgrad" and ident.two_stage: - # 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. - from kernels.common.conv_wgrad_workspace_reduce import ( - wgrad_reduce_grid, - ) - - s2_launcher, s2_spec = _stage2_reduce( - problem, - ident.dtype_d, - arch, - stage2_kernels, - KernelLauncher, - _ws_replicas_of(ident), - ) - s2_values = { - "ws_ptr": extras["ws_ptr"], - "dw_ptr": D_dev, - "wg_M": s2_spec.wg_M, - "wg_N": s2_spec.wg_N, - "ws_bytes": extras["ws_bytes"], - "dw_bytes": D_t.nbytes, - "groups": problem.groups, - } - s2_cfg = LaunchConfig( - grid=wgrad_reduce_grid(s2_spec), - block=(s2_spec.block_size, 1, 1), - stream=0, - ) - - def run( - values=values, - cfg=cfg, - extras=extras, - s2_launcher=s2_launcher, - s2_values=s2_values, - s2_cfg=s2_cfg, - ): - rt.memset(extras["ws_ptr"], 0, extras["ws_bytes"]) - launcher(values, config=cfg) - s2_launcher(s2_values, config=s2_cfg) - - label = ident.label() - if split_k is not None: - label += f" @split_k={split_k}" + if bound is None: + continue + run, label = bound if ref_out is not None: # The output is zeroed first: the atomic split-K # epilogues accumulate into it. A kernel that fails is @@ -3933,6 +4185,9 @@ def run( flush=True, ) + if sk_free is not None: + sk_free() + synchronize_and_release() if ref_out is not None: print( @@ -3991,6 +4246,8 @@ def kind(ident): ident.split_k, ident.two_stage, ident.group_merge, + ident.streamk, + ident.streamk_reduction, ) def rank(ident): @@ -4007,7 +4264,14 @@ def rank(ident): # kernels are dropped. Where nothing stores wider (two-stage wgrad is # default-only, or the problem's output run is odd) they stay. def store_kind(ident): - return (ident.direction, ident.split_k, ident.two_stage, ident.group_merge) + return ( + ident.direction, + ident.split_k, + ident.two_stage, + ident.group_merge, + ident.streamk, + ident.streamk_reduction, + ) widest_store = {} for ident, *_ in candidates: @@ -4059,7 +4323,11 @@ def _signature_for_identity(ident, dtype): if ident.direction == "wgrad": return conv_args_signature( - dtype, direction="wgrad", is_3d=ident.is_3d, two_stage=ident.two_stage + dtype, + direction="wgrad", + is_3d=ident.is_3d, + two_stage=ident.two_stage, + streamk=ident.streamk_reduction if ident.streamk != "off" else None, ) if ident.direction == "dgrad": return conv_args_signature(dtype, direction="dgrad", is_3d=ident.is_3d) @@ -4112,6 +4380,189 @@ def _stage2_reduce(problem, dtype_d, arch, compiled, KernelLauncher, ws_replicas return compiled[key], spec +def _stage2_launch(s2_spec, *, ws_ptr, ws_bytes, dw_ptr, dw_bytes): + """``(values, config)`` of the Stage 2 fold of ``s2_spec`` into dW. + + Stage 2 (:class:`WgradReduceSpec`) folds the f32 scratch at ``ws_ptr`` + into dW for both the two-stage split-K path and the stream-K + ``workspace`` reduction; every leg launches it through here. + """ + from rocke.runtime.launcher import LaunchConfig + from kernels.common.conv_wgrad_workspace_reduce import wgrad_reduce_grid + + values = { + "ws_ptr": ws_ptr, + "dw_ptr": dw_ptr, + "wg_M": s2_spec.wg_M, + "wg_N": s2_spec.wg_N, + "ws_bytes": ws_bytes, + "dw_bytes": dw_bytes, + "groups": s2_spec.groups, + } + config = LaunchConfig( + grid=wgrad_reduce_grid(s2_spec), block=(s2_spec.block_size, 1, 1), stream=0 + ) + return values, config + + +def _streamk_launch( + *, + rt, + launcher, + spec, + problem, + arch: str, + ptrs, + sizes, + stage2: dict, + KernelLauncher, + LaunchConfig, +): + """Set up one stream-K wgrad launch per the kernel's zeroing contract. + + ``spec`` is the stream-K :class:`WgradConvSpec` (its ``streamk_ctas`` is + the CTA pool, -1 = the arch default); ``problem`` the launch shape; + ``ptrs`` / ``sizes`` the dY, X, dW buffers. Allocates the reduction's + workspace and returns ``(run, free, pool)``: ``run()`` zeroes what the + reduction accumulates into -- dW for ``atomic``, the f32 scratch for + ``workspace``, only the flag table for ``linear``/``tree`` -- launches the + kernel and, for ``workspace``, the Stage 2 reduce into dW; ``free()`` + releases the workspace; ``pool`` is the resolved CTA pool. + """ + from kernels.common.conv_implicit_gemm_wgrad import ( + wgrad_streamk_grid, + wgrad_streamk_launch, + wgrad_streamk_launch_values, + wgrad_streamk_workspace_layout, + wgrad_streamk_workspace_nbytes, + ) + + reduction = spec.streamk_reduction + ws_bytes = wgrad_streamk_workspace_nbytes(spec, arch=arch, problem=problem) + flags_bytes, _ = wgrad_streamk_workspace_layout(spec, arch=arch, problem=problem) + grid = wgrad_streamk_grid(spec, arch=arch, problem=problem) + pool = wgrad_streamk_launch(spec, arch=arch).pool + # An all-data-parallel partition (no stream-K CTAs) needs no workspace, + # but the kernargs still carry the pointers, so a 4-byte stub stands in + # for the empty allocation. + ws = rt.alloc(max(ws_bytes, 4)) + try: + extra = {} + if reduction == "workspace": + extra = {"ws_ptr": ws, "ws_bytes": ws_bytes} + elif reduction in ("linear", "tree"): + extra = {"sk_flags": ws, "sk_partials": int(ws) + flags_bytes} + values = wgrad_streamk_launch_values( + spec, *ptrs, *sizes, arch=arch, problem=problem, **extra + ) + cfg = LaunchConfig(grid=grid, block=(spec.block_size, 1, 1), stream=0) + s2 = None + if reduction == "workspace": + s2_launcher, s2_spec = _stage2_reduce( + problem, + spec.data.dtype_d, + arch, + stage2, + KernelLauncher, + spec.ws_replicas, + ) + s2 = ( + s2_launcher, + *_stage2_launch( + s2_spec, + ws_ptr=ws, + ws_bytes=ws_bytes, + dw_ptr=ptrs[2], + dw_bytes=sizes[2], + ), + ) + except Exception: + rt.free(ws) + raise + + def run(): + if reduction == "atomic": + rt.memset(ptrs[2], 0, sizes[2]) + elif reduction == "workspace": + rt.memset(ws, 0, ws_bytes) + elif flags_bytes: + rt.memset(ws, 0, flags_bytes) + launcher(values, config=cfg) + if s2 is not None: + s2[0](s2[1], config=s2[2]) + + def free(): + rt.free(ws) + + return run, free, pool + + +def _streamk_mode_label(mode: str, reduction: str) -> str: + """Short ranked-table label of a stream-K kernel, e.g. ``skdp-tree``.""" + m = {"dp_sk": "dp", "persistent": "pers"}.get(mode, mode) + r = {"linear": "lin", "workspace": "ws", "atomic": "atom"}.get(reduction, reduction) + return f"sk{m}-{r}" + + +def _cached_streamk_run( + args, + problem, + ident, + rt, + arch, + launcher, + ptrs, + sizes, + stage2, + KernelLauncher, + LaunchConfig, +): + """``(run, free, label)`` for a cached stream-K wgrad binary. + + The launch helpers take a spec; this one describes the cached binary on + the launch shape. Nothing is built from it -- the partition, workspace and + kernargs depend only on the tile, the warp layout, the stream-K mode and + reduction, the scratch replicas and the launch shape. + """ + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + from kernels.common.conv_implicit_gemm_wgrad import WgradConvSpec + + spec = WgradConvSpec( + problem=problem, + data=ConvDataSpec( + dtype_a=ident.dtype_a, dtype_b=ident.dtype_b, dtype_d=ident.dtype_d + ), + tile_m=ident.tile_m, + tile_n=ident.tile_n, + tile_k=ident.tile_k, + warp_m=ident.warp_m, + warp_n=ident.warp_n, + warp_tile_m=ident.warp_tile_m, + warp_tile_n=ident.warp_tile_n, + warp_tile_k=ident.warp_tile_k, + wave_size=ident.wave_size, + pipeline=ident.pipeline, + epilogue=ident.epilogue, + streamk=ident.streamk, + streamk_reduction=ident.streamk_reduction, + streamk_ctas=args.streamk_ctas, + ws_replicas=_ws_replicas_of(ident), + ) + run, free, pool = _streamk_launch( + rt=rt, + launcher=launcher, + spec=spec, + problem=problem, + arch=arch, + ptrs=ptrs, + sizes=sizes, + stage2=stage2, + KernelLauncher=KernelLauncher, + LaunchConfig=LaunchConfig, + ) + return run, free, f"{ident.label()} @streamk_ctas={pool}" + + def _split_k_degrees(problem, tile_m: int, tile_n: int, tile_k: int) -> tuple: """Split-K degrees to launch a wgrad split-K kernel with. @@ -4143,6 +4594,85 @@ def _runtime_split_ks(problem, ident) -> tuple: return _split_k_degrees(problem, ident.tile_m, ident.tile_n, ident.tile_k) +def _cached_bind( + direction, + problem, + ident, + meta, + rt, + launcher, + case_dtype, + split_k, + ptrs, + sizes, + arch, + stage2_kernels, + KernelLauncher, + LaunchConfig, +): + """``(run, label)`` of a cached (non-stream-K) binary at ``split_k``. + + ``ptrs`` / ``sizes`` are the A, B, D buffers. Returns None, after saying + why, when the binary cannot launch ``problem``. + """ + from benchmarks.common.kernel_sweep import _launch_values_for + + try: + extras, grid = _extras_and_grid( + direction, problem, ident, rt, case_dtype, split_k=split_k + ) + except ValueError as e: + print(f" [skip] {ident.label()}: {e}", flush=True) + return None + values = _launch_values_for(direction, problem, ident, ptrs, sizes, extras) + cfg = LaunchConfig( + grid=grid, + block=( + meta.get("block_size") or (ident.warp_m * ident.warp_n * ident.wave_size), + 1, + 1, + ), + stream=0, + ) + run = launcher.bind(values, config=cfg) + if direction == "wgrad" and ident.two_stage: + # 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. + s2_launcher, s2_spec = _stage2_reduce( + problem, + ident.dtype_d, + arch, + stage2_kernels, + KernelLauncher, + _ws_replicas_of(ident), + ) + s2_values, s2_cfg = _stage2_launch( + s2_spec, + ws_ptr=extras["ws_ptr"], + ws_bytes=extras["ws_bytes"], + dw_ptr=ptrs[2], + dw_bytes=sizes[2], + ) + + def run( + values=values, + cfg=cfg, + extras=extras, + s2_launcher=s2_launcher, + s2_values=s2_values, + s2_cfg=s2_cfg, + ): + rt.memset(extras["ws_ptr"], 0, extras["ws_bytes"]) + launcher(values, config=cfg) + s2_launcher(s2_values, config=s2_cfg) + + label = ident.label() + if split_k is not None: + label += f" @split_k={split_k}" + return run, label + + def _extras_and_grid(direction, problem, ident, rt, dtype, split_k=None): """Trailing kernargs and the launch grid for one cached kernel. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py index ad6ecaf7cca9..f21025f7ccdd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py @@ -200,6 +200,12 @@ class KernelIdentity: # binary's identity rather than a default the two could drift apart on. # 0 = no scratch (not two-stage). ws_replicas: int = 0 + # wgrad stream-K (WgradConvSpec.streamk / streamk_reduction): the launch + # mode ("dp_sk" or "persistent") and the partial-sum reduction are baked + # into the ISA; the CTA pool is a kernarg, so it is not here. "off" / "" + # for every other kernel. + streamk: str = "off" + streamk_reduction: str = "" group_merge: int = 1 num_load_waves: int = 0 cshuffle_no_alias: bool = False @@ -271,6 +277,9 @@ def stable_hash(self) -> str: # field keep their hash (and their cache entries). if not fields["ws_replicas"]: del fields["ws_replicas"] + # 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"] payload = json.dumps(fields, sort_keys=True, separators=(",", ":")) return hashlib.sha1(payload.encode()).hexdigest() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_label.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_label.py index b1fb7cfe491e..d4195736f94d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_label.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_label.py @@ -68,6 +68,8 @@ split_k=1, two_stage=False, ws_replicas=0, + streamk="off", + streamk_reduction="", group_merge=1, num_load_waves=0, cshuffle_no_alias=False, 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 f438c607b590..73e2d931cd5e 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 @@ -51,6 +51,14 @@ ``ws_replicas`` = 8). Plus group-merged depthwise kernels (``group_merge`` in 2..64, two-stage, load widths derived by the builder), served only to depthwise problems whose merged GEMM fits the tile. + Opt-in (``streamk_modes``; ``--compile-all --streamk``): stream-K + kernels, default epilogue only, ``split_k`` = 1, c = 1. At the widest + vector width W of the dtype (a = W; b in {W, 1}): one per launch mode + (``dp_sk``, ``persistent``) x reduction (``linear``, ``workspace`` with + ``ws_replicas`` = 8, and ``atomic`` where dW is fp32). Fully scalar + (a = b = 1), on tiles with tile_n >= 128 only: ``dp_sk`` x (``tree``, + ``workspace``), for large-filter depthwise layers. The CTA pool is a + kernarg, so it is not swept here. * **dgrad** -- vector widths: a in {1, 2, 4, 8} (dY along kpg), b = c in {1, 2, 4, 8} (W and dX along cpg); with the default epilogue c = 1 and b in {1, 2, 4, 8}. stride and dilation are baked into the @@ -175,6 +183,40 @@ # separately if you need it; the wgrad benchmark's JIT sweep still covers it. CACHE_WGRAD_SPLIT_KS: Tuple[int, ...] = (2,) +# ---- 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 +# partition -- is a kernarg, so one binary per (mode, reduction) serves every +# shape and pool. It is off unless a caller asks for it (--streamk): it adds +# a default-epilogue kernel per (mode, reduction) to every wgrad geometry, +# which more than doubles the wgrad grid. "atomic" accumulates into dW with +# f32 atomics and the validator keeps it to fp32 dW only. +# +# Only a slice of the (mode, reduction) x vector-width grid is cached (see +# _wgrad_streamk_builds). Stream-K pays off where the shape loads at the +# widest vector width, so the requested variants are built there, minus +# "tree" ("linear" is the other deterministic fixup). Narrower widths are left +# to the split-K kernels, which are cached at every width and serve those +# shapes as well -- except for one case: large-filter depthwise layers, whose +# per-group reduction is long and whose GEMM is wide. They load scalar, and +# stream-K wins there, so a scalar dp_sk slice is kept on the wide tiles. +CACHE_WGRAD_STREAMK_MODES: Tuple[str, ...] = ("dp_sk", "persistent") +CACHE_WGRAD_STREAMK_REDUCTIONS: Tuple[str, ...] = ( + "linear", + "tree", + "workspace", + "atomic", +) +# Bytes per stream-K operand load at the widest vector width: 8 elements for +# fp16/bf16, 4 for fp32. +CACHE_WGRAD_STREAMK_LOAD_BYTES = 16 +# Reductions not cached at the widest vector widths. +CACHE_WGRAD_STREAMK_WIDE_EXCLUDED_REDUCTIONS: Tuple[str, ...] = ("tree",) +# The scalar (vector width 1) slice for large-filter depthwise layers. +CACHE_WGRAD_STREAMK_SCALAR_MODES: Tuple[str, ...] = ("dp_sk",) +CACHE_WGRAD_STREAMK_SCALAR_REDUCTIONS: Tuple[str, ...] = ("tree", "workspace") +CACHE_WGRAD_STREAMK_SCALAR_MIN_TILE_N = 128 + # Private aliases for internal use -- the generators below refer to these. _TILE_MN = CACHE_TILE_MN _TILE_K = CACHE_TILE_K @@ -418,6 +460,56 @@ def _async_chunks(cfg: dict, dtypes, caps: dict, arch: str) -> dict: ) +def _wgrad_streamk_variants( + epilogue: str, + streamk_modes: Sequence[str], + streamk_reductions: Sequence[str], +) -> List[Tuple[str, str]]: + """``(mode, reduction)`` stream-K variants built per wgrad geometry. + + Stream-K runs its epilogue once per tile inside the stream-K tile loop and + only the default epilogue supports that, so cshuffle geometries get none. + """ + if epilogue != "default": + return [] + return list(itertools.product(streamk_modes, streamk_reductions)) + + +def _wgrad_streamk_builds( + epilogue: str, + tile_n: int, + dtype: str, + streamk_modes: Sequence[str], + streamk_reductions: Sequence[str], +) -> List[Tuple[str, str, int, int]]: + """``(mode, reduction, vector_size_a, vector_size_b)`` built per geometry. + + Of the requested variants: every one but the + ``CACHE_WGRAD_STREAMK_WIDE_EXCLUDED_REDUCTIONS`` at the widest width W of + the dtype on both operands, and at W on dY with a scalar X for the + narrow-channel layers (C = 3 stems) that cannot load X wide; plus the + ``CACHE_WGRAD_STREAMK_SCALAR_*`` slice, fully scalar, on tiles at least + ``CACHE_WGRAD_STREAMK_SCALAR_MIN_TILE_N`` wide. Every other combination + falls to the split-K kernels. + """ + variants = _wgrad_streamk_variants(epilogue, streamk_modes, streamk_reductions) + w = CACHE_WGRAD_STREAMK_LOAD_BYTES // {"fp16": 2, "bf16": 2, "fp32": 4}[dtype] + builds = [ + (mode, reduction, vec_a, vec_b) + for mode, reduction in variants + if reduction not in CACHE_WGRAD_STREAMK_WIDE_EXCLUDED_REDUCTIONS + for vec_a, vec_b in ((w, w), (w, 1)) + ] + if tile_n >= CACHE_WGRAD_STREAMK_SCALAR_MIN_TILE_N: + builds += [ + (mode, reduction, 1, 1) + for mode, reduction in variants + if mode in CACHE_WGRAD_STREAMK_SCALAR_MODES + and reduction in CACHE_WGRAD_STREAMK_SCALAR_REDUCTIONS + ] + return builds + + def _wgrad_jobs( arch, dtype, @@ -428,6 +520,8 @@ def _wgrad_jobs( *, geometries: Optional[Sequence[tuple]] = None, llvm_flavor: Optional[str] = None, + streamk_modes: Sequence[str] = (), + streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, ) -> Iterator[BuildJob]: from kernels.common.conv_implicit_gemm_wgrad import _DEFAULT_WS_REPLICAS @@ -496,6 +590,52 @@ def _wgrad_jobs( spec_kwargs=cfg, caps=caps, ) + # Stream-K kernels (opt-in). Default epilogue only, so the store is + # scalar and B takes its own width; split_k stays 1 (stream-K owns the + # K partition). The workspace reduction reuses the two-stage scratch, + # so its replica count is pinned exactly as for two-stage. + for (mode, reduction, vec_a, vec_b), caps in itertools.product( + _wgrad_streamk_builds( + epilogue, tile_n, dtype, streamk_modes, streamk_reductions + ), + _layout_caps(), + ): + 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_a, + vector_size_b=vec_b, + vector_size_c=1, + streamk=mode, + streamk_reduction=reduction, + ) + if reduction == "workspace": + cfg["ws_replicas"] = _DEFAULT_WS_REPLICAS + yield BuildJob( + identity=KernelIdentity( + arch=arch, + direction="wgrad", + algorithm="implicit_gemm", + llvm_flavor=flavor, + dtype_a=da, + dtype_b=db, + dtype_d=dd, + grouped=caps["grouped"], + **cfg, + ), + direction="wgrad", + spec_kwargs=cfg, + caps=caps, + ) # Group-merged depthwise kernels. Merged split-K must take the # two-stage path (the packed atomic cannot drop off-diagonal pairs). # The vector widths are left to the builder: it derives them from the @@ -677,6 +817,14 @@ def job_for_identity(identity: KernelIdentity) -> BuildJob: atom, ) want = identity.stable_hash() + streamk = ( + { + "streamk_modes": (identity.streamk,), + "streamk_reductions": (identity.streamk_reduction,), + } + if identity.direction == "wgrad" and identity.streamk != "off" + else {} + ) for job in _direction_jobs( identity.direction, identity.arch, @@ -686,6 +834,7 @@ def job_for_identity(identity: KernelIdentity) -> BuildJob: identity.max_sub_gemms, geometries=[geometry], llvm_flavor=identity.llvm_flavor, + **streamk, ): if job.identity.stable_hash() == want: return job @@ -745,6 +894,8 @@ def count_jobs( target, directions: Sequence[str], split_ks: Sequence[int] = CACHE_WGRAD_SPLIT_KS, + streamk_modes: Sequence[str] = (), + streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, ) -> int: """Raw candidate count :func:`enumerate_jobs` walks, before dedup/validation. @@ -766,6 +917,11 @@ def count_jobs( if "wgrad" in directions and _aot_pipeline("wgrad", pipeline): total += vecs * sum(len(_wgrad_two_stages(sk)) for sk in split_ks) * layouts total += merged + total += layouts * len( + _wgrad_streamk_builds( + geo[7], geo[1], dtype, streamk_modes, streamk_reductions + ) + ) if "dgrad" in directions and _aot_pipeline("dgrad", pipeline): total += vecs * dgrad_caps return total @@ -778,6 +934,9 @@ def _direction_jobs( target, split_ks: Sequence[int], max_sub_gemms: int, + *, + streamk_modes: Sequence[str] = (), + streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, **kw, ) -> Iterator[BuildJob]: wave_size = target.wave_size @@ -785,7 +944,17 @@ def _direction_jobs( if direction == "fwd": return _fwd_jobs(arch, dtype, wave_size, mma_family, target, **kw) if direction == "wgrad": - return _wgrad_jobs(arch, dtype, wave_size, mma_family, target, split_ks, **kw) + return _wgrad_jobs( + arch, + dtype, + wave_size, + mma_family, + target, + split_ks, + streamk_modes=streamk_modes, + streamk_reductions=streamk_reductions, + **kw, + ) if direction == "dgrad": return _dgrad_jobs( arch, dtype, wave_size, mma_family, target, max_sub_gemms, **kw @@ -809,6 +978,7 @@ def _enumerate_chunk( validate: bool, llvm_flavor: str, geometries: Sequence[tuple], + streamk: Tuple[Sequence[str], Sequence[str]] = ((), ()), ) -> Tuple[int, List[Tuple[str, BuildJob]]]: """Walk one direction over ``geometries``. @@ -829,6 +999,8 @@ def _enumerate_chunk( max_sub_gemms, geometries=geometries, llvm_flavor=llvm_flavor, + streamk_modes=streamk[0], + streamk_reductions=streamk[1], ): n_raw += 1 key = job.identity.stable_hash() @@ -849,9 +1021,18 @@ def _enumerate_chunk( def _enumerate_worker(payload): - (direction, arch, dtype, split_ks, max_sub_gemms, validate, flavor, lo, hi) = ( - payload - ) + ( + direction, + arch, + dtype, + split_ks, + max_sub_gemms, + validate, + flavor, + streamk, + lo, + hi, + ) = payload from rocke.core.arch import ArchTarget target = ArchTarget.from_gfx(arch) @@ -868,6 +1049,7 @@ def _enumerate_worker(payload): validate, flavor, geos[lo:hi], + streamk, ) @@ -902,9 +1084,14 @@ def enumerate_jobs( jobs: int = 1, log=None, log_every_s: float = 5.0, + streamk_modes: Sequence[str] = (), + streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, ) -> List[BuildJob]: """All buildable jobs for the requested directions, deduped by identity. + ``streamk_modes`` (empty by default) adds the opt-in wgrad stream-K + kernels for those launch modes x ``streamk_reductions``. + With ``validate=True`` (the default) each candidate is run through its direction's spec validator first, so the returned list is what will actually compile rather than the raw cross product. @@ -931,8 +1118,18 @@ def enumerate_jobs( (len(geos) * i // n_chunks, len(geos) * (i + 1) // n_chunks) for i in range(n_chunks) ] + streamk = (tuple(streamk_modes), tuple(streamk_reductions)) payloads = [ - (direction, arch, dtype, tuple(split_ks), max_sub_gemms, validate, flavor) + ( + direction, + arch, + dtype, + tuple(split_ks), + max_sub_gemms, + validate, + flavor, + streamk, + ) + bound for direction in directions for bound in bounds @@ -941,7 +1138,12 @@ def enumerate_jobs( total = 0 if log is not None: total = count_jobs( - dtype=dtype, target=target, directions=directions, split_ks=split_ks + dtype=dtype, + target=target, + directions=directions, + split_ks=split_ks, + streamk_modes=streamk_modes, + streamk_reductions=streamk_reductions, ) log(f" {total} candidates to check with {jobs} process(es)") @@ -981,6 +1183,7 @@ def _done(idx: int, chunk_raw: int, chunk_out) -> None: validate, flavor, geos[lo:hi], + streamk, ), ) else: @@ -1273,9 +1476,14 @@ def compile_all( jobs: int = 1, limit: Optional[int] = None, log=print, + streamk_modes: Sequence[str] = (), + streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, ) -> int: """Populate ``cache`` with every implicit-GEMM variant for ``arch``/``dtype``. + ``streamk_modes`` opts the wgrad stream-K kernels in (see + :data:`CACHE_WGRAD_STREAMK_MODES`). + Compilation is embarrassingly parallel and dominated by LLVM, so it fans out over processes; ``jobs`` defaults to one but the caller normally passes ``os.cpu_count()``. @@ -1294,6 +1502,8 @@ def compile_all( directions=directions, jobs=jobs, log=log, + streamk_modes=streamk_modes, + streamk_reductions=streamk_reductions, ), build=build_kernel, arch=arch, @@ -1321,6 +1531,10 @@ def compile_all( lambda i: ("two-stage" if i.two_stage else "atomic") if i.split_k > 1 else "no", ), ("group_merge", lambda i: i.group_merge), + ( + "stream-K", + lambda i: f"{i.streamk}/{i.streamk_reduction}" if i.streamk != "off" else "no", + ), ("stride x dilation", lambda i: f"{i.stride_h}x{i.dilation_h}"), ("filter", lambda i: f"{i.filter_h}x{i.filter_w}"), ("pad", lambda i: i.pad_h), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm.md index a02ac3da23ab..ccf11e0897f4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm.md @@ -233,6 +233,13 @@ the K loop across CTAs (similar to split-K but without the atomic reduction overhead), potentially delivering better utilisation for small or irregular shapes. +This has landed for **wgrad** (`streamk` on `WgradConvSpec`; see +`README_conv_implicit_gemm_wgrad.md`), where the lopsided `K_wg = N*Ho*Wo` +reduction makes it most valuable. The forward direction still maps one CTA per +output tile; the partitioner (`rocke.helpers.streamk.StreamKIterPartition`) and +the flag-based fixups are family-agnostic, so porting them here is wiring rather +than new machinery. + ### Wider vector loads for A The A descriptor emits one element offset per thread per load. Enabling vector diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index 7700ef1b4515..a05d54d52ff9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -246,6 +246,101 @@ Both loops that Python-unroll the K iteration — `pipeline="basic"` and `async_dma` — are bounded by `_MAX_UNROLLED_K_ITERS` (128). Over the cap the spec is rejected; raise `split_k` or `tile_k` rather than the constant. +### Stream-K (`streamk`, `streamk_reduction`, `streamk_ctas`) + +Split-K cuts `K_wg` into `split_k` equal slices and gives every output tile the +same number of them. That is the right shape when the tiles outnumber the CUs, +but a wgrad is usually the opposite: `wg_M x wg_N` is filter-sized while +`K_wg = N*Ho*Wo` is huge, so the tile count can be a handful and the leftover +CUs idle no matter how the slices are cut. + +Stream-K (CK Tile's `StreamKTilePartitioner`) balances *MAC iterations* instead. +Output tiles that divide the CTA pool evenly run data-parallel — one CTA, the +whole K loop, no reduction at all — and only the remainder is spread over +`streamk_ctas` stream-K CTAs, each owning a contiguous, balanced range of +iterations that may straddle tile boundaries. The partition is resolved +host-side (`rocke.helpers.streamk.StreamKIterPartition`, reached through +`ConvArgs.streamk_partition`) and passed in as the `sk_*` kernargs; the device +only maps `blockIdx.x` onto its iteration range. Like the split-K degree, the +CTA pool is a launch parameter rather than a build-time constant, so one AOT +binary serves any shape and any pool. + +`streamk="dp_sk"` launches `dp_tiles + sk_ctas` CTAs. `streamk="persistent"` +launches only the pool: each CTA round-robins the data-parallel tiles +(`tile = cta + t*pool`, CK's static `StreamKDispatch`, no work queue) and then +runs its stream-K range. Both walk their work with one bounded `scf.for` whose +trip count is CTA-uniform, so the barriers inside the GEMM body stay uniform and +the body is emitted once. + +`streamk_reduction` picks how the CTAs sharing a tile combine partial sums: + +| | Determinism | dW zeroing | Workspace | +|---|---|---|---| +| `linear` | yes | not needed | flags + one f32 tile per SK CTA | +| `tree` | yes | not needed | same | +| `atomic` | no | required | none (fp32 dW only) | +| `workspace` | no | not needed | the two-stage f32 scratch + Stage 2 reduce | + +`linear` and `tree` hand partials over explicitly: a non-owner stores its tile +to its own slot, issues a release fence and sets a per-wave flag; the consumer +spins on an acquire load of that flag and then adds. Both are deterministic +because the fold order is fixed by the partition, not by arrival. `linear` has +the owner walk its later contributors in turn; `tree` fans in pairwise over +`ceil(log2)` rounds. The flag protocol is three core IR ops (`memref.fence`, +`memref.global_flag_store`, `memref.global_flag_wait_eq`) rather than raw cache +hints, so the AMDGPU memory model supplies the L2 writeback/invalidate that +makes a partial visible across XCDs. + +Size the buffer with `wgrad_streamk_workspace_nbytes`; split it with +`wgrad_streamk_workspace_layout` (`sk_flags` at the base, `sk_partials` at +`base + flags_bytes`). **Only the flag region has to be zeroed before each +launch** — partials are always written before they are read. + +Constraints, all enforced by `validate()` / `is_valid_wgrad_spec` in both +engines: gfx942/gfx950 (MFMA wave64) only; `split_k=1` and no `two_stage` +(stream-K owns the K range); `epilogue="default"`; no `async_dma` / `unroll_k` +(they pipeline the K loop across work items, which stream-K re-enters per tile); no `chiplet_swizzle` or +`group_merge` (both remap tiles, which stream-K owns). Conv groups are folded +into GEMM-M instead of riding `block_id_z`: `m_tiles = ceil(wg_M/tile_m) * +groups`, decoded group-fastest. Defining the tile count per group makes the fold +exact for any `wg_M`, avoiding CK's `GemmM % MPerBlock` caveat. + +`streamk_ctas` sizes the host partition only; `-1` resolves per reduction +(`wgrad_streamk_default_ctas`): + +* `workspace` / `atomic` never wait on another CTA, so their pool is sized for + throughput: four CTAs per CU (`_STREAMK_CTAS_PER_CU`). Each resident CTA + streams its own K range, so this counts independent reduction streams per CU, + and is counted in CTAs rather than waves: a wider CTA adds no streams. One + CTA per CU leaves a single stream per CU with nothing to hide memory latency + behind, which is what made a CU-sized pool lose to split-K on most shapes; + past four the CTAs mostly contend for the same caches. +* `linear` / `tree` stay at one CTA per CU, because these fixups get slower + with a larger pool: the tile owner folds every additional contributor. The + pool does not have to be resident at once. A CTA waits only on CTAs with + higher indices and publishes its own partial before it waits, so with + workgroups dispatched in ascending order only the CTAs of the one tile that + straddles the dispatch front (its owner for `linear`, about one per round + for `tree`) can wait on a CTA that has not started yet. + +Not selected by the grouped-convolution dispatcher: `ConvGroupedRequest` has +no stream-K field and the dispatcher only builds the grid-per-tile kernels, so +stream-K is reached by building a `WgradConvSpec` with `streamk` set. To +measure it, use the implicit-GEMM sweep +(`library/benchmarks/common/benchmark_implicit_gemm_conv.py`, also driven by +`benchmark_conv_compare.py`) with `--direction wgrad --streamk {dp_sk,persistent,all}`, +optionally narrowed with `--streamk-reductions` and pinned to a pool with +`--streamk-ctas` (default `-1`, the per-reduction pool above). The JIT sweep then +adds a stream-K leg, one kernel per geometry x pipeline x mode x reduction +(`atomic` only for fp32 dW); `--compile-all --streamk ...` adds stream-K kernels +to the AOT cache (`kernel_sweep.py`, opt-in) and `--run-from-cache --streamk ...` +runs the cached ones. The cache keeps a subset: the widest vector width of the +dtype without `tree`, plus a scalar `dp_sk` `tree`/`workspace` slice on wide +tiles for large-filter depthwise layers; other narrow-channel shapes fall to the +split-K kernels. Every +launch follows the zeroing contract above, and `workspace` launches are timed +together with their Stage 2 reduce. + ## Next steps ### K0-M-K1 LDS layout 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 a214c9528c89..a790b045e050 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 @@ -37,6 +37,7 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple __all__ = [ + "STREAMK_REDUCTIONS", "ArgSpec", "conv_arg_names", "conv_args_signature", @@ -151,10 +152,28 @@ def conv_fwd_problem_block(*, is_3d: bool = False) -> List[ArgSpec]: return _fwd_arg_names(is_3d=is_3d)[_LEADING_ARGS:] +# Stream-K reductions of the wgrad kernel, as ABI selectors. +STREAMK_REDUCTIONS = ("atomic", "workspace", "linear", "tree") + +# The stream-K launch block: the host-computed iteration partition. The kernel +# is shape- and pool-agnostic, so every value the partition derives from the +# launch shape and the CTA pool arrives here rather than being baked in. +_STREAMK_SCALARS = ( + "sk_iters_per_tile", # ceil(wg_K / tile_k): MAC iterations per output tile + "sk_dp_tiles", # tiles owned whole by data-parallel CTAs / the DP sweep + "sk_ctas", # stream-K CTAs (0 for an all-data-parallel partition) + "sk_iters_per_cta", # base iterations per stream-K CTA; never 0 (a divisor) + "sk_extra_iters", # the first sk_extra_iters stream-K CTAs own one more + "sk_pool", # the CTA pool; the persistent launch's DP-sweep stride + "sk_tree_rounds", # pairwise rounds the tree fixup runs +) + + def _wgrad_arg_names( *, is_3d: bool = False, two_stage: bool = False, + streamk: Optional[str] = None, ) -> List[ArgSpec]: """Weight-gradient conv: ``dY`` (N[D]HWK), ``X`` (N[D]HWC), ``dW`` (K[Z]YXC). @@ -181,6 +200,13 @@ def _wgrad_arg_names( two cases one removes the pair of ABI flags that used to select between them -- and with them the chance of building a kernel against one variant and launching it with the other. + + ``streamk`` names the stream-K reduction of a stream-K kernel, ``None`` + otherwise. The ``workspace`` reduction reuses the two-stage scratch pair; + ``linear`` and ``tree`` add their flag and partial buffers after + ``ks_count``, and every reduction then appends the partition block + (``_STREAMK_SCALARS``). A stream-K kernel ignores ``ks``/``ks_count`` -- + it partitions K itself -- but keeps them so the prefix is unchanged. """ out: List[ArgSpec] = [ ("dY", "a"), @@ -210,12 +236,16 @@ def _wgrad_arg_names( out += _magic("n_", (["Y"] if is_3d else []) + ["X", "cpg"]) out += _grid() # ---- variant-specific extras ---- - if two_stage: + if two_stage or streamk == "workspace": out += [("ws_ptr", "f32*"), ("ws_bytes", "i32")] # The split-K slice width and degree, always. `ks` alone is not enough: # decoding block_id_z needs the number of slices, which is a different # quantity from how wide each one is. out += [("ks", "i32"), ("ks_count", "i32")] + if streamk in ("linear", "tree"): + out += [("sk_flags", "i32*"), ("sk_partials", "f32*")] + if streamk is not None: + out += [(name, "i32") for name in _STREAMK_SCALARS] return out @@ -271,6 +301,7 @@ def conv_arg_names( direction: str = "fwd", is_3d: bool = False, two_stage: bool = False, + streamk: Optional[str] = None, ) -> List[ArgSpec]: """The implicit-GEMM conv kernarg ABI for one direction. @@ -284,15 +315,31 @@ def conv_arg_names( ``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. + + ``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": raise ValueError( f"two_stage is a wgrad-only variant (got direction={direction!r})" ) + if streamk is not None: + if direction != "wgrad": + raise ValueError( + f"streamk is a wgrad-only variant (got direction={direction!r})" + ) + if streamk not in STREAMK_REDUCTIONS: + raise ValueError( + f"unknown stream-K reduction {streamk!r}; expected one of " + f"{STREAMK_REDUCTIONS}" + ) + 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) if direction == "wgrad": - return _wgrad_arg_names(is_3d=is_3d, two_stage=two_stage) + return _wgrad_arg_names(is_3d=is_3d, two_stage=two_stage, streamk=streamk) if direction == "dgrad": return _dgrad_arg_names(is_3d=is_3d) raise ValueError( @@ -431,20 +478,24 @@ def conv_args_signature( dtype_d: Optional[str] = None, is_3d: bool = False, two_stage: bool = False, + streamk: Optional[str] = None, ) -> List[Dict[str, Any]]: """Implicit-GEMM conv launch signature for one direction. ``direction`` is ``"fwd"``, ``"wgrad"`` or ``"dgrad"``; ``dtype`` sets all three operand types unless ``dtype_b`` / ``dtype_d`` override them, which mixed-precision configs need. ``two_stage`` adds wgrad's f32 workspace - pair and is rejected for the other two directions. + pair and is rejected for the other two directions; ``streamk`` selects a + stream-K wgrad kernel's ABI by its reduction. One entry point rather than three: the direction is already a value every caller has in hand, so making it an argument removes the per-direction branch each of them used to carry. """ return _conv_signature( - conv_arg_names(direction=direction, is_3d=is_3d, two_stage=two_stage), + conv_arg_names( + direction=direction, is_3d=is_3d, two_stage=two_stage, streamk=streamk + ), dtype, dtype_b or dtype, dtype_d or dtype, 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 879aecf2bd93..3b89b31f031a 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 @@ -25,12 +25,19 @@ from dataclasses import dataclass from typing import Dict, List, Optional, Sequence, Tuple -from kernels.common.conv_abi import ArgSpec, conv_arg_names, conv_direct_arg_names +from kernels.common.conv_abi import ( + STREAMK_REDUCTIONS, + ArgSpec, + conv_arg_names, + conv_direct_arg_names, +) +from rocke.helpers.streamk import StreamKIterPartition, streamk_iter_partition from rocke.helpers.transforms import calculate_magic_numbers __all__ = [ "ConvArgs", "ConvGeometry", + "StreamKLaunch", ] @@ -330,7 +337,39 @@ def _finalize( # Host state the kernel reads but that cannot be derived from the problem: # the caller has to hand it over, and forgetting to is an error, not a zero. -_CALLER_SUPPLIED = ("ws_ptr", "ws_bytes", "sub_gemm_buf", "num_sub_gemms") +_CALLER_SUPPLIED = ( + "ws_ptr", + "ws_bytes", + "sub_gemm_buf", + "num_sub_gemms", + "sk_flags", + "sk_partials", +) + + +@dataclass(frozen=True) +class StreamKLaunch: + """How a stream-K wgrad kernel is launched. + + ``reduction`` is baked into the kernel (it picks the ABI and the fixup); + ``pool`` and ``persistent`` are not -- the pool is the CTA budget the host + partitions the MAC iterations over, and ``persistent`` must match the + kernel's launch mode, which decides the grid. Like the split-K degree, the + pool is a launch parameter, so one binary serves any pool. + """ + + pool: int + persistent: bool = False + reduction: str = "linear" + + def __post_init__(self) -> None: + if self.pool <= 0: + raise ValueError(f"stream-K pool must be > 0 (got {self.pool})") + if self.reduction not in STREAMK_REDUCTIONS: + raise ValueError( + f"unknown stream-K reduction {self.reduction!r}; expected one " + f"of {STREAMK_REDUCTIONS}" + ) # Which directions each algorithm has a kernel for. @@ -441,14 +480,22 @@ def from_problem( # ---- the ABI ---- - def arg_names(self, *, two_stage: bool = False) -> List[ArgSpec]: - """The ordered kernarg ABI of the kernel these args launch.""" + def arg_names( + self, *, two_stage: bool = False, streamk: Optional[str] = None + ) -> List[ArgSpec]: + """The ordered kernarg ABI of the kernel these args launch. + + ``streamk`` is the stream-K reduction of a stream-K wgrad kernel. + """ if self.algorithm == "direct": - if two_stage: - raise ValueError("direct conv has no two-stage variant") + if two_stage or streamk is not None: + raise ValueError("direct conv has no two-stage or stream-K variant") return conv_direct_arg_names(direction=self.direction) return conv_arg_names( - direction=self.direction, is_3d=self.geom.is_3d, two_stage=two_stage + direction=self.direction, + is_3d=self.geom.is_3d, + two_stage=two_stage, + streamk=streamk, ) # ---- GEMM view (implicit GEMM only) ---- @@ -497,6 +544,52 @@ def grid(self, split_k: int = 1) -> Tuple[int, int, int]: """ return (self.num_pid_n, self.num_pid_m, self.geom.groups * split_k) + # ---- stream-K (implicit-GEMM wgrad only) ---- + + def streamk_partition( + self, pool: int, *, persistent: bool = False + ) -> StreamKIterPartition: + """The stream-K iteration partition of this launch over ``pool`` CTAs. + + Output tiles are ``num_pid_m * groups`` by ``num_pid_n``: the conv + groups fold into GEMM-M per group, so a partial last M tile of one + group never spills into the next. A K iteration is one ``tile_k``. + """ + self._require_gemm() + if self.direction != "wgrad": + raise ValueError( + f"stream-K is implemented for wgrad only (got {self.direction!r})" + ) + return streamk_iter_partition( + m_tiles=self.num_pid_m * self.geom.groups, + n_tiles=self.num_pid_n, + iters_per_tile=_ceil_div(self.gemm_k, self.tile_k), + max_active_wgs=pool, + persistent=persistent, + ) + + def streamk_grid(self, sk: StreamKLaunch) -> Tuple[int, int, int]: + """One-dimensional grid of a stream-K launch.""" + return ( + self.streamk_partition(sk.pool, persistent=sk.persistent).grid_size, + 1, + 1, + ) + + def _streamk_values(self, sk: StreamKLaunch) -> Dict[str, int]: + part = self.streamk_partition(sk.pool, persistent=sk.persistent) + return { + "sk_iters_per_tile": part.iters_per_tile, + "sk_dp_tiles": part.dp_tiles, + "sk_ctas": part.sk_ctas, + # A divisor in the kernel's owner lookup, which an all-DP launch + # still evaluates; the partition clamps it to >= 1 for that case. + "sk_iters_per_cta": part.iters_per_sk_cta, + "sk_extra_iters": part.extra_iters, + "sk_pool": part.max_active_wgs, + "sk_tree_rounds": part.tree_rounds, + } + # ---- values ---- def _named_values(self) -> Dict[str, int]: @@ -607,6 +700,9 @@ def to_launch_values( ws_bytes: Optional[int] = None, sub_gemm_buf: Optional[int] = None, num_sub_gemms: Optional[int] = None, + streamk: Optional[StreamKLaunch] = None, + sk_flags: Optional[int] = None, + sk_partials: Optional[int] = None, ) -> Dict[str, int]: """Return the ``values`` dict for ``KernelLauncher.__call__()``. @@ -621,6 +717,10 @@ def to_launch_values( them selects the two-stage ABI. * ``sub_gemm_buf`` / ``num_sub_gemms`` -- dgrad's tilde record table, which the kernel always reads. + * ``streamk`` -- a stream-K wgrad launch; selects the stream-K ABI and + fills its partition block. The ``workspace`` reduction takes + ``ws_ptr``/``ws_bytes``; ``linear``/``tree`` take ``sk_flags`` and + ``sk_partials`` (see ``wgrad_streamk_workspace_layout``). Anything the ABI needs but was not passed, and anything passed that the ABI does not have, is an error here rather than a misread on the @@ -637,7 +737,15 @@ def to_launch_values( f"implicit-GEMM fwd needs a reduction extent (Z*Y*X*cpg) below " f"2**23 for its 24-bit address products; got {self.gemm_k}" ) - abi = self.arg_names(two_stage=ws_ptr is not None) + if streamk is not None and split_k != 1: + raise ValueError( + f"stream-K partitions K itself; launch it with split_k=1 " + f"(got {split_k})" + ) + abi = self.arg_names( + two_stage=ws_ptr is not None and streamk is None, + streamk=streamk.reduction if streamk is not None else None, + ) names = [n for n, _ in abi] values: Dict[str, int] = dict( zip(names[:6], (a_ptr, b_ptr, d_ptr, a_bytes, b_bytes, d_bytes)) @@ -647,10 +755,14 @@ def to_launch_values( "ws_bytes": (ws_bytes or 0) if ws_ptr is not None else ws_bytes, "sub_gemm_buf": sub_gemm_buf, "num_sub_gemms": num_sub_gemms, + "sk_flags": sk_flags, + "sk_partials": sk_partials, } 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 streamk is not None: + named.update(self._streamk_values(streamk)) for name in names[6:]: values[name] = ( supplied.pop(name) if name in supplied else self._resolve(name, named) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index 006244df3b39..7db2396d5000 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -90,6 +90,7 @@ BF16, F16, F32, + I32, IRBuilder, KernelDef, PtrType, @@ -567,6 +568,37 @@ class WgradConvSpec: # is a knob, not a constant: another arch or a different filter shape can # want a different contention-versus-footprint tradeoff. ws_replicas: int = _DEFAULT_WS_REPLICAS + # Stream-K execution (CK Tile StreamKTilePartitioner). "off" keeps the + # classic grid-per-tile kernel. "dp_sk" launches one CTA per data-parallel + # tile plus ``streamk_ctas`` stream-K CTAs, each owning a balanced, + # contiguous range of MAC iterations that may straddle output tiles. + # "persistent" launches exactly ``streamk_ctas`` CTAs that round-robin the + # data-parallel tiles and then run their stream-K range. + # + # Default "off": strictly additive, so every existing config emits + # byte-identical IR and keeps its kernel name. + streamk: str = "off" + # How the CTAs that share an output tile combine their partial K-sums. + # "linear" -- the tile owner waits on each later contributor's flag + # and folds its partial in order, then stores dW once. + # Deterministic; dW needs no pre-zeroing. + # "tree" -- pairwise fan-in with doubling stride; the owner stores. + # Deterministic; dW needs no pre-zeroing. + # "atomic" -- every contributor f32-atomic-adds into dW (dtype_d fp32 + # only). The caller zeroes dW before every launch. + # "workspace" -- every contributor f32-atomic-adds into the two-stage + # scratch; the Stage 2 reduce kernel casts it into dW. + # linear/tree hand partials over a flag table; their workspace is + # ``wgrad_streamk_workspace_nbytes`` and only its flag region has to be + # zeroed before each launch. + streamk_reduction: str = "linear" + # The CTA pool the stream-K remainder is spread over (CK max_active_wgs). + # -1 = auto, resolved by the builder via wgrad_streamk_default_ctas: sized + # to the occupancy target for the reductions that never wait, one per CU + # for linear/tree. Any explicit value > 0 is safe: those two wait only on + # later CTAs, which ascending dispatch always starts (see + # wgrad_streamk_default_ctas). + streamk_ctas: int = -1 @property def block_size(self) -> int: @@ -701,6 +733,14 @@ def kernel_name(self) -> str: # body under the same name otherwise. Only tagged when set, so # every other kernel keeps its name. "unroll": self.unroll_k, + # Stream-K changes the grid decode, the K ranges and the + # epilogue, and the launch mode and the reduction both change + # the emitted body. The CTA pool does not: like the split-K + # degree it is a launch parameter, so it is not in the name. + "sk": self.streamk != "off", + "skpers": self.streamk == "persistent", + f"skr{self.streamk_reduction}": self.streamk != "off" + and self.streamk_reduction != "linear", }, ) @@ -736,6 +776,9 @@ def validate(self) -> None: f"ws_replicas must be >= 1 (got {self.ws_replicas}); it is the " f"number of scratch slabs a group's K-slices spread over" ) + _sk_ok, _sk_why = wgrad_streamk_available(self) + if not _sk_ok: + raise ValueError(_sk_why) # Delegate rather than re-derive -- is_valid_wgrad_spec calls the same # predicate, so the two cannot drift apart. _gm_ok, _gm_why = wgrad_group_merge_available(self) @@ -1068,6 +1111,329 @@ def dw_addr(b_: IRBuilder, m_val: Value, n_val: Value): return dw_addr +# Stream-K is MFMA-only: the flag/fence fixup and the per-CTA partial slots +# are validated on these targets. +_STREAMK_ARCHES = ("gfx942", "gfx950") +_STREAMK_MODES = ("off", "dp_sk", "persistent") +_STREAMK_REDUCTIONS = ("linear", "tree", "atomic", "workspace") + + +def wgrad_streamk_available(spec: "WgradConvSpec") -> Tuple[bool, str]: + """Arch-independent stream-K gates, shared by ``validate()`` and + :func:`is_valid_wgrad_spec` so the two cannot drift apart. + + Stream-K owns the K range and the tile mapping, so every knob that also + partitions K, remaps tiles, or pipelines the K loop is rejected alongside + it. + """ + if spec.streamk not in _STREAMK_MODES: + return False, ( + f"streamk must be 'off', 'dp_sk' or 'persistent' (got {spec.streamk!r})" + ) + if spec.streamk == "off": + return True, "ok" + if spec.streamk_reduction not in _STREAMK_REDUCTIONS: + return False, ( + "streamk_reduction must be 'linear', 'tree', 'atomic' or 'workspace' " + f"(got {spec.streamk_reduction!r})" + ) + if spec.streamk_ctas == 0 or spec.streamk_ctas < -1: + return False, f"streamk_ctas must be -1 (auto) or > 0 (got {spec.streamk_ctas})" + if spec.wave_size != 64: + return ( + False, + f"stream-K wgrad is MFMA wave64 only (got wave_size={spec.wave_size})", + ) + if spec.split_k != 1: + return False, ( + f"streamk and split_k both partition K_wg; use split_k=1 with streamk " + f"(got split_k={spec.split_k})" + ) + if spec.two_stage: + return False, ( + "two_stage is the split-K scratch path; use " + "streamk_reduction='workspace' for the stream-K equivalent" + ) + if spec.async_dma or spec.unroll_k: + return False, ( + "streamk runs one plain K loop per work item; async_dma/unroll_k " + "pipeline the K loop across items" + ) + if spec.epilogue != "default": + return False, ( + f"streamk requires epilogue='default' (got {spec.epilogue!r}): the " + "epilogue runs once per tile inside the stream-K tile loop" + ) + if spec.chiplet_swizzle: + return False, ( + "chiplet_swizzle remaps the 2-D tile grid; stream-K decodes its " + "tiles from a linear CTA index instead" + ) + if spec.group_merge > 1: + return False, ( + f"streamk folds conv groups into GEMM-M; use group_merge=1 " + f"(got {spec.group_merge})" + ) + if spec.streamk_reduction == "atomic" and spec.data.dtype_d != "fp32": + return False, ( + f"streamk_reduction='atomic' needs dtype_d='fp32' (got " + f"{spec.data.dtype_d!r}); use 'linear', 'tree' or 'workspace'" + ) + if not spec.acc_epilogue.is_identity(): + return False, ( + "streamk supports only the identity acc_epilogue: partial K-sums " + "are combined after the per-tile accumulation" + ) + return True, "ok" + + +# CTAs per CU in the default pool of the stream-K reductions that never wait. +# Each resident CTA streams its own K range, so this is the number of +# independent reduction streams a CU keeps in flight to hide memory latency. +# Counted in CTAs rather than waves: a wider CTA does not add streams. One per +# CU leaves a single stream per CU; past four the CTAs mostly contend for the +# same caches. +_STREAMK_CTAS_PER_CU = 4 + + +def wgrad_streamk_default_ctas(spec: "WgradConvSpec", arch: str) -> int: + """Default stream-K CTA pool for ``spec`` on ``arch`` (``streamk_ctas=-1``). + + The reductions that never wait on another CTA (``workspace``, ``atomic``) + are sized for throughput: ``_STREAMK_CTAS_PER_CU`` CTAs per CU. Nothing + there requires the pool to be resident at once, so a CTA that does not fit + only runs later. + + ``linear`` and ``tree`` get one CTA per CU: those fixups get slower, not + faster, with more contributors per tile (the owner folds every partial). + Their pool need not be resident at once. A CTA waits only on higher-index + CTAs and publishes its own partial before it waits, so with workgroups + dispatched in ascending order only the CTAs of the tile straddling the + dispatch front can wait on one that has not started, and they hold a + bounded number of slots (one for linear, about one per round for tree). + """ + return wgrad_streamk_default_pool(spec.streamk_reduction, arch) + + +def wgrad_streamk_default_pool(reduction: str, arch: str) -> int: + """:func:`wgrad_streamk_default_ctas` for a bare ``reduction`` name. + + Lets selection policy size the pool a launch will use before any spec + exists. + """ + from rocke.helpers.split_k import _ARCH_NUM_CUS, _DEFAULT_NUM_CUS + + num_cus = _ARCH_NUM_CUS.get(arch, _DEFAULT_NUM_CUS) + if reduction in ("linear", "tree"): + return num_cus + return num_cus * _STREAMK_CTAS_PER_CU + + +def wgrad_streamk_launch(spec: "WgradConvSpec", *, arch: str = "gfx950"): + """The :class:`~kernels.common.conv_args.StreamKLaunch` of ``spec``. + + The pool is ``spec.streamk_ctas``, or the arch default for ``-1``. It is + a launch parameter: the kernel reads the partition it implies from its + kernargs, so the same binary runs with any pool. + """ + from kernels.common.conv_args import StreamKLaunch + + if spec.streamk == "off": + raise ValueError("wgrad_streamk_launch requires streamk != 'off'") + # Same contract as the validator and the C port (rocke_wgrad_conv_streamk_ctas): + # -1 is auto, anything else not > 0 is an error rather than the auto pool. + if spec.streamk_ctas == 0 or spec.streamk_ctas < -1: + raise ValueError( + f"streamk_ctas must be -1 (auto) or > 0 (got {spec.streamk_ctas})" + ) + pool = ( + spec.streamk_ctas + if spec.streamk_ctas > 0 + else wgrad_streamk_default_ctas(spec, arch) + ) + return StreamKLaunch( + pool=pool, + persistent=spec.streamk == "persistent", + reduction=spec.streamk_reduction, + ) + + +def _wgrad_streamk_launch_problem(spec: WgradConvSpec, problem=None) -> ConvProblem: + """The launch shape, checked against what the kernel baked in. + + One stream-K binary serves any shape that keeps its build-time modes: + grouping, 2-D/3-D (the kernarg list), pointwise addressing, and channel + runs at least as wide as the built load/store vectors. Anything else would + launch the wrong code, so it is rejected rather than silently mis-run. + """ + built = spec.problem + if problem is None: + return built + + def _vec(chan: int, dtype: str, cap: int | None) -> int: + widths = (8, 4, 2, 1) if dtype != "fp32" else (4, 2, 1) + auto = next(v for v in widths if chan % v == 0) + return auto if cap is None else min(auto, cap) + + checks = ( + ("grouping (groups > 1)", built.groups > 1, problem.groups > 1), + ("is_3d", built.is_3d, problem.is_3d), + ("is_pointwise", built.is_pointwise, problem.is_pointwise), + ) + for what, want, got in checks: + if want != got: + raise ValueError( + f"launch problem is incompatible with the built kernel: {what} " + f"differs (built {want}, launch {got}); build a kernel for this shape" + ) + # Both X (operand B) and dW (operand D) run along cpg, with their own + # vector widths, so each check names the operand it guards. + chans = ( + ("dY", "kpg", built.kpg, problem.kpg, spec.data.dtype_a, spec.vector_size_a), + ("X", "cpg", built.cpg, problem.cpg, spec.data.dtype_b, spec.vector_size_b), + ("dW", "cpg", built.cpg, problem.cpg, spec.data.dtype_d, spec.vector_size_c), + ) + for operand, what, want, got, dtype, cap in chans: + width = _vec(want, dtype, cap) + if got % width: + raise ValueError( + f"launch problem is incompatible with the built kernel: {what}=" + f"{got} is not a multiple of the built {width}-wide {operand} " + f"vector ({what}={want}); build a kernel for this shape" + ) + return problem + + +def _wgrad_streamk_args(spec: "WgradConvSpec", problem=None): + from kernels.common.conv_args import ConvArgs + + return ConvArgs.from_problem( + _wgrad_streamk_launch_problem(spec, problem), + direction="wgrad", + tile_m=spec.tile_m, + tile_n=spec.tile_n, + tile_k=spec.tile_k, + ) + + +def wgrad_streamk_partition( + spec: "WgradConvSpec", *, arch: str = "gfx950", problem=None +): + """The :class:`~rocke.helpers.streamk.StreamKIterPartition` of a launch. + + ``problem`` is the launch shape and defaults to ``spec.problem``. Output + tiles are ``ceil(wg_M / tile_m) * groups`` by ``ceil(wg_N / tile_n)``; the + conv groups fold into GEMM-M per group, so a partial last M tile of one + group never spills into the next. + """ + sk = wgrad_streamk_launch(spec, arch=arch) + return _wgrad_streamk_args(spec, problem).streamk_partition( + sk.pool, persistent=sk.persistent + ) + + +def wgrad_streamk_grid( + spec: "WgradConvSpec", *, arch: str = "gfx950", problem=None +) -> Tuple[int, int, int]: + """Launch grid of a stream-K wgrad kernel (one-dimensional).""" + return (wgrad_streamk_partition(spec, arch=arch, problem=problem).grid_size, 1, 1) + + +def wgrad_streamk_workspace_layout( + spec: "WgradConvSpec", *, arch: str = "gfx950", problem=None +) -> Tuple[int, int]: + """``(flags_bytes, partials_bytes)`` of a linear/tree stream-K launch. + + One i32 flag per (stream-K CTA, wave), padded to 256 bytes, then one f32 + tile of partial sums per stream-K CTA. Pass the base as ``sk_flags`` and + ``base + flags_bytes`` as ``sk_partials``. Only the flag region has to be + zeroed before each launch; partials are always written before being read. + """ + if spec.streamk_reduction not in ("linear", "tree"): + return 0, 0 + part = wgrad_streamk_partition(spec, arch=arch, problem=problem) + flags_bytes = part.flags_bytes(spec.warp_m * spec.warp_n) + partials_bytes = part.sk_ctas * spec.tile_m * spec.tile_n * 4 + return flags_bytes, partials_bytes + + +def wgrad_streamk_workspace_nbytes( + spec: "WgradConvSpec", *, arch: str = "gfx950", problem=None +) -> int: + """Device workspace a stream-K wgrad launch needs, in bytes. + + linear/tree: the flag table plus the partial slots. workspace: the + two-stage f32 scratch (``groups * ws_replicas * wg_M * wg_N * 4``) that the + Stage 2 reduce kernel consumes. atomic: none. + """ + if spec.streamk == "off" or spec.streamk_reduction == "atomic": + return 0 + if spec.streamk_reduction == "workspace": + p = _wgrad_streamk_launch_problem(spec, problem) + return p.groups * spec.ws_replicas * _wg_M(p) * _wg_N(p) * 4 + flags_bytes, partials_bytes = wgrad_streamk_workspace_layout( + spec, arch=arch, problem=problem + ) + return flags_bytes + partials_bytes + + +def wgrad_streamk_signature(spec: "WgradConvSpec") -> list: + """Launch signature of a stream-K wgrad kernel. + + The wgrad conv ABI with the stream-K extras of ``spec``'s reduction: the + ``ws_ptr``/``ws_bytes`` scratch for ``workspace``, ``sk_flags``/ + ``sk_partials`` for ``linear``/``tree``, then the partition block. + Binding is positional, as for every AOT conv kernel. + """ + from kernels.common.conv_abi import conv_args_signature + + return conv_args_signature( + spec.data.dtype_a, + direction="wgrad", + dtype_b=spec.data.dtype_b, + dtype_d=spec.data.dtype_d, + is_3d=spec.problem.is_3d, + streamk=spec.streamk_reduction, + ) + + +def wgrad_streamk_launch_values( + spec: "WgradConvSpec", + dY_ptr: int, + X_ptr: int, + dW_ptr: int, + dY_bytes: int, + X_bytes: int, + dW_bytes: int, + *, + arch: str = "gfx950", + problem=None, + ws_ptr: Optional[int] = None, + ws_bytes: Optional[int] = None, + sk_flags: Optional[int] = None, + sk_partials: Optional[int] = None, +) -> Dict[str, int]: + """The kernarg ``values`` of a stream-K wgrad launch. + + ``problem`` is the launch shape (default ``spec.problem``); the partition + block is computed from it and from the pool of + :func:`wgrad_streamk_launch`. + """ + return _wgrad_streamk_args(spec, problem).to_launch_values( + dY_ptr, + X_ptr, + dW_ptr, + dY_bytes, + X_bytes, + dW_bytes, + streamk=wgrad_streamk_launch(spec, arch=arch), + ws_ptr=ws_ptr, + ws_bytes=ws_bytes, + sk_flags=sk_flags, + sk_partials=sk_partials, + ) + + def wgrad_group_merge_available( spec: WgradConvSpec, arch: str = "gfx950" ) -> Tuple[bool, str]: @@ -1183,6 +1549,14 @@ def is_valid_wgrad_spec(spec: WgradConvSpec, arch: str = "gfx950") -> Tuple[bool return False, ( f"spec wave_size {spec.wave_size} != {arch} wave_size {target.wave_size}" ) + if spec.streamk != "off" and arch not in _STREAMK_ARCHES: + return False, ( + f"stream-K wgrad supports only {', '.join(_STREAMK_ARCHES)} " + f"(got {arch})" + ) + _sk_ok, _sk_why = wgrad_streamk_available(spec) + if not _sk_ok: + return False, _sk_why sk = spec.split_k if sk < -1 or sk == 0: @@ -1535,6 +1909,9 @@ def build_implicit_gemm_conv_wgrad( # Two-stage: the scratch slab index is group * R + z % R and the group # decode reads ks_count, so nothing in the body depends on the degree. _is_two_stage = _is_split_k and spec.two_stage + _is_streamk = spec.streamk != "off" + # Stream-K reductions that accumulate atomically instead of storing once. + _streamk_atomic = _is_streamk and spec.streamk_reduction in ("atomic", "workspace") # At group_merge == groups the merged problem has a single group, but the # kernel still has to decode the K-slice off z and the epilogue still has to # mask off-diagonal pairs -- so the grouped path stays engaged. Without this @@ -1560,7 +1937,10 @@ def build_implicit_gemm_conv_wgrad( # split_k > 1 two-stage: Stage 1 never touches dW (Stage 2 writes it), # but dW stays in the ABI so the atomic and # two-stage signatures line up. - _dw_writeonly = not _is_split_k + # stream-K linear/tree: the tile owner stores dW once (writeonly); the + # atomic and workspace reductions read-modify-write + # their target. + _dw_writeonly = not _is_split_k and not _streamk_atomic def _declare_wgrad_ptr(name: str, kind: str) -> Value: if kind == "a": @@ -1587,14 +1967,22 @@ def _declare_wgrad_ptr(name: str, kind: str) -> Value: writeonly=_dw_writeonly, align=16, ) - # Two-stage scratch: f32 accumulator every K-slice atomic-adds into, - # size groups*ws_replicas*wg_M*wg_N*4. Not ``writeonly``: an - # atomicrmw reads its target. + if kind == "i32*": + # Stream-K linear/tree ready flags, read back by the tile owner. + return b.param(name, PtrType(I32, "global"), noalias=True, align=16) + # Two-stage / stream-K workspace scratch: f32 accumulator every K-slice + # atomic-adds into, size groups*ws_replicas*wg_M*wg_N*4; or the stream-K + # linear/tree partial slots. Not ``writeonly``: both are read back. return b.param(name, PtrType(F32, "global"), noalias=True, align=16) params = emit_param_block( b, - conv_arg_names(direction="wgrad", is_3d=p.is_3d, two_stage=_is_two_stage), + conv_arg_names( + direction="wgrad", + is_3d=p.is_3d, + two_stage=_is_two_stage, + streamk=spec.streamk_reduction if _is_streamk else None, + ), declare_ptr=_declare_wgrad_ptr, ) dY = params["dY"] @@ -1622,6 +2010,8 @@ def _declare_wgrad_ptr(name: str, kind: str) -> Value: p_num_pid_n = params["p_num_pid_n"] # Variant-specific extras; absent from the parameter list when unused. ws_ptr = params.get("ws_ptr") + sk_flags = params.get("sk_flags") + sk_partials = params.get("sk_partials") # Slice width and slice count, always kernargs. ks_count is what decodes # block_id_z = group*ks_count + slice; ks is how wide each slice is. They # are different numbers, so one cannot stand in for the other. @@ -1715,17 +2105,21 @@ def _declare_wgrad_ptr(name: str, kind: str) -> Value: # bounds, so the padding contributes nothing -- which is what lets the split # and unsplit cases share one bound computation instead of two branches # that have to be kept in step. + # + # Stream-K launches a one-dimensional grid and folds the group into its + # linear tile index, so it skips this decode and binds both per work item. c_ks = _ks_param - z_id = b.block_id_z() - if grouped: - group_v = b.to_sgpr_u32(b.div(z_id, _ks_count_param)) - slice_id = b.mod(z_id, _ks_count_param) - if spec.group_merge > 1: - slice_v = slice_id - else: - slice_id = b.mod(z_id, _ks_count_param) - k_lo = b.to_sgpr_u32(b.mul(slice_id, c_ks)) - k_hi = b.to_sgpr_u32(b.add(k_lo, c_ks)) + if not _is_streamk: + z_id = b.block_id_z() + if grouped: + group_v = b.to_sgpr_u32(b.div(z_id, _ks_count_param)) + slice_id = b.mod(z_id, _ks_count_param) + if spec.group_merge > 1: + slice_v = slice_id + else: + slice_id = b.mod(z_id, _ks_count_param) + k_lo = b.to_sgpr_u32(b.mul(slice_id, c_ks)) + k_hi = b.to_sgpr_u32(b.add(k_lo, c_ks)) # Chiplet swizzle (same logic as forward; tile counts from wgrad GEMM dims). # The tile counts are runtime kernargs — the host computes them from the @@ -2277,6 +2671,287 @@ def emit_mfma_phase( return new_accs + def _emit_streamk_tiles() -> None: + """Stream-K body: a loop over this CTA's work items. + + A work item is one output tile plus the slice of its MAC iterations this + CTA owns. A non-persistent CTA walks the tiles its iteration range + touches (exactly one whole tile for a data-parallel CTA); a persistent + CTA first round-robins the data-parallel tiles and then walks its + stream-K range. Every item runs the ordinary K loop over its slice and + then the configured fixup, so the GEMM body is emitted once. + + The kernel is shape- and pool-agnostic: the partition arrives as the + ``sk_*`` kernargs (see ``ConvArgs.streamk_partition``), so every loop + bound below is a CTA-uniform runtime value. + + The tile mapping and group fold rebind the offsets the descriptor + closures read, exactly as the split-K path binds its K slice. + """ + nonlocal block_m_off_v, block_n_off_v, group_v, grid + from rocke.helpers.streamk import ( + StreamKIterArgs, + emit_streamk_iter_range, + emit_streamk_sk_start_iter, + ) + + reduction = spec.streamk_reduction + persistent = spec.streamk == "persistent" + c_one = b.const_i32(1) + c_ipt = params["sk_iters_per_tile"] + c_dp_tiles = params["sk_dp_tiles"] + c_sk_ctas = params["sk_ctas"] + c_per_cta = params["sk_iters_per_cta"] + c_extra = params["sk_extra_iters"] + sk_args = StreamKIterArgs( + iters_per_tile=c_ipt, + dp_tiles=c_dp_tiles, + total_dp_iters=b.to_sgpr_u32(b.mul(c_dp_tiles, c_ipt)), + iters_per_sk_cta=c_per_cta, + extra_iters=c_extra, + ) + cta = b.to_sgpr_u32(b.block_id_x()) + if persistent: + c_pool = params["sk_pool"] + # DP tiles cta, cta + pool, ... : ceil((dp_tiles - cta) / pool). + # cta < pool, so the numerator is never negative. + n_dp = b.div(b.add(b.sub(c_dp_tiles, cta), b.sub(c_pool, c_one)), c_pool) + rng = emit_streamk_iter_range(b, cta, sk_args, dp=False) + first_tile = b.div(rng.start, c_ipt) + last_tile = b.div(b.sub(rng.end, c_one), c_ipt) + n_sk = b.add(b.sub(last_tile, first_tile), c_one) + # An all-DP partition has no stream-K CTAs; persistent CTAs are + # either all stream-K CTAs or none of them. + n_sk = b.select(b.cmp_lt(cta, c_sk_ctas), n_sk, c0) + n_items = b.to_sgpr_u32(b.add(n_dp, n_sk)) + sk_cta = cta + else: + rng = emit_streamk_iter_range(b, cta, sk_args, dp=True) + first_tile = b.div(rng.start, c_ipt) + last_tile = b.div(b.sub(rng.end, c_one), c_ipt) + n_items = b.to_sgpr_u32(b.add(b.sub(last_tile, first_tile), c_one)) + sk_cta = rng.sk_cta + + n_waves = spec.warp_m * spec.warp_n + n_slots = len(accs) * c_per_lane + + def _flag_idx(sk_idx: Value) -> Value: + # One flag per (stream-K CTA, wave): a consumer lane only ever + # reads the partial its own wave's lane wrote, so each wave can + # release and acquire independently with no workgroup barrier. + return b.add(b.mul(sk_idx, b.const_i32(n_waves)), warp_id) + + def _slot_base(sk_idx: Value) -> Value: + # Partials are [sk_cta][slot][tid] f32: lanes of one slot are + # adjacent, so every partial store and load is fully coalesced. + return b.add(b.mul(sk_idx, b.const_i32(n_slots * threads)), tid) + + def _store_partial(sk_idx: Value, vals: Sequence[Value]) -> None: + base = _slot_base(sk_idx) + for ai, acc in enumerate(vals): + for i in range(c_per_lane): + off = b.add(base, b.const_i32((ai * c_per_lane + i) * threads)) + b.global_store(sk_partials, off, b.vec_extract(acc, i), align=4) + b.fence(scope="agent", ordering="release") + b.global_flag_store(sk_flags, _flag_idx(sk_idx), c_one) + + def _load_partial(sk_idx: Value) -> List[Value]: + base = _slot_base(sk_idx) + out = [] + for ai in range(len(accs)): + comps = [ + b.global_load_f32( + sk_partials, + b.add(base, b.const_i32((ai * c_per_lane + i) * threads)), + ) + for i in range(c_per_lane) + ] + out.append(b.vec_pack(comps, F32)) + return out + + def _sk_cta_of(it: Value) -> Value: + # The stream-K CTA owning global iteration ``it``: CTAs below + # extra_iters own iters_per_sk_cta + 1 iterations, the rest one fewer. + rel = b.sub(it, sk_args.total_dp_iters) + wide = b.add(c_per_cta, c_one) + boundary = b.mul(c_extra, wide) + lead = b.div(rel, wide) + tail = b.add(c_extra, b.div(b.sub(rel, boundary), c_per_cta)) + return b.select(b.cmp_lt(rel, boundary), lead, tail) + + def _store_tile(vals: Sequence[Value]) -> None: + _emit_wgrad_direct_epilogue( + b, spec, vals, grid, dw_rsrc, params, group=group_v + ) + + loop = b.scf_for(c0, n_items, c_one, iv_name="sk_item") + with loop as item: + if persistent: + is_dp = b.cmp_lt(item, n_dp) + dp_tile = b.add(cta, b.mul(item, c_pool)) + sk_tile = b.add(first_tile, b.sub(item, n_dp)) + tile = b.to_sgpr_u32(b.select(is_dp, dp_tile, sk_tile)) + dp_start = b.mul(tile, c_ipt) + item_start = b.select(is_dp, dp_start, rng.start) + item_end = b.select(is_dp, b.add(dp_start, c_ipt), rng.end) + else: + tile = b.to_sgpr_u32(b.add(first_tile, item)) + item_start = rng.start + item_end = rng.end + tile_start = b.mul(tile, c_ipt) + tile_end = b.add(tile_start, c_ipt) + lo = b.sub(b.smax(item_start, tile_start), tile_start) + hi = b.sub(b.smin(item_end, tile_end), tile_start) + + # Tile -> (group, m_tile, n_tile). Groups are folded into GEMM-M + # group-fastest, CK's i_g = tile_m % GemmBatch; m_tiles is defined + # per group, so the fold is exact for any wg_M. + n_tile = b.mod(tile, p_num_pid_n) + m_idx = b.div(tile, p_num_pid_n) + if grouped: + group_v = b.to_sgpr_u32(b.mod(m_idx, p_groups)) + m_tile = b.div(m_idx, p_groups) + else: + m_tile = m_idx + block_m_off_v = b.to_sgpr_u32(b.mul(m_tile, b.const_i32(block_m))) + block_n_off_v = b.to_sgpr_u32(b.mul(n_tile, b.const_i32(block_n))) + grid = dc_replace( + grid, block_m_off=block_m_off_v, block_n_off=block_n_off_v + ) + + k_lo_t = b.to_sgpr_u32(b.mul(lo, c_block_k)) + k_hi_t = b.to_sgpr_u32(b.mul(hi, c_block_k)) + for_op = b.scf_for_iter(k_lo_t, k_hi_t, c_block_k, accs, iv_name="k0") + with for_op as (k0, iter_vars): + emit_load_phase(k0, A_smem, B_smem) + b.sync() + new_accs = emit_mfma_phase(A_smem, B_smem, iter_vars) + b.sync() + b.scf_yield(*new_accs) + tile_accs = for_op.results + + if reduction in ("linear", "tree"): + # Whether this item covers the start / the end of its tile. A + # data-parallel tile (and every tile of an all-DP partition) + # has both. Built only for the flag-handoff fixups, so the + # atomic and workspace bodies carry no dead compares. + started = b.cmp_eq(lo, c0) + ended = b.cmp_eq(hi, c_ipt) + + if reduction == "atomic": + _emit_wgrad_split_k_epilogue( + b, + spec, + atom, + tile_accs, + warp_m_idx, + warp_n_idx, + lane, + block_m_off_v, + block_n_off_v, + dW, + c_per_lane, + params, + group=group_v, + ) + elif reduction == "workspace": + _emit_wgrad_workspace_store_epilogue( + b, + spec, + atom, + tile_accs, + warp_m_idx, + warp_n_idx, + lane, + block_m_off_v, + block_n_off_v, + ws_ptr, + c_per_lane, + params, + gm_group=group_v, + replica_sel=cta, + ) + elif reduction == "linear": + # A data-parallel tile has started and ended, so it folds no + # partners and stores. + with b.scf_if_else(started) as (owner, contributor): + with owner: + # Fold every later contributor in order, then store once. + last = _sk_cta_of(b.sub(tile_end, c_one)) + n_partners = b.select(ended, c0, b.sub(last, sk_cta)) + fold = b.scf_for_iter( + c_one, + b.to_sgpr_u32(b.add(n_partners, c_one)), + c_one, + [(f"sk_{nm}", v) for (nm, _), v in zip(accs, tile_accs)], + iv_name="sk_p", + ) + with fold as (pi, folded): + partner = b.add(sk_cta, pi) + b.global_flag_wait_eq(sk_flags, _flag_idx(partner), c_one) + parts = _load_partial(partner) + b.scf_yield( + *[b.vector_add(v, q) for v, q in zip(folded, parts)] + ) + _store_tile(fold.results) + with contributor: + _store_partial(sk_cta, tile_accs) + else: + # Tree: round r pairs CTAs 2^r apart within the tile. A receiver + # folds its partner in; the first round in which a contributor + # is not a receiver it publishes its sum and drops out. The + # round count is a kernarg (0 for an all-DP partition), so the + # rounds are a runtime loop carrying the partial sums, the pair + # stride and a dropped-out flag; every predicate is uniform. + not_started = b.lnot(started) + local = b.sub(sk_cta, _sk_cta_of(tile_start)) + zero_acc = b.zero_vec_f32(c_per_lane) + c_two = b.const_i32(2) + n_acc = len(accs) + rounds = b.scf_for_iter( + c0, + params["sk_tree_rounds"], + c_one, + [(f"skt_{nm}", v) for (nm, _), v in zip(accs, tile_accs)] + + [("skt_stride", c_one), ("skt_done", c0)], + iv_name="sk_r", + ) + with rounds as (_r, carried): + vals = list(carried[:n_acc]) + stride = carried[n_acc] + live = b.cmp_eq(carried[n_acc + 1], c0) + partner = b.add(sk_cta, stride) + in_tile = b.land( + b.lnot(ended), + b.cmp_lt( + emit_streamk_sk_start_iter(b, partner, sk_args), tile_end + ), + ) + is_recv = b.cmp_eq(b.mod(local, b.mul(stride, c_two)), c0) + recv = b.land(b.land(is_recv, in_tile), live) + with b.scf_if(recv): + b.global_flag_wait_eq(sk_flags, _flag_idx(partner), c_one) + # Non-receivers read slot 0 (always in bounds) and discard it. + parts = _load_partial(b.select(recv, partner, c0)) + vals = [ + b.vector_add(v, b.select(recv, q, zero_acc)) + for v, q in zip(vals, parts) + ] + send = b.land(b.land(not_started, b.lnot(is_recv)), live) + with b.scf_if(send): + _store_partial(sk_cta, vals) + b.scf_yield( + *vals, + b.mul(stride, c_two), + b.select(send, c_one, carried[n_acc + 1]), + ) + with b.scf_if(started): + _store_tile(list(rounds.results[:n_acc])) + + if _is_streamk: + _emit_streamk_tiles() + return b.kernel + # ---- K loop ---- # k_lo / k_hi select the slice this CTA processes: k_lo = slice*ks, # k_hi = k_lo + ks (SGPR-pinned, scalar arith); an unsplit launch passes @@ -3063,6 +3738,7 @@ def _emit_wgrad_workspace_store_epilogue( c_per_lane: int, params: Dict[str, Value], gm_group: Optional[Value] = None, + replica_sel: Optional[Value] = None, ) -> None: """Two-stage Stage 1 epilogue: f32 atomic-add into the scratch accumulator. @@ -3116,17 +3792,24 @@ def _emit_wgrad_workspace_store_epilogue( # downstream SSA value and break byte-identity on the default path. # Scratch total size = groups * R * wg_M * wg_N (f32 elements); R is a # build-time knob, the slab size is the runtime wg_M * wg_N. + # + # ``replica_sel`` replaces blockIdx.z as the hash input: a stream-K launch + # is one-dimensional, so it passes its linear CTA index instead. reps = spec.ws_replicas _slab_idx = None + + def _replica_src() -> Value: + return replica_sel if replica_sel is not None else b.block_id_z() + if gm_group is not None and reps > 1: _slab_idx = b.add( b.mul(gm_group, b.const_i32(reps)), - b.mod(b.block_id_z(), b.const_i32(reps)), + b.mod(_replica_src(), b.const_i32(reps)), ) elif gm_group is not None: _slab_idx = gm_group elif reps > 1: - _slab_idx = b.mod(b.block_id_z(), b.const_i32(reps)) + _slab_idx = b.mod(_replica_src(), b.const_i32(reps)) slab_off = b.mul(_slab_idx, slab_v) if _slab_idx is not None else None # Merged-path setup. Only materialised when gm > 1, for the same @@ -3203,7 +3886,7 @@ def _emit_wgrad_workspace_store_epilogue( if reps > 1: group_true = b.add( b.mul(group_true, c_reps_v), - b.mod(b.block_id_z(), c_reps_v), + b.mod(_replica_src(), c_reps_v), ) n_true = b.mul(yx, c_cpg_v) if p.cpg > 1 else yx if p.cpg > 1: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c index 0e8279535bad..ce94560db390 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c @@ -31,6 +31,18 @@ * 17 N8H56W56C64_K64Y3X3 pad1, t32x32x32, w1x1, a16x16x32, mem/default, gfx1250 (WMMA w32), lds_k_outer fp32 out * 18 N8H56W56C64_K64Y3X3, t64x64x64, w2x2, a32x32x16, mem/default, gfx950, unroll_k, split_k=4 fp32 * 19 N8H56W56C3_K64Y3X3, t64x64x64, w2x2, a32x32x16, mem/default, gfx950, split_k=2 two_stage fp16 + * 20 N10H8W8C64_K64Y3X3, t64x64x64, w2x2, a32x32x16, mem/cshuffle, gfx950, lds_k_outer + async_dma, split_k=2 + * 21 N8H56W56C64_K64Y3X3, a32x32x16, gfx950, streamk dp_sk linear fp16 + * 22 N8H56W56C64_K64Y3X3, a32x32x16, gfx950, streamk dp_sk tree fp16 + * 23 N8H56W56C64_K64Y3X3, a32x32x16, gfx950, streamk dp_sk atomic fp32 + * 24 N8H56W56C64_K64Y3X3, a32x32x16, gfx950, streamk dp_sk workspace fp16 + * 25 N8H56W56C64_K64Y3X3, a16x16x16, gfx942, streamk persistent linear fp16 + * 26 N4H28W28C96_K120Y3X3G3 pad1, a32x32x16, gfx950, streamk dp_sk linear fp16 + * 27 N4H28W28C64_K128Y3X3G2 pad1, a32x32x16, gfx950, streamk persistent tree bf16 dW + * 28 N4H28W28C64_K128Y3X3G2 pad1, a32x32x16, gfx950, streamk persistent atomic fp32 + * 29 N8H56W56C64_K64Y3X3, a16x16x16, gfx942, streamk persistent workspace fp16 + * 30 N4H14W14C64_K96Y3X3G2 pad1, a32x32x16, gfx950, streamk dp_sk linear fp16, lds_k_outer + * (21-30 carry a CTA pool, but it is a launch parameter and leaves the IR as is.) */ #include #include @@ -42,6 +54,41 @@ #include "rocke/lower_llvm.h" #include "rocke/verify.h" +/* Stream-K fields of configs 21-29 (Python _STREAMK_CFGS). */ +static int streamk_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, const char** arch) +{ + /* idx: dtype a/b, dtype d, atom edge, streamk, reduction, ctas, arch */ + static const struct + { + const char *dt_ab, *dt_d; + int edge; + const char *mode, *reduction; + int ctas; + const char* arch; + } cfgs[9] = { + {"fp16", "fp16", 32, "dp_sk", "linear", 16, "gfx950"}, + {"fp16", "fp16", 32, "dp_sk", "tree", 16, "gfx950"}, + {"fp16", "fp32", 32, "dp_sk", "atomic", 4, "gfx950"}, + {"fp16", "fp16", 32, "dp_sk", "workspace", 4, "gfx950"}, + {"fp16", "fp16", 16, "persistent", "linear", 4, "gfx942"}, + {"fp16", "fp16", 32, "dp_sk", "linear", 6, "gfx950"}, + {"fp16", "bf16", 32, "persistent", "tree", 5, "gfx950"}, + {"fp16", "fp32", 32, "persistent", "atomic", 3, "gfx950"}, + {"fp16", "fp16", 16, "persistent", "workspace", 6, "gfx942"}, + }; + const int i = idx - 21; + spec->dtype_a = cfgs[i].dt_ab; + spec->dtype_b = cfgs[i].dt_ab; + spec->dtype_d = cfgs[i].dt_d; + spec->warp_tile_m = cfgs[i].edge; + spec->warp_tile_n = cfgs[i].edge; + spec->streamk = cfgs[i].mode; + spec->streamk_reduction = cfgs[i].reduction; + spec->streamk_ctas = cfgs[i].ctas; + *arch = cfgs[i].arch; + return 0; +} + /* Fill the config for index `idx`. Returns 0 on success, -1 if unknown. * On success sets *spec and *arch. */ static int make_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, const char** arch) @@ -242,6 +289,36 @@ static int make_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, const spec->split_k = 2; *arch = "gfx950"; return 0; + case 21: + case 22: + case 23: + case 24: + case 25: + case 29: + spec->problem = rocke_conv_problem_default(8, 56, 56, 64, 64, 3, 3); + return streamk_cfg(idx, spec, arch); + case 26: + spec->problem = rocke_conv_problem_default(4, 28, 28, 96, 120, 3, 3); + spec->problem.pH = 1; + spec->problem.pW = 1; + spec->problem.groups = 3; + return streamk_cfg(idx, spec, arch); + case 27: + case 28: + spec->problem = rocke_conv_problem_default(4, 28, 28, 64, 128, 3, 3); + spec->problem.pH = 1; + spec->problem.pW = 1; + spec->problem.groups = 2; + return streamk_cfg(idx, spec, arch); + case 30: + spec->problem = rocke_conv_problem_default(4, 14, 14, 64, 96, 3, 3); + spec->problem.pH = 1; + spec->problem.pW = 1; + spec->problem.groups = 2; + streamk_cfg(26, spec, arch); + spec->streamk_ctas = 7; + spec->lds_k_outer = true; + return 0; default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py index fca631e37854..bd725f526ec6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py @@ -33,15 +33,32 @@ # 20 -- K-outer + async_dma + split-K=2 with the fp16 cshuffle atomic epilogue, # 3 K tiles per slice (odd: the phase-B prefetch is redirected to the # zero-fill offset), gfx950 -# (no grouped or bf16 case: the C++ engine does not build grouped wgrad, and -# its tile loader has no elem_dtype, so bf16 operands load as half) +# 21 -- stream-K dp_sk, linear fixup, fp16, gfx950 +# 22 -- stream-K dp_sk, tree fixup, fp16, gfx950 +# 23 -- stream-K dp_sk, atomic reduction, fp32, gfx950 +# 24 -- stream-K dp_sk, workspace reduction, fp16, gfx950 +# 25 -- stream-K persistent, linear fixup, 16x16x16 atom, gfx942 +# 26 -- stream-K dp_sk, linear, groups=3 with a partial last M tile, gfx950 +# 27 -- stream-K persistent, tree, bf16 dW, groups=2, gfx950 +# 28 -- stream-K persistent, atomic, fp32, groups=2, gfx950 +# 29 -- stream-K persistent, workspace, fp16, gfx942 +# 30 -- stream-K dp_sk, linear, K-outer LDS tile, groups=2, gfx950 +# (Grouped wgrad reaches the C++ engine only through stream-K, and its tile +# loader has no elem_dtype, so no case loads bf16 operands. The stream-K +# partition and CTA pool are kernargs, so no case varies them. async_dma is +# rejected with stream-K.) # # Negative cases (configs 100+) verify that invalid specs are rejected: # 100 -- odd C with fp16 split-K (must raise ValueError) # 102 -- split_k > 1 on RDNA gfx1151 (must raise ValueError) # 103 -- two_stage=True with split_k=1 (must raise ValueError) +# 104 -- streamk with split_k=4 (must raise ValueError) +# 105 -- streamk on gfx1250 (must raise ValueError) +# 106 -- streamk with the cshuffle epilogue (must raise ValueError) +# 107 -- streamk with unroll_k (must raise ValueError) +# 108 -- streamk atomic reduction into fp16 dW (must raise ValueError) # (These illustrate the validator contract. The C emitter defines only cases -# 0-19, so run_diff.py stops at the shared END before reaching 100+; these +# 0-30, so run_diff.py stops at the shared END before reaching 100+; these # configs are not exercised by the differential gate.) from kernels.common.conv_implicit_gemm_wgrad import ( WgradConvSpec, @@ -538,6 +555,24 @@ def _spec(idx: int): "gfx950", ) + if 21 <= idx <= 29: + return _streamk_spec(idx) + if idx == 30: + from dataclasses import replace + + spec, arch = _streamk_spec(26) + return ( + replace( + spec, + problem=ConvProblem( + N=4, Hi=14, Wi=14, C=64, K=96, Y=3, X=3, pH=1, pW=1, groups=2 + ), + streamk_ctas=7, + lds_k_outer=True, + ), + arch, + ) + # ---------------------------------------------------------------- # Negative cases: these specs must be REJECTED by the validator. # The harness (run_emit) expects a ValueError / SystemExit when @@ -605,9 +640,136 @@ def _spec(idx: int): ), "gfx950", ) + if 104 <= idx <= 108: + from dataclasses import replace + + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + + base, arch = _streamk_spec(21) + bad = { + 104: (dict(split_k=4), arch), + 105: (dict(), "gfx1250"), + 106: (dict(epilogue="cshuffle"), arch), + 107: (dict(unroll_k=True), arch), + 108: (dict(streamk_reduction="atomic", data=ConvDataSpec()), arch), + }[idx] + return replace(base, **bad[0]), bad[1] raise SystemExit(f"unknown config index {idx}") +# (problem kwargs, data dtypes (a/b, d), atom edge, streamk, reduction, ctas, arch) +_STREAMK_CFGS = { + 21: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 32, + "dp_sk", + "linear", + 16, + "gfx950", + ), + 22: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 32, + "dp_sk", + "tree", + 16, + "gfx950", + ), + 23: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp32"), + 32, + "dp_sk", + "atomic", + 4, + "gfx950", + ), + 24: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 32, + "dp_sk", + "workspace", + 4, + "gfx950", + ), + 25: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 16, + "persistent", + "linear", + 4, + "gfx942", + ), + 26: ( + dict(N=4, Hi=28, Wi=28, C=96, K=120, Y=3, X=3, pH=1, pW=1, groups=3), + ("fp16", "fp16"), + 32, + "dp_sk", + "linear", + 6, + "gfx950", + ), + 27: ( + dict(N=4, Hi=28, Wi=28, C=64, K=128, Y=3, X=3, pH=1, pW=1, groups=2), + ("fp16", "bf16"), + 32, + "persistent", + "tree", + 5, + "gfx950", + ), + 28: ( + dict(N=4, Hi=28, Wi=28, C=64, K=128, Y=3, X=3, pH=1, pW=1, groups=2), + ("fp16", "fp32"), + 32, + "persistent", + "atomic", + 3, + "gfx950", + ), + 29: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 16, + "persistent", + "workspace", + 6, + "gfx942", + ), +} + + +def _streamk_spec(idx: int): + """Stream-K configs; mirrored field-for-field by the C emitter.""" + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + + prob, (dt_ab, dt_d), edge, mode, reduction, ctas, arch = _STREAMK_CFGS[idx] + return ( + WgradConvSpec( + problem=ConvProblem(**prob), + data=ConvDataSpec(dtype_a=dt_ab, dtype_b=dt_ab, dtype_d=dt_d), + tile_m=64, + tile_n=64, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=edge, + warp_tile_n=edge, + warp_tile_k=16, + pipeline="mem", + epilogue="default", + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=ctas, + ), + arch, + ) + + def main() -> int: return run_emit( _spec, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py new file mode 100644 index 000000000000..5c7127ab7ab0 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py @@ -0,0 +1,713 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Host-side tests for stream-K implicit-GEMM backward-weight convolution. + +No GPU needed. Covers the spec surface (defaults, kernel-name tags), the +validator gates and their agreement between ``validate()``, +``is_valid_wgrad_spec`` and the C++ port, admits-implies-builds over the +mode x reduction x arch x groups matrix, the AOT kernarg ABI and launch values +(the partition block), the host-side launch plan (grid, workspace layout, +determinism) against the C port, and byte-identical lowering between the two +engines through the ``rocke_engine`` binding. + +GPU numerics live in ``test_conv_wgrad_streamk_gpu.py``. +""" + +from __future__ import annotations + +import dataclasses +import unittest +from dataclasses import replace + +from kernels.common._conv_implicit_gemm_common import ConvDataSpec, ConvProblem +from kernels.common.conv_abi import _STREAMK_SCALARS, conv_arg_names +from kernels.common.conv_args import ConvArgs, StreamKLaunch +from kernels.common.conv_implicit_gemm_wgrad import ( + WgradConvSpec, + build_implicit_gemm_conv_wgrad, + is_valid_wgrad_spec, + wgrad_streamk_grid, + wgrad_streamk_partition, + wgrad_streamk_signature, + wgrad_streamk_workspace_layout, + wgrad_streamk_workspace_nbytes, +) + +_DENSE = ConvProblem(N=2, Hi=12, Wi=12, C=32, K=96, Y=3, X=3, pH=1, pW=1) +# kpg = 40 is not a multiple of tile_m, so the last M tile of every group is +# partial: the case CK's own GemmM * GemmBatch fold leaves untested. +_GROUPED = ConvProblem(N=2, Hi=10, Wi=10, C=48, K=120, Y=3, X=3, pH=1, pW=1, groups=3) +# Depthwise: the only shape group merging admits, so its stream-K gate is reached. +_DW = ConvProblem(N=2, Hi=8, Wi=8, C=16, K=16, Y=3, X=3, pH=1, pW=1, groups=16) +# 17 x 27 = 459 output tiles at 64x64: more than the default pool of either +# stream-K arch (256 CTAs on gfx950, 304 on gfx942 for linear/tree), so an +# auto-pool launch really splits tiles and needs flags and partial slots. +_SPLITS = ConvProblem(N=2, Hi=12, Wi=12, C=192, K=1056, Y=3, X=3, pH=1, pW=1) + + +def _spec(problem=_DENSE, *, dtype_d="fp16", edge=32, **kw) -> WgradConvSpec: + base = dict( + problem=problem, + data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d=dtype_d), + tile_m=64, + tile_n=64, + tile_k=32, + warp_m=2, + warp_n=2, + warp_tile_m=edge, + warp_tile_n=edge, + warp_tile_k=16, + streamk="dp_sk", + streamk_ctas=6, + ) + base.update(kw) + return WgradConvSpec(**base) + + +def _engine(): + """The ``rocke_engine`` binding, or skip the calling test.""" + from rocke.core.backend import BackendError, _import_engine + + try: + return _import_engine() + except BackendError as e: + raise unittest.SkipTest(f"C++ engine not importable: {e}") from None + + +def _spec_dict(spec: WgradConvSpec) -> dict: + """Flatten a WgradConvSpec into the dict ``conv_wgrad_build_spec`` reads.""" + d = { + f.name: getattr(spec, f.name) + for f in dataclasses.fields(spec) + if f.name not in ("problem", "data", "acc_epilogue", "lds_layout") + } + d = {k: v for k, v in d.items() if v is not None} + d["problem"] = { + k: v for k, v in dataclasses.asdict(spec.problem).items() if v is not None + } + d["dtype_a"] = spec.data.dtype_a + d["dtype_b"] = spec.data.dtype_b + d["dtype_d"] = spec.data.dtype_d + return d + + +class TestStreamKSpecSurface(unittest.TestCase): + def test_default_is_off_and_name_unchanged(self): + plain = replace(_spec(), streamk="off", streamk_ctas=-1) + self.assertNotIn("_sk", plain.kernel_name()) + self.assertEqual( + plain.kernel_name(), + WgradConvSpec( + **{ + f.name: getattr(plain, f.name) + for f in dataclasses.fields(plain) + if not f.name.startswith("streamk") + } + ).kernel_name(), + ) + + def test_name_tags_every_knob_that_changes_the_body(self): + names = { + _spec().kernel_name(), + _spec(streamk="persistent").kernel_name(), + _spec(streamk_reduction="tree").kernel_name(), + _spec(streamk_reduction="workspace").kernel_name(), + _spec(dtype_d="fp32", streamk_reduction="atomic").kernel_name(), + } + self.assertEqual(len(names), 5, names) + self.assertTrue(_spec().kernel_name().endswith("_sk")) + + def test_pool_is_a_launch_parameter(self): + # The pool only sizes the host partition: the name and the kernel body + # are the same for every pool, like the split-K degree. + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + def ll(spec): + k = build_implicit_gemm_conv_wgrad(spec, arch="gfx950") + return _lower_kernel_to_llvm_python(k, arch="gfx950") + + base = ll(_spec()) + for ctas in (8, -1): + spec = _spec(streamk_ctas=ctas) + self.assertEqual(spec.kernel_name(), _spec().kernel_name()) + self.assertEqual(ll(spec), base) + + def test_auto_pool_depends_on_the_reduction(self): + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_default_ctas + + # linear/tree wait on later CTAs: one per CU, always resident. + for red in ("linear", "tree"): + spec = _spec(streamk_ctas=-1, streamk_reduction=red) + self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx950"), 256) + self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx942"), 304) + # workspace/atomic never wait: four CTAs per CU. + for red, dd in (("workspace", "fp16"), ("atomic", "fp32")): + spec = _spec(streamk_ctas=-1, streamk_reduction=red, dtype_d=dd) + self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx950"), 1024) + self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx942"), 1216) + # Counted in CTAs, not waves: an 8-wave CTA gets the same pool. + wide = _spec(streamk_ctas=-1, streamk_reduction="workspace", warp_m=4) + self.assertEqual(wgrad_streamk_default_ctas(wide, "gfx950"), 1024) + + def test_problem_is_not_baked_in(self): + # AOT: the problem shape is a kernarg, so two problems with the same + # build-time knobs (here: grouping and 2-D) lower to the same body. + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + other = ConvProblem(N=8, Hi=20, Wi=6, C=48, K=24, Y=3, X=3, pH=1, pW=1) + lls = [] + for problem in (_DENSE, other): + k = build_implicit_gemm_conv_wgrad(_spec(problem), arch="gfx950") + body = _lower_kernel_to_llvm_python(k, arch="gfx950") + lls.append(body.replace(k.name, "KERNEL")) + self.assertEqual(lls[0], lls[1]) + + +class TestStreamKValidator(unittest.TestCase): + """Every gate rejects in ``validate()`` and ``is_valid_wgrad_spec`` alike.""" + + _REJECTS = [ + (dict(streamk="sideways"), "streamk must be"), + (dict(streamk_reduction="ring"), "streamk_reduction must be"), + (dict(streamk_ctas=0), "streamk_ctas must be"), + (dict(streamk_ctas=-3), "streamk_ctas must be"), + (dict(split_k=4), "both partition K_wg"), + (dict(split_k=-1), "both partition K_wg"), + (dict(two_stage=True, split_k=2), "both partition K_wg"), + (dict(unroll_k=True), "one plain K loop"), + (dict(epilogue="cshuffle"), "epilogue='default'"), + (dict(chiplet_swizzle=True), "chiplet_swizzle"), + (dict(streamk_reduction="atomic"), "needs dtype_d='fp32'"), + ] + + def test_rejects(self): + for kw, needle in self._REJECTS: + spec = _spec(**kw) + with self.subTest(kw=kw): + ok, why = is_valid_wgrad_spec(spec, "gfx950") + self.assertFalse(ok) + self.assertIn(needle, why) + with self.assertRaises(ValueError) as cm: + spec.validate() + self.assertEqual(str(cm.exception), why) + + def test_group_merge_and_acc_epilogue_rejected(self): + from kernels.common._conv_implicit_gemm_common import ConvAccumulatorEpilogue + + for spec, needle in ( + (_spec(_DW, group_merge=2), "group_merge=1"), + (_spec(acc_epilogue=ConvAccumulatorEpilogue(relu=True)), "acc_epilogue"), + ): + with self.subTest(needle=needle): + ok, why = is_valid_wgrad_spec(spec, "gfx950") + self.assertFalse(ok) + self.assertIn(needle, why) + + def test_arch_gate(self): + ok, why = is_valid_wgrad_spec(_spec(), "gfx90a") + self.assertFalse(ok) + self.assertIn("stream-K wgrad supports only gfx942, gfx950", why) + # Wave32 targets already fail the wave-size check; either way, rejected. + for arch in ("gfx1151", "gfx1201", "gfx1250"): + with self.subTest(arch=arch): + self.assertFalse(is_valid_wgrad_spec(_spec(), arch)[0]) + + def test_cpp_rejects_with_the_same_reason(self): + eng = _engine() + for kw, _needle in self._REJECTS: + spec = _spec(**kw) + _ok, why = is_valid_wgrad_spec(spec, "gfx950") + with self.subTest(kw=kw), self.assertRaises(RuntimeError) as cm: + eng.conv_wgrad_lower_llvm(_spec_dict(spec), "gfx950") + self.assertIn(why, str(cm.exception)) + + # (label, problem, arch, spec overrides). Each rejected config trips one + # gate only; _MULTI below covers configs that trip several. + _AGREEMENT = [ + ("ok_dp_sk_linear", _DENSE, "gfx950", {}), + ( + "ok_pers_tree", + _GROUPED, + "gfx950", + dict(streamk="persistent", streamk_reduction="tree"), + ), + ( + "ok_atomic_fp32", + _DENSE, + "gfx950", + dict(dtype_d="fp32", streamk_reduction="atomic"), + ), + ("ok_workspace", _GROUPED, "gfx950", dict(streamk_reduction="workspace")), + ("ok_gfx942", _DENSE, "gfx942", dict(edge=16, streamk="persistent")), + ("ok_auto_ctas", _DENSE, "gfx950", dict(streamk_ctas=-1)), + ("ok_off", _DENSE, "gfx950", dict(streamk="off")), + ("split_k", _DENSE, "gfx950", dict(split_k=4)), + ("split_k_auto", _DENSE, "gfx950", dict(split_k=-1)), + ("two_stage", _DENSE, "gfx950", dict(two_stage=True, split_k=2)), + ("cshuffle", _DENSE, "gfx950", dict(epilogue="cshuffle")), + ("unroll_k", _DENSE, "gfx950", dict(unroll_k=True)), + ("chiplet", _DENSE, "gfx950", dict(chiplet_swizzle=True)), + ("atomic_fp16", _DENSE, "gfx950", dict(streamk_reduction="atomic")), + ( + "atomic_bf16", + _DENSE, + "gfx950", + dict(dtype_d="bf16", streamk_reduction="atomic"), + ), + ("ctas_zero", _DENSE, "gfx950", dict(streamk_ctas=0)), + ("ctas_minus_two", _DENSE, "gfx950", dict(streamk_ctas=-2)), + ("bad_mode", _DENSE, "gfx950", dict(streamk="sideways")), + ("bad_reduction", _DENSE, "gfx950", dict(streamk_reduction="sideways")), + ("gfx90a", _DENSE, "gfx90a", {}), + # wave32 so Python's wave-size check passes and the arch gate decides. + ("gfx1250", _DENSE, "gfx1250", dict(wave_size=32, warp_tile_k=32)), + ("group_merge", _DW, "gfx950", dict(group_merge=2)), + ] + + def test_cpp_validity_agrees(self): + eng = _engine() + for label, problem, arch, kw in self._AGREEMENT: + spec = _spec(problem, **kw) + want = is_valid_wgrad_spec(spec, arch) + with self.subTest(label): + self.assertEqual(label.startswith("ok_"), want[0], want[1]) + got = eng.conv_wgrad_is_valid_spec(_spec_dict(spec), arch) + self.assertEqual(tuple(got), want) + + # Configs that fail more than one gate: the C port runs the gates in + # Python's order, so both report the first one Python reaches. + _MULTI = [ + ("ctas_zero+tile_m", _DENSE, "gfx950", dict(streamk_ctas=0, tile_m=96)), + ( + "block_4096+ctas_zero", + _DENSE, + "gfx950", + dict(warp_m=8, warp_n=8, tile_m=256, tile_n=256, streamk_ctas=0), + ), + ("unknown_arch", _DENSE, "gfx9999", {}), + ("unknown_arch+tile_m", _DENSE, "gfx9999", dict(tile_m=96)), + ( + "wave32_on_gfx950+ctas", + _DENSE, + "gfx950", + dict(wave_size=32, streamk_ctas=-5), + ), + ("wave64_on_gfx1250", _DENSE, "gfx1250", {}), + ("async_dma", _DENSE, "gfx950", dict(async_dma=True)), + ("async_dma_off", _DENSE, "gfx950", dict(async_dma=True, streamk="off")), + ("vec_c+ctas", _DENSE, "gfx950", dict(vector_size_c=4, streamk_ctas=-5)), + ("vec_c_off", _DENSE, "gfx950", dict(vector_size_c=4, streamk="off")), + ("group_merge_degree", _DW, "gfx950", dict(group_merge=3)), + ("group_merge+tile_m", _DW, "gfx950", dict(group_merge=2, tile_m=96)), + ("gfx90a+split_k", _DENSE, "gfx90a", dict(split_k=4)), + ] + + def test_cpp_validity_agrees_multi_failure(self): + eng = _engine() + for label, problem, arch, kw in self._MULTI: + spec = _spec(problem, **kw) + want = is_valid_wgrad_spec(spec, arch) + with self.subTest(label): + self.assertFalse(want[0]) + got = eng.conv_wgrad_is_valid_spec(_spec_dict(spec), arch) + self.assertEqual(tuple(got), want) + + def test_cpp_rejects_grouped_non_streamk(self): + # Intentional divergence: grouped wgrad is only ported on the stream-K + # path, so the C port rejects a grouped spec Python accepts -- but only + # after every gate the two share. + eng = _engine() + spec = _spec(_GROUPED, streamk="off") + self.assertEqual(is_valid_wgrad_spec(spec, "gfx950"), (True, "ok")) + ok, why = eng.conv_wgrad_is_valid_spec(_spec_dict(spec), "gfx950") + self.assertFalse(ok) + self.assertIn("grouped convolution (groups=3 > 1)", why) + bad = _spec(_GROUPED, streamk="off", tile_m=96) + self.assertEqual( + tuple(eng.conv_wgrad_is_valid_spec(_spec_dict(bad), "gfx950")), + is_valid_wgrad_spec(bad, "gfx950"), + ) + + def test_cpp_rejects_group_merge_with_the_same_reason(self): + eng = _engine() + spec = _spec(_DW, group_merge=2) + _ok, why = is_valid_wgrad_spec(spec, "gfx950") + with self.assertRaises(RuntimeError) as cm: + eng.conv_wgrad_lower_llvm(_spec_dict(spec), "gfx950") + self.assertIn(why, str(cm.exception)) + + +class TestStreamKAdmitsImpliesBuilds(unittest.TestCase): + def test_matrix(self): + built = 0 + for arch, edge in (("gfx950", 32), ("gfx942", 16)): + for problem in (_DENSE, _GROUPED): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in ( + ("linear", "fp16"), + ("tree", "bf16"), + ("atomic", "fp32"), + ("workspace", "fp16"), + ): + for ctas in (1, 4, 7, 64): + spec = _spec( + problem, + dtype_d=dtype_d, + edge=edge, + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=ctas, + ) + ok, why = is_valid_wgrad_spec(spec, arch) + with self.subTest( + arch=arch, + g=problem.groups, + mode=mode, + r=reduction, + w=ctas, + ): + self.assertTrue(ok, why) + build_implicit_gemm_conv_wgrad(spec, arch=arch) + built += 1 + self.assertEqual(built, 2 * 2 * 2 * 4 * 4) + + +class TestStreamKHostPlan(unittest.TestCase): + def test_grid_matches_partition(self): + for mode in ("dp_sk", "persistent"): + spec = _spec(streamk=mode, streamk_ctas=4) + part = wgrad_streamk_partition(spec) + with self.subTest(mode=mode): + want = 4 if mode == "persistent" else part.dp_tiles + part.sk_ctas + self.assertEqual(wgrad_streamk_grid(spec), (want, 1, 1)) + + def test_group_fold_tile_count(self): + # m_tiles is per group: ceil(kpg / tile_m) * groups, never + # ceil(kpg * groups / tile_m). + part = wgrad_streamk_partition(_spec(_GROUPED)) + self.assertEqual(part.m_tiles, 3) # ceil(40 / 64) * 3 + self.assertEqual(part.n_tiles, 3) # ceil(3*3*16 / 64) + + def test_workspace_layout(self): + spec = _spec(streamk_ctas=7) + self.assertEqual(wgrad_streamk_partition(spec).sk_ctas, 7) + flags, partials = wgrad_streamk_workspace_layout(spec) + self.assertEqual(flags, 256) # 7 CTAs * 4 waves * 4 B, 256-aligned + self.assertEqual(partials, 7 * 64 * 64 * 4) + self.assertEqual(wgrad_streamk_workspace_nbytes(spec), flags + partials) + # 10 tiles over a pool of 5 split evenly: all data-parallel, no slots. + self.assertEqual(wgrad_streamk_workspace_layout(_spec(streamk_ctas=5)), (0, 0)) + atomic = _spec(dtype_d="fp32", streamk_reduction="atomic") + self.assertEqual(wgrad_streamk_workspace_nbytes(atomic), 0) + ws = _spec(streamk_reduction="workspace") + self.assertEqual(wgrad_streamk_workspace_nbytes(ws), 8 * 96 * 288 * 4) + + def test_launch_rejects_bad_pool(self): + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_launch + + for ctas in (0, -2, -3): + with self.subTest(ctas=ctas), self.assertRaisesRegex( + ValueError, "streamk_ctas must be -1" + ): + wgrad_streamk_launch(_spec(streamk_ctas=ctas)) + self.assertEqual(wgrad_streamk_launch(_spec(streamk_ctas=-1)).pool, 256) + self.assertEqual( + wgrad_streamk_launch(_spec(streamk_ctas=-1), arch="gfx942").pool, 304 + ) + + def test_cpp_host_rejects_bad_pool(self): + eng = _engine() + for ctas in (0, -3): + with self.subTest(ctas=ctas), self.assertRaises(ValueError): + eng.conv_wgrad_streamk_host(_spec_dict(_spec(streamk_ctas=ctas))) + with self.assertRaises(ValueError): + eng.conv_wgrad_streamk_host(_spec_dict(_spec(streamk="off"))) + + def test_signature_tracks_reduction(self): + names = lambda s: [p["name"] for p in wgrad_streamk_signature(s)] # noqa: E731 + plain = [n for n, _ in conv_arg_names(direction="wgrad")] + block = list(_STREAMK_SCALARS) + self.assertEqual(names(_spec()), plain + ["sk_flags", "sk_partials"] + block) + ws = plain[:-2] + ["ws_ptr", "ws_bytes"] + plain[-2:] + block + self.assertEqual(names(_spec(streamk_reduction="workspace")), ws) + self.assertEqual( + names(_spec(dtype_d="fp32", streamk_reduction="atomic")), plain + block + ) + + def test_cpp_host_plan_agrees(self): + eng = _engine() + for problem in (_DENSE, _GROUPED): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in ( + ("linear", "fp16"), + ("tree", "fp16"), + ("atomic", "fp32"), + ("workspace", "fp16"), + ): + spec = _spec( + problem, + dtype_d=dtype_d, + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=7, + ) + got = eng.conv_wgrad_streamk_host(_spec_dict(spec)) + flags, partials = wgrad_streamk_workspace_layout(spec) + with self.subTest(g=problem.groups, mode=mode, r=reduction): + self.assertEqual(tuple(got["grid"]), wgrad_streamk_grid(spec)) + self.assertEqual(got["flags_bytes"], flags) + self.assertEqual(got["partials_bytes"], partials) + self.assertEqual( + got["workspace_bytes"], wgrad_streamk_workspace_nbytes(spec) + ) + self.assertEqual( + got["is_deterministic"], reduction in ("linear", "tree") + ) + + def test_cpp_default_pool_agrees(self): + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_default_pool + from rocke.helpers.split_k import _ARCH_NUM_CUS + + from rocke.core.arch.target import known_arches + + eng = _engine() + # Both directions: every arch of the Python CU table, and every arch + # rocke knows of plus a few it does not -- a key only the C table has + # would make the C pool differ from Python's 256-CU fallback there. + arches = {*_ARCH_NUM_CUS, *known_arches(), "gfx90a", "gfx908", "gfx1100"} + arches |= {"gfx1200", "gfx9999"} + for arch in sorted(arches): + for reduction in ("linear", "tree", "atomic", "workspace"): + with self.subTest(arch=arch, r=reduction): + self.assertEqual( + eng.conv_wgrad_streamk_default_pool(reduction, arch), + wgrad_streamk_default_pool(reduction, arch), + ) + + def test_cpp_host_plan_agrees_auto_pool(self): + # streamk_ctas=-1 (the spec default) resolves to the arch's CU-based + # pool in both engines. + eng = _engine() + split = 0 + for arch, edge in (("gfx950", 32), ("gfx942", 16)): + for problem in (_DENSE, _GROUPED, _SPLITS): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in ( + ("linear", "fp16"), + ("tree", "fp16"), + ("atomic", "fp32"), + ("workspace", "fp16"), + ): + spec = _spec( + problem, + dtype_d=dtype_d, + edge=edge, + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=-1, + ) + got = eng.conv_wgrad_streamk_host(_spec_dict(spec), arch) + flags, partials = wgrad_streamk_workspace_layout( + spec, arch=arch + ) + with self.subTest( + arch=arch, g=problem.groups, mode=mode, r=reduction + ): + self.assertEqual( + tuple(got["grid"]), wgrad_streamk_grid(spec, arch=arch) + ) + self.assertEqual(got["flags_bytes"], flags) + self.assertEqual(got["partials_bytes"], partials) + self.assertEqual( + got["workspace_bytes"], + wgrad_streamk_workspace_nbytes(spec, arch=arch), + ) + if problem is _SPLITS and reduction in ("linear", "tree"): + # 459 tiles over the default pool: every CTA is a + # stream-K CTA, so both regions are non-empty. + pool = 256 if arch == "gfx950" else 304 + self.assertEqual(got["grid"][0], pool) + self.assertEqual(flags, pool * 4 * 4) + self.assertEqual(partials, pool * 64 * 64 * 4) + split += 1 + self.assertEqual(split, 2 * 2 * 2) + # Without an arch both engines default to gfx950. + spec = _spec(streamk="persistent", streamk_ctas=-1) + got = eng.conv_wgrad_streamk_host(_spec_dict(spec)) + self.assertEqual(tuple(got["grid"]), wgrad_streamk_grid(spec)) + + +class TestStreamKLaunchAbi(unittest.TestCase): + """The stream-K kernarg ABI and the launch values that fill it.""" + + def _args(self, problem=_DENSE): + return ConvArgs.from_problem( + problem, direction="wgrad", tile_m=64, tile_n=64, tile_k=32 + ) + + def test_extras_follow_the_unchanged_prefix(self): + plain = conv_arg_names(direction="wgrad", is_3d=True) + for red in ("atomic", "workspace", "linear", "tree"): + sk = conv_arg_names(direction="wgrad", is_3d=True, streamk=red) + with self.subTest(reduction=red): + head = plain[:-2] + if red == "workspace": + head = head + [("ws_ptr", "f32*"), ("ws_bytes", "i32")] + self.assertEqual(sk[: len(head) + 2], head + plain[-2:]) + tail = [n for n, _ in sk[len(head) + 2 :]] + bufs = ["sk_flags", "sk_partials"] if red in ("linear", "tree") else [] + self.assertEqual(tail, bufs + list(_STREAMK_SCALARS)) + + def test_abi_rejects(self): + with self.assertRaises(ValueError): + conv_arg_names(direction="fwd", streamk="linear") + with self.assertRaises(ValueError): + conv_arg_names(direction="wgrad", streamk="sideways") + with self.assertRaises(ValueError): + conv_arg_names(direction="wgrad", two_stage=True, streamk="workspace") + + def test_values_are_the_partition(self): + for problem in (_DENSE, _GROUPED): + for persistent in (False, True): + for pool in (5, 6, 7, 1024): + sk = StreamKLaunch(pool=pool, persistent=persistent) + args = self._args(problem) + v = args.to_launch_values( + 1, 2, 3, 4, 5, 6, streamk=sk, sk_flags=7, sk_partials=8 + ) + part = args.streamk_partition(pool, persistent=persistent) + with self.subTest(g=problem.groups, pers=persistent, pool=pool): + self.assertEqual(v["sk_iters_per_tile"], part.iters_per_tile) + self.assertEqual(v["sk_dp_tiles"], part.dp_tiles) + self.assertEqual(v["sk_ctas"], part.sk_ctas) + self.assertEqual(v["sk_iters_per_cta"], part.iters_per_sk_cta) + self.assertEqual(v["sk_extra_iters"], part.extra_iters) + self.assertEqual(v["sk_pool"], pool) + self.assertEqual(v["sk_tree_rounds"], part.tree_rounds) + self.assertEqual((v["sk_flags"], v["sk_partials"]), (7, 8)) + self.assertEqual(args.streamk_grid(sk), (part.grid_size, 1, 1)) + + def test_launch_value_contract(self): + args = self._args() + with self.assertRaisesRegex(ValueError, "sk_flags must be passed"): + args.to_launch_values(1, 2, 3, 4, 5, 6, streamk=StreamKLaunch(pool=6)) + with self.assertRaisesRegex(ValueError, "ws_ptr must be passed"): + args.to_launch_values( + 1, 2, 3, 4, 5, 6, streamk=StreamKLaunch(pool=6, reduction="workspace") + ) + with self.assertRaisesRegex(ValueError, "split_k=1"): + args.to_launch_values( + 1, + 2, + 3, + 4, + 5, + 6, + split_k=2, + streamk=StreamKLaunch(pool=6, reduction="atomic"), + ) + with self.assertRaisesRegex(ValueError, "not arguments"): + args.to_launch_values( + 1, + 2, + 3, + 4, + 5, + 6, + sk_flags=7, + streamk=StreamKLaunch(pool=6, reduction="atomic"), + ) + with self.assertRaises(ValueError): + StreamKLaunch(pool=0) + + def test_launch_shape_overrides_the_spec_problem(self): + # One binary, many shapes: the helpers size the launch from `problem`. + other = ConvProblem(N=8, Hi=16, Wi=16, C=32, K=96, Y=3, X=3, pH=1, pW=1) + spec = _spec() + self.assertEqual( + wgrad_streamk_partition(spec, problem=other), + self._args(other).streamk_partition(6), + ) + self.assertNotEqual( + wgrad_streamk_partition(spec, problem=other), wgrad_streamk_partition(spec) + ) + + def test_launch_shape_must_keep_the_built_modes(self): + # The kernel bakes grouping, 2-D/3-D, pointwise addressing and its + # channel vector widths; a launch shape that changes any is rejected. + base = {"N": 2, "Hi": 12, "Wi": 12, "Y": 3, "X": 3, "pH": 1, "pW": 1} + spec = _spec() # _DENSE: groups=1, 2-D, 3x3, kpg=96, cpg=32 + for needle, other in ( + ("grouping", ConvProblem(C=32, K=96, groups=2, **base)), + ("is_3d", ConvProblem(C=32, K=96, Di=4, Z=1, sD=1, pD=0, dD=1, **base)), + ("is_pointwise", ConvProblem(N=2, Hi=12, Wi=12, C=32, K=96, Y=1, X=1)), + ( + "kpg=36 is not a multiple of the built 8-wide dY", + ConvProblem(C=32, K=36, **base), + ), + ( + "cpg=20 is not a multiple of the built 8-wide X", + ConvProblem(C=20, K=96, **base), + ), + ): + with self.subTest(needle=needle): + for helper in (wgrad_streamk_partition, wgrad_streamk_workspace_nbytes): + with self.assertRaises(ValueError) as cm: + helper(spec, problem=other) + self.assertIn(needle, str(cm.exception)) + # X and dW both run along cpg; the error names the operand that failed. + narrow_b = _spec(vector_size_b=4) + with self.assertRaises(ValueError) as cm: + wgrad_streamk_partition(narrow_b, problem=ConvProblem(C=20, K=96, **base)) + self.assertIn( + "cpg=20 is not a multiple of the built 8-wide dW", str(cm.exception) + ) + # Wider channel runs than the build keep the same code path. + wgrad_streamk_partition(spec, problem=ConvProblem(C=64, K=192, **base)) + + +class TestStreamKDualEngine(unittest.TestCase): + """The C++ builder lowers every stream-K variant to the Python engine's bytes.""" + + def test_lowered_ir_is_byte_identical(self): + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + eng = _engine() + cases = [ + ("gfx950", _DENSE, "dp_sk", "linear", "fp16", 32), + ("gfx950", _DENSE, "dp_sk", "tree", "fp16", 32), + ("gfx950", _GROUPED, "persistent", "linear", "bf16", 32), + ("gfx950", _GROUPED, "dp_sk", "atomic", "fp32", 32), + ("gfx942", _GROUPED, "persistent", "tree", "fp16", 16), + ("gfx942", _DENSE, "persistent", "workspace", "fp16", 16), + ] + for arch, problem, mode, reduction, dtype_d, edge in cases: + spec = _spec( + problem, + dtype_d=dtype_d, + edge=edge, + streamk=mode, + streamk_reduction=reduction, + ) + with self.subTest(arch=arch, g=problem.groups, mode=mode, r=reduction): + py_ll = _lower_kernel_to_llvm_python( + build_implicit_gemm_conv_wgrad(spec, arch=arch), arch=arch + ) + self.assertEqual( + eng.conv_wgrad_lower_llvm(_spec_dict(spec), arch), py_ll + ) + + def test_cpp_builds_auto_pool(self): + # The pool is not part of the kernel, so the C++ builder takes -1 too. + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + eng = _engine() + spec = _spec(streamk_ctas=-1) + py_ll = _lower_kernel_to_llvm_python( + build_implicit_gemm_conv_wgrad(spec, arch="gfx950"), arch="gfx950" + ) + self.assertEqual(eng.conv_wgrad_lower_llvm(_spec_dict(spec), "gfx950"), py_ll) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py new file mode 100644 index 000000000000..25e70741d8d7 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py @@ -0,0 +1,694 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""GPU numerics for stream-K implicit-GEMM backward-weight convolution. + +Every case builds, compiles and launches a stream-K wgrad kernel, then checks +dW against a float32 ``torch.nn.grad.conv2d_weight`` reference computed on the +CPU. Before the tolerance check a row-coverage guard reports whole output +channels that were never written (a dropped group or tile reports as such +rather than as an opaque numeric miss). + +The shapes are chosen to exercise the partitioner, not the GEMM: several +stream-K CTAs per tile (multi-round linear/tree fixups), a data-parallel plus +stream-K mix, a grouped problem whose last M tile is partial, and a small CTA +pool so the stream-K path is taken at these sizes. A shape large enough to +split under the default (per-CU) pool runs every reduction too. + +The linear and tree fixups are also run repeatedly and must be bit-identical +across launches. Before every launch the partial slots and dW are poisoned +with NaN bytes (only the flags are zeroed, as the contract requires), so a +consumer that reads a slot before its producer's partial is visible sees NaN +rather than the previous launch's bit-identical partial. That catches a +missing *release* on the producer side. It does not catch a missing +*acquire*: the consumer's partial loads are data- and control-dependent on +the spin exit and no stale cache line exists to observe, so a relaxed spin +still reads the right bytes on this hardware. The acquire side is pinned by +codegen instead -- ``TestConvWgradStreamKFlagOrderingIR`` checks the lowered +``.ll`` (host-only), ``TestConvWgradStreamKFlagOrderingIsa`` the gfx942/gfx950 +ISA (COMGR + ``llvm-objdump``, no GPU; skips when either is missing). + +The numeric cases require a gfx942/gfx950 GPU and torch and skip cleanly +otherwise. Host-side coverage lives in ``test_conv_wgrad_streamk.py``. +""" + +from __future__ import annotations + +import ctypes +import importlib.util +import os +import re +import shutil +import subprocess +import tempfile +import unittest +from pathlib import Path + +from rocke.runtime.hip_module import get_device_arch + +_HAS_TORCH = importlib.util.find_spec("torch") is not None +GPU_ARCH = get_device_arch(0) + + +def _skip_reason() -> str: + if not GPU_ARCH: + return "no ROCm GPU detected" + if not _HAS_TORCH: + return "torch not importable" + if GPU_ARCH not in ("gfx942", "gfx950"): + return f"stream-K wgrad is gfx942/gfx950 only; got {GPU_ARCH}" + return "" + + +_SKIP_REASON = _skip_reason() + +# Relative max-error bound by dW dtype: the reference is fp32 over fp16/bf16 +# inputs, so the bound is set by the output rounding. +_TOL = {"fp32": 1e-4, "fp16": 2e-3, "bf16": 2e-2} +_REDUCTIONS = ( + ("linear", "fp16"), + ("tree", "fp16"), + ("atomic", "fp32"), + ("workspace", "bf16"), +) + + +def _u8(t): + return (ctypes.c_uint8 * t.nbytes).from_address(t.data_ptr()) + + +def _reference(X, dY, p): + import torch + + w = torch.nn.grad.conv2d_weight( + X.permute(0, 3, 1, 2).contiguous(), + weight_size=(p.K, p.C // p.groups, p.Y, p.X), + grad_output=dY.permute(0, 3, 1, 2).contiguous(), + stride=(p.sH, p.sW), + padding=(p.pH, p.pW), + dilation=(p.dH, p.dW), + groups=p.groups, + ) + return w.permute(0, 2, 3, 1).contiguous() + + +def _run(spec, *, launches: int = 1, problem=None, artifact=None): + """Build, launch ``launches`` times, return (dW outputs, fp32 reference). + + ``problem`` is the launch shape (default ``spec.problem``): the kernel is + AOT, so one binary -- passed in as ``artifact`` to reuse it -- serves any + shape its spec admits. + """ + import torch + + from rocke import compile_kernel + from rocke.runtime.hip_module import Runtime + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + from kernels.common.conv_implicit_gemm_wgrad import ( + build_implicit_gemm_conv_wgrad, + wgrad_streamk_grid, + wgrad_streamk_launch_values, + wgrad_streamk_signature, + wgrad_streamk_workspace_layout, + wgrad_streamk_workspace_nbytes, + ) + from kernels.common.conv_wgrad_workspace_reduce import ( + WgradReduceSpec, + build_conv_wgrad_workspace_reduce, + wgrad_reduce_grid, + wgrad_reduce_signature, + ) + + arch = GPU_ARCH + p = problem if problem is not None else spec.problem + art = artifact or compile_kernel( + build_implicit_gemm_conv_wgrad(spec, arch=arch), arch=arch + ) + launcher = KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=wgrad_streamk_signature(spec), + ) + dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} + torch.manual_seed(0) + X = torch.empty(p.N, p.Hi, p.Wi, p.C).uniform_(-1, 1).to(dtypes[spec.data.dtype_b]) + dY = torch.empty(p.N, p.Ho, p.Wo, p.K).uniform_(-1, 1).to(dtypes[spec.data.dtype_a]) + dW = torch.empty(p.K, p.Y, p.X, p.C // p.groups, dtype=dtypes[spec.data.dtype_d]) + ref = _reference(X.float(), dY.float(), p) + + rt = Runtime() + ws_bytes = wgrad_streamk_workspace_nbytes(spec, arch=arch, problem=p) + bufs = [rt.alloc(dY.nbytes), rt.alloc(X.nbytes), rt.alloc(dW.nbytes)] + ws = rt.alloc(max(ws_bytes, 4)) + bufs.append(ws) + rt.memcpy_h2d(bufs[0], _u8(dY), dY.nbytes) + rt.memcpy_h2d(bufs[1], _u8(X), X.nbytes) + reduction = spec.streamk_reduction + flags_bytes, partials_bytes = wgrad_streamk_workspace_layout( + spec, arch=arch, problem=p + ) + extra = {} + if reduction == "workspace": + extra = dict(ws_ptr=ws, ws_bytes=ws_bytes) + elif reduction in ("linear", "tree"): + extra = dict(sk_flags=ws, sk_partials=ws + flags_bytes) + values = wgrad_streamk_launch_values( + spec, + bufs[0], + bufs[1], + bufs[2], + dY.nbytes, + X.nbytes, + dW.nbytes, + arch=arch, + problem=p, + **extra, + ) + if reduction == "workspace": + rspec = WgradReduceSpec( + problem=p, + dtype_d=spec.data.dtype_d, + groups=p.groups, + ws_replicas=spec.ws_replicas, + ) + rart = compile_kernel( + build_conv_wgrad_workspace_reduce(rspec, arch=arch), arch=arch + ) + reducer = KernelLauncher( + hsaco=rart.hsaco, + kernel_name=rart.kernel_name, + signature=wgrad_reduce_signature(rspec), + ) + rvalues = { + "ws_ptr": ws, + "dw_ptr": bufs[2], + "wg_M": p.K // p.groups, + "wg_N": p.Y * p.X * (p.C // p.groups), + "ws_bytes": ws_bytes, + "dw_bytes": dW.nbytes, + "groups": p.groups, + } + + outs = [] + try: + for _ in range(launches): + # The documented zeroing contract, per reduction: atomic accumulates + # into dW, workspace into the scratch, linear/tree only reset flags. + # Everything the contract leaves alone is poisoned (0xFF bytes are + # NaN in every float dtype): a stale value from the previous + # launch would otherwise be bit-identical to the right one and + # hide an early read of a partial or an unwritten dW element. + if reduction == "atomic": + rt.memset(bufs[2], 0, dW.nbytes) + else: + rt.memset(bufs[2], 0xFF, dW.nbytes) + if reduction == "workspace": + rt.memset(ws, 0, ws_bytes) + elif reduction in ("linear", "tree"): + if partials_bytes: + rt.memset(ws + flags_bytes, 0xFF, partials_bytes) + if flags_bytes: + rt.memset(ws, 0, flags_bytes) + launcher( + values, + config=LaunchConfig( + grid=wgrad_streamk_grid(spec, arch=arch, problem=p), + block=(spec.block_size, 1, 1), + fence=True, + ), + ) + if reduction == "workspace": + reducer( + rvalues, + config=LaunchConfig( + grid=wgrad_reduce_grid(rspec), + block=(rspec.block_size, 1, 1), + fence=True, + ), + ) + out = torch.empty_like(dW) + rt.memcpy_d2h(_u8(out), bufs[2], dW.nbytes) + outs.append(out) + finally: + for buf in bufs: + rt.free(buf) + return outs, ref + + +def _check(test, spec, *, launches: int = 1, problem=None, artifact=None): + outs, ref = _run(spec, launches=launches, problem=problem, artifact=artifact) + p = problem if problem is not None else spec.problem + out = outs[0].float() + written = out.reshape(p.K, -1).abs().sum(dim=1) > 0 + expected = ref.reshape(p.K, -1).abs().sum(dim=1) > 0 + missing = int((expected & ~written).sum()) + test.assertEqual( + missing, 0, f"{missing}/{p.K} dW output-channel rows were never written" + ) + rel = float((out - ref).abs().max() / ref.abs().max().clamp(min=1.0)) + test.assertLess(rel, _TOL[spec.data.dtype_d], f"rel_err={rel:.3e}") + return outs + + +def _spec(problem, *, mode, reduction, dtype_d, ctas): + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + from kernels.common.conv_implicit_gemm_wgrad import WgradConvSpec + + ab = "bf16" if dtype_d == "bf16" else "fp16" + edge = 32 if GPU_ARCH == "gfx950" else 16 + return WgradConvSpec( + problem=problem, + data=ConvDataSpec(dtype_a=ab, dtype_b=ab, dtype_d=dtype_d), + tile_m=64, + tile_n=64, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=edge, + warp_tile_n=edge, + warp_tile_k=16, + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=ctas, + ) + + +def _problems(): + from kernels.common._conv_implicit_gemm_common import ConvProblem + + return { + # 6 tiles, 32 iterations each, over 16 CTAs: up to 3 partners per tile, + # two tree rounds. + "multi_partner": ( + ConvProblem(N=8, Hi=16, Wi=16, C=16, K=96, Y=3, X=3, pH=1, pW=1), + 16, + ), + # 25 tiles over 7 CTAs: 14 data-parallel tiles plus 11 stream-K tiles. + "dp_plus_sk": ( + ConvProblem(N=2, Hi=8, Wi=8, C=32, K=320, Y=3, X=3, pH=1, pW=1), + 7, + ), + "strided": ( + ConvProblem( + N=3, Hi=13, Wi=11, C=24, K=72, Y=3, X=3, sH=2, sW=2, pH=1, pW=1 + ), + 5, + ), + } + + +@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") +class TestConvWgradStreamKNumerics(unittest.TestCase): + def test_modes_and_reductions(self): + for name, (problem, ctas) in _problems().items(): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in _REDUCTIONS: + spec = _spec( + problem, + mode=mode, + reduction=reduction, + dtype_d=dtype_d, + ctas=ctas, + ) + with self.subTest(shape=name, mode=mode, reduction=reduction): + _check(self, spec) + + def test_default_pool(self): + # streamk_ctas=-1: one CTA per CU for linear/tree, the occupancy-sized + # pool (several CTAs per CU) for the reductions that never wait. + # dp_plus_sk / multi_partner have far fewer tiles than that pool, so + # they run all-DP; "splits" (17 x 27 = 459 tiles) is the one that + # actually takes a linear/tree/workspace/atomic fixup at the default + # pool (on gfx950: 256 linear/tree CTAs, 1024 atomic/workspace CTAs). + from kernels.common._conv_implicit_gemm_common import ConvProblem + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_partition + + shapes = { + name: _problems()[name][0] for name in ("dp_plus_sk", "multi_partner") + } + shapes["splits"] = ConvProblem( + N=2, Hi=12, Wi=12, C=192, K=1056, Y=3, X=3, pH=1, pW=1 + ) + for name, problem in shapes.items(): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in _REDUCTIONS: + spec = _spec( + problem, + mode=mode, + reduction=reduction, + dtype_d=dtype_d, + ctas=-1, + ) + with self.subTest(shape=name, mode=mode, reduction=reduction): + if name == "splits": + part = wgrad_streamk_partition(spec, arch=GPU_ARCH) + self.assertGreater(part.sk_ctas, 0, part) + self.assertGreater(part.sk_tiles, 0, part) + _check(self, spec) + + +@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") +class TestConvWgradStreamKGrouped(unittest.TestCase): + """Groups folded into GEMM-M, including a partial last M tile per group.""" + + def test_grouped(self): + from kernels.common._conv_implicit_gemm_common import ConvProblem + + problems = { + "g3_partial_m": ( + ConvProblem( + N=4, Hi=10, Wi=10, C=48, K=120, Y=3, X=3, pH=1, pW=1, groups=3 + ), + 6, + ), + "g8": ( + ConvProblem( + N=2, Hi=12, Wi=12, C=64, K=64, Y=3, X=3, pH=1, pW=1, groups=8 + ), + 5, + ), + } + for name, (problem, ctas) in problems.items(): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in _REDUCTIONS: + spec = _spec( + problem, + mode=mode, + reduction=reduction, + dtype_d=dtype_d, + ctas=ctas, + ) + with self.subTest(shape=name, mode=mode, reduction=reduction): + _check(self, spec) + + +@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") +class TestConvWgradStreamKAot(unittest.TestCase): + """One binary per spec serves every launch shape. + + The partition is a kernarg block, so a kernel built for one problem must + be correct on others -- with a different tile count, iteration count, + data-parallel/stream-K split and group count. + """ + + def test_one_binary_many_shapes(self): + from rocke import compile_kernel + from kernels.common._conv_implicit_gemm_common import ConvProblem + from kernels.common.conv_implicit_gemm_wgrad import ( + build_implicit_gemm_conv_wgrad, + ) + + build = ConvProblem(N=2, Hi=8, Wi=8, C=48, K=96, Y=3, X=3, pH=1, pW=1, groups=3) + launches = [ + ConvProblem(N=4, Hi=10, Wi=10, C=48, K=120, Y=3, X=3, pH=1, pW=1, groups=3), + ConvProblem(N=8, Hi=16, Wi=16, C=24, K=96, Y=3, X=3, pH=1, pW=1, groups=3), + ConvProblem( + N=3, + Hi=13, + Wi=11, + C=72, + K=72, + Y=3, + X=3, + sH=2, + sW=2, + pH=1, + pW=1, + groups=3, + ), + ] + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in _REDUCTIONS: + spec = _spec( + build, mode=mode, reduction=reduction, dtype_d=dtype_d, ctas=6 + ) + art = compile_kernel( + build_implicit_gemm_conv_wgrad(spec, arch=GPU_ARCH), arch=GPU_ARCH + ) + for i, problem in enumerate(launches): + with self.subTest(mode=mode, reduction=reduction, shape=i): + _check(self, spec, problem=problem, artifact=art) + + +@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") +class TestConvWgradStreamKDeterminism(unittest.TestCase): + """Linear and tree fold partials in a fixed order: dW is bit-exact. + + With the partial slots NaN-poisoned this also catches a missing release on + the partial hand-off, but not a missing acquire -- see the module + docstring and ``TestConvWgradStreamKFlagOrderingIR``/``...Isa``. + """ + + def test_repeat_launches_are_bit_identical(self): + import torch + + problem, ctas = _problems()["multi_partner"] + for mode in ("dp_sk", "persistent"): + for reduction in ("linear", "tree"): + spec = _spec( + problem, mode=mode, reduction=reduction, dtype_d="fp16", ctas=ctas + ) + with self.subTest(mode=mode, reduction=reduction): + outs = _check(self, spec, launches=20) + for i, out in enumerate(outs[1:], start=1): + self.assertTrue( + torch.equal(out, outs[0]), + f"launch {i} differs from launch 0", + ) + + +# -------------------------------------------------------------------------- +# Flag hand-off ordering, pinned in codegen. +# +# The repeat-launch test above catches a missing release but cannot catch a +# missing acquire (see the module docstring), so the acquire -- and, for +# symmetry, the release -- is asserted on the code a real linear and tree +# kernel lowers to. Host-only: no GPU is needed for either level. + +_FLAG_ARCHES = ("gfx950", "gfx942") +_FLAG_KERNELS = ( + ("dp_sk", "linear"), + ("dp_sk", "tree"), + ("persistent", "linear"), + ("persistent", "tree"), +) + + +def _flag_comgr_input(arch: str, mode: str, reduction: str): + """The lowered ``.ll`` + COMGR target of a real stream-K fixup kernel. + + Always the Python engine (the C++ engine is byte-identical by contract and + its extension may not be importable), and lowered with ``arch`` as the + target rather than the detected GPU, so the gfx942 code is checked on a + gfx950 host and vice versa. + """ + from rocke.helpers.compile import lower_kernel_for_comgr + + from kernels.common._conv_implicit_gemm_common import ConvDataSpec, ConvProblem + from kernels.common.conv_implicit_gemm_wgrad import ( + WgradConvSpec, + build_implicit_gemm_conv_wgrad, + ) + + edge = 32 if arch == "gfx950" else 16 + spec = WgradConvSpec( + problem=ConvProblem(N=8, Hi=16, Wi=16, C=16, K=96, Y=3, X=3, pH=1, pW=1), + data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d="fp16"), + tile_m=64, + tile_n=64, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=edge, + warp_tile_n=edge, + warp_tile_k=16, + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=16, + ) + kernel = build_implicit_gemm_conv_wgrad(spec, arch=arch) + return lower_kernel_for_comgr(kernel, arch=arch, backend="python") + + +def _ll_blocks(ll: str) -> dict: + """``{label: [instruction lines]}`` of every basic block in ``ll``.""" + blocks, cur = {}, None + for line in ll.splitlines(): + m = re.match(r"^([A-Za-z0-9_.$-]+):", line) + if m: + cur = blocks.setdefault(m.group(1), []) + elif cur is not None and line.startswith(" "): + cur.append(line.strip()) + return blocks + + +def _gep_base(ll: str, ptr: str) -> str: + """Base pointer of the ``getelementptr`` that defines ``ptr``.""" + m = re.search( + rf"^\s*{re.escape(ptr)} = getelementptr\b[^,]*, ptr addrspace\(1\) (%[\w.]+),", + ll, + re.MULTILINE, + ) + return m.group(1) if m else "" + + +def _assert_flag_ll(test, ll: str) -> None: + """Acquire on every flag spin, fence + release on every flag publish.""" + blocks = _ll_blocks(ll) + load_re = re.compile( + r"^(%[\w.]+) = load atomic i32, ptr addrspace\(1\) (%[\w.]+) (.*), align 4$" + ) + store_re = re.compile( + r"^store atomic i32 1, ptr addrspace\(1\) (%[\w.]+) (.*), align 4$" + ) + spins = publishes = 0 + for label, insts in blocks.items(): + for i, inst in enumerate(insts): + m = load_re.match(inst) + if m: + spins += 1 + test.assertEqual(_gep_base(ll, m.group(2)), "%sk_flags", inst) + test.assertEqual(m.group(3), 'syncscope("agent") acquire', inst) + # A spin: the block re-branches to itself until the flag hits. + test.assertRegex( + insts[-1], rf"^br i1 .*label %{re.escape(label)}\b", label + ) + continue + m = store_re.match(inst) + if m: + publishes += 1 + test.assertEqual(_gep_base(ll, m.group(1)), "%sk_flags", inst) + test.assertEqual(m.group(2), 'syncscope("agent") release', inst) + # The nearest preceding memory op is the agent release fence, + # and every partial store of the block is ordered before it. + mem = [ + j + for j in range(i) + if re.match( + r"^(fence|store|%[\w.]+ = (load|atomicrmw))\b", insts[j] + ) + ] + test.assertTrue(mem, f"{label}: flag store with no fence before it") + test.assertEqual( + insts[mem[-1]], 'fence syncscope("agent") release', label + ) + test.assertTrue( + any(re.match(r"^store float ", insts[j]) for j in mem[:-1]), + f"{label}: no partial store before the release fence", + ) + test.assertGreater(spins, 0, "no flag-wait spin loop in the kernel") + test.assertGreater(publishes, 0, "no flag publish in the kernel") + + +def _objdump() -> str | None: + for candidate in ( + os.environ.get("LLVM_OBJDUMP"), + os.path.join(os.environ.get("ROCM_PATH", "/opt/rocm"), "llvm/bin/llvm-objdump"), + shutil.which("llvm-objdump"), + ): + if candidate and os.path.isfile(candidate): + return candidate + return None + + +def _isa_lines(objdump: str, comgr_input, arch: str) -> list: + from rocke.runtime.comgr import build_hsaco_from_llvm_ir + + hsaco, _ = build_hsaco_from_llvm_ir( + comgr_input.llvm_text, + isa=comgr_input.isa, + options=list(comgr_input.options), + ) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "streamk.hsaco" + path.write_bytes(hsaco) + text = subprocess.run( + [objdump, "-d", f"--mcpu={arch}", str(path)], + check=True, + capture_output=True, + text=True, + ).stdout + lines = (ln.split("//")[0].strip() for ln in text.splitlines()) + return [ln for ln in lines if ln and not ln.endswith(">:")] + + +_SC1_LOAD = re.compile(r"^global_load_dword\b.*\bsc1\b") +_SC1_STORE = re.compile(r"^global_store_dword\b.*\bsc1\b") + + +def _assert_flag_isa(test, isa: list) -> None: + """gfx94x agent-scope acquire/release: ``buffer_inv sc1`` after each spin + load, ``buffer_wbl2 sc1`` + ``s_waitcnt vmcnt(0)`` before each flag store. + """ + spins = [i for i, ln in enumerate(isa) if _SC1_LOAD.match(ln)] + stores = [i for i, ln in enumerate(isa) if _SC1_STORE.match(ln)] + test.assertTrue(spins, "no sc1 flag load in the ISA") + test.assertTrue(stores, "no sc1 flag store in the ISA") + for i in spins: + # Up to the spin's back-branch: wait for the load, then invalidate so + # the partial loads after the exit cannot hit a stale line. + end = next( + (j for j in range(i + 1, len(isa)) if isa[j].startswith("s_cbranch")), + len(isa), + ) + body = isa[i + 1 : end] + test.assertIn("buffer_inv sc1", body, f"spin at {i}: no acquire invalidate") + test.assertLess( + body.index("s_waitcnt vmcnt(0)"), + body.index("buffer_inv sc1"), + f"spin at {i}: buffer_inv before the flag load completed", + ) + for i in stores: + # Back to the last partial store: write back L2 and drain every + # outstanding store before the flag goes out. + start = max( + ( + j + for j in range(i) + if isa[j].startswith(("global_store", "buffer_store", "global_atomic")) + ), + default=-1, + ) + window = isa[start + 1 : i] + test.assertIn("buffer_wbl2 sc1", window, f"flag store at {i}: no writeback") + test.assertIn( + "s_waitcnt vmcnt(0)", window, f"flag store at {i}: stores not drained" + ) + + +class TestConvWgradStreamKFlagOrderingIR(unittest.TestCase): + """The lowered ``.ll`` of real linear/tree kernels orders the hand-off.""" + + def test_acquire_spin_and_release_publish(self): + for arch in _FLAG_ARCHES: + for mode, reduction in _FLAG_KERNELS: + with self.subTest(arch=arch, mode=mode, reduction=reduction): + ll = _flag_comgr_input(arch, mode, reduction).llvm_text + _assert_flag_ll(self, ll) + + +class TestConvWgradStreamKFlagOrderingIsa(unittest.TestCase): + """The same kernels' gfx942/gfx950 ISA keeps the cache maintenance. + + Offline compile (COMGR) + ``llvm-objdump``; no GPU. Skips when either + tool is unavailable. + """ + + def test_acquire_invalidate_and_release_writeback(self): + from rocke.runtime import comgr + + objdump = _objdump() + if objdump is None: + self.skipTest("llvm-objdump not found (LLVM_OBJDUMP / ROCM_PATH / PATH)") + try: + comgr._resolve_lib() + except (comgr.ComgrError, OSError) as e: + self.skipTest(f"COMGR not loadable: {str(e)[:200]}") + for arch in _FLAG_ARCHES: + for mode, reduction in _FLAG_KERNELS: + with self.subTest(arch=arch, mode=mode, reduction=reduction): + c = _flag_comgr_input(arch, mode, reduction) + _assert_flag_isa(self, _isa_lines(objdump, c, arch)) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py new file mode 100644 index 000000000000..15bc917bd8c6 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py @@ -0,0 +1,277 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""The implicit-GEMM sweep's stream-K wgrad launch helper. + +``benchmark_implicit_gemm_conv._streamk_launch`` sets up one stream-K launch +for the JIT sweep and for ``--run-from-cache``: it sizes and allocates the +reduction's workspace, fills the kernargs and wraps the per-reduction zeroing +contract (and the Stage 2 fold for ``workspace``) in a ``run()`` closure. + +The case pinned here is a partition with no stream-K CTAs at all: when every +output tile is data-parallel the workspace is zero bytes, and the helper must +still hand the kernel valid flag / partial pointers instead of failing (an +uncaught error there aborts the whole sweep). The host tests drive the helper +with a recording runtime; the GPU test launches the real kernel through it +and checks dW, and skips without a gfx942/gfx950 GPU and torch. +""" + +from __future__ import annotations + +import ctypes +import importlib.util +import unittest + +from rocke.runtime.hip_module import get_device_arch + +from benchmarks.common.benchmark_implicit_gemm_conv import _streamk_launch + +_HAS_TORCH = importlib.util.find_spec("torch") is not None +GPU_ARCH = get_device_arch(0) + +# Every output tile is data-parallel under the default gfx950 pool: 1 x 9 +# tiles of 64x64 over K=64, Y*X*C=576, so the stream-K CTA count is 0. +_ALL_DP = dict(N=2, Hi=14, Wi=14, C=64, K=64, Y=3, X=3, pH=1, pW=1) +# 459 tiles: more than the default pool, so stream-K CTAs share some of them. +_SPLITS = dict(N=2, Hi=12, Wi=12, C=192, K=1056, Y=3, X=3, pH=1, pW=1) +# (reduction, dW dtype): "atomic" needs fp32 dW. +_REDUCTIONS = ( + ("linear", "fp16"), + ("tree", "fp16"), + ("atomic", "fp32"), + ("workspace", "bf16"), +) + + +def _spec(shape, *, arch, reduction, dtype_d, mode="dp_sk"): + from kernels.common._conv_implicit_gemm_common import ConvDataSpec, ConvProblem + from kernels.common.conv_implicit_gemm_wgrad import WgradConvSpec + + ab = "bf16" if dtype_d == "bf16" else "fp16" + edge = 32 if arch == "gfx950" else 16 + return WgradConvSpec( + problem=ConvProblem(**shape), + data=ConvDataSpec(dtype_a=ab, dtype_b=ab, dtype_d=dtype_d), + tile_m=64, + tile_n=64, + tile_k=32, + warp_m=2, + warp_n=2, + warp_tile_m=edge, + warp_tile_n=edge, + warp_tile_k=16, + streamk=mode, + streamk_reduction=reduction, + ) + + +class _RecordingRuntime: + """The subset of ``Runtime`` the helper uses, recording every call.""" + + def __init__(self): + self.allocs: dict = {} + self.freed: list = [] + self.memsets: list = [] + self._next = 0x10000 + + def alloc(self, nbytes): + ptr = self._next + self._next += max(nbytes, 1) + 0x1000 + self.allocs[ptr] = nbytes + return ptr + + def free(self, ptr): + self.freed.append(ptr) + + def memset(self, ptr, value, nbytes): + self.memsets.append((ptr, value, nbytes)) + + +class _RecordingLauncher: + def __init__(self): + self.calls: list = [] + + def __call__(self, values, *, config): + self.calls.append((values, config)) + + +def _launch_host(spec, arch): + from rocke.runtime.launcher import LaunchConfig + + rt = _RecordingRuntime() + launcher = _RecordingLauncher() + stage2 = _RecordingLauncher() + # Pre-seed the Stage 2 memo so the workspace reduction needs no compiler. + stage2_memo = {(spec.data.dtype_d, spec.ws_replicas): stage2} + run, free, pool = _streamk_launch( + rt=rt, + launcher=launcher, + spec=spec, + problem=spec.problem, + arch=arch, + ptrs=(0x100, 0x200, 0x300), + sizes=(64, 64, 64), + stage2=stage2_memo, + KernelLauncher=None, + LaunchConfig=LaunchConfig, + ) + return rt, launcher, stage2, run, free, pool + + +class TestStreamKSweepLaunchHost(unittest.TestCase): + """Host-only: the helper's workspace handling, with a recording runtime.""" + + ARCH = "gfx950" + + def test_all_dp_partition_has_no_workspace(self): + # The premise of the other cases: this shape takes no stream-K CTAs. + from kernels.common.conv_implicit_gemm_wgrad import ( + wgrad_streamk_partition, + wgrad_streamk_workspace_nbytes, + ) + + spec = _spec(_ALL_DP, arch=self.ARCH, reduction="linear", dtype_d="fp16") + part = wgrad_streamk_partition(spec, arch=self.ARCH) + self.assertEqual(part.sk_ctas, 0) + self.assertEqual(wgrad_streamk_workspace_nbytes(spec, arch=self.ARCH), 0) + + def test_all_dp_partition_launches(self): + for reduction, dtype_d in _REDUCTIONS: + with self.subTest(reduction=reduction): + spec = _spec( + _ALL_DP, arch=self.ARCH, reduction=reduction, dtype_d=dtype_d + ) + rt, launcher, stage2, run, free, _ = _launch_host(spec, self.ARCH) + # One (stub) allocation, never a zero-byte one. + self.assertEqual(len(rt.allocs), 1) + (ws,) = rt.allocs + self.assertGreaterEqual(rt.allocs[ws], 4) + run() + self.assertEqual(len(launcher.calls), 1) + values, _ = launcher.calls[0] + if reduction in ("linear", "tree"): + self.assertEqual(values["sk_flags"], ws) + self.assertEqual(values["sk_partials"], ws) + self.assertEqual(len(stage2.calls), reduction == "workspace") + free() + self.assertEqual(rt.freed, [ws]) + + def test_split_partition_lays_out_flags_then_partials(self): + from kernels.common.conv_implicit_gemm_wgrad import ( + wgrad_streamk_partition, + wgrad_streamk_workspace_layout, + ) + + spec = _spec(_SPLITS, arch=self.ARCH, reduction="linear", dtype_d="fp16") + self.assertGreater(wgrad_streamk_partition(spec, arch=self.ARCH).sk_ctas, 0) + flags_bytes, partials_bytes = wgrad_streamk_workspace_layout( + spec, arch=self.ARCH + ) + self.assertGreater(flags_bytes, 0) + rt, launcher, _, run, free, _ = _launch_host(spec, self.ARCH) + (ws,) = rt.allocs + self.assertEqual(rt.allocs[ws], flags_bytes + partials_bytes) + run() + values, _ = launcher.calls[0] + self.assertEqual(values["sk_flags"], ws) + self.assertEqual(values["sk_partials"], ws + flags_bytes) + # linear/tree reset only the flag table before each launch. + self.assertEqual(rt.memsets, [(ws, 0, flags_bytes)]) + free() + self.assertEqual(rt.freed, [ws]) + + +def _gpu_skip_reason() -> str: + if not GPU_ARCH: + return "no ROCm GPU detected" + if not _HAS_TORCH: + return "torch not importable" + if GPU_ARCH not in ("gfx942", "gfx950"): + return f"stream-K wgrad is gfx942/gfx950 only; got {GPU_ARCH}" + return "" + + +_GPU_SKIP_REASON = _gpu_skip_reason() + + +@unittest.skipUnless(not _GPU_SKIP_REASON, _GPU_SKIP_REASON or "no GPU") +class TestStreamKSweepLaunchGpu(unittest.TestCase): + """The helper's ``run()`` on the GPU for the all-data-parallel shape.""" + + def _check(self, reduction, dtype_d): + import torch + + from rocke import compile_kernel + from rocke.runtime import synchronize_and_release + from rocke.runtime.hip_module import Runtime + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + from kernels.common.conv_implicit_gemm_wgrad import ( + build_implicit_gemm_conv_wgrad, + wgrad_streamk_signature, + ) + + arch = GPU_ARCH + spec = _spec(_ALL_DP, arch=arch, reduction=reduction, dtype_d=dtype_d) + p = spec.problem + art = compile_kernel(build_implicit_gemm_conv_wgrad(spec, arch=arch), arch=arch) + launcher = KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=wgrad_streamk_signature(spec), + ) + dt = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} + torch.manual_seed(0) + X = torch.empty(p.N, p.Hi, p.Wi, p.C).uniform_(-1, 1).to(dt[spec.data.dtype_b]) + dY = torch.empty(p.N, p.Ho, p.Wo, p.K).uniform_(-1, 1).to(dt[spec.data.dtype_a]) + dW = torch.empty(p.K, p.Y, p.X, p.C, dtype=dt[dtype_d]) + ref = torch.nn.grad.conv2d_weight( + X.float().permute(0, 3, 1, 2), + weight_size=(p.K, p.C, p.Y, p.X), + grad_output=dY.float().permute(0, 3, 1, 2), + padding=(p.pH, p.pW), + ).permute(0, 2, 3, 1) + + def u8(t): + return (ctypes.c_uint8 * t.nbytes).from_address(t.data_ptr()) + + rt = Runtime() + bufs = [rt.alloc(t.nbytes) for t in (dY, X, dW)] + try: + rt.memcpy_h2d(bufs[0], u8(dY), dY.nbytes) + rt.memcpy_h2d(bufs[1], u8(X), X.nbytes) + # Poison dW: a reduction that leaves it alone must overwrite it. + rt.memset(bufs[2], 0xFF if reduction != "atomic" else 0, dW.nbytes) + run, free, _ = _streamk_launch( + rt=rt, + launcher=launcher, + spec=spec, + problem=p, + arch=arch, + ptrs=tuple(int(b) for b in bufs), + sizes=(dY.nbytes, X.nbytes, dW.nbytes), + stage2={}, + KernelLauncher=KernelLauncher, + LaunchConfig=LaunchConfig, + ) + try: + run() + synchronize_and_release(0) + finally: + free() + rt.memcpy_d2h(u8(dW), bufs[2], dW.nbytes) + finally: + for b in bufs: + rt.free(b) + out = dW.float() + rel = float((out - ref).abs().max() / ref.abs().max().clamp(min=1.0)) + tol = {"fp32": 1e-4, "fp16": 2e-3, "bf16": 2e-2}[dtype_d] + self.assertLess(rel, tol, f"{reduction}: rel_err={rel:.3e}") + + def test_all_dp_linear(self): + self._check("linear", "fp16") + + def test_all_dp_workspace(self): + self._check("workspace", "bf16") + + +if __name__ == "__main__": + unittest.main() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index 8b584268bc22..a7db23e478bf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -437,6 +437,8 @@ if(ROCKE_INSTALL_TESTS) # prefix on each, matching the other installed tests. set(_ROCKE_LIBRARY_CONV_TEST_RELPATHS "tests/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py" + "tests/library/tests/test_conv_wgrad_streamk.py" + "tests/library/tests/test_streamk_sweep_launch.py" "tests/library/tests/test_conv_build_signatures.py" "tests/library/tests/test_builder_signature_contract.py" "tests/library/tests/test_library_layering.py") @@ -495,6 +497,7 @@ if(ROCKE_INSTALL_TESTS) set(_ROCKE_LIBRARY_GPU_TEST_RELPATHS "tests/library/tests/test_conv_fwd_correctness.py" "tests/library/tests/test_conv_wgrad_correctness.py" + "tests/library/tests/test_conv_wgrad_streamk_gpu.py" "tests/library/tests/test_conv_dgrad_correctness.py" "tests/library/tests/test_direct_conv_correctness.py" "tests/library/tests/test_attention_dense_gfx942_numeric.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 dc7aeabb836b..1844de8c77ef 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 @@ -1462,6 +1462,7 @@ void fill_conv_problem(rocke_conv_problem_t* p, const py::dict& d) p->pW = dict_int(d, "pW", p->pW); p->dH = dict_int(d, "dH", p->dH); p->dW = dict_int(d, "dW", p->dW); + p->groups = dict_int(d, "groups", p->groups); p->is_3d = dict_bool(d, "is_3d", p->is_3d); p->Di = dict_int(d, "Di", p->Di); p->Z = dict_int(d, "Z", p->Z); @@ -2271,10 +2272,37 @@ rocke_implicit_gemm_conv_wgrad_spec_t conv_wgrad_build_spec(const py::dict& d, s.split_k = dict_int(d, "split_k", s.split_k); s.two_stage = dict_bool(d, "two_stage", s.two_stage); s.ws_replicas = dict_int(d, "ws_replicas", s.ws_replicas); + s.streamk_ctas = dict_int(d, "streamk_ctas", s.streamk_ctas); + s.group_merge = dict_int(d, "group_merge", s.group_merge); + s.async_dma = dict_bool(d, "async_dma", s.async_dma); + s.unroll_k = dict_bool(d, "unroll_k", s.unroll_k); + s.lds_k_outer = dict_bool(d, "lds_k_outer", s.lds_k_outer); + s.chiplet_swizzle = dict_bool(d, "chiplet_swizzle", s.chiplet_swizzle); + s.chiplet_wgm = dict_int(d, "chiplet_wgm", s.chiplet_wgm); + s.chiplet_num_xcds = dict_int(d, "chiplet_num_xcds", s.chiplet_num_xcds); + s.chiplet_chunk_size = dict_int(d, "chiplet_chunk_size", s.chiplet_chunk_size); + /* Optional ints: absent or None keeps Python's None (has_* = false). + * Only the fields the C builder or validator honours are read: the C wgrad + * builder does not apply lds_k_pad or vector_size_a/b, so reading them + * would tag the kernel name (e.g. "pad") without changing its body. */ + if(d.contains("waves_per_eu") && !d["waves_per_eu"].is_none()) + { + s.has_waves_per_eu = true; + s.waves_per_eu = d["waves_per_eu"].cast(); + } + if(d.contains("vector_size_c") && !d["vector_size_c"].is_none()) + { + s.has_vector_size_c = true; + s.vector_size_c = d["vector_size_c"].cast(); + } { std::string v; if(dict_str(d, "name", v)) s.name = keep(v); + if(dict_str(d, "streamk", v)) + s.streamk = keep(v); + if(dict_str(d, "streamk_reduction", v)) + s.streamk_reduction = keep(v); if(dict_str(d, "dtype_a", v)) s.dtype_a = keep(v); if(dict_str(d, "dtype_b", v)) @@ -3890,6 +3918,59 @@ PYBIND11_MODULE(rocke_engine, m) "Return workspace bytes for the two-stage wgrad path.\n" "Formula: groups * ws_replicas * wg_M * wg_N * 4 (always f32).\n" "Returns 0 when two_stage=false, or split_k <= 1."); + m.def( + "conv_wgrad_is_valid_spec", + [](const py::dict& d, const std::string& arch) -> py::tuple { + std::deque store; + rocke_implicit_gemm_conv_wgrad_spec_t s = conv_wgrad_build_spec(d, store); + char reason[1024] = {0}; + const bool ok = rocke_implicit_gemm_conv_wgrad_is_valid_spec( + &s, arch_or_default(arch), reason, sizeof(reason)); + return py::make_tuple(ok, ok ? std::string("ok") : std::string(reason)); + }, + py::arg("spec"), + py::arg("arch") = "gfx950", + "C port of is_valid_wgrad_spec: (ok, reason), reason 'ok' when accepted."); + m.def( + "conv_wgrad_streamk_default_pool", + [](const std::string& reduction, const std::string& arch) -> int { + return rocke_wgrad_conv_streamk_default_pool(reduction.c_str(), arch_or_default(arch)); + }, + py::arg("reduction"), + py::arg("arch") = "gfx950", + "C port of wgrad_streamk_default_pool: the CTA pool of a streamk_ctas=-1 launch."); + m.def( + "conv_wgrad_streamk_host", + [](const py::dict& d, const std::string& arch) -> py::dict { + std::deque store; + rocke_implicit_gemm_conv_wgrad_spec_t s = conv_wgrad_build_spec(d, store); + /* Every helper resolves an auto pool (-1) for `arch` itself, as the + * Python wgrad_streamk_*(spec, arch=arch) twins do. Errors are + * ValueError, matching Python. */ + const char* a = arch_or_default(arch); + int grid[3] = {0, 0, 0}; + if(rocke_wgrad_conv_streamk_grid(&s, a, grid) != ROCKE_OK) + throw py::value_error( + "rocke_engine.conv_wgrad_streamk_host: needs streamk != 'off' and " + "streamk_ctas == -1 or > 0"); + size_t flags_bytes = 0; + size_t partials_bytes = 0; + if(rocke_wgrad_conv_streamk_workspace_layout(&s, a, &flags_bytes, &partials_bytes) + != ROCKE_OK) + throw py::value_error( + "rocke_engine.conv_wgrad_streamk_host: invalid stream-K workspace layout"); + py::dict out; + out["grid"] = py::make_tuple(grid[0], grid[1], grid[2]); + out["flags_bytes"] = flags_bytes; + out["partials_bytes"] = partials_bytes; + out["workspace_bytes"] = rocke_wgrad_conv_streamk_workspace_bytes(&s, a); + out["is_deterministic"] = rocke_wgrad_conv_spec_is_deterministic(&s); + return out; + }, + py::arg("spec"), + py::arg("arch") = "gfx950", + "Stream-K wgrad host plan of the C port: launch grid, workspace layout\n" + "and determinism, for cross-checking the Python wgrad_streamk_* helpers."); /* ---- attention families (separate TU; shared fmha/tiled struct tags) ---- */ register_attention(m); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp index cd34e3b984c0..7c2230c5c50e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp @@ -668,7 +668,12 @@ static const char* const rocke_opcode_names[ROCKE_OP__COUNT] = { "scf.if", "scf.if_else", "scf.yield", - "cf.return"}; + "cf.return", + + /* memref.* cross-workgroup synchronisation */ + "memref.fence", + "memref.global_flag_store", + "memref.global_flag_wait_eq"}; const char* rocke_opcode_name(rocke_opcode_t op) { @@ -914,7 +919,12 @@ static const bool rocke_opcode_pure[ROCKE_OP__COUNT] = { /* scf.if */ false, /* scf.if_else */ false, /* scf.yield */ false, - /* cf.return */ false}; + /* cf.return */ false, + + /* memref.* cross-workgroup synchronisation (effectful) */ + /* memref.fence */ false, + /* memref.global_flag_store */ false, + /* memref.global_flag_wait_eq */ false}; bool rocke_opcode_is_pure(rocke_opcode_t op) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp index fefbf7c9e429..fa16a53af676 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp @@ -503,6 +503,147 @@ static bool rocke_i_ordering_ok(const char* ordering) || !strcmp(ordering, "seq_cst")); } +/* IRBuilder._FENCE_SCOPES */ +static bool rocke_i_fence_scope_ok(const char* scope) +{ + return scope + && (!strcmp(scope, "workgroup") || !strcmp(scope, "agent") || !strcmp(scope, "system")); +} + +void rocke_b_fence(rocke_ir_builder_t* b, const char* scope, const char* ordering) +{ + rocke_attr_map_t a; + if(!rocke_i_live(b)) + return; + if(scope == NULL) + scope = "agent"; + if(ordering == NULL) + ordering = "acq_rel"; + if(!rocke_i_fence_scope_ok(scope)) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "unknown fence scope '%s'", scope); + return; + } + if(strcmp(ordering, "acquire") && strcmp(ordering, "release") && strcmp(ordering, "acq_rel") + && strcmp(ordering, "seq_cst")) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "unknown fence ordering '%s'", ordering); + return; + } + a = rocke_i_attrs(b); + rocke_attr_set_str(b, &a, "scope", scope); + rocke_attr_set_str(b, &a, "ordering", ordering); + (void)rocke_i_op0(b, ROCKE_OP_MEMREF_FENCE, NULL, 0, &a); +} + +/* ir._is_flag_ptr: the ptr a cross-workgroup flag op addresses. */ +static bool rocke_i_flag_ptr_ok(const rocke_type_t* t) +{ + return t && t->kind == ROCKE_TYPE_PTR && t->space && !strcmp(t->space, "global") && t->pointee + && rocke_i_type_is(t->pointee, "i32"); +} + +void rocke_b_global_flag_store(rocke_ir_builder_t* b, + rocke_value_t* ptr, + rocke_value_t* idx, + rocke_value_t* value, + const char* scope, + const char* ordering) +{ + rocke_value_t* ops[3]; + rocke_attr_map_t a; + if(!rocke_i_live(b)) + return; + if(!ptr || !idx || !value) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "global_flag_store: null operand"); + return; + } + if(scope == NULL) + scope = "agent"; + if(ordering == NULL) + ordering = "release"; + if(!rocke_i_flag_ptr_ok(ptr->type)) + { + (void)rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "global_flag_store expects a global i32 pointer, got %s", + ptr->type ? ptr->type->name : "?"); + return; + } + if(!rocke_i_type_is(value->type, "i32")) + { + (void)rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "global_flag_store expects an i32 value, got %s", + value->type->name); + return; + } + if(!rocke_i_fence_scope_ok(scope)) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "unknown flag scope '%s'", scope); + return; + } + if(strcmp(ordering, "monotonic") && strcmp(ordering, "release") && strcmp(ordering, "seq_cst")) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "unknown flag store ordering '%s'", ordering); + return; + } + ops[0] = ptr; + ops[1] = idx; + ops[2] = value; + a = rocke_i_attrs(b); + rocke_attr_set_str(b, &a, "scope", scope); + rocke_attr_set_str(b, &a, "ordering", ordering); + (void)rocke_i_op0(b, ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, ops, 3, &a); +} + +void rocke_b_global_flag_wait_eq(rocke_ir_builder_t* b, + rocke_value_t* ptr, + rocke_value_t* idx, + rocke_value_t* expect, + const char* scope) +{ + rocke_value_t* ops[3]; + rocke_attr_map_t a; + if(!rocke_i_live(b)) + return; + if(!ptr || !idx || !expect) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "global_flag_wait_eq: null operand"); + return; + } + if(scope == NULL) + scope = "agent"; + if(!rocke_i_flag_ptr_ok(ptr->type)) + { + (void)rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "global_flag_wait_eq expects a global i32 pointer, got %s", + ptr->type ? ptr->type->name : "?"); + return; + } + if(!rocke_i_type_is(expect->type, "i32")) + { + (void)rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "global_flag_wait_eq expects an i32 value, got %s", + expect->type->name); + return; + } + if(!rocke_i_fence_scope_ok(scope)) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "unknown flag scope '%s'", scope); + return; + } + ops[0] = ptr; + ops[1] = idx; + ops[2] = expect; + a = rocke_i_attrs(b); + rocke_attr_set_str(b, &a, "scope", scope); + (void)rocke_i_op0(b, ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, ops, 3, &a); +} + rocke_value_t* rocke_b_global_atomic_add(rocke_ir_builder_t* b, rocke_value_t* ptr, rocke_value_t* idx, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp index 7e6f6364b9ad..8023901a6163 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp @@ -762,6 +762,129 @@ static rocke_status_t _op_memref_global_atomic_add(rocke_h_lowerer_t* lw, const return lw->status; } +/* ir._FENCE_SCOPES / _FENCE_ORDERINGS / _FLAG_STORE_ORDERINGS. */ +static const char* const k_hip_sync_scopes[] = {"workgroup", "agent", "system", NULL}; +static const char* const k_hip_fence_orderings[] + = {"acquire", "release", "acq_rel", "seq_cst", NULL}; +static const char* const k_hip_flag_store_orderings[] = {"monotonic", "release", "seq_cst", NULL}; + +/* Python ir._sync_attr: the lowerer may run on unverified IR, so a missing or + * unknown scope / ordering is rejected rather than defaulted. */ +static const char* hip_sync_attr(rocke_h_lowerer_t* lw, + const rocke_op_t* op, + const char* key, + const char* const* allowed) +{ + const char* s = rocke_attr_get_str(&op->attrs, key); + if(!s) + { + rocke_h_fail( + lw, ROCKE_ERR_VALUE, "%s: %s attribute missing or not a string", op->name, key); + } + for(const char* const* a = allowed; *a; ++a) + { + if(!strcmp(s, *a)) + { + return s; + } + } + rocke_h_fail(lw, ROCKE_ERR_VALUE, "%s: unknown %s '%s'", op->name, key, s); +} + +/* Python _HIP_ATOMIC_ORDER / _HIP_FENCE_SCOPE / _HIP_MEMORY_SCOPE, applied to + * values hip_sync_attr already accepted. */ +static const char* hip_atomic_order(const char* o) +{ + if(!strcmp(o, "monotonic")) + return "__ATOMIC_RELAXED"; + if(!strcmp(o, "acquire")) + return "__ATOMIC_ACQUIRE"; + if(!strcmp(o, "release")) + return "__ATOMIC_RELEASE"; + if(!strcmp(o, "acq_rel")) + return "__ATOMIC_ACQ_REL"; + return "__ATOMIC_SEQ_CST"; +} + +static const char* hip_fence_scope(const char* s) +{ + if(!strcmp(s, "workgroup")) + return "workgroup"; + if(!strcmp(s, "agent")) + return "agent"; + return ""; +} + +static const char* hip_memory_scope(const char* s) +{ + if(!strcmp(s, "workgroup")) + return "__HIP_MEMORY_SCOPE_WORKGROUP"; + if(!strcmp(s, "agent")) + return "__HIP_MEMORY_SCOPE_AGENT"; + return "__HIP_MEMORY_SCOPE_SYSTEM"; +} + +/* Python _op_memref_fence */ +static rocke_status_t _op_memref_fence(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + const char *scp, *ord; + if(!rocke_h_live(lw)) + { + return lw->status; + } + scp = hip_sync_attr(lw, op, "scope", k_hip_sync_scopes); + ord = hip_sync_attr(lw, op, "ordering", k_hip_fence_orderings); + rocke_h_emitf( + lw, "__builtin_amdgcn_fence(%s, \"%s\");", hip_atomic_order(ord), hip_fence_scope(scp)); + return lw->status; +} + +/* Python _op_memref_global_flag_store */ +static rocke_status_t _op_memref_global_flag_store(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + const char *scp, *ord; + if(!rocke_h_live(lw)) + { + return lw->status; + } + if(op->num_operands < 3) + { + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.global_flag_store: too few operands"); + } + scp = hip_sync_attr(lw, op, "scope", k_hip_sync_scopes); + ord = hip_sync_attr(lw, op, "ordering", k_hip_flag_store_orderings); + rocke_h_emitf(lw, + "__hip_atomic_store(&%s[%s], %s, %s, %s);", + rocke_h_name(lw, op->operands[0]), + rocke_h_name(lw, op->operands[1]), + rocke_h_name(lw, op->operands[2]), + hip_atomic_order(ord), + hip_memory_scope(scp)); + return lw->status; +} + +/* Python _op_memref_global_flag_wait_eq */ +static rocke_status_t _op_memref_global_flag_wait_eq(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + const char* scp; + if(!rocke_h_live(lw)) + { + return lw->status; + } + if(op->num_operands < 3) + { + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.global_flag_wait_eq: too few operands"); + } + scp = hip_sync_attr(lw, op, "scope", k_hip_sync_scopes); + rocke_h_emitf(lw, + "while (__hip_atomic_load(&%s[%s], __ATOMIC_ACQUIRE, %s) != %s) {}", + rocke_h_name(lw, op->operands[0]), + rocke_h_name(lw, op->operands[1]), + hip_memory_scope(scp), + rocke_h_name(lw, op->operands[2])); + return lw->status; +} + /* Python _op_memref_global_atomic_add_f32 */ static rocke_status_t _op_memref_global_atomic_add_f32(rocke_h_lowerer_t* lw, const rocke_op_t* op) { @@ -1582,6 +1705,9 @@ const rocke_h_handler_entry_t* rocke_h_handlers_mem(void) {ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_F32, _op_memref_global_atomic_add_f32}, {ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_PK_BF16, _op_memref_global_atomic_add_pk_bf16}, {ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_PK_F16, _op_memref_global_atomic_add_pk_f16}, + {ROCKE_OP_MEMREF_FENCE, _op_memref_fence}, + {ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, _op_memref_global_flag_store}, + {ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, _op_memref_global_flag_wait_eq}, {ROCKE_OP_MEMREF_COOPERATIVE_GLOBAL_STORE, _op_memref_cooperative_global_store}, /* global pointer arithmetic + buffer rsrc */ {ROCKE_OP_TILE_GLOBAL_PTR_ADD, _op_tile_global_ptr_add}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index 9af3b8f3eafa..114e8ff853c6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -183,6 +183,115 @@ static void op_memref_global_store_typed(rocke_lower_t* L, const rocke_op_t* op) (long long)align); } +/* ir._FENCE_SCOPES / _FENCE_ORDERINGS / _FLAG_STORE_ORDERINGS. */ +static const char* const k_ll_sync_scopes[] = {"workgroup", "agent", "system", NULL}; +static const char* const k_ll_fence_orderings[] + = {"acquire", "release", "acq_rel", "seq_cst", NULL}; +static const char* const k_ll_flag_store_orderings[] = {"monotonic", "release", "seq_cst", NULL}; + +/* Python ir._sync_attr: the lowerer may run on unverified IR, so a missing or + * unknown scope / ordering is rejected rather than defaulted or rendered into + * invalid IR. */ +static const char* ll_sync_attr(rocke_lower_t* L, + const rocke_op_t* op, + const char* key, + const char* const* allowed) +{ + const char* s = rocke_attr_get_str(&op->attrs, key); + if(!s) + rocke_ll_fail( + L, ROCKE_ERR_VALUE, "%s: %s attribute missing or not a string", op->name, key); + for(const char* const* a = allowed; *a; ++a) + if(strcmp(s, *a) == 0) + return s; + rocke_ll_fail(L, ROCKE_ERR_VALUE, "%s: unknown %s '%s'", op->name, key, s); +} + +/* Python _llvm_syncscope: "system" is LLVM's default scope (no clause). */ +static const char* ll_syncscope(rocke_lower_t* L, const char* scope) +{ + if(strcmp(scope, "system") == 0) + return ""; + return rocke_arena_printf(&L->arena, " syncscope(\"%s\")", scope); +} + +/* Python _op_memref_fence. */ +static void op_memref_fence(rocke_lower_t* L, const rocke_op_t* op) +{ + const char* scope = ll_sync_attr(L, op, "scope", k_ll_sync_scopes); + const char* ordering = ll_sync_attr(L, op, "ordering", k_ll_fence_orderings); + rocke_ll_emitf(L, " fence%s %s", ll_syncscope(L, scope), ordering); +} + +/* Python _op_memref_global_flag_store. */ +static void op_memref_global_flag_store(rocke_lower_t* L, const rocke_op_t* op) +{ + if(op->num_operands != 3) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "memref.global_flag_store expects three operands"); + const rocke_value_t* ptr = op->operands[0]; + const rocke_value_t* idx = op->operands[1]; + const rocke_value_t* val = op->operands[2]; + const char* scope = ll_sync_attr(L, op, "scope", k_ll_sync_scopes); + const char* ordering = ll_sync_attr(L, op, "ordering", k_ll_flag_store_orderings); + const char* gep = rocke_ll_fresh(L, "gep"); + rocke_ll_emitf(L, + " %s = getelementptr inbounds i32, ptr addrspace(1) %s, i32 %s", + gep, + rocke_ll_operand(L, ptr), + rocke_ll_operand(L, idx)); + rocke_ll_emitf(L, + " store atomic i32 %s, ptr addrspace(1) %s%s %s, align 4", + rocke_ll_operand(L, val), + gep, + ll_syncscope(L, scope), + ordering); +} + +/* Python _op_memref_global_flag_wait_eq: a self-contained spin loop around an + * acquire atomic load; the fall-through block becomes current. */ +static void op_memref_global_flag_wait_eq(rocke_lower_t* L, const rocke_op_t* op) +{ + if(op->num_operands != 3) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "memref.global_flag_wait_eq expects three operands"); + const rocke_value_t* ptr = op->operands[0]; + const rocke_value_t* idx = op->operands[1]; + const rocke_value_t* expect = op->operands[2]; + const char* scope = ll_sync_attr(L, op, "scope", k_ll_sync_scopes); + const char* gep = rocke_ll_fresh(L, "gep"); + rocke_ll_block_t* cur = rocke_ll_current(L); + rocke_ll_block_t* spin; + rocke_ll_block_t* done; + const char* val; + const char* hit; + rocke_ll_block_emitf(L, + cur, + " %s = getelementptr inbounds i32, ptr addrspace(1) %s, i32 %s", + gep, + rocke_ll_operand(L, ptr), + rocke_ll_operand(L, idx)); + spin = rocke_ll_new_block(L, "flag.wait"); + if(!cur || !spin) + return; + rocke_ll_block_emitf(L, cur, " br label %%%s", spin->label); + cur->terminated = true; + val = rocke_ll_fresh(L, "flag"); + hit = rocke_ll_fresh(L, "flag.hit"); + rocke_ll_block_emitf(L, + spin, + " %s = load atomic i32, ptr addrspace(1) %s%s acquire, align 4", + val, + gep, + ll_syncscope(L, scope)); + rocke_ll_block_emitf( + L, spin, " %s = icmp eq i32 %s, %s", hit, val, rocke_ll_operand(L, expect)); + done = rocke_ll_new_block(L, "flag.done"); + if(!done) + return; + rocke_ll_block_emitf( + L, spin, " br i1 %s, label %%%s, label %%%s", hit, done->label, spin->label); + spin->terminated = true; +} + static void op_memref_global_atomic_add(rocke_lower_t* L, const rocke_op_t* op) { const rocke_value_t* ptr = op->operands[0]; @@ -1672,6 +1781,9 @@ void rocke_ll_register_mem(void) rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_STORE_TYPED, op_memref_global_store_typed); rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_STORE_VN, op_memref_global_store_vN); rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD, op_memref_global_atomic_add); + rocke_ll_set_handler(ROCKE_OP_MEMREF_FENCE, op_memref_fence); + rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, op_memref_global_flag_store); + rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, op_memref_global_flag_wait_eq); rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_F32, op_memref_global_atomic_add_f32); rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_PK_BF16, op_memref_global_atomic_add_pk_bf16); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp index d30fcbb53c6a..62afcbad52fe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp @@ -291,6 +291,8 @@ static const char* const* required_attrs(rocke_opcode_t o, int* n) static const char* mma_keys[] = {"op_id"}; static const char* yield_keys[] = {"num"}; static const char* asm_keys[] = {"template", "constraints"}; + static const char* scope_order_keys[] = {"scope", "ordering"}; + static const char* scope_keys[] = {"scope"}; switch(o) { case ROCKE_OP_ARITH_CONSTANT: @@ -309,6 +311,13 @@ static const char* const* required_attrs(rocke_opcode_t o, int* n) case ROCKE_OP_TILE_INLINE_ASM: *n = 2; return asm_keys; + case ROCKE_OP_MEMREF_FENCE: + case ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE: + *n = 2; + return scope_order_keys; + case ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ: + *n = 1; + return scope_keys; default: *n = 0; return NULL; @@ -320,6 +329,31 @@ static void check_op(verifier_t* v, const rocke_op_t* op); /* ---- contracts ---- */ +/* IRBuilder._FENCE_SCOPES / _FENCE_ORDERINGS / _FLAG_STORE_ORDERINGS. */ +static const char* const k_fence_scopes[] = {"workgroup", "agent", "system", NULL}; +static const char* const k_fence_orderings[] = {"acquire", "release", "acq_rel", "seq_cst", NULL}; +static const char* const k_flag_store_orderings[] = {"monotonic", "release", "seq_cst", NULL}; + +static bool in_set(const char* s, const char* const* set) +{ + for(; *set; ++set) + if(!strcmp(s, *set)) + return true; + return false; +} + +/* Python Verifier._check_sync_attrs: scope / ordering values of a fence or + * flag op (presence is checked by the required-attrs pass). */ +static void check_sync_attrs(verifier_t* v, const rocke_op_t* op, const char* const* orderings) +{ + const char* scope = rocke_attr_get_str(&op->attrs, "scope"); + const char* ordering = rocke_attr_get_str(&op->attrs, "ordering"); + if(scope && !in_set(scope, k_fence_scopes)) + v_errf(v, op, "%s: unknown scope '%s'", op->name, scope); + if(orderings && ordering && !in_set(ordering, orderings)) + v_errf(v, op, "%s: unknown ordering '%s'", op->name, ordering); +} + static void check_contract(verifier_t* v, const rocke_op_t* op) { const char* tf32_error = rocke_tf32_op_error(op); @@ -474,6 +508,45 @@ static void check_contract(verifier_t* v, const rocke_op_t* op) v_errf(v, op, "%s: must have no operands and no results", name); } } + else if(o == ROCKE_OP_MEMREF_FENCE) + { + if(op->num_operands || op->num_results) + { + v_errf(v, op, "%s: must have no operands and no results", name); + } + check_sync_attrs(v, op, k_fence_orderings); + } + else if(o == ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE || o == ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ) + { + if(op->num_operands != 3 || op->num_results) + { + v_errf(v, + op, + "%s: expected 3 operands / 0 results, got %d / %d", + name, + op->num_operands, + op->num_results); + return; + } + const rocke_type_t* pt = op->operands[0]->type; + /* ir._is_flag_ptr: both lowerers address the flag as an i32. */ + if(!pt || pt->kind != ROCKE_TYPE_PTR || !pt->space || strcmp(pt->space, "global") + || strcmp(tn(pt->pointee), "i32")) + { + v_errf(v, op, "%s: expected a global i32 pointer, got %s", name, tn(pt)); + } + const char* what[] = {"index", "value"}; + for(int i = 0; i < 2; ++i) + { + const rocke_type_t* t = op->operands[i + 1]->type; + if(strcmp(tn(t), "i32")) + { + v_errf(v, op, "%s: %s must be i32, got %s", name, what[i], tn(t)); + } + } + check_sync_attrs( + v, op, o == ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE ? k_flag_store_orderings : NULL); + } } /* ---- scf.for ---- */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp index 60abbab8d85b..3f5011d01cf6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp @@ -138,3 +138,173 @@ rocke_streamk_decoded_tile_t rocke_emit_streamk_decode(rocke_ir_builder_t* b, return res; } + +/* ------------------------------------------------------------------------ + * Iteration-balanced stream-K (Python StreamKIterPartition and friends). + * ------------------------------------------------------------------------ */ +static int sk_ceil_div(int a, int c) +{ + return (a + c - 1) / c; +} + +rocke_streamk_iter_partition_t rocke_streamk_iter_partition_make(int m_tiles, + int n_tiles, + int iters_per_tile, + int max_active_wgs, + bool persistent, + rocke_status_t* out_status) +{ + rocke_streamk_iter_partition_t p; + p.m_tiles = 0; + p.n_tiles = 0; + p.iters_per_tile = 0; + p.max_active_wgs = 0; + p.persistent = false; + if(m_tiles <= 0 || n_tiles <= 0 || iters_per_tile <= 0 || max_active_wgs <= 0) + { + if(out_status) + *out_status = ROCKE_ERR_VALUE; + return p; + } + p.m_tiles = m_tiles; + p.n_tiles = n_tiles; + p.iters_per_tile = iters_per_tile; + p.max_active_wgs = max_active_wgs; + p.persistent = persistent; + if(out_status) + *out_status = ROCKE_OK; + return p; +} + +rocke_streamk_iter_plan_t rocke_streamk_iter_plan(const rocke_streamk_iter_partition_t* p) +{ + rocke_streamk_iter_plan_t r; + int rem; + int contributors; + int rest; + r.num_tiles = p->m_tiles * p->n_tiles; + /* sk_tiles: CK "DP + 2-tile SK", reverting to all-DP when the remainder + * cannot give every stream-K CTA one iteration. */ + rem = p->max_active_wgs > 0 ? r.num_tiles % p->max_active_wgs : 0; + if(rem == 0) + r.sk_tiles = 0; + else + { + r.sk_tiles = r.num_tiles > p->max_active_wgs ? p->max_active_wgs + rem : r.num_tiles; + if(r.sk_tiles * p->iters_per_tile < p->max_active_wgs) + r.sk_tiles = 0; + } + r.sk_ctas = r.sk_tiles ? p->max_active_wgs : 0; + r.total_sk_iters = r.sk_tiles * p->iters_per_tile; + /* Clamped to 1 for an all-DP plan: kernels divide by it unconditionally. */ + r.iters_per_sk_cta = r.sk_ctas ? r.total_sk_iters / r.sk_ctas : 1; + r.extra_iters = r.sk_ctas ? r.total_sk_iters % r.sk_ctas : 0; + r.dp_tiles = r.num_tiles - r.sk_tiles; + r.total_dp_iters = r.dp_tiles * p->iters_per_tile; + r.grid_size = p->persistent ? p->max_active_wgs : r.dp_tiles + r.sk_ctas; + if(!r.sk_ctas) + r.max_linear_partners = 0; + else + r.max_linear_partners = sk_ceil_div(p->iters_per_tile - 1, r.iters_per_sk_cta); + /* tree_rounds = (contributors - 1).bit_length() */ + contributors = r.max_linear_partners + 1; + r.tree_rounds = 0; + for(rest = contributors - 1; rest > 0; rest >>= 1) + ++r.tree_rounds; + return r; +} + +/* Python StreamKIterPartition.flags_bytes(n_waves): one i32 flag per + * (stream-K CTA, wave), padded to 256 bytes. */ +size_t rocke_streamk_flags_bytes(const rocke_streamk_iter_plan_t* r, + int n_waves, + rocke_status_t* out_status) +{ + if(r == NULL || n_waves <= 0) + { + if(out_status) + *out_status = ROCKE_ERR_VALUE; + return 0; + } + if(out_status) + *out_status = ROCKE_OK; + const size_t flags = (size_t)4 * (size_t)r->sk_ctas * (size_t)n_waves; + return (flags + 255) / 256 * 256; +} + +int rocke_streamk_start_iter(const rocke_streamk_iter_partition_t* p, int sk_cta) +{ + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(p); + return r.total_dp_iters + sk_cta * r.iters_per_sk_cta + + (sk_cta < r.extra_iters ? sk_cta : r.extra_iters); +} + +int rocke_streamk_end_iter(const rocke_streamk_iter_partition_t* p, int sk_cta) +{ + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(p); + return rocke_streamk_start_iter(p, sk_cta) + r.iters_per_sk_cta + + (sk_cta < r.extra_iters ? 1 : 0); +} + +/* Python emit_streamk_sk_start_iter: + * c_dp_iters = b.const_i32(part.total_dp_iters) + * c_per_cta = b.const_i32(part.iters_per_sk_cta) + * c_extra = b.const_i32(part.extra_iters) + * body = b.mul(sk_cta, c_per_cta) + * lead = b.smin(sk_cta, c_extra) + * return b.add(c_dp_iters, b.add(body, lead)) + */ +rocke_value_t* rocke_emit_streamk_sk_start_iter(rocke_ir_builder_t* b, + rocke_value_t* sk_cta, + const rocke_streamk_iter_args_t* a) +{ + rocke_value_t* body = rocke_b_mul(b, sk_cta, a->iters_per_sk_cta); + rocke_value_t* lead = rocke_b_smin(b, sk_cta, a->extra_iters); + rocke_value_t* inner = rocke_b_add(b, body, lead); + return rocke_b_add(b, a->total_dp_iters, inner); +} + +/* Python emit_streamk_iter_range; every nested builder call is pinned to a + * temporary in the Python evaluation order. */ +rocke_streamk_iter_range_t rocke_emit_streamk_iter_range(rocke_ir_builder_t* b, + rocke_value_t* cta, + const rocke_streamk_iter_args_t* a, + bool dp) +{ + rocke_streamk_iter_range_t res; + rocke_value_t* sk_cta; + rocke_value_t* sk_start; + rocke_value_t* c_one; + rocke_value_t* c_zero; + rocke_value_t* has_extra; + rocke_value_t* extra; + rocke_value_t* sk_len; + rocke_value_t* start; + rocke_value_t* length; + if(dp) + sk_cta = rocke_b_sub(b, cta, a->dp_tiles); + else + sk_cta = cta; + sk_start = rocke_emit_streamk_sk_start_iter(b, sk_cta, a); + c_one = rocke_b_const_i32(b, 1); + c_zero = rocke_b_const_i32(b, 0); + has_extra = rocke_b_cmp_lt(b, sk_cta, a->extra_iters); + extra = rocke_b_select(b, has_extra, c_one, c_zero); + sk_len = rocke_b_add(b, a->iters_per_sk_cta, extra); + if(dp) + { + rocke_value_t* is_dp = rocke_b_cmp_lt(b, cta, a->dp_tiles); + rocke_value_t* dp_start = rocke_b_mul(b, cta, a->iters_per_tile); + start = rocke_b_select(b, is_dp, dp_start, sk_start); + length = rocke_b_select(b, is_dp, a->iters_per_tile, sk_len); + } + else + { + start = sk_start; + length = sk_len; + } + res.start = rocke_b_to_sgpr_u32(b, start); + res.end = rocke_b_to_sgpr_u32(b, rocke_b_add(b, res.start, length)); + res.sk_cta = sk_cta; + return res; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h index 9a43f6ef2587..71d29644048f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h @@ -141,6 +141,104 @@ rocke_streamk_decoded_tile_t rocke_emit_streamk_decode(rocke_ir_builder_t* b, rocke_value_t* linear_id, const rocke_streamk_partition_t* spec); +/* ------------------------------------------------------------------ * + * Iteration-balanced stream-K (CK Tile StreamKTilePartitioner) + * ------------------------------------------------------------------ * + * Python C99 (this header) + * ----------------------------------- ----------------------------------- + * class StreamKIterPartition rocke_streamk_iter_partition_t + * (its @property values) + rocke_streamk_iter_plan_t, resolved + * once by rocke_streamk_iter_plan() + * streamk_iter_partition(...) rocke_streamk_iter_partition_make(...) + * StreamKIterPartition.flags_bytes() rocke_streamk_flags_bytes(...) + * streamk_start_iter / streamk_end_iter rocke_streamk_start_iter / _end_iter + * StreamKIterArgs rocke_streamk_iter_args_t + * emit_streamk_sk_start_iter(...) rocke_emit_streamk_sk_start_iter(...) + * emit_streamk_iter_range(...) rocke_emit_streamk_iter_range(...) + * + * The derived quantities are returned as one resolved plan by value rather + * than as per-property accessors, so the C side cannot recompute any of them + * with different rounding than the Python @property it mirrors. + */ +typedef struct rocke_streamk_iter_partition +{ + int m_tiles; + int n_tiles; + int iters_per_tile; + int max_active_wgs; + bool persistent; +} rocke_streamk_iter_partition_t; + +typedef struct rocke_streamk_iter_plan +{ + int num_tiles; + int sk_tiles; + int sk_ctas; + int total_sk_iters; + int iters_per_sk_cta; /* >= 1; clamped to 1 for an all-DP plan (a divisor) */ + int extra_iters; + int dp_tiles; + int total_dp_iters; + int grid_size; + int max_linear_partners; + int tree_rounds; +} rocke_streamk_iter_plan_t; + +/* streamk_iter_partition: every count must be > 0. On violation *out_status + * (if non-NULL) is ROCKE_ERR_VALUE and the returned partition is zeroed. */ +rocke_streamk_iter_partition_t rocke_streamk_iter_partition_make(int m_tiles, + int n_tiles, + int iters_per_tile, + int max_active_wgs, + bool persistent, + rocke_status_t* out_status); + +/* Every StreamKIterPartition @property, resolved once. */ +rocke_streamk_iter_plan_t rocke_streamk_iter_plan(const rocke_streamk_iter_partition_t* p); + +/* StreamKIterPartition.flags_bytes(n_waves): bytes of the stream-K fixup flag + * table, one i32 flag per (stream-K CTA, wave) padded to 256 bytes; 0 for an + * all-DP plan. n_waves is the waves per CTA that signal independently + * (warp_m * warp_n for wgrad). n_waves <= 0 (Python ValueError) sets + * *out_status (if non-NULL) to ROCKE_ERR_VALUE and returns 0. */ +size_t rocke_streamk_flags_bytes(const rocke_streamk_iter_plan_t* r, + int n_waves, + rocke_status_t* out_status); + +int rocke_streamk_start_iter(const rocke_streamk_iter_partition_t* p, int sk_cta); +int rocke_streamk_end_iter(const rocke_streamk_iter_partition_t* p, int sk_cta); + +/* The runtime stream-K partition a shape-agnostic kernel reads (Python + * StreamKIterArgs). The host computes a partition per launch and passes these + * as kernel arguments; the emitters below only see them as SSA values. */ +typedef struct rocke_streamk_iter_args +{ + rocke_value_t* iters_per_tile; /* i32 MAC iterations per output tile */ + rocke_value_t* dp_tiles; /* i32 tiles owned whole by data-parallel CTAs */ + rocke_value_t* total_dp_iters; /* i32 dp_tiles * iters_per_tile */ + rocke_value_t* iters_per_sk_cta; /* i32 base iterations per stream-K CTA (>= 1) */ + rocke_value_t* extra_iters; /* i32 stream-K CTAs that own one more */ +} rocke_streamk_iter_args_t; + +/* emit_streamk_sk_start_iter: SSA start iteration of a runtime SK CTA index. */ +rocke_value_t* rocke_emit_streamk_sk_start_iter(rocke_ir_builder_t* b, + rocke_value_t* sk_cta, + const rocke_streamk_iter_args_t* a); + +/* emit_streamk_iter_range result (Python _IterRange). */ +typedef struct rocke_streamk_iter_range +{ + rocke_value_t* start; /* i32 global first iteration (SGPR) */ + rocke_value_t* end; /* i32 one past the last iteration (SGPR) */ + rocke_value_t* sk_cta; /* i32 stream-K CTA index; meaningless for a DP CTA */ +} rocke_streamk_iter_range_t; + +/* emit_streamk_iter_range(b, cta, args, *, dp). */ +rocke_streamk_iter_range_t rocke_emit_streamk_iter_range(rocke_ir_builder_t* b, + rocke_value_t* cta, + const rocke_streamk_iter_args_t* a, + bool dp); + #ifdef __cplusplus } /* extern "C" */ #endif 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 f00fff47da2d..21e7b3fd2747 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 @@ -68,6 +68,18 @@ bool rocke_conv_arg_names(const char* direction, bool two_stage, rocke_conv_arg_list_t* out); +/* As above, plus wgrad's stream-K variant: `streamk` names the reduction + * ("atomic", "workspace", "linear" or "tree") of a stream-K kernel, or is NULL + * for a kernel without stream-K. A stream-K list carries the reduction's + * buffers after ks_count and then the partition block. Returns false for a + * non-wgrad direction, an unknown reduction, or stream-K combined with + * two_stage -- the cases conv_arg_names() raises on. */ +bool rocke_conv_arg_names_ex(const char* direction, + bool is_3d, + bool two_stage, + const char* streamk, + rocke_conv_arg_list_t* out); + /* Direct grouped conv ABI ("fwd", "dgrad" or "wgrad"); false for any other * direction. */ bool rocke_conv_direct_arg_names(const char* direction, rocke_conv_arg_list_t* out); 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 fb851dd7dd03..5524f8e4b66f 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 @@ -336,6 +336,13 @@ typedef struct rocke_conv_build_ctx rocke_value_t* group_idx; rocke_value_t* k_out_group_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 + * place of block_id_z (Python replica_sel). */ + rocke_value_t* wgrad_c_kpg; + rocke_value_t* wgrad_replica_sel; + /* ---- LDS plan + smem handles ---- */ rocke_conv_lds_layout_t lds_layout; /* spec.effective_lds_layout() */ bool double_buffer; /* compv4 || async_dma || unroll_k */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h index 3c5359ed7e2d..f28e23a747a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h @@ -63,6 +63,7 @@ #define ROCKE_WGRAD_LDS_K_OUTER_ARCH "gfx950" #include +#include "rocke/helper_rocke.helpers.streamk.h" /* rocke_streamk_iter_partition_t */ #include "rocke/helper_rocke.instances.common.conv_implicit_gemm.h" /* rocke_conv_problem_t */ #include "rocke/ir.h" #include "rocke/lower_llvm.h" @@ -176,6 +177,27 @@ typedef struct rocke_implicit_gemm_conv_wgrad_spec * cut that R-fold and Stage 2 folds them back with a fixed unrolled add. * Scratch size is groups * R * wg_M * wg_N * 4 bytes. */ int ws_replicas; /* default 8 */ + + /* Stream-K execution (CK Tile StreamKTilePartitioner). Mirrors + * WgradConvSpec.streamk / streamk_reduction / streamk_ctas: + * streamk "off" | "dp_sk" | "persistent" + * streamk_reduction "linear" | "tree" | "atomic" | "workspace" + * streamk_ctas CTA pool the stream-K remainder is spread over; + * a launch parameter, not part of the kernel, so the + * builder accepts -1 (auto) (unlike split_k=-1); the + * rocke_wgrad_conv_streamk_* host helpers resolve -1 + * to the per-arch default pool of the arch they are + * given (rocke_wgrad_conv_streamk_default_pool). + * "off" keeps every existing config byte-identical. Stream-K folds the + * conv groups into GEMM-M, so it is the one path on which this port + * accepts groups > 1. */ + const char* streamk; /* default "off" */ + const char* streamk_reduction; /* default "linear" */ + int streamk_ctas; /* default -1 */ + + /* WgradConvSpec.group_merge. Only 1 is ported: stream-K folds the + * groups into GEMM-M itself, and the grid-per-group path is Python-only. */ + int group_merge; /* default 1 */ } rocke_implicit_gemm_conv_wgrad_spec_t; /* Default-constructed spec (every field == Python dataclass default). */ @@ -212,6 +234,10 @@ int rocke_wgrad_conv_spec_wg_N(const rocke_implicit_gemm_conv_wgrad_spec_t* s); * Stage 2's ordered fold over those slabs cannot un-reorder * sums that were already reordered inside one. * + * streamk "linear"/"tree" -> true. The owner of each tile folds the other + * contributors' partials in a fixed order and stores once. + * streamk "atomic"/"workspace" -> false, for the split-K reasons above. + * * So a deterministic wgrad is still available -- ask for split_k <= 1 -- but a * split-K wgrad, two-stage or not, is not one. Hosts that need bit-exactness * should gate on this predicate rather than on two_stage. */ @@ -224,6 +250,53 @@ bool rocke_wgrad_conv_spec_is_deterministic(const rocke_implicit_gemm_conv_wgrad * Analogous to rocke_streamk_gemm_workspace_bytes / rocke_moe_fused_workspace_bytes. */ size_t rocke_wgrad_conv_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s); +/* ---- stream-K host-side helpers (Python wgrad_streamk_*) ---- */ + +/* Python wgrad_streamk_default_pool(reduction, arch): the CTA pool an auto + * (streamk_ctas == -1) launch uses -- one CTA per CU for linear/tree, four per + * CU for atomic/workspace. CU counts come from a static per-arch table that + * mirrors rocke.helpers.split_k._ARCH_NUM_CUS (256 for an unknown arch). + * A NULL arch is gfx950, Python's default. */ +int rocke_wgrad_conv_streamk_default_pool(const char* reduction, const char* arch); + +/* Python wgrad_streamk_launch(spec, arch=arch).pool: streamk_ctas when > 0, + * the arch default pool when -1, and 0 for streamk == "off" or an invalid + * count. A NULL arch is gfx950. */ +int rocke_wgrad_conv_streamk_ctas(const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch); + +/* The helpers below take the target arch because an auto pool + * (streamk_ctas == -1) resolves to a per-arch CU count; every one of them + * resolves it the same way, so the grid, the workspace and the kernarg + * partition block of one launch always agree. A NULL arch is gfx950, the + * default of the Python twins (wgrad_streamk_*(spec, arch="gfx950")). */ + +/* Python wgrad_streamk_partition(spec, arch=arch): the iteration partition the + * kernel and its launch grid share. Requires a non-NULL spec with + * streamk != "off" and streamk_ctas == -1 (auto) or > 0; otherwise + * *out_status (if non-NULL) is ROCKE_ERR_VALUE and the result is zeroed. */ +rocke_streamk_iter_partition_t rocke_wgrad_conv_streamk_partition( + const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch, rocke_status_t* out_status); + +/* Python wgrad_streamk_grid(spec, arch=arch): one-dimensional launch grid. + * ROCKE_ERR_VALUE when the partition is invalid or out is NULL. */ +rocke_status_t rocke_wgrad_conv_streamk_grid(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, + int out[3]); + +/* Python wgrad_streamk_workspace_layout(spec, arch=arch): (flags_bytes, + * partials_bytes) of a linear/tree kernel; both 0 for the atomic and workspace + * reductions. ROCKE_ERR_VALUE for a NULL argument or an invalid partition. */ +rocke_status_t + rocke_wgrad_conv_streamk_workspace_layout(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, + size_t* flags_bytes, + size_t* partials_bytes); + +/* Python wgrad_streamk_workspace_nbytes(spec, arch=arch); 0 for a NULL spec, + * streamk == "off", or an invalid partition. */ +size_t rocke_wgrad_conv_streamk_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch); + /* spec.wg_K: output spatial positions (N * Ho * Wo [* Do]). */ int rocke_wgrad_conv_spec_wg_K(const rocke_implicit_gemm_conv_wgrad_spec_t* s); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h index 6ff0f81245f2..ad64e60191fb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h @@ -423,6 +423,12 @@ typedef enum rocke_opcode ROCKE_OP_SCF_YIELD, ROCKE_OP_CF_RETURN, + /* memref.* cross-workgroup synchronisation (stream-K fixup). Appended + * after the control-flow block so no existing opcode index shifts. */ + ROCKE_OP_MEMREF_FENCE, + ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, + ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, + ROCKE_OP__COUNT } rocke_opcode_t; @@ -831,6 +837,27 @@ rocke_value_t* rocke_b_cvt_pk_fp8_f32x4(rocke_ir_builder_t* b, rocke_value_t* v) rocke_value_t* rocke_b_cvt_pk_bf8_f32x4(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_pk_i8_f32x4(rocke_ir_builder_t* b, rocke_value_t* v); +/* ----- cross-workgroup synchronisation (stream-K fixup) ----- */ +/* IRBuilder.fence: scope in {workgroup, agent, system}; ordering in + * {acquire, release, acq_rel, seq_cst}. NULL selects the Python defaults + * (agent / acq_rel). */ +void rocke_b_fence(rocke_ir_builder_t* b, const char* scope, const char* ordering); +/* IRBuilder.global_flag_store: atomic i32 ptr[idx] = value. NULL scope / + * ordering select agent / release. */ +void rocke_b_global_flag_store(rocke_ir_builder_t* b, + rocke_value_t* ptr, + rocke_value_t* idx, + rocke_value_t* value, + const char* scope, + const char* ordering); +/* IRBuilder.global_flag_wait_eq: spin until ptr[idx] == expect (acquire). + * NULL scope selects agent. */ +void rocke_b_global_flag_wait_eq(rocke_ir_builder_t* b, + rocke_value_t* ptr, + rocke_value_t* idx, + rocke_value_t* expect, + const char* scope); + /* ----- atomics ----- */ rocke_value_t* rocke_b_global_atomic_add(rocke_ir_builder_t* b, rocke_value_t* ptr, 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 2ddd1c83cf31..76ae1370edcb 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 @@ -152,7 +152,28 @@ void fwd_arg_names(rocke_conv_arg_list_t* out, bool is_3d) grid(out); } -void wgrad_arg_names(rocke_conv_arg_list_t* out, bool is_3d, bool two_stage) +/* The stream-K launch block: the host-computed iteration partition. The + * kernel is shape- and pool-agnostic, so every value the partition derives + * from the launch shape and the CTA pool arrives here. Mirrors the Python + * _STREAMK_SCALARS, in order. */ +const char* const kStreamKScalars[] = { + "sk_iters_per_tile", + "sk_dp_tiles", + "sk_ctas", + "sk_iters_per_cta", + "sk_extra_iters", + "sk_pool", + "sk_tree_rounds", +}; +const int kNumStreamKScalars = (int)(sizeof(kStreamKScalars) / sizeof(kStreamKScalars[0])); + +bool streamk_reduction_known(const char* red) +{ + return strcmp(red, "atomic") == 0 || strcmp(red, "workspace") == 0 || strcmp(red, "linear") == 0 + || strcmp(red, "tree") == 0; +} + +void wgrad_arg_names(rocke_conv_arg_list_t* out, bool is_3d, bool two_stage, const char* streamk) { leading(out, "dY", "X", "dW"); dims(out, is_3d); @@ -180,7 +201,7 @@ void wgrad_arg_names(rocke_conv_arg_list_t* out, bool is_3d, bool two_stage) magic_2d3d(out, "n_", is_3d, "Y", "X", "cpg"); grid(out); /* ---- variant-specific extras ---- */ - if(two_stage) + if(two_stage || (streamk != NULL && strcmp(streamk, "workspace") == 0)) { push(out, "ws_ptr", ROCKE_CONV_ARG_F32_PTR); push_i32(out, "ws_bytes"); @@ -188,6 +209,16 @@ void wgrad_arg_names(rocke_conv_arg_list_t* out, bool is_3d, bool two_stage) /* The split-K slice width and degree, always. */ push_i32(out, "ks"); push_i32(out, "ks_count"); + if(streamk != NULL) + { + if(strcmp(streamk, "linear") == 0 || strcmp(streamk, "tree") == 0) + { + push(out, "sk_flags", ROCKE_CONV_ARG_I32_PTR); + push(out, "sk_partials", ROCKE_CONV_ARG_F32_PTR); + } + for(int i = 0; i < kNumStreamKScalars; ++i) + push_i32(out, kStreamKScalars[i]); + } } void dgrad_arg_names(rocke_conv_arg_list_t* out, bool is_3d) @@ -248,18 +279,34 @@ bool rocke_conv_arg_names(const char* direction, bool is_3d, bool two_stage, rocke_conv_arg_list_t* out) +{ + return rocke_conv_arg_names_ex(direction, is_3d, two_stage, NULL, out); +} + +bool rocke_conv_arg_names_ex(const char* direction, + bool is_3d, + bool two_stage, + const char* streamk, + rocke_conv_arg_list_t* out) { 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. */ + * 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) return false; + if(streamk != NULL) + { + if(strcmp(direction, "wgrad") != 0 || !streamk_reduction_known(streamk) || two_stage) + return false; + } if(strcmp(direction, "fwd") == 0) fwd_arg_names(out, is_3d); else if(strcmp(direction, "wgrad") == 0) - wgrad_arg_names(out, is_3d, two_stage); + wgrad_arg_names(out, is_3d, two_stage, streamk); else if(strcmp(direction, "dgrad") == 0) dgrad_arg_names(out, is_3d); else diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index d8fb21204d52..ee673f2eb732 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -36,7 +36,7 @@ #include /* strcmp, memset, memcpy */ #include "rocke/error_boundary.hpp" -#include "rocke/helper_rocke.core.arch.h" /* rocke_archtarget_async_lds_max_dwords */ +#include "rocke/helper_rocke.core.arch.h" /* rocke_archtarget_* */ #include "rocke/helper_rocke.helpers.spec.h" #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/instance_conv_abi.h" @@ -75,9 +75,24 @@ rocke_implicit_gemm_conv_wgrad_spec_t rocke_implicit_gemm_conv_wgrad_spec_defaul s.chiplet_chunk_size = 64; s.split_k = 1; s.ws_replicas = ROCKE_WGRAD_DEFAULT_WS_REPLICAS; + s.streamk = "off"; + s.streamk_reduction = "linear"; + s.streamk_ctas = -1; + s.group_merge = 1; return s; } +/* Python: spec.streamk != "off". A NULL field is the default. */ +static bool wgrad_streamk_on(const rocke_implicit_gemm_conv_wgrad_spec_t* s) +{ + return s->streamk != NULL && strcmp(s->streamk, "off") != 0; +} + +static const char* wgrad_streamk_reduction(const rocke_implicit_gemm_conv_wgrad_spec_t* s) +{ + return s->streamk_reduction ? s->streamk_reduction : "linear"; +} + int rocke_wgrad_conv_spec_block_size(const rocke_implicit_gemm_conv_wgrad_spec_t* s) { return s->warp_m * s->warp_n * s->wave_size; @@ -98,18 +113,18 @@ int rocke_wgrad_conv_spec_mfmas_per_warp_n(const rocke_implicit_gemm_conv_wgrad_ return s->tile_n / (s->warp_n * s->warp_tile_n); } -/* wg_M = K (output channels, groups=1 always for wgrad) */ +/* wg_M = kpg (output channels per group; K when groups == 1) */ int rocke_wgrad_conv_spec_wg_M(const rocke_implicit_gemm_conv_wgrad_spec_t* s) { - return s->problem.K; + return rocke_conv_problem_kpg(&s->problem); } -/* wg_N = Z*Y*X*C (filter spatial x input channel) */ +/* wg_N = Z*Y*X*cpg (filter spatial x input channels per group) */ int rocke_wgrad_conv_spec_wg_N(const rocke_implicit_gemm_conv_wgrad_spec_t* s) { const rocke_conv_problem_t* p = &s->problem; int z = p->is_3d ? p->Z : 1; - return z * p->Y * p->X * p->C; + return z * p->Y * p->X * rocke_conv_problem_cpg(p); } /* wg_K = N*Ho*Wo (output spatial positions) */ @@ -136,10 +151,151 @@ bool rocke_wgrad_conv_spec_is_deterministic(const rocke_implicit_gemm_conv_wgrad * f32-atomic-adds into ws_replicas shared slabs, so the order in which a * group's slices land in a slab is scheduler-dependent and f32 addition * is not associative. Stage 2's fold over the replicas is ordered, which - * does nothing for partial sums that were already reordered. */ + * does nothing for partial sums that were already reordered. + * streamk: the linear and tree fixups fold the partials of a tile in a + * fixed order, so they are reproducible; atomic and workspace are the + * split-K cases above by another name. */ + if(wgrad_streamk_on(s)) + { + const char* red = wgrad_streamk_reduction(s); + return strcmp(red, "linear") == 0 || strcmp(red, "tree") == 0; + } return s->split_k <= 1; } +/* Per-arch CU counts. Mirrors the static table rocke.helpers.split_k + * _ARCH_NUM_CUS / _DEFAULT_NUM_CUS that Python wgrad_streamk_default_pool + * reads; test_conv_wgrad_streamk.py checks the two agree key by key. */ +static int wgrad_arch_num_cus(const char* arch) +{ + static const struct + { + const char* gfx; + int num_cus; + } table[] = { + {"gfx942", 304}, /* MI300X: 8 XCDs * 38 CUs */ + {"gfx950", 256}, /* MI350X: 8 XCDs * 32 CUs */ + {"gfx1151", 64}, /* RDNA gfx1151 (WMMA) */ + }; + if(arch != NULL) + for(const auto& e : table) + if(strcmp(arch, e.gfx) == 0) + return e.num_cus; + return 256; /* _DEFAULT_NUM_CUS */ +} + +/* Python _STREAMK_CTAS_PER_CU: CTAs per CU in the default pool of the + * reductions that never wait (atomic, workspace). */ +#define ROCKE_WGRAD_STREAMK_CTAS_PER_CU 4 + +/* Python wgrad_streamk_default_pool. */ +int rocke_wgrad_conv_streamk_default_pool(const char* reduction, const char* arch) +{ + const int num_cus = wgrad_arch_num_cus(arch ? arch : "gfx950"); + const char* red = reduction ? reduction : "linear"; + if(strcmp(red, "linear") == 0 || strcmp(red, "tree") == 0) + return num_cus; + return num_cus * ROCKE_WGRAD_STREAMK_CTAS_PER_CU; +} + +/* Python wgrad_streamk_launch(spec, arch=arch).pool. */ +int rocke_wgrad_conv_streamk_ctas(const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch) +{ + if(s == NULL || !wgrad_streamk_on(s)) + return 0; + if(s->streamk_ctas > 0) + return s->streamk_ctas; + if(s->streamk_ctas == -1) + return rocke_wgrad_conv_streamk_default_pool(wgrad_streamk_reduction(s), arch); + return 0; +} + +/* Python wgrad_streamk_partition (a NULL arch is gfx950, Python's default). */ +rocke_streamk_iter_partition_t rocke_wgrad_conv_streamk_partition( + const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch, rocke_status_t* out_status) +{ + rocke_streamk_iter_partition_t zero; + memset(&zero, 0, sizeof(zero)); + const int pool = rocke_wgrad_conv_streamk_ctas(s, arch); + if(pool <= 0) + { + if(out_status) + *out_status = ROCKE_ERR_VALUE; + return zero; + } + const int groups = s->problem.groups > 0 ? s->problem.groups : 1; + const int wg_M = rocke_wgrad_conv_spec_wg_M(s); + const int wg_N = rocke_wgrad_conv_spec_wg_N(s); + const int wg_K = rocke_wgrad_conv_spec_wg_K(s); + return rocke_streamk_iter_partition_make((wg_M + s->tile_m - 1) / s->tile_m * groups, + (wg_N + s->tile_n - 1) / s->tile_n, + (wg_K + s->tile_k - 1) / s->tile_k, + pool, + strcmp(s->streamk, "persistent") == 0, + out_status); +} + +rocke_status_t rocke_wgrad_conv_streamk_grid(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, + int out[3]) +{ + rocke_status_t st = ROCKE_OK; + rocke_streamk_iter_partition_t part = rocke_wgrad_conv_streamk_partition(s, arch, &st); + if(st != ROCKE_OK || out == NULL) + return ROCKE_ERR_VALUE; + out[0] = rocke_streamk_iter_plan(&part).grid_size; + out[1] = 1; + out[2] = 1; + return ROCKE_OK; +} + +rocke_status_t + rocke_wgrad_conv_streamk_workspace_layout(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, + size_t* flags_bytes, + size_t* partials_bytes) +{ + if(s == NULL || flags_bytes == NULL || partials_bytes == NULL) + return ROCKE_ERR_VALUE; + *flags_bytes = 0; + *partials_bytes = 0; + const char* red = wgrad_streamk_reduction(s); + if(strcmp(red, "linear") != 0 && strcmp(red, "tree") != 0) + return ROCKE_OK; + rocke_status_t st = ROCKE_OK; + rocke_streamk_iter_partition_t part = rocke_wgrad_conv_streamk_partition(s, arch, &st); + if(st != ROCKE_OK) + return st; + const rocke_streamk_iter_plan_t plan = rocke_streamk_iter_plan(&part); + *flags_bytes = rocke_streamk_flags_bytes(&plan, s->warp_m * s->warp_n, &st); + if(st != ROCKE_OK) + return st; + *partials_bytes = (size_t)plan.sk_ctas * (size_t)s->tile_m * (size_t)s->tile_n * sizeof(float); + return ROCKE_OK; +} + +size_t rocke_wgrad_conv_streamk_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch) +{ + if(s == NULL || !wgrad_streamk_on(s)) + return 0; + const char* red = wgrad_streamk_reduction(s); + if(strcmp(red, "atomic") == 0) + return 0; + if(strcmp(red, "workspace") == 0) + { + const int groups = s->problem.groups > 0 ? s->problem.groups : 1; + return (size_t)groups * (size_t)s->ws_replicas * (size_t)rocke_wgrad_conv_spec_wg_M(s) + * (size_t)rocke_wgrad_conv_spec_wg_N(s) * sizeof(float); + } + size_t flags_bytes = 0; + size_t partials_bytes = 0; + if(rocke_wgrad_conv_streamk_workspace_layout(s, arch, &flags_bytes, &partials_bytes) + != ROCKE_OK) + return 0; + return flags_bytes + partials_bytes; +} + size_t rocke_wgrad_conv_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s) { if(!s->two_stage) @@ -217,8 +373,8 @@ rocke_status_t rocke_wgrad_conv_spec_kernel_name(const rocke_implicit_gemm_conv_ /* flags: async, kouter, pad{N}, spk/spkauto, twostage, wsr{N}, unroll */ char pad_flag[32] = {0}; - const char* flag_names[7]; - int flag_on[7]; + const char* flag_names[10]; + int flag_on[10]; int n_flags = 0; flag_names[n_flags] = "async"; @@ -291,6 +447,31 @@ rocke_status_t rocke_wgrad_conv_spec_kernel_name(const rocke_implicit_gemm_conv_ flag_on[n_flags] = s->unroll_k ? 1 : 0; n_flags++; + /* Stream-K tags, in the Python flag order: sk, skpers, skr + * (the last only when not the default "linear"). The CTA pool is a launch + * parameter and so is NOT in the name, same as the split-K degree. */ + char skr_buf[48]; + if(wgrad_streamk_on(s)) + { + const char* red = wgrad_streamk_reduction(s); + flag_names[n_flags] = "sk"; + flag_on[n_flags] = 1; + n_flags++; + if(strcmp(s->streamk, "persistent") == 0) + { + flag_names[n_flags] = "skpers"; + flag_on[n_flags] = 1; + n_flags++; + } + if(strcmp(red, "linear") != 0) + { + snprintf(skr_buf, sizeof(skr_buf), "skr%s", red); + flag_names[n_flags] = skr_buf; + flag_on[n_flags] = 1; + n_flags++; + } + } + return rocke_kernel_name_join( s->name, parts, 5, flag_names, flag_on, n_flags, out, out_cap, NULL); } @@ -299,6 +480,117 @@ rocke_status_t rocke_wgrad_conv_spec_kernel_name(const rocke_implicit_gemm_conv_ // is_valid_wgrad_spec // --------------------------------------------------------------------------- +/* Python is_valid_wgrad_spec's arch gate plus wgrad_streamk_available: the + * stream-K gates, with byte-identical reason strings. */ +static bool wgrad_streamk_gates(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, + char* reason, + size_t reason_cap) +{ +#define SK_REJECT(...) \ + do \ + { \ + if(reason && reason_cap) \ + snprintf(reason, reason_cap, __VA_ARGS__); \ + return false; \ + } while(0) + const char* mode = s->streamk ? s->streamk : "off"; + if(strcmp(mode, "off") != 0 && strcmp(arch, "gfx942") != 0 && strcmp(arch, "gfx950") != 0) + SK_REJECT("stream-K wgrad supports only gfx942, gfx950 (got %s)", arch); + if(strcmp(mode, "off") != 0 && strcmp(mode, "dp_sk") != 0 && strcmp(mode, "persistent") != 0) + SK_REJECT("streamk must be 'off', 'dp_sk' or 'persistent' (got '%s')", mode); + if(strcmp(mode, "off") == 0) + return true; + const char* red = wgrad_streamk_reduction(s); + if(strcmp(red, "linear") != 0 && strcmp(red, "tree") != 0 && strcmp(red, "atomic") != 0 + && strcmp(red, "workspace") != 0) + SK_REJECT("streamk_reduction must be 'linear', 'tree', 'atomic' or 'workspace' " + "(got '%s')", + red); + if(s->streamk_ctas == 0 || s->streamk_ctas < -1) + SK_REJECT("streamk_ctas must be -1 (auto) or > 0 (got %d)", s->streamk_ctas); + if(s->wave_size != 64) + SK_REJECT("stream-K wgrad is MFMA wave64 only (got wave_size=%d)", s->wave_size); + if(s->split_k != 1) + SK_REJECT("streamk and split_k both partition K_wg; use split_k=1 with streamk " + "(got split_k=%d)", + s->split_k); + if(s->two_stage) + SK_REJECT("two_stage is the split-K scratch path; use " + "streamk_reduction='workspace' for the stream-K equivalent"); + if(s->async_dma || s->unroll_k) + SK_REJECT("streamk runs one plain K loop per work item; async_dma/unroll_k " + "pipeline the K loop across items"); + const char* epi = s->epilogue ? s->epilogue : "default"; + if(strcmp(epi, "default") != 0) + SK_REJECT("streamk requires epilogue='default' (got '%s'): the epilogue runs once " + "per tile inside the stream-K tile loop", + epi); + if(s->chiplet_swizzle) + SK_REJECT("chiplet_swizzle remaps the 2-D tile grid; stream-K decodes its tiles " + "from a linear CTA index instead"); + if(s->group_merge > 1) + SK_REJECT("streamk folds conv groups into GEMM-M; use group_merge=1 (got %d)", + s->group_merge); + const char* dd = s->dtype_d ? s->dtype_d : "fp16"; + if(strcmp(red, "atomic") == 0 && strcmp(dd, "fp32") != 0) + SK_REJECT("streamk_reduction='atomic' needs dtype_d='fp32' (got '%s'); use " + "'linear', 'tree' or 'workspace'", + dd); + return true; +#undef SK_REJECT +} + +/* Python wgrad_group_merge_available, with byte-identical reason strings. + * Python checks it right after the arch lookup, before any other gate. */ +static bool wgrad_group_merge_available(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + char* reason, + size_t reason_cap) +{ +#define GM_REJECT(...) \ + do \ + { \ + if(reason && reason_cap) \ + snprintf(reason, reason_cap, __VA_ARGS__); \ + return false; \ + } while(0) + const int gm = s->group_merge; + if(gm == 1) + return true; + /* Python _GROUP_MERGE_DEGREES */ + if(gm != 2 && gm != 4 && gm != 8 && gm != 16 && gm != 32 && gm != 64) + GM_REJECT("group_merge must be 1 or one of (2, 4, 8, 16, 32, 64); got %d", gm); + const rocke_conv_problem_t* p = &s->problem; + if(p->groups % gm != 0) + GM_REJECT("group_merge %d does not divide groups %d: a workgroup would own a " + "partial group", + gm, + p->groups); + const int cpg = rocke_conv_problem_cpg(p); + const int kpg = rocke_conv_problem_kpg(p); + if(cpg != 1 || kpg != 1) + GM_REJECT("group_merge is implemented for depthwise only (cpg == kpg == 1); got " + "cpg=%d kpg=%d. The merged index math is correct for wider groups but is " + "not covered, so it is gated off.", + cpg, + kpg); + if(s->wave_size != 64) + GM_REJECT("group_merge is MFMA-only (wave_size 64); got wave_size=%d", s->wave_size); + if(s->split_k > 1 && !s->two_stage) + GM_REJECT("group_merge with split_k=%d needs the two-stage path (two_stage=True); " + "the packed-atomic split-K epilogue cannot drop off-diagonal group pairs", + s->split_k); + /* grid_M / grid_N: the merged GEMM extents, kpg*Gm and Z*Y*X*cpg*Gm. */ + const int grid_M = kpg * gm; + const int grid_N = (p->is_3d ? p->Z : 1) * p->Y * p->X * cpg * gm; + if(grid_M > s->tile_m) + GM_REJECT("merged GEMM-M %d (kpg*%d) exceeds tile_m %d", grid_M, gm, s->tile_m); + if(grid_N > s->tile_n) + GM_REJECT("merged GEMM-N %d (Z*Y*X*cpg*%d) exceeds tile_n %d", grid_N, gm, s->tile_n); + return true; +#undef GM_REJECT +} + bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch, char* reason, @@ -320,17 +612,24 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv if(arch == NULL) arch = "gfx950"; - /* groups > 1 is not supported for wgrad (Python validate() raises on it). */ - if(s->problem.groups != 1) + /* The gates below run in Python is_valid_wgrad_spec's order -- arch lookup, + * group_merge, geometry, block size, vec_c, wave size, stream-K -- so a + * spec that fails several of them is rejected with the same reason by + * both engines. */ + + /* Python: ArchTarget.from_gfx(arch) -> KeyError. The forward validator + * checks the arch first and reproduces str(KeyError) verbatim, so an + * unknown arch is handed to it for the reason text. */ + const rocke_archtarget_t* target = rocke_archtarget_from_gfx(arch); + if(target == NULL) { - if(reason && reason_cap) - snprintf(reason, - reason_cap, - "grouped convolution (groups=%d > 1) is not supported for wgrad", - s->problem.groups); - return false; + rocke_implicit_gemm_conv_spec_t probe = rocke_implicit_gemm_conv_spec_default(); + return rocke_implicit_gemm_conv_is_valid_spec(&probe, arch, reason, reason_cap); } + if(!wgrad_group_merge_available(s, reason, reason_cap)) + return false; + /* geometry */ if(s->tile_m % (s->warp_m * s->warp_tile_m)) { @@ -350,11 +649,72 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv snprintf(reason, reason_cap, "tile_k not divisible by warp_tile_k"); return false; } - int block_size = rocke_wgrad_conv_spec_block_size(s); - if(block_size > 1024) + const int block_size = rocke_wgrad_conv_spec_block_size(s); + const int max_tpb = rocke_archtarget_max_threads_per_block(target); + if(block_size > max_tpb) { if(reason && reason_cap) - snprintf(reason, reason_cap, "block_size %d > 1024", block_size); + snprintf(reason, + reason_cap, + "block_size %d > %d (hardware cap) on %s", + block_size, + max_tpb, + arch); + return false; + } + + /* vector_size_c > 1 needs the cshuffle epilogue. Mirrors Python + * is_valid_wgrad_spec (and the dgrad port's gate). */ + { + const char* epi = s->epilogue ? s->epilogue : "default"; + if(s->has_vector_size_c && s->vector_size_c > 1 && strcmp(epi, "default") == 0) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "default epilogue is not supported with vector size c: %d", + s->vector_size_c); + return false; + } + } + + if(s->wave_size != target->wave_size) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "spec wave_size %d != %s wave_size %d", + s->wave_size, + arch, + target->wave_size); + return false; + } + + if(!wgrad_streamk_gates(s, arch, reason, reason_cap)) + return false; + + /* C-port-only gates, after every gate Python shares: grouped wgrad + * (groups > 1) and group merging are only ported on the stream-K path, + * which folds the group into its linear tile index; the grid-per-group + * path stays Python-only. This is an intentional divergence -- Python + * accepts these specs, the C port rejects them. */ + const bool sk_on = wgrad_streamk_on(s); + if(s->problem.groups != 1 && !sk_on) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "grouped convolution (groups=%d > 1) is not supported for wgrad", + s->problem.groups); + return false; + } + if(s->group_merge != 1 && !sk_on) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "group_merge=%d is not supported for wgrad in the C port", + s->group_merge); return false; } /* Explicit vector widths fit one 16-byte per-lane access (Python: @@ -432,9 +792,8 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv * any grouped conv it disagreed with Python (which tests cpg); * - it did not exempt the two-stage path, which stores f32 to a * workspace and emits no atomic at all. - * The local per-group computation is deliberate: the shared - * rocke_wgrad_conv_spec_wg_N() helper still returns the dense Z*Y*X*C - * and is used for workspace sizing, so it is not interchangeable here. */ + * The product is recomputed locally in 64-bit rather than through + * rocke_wgrad_conv_spec_wg_N() (the same per-group value, as an int). */ const bool effective_two_stage_gate = s->two_stage && sk > 1; const char* dt = s->dtype_d ? s->dtype_d : "fp16"; if(!effective_two_stage_gate && (strcmp(dt, "fp16") == 0 || strcmp(dt, "bf16") == 0)) @@ -678,6 +1037,16 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv } } + /* Python checks this between the MMA atom and the LDS budget, both of + * which the forward validator below owns; checking it just before is the + * closest order the delegation allows. */ + if(s->problem.groups > 1 && rocke_conv_problem_is_pointwise(&s->problem)) + { + if(reason && reason_cap) + snprintf(reason, reason_cap, "grouped pointwise (1x1) wgrad is not yet supported"); + return false; + } + /* Delegate the MMA-atom + LDS + WMMA gates to the forward validator via an * adapter spec. The forward validator only reads: tile_m/n/k, warp_m/n, * warp_tile_m/n/k, wave_size, pipeline, epilogue, async_dma, unroll_k, @@ -805,6 +1174,17 @@ static rocke_tensor_descriptor_t* wgrad_make_dy_descriptor(rocke_ir_builder_t* b * Simple approach: build the DAG manually mirroring make_a_descriptor but * substituting "m"->"k_gemm" and "k"->"k_out" throughout. */ +/* Grouped X channel decode (Python _a_channel_decode, groups > 1): + * embed(["group", "c_in_group"] -> "c", strides=[cpg, 1], offset=0, lo=0, hi=C) + * so the group selects its [g*cpg, (g+1)*cpg) slab of the NHWC channel axis. */ +static const rocke_transform_t* wgrad_group_channel_embed(rocke_ir_builder_t* b, + const rocke_conv_problem_t* p) +{ + const char* up[2] = {"group", "c_in_group"}; + int strides[2] = {rocke_conv_problem_cpg(p), 1}; + return rocke_embed_bounded(b, up, 2, "c", strides, 0, 0, p->C); +} + static rocke_tensor_descriptor_t* wgrad_make_x_descriptor(rocke_ir_builder_t* b, const rocke_conv_problem_t* p) { @@ -866,13 +1246,21 @@ static rocke_tensor_descriptor_t* wgrad_make_x_descriptor(rocke_ir_builder_t* b, return NULL; n_x++; - /* unmerge_magic("k_out" -> [z,y,x,c]) */ - const char* into_k[4] = {"z", "y", "x", "c"}; - int dims_k[4] = {p->Z, p->Y, p->X, p->C}; + /* unmerge_magic("k_out" -> [z,y,x,c]); grouped: Python _a_channel_decode */ + const bool grouped3 = p->groups > 1; + const char* into_k[4] = {"z", "y", "x", grouped3 ? "c_in_group" : "c"}; + int dims_k[4] = {p->Z, p->Y, p->X, grouped3 ? rocke_conv_problem_cpg(p) : p->C}; xforms[n_x] = rocke_unmerge_magic(b, "k_out", into_k, 4, dims_k); if(xforms[n_x] == NULL) return NULL; n_x++; + if(grouped3) + { + xforms[n_x] = wgrad_group_channel_embed(b, p); + if(xforms[n_x] == NULL) + return NULL; + n_x++; + } xforms[n_x] = rocke_pad(b, "z", 0, p->Z); if(xforms[n_x] == NULL) @@ -922,13 +1310,21 @@ static rocke_tensor_descriptor_t* wgrad_make_x_descriptor(rocke_ir_builder_t* b, return NULL; n_x++; - /* unmerge_magic("k_out" -> [y, x, c]) */ - const char* into_k[3] = {"y", "x", "c"}; - int dims_k[3] = {p->Y, p->X, p->C}; + /* unmerge_magic("k_out" -> [y, x, c]); grouped: Python _a_channel_decode */ + const bool grouped = p->groups > 1; + const char* into_k[3] = {"y", "x", grouped ? "c_in_group" : "c"}; + int dims_k[3] = {p->Y, p->X, grouped ? rocke_conv_problem_cpg(p) : p->C}; xforms[n_x] = rocke_unmerge_magic(b, "k_out", into_k, 3, dims_k); if(xforms[n_x] == NULL) return NULL; n_x++; + if(grouped) + { + xforms[n_x] = wgrad_group_channel_embed(b, p); + if(xforms[n_x] == NULL) + return NULL; + n_x++; + } xforms[n_x] = rocke_pad(b, "y", 0, p->Y); if(xforms[n_x] == NULL) @@ -957,17 +1353,19 @@ static rocke_tensor_descriptor_t* wgrad_make_x_descriptor(rocke_ir_builder_t* b, static rocke_tensor_descriptor_t* wgrad_make_dw_descriptor(rocke_ir_builder_t* b, const rocke_conv_problem_t* p) { + /* Grouped dW is packed per group: [K, [Z,] Y, X, cpg] (Python cdim). */ + const int cdim = p->groups > 1 ? rocke_conv_problem_cpg(p) : p->C; if(p->is_3d) { /* naive coords: first dim "m" = K (output channels), rest are spatial */ - int lengths[5] = {p->K, p->Z, p->Y, p->X, p->C}; + int lengths[5] = {p->K, p->Z, p->Y, p->X, cdim}; const char* coords[5] = {"k_out", "z", "y", "x", "c"}; rocke_tensor_descriptor_t* desc = rocke_tensor_descriptor_naive(b, "dW_kzyxc", lengths, 5, NULL, coords, 5); if(desc == NULL) return NULL; const char* into[4] = {"z", "y", "x", "c"}; - int dims[4] = {p->Z, p->Y, p->X, p->C}; + int dims[4] = {p->Z, p->Y, p->X, cdim}; const rocke_transform_t* xforms[4]; int n_x = 0; xforms[n_x] = rocke_unmerge_magic(b, "n_wg", into, 4, dims); @@ -990,14 +1388,14 @@ static rocke_tensor_descriptor_t* wgrad_make_dw_descriptor(rocke_ir_builder_t* b } /* 2-D: naive("dW_kyxc", [K,Y,X,C], coords=["k_out","y","x","c"]) */ - int lengths[4] = {p->K, p->Y, p->X, p->C}; + int lengths[4] = {p->K, p->Y, p->X, cdim}; const char* coords[4] = {"k_out", "y", "x", "c"}; rocke_tensor_descriptor_t* desc = rocke_tensor_descriptor_naive(b, "dW_kyxc", lengths, 4, NULL, coords, 4); if(desc == NULL) return NULL; const char* into[3] = {"y", "x", "c"}; - int dims[3] = {p->Y, p->X, p->C}; + int dims[3] = {p->Y, p->X, cdim}; const rocke_transform_t* xforms[4]; int n_x = 0; xforms[n_x] = rocke_unmerge_magic(b, "n_wg", into, 3, dims); @@ -1056,6 +1454,13 @@ static rocke_value_t* wgrad_dy_descriptor(rocke_ir_builder_t* b, rocke_conv_build_ctx_t* ctx = (rocke_conv_build_ctx_t*)ctx_user; /* k_out = block_m_off + row (= output channel, m_val) -- computed FIRST */ rocke_value_t* m_val = rocke_b_add(b, ctx->block_m_off_v, row); + if(ctx->wgrad_c_kpg != NULL) + { + /* Grouped: group g owns the dY output-channel slab [g*kpg, (g+1)*kpg). + * Python: k_out = b_.add(k_out, b_.mul(group_v, c_kpg)) */ + rocke_value_t* g_base = rocke_b_mul(b, ctx->group_idx, ctx->wgrad_c_kpg); + m_val = rocke_b_add(b, m_val, g_base); + } /* k_wg_red = k_off + col (= output position, k_val) -- computed SECOND */ rocke_value_t* k_val = rocke_b_add(b, ctx->k_off_capture, col); @@ -1122,13 +1527,15 @@ static rocke_value_t* wgrad_x_descriptor(rocke_ir_builder_t* b, return off; } - /* Full descriptor path: X_desc.offset(k_gemm=m_val, k_out=k_val) - * X descriptor top-level coords: "k_gemm" (K_wg reduction) and "k_out" (N_wg). */ - const char* names[2] = {"k_gemm", "k_out"}; - rocke_value_t* vals[2] = {m_val, k_val}; + /* Full descriptor path: X_desc.offset(k_gemm=m_val, k_out=k_val[, group=g]) + * X descriptor top-level coords: "k_gemm" (K_wg reduction) and "k_out" (N_wg); + * grouped adds "group", which selects the channel slab. */ + const char* names[3] = {"k_gemm", "k_out", "group"}; + rocke_value_t* vals[3] = {m_val, k_val, ctx->group_idx}; + const int n_coords = ctx->wgrad_c_kpg != NULL ? 3 : 2; rocke_value_t* off = NULL; rocke_value_t* valid = NULL; - rocke_transforms_descriptor_offset(b, ctx->B_desc, names, vals, 2, &off, &valid); + rocke_transforms_descriptor_offset(b, ctx->B_desc, names, vals, n_coords, &off, &valid); if(out_valid) *out_valid = valid; return off; @@ -1243,11 +1650,13 @@ static void wgrad_emit_workspace_store_epilogue(rocke_ir_builder_t* b, const int reps = spec->ws_replicas; rocke_value_t* group_v = ctx->group_idx; rocke_value_t* slab_idx = NULL; + /* Python replica_sel: a stream-K launch is one-dimensional and hashes its + * linear CTA index into the replica instead of block_id_z. */ if(group_v != NULL && reps > 1) { rocke_value_t* c_reps_a = rocke_b_const_i32(b, reps); rocke_value_t* grp_term = rocke_b_mul(b, group_v, c_reps_a); - rocke_value_t* z = rocke_b_block_id_z(b); + rocke_value_t* z = ctx->wgrad_replica_sel ? ctx->wgrad_replica_sel : rocke_b_block_id_z(b); rocke_value_t* c_reps_b = rocke_b_const_i32(b, reps); rocke_value_t* rep_term = rocke_b_mod(b, z, c_reps_b); slab_idx = rocke_b_add(b, grp_term, rep_term); @@ -1258,7 +1667,7 @@ static void wgrad_emit_workspace_store_epilogue(rocke_ir_builder_t* b, } else if(reps > 1) { - rocke_value_t* z = rocke_b_block_id_z(b); + rocke_value_t* z = ctx->wgrad_replica_sel ? ctx->wgrad_replica_sel : rocke_b_block_id_z(b); rocke_value_t* c_reps = rocke_b_const_i32(b, reps); slab_idx = rocke_b_mod(b, z, c_reps); } @@ -1362,6 +1771,18 @@ static void wgrad_emit_workspace_store_epilogue(rocke_ir_builder_t* b, * so accumulation error grows with split_k. For large split_k consider using * an fp32 workspace and downcasting after the kernel completes. */ +/* Python _row_addr: c_m, or c_m + group*kpg on the grouped path. */ +static rocke_value_t* wgrad_row_addr(rocke_ir_builder_t* b, + rocke_value_t* c_m, + rocke_value_t* group, + rocke_value_t* c_kpg) +{ + if(group == NULL) + return c_m; + rocke_value_t* g_base = rocke_b_mul(b, group, c_kpg); + return rocke_b_add(b, c_m, g_base); +} + static void wgrad_emit_split_k_epilogue_f32(rocke_ir_builder_t* b, const rocke_conv_build_ctx_t* ctx, const rocke_implicit_gemm_conv_wgrad_spec_t* spec, @@ -1427,6 +1848,10 @@ static void wgrad_emit_split_k_epilogue_f32(rocke_ir_builder_t* b, /* AOT: the GEMM extents are kernargs, not folded constants. */ rocke_value_t* wg_M_v = ctx->p_wg_M; rocke_value_t* wg_N_v = ctx->p_wg_N; + /* Grouped: Python _c_kpg = params["p_kpg"]; the atomic row lands in the + * group's absolute output-channel slab (_row_addr) while the bounds stay on + * the per-group row. */ + rocke_value_t* c_kpg_e = ctx->group_idx != NULL ? ctx->p_kpg : NULL; int flat = 0; for(int mi = 0; mi < mfmas_m; ++mi) @@ -1453,7 +1878,8 @@ static void wgrad_emit_split_k_epilogue_f32(rocke_ir_builder_t* b, * c_off = add(mul(c_m, wg_N_v), c_n) * ok = land(cmp_lt(c_m, wg_M_v), cmp_lt(c_n, wg_N_v)) * scf_if(ok): global_atomic_add(ptr, c_off, val) */ - rocke_value_t* c_off = rocke_b_add(b, rocke_b_mul(b, c_m, wg_N_v), c_n); + rocke_value_t* row = wgrad_row_addr(b, c_m, ctx->group_idx, c_kpg_e); + rocke_value_t* c_off = rocke_b_add(b, rocke_b_mul(b, row, wg_N_v), c_n); rocke_value_t* m_ok = rocke_b_cmp_lt(b, c_m, wg_M_v); rocke_value_t* n_ok = rocke_b_cmp_lt(b, c_n, wg_N_v); rocke_value_t* ok = rocke_b_land(b, m_ok, n_ok); @@ -1493,8 +1919,9 @@ static void wgrad_emit_split_k_epilogue_f32(rocke_ir_builder_t* b, rocke_value_t* is_odd = rocke_b_cmp_ne(b, c_n_is_odd, rocke_b_const_i32(b, 0)); rocke_value_t* c_n_even = rocke_b_sub(b, c_n, c_n_is_odd); + rocke_value_t* row = wgrad_row_addr(b, c_m, ctx->group_idx, c_kpg_e); rocke_value_t* c_off_even - = rocke_b_add(b, rocke_b_mul(b, c_m, wg_N_v), c_n_even); + = rocke_b_add(b, rocke_b_mul(b, row, wg_N_v), c_n_even); rocke_value_t* v_even = rocke_b_select(b, is_odd, zero, val); rocke_value_t* v_odd = rocke_b_select(b, is_odd, val, zero); rocke_value_t* elems[2] = {v_even, v_odd}; @@ -1582,6 +2009,9 @@ static void wgrad_emit_split_k_cshuffle_epilogue(rocke_ir_builder_t* b, struct WgradDwAddrCtx { rocke_tensor_descriptor_t* dW_desc; + /* Grouped: the per-group row m_val gains group*kpg (Python dw_addr). */ + rocke_value_t* group; + rocke_value_t* c_kpg; }; static rocke_value_t* wgrad_dw_addr(rocke_ir_builder_t* b, @@ -1591,6 +2021,12 @@ static rocke_value_t* wgrad_dw_addr(rocke_ir_builder_t* b, void* user) { WgradDwAddrCtx* wc = static_cast(user); + if(wc->group != NULL) + { + /* Python: m_g = b_.add(m_val, b_.mul(group, _c_kpg)) */ + rocke_value_t* g_base = rocke_b_mul(b, wc->group, wc->c_kpg); + m_global = rocke_b_add(b, m_global, g_base); + } /* Python: dW_desc.offset(b_, k_out=m_val, n_wg=n_val) * dW_desc top-level coords: ("k_out" = output channel, "n_wg" = filter+chan). * The epilogue calls addr_fn(m_global=output_channel, n_global=filter+channel). */ @@ -1689,6 +2125,8 @@ static void wgrad_emit_direct_epilogue(rocke_ir_builder_t* b, /* Built here, where Python builds it: its stride constant must * land at this point in the SSA sequence. */ addr_ctx.dW_desc = wgrad_build_dw_descriptor(b, ctx); + addr_ctx.group = ctx->group_idx; + addr_ctx.c_kpg = ctx->group_idx != NULL ? ctx->p_kpg : NULL; rocke_value_t* bound_m = ctx->p_wg_M; rocke_value_t* bound_n = ctx->p_wg_N; rocke_direct_epilogue_store(b, @@ -1855,10 +2293,15 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, return false; } { - int wg_M_val = spec->problem.K; /* groups==1 */ + int wg_M_val = rocke_wgrad_conv_spec_wg_M(spec); /* kpg */ int wg_N_val = rocke_wgrad_conv_spec_wg_N(spec); /* N=wg_M, Hi=1, Wi=1, C=wg_K, K=wg_N, Y=1, X=1 -> Ho=Wo=1, all positive */ *stub_p = rocke_conv_problem_default(wg_M_val, 1, 1, wg_K, wg_N_val, 1, 1); + /* The shared descriptor builders read the group count off ctx->p to + * decide whether the X channel axis carries the per-group slab embed + * (_a_channel_decode). The stub stands in for the GEMM shape only, so + * the true count has to be carried across. */ + stub_p->groups = spec->problem.groups > 0 ? spec->problem.groups : 1; } ctx->p = stub_p; @@ -2022,6 +2465,11 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, ctx->c_block_k = rocke_b_const_i32(b, ctx->block_k); /* AOT: the reduction extent is the p_wg_K kernarg, not a folded constant. */ rocke_value_t* c_wg_K = ctx->p_wg_K; + /* Grouped (stream-K only in this port): Python + * if grouped: c_kpg = p_kpg + * The group index itself comes from the stream-K tile decode. */ + if(spec->problem.groups > 1) + ctx->wgrad_c_kpg = ctx->p_kpg; /* Split-K K-slice bounds. The degree is a launch parameter (ks_count), * never a compile-time constant, so there is exactly one shape of decode: @@ -2039,12 +2487,20 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, * * C leaves argument evaluation order unspecified, so each subexpression is * bound to a temp in Python's left-to-right order. */ - rocke_value_t* z_id = rocke_b_block_id_z(b); - rocke_value_t* slice_id = rocke_b_mod(b, z_id, ks_count_param); - rocke_value_t* mul_lo = rocke_b_mul(b, slice_id, ks_param); - rocke_value_t* k_lo = rocke_b_to_sgpr_u32(b, mul_lo); - rocke_value_t* add_hi = rocke_b_add(b, k_lo, ks_param); - rocke_value_t* k_hi_v = rocke_b_to_sgpr_u32(b, add_hi); + /* Stream-K launches a one-dimensional grid and folds the group into its + * linear tile index, so it skips this decode and binds both per work + * item (wgrad_emit_streamk). */ + rocke_value_t* k_lo = c0_node; + rocke_value_t* k_hi_v = c_wg_K; + if(!wgrad_streamk_on(spec)) + { + rocke_value_t* z_id = rocke_b_block_id_z(b); + rocke_value_t* slice_id = rocke_b_mod(b, z_id, ks_count_param); + rocke_value_t* mul_lo = rocke_b_mul(b, slice_id, ks_param); + k_lo = rocke_b_to_sgpr_u32(b, mul_lo); + rocke_value_t* add_hi = rocke_b_add(b, k_lo, ks_param); + k_hi_v = rocke_b_to_sgpr_u32(b, add_hi); + } /* ctx->c0 is the literal zero the epilogue passes as a buffer soffset. * It used to be aliased to k_lo, which only worked while the unsplit path * made k_lo the constant 0; now that every launch computes k_lo from the @@ -2331,6 +2787,403 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, return rocke_ir_builder_ok(b); } +// --------------------------------------------------------------------------- +// Stream-K body (Python _emit_streamk_tiles in build_implicit_gemm_conv_wgrad) +// --------------------------------------------------------------------------- +// +// Every helper below sequences each builder call into its own local: C++ leaves +// the evaluation order of call arguments unspecified and the builder numbers +// SSA values in call order, so the statement order here is Python's +// left-to-right argument evaluation, call for call. + +/* The stream-K kernargs (Python params["sk_*"]), in ABI order. */ +typedef struct wgrad_streamk_args +{ + rocke_value_t* iters_per_tile; + rocke_value_t* dp_tiles; + rocke_value_t* ctas; + rocke_value_t* iters_per_cta; + rocke_value_t* extra_iters; + rocke_value_t* pool; + rocke_value_t* tree_rounds; +} wgrad_streamk_args_t; + +typedef struct wgrad_streamk_env +{ + rocke_conv_build_ctx_t* ctx; + const rocke_implicit_gemm_conv_wgrad_spec_t* spec; + wgrad_streamk_args_t k; + rocke_streamk_iter_args_t iter; + rocke_value_t* sk_flags; + rocke_value_t* sk_partials; + rocke_value_t* c_one; + int n_waves; + int n_slots; + int threads; +} wgrad_streamk_env_t; + +/* _flag_idx: add(mul(sk_idx, const(n_waves)), warp_id) */ +static rocke_value_t* wgrad_sk_flag_idx(wgrad_streamk_env_t* e, rocke_value_t* sk_idx) +{ + rocke_ir_builder_t* b = e->ctx->b; + rocke_value_t* c_waves = rocke_b_const_i32(b, e->n_waves); + rocke_value_t* scaled = rocke_b_mul(b, sk_idx, c_waves); + return rocke_b_add(b, scaled, e->ctx->warp_id); +} + +/* _slot_base: add(mul(sk_idx, const(n_slots * threads)), tid) */ +static rocke_value_t* wgrad_sk_slot_base(wgrad_streamk_env_t* e, rocke_value_t* sk_idx) +{ + rocke_ir_builder_t* b = e->ctx->b; + rocke_value_t* c_stride = rocke_b_const_i32(b, e->n_slots * e->threads); + rocke_value_t* scaled = rocke_b_mul(b, sk_idx, c_stride); + return rocke_b_add(b, scaled, e->ctx->tid); +} + +/* _store_partial: coalesced scalar stores, release fence, per-wave flag. */ +static void wgrad_sk_store_partial(wgrad_streamk_env_t* e, + rocke_value_t* sk_idx, + rocke_value_t* const* vals) +{ + rocke_ir_builder_t* b = e->ctx->b; + const int c_per_lane = e->ctx->c_per_lane; + rocke_value_t* base = wgrad_sk_slot_base(e, sk_idx); + for(int ai = 0; ai < e->ctx->num_accs; ++ai) + { + for(int i = 0; i < c_per_lane; ++i) + { + rocke_value_t* c_slot = rocke_b_const_i32(b, (ai * c_per_lane + i) * e->threads); + rocke_value_t* off = rocke_b_add(b, base, c_slot); + rocke_value_t* v = rocke_b_vec_extract(b, vals[ai], i); + rocke_b_global_store(b, e->sk_partials, off, v, 4); + } + } + rocke_b_fence(b, "agent", "release"); + rocke_value_t* flag = wgrad_sk_flag_idx(e, sk_idx); + rocke_b_global_flag_store(b, e->sk_flags, flag, e->c_one, "agent", "release"); +} + +/* _load_partial: one per accumulator. */ +static void + wgrad_sk_load_partial(wgrad_streamk_env_t* e, rocke_value_t* sk_idx, rocke_value_t** out) +{ + rocke_ir_builder_t* b = e->ctx->b; + const int c_per_lane = e->ctx->c_per_lane; + rocke_value_t* base = wgrad_sk_slot_base(e, sk_idx); + rocke_value_t* comps[64]; + for(int ai = 0; ai < e->ctx->num_accs; ++ai) + { + for(int i = 0; i < c_per_lane; ++i) + { + rocke_value_t* c_slot = rocke_b_const_i32(b, (ai * c_per_lane + i) * e->threads); + rocke_value_t* off = rocke_b_add(b, base, c_slot); + comps[i] = rocke_b_global_load_f32(b, e->sk_partials, off, 4); + } + out[ai] = rocke_b_vec_pack(b, comps, c_per_lane, rocke_f32()); + } +} + +/* _sk_cta_of: the stream-K CTA owning global iteration `it`. */ +static rocke_value_t* wgrad_sk_cta_of(wgrad_streamk_env_t* e, rocke_value_t* it) +{ + rocke_ir_builder_t* b = e->ctx->b; + rocke_value_t* rel = rocke_b_sub(b, it, e->iter.total_dp_iters); + rocke_value_t* wide = rocke_b_add(b, e->k.iters_per_cta, e->c_one); + rocke_value_t* boundary = rocke_b_mul(b, e->k.extra_iters, wide); + rocke_value_t* lead = rocke_b_div(b, rel, wide); + rocke_value_t* past = rocke_b_sub(b, rel, boundary); + rocke_value_t* past_ctas = rocke_b_div(b, past, e->k.iters_per_cta); + rocke_value_t* tail = rocke_b_add(b, e->k.extra_iters, past_ctas); + rocke_value_t* in_lead = rocke_b_cmp_lt(b, rel, boundary); + return rocke_b_select(b, in_lead, lead, tail); +} + +/* _store_tile: the plain dW store of a finished tile. */ +static void + wgrad_sk_store_tile(wgrad_streamk_env_t* e, rocke_value_t* const* vals, int wg_M, int wg_N) +{ + rocke_conv_build_ctx_t* ctx = e->ctx; + for(int i = 0; i < ctx->num_accs; ++i) + ctx->final_accs[i] = vals[i]; + ctx->num_final_accs = ctx->num_accs; + wgrad_emit_direct_epilogue(ctx->b, ctx, e->spec, ctx->d_rsrc, wg_M, wg_N); +} + +static void wgrad_emit_streamk(rocke_conv_build_ctx_t* ctx, + const rocke_implicit_gemm_conv_wgrad_spec_t* spec, + rocke_value_t* dW, + rocke_value_t* ws_ptr, + rocke_value_t* sk_flags, + rocke_value_t* sk_partials, + const wgrad_streamk_args_t* sk_args, + int wg_M, + int wg_N) +{ + rocke_ir_builder_t* b = ctx->b; + wgrad_streamk_env_t e; + memset(&e, 0, sizeof(e)); + e.ctx = ctx; + e.spec = spec; + e.k = *sk_args; + e.sk_flags = sk_flags; + e.sk_partials = sk_partials; + e.n_waves = spec->warp_m * spec->warp_n; + e.n_slots = ctx->num_accs * ctx->c_per_lane; + e.threads = ctx->threads; + if(ctx->c_per_lane > 64) + { + rocke_i_set_err(b, ROCKE_ERR_VALUE, "wgrad: stream-K partial too wide"); + return; + } + const char* reduction = wgrad_streamk_reduction(spec); + const bool persistent = strcmp(spec->streamk, "persistent") == 0; + const bool grouped = spec->problem.groups > 1; + + e.c_one = rocke_b_const_i32(b, 1); + rocke_value_t* c_one = e.c_one; + rocke_value_t* c_ipt = e.k.iters_per_tile; + rocke_value_t* c_pool = e.k.pool; + e.iter.iters_per_tile = e.k.iters_per_tile; + e.iter.dp_tiles = e.k.dp_tiles; + e.iter.total_dp_iters + = rocke_b_to_sgpr_u32(b, rocke_b_mul(b, e.k.dp_tiles, e.k.iters_per_tile)); + e.iter.iters_per_sk_cta = e.k.iters_per_cta; + e.iter.extra_iters = e.k.extra_iters; + rocke_value_t* cta = rocke_b_to_sgpr_u32(b, rocke_b_block_id_x(b)); + rocke_value_t* n_dp = NULL; + rocke_value_t* first_tile = NULL; + rocke_value_t* n_items = NULL; + rocke_value_t* sk_cta = NULL; + rocke_value_t* c0 = ctx->c0; /* k_lo of the classic path == const 0 */ + rocke_streamk_iter_range_t rng; + memset(&rng, 0, sizeof(rng)); + if(persistent) + { + rocke_value_t* left = rocke_b_sub(b, e.k.dp_tiles, cta); + rocke_value_t* c_round = rocke_b_sub(b, c_pool, c_one); + rocke_value_t* rounded = rocke_b_add(b, left, c_round); + n_dp = rocke_b_div(b, rounded, c_pool); + rng = rocke_emit_streamk_iter_range(b, cta, &e.iter, false); + first_tile = rocke_b_div(b, rng.start, c_ipt); + rocke_value_t* last_it = rocke_b_sub(b, rng.end, c_one); + rocke_value_t* last_tile = rocke_b_div(b, last_it, c_ipt); + rocke_value_t* span = rocke_b_sub(b, last_tile, first_tile); + rocke_value_t* n_sk_raw = rocke_b_add(b, span, c_one); + rocke_value_t* is_sk = rocke_b_cmp_lt(b, cta, e.k.ctas); + rocke_value_t* n_sk = rocke_b_select(b, is_sk, n_sk_raw, c0); + n_items = rocke_b_to_sgpr_u32(b, rocke_b_add(b, n_dp, n_sk)); + sk_cta = cta; + } + else + { + rng = rocke_emit_streamk_iter_range(b, cta, &e.iter, true); + first_tile = rocke_b_div(b, rng.start, c_ipt); + rocke_value_t* last_it = rocke_b_sub(b, rng.end, c_one); + rocke_value_t* last_tile = rocke_b_div(b, last_it, c_ipt); + rocke_value_t* span = rocke_b_sub(b, last_tile, first_tile); + n_items = rocke_b_to_sgpr_u32(b, rocke_b_add(b, span, c_one)); + sk_cta = rng.sk_cta; + } + + rocke_for_t loop = rocke_b_scf_for(b, c0, n_items, c_one, "sk_item"); + rocke_b_region_enter(b, loop.body); + { + rocke_value_t* item = loop.iv; + rocke_value_t* tile; + rocke_value_t* item_start; + rocke_value_t* item_end; + if(persistent) + { + rocke_value_t* is_dp = rocke_b_cmp_lt(b, item, n_dp); + rocke_value_t* stride = rocke_b_mul(b, item, c_pool); + rocke_value_t* dp_tile = rocke_b_add(b, cta, stride); + rocke_value_t* sk_off = rocke_b_sub(b, item, n_dp); + rocke_value_t* sk_tile = rocke_b_add(b, first_tile, sk_off); + tile = rocke_b_to_sgpr_u32(b, rocke_b_select(b, is_dp, dp_tile, sk_tile)); + rocke_value_t* dp_start = rocke_b_mul(b, tile, c_ipt); + item_start = rocke_b_select(b, is_dp, dp_start, rng.start); + rocke_value_t* dp_end = rocke_b_add(b, dp_start, c_ipt); + item_end = rocke_b_select(b, is_dp, dp_end, rng.end); + } + else + { + tile = rocke_b_to_sgpr_u32(b, rocke_b_add(b, first_tile, item)); + item_start = rng.start; + item_end = rng.end; + } + rocke_value_t* tile_start = rocke_b_mul(b, tile, c_ipt); + rocke_value_t* tile_end = rocke_b_add(b, tile_start, c_ipt); + rocke_value_t* lo_max = rocke_b_smax(b, item_start, tile_start); + rocke_value_t* lo = rocke_b_sub(b, lo_max, tile_start); + rocke_value_t* hi_min = rocke_b_smin(b, item_end, tile_end); + rocke_value_t* hi = rocke_b_sub(b, hi_min, tile_start); + + /* Tile -> (group, m_tile, n_tile), groups folded into GEMM-M. */ + rocke_value_t* n_tile = rocke_b_mod(b, tile, ctx->p_num_pid_n); + rocke_value_t* m_idx = rocke_b_div(b, tile, ctx->p_num_pid_n); + rocke_value_t* m_tile = m_idx; + if(grouped) + { + ctx->group_idx = rocke_b_to_sgpr_u32(b, rocke_b_mod(b, m_idx, ctx->p_groups)); + m_tile = rocke_b_div(b, m_idx, ctx->p_groups); + } + rocke_value_t* c_bm = rocke_b_const_i32(b, ctx->block_m); + ctx->block_m_off_v = rocke_b_to_sgpr_u32(b, rocke_b_mul(b, m_tile, c_bm)); + rocke_value_t* c_bn = rocke_b_const_i32(b, ctx->block_n); + ctx->block_n_off_v = rocke_b_to_sgpr_u32(b, rocke_b_mul(b, n_tile, c_bn)); + ctx->grid.block_m_off = ctx->block_m_off_v; + ctx->grid.block_n_off = ctx->block_n_off_v; + + /* The item's K slice. kloop_k_lo is what the loop driver reads; + * ctx->c0 stays the literal zero the epilogue uses as a soffset. */ + ctx->kloop_k_lo = rocke_b_to_sgpr_u32(b, rocke_b_mul(b, lo, ctx->c_block_k)); + ctx->c_K_gemm = rocke_b_to_sgpr_u32(b, rocke_b_mul(b, hi, ctx->c_block_k)); + ctx->kloop_k_hi = ctx->c_K_gemm; + rocke_conv_emit_kloop_simple(ctx); + rocke_value_t* tile_accs[ROCKE_CONV_MAX_ACCS]; + for(int i = 0; i < ctx->num_accs; ++i) + tile_accs[i] = ctx->final_accs[i]; + + /* Whether this item covers the start / the end of its tile. Built only + * for the flag-handoff fixups, so the atomic and workspace bodies carry + * no dead compares. Mirrors Python. */ + rocke_value_t* started = NULL; + rocke_value_t* ended = NULL; + if(strcmp(reduction, "linear") == 0 || strcmp(reduction, "tree") == 0) + { + started = rocke_b_cmp_eq(b, lo, c0); + ended = rocke_b_cmp_eq(b, hi, c_ipt); + } + + if(strcmp(reduction, "atomic") == 0) + { + wgrad_emit_split_k_epilogue_f32(b, ctx, spec, dW, wg_M, wg_N); + } + else if(strcmp(reduction, "workspace") == 0) + { + ctx->wgrad_replica_sel = cta; + wgrad_emit_workspace_store_epilogue(b, ctx, spec, ws_ptr, wg_M, wg_N); + } + else if(strcmp(reduction, "linear") == 0) + { + rocke_if_else_t ie = rocke_b_scf_if_else(b, started); + rocke_b_region_enter(b, ie.then_region); + { + rocke_value_t* last_it = rocke_b_sub(b, tile_end, c_one); + rocke_value_t* last = wgrad_sk_cta_of(&e, last_it); + rocke_value_t* span = rocke_b_sub(b, last, sk_cta); + rocke_value_t* n_partners = rocke_b_select(b, ended, c0, span); + rocke_iter_arg_t iter_args[ROCKE_CONV_MAX_ACCS]; + for(int i = 0; i < ctx->num_accs; ++i) + { + iter_args[i].name = rocke_arena_printf(&b->arena, "sk_%s", ctx->acc_names[i]); + iter_args[i].init = tile_accs[i]; + } + rocke_value_t* fold_hi = rocke_b_to_sgpr_u32(b, rocke_b_add(b, n_partners, c_one)); + rocke_for_t fold = rocke_b_scf_for_iter( + b, c_one, fold_hi, c_one, iter_args, ctx->num_accs, "sk_p", false, true); + rocke_b_region_enter(b, fold.body); + { + rocke_value_t* partner = rocke_b_add(b, sk_cta, fold.iv); + rocke_value_t* flag = wgrad_sk_flag_idx(&e, partner); + rocke_b_global_flag_wait_eq(b, sk_flags, flag, c_one, "agent"); + rocke_value_t* parts[ROCKE_CONV_MAX_ACCS]; + wgrad_sk_load_partial(&e, partner, parts); + rocke_value_t* sums[ROCKE_CONV_MAX_ACCS]; + for(int i = 0; i < ctx->num_accs; ++i) + sums[i] = rocke_b_vector_add(b, fold.iter_vars[i], parts[i]); + rocke_b_scf_yield(b, sums, ctx->num_accs); + } + rocke_b_region_leave(b); + wgrad_sk_store_tile(&e, fold.op->results, wg_M, wg_N); + } + rocke_b_region_leave(b); + rocke_b_region_enter(b, ie.else_region); + wgrad_sk_store_partial(&e, sk_cta, tile_accs); + rocke_b_region_leave(b); + } + else + { + /* Tree: the round count is a kernarg, so the rounds are a loop + * carrying the partial sums, the pair stride and a dropped-out + * flag (see the Python comment). */ + rocke_value_t* not_started = rocke_b_lnot(b, started); + rocke_value_t* owner = wgrad_sk_cta_of(&e, tile_start); + rocke_value_t* local = rocke_b_sub(b, sk_cta, owner); + rocke_value_t* zero_acc = rocke_b_zero_vec_f32(b, ctx->c_per_lane); + rocke_value_t* c_two = rocke_b_const_i32(b, 2); + const int n_acc = ctx->num_accs; + rocke_iter_arg_t round_args[ROCKE_CONV_MAX_ACCS + 2]; + for(int i = 0; i < n_acc; ++i) + { + round_args[i].name = rocke_arena_printf(&b->arena, "skt_%s", ctx->acc_names[i]); + round_args[i].init = tile_accs[i]; + } + round_args[n_acc].name = "skt_stride"; + round_args[n_acc].init = c_one; + round_args[n_acc + 1].name = "skt_done"; + round_args[n_acc + 1].init = c0; + rocke_for_t rounds = rocke_b_scf_for_iter( + b, c0, e.k.tree_rounds, c_one, round_args, n_acc + 2, "sk_r", false, true); + rocke_b_region_enter(b, rounds.body); + { + rocke_value_t* vals[ROCKE_CONV_MAX_ACCS]; + for(int i = 0; i < n_acc; ++i) + vals[i] = rounds.iter_vars[i]; + rocke_value_t* stride = rounds.iter_vars[n_acc]; + rocke_value_t* live = rocke_b_cmp_eq(b, rounds.iter_vars[n_acc + 1], c0); + rocke_value_t* partner = rocke_b_add(b, sk_cta, stride); + rocke_value_t* not_ended = rocke_b_lnot(b, ended); + rocke_value_t* p_start = rocke_emit_streamk_sk_start_iter(b, partner, &e.iter); + rocke_value_t* p_before_end = rocke_b_cmp_lt(b, p_start, tile_end); + rocke_value_t* in_tile = rocke_b_land(b, not_ended, p_before_end); + rocke_value_t* pair = rocke_b_mul(b, stride, c_two); + rocke_value_t* rem = rocke_b_mod(b, local, pair); + rocke_value_t* is_recv = rocke_b_cmp_eq(b, rem, c0); + rocke_value_t* recv_in = rocke_b_land(b, is_recv, in_tile); + rocke_value_t* recv = rocke_b_land(b, recv_in, live); + rocke_if_t wait_if = rocke_b_scf_if(b, recv); + rocke_b_region_enter(b, wait_if.then_region); + { + rocke_value_t* flag = wgrad_sk_flag_idx(&e, partner); + rocke_b_global_flag_wait_eq(b, sk_flags, flag, c_one, "agent"); + } + rocke_b_region_leave(b); + rocke_value_t* src = rocke_b_select(b, recv, partner, c0); + rocke_value_t* parts[ROCKE_CONV_MAX_ACCS]; + wgrad_sk_load_partial(&e, src, parts); + for(int i = 0; i < n_acc; ++i) + { + rocke_value_t* q = rocke_b_select(b, recv, parts[i], zero_acc); + vals[i] = rocke_b_vector_add(b, vals[i], q); + } + rocke_value_t* not_recv = rocke_b_lnot(b, is_recv); + rocke_value_t* send_ns = rocke_b_land(b, not_started, not_recv); + rocke_value_t* send = rocke_b_land(b, send_ns, live); + rocke_if_t send_if = rocke_b_scf_if(b, send); + rocke_b_region_enter(b, send_if.then_region); + wgrad_sk_store_partial(&e, sk_cta, vals); + rocke_b_region_leave(b); + rocke_value_t* yields[ROCKE_CONV_MAX_ACCS + 2]; + for(int i = 0; i < n_acc; ++i) + yields[i] = vals[i]; + yields[n_acc] = rocke_b_mul(b, stride, c_two); + yields[n_acc + 1] = rocke_b_select(b, send, c_one, rounds.iter_vars[n_acc + 1]); + rocke_b_scf_yield(b, yields, n_acc + 2); + } + rocke_b_region_leave(b); + rocke_value_t* vals[ROCKE_CONV_MAX_ACCS]; + for(int i = 0; i < n_acc; ++i) + vals[i] = rounds.op->results[i]; + rocke_if_t store_if = rocke_b_scf_if(b, started); + rocke_b_region_enter(b, store_if.then_region); + wgrad_sk_store_tile(&e, vals, wg_M, wg_N); + rocke_b_region_leave(b); + } + } + rocke_b_region_leave(b); +} + // --------------------------------------------------------------------------- // rocke_build_implicit_gemm_conv_wgrad // --------------------------------------------------------------------------- @@ -2363,8 +3216,11 @@ static rocke_value_t* wgrad_declare_ptr(rocke_ir_builder_t* b, case ROCKE_CONV_ARG_D: return rocke_b_param(b, name, d->d_type, d->d_opts); case ROCKE_CONV_ARG_F32_PTR: - /* Two-stage: f32 workspace that receives the partial sums. */ + /* Two-stage workspace, or the stream-K partial slots. */ return rocke_b_param(b, name, rocke_ptr_type(b, rocke_f32(), "global"), d->ws_opts); + case ROCKE_CONV_ARG_I32_PTR: + /* Stream-K linear/tree ready flags. */ + return rocke_b_param(b, name, rocke_ptr_type(b, rocke_i32(), "global"), d->ws_opts); default: return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "wgrad: unexpected pointer kind for %s", name); @@ -2403,6 +3259,13 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( return NULL; } + /* streamk_ctas needs no resolution here: the CTA pool is a launch + * parameter, so -1 (auto) and any explicit pool build the same kernel. */ + const bool is_streamk = wgrad_streamk_on(spec); + const char* sk_red = wgrad_streamk_reduction(spec); + const bool streamk_atomic + = is_streamk && (strcmp(sk_red, "atomic") == 0 || strcmp(sk_red, "workspace") == 0); + bool effective_two_stage = spec->two_stage; bool is_split_k = split_k > 1; @@ -2449,7 +3312,7 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( /* split_k>1: dW is read+write (atomic); split_k=1: writeonly. * Caller MUST zero-init dW before launch for atomic paths -- the kernel only * issues atomic-adds. See the header contract note for details. */ - d_opts.writeonly = !is_split_k; + d_opts.writeonly = !is_split_k && !streamk_atomic; d_opts.writeonly_set = true; d_opts.align = 16; d_opts.align_set = true; @@ -2476,6 +3339,13 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( rocke_value_t* X_bytes = NULL; rocke_value_t* dW_bytes = NULL; + /* Stream-K linear/tree: per-wave ready flags and per-CTA partial slots. + * Both are read back by the tile owner, so neither is writeonly. */ + rocke_value_t* sk_flags = NULL; + rocke_value_t* sk_partials = NULL; + wgrad_streamk_args_t sk_args; + memset(&sk_args, 0, sizeof(sk_args)); + /* ---- the whole kernarg list ---- * Pointers, byte sizes, the shared extent block, wgrad's own GEMM dims, * strides and magic pairs, then the variant-specific extras -- emitted @@ -2562,10 +3432,19 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( {"ws_ptr", &ws_ptr}, {"ks", &ks_param}, {"ks_count", &ks_count_param}, + {"sk_flags", &sk_flags}, + {"sk_partials", &sk_partials}, + {"sk_iters_per_tile", &sk_args.iters_per_tile}, + {"sk_dp_tiles", &sk_args.dp_tiles}, + {"sk_ctas", &sk_args.ctas}, + {"sk_iters_per_cta", &sk_args.iters_per_cta}, + {"sk_extra_iters", &sk_args.extra_iters}, + {"sk_pool", &sk_args.pool}, + {"sk_tree_rounds", &sk_args.tree_rounds}, }; rocke_conv_arg_list_t abi; ctx.params_is_3d = is_3d; - if(!rocke_conv_arg_names("wgrad", is_3d, is_two_stage, &abi)) + if(!rocke_conv_arg_names_ex("wgrad", is_3d, is_two_stage, is_streamk ? sk_red : NULL, &abi)) { rocke_i_set_err(b, ROCKE_ERR_VALUE, "wgrad: no kernarg ABI for this variant"); return NULL; @@ -2576,7 +3455,6 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( return NULL; } } - if(!wgrad_build_ctx_init(&ctx, b, spec, arch, wg_K, ks_param, ks_count_param)) return NULL; @@ -2750,6 +3628,14 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( if(!rocke_ir_builder_ok(b)) return NULL; + if(is_streamk) + { + wgrad_emit_streamk(&ctx, spec, dW, ws_ptr, sk_flags, sk_partials, &sk_args, wg_M, wg_N); + if(!rocke_ir_builder_ok(b)) + return NULL; + return b->kernel; + } + /* --- K-loop --- */ if(spec->unroll_k) rocke_conv_emit_kloop_unroll(&ctx); @@ -2806,6 +3692,8 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( * emit its stride constant for both and in the wrong place. */ WgradDwAddrCtx dw_addr_ctx; dw_addr_ctx.dW_desc = NULL; + dw_addr_ctx.group = NULL; + dw_addr_ctx.c_kpg = NULL; if(spec->epilogue && strcmp(spec->epilogue, "cshuffle") == 0) { 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 60814886d340..15ae339a07d8 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 @@ -1991,6 +1991,10 @@ Per-arch facts belong in the arch references — see §21. | `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 | | 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 | +| `streamk_ctas` | `WgradConvSpec` | `-1` | The CTA pool the stream-K remainder is spread over (CK `max_active_wgs`). A launch parameter, not part of the kernel: it sizes the host partition that arrives as the `sk_*` kernargs. `-1` resolves per reduction: `workspace`/`atomic` (no waiting) get four CTAs per CU (`_STREAMK_CTAS_PER_CU`, counted in CTAs, not waves); `linear`/`tree` get one CTA per CU, because their fixups get slower as more CTAs share a tile. The pool need not be resident: a CTA waits only on higher-index CTAs and publishes its own partial first, so ascending workgroup dispatch keeps the waits from starving their producers | +| `--streamk` (sweep) | `benchmark_implicit_gemm_conv.py` / `kernel_sweep.py` | `off` | Opt-in wgrad stream-K candidates (`dp_sk`/`persistent`/`all` × `--streamk-reductions`); the dispatcher does not select stream-K | #### 12.1.J Quantization diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md index 66d0dfe434bc..e6d2dbd2f2b7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md @@ -311,6 +311,63 @@ b.global_atomic_add(Hist, eid, b.const_i32(1)) local_off = b.global_atomic_add(Counter, eid, b.const_i32(1)) ``` +### Cross-workgroup flags and fences + +Stream-K fixups hand a partial tile from one workgroup to another through +global memory, guarded by i32 flags. Three leaf ops cover the protocol: + +```text +fence(*, scope="agent", ordering="acq_rel") +global_flag_store(ptr, idx, value, *, scope="agent", ordering="release") +global_flag_wait_eq(ptr, idx, expect, *, scope="agent") +``` + +`scope` is one of `workgroup`, `agent`, `system` (LLVM `syncscope`; +`system` has no clause). `ptr` must be a `ptr`: the builders and +the verifier reject any other element type, because the LLVM lowering +addresses the flag as an i32 while the HIP lowering indexes the pointee. +Both lowerers also reject a missing or unknown `scope` / `ordering` instead +of defaulting it, so IR that skipped the verifier cannot lower to an invalid +`fence` or atomic. + +The wgrad stream-K fixups (`linear` and `tree` in +`library/kernels/common/conv_implicit_gemm_wgrad.py`) keep one flag per +(stream-K CTA, wave), at index `sk_cta * n_waves + warp_id`. A consumer lane +only reads the partial written by the same lane of the same wave in the +producer CTA, so each wave hands off on its own, with no workgroup barrier +on either side: + +- Producer wave: store this wave's slice of the partial tile, issue + `fence(scope="agent", ordering="release")`, then + `global_flag_store(flags, sk_cta * n_waves + warp_id, 1)`. +- Consumer wave: `global_flag_wait_eq(flags, partner * n_waves + warp_id, 1)`, + which lowers to a self-contained loop around an `acquire` atomic load, + then load the partner's partial. The acquire orders those loads after the + observed flag, so they see what the producer wave released. + +At `agent` scope on multi-XCD CDNA parts the release writes back and the +acquire invalidates the per-XCD L2. The flag region must be zeroed before +each launch; the partials are always written before they are read. + +Forward progress: a CTA only waits on a higher-index stream-K CTA (the +`linear` owner folds `sk_cta + 1 ...`, a `tree` receiver folds +`sk_cta + stride`), never on a lower one. If the whole CTA pool is resident +at once, every wait terminates. When the pool is larger than what can be +resident, the protocol relies on workgroups being dispatched in ascending +index order, which the hardware does in practice but the programming model +does not guarantee. Under that order only the CTAs of the tile straddling +the dispatch front can wait on a CTA that has not started, and they hold a +bounded number of slots (one for `linear`, about one per round for +`tree`), so the other resident CTAs finish and free slots for the producers. +Do not add a wait on a lower-index CTA, and do not reorder the grid (for +example by remapping `block_id_x`) without revisiting this. + +| IR op | LLVM | HIP | +|---|---|---| +| `memref.fence` | `fence syncscope("agent") release` | `__builtin_amdgcn_fence` | +| `memref.global_flag_store` | `store atomic i32 ... release` | `__hip_atomic_store` | +| `memref.global_flag_wait_eq` | `load atomic i32 ... acquire` spin | `while (__hip_atomic_load(...) != e)` | + Block-level histogram + scan helpers in `helpers/scan.py` wrap the common patterns: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md index 43d968991cae..2e801ce84229 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md @@ -353,6 +353,11 @@ b.clamp_f32(v, lo, hi) # fmin(hi, fmax(lo, v)) -- folds to v_med3_f32 b.global_atomic_add(ptr, idx, val, *, ordering="monotonic") b.lds_atomic_add(smem, [idx], val, *, ordering="monotonic") +# : cross-workgroup synchronisation (stream-K fixup) +b.fence(*, scope="agent", ordering="acq_rel") # memref.fence +b.global_flag_store(ptr, idx, val_i32, *, scope="agent", ordering="release") +b.global_flag_wait_eq(ptr, idx, expect_i32, *, scope="agent") # acquire spin + # : bitwise integer helpers (also used by i4 unpack) b.xor(a, b) # arith.xor; same type on both operands b.shl(a, b) # arith.shl; logical left shift diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md index da95fd84bb0d..4730a9f37f04 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md @@ -94,6 +94,9 @@ Complete reference of operations recognized by `core/ir.py` and lowered to AMDGP | `memref.global_store_typed` | `global_store` | typed `store addrspace(1)` | | `memref.global_store_vN` | `global_store_vN` | vector `store addrspace(1)` | | `memref.global_atomic_add_f32`| `global_atomic_add_f32` | `atomicrmw fadd addrspace(1)` | +| `memref.fence` | `fence` | `fence syncscope(...) ` | +| `memref.global_flag_store` | `global_flag_store` | `store atomic i32 ... syncscope(...) ` | +| `memref.global_flag_wait_eq` | `global_flag_wait_eq` | spin loop around `load atomic i32 ... acquire` | `masked_global_load(ptr, idx, mask, other, dtype)` clamps the false-lane index to 0 before loading, then selects `other` afterward. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py index 88beb2528cf4..4880390aa68d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py @@ -144,6 +144,44 @@ def _streaming(temporal_hint: TemporalHint) -> bool: return temporal_hint is TemporalHint.STREAMING +# LLVM ``syncscope`` values accepted by the cross-workgroup fence / flag ops. +# ``system`` renders as the empty (default) scope. +_FENCE_SCOPES = ("workgroup", "agent", "system") +# Orderings each op accepts; the verifier checks serialized IR against these. +_FENCE_ORDERINGS = ("acquire", "release", "acq_rel", "seq_cst") +_FLAG_STORE_ORDERINGS = ("monotonic", "release", "seq_cst") + + +def _is_flag_ptr(t: "Type") -> bool: + """True for the ``ptr`` a cross-workgroup flag op addresses. + + Both lowerers hard-code an i32 element (LLVM GEPs ``i32``; HIP indexes + the pointer's own element type), so any other pointee would make the two + disagree on the flag's address and width. + """ + return isinstance(t, PtrType) and t.space == "global" and t.pointee.name == "i32" + + +def _check_flag_ptr(op_name: str, ptr: "Value") -> None: + if not _is_flag_ptr(ptr.type): + raise ValueError(f"{op_name} expects a global i32 pointer, got {ptr.type.name}") + + +def _sync_attr(op: "Op", key: str, allowed: Sequence[str]) -> str: + """Scope / ordering attr of a fence or flag op, checked for the lowerers. + + The lowerers can run on IR that never went through :func:`verify`, so a + missing or unknown value is rejected here instead of being defaulted or + rendered into invalid IR. The C++ engine raises the same messages. + """ + value = op.attrs.get(key) + if not isinstance(value, str): + raise ValueError(f"{op.name}: {key} attribute missing or not a string") + if value not in allowed: + raise ValueError(f"{op.name}: unknown {key} '{value}'") + return value + + # ----- target-neutral MMA metadata --------------------------------------- # # ``IRBuilder.mma`` emits a single ``tile.mma`` op keyed by ``op_id``; the ISA @@ -1466,6 +1504,82 @@ def global_atomic_add_pk_f16( result_name_hint="atom_f16", ).result + # ----- cross-workgroup synchronisation (stream-K fixup) ----- + + def fence(self, *, scope: str = "agent", ordering: str = "acq_rel") -> None: + """Memory fence at ``scope`` (LLVM ``fence syncscope(...) ``). + + A ``release`` fence orders every prior global store of the issuing + wave before any later flag store; an ``acquire`` fence orders a prior + flag load before every later global load. At ``agent`` scope on + multi-XCD CDNA parts the backend also writes back / invalidates the + per-XCD L2, which is what makes a producer's partial tile visible to + a consumer workgroup running on a different XCD. + """ + if scope not in _FENCE_SCOPES: + raise ValueError(f"unknown fence scope {scope!r}") + if ordering not in _FENCE_ORDERINGS: + raise ValueError(f"unknown fence ordering {ordering!r}") + self._op("memref.fence", [], attrs={"scope": scope, "ordering": ordering}) + + def global_flag_store( + self, + ptr: Value, + idx: Value, + value: Value, + *, + scope: str = "agent", + ordering: str = "release", + ) -> None: + """Atomic i32 store of a cross-workgroup flag (``ptr[idx] = value``). + + Lowers to ``store atomic i32 ... syncscope(scope) ordering``. Pair + with :meth:`global_flag_wait_eq` on the consumer side. + """ + _check_flag_ptr("global_flag_store", ptr) + if value.type.name != "i32": + raise ValueError( + f"global_flag_store expects an i32 value, got {value.type.name}" + ) + if scope not in _FENCE_SCOPES: + raise ValueError(f"unknown flag scope {scope!r}") + if ordering not in _FLAG_STORE_ORDERINGS: + raise ValueError(f"unknown flag store ordering {ordering!r}") + self._op( + "memref.global_flag_store", + [ptr, idx, value], + attrs={"scope": scope, "ordering": ordering}, + ) + + def global_flag_wait_eq( + self, + ptr: Value, + idx: Value, + expect: Value, + *, + scope: str = "agent", + ) -> None: + """Spin until the i32 flag ``ptr[idx]`` equals ``expect``. + + Lowers to a self-contained loop around an ``acquire`` atomic load at + ``scope``, so global loads issued after the wait observe every store + the producer released before setting the flag. The loop is the only + data-dependent trip count in the IR; it is a leaf op so structured + control flow stays ``scf.for`` / ``scf.if`` only. + """ + _check_flag_ptr("global_flag_wait_eq", ptr) + if expect.type.name != "i32": + raise ValueError( + f"global_flag_wait_eq expects an i32 value, got {expect.type.name}" + ) + if scope not in _FENCE_SCOPES: + raise ValueError(f"unknown flag scope {scope!r}") + self._op( + "memref.global_flag_wait_eq", + [ptr, idx, expect], + attrs={"scope": scope}, + ) + def fp16_zero(self) -> Value: return self._op( "arith.constant", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py index 8108ba4359f4..a18b3d9a6fcb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py @@ -20,6 +20,9 @@ from .arch.wmma_scale import gfx1250_scaled_wmma from .dtypes import dtype_info from .ir import ( + _FENCE_ORDERINGS, + _FENCE_SCOPES, + _FLAG_STORE_ORDERINGS, KernelDef, Op, PtrType, @@ -28,6 +31,7 @@ Type, Value, VectorType, + _sync_attr, require_streaming_arch, ) @@ -269,6 +273,21 @@ def _encode_waitcnt_gfx11(vmcnt: int, expcnt: int, lgkmcnt: int) -> int: return (ec_b & 0x7) | ((lk_b & 0x3F) << 4) | ((vm_b & 0x3F) << 10) +# Fence / flag-op scope and ordering spellings for the HIP atomic builtins. +_HIP_ATOMIC_ORDER = { + "monotonic": "__ATOMIC_RELAXED", + "acquire": "__ATOMIC_ACQUIRE", + "release": "__ATOMIC_RELEASE", + "acq_rel": "__ATOMIC_ACQ_REL", + "seq_cst": "__ATOMIC_SEQ_CST", +} +_HIP_FENCE_SCOPE = {"workgroup": "workgroup", "agent": "agent", "system": ""} +_HIP_MEMORY_SCOPE = { + "workgroup": "__HIP_MEMORY_SCOPE_WORKGROUP", + "agent": "__HIP_MEMORY_SCOPE_AGENT", + "system": "__HIP_MEMORY_SCOPE_SYSTEM", +} + # Default arch for the HIP path. gfx950 is the historical CDNA target whose # emitted source is the byte-identical baseline; callers that don't pass an # arch (e.g. the in-tree coverage tests) keep getting exactly that output. @@ -1731,6 +1750,34 @@ def _op_memref_global_atomic_add(self, op: Op) -> None: f"{_name(val)});" ) + def _op_memref_fence(self, op: Op) -> None: + """Lower ``memref.fence`` via ``__builtin_amdgcn_fence``.""" + scope = _sync_attr(op, "scope", _FENCE_SCOPES) + ordering = _sync_attr(op, "ordering", _FENCE_ORDERINGS) + self._emit( + f"__builtin_amdgcn_fence({_HIP_ATOMIC_ORDER[ordering]}, " + f'"{_HIP_FENCE_SCOPE[scope]}");' + ) + + def _op_memref_global_flag_store(self, op: Op) -> None: + """Lower ``memref.global_flag_store`` via ``__hip_atomic_store``.""" + ptr, idx, val = op.operands + scope = _sync_attr(op, "scope", _FENCE_SCOPES) + ordering = _sync_attr(op, "ordering", _FLAG_STORE_ORDERINGS) + self._emit( + f"__hip_atomic_store(&{_name(ptr)}[{_name(idx)}], {_name(val)}, " + f"{_HIP_ATOMIC_ORDER[ordering]}, {_HIP_MEMORY_SCOPE[scope]});" + ) + + def _op_memref_global_flag_wait_eq(self, op: Op) -> None: + """Lower ``memref.global_flag_wait_eq`` to an acquire-load spin loop.""" + ptr, idx, expect = op.operands + scope = _sync_attr(op, "scope", _FENCE_SCOPES) + self._emit( + f"while (__hip_atomic_load(&{_name(ptr)}[{_name(idx)}], __ATOMIC_ACQUIRE, " + f"{_HIP_MEMORY_SCOPE[scope]}) != {_name(expect)}) {{}}" + ) + def _op_memref_global_atomic_add_pk_bf16(self, op: Op) -> None: """Lower the packed-bf16 atomic add via the AMDGPU builtin. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py index 62ef60532550..34dac431b950 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py @@ -39,6 +39,9 @@ from .codegen_policy import codegen_policy_for_kernel from .ir import ( + _FENCE_ORDERINGS, + _FENCE_SCOPES, + _FLAG_STORE_ORDERINGS, BF16, F32, F16, @@ -54,6 +57,7 @@ Type, Value, VectorType, + _sync_attr, require_streaming_arch, split_loc, ) @@ -1136,6 +1140,16 @@ def _param_llvm_type(p: Param) -> str: return _llvm_type(p.type) +def _llvm_syncscope(scope: str) -> str: + """Render a fence / flag-op scope as an LLVM ``syncscope`` clause. + + ``system`` is LLVM's default scope and has no textual clause. + """ + if scope == "system": + return "" + return f' syncscope("{scope}")' + + def _escape_llvm_asm_string(s: str) -> str: r"""Escape a Python string for an LLVM IR asm/string literal. @@ -2945,6 +2959,56 @@ def _op_memref_global_atomic_add(self, op: Op) -> None: f"{elem_ty} {self._operand(val)} {ordering}{md}" ) + def _op_memref_fence(self, op: Op) -> None: + """Lower ``memref.fence`` to an LLVM ``fence`` at the op's scope.""" + scope = _sync_attr(op, "scope", _FENCE_SCOPES) + ordering = _sync_attr(op, "ordering", _FENCE_ORDERINGS) + self._current().emit(f" fence{_llvm_syncscope(scope)} {ordering}") + + def _op_memref_global_flag_store(self, op: Op) -> None: + """Lower ``memref.global_flag_store`` to a scoped ``store atomic i32``.""" + ptr, idx, val = op.operands + scope = _sync_attr(op, "scope", _FENCE_SCOPES) + ordering = _sync_attr(op, "ordering", _FLAG_STORE_ORDERINGS) + gep = self._fresh("gep") + self._current().emit( + f" {gep} = getelementptr inbounds i32, ptr addrspace(1) " + f"{self._operand(ptr)}, i32 {self._operand(idx)}" + ) + self._current().emit( + f" store atomic i32 {self._operand(val)}, ptr addrspace(1) {gep}" + f"{_llvm_syncscope(scope)} {ordering}, align 4" + ) + + def _op_memref_global_flag_wait_eq(self, op: Op) -> None: + """Lower ``memref.global_flag_wait_eq`` to a self-contained spin loop. + + The loop re-issues an ``acquire`` atomic load until the flag matches; + an atomic load cannot be hoisted out of the loop, and the acquire + orders every later global load after the observed flag value. + """ + ptr, idx, expect = op.operands + scope = _sync_attr(op, "scope", _FENCE_SCOPES) + gep = self._fresh("gep") + cur = self._current() + cur.emit( + f" {gep} = getelementptr inbounds i32, ptr addrspace(1) " + f"{self._operand(ptr)}, i32 {self._operand(idx)}" + ) + spin = self._new_block("flag.wait") + cur.emit(f" br label %{spin.label}") + cur.terminated = True + val = self._fresh("flag") + hit = self._fresh("flag.hit") + spin.emit( + f" {val} = load atomic i32, ptr addrspace(1) {gep}" + f"{_llvm_syncscope(scope)} acquire, align 4" + ) + spin.emit(f" {hit} = icmp eq i32 {val}, {self._operand(expect)}") + done = self._new_block("flag.done") + spin.emit(f" br i1 {hit}, label %{done.label}, label %{spin.label}") + spin.terminated = True + def _op_memref_global_atomic_add_pk_bf16(self, op: Op) -> None: """Lower the packed-bf16 atomic add via a generic ``atomicrmw fadd`` on a ``<2 x bfloat>``. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py index a099becd169a..8a2dbdca93c8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py @@ -35,6 +35,10 @@ from typing import Dict, List, Optional, Set from .ir import ( + _FENCE_ORDERINGS, + _FENCE_SCOPES, + _FLAG_STORE_ORDERINGS, + _is_flag_ptr, KernelDef, Op, Region, @@ -120,6 +124,9 @@ def __str__(self) -> str: "tile.mma": ["op_id"], "scf.yield": ["num"], "tile.inline_asm": ["template", "constraints"], + "memref.fence": ["scope", "ordering"], + "memref.global_flag_store": ["scope", "ordering"], + "memref.global_flag_wait_eq": ["scope"], } @@ -381,6 +388,41 @@ def _check_contract(self, op: Op, scope: Dict[str, Value]) -> None: elif name == "cf.return": if op.operands or op.results: self.err(f"{name}: must have no operands and no results", op) + elif name == "memref.fence": + if op.operands or op.results: + self.err(f"{name}: must have no operands and no results", op) + self._check_sync_attrs(op, _FENCE_ORDERINGS) + elif name in ("memref.global_flag_store", "memref.global_flag_wait_eq"): + if len(op.operands) != 3 or op.results: + self.err( + f"{name}: expected 3 operands / 0 results, got " + f"{len(op.operands)} / {len(op.results)}", + op, + ) + return + ptr, idx, val = op.operands + if not _is_flag_ptr(ptr.type): + self.err( + f"{name}: expected a global i32 pointer, got {ptr.type.name}", op + ) + for what, v in (("index", idx), ("value", val)): + if v.type.name != "i32": + self.err(f"{name}: {what} must be i32, got {v.type.name}", op) + orderings = ( + _FLAG_STORE_ORDERINGS if name == "memref.global_flag_store" else None + ) + self._check_sync_attrs(op, orderings) + + def _check_sync_attrs(self, op: Op, orderings) -> None: + """Scope / ordering values of a fence or flag op (presence is checked + by the required-attrs pass, so a missing key is not reported twice).""" + scope = op.attrs.get("scope") + if isinstance(scope, str) and scope not in _FENCE_SCOPES: + self.err(f"{op.name}: unknown scope '{scope}'", op) + ordering = op.attrs.get("ordering") + if orderings is not None and isinstance(ordering, str): + if ordering not in orderings: + self.err(f"{op.name}: unknown ordering '{ordering}'", op) # ---- scf.for ---- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py index 87cd8796f8c7..42c27763368d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py @@ -346,13 +346,18 @@ ) from .schedule import SchedulePolicy from .streamk import ( + StreamKIterArgs, + StreamKIterPartition, StreamKPartition, StreamKReductionStrategy, compute_streamk_grid_size, emit_streamk_decode, - emit_streamk_partial_load_accumulate, - emit_streamk_partial_store, + emit_streamk_iter_range, + emit_streamk_sk_start_iter, + streamk_end_iter, + streamk_iter_partition, streamk_num_macro_tiles, + streamk_start_iter, ) from .spec import ( IOSpecRule, @@ -611,12 +616,17 @@ "load_sorted_topk_weight", "scatter_token_offset", # StreamK partitioner ( StreamK GEMM) + "StreamKIterArgs", + "StreamKIterPartition", "StreamKPartition", "StreamKReductionStrategy", "compute_streamk_grid_size", "emit_streamk_decode", - "emit_streamk_partial_load_accumulate", - "emit_streamk_partial_store", + "emit_streamk_iter_range", + "emit_streamk_sk_start_iter", + "streamk_end_iter", + "streamk_iter_partition", + "streamk_start_iter", "mfma_attention_bwd_dot_do_o_inner_body", "mfma_attention_bwd_dq_dk_dv_inner_body", "streamk_num_macro_tiles", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py index 587bffd19c93..3fb1428fdae6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py @@ -21,11 +21,12 @@ and is the path CK Tile defaults to. This module ships the partitioner math (a pure-Python view) and the -IR-side glue that -``instances/streamk_gemm.py`` uses. The kernel-shape primitives stay +IR-side glue that ``instances/streamk_gemm.py`` and the stream-K +implicit-GEMM wgrad conv kernel use. The kernel-shape primitives stay in :mod:`rocke.helpers.atoms` / :mod:`rocke.helpers.loads` / :mod:`rocke.helpers.epilogues`; this file is purely about -*partitioning* the work, not about running the per-tile GEMM. +*partitioning* the work, not about running the per-tile GEMM or its +cross-CTA fixup. What we ship today: @@ -37,12 +38,17 @@ * :func:`emit_streamk_decode` -- IR-side decode of a linear partition id into the SSA ``(m_tile, n_tile, k_iter, is_first, is_last)`` bundle the kernel needs. - -v1 implements the ``Atomic`` strategy end-to-end and the -``Reduction`` strategy's *decode* surface (so the spec layer can -emit the right kernel shape); the per-tile reduction pass that the -``Reduction`` strategy needs lands with the StreamK GEMM kernel -itself in . +* :class:`StreamKIterPartition` and friends -- CK Tile's + iteration-balanced (DP + stream-K) partitioner, plus the SSA emitters + that map a CTA onto its MAC-iteration range. + +The ``Reduction`` strategy's cross-CTA fixup (partial slots + a +per-(CTA, wave) flag table, Linear or Tree fan-in) is kernel-specific +and lives with its kernel; the reference implementation is the nested +``_emit_streamk_tiles`` of ``build_implicit_gemm_conv_wgrad`` in +``library/kernels/common/conv_implicit_gemm_wgrad.py`` (C++ twin: +``wgrad_emit_streamk`` in ``conv_implicit_gemm_wgrad.cpp``). +:meth:`StreamKIterPartition.flags_bytes` sizes its flag table. """ from __future__ import annotations @@ -55,13 +61,18 @@ __all__ = [ + "StreamKIterArgs", + "StreamKIterPartition", "StreamKPartition", "StreamKReductionStrategy", "compute_streamk_grid_size", "emit_streamk_decode", - "emit_streamk_partial_load_accumulate", - "emit_streamk_partial_store", + "emit_streamk_iter_range", + "emit_streamk_sk_start_iter", + "streamk_end_iter", + "streamk_iter_partition", "streamk_num_macro_tiles", + "streamk_start_iter", ] @@ -174,62 +185,245 @@ def emit_streamk_decode( ) -def emit_streamk_partial_store( - b: IRBuilder, - *, - workspace: Value, - flag_table: Value, - workspace_off: Value, - flag_off: Value, - value: Value, - is_first: Value, -) -> None: - """Cooperative ``Reduction`` strategy partial-store + signal. - - P37: each contributing CTA stores its partial K-sum at the - ``(m_tile, n_tile, k_iter)`` slot in the workspace and bumps the - flag-table counter for ``(m_tile, n_tile)`` via ``atomic_add(+1)``. - The first contributor's store is unconditionally a fresh write - (``is_first=True``); later contributors atomic-add into the - workspace too so the order doesn't matter. - - Reference: CK Tile ``streamk_common.hpp::SignalStorePartialDone`` - + ``streamk_gemm_kernel.hpp:448-504`` (Tree-reduction). +# --------------------------------------------------------------------------- +# Iteration-balanced stream-K (CK Tile StreamKTilePartitioner) +# --------------------------------------------------------------------------- +# +# The partition above hands every CTA exactly one ``tile_k`` slice. CK Tile's +# production partitioner instead balances *iterations*: output tiles that +# divide evenly across the CTA pool run data-parallel (one CTA per tile, the +# full K loop), and only the remainder is spread over ``max_active_wgs`` +# stream-K CTAs, each owning a contiguous range of MAC-loop iterations that +# may straddle output-tile boundaries. rocKE compiles one kernel per shape, so +# the whole partition is resolved host-side and folded into the IR as +# constants; the device only maps its linear CTA id onto an iteration range. + + +def _ceil_div(a: int, b: int) -> int: + return -(-a // b) + + +@dataclass(frozen=True) +class StreamKIterPartition: + """CK Tile ``StreamKTilePartitionerBase``, resolved at build time. + + ``m_tiles`` / ``n_tiles`` count output tiles (a caller folding a batch + dimension into M passes the folded count). ``iters_per_tile`` is the + number of ``tile_k`` MAC iterations per output tile. ``max_active_wgs`` + is the CTA pool the stream-K remainder is spread over. It need not fit on + the device at once: a CTA only waits on higher-index CTAs and publishes + its own partial before it waits, so with workgroups dispatched in + ascending order the waiting CTAs cannot starve their producers. """ - # Initial store: lane-0-only write (the workspace is per-lane sized - # already in the caller's view) — the consumer in - # :func:`emit_streamk_partial_load_accumulate` reads each slot once. - with b.scf_if(is_first): - b.global_store(workspace, workspace_off, value, align=4) - # Non-first contributors: atomic add into the same f32 slot (the - # ``is_last`` reducer reads the converged value). - with b.scf_if(b.lnot(is_first)): - b.global_atomic_add(workspace, workspace_off, value) - # Bump the flag counter so the last-finishing CTA can detect "I'm - # the reducer" via flag == k_iters_per_tile. - b.global_atomic_add(flag_table, flag_off, b.const_i32(1)) - - -def emit_streamk_partial_load_accumulate( - b: IRBuilder, + + m_tiles: int + n_tiles: int + iters_per_tile: int + max_active_wgs: int + persistent: bool = False + + @property + def num_tiles(self) -> int: + return self.m_tiles * self.n_tiles + + @property + def sk_tiles(self) -> int: + """Output tiles handled by stream-K CTAs (0 means all-DP). + + CK's "DP + 2-tile SK": when the tiles do not divide the CTA pool, + one full wave of tiles plus the remainder go stream-K so the last + wave is balanced; a remainder too small to give every stream-K CTA + at least one iteration falls back to all-DP. + """ + rem = self.num_tiles % self.max_active_wgs + if rem == 0: + return 0 + if self.num_tiles > self.max_active_wgs: + tiles = self.max_active_wgs + rem + else: + tiles = self.num_tiles + if tiles * self.iters_per_tile < self.max_active_wgs: + return 0 + return tiles + + @property + def sk_ctas(self) -> int: + return self.max_active_wgs if self.sk_tiles else 0 + + @property + def total_sk_iters(self) -> int: + return self.sk_tiles * self.iters_per_tile + + @property + def iters_per_sk_cta(self) -> int: + """Base MAC iterations per stream-K CTA; always >= 1. + + Kernels divide by this value (the owner lookup) on every path, + including an all-DP launch, so it is clamped to 1 when there are no + stream-K CTAs. With ``sk_ctas > 0`` it is already >= 1 because + :attr:`sk_tiles` falls back to all-DP otherwise. + """ + if not self.sk_ctas: + return 1 + return self.total_sk_iters // self.sk_ctas + + @property + def extra_iters(self) -> int: + """Stream-K CTAs ``[0, extra_iters)`` own one extra iteration.""" + return self.total_sk_iters % self.sk_ctas if self.sk_ctas else 0 + + @property + def dp_tiles(self) -> int: + return self.num_tiles - self.sk_tiles + + @property + def total_dp_iters(self) -> int: + return self.dp_tiles * self.iters_per_tile + + @property + def grid_size(self) -> int: + """Launch grid (x): the CTA pool when persistent, else DP + SK CTAs.""" + if self.persistent: + return self.max_active_wgs + return self.dp_tiles + self.sk_ctas + + @property + def max_linear_partners(self) -> int: + """Upper bound on the CTAs a tile owner has to fold in (Linear). + + The owner is whichever CTA covers the tile's first iteration; it may + have started in the previous tile, so it is only guaranteed one + iteration here. Every later contributor covers at least + ``iters_per_sk_cta`` iterations until the tile ends. + """ + if not self.sk_ctas: + return 0 + return _ceil_div(self.iters_per_tile - 1, self.iters_per_sk_cta) + + @property + def tree_rounds(self) -> int: + """Pairwise fan-in rounds bounding the Tree fixup: ``ceil(log2)`` of + the most CTAs that can contribute to one tile.""" + contributors = self.max_linear_partners + 1 + return (contributors - 1).bit_length() + + def flags_bytes(self, n_waves: int) -> int: + """Bytes of the fixup flag table: one i32 flag per (stream-K CTA, + wave), padded to 256 bytes (CK layout). + + ``n_waves`` is the waves per CTA that publish/wait independently + (``warp_m * warp_n`` for the wgrad kernel); pass 1 for one flag per + CTA. 0 for an all-DP partition. + """ + if n_waves <= 0: + raise ValueError(f"stream-K flags: n_waves must be > 0 (got {n_waves})") + return _ceil_div(4 * self.sk_ctas * n_waves, 256) * 256 + + +def streamk_iter_partition( *, - workspace: Value, - workspace_off: Value, - is_last: Value, + m_tiles: int, + n_tiles: int, + iters_per_tile: int, + max_active_wgs: int, + persistent: bool = False, +) -> StreamKIterPartition: + """Validated constructor for :class:`StreamKIterPartition`.""" + for name, v in ( + ("m_tiles", m_tiles), + ("n_tiles", n_tiles), + ("iters_per_tile", iters_per_tile), + ("max_active_wgs", max_active_wgs), + ): + if v <= 0: + raise ValueError(f"stream-K partition: {name} must be > 0 (got {v})") + return StreamKIterPartition( + m_tiles=m_tiles, + n_tiles=n_tiles, + iters_per_tile=iters_per_tile, + max_active_wgs=max_active_wgs, + persistent=persistent, + ) + + +def streamk_start_iter(part: StreamKIterPartition, sk_cta: int) -> int: + """Global first MAC iteration of stream-K CTA ``sk_cta``.""" + return ( + part.total_dp_iters + + sk_cta * part.iters_per_sk_cta + + min(sk_cta, part.extra_iters) + ) + + +def streamk_end_iter(part: StreamKIterPartition, sk_cta: int) -> int: + """One past the last MAC iteration of stream-K CTA ``sk_cta``.""" + return ( + streamk_start_iter(part, sk_cta) + + part.iters_per_sk_cta + + (1 if sk_cta < part.extra_iters else 0) + ) + + +class StreamKIterArgs(NamedTuple): + """The runtime stream-K partition a shape-agnostic kernel reads. + + The host computes a :class:`StreamKIterPartition` per launch and passes + these fields as kernel arguments; the device emitters below only ever + see them as SSA values. + """ + + iters_per_tile: Value # i32 MAC iterations per output tile + dp_tiles: Value # i32 tiles owned whole by data-parallel CTAs + total_dp_iters: Value # i32 dp_tiles * iters_per_tile + iters_per_sk_cta: Value # i32 base iterations per stream-K CTA (>= 1) + extra_iters: Value # i32 stream-K CTAs that own one more iteration + + +def emit_streamk_sk_start_iter( + b: IRBuilder, sk_cta: Value, args: StreamKIterArgs ) -> Value: - """Read the converged f32 partial sum at ``workspace_off`` once - every contributor has signalled (``is_last == True`` is the - caller's responsibility). - - The flag-table-based wait protocol is a busy-loop on - ``atomic_load(flag_table[tile_id]) == k_iters_per_tile``; we expose - it as a separate helper because the busy-loop shape differs - between Linear / Tree reduction. For now the Atomic strategy is - the canonical path; this helper plus - :func:`emit_streamk_partial_store` give callers the surface - needed to opt into the Reduction strategy. - - Reference: CK Tile ``streamk_common.hpp:34-73`` - (``WaitStorePartialDone``). + """SSA :func:`streamk_start_iter` for a runtime stream-K CTA index.""" + body = b.mul(sk_cta, args.iters_per_sk_cta) + lead = b.smin(sk_cta, args.extra_iters) + return b.add(args.total_dp_iters, b.add(body, lead)) + + +class _IterRange(NamedTuple): + """The SSA bundle returned by :func:`emit_streamk_iter_range`.""" + + start: Value # i32 global first iteration (SGPR) + end: Value # i32 one past the last iteration (SGPR) + sk_cta: Value # i32 stream-K CTA index; meaningless for a DP CTA + + +def emit_streamk_iter_range( + b: IRBuilder, cta: Value, args: StreamKIterArgs, *, dp: bool +) -> _IterRange: + """Map a CTA onto its contiguous range of global MAC iterations. + + ``dp=True`` (non-persistent launch): CTAs ``[0, dp_tiles)`` each own one + whole tile and the rest are stream-K CTAs, selected branch-free so the + GEMM body is emitted once. ``dp=False``: ``cta`` already is a stream-K + CTA index (the persistent launch runs its DP sweep separately). """ - return b.global_load_f32(workspace, workspace_off) + if dp: + sk_cta = b.sub(cta, args.dp_tiles) + else: + sk_cta = cta + sk_start = emit_streamk_sk_start_iter(b, sk_cta, args) + c_one = b.const_i32(1) + c_zero = b.const_i32(0) + sk_len = b.add( + args.iters_per_sk_cta, + b.select(b.cmp_lt(sk_cta, args.extra_iters), c_one, c_zero), + ) + if dp: + is_dp = b.cmp_lt(cta, args.dp_tiles) + start = b.select(is_dp, b.mul(cta, args.iters_per_tile), sk_start) + length = b.select(is_dp, args.iters_per_tile, sk_len) + else: + start = sk_start + length = sk_len + start = b.to_sgpr_u32(start) + end = b.to_sgpr_u32(b.add(start, length)) + return _IterRange(start=start, end=end, sk_cta=sk_cta) 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 8ae2792ef2a4..9c3377b3c2f6 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 @@ -73,6 +73,8 @@ See the [packed FP6 input contract](../examples/gfx1250/gemm/FP6.md). | `conv_implicit_gemm_auto` | ✅ | ✅ | ❌ | MFMA-specialized autotuned path (raw `MfmaAtom`, K=32 kpack); not ported to WMMA | | `direct_conv_16c` | ❌ | ✅ | ❌ | `fold_k32` needs 16x16x32 atom (CDNA4) | | `direct_conv_4c` | ✅ | ✅ | ❌ | 4x4x4 MFMA atom not in WMMA catalog | +| `conv_implicit_gemm_wgrad` | ✅ | ✅ | ✅ | gfx1151/gfx1201/gfx1250: WMMA, `default` epilogue, `split_k=1` | +| `conv_implicit_gemm_wgrad` stream-K | ✅ | ✅ | ❌ | CK iteration-balanced partition; `linear`/`tree` deterministic fixups need the cross-workgroup flag 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_depthwise_tiled` | ✅ | ✅ | ❌ | `cpg=kpg=1`, filters up to 31x31, any padding; scalar FMA, wave64. Built from `library/kernels`; not yet wired into dispatch | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt index 081bedc9f2c1..246352ff980a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt @@ -279,6 +279,47 @@ if(BUILD_TESTING) rocke_moe_fused_mega_spec_guards) endif() + # Host-side stream-K partition math: the launch grid it sizes is invisible to + # the byte-identity gate, so the C port is pinned to the Python reference + # vectors here (mirrors tests/instances/test_streamk_iter_partition.py). + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/instances/streamk_iter_partition.cpp") + add_executable(rocke_streamk_iter_partition instances/streamk_iter_partition.cpp) + target_link_libraries(rocke_streamk_iter_partition PRIVATE rocke_core) + if(UNIX) + target_link_libraries(rocke_streamk_iter_partition PRIVATE m) + endif() + add_test(NAME rocke_streamk_iter_partition COMMAND rocke_streamk_iter_partition) + set_tests_properties(rocke_streamk_iter_partition PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_streamk_iter_partition) + set_target_properties(rocke_streamk_iter_partition + PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_streamk_iter_partition + RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS + rocke_streamk_iter_partition) + endif() + + # Stream-K wgrad host helpers and validator gates (auto CTA pool per arch, + # grid, workspace layout, reject reasons): also invisible to the byte-identity + # gate, and the Python-vs-C pytest needs the pybind module CI does not build, + # so the C port is pinned to Python reference vectors here. + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/instances/wgrad_streamk_host.cpp") + add_executable(rocke_wgrad_streamk_host instances/wgrad_streamk_host.cpp) + target_link_libraries(rocke_wgrad_streamk_host PRIVATE rocke_core) + if(UNIX) + target_link_libraries(rocke_wgrad_streamk_host PRIVATE m) + endif() + add_test(NAME rocke_wgrad_streamk_host COMMAND rocke_wgrad_streamk_host) + set_tests_properties(rocke_wgrad_streamk_host PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_wgrad_streamk_host) + set_target_properties(rocke_wgrad_streamk_host + PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_wgrad_streamk_host + RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS + rocke_wgrad_streamk_host) + endif() + # Apply YAML-driven quick/standard/comprehensive/full tier labels to the # tests above. Must run in THIS directory scope: CMake (pre-3.28, # which this project's cmake_minimum_required predates) only allows diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp index d2f9209d7039..2e715e38b0e0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include "rocke/arch_target.h" @@ -24,6 +25,7 @@ #include "rocke/lower_hip.h" #include "rocke/lower_llvm.h" #include "rocke/strbuf.h" +#include "rocke/verify.h" namespace { @@ -607,6 +609,285 @@ void case_s_wqm() EXPECT_IR(i64_ir, "call i64 @llvm.amdgcn.s.wqm.i64.i64(i64 15)"); } +/* ---- cross-workgroup fence / flag ops (stream-K fixup) ---- */ +void case_flag_ops() +{ + const std::string ir = lower_one("flagwt", [](rocke_ir_builder_t* b) { + rocke_value_t* f = global_ptr_param(b, "flags", rocke_i32()); + rocke_b_global_flag_store( + b, f, rocke_b_const_i32(b, 3), rocke_b_const_i32(b, 1), NULL, NULL); + rocke_b_fence(b, "agent", "release"); + rocke_b_global_flag_wait_eq(b, f, rocke_b_const_i32(b, 2), rocke_b_const_i32(b, 1), NULL); + }); + EXPECT_IR(ir, + " %gep.1 = getelementptr inbounds i32, ptr addrspace(1) %flags, i32 3\n" + " store atomic i32 1, ptr addrspace(1) %gep.1 syncscope(\"agent\") release, " + "align 4\n" + " fence syncscope(\"agent\") release\n" + " %gep.2 = getelementptr inbounds i32, ptr addrspace(1) %flags, i32 2\n" + " br label %flag.wait.1\n" + "flag.wait.1:\n" + " %flag.3 = load atomic i32, ptr addrspace(1) %gep.2 syncscope(\"agent\") " + "acquire, align 4\n" + " %flag.hit.4 = icmp eq i32 %flag.3, 1\n" + " br i1 %flag.hit.4, label %flag.done.2, label %flag.wait.1\n" + "flag.done.2:\n"); +} + +void case_fence_system_scope() +{ + const std::string ir = lower_one( + "fencesys", [](rocke_ir_builder_t* b) { rocke_b_fence(b, "system", "acquire"); }); + EXPECT_IR(ir, " fence acquire\n"); + EXPECT_NO_IR(ir, "syncscope"); +} + +void case_flag_ops_hip() +{ + const std::string hip = lower_one_hip( + "flaghip", + [](rocke_ir_builder_t* b) { + rocke_value_t* f = global_ptr_param(b, "flags", rocke_i32()); + rocke_b_global_flag_store( + b, f, rocke_b_const_i32(b, 0), rocke_b_const_i32(b, 1), NULL, NULL); + rocke_b_fence(b, "agent", "acquire"); + rocke_b_global_flag_wait_eq( + b, f, rocke_b_const_i32(b, 0), rocke_b_const_i32(b, 1), NULL); + }, + "gfx950"); + EXPECT_IR(hip, "__ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_AGENT);"); + EXPECT_IR(hip, "__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, \"agent\");"); + EXPECT_IR(hip, "while (__hip_atomic_load("); +} + +/* Raw op with a bad or missing scope/ordering, as serialized IR could carry + * past the builder's checks: both lowerers must reject it (ROCKE_ERR_VALUE), + * never pick a default or render it into invalid IR. Every other attr is + * valid, so the rejection is attributable to `key`; value == nullptr drops + * the key. Mirrors test_rocke.py test_flag_ops_lowerers_reject_bad_attrs. */ +void expect_flag_attr_rejected(rocke_opcode_t kind, const char* key, const char* value) +{ + rocke_ir_builder_t b; + rocke_ir_builder_init(&b, "flag_bad_attr"); + rocke_value_t* operands[3]; + int n = 0; + if(kind != ROCKE_OP_MEMREF_FENCE) + { + operands[n++] = global_ptr_param(&b, "flags", rocke_i32()); + operands[n++] = rocke_b_const_i32(&b, 0); + operands[n++] = rocke_b_const_i32(&b, 1); + } + const char* scope = "agent"; + const char* ordering = kind == ROCKE_OP_MEMREF_FENCE ? "acq_rel" + : kind == ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE ? "release" + : nullptr; + if(!strcmp(key, "scope")) + scope = value; + else + ordering = value; + rocke_attr_map_t attrs; + rocke_attr_map_init(&attrs); + if(scope) + rocke_attr_set_str(&b, &attrs, "scope", scope); + if(ordering) + rocke_attr_set_str(&b, &attrs, "ordering", ordering); + rocke_b_op(&b, kind, operands, n, nullptr, 0, &attrs, nullptr, 0, nullptr, nullptr); + rocke_b_ret(&b); + + rocke_strbuf_t out; + rocke_strbuf_init(&out, 256); + rocke_lower_hip_opts_t opts{}; + opts.arch = "gfx950"; + const rocke_status_t st + = rocke_lower_kernel_to_hip(&b, rocke_ir_builder_kernel(&b), &opts, &out); + if(st != ROCKE_ERR_VALUE) + fail("HIP lowering must reject a missing/unknown flag/fence scope or ordering", __LINE__); + rocke_strbuf_free(&out); + + char* ll = nullptr; + char err[ROCKE_ERR_MSG_CAP]; + err[0] = '\0'; + const rocke_status_t lst = rocke_lower_kernel_to_llvm_ex( + rocke_ir_builder_kernel(&b), ROCKE_LLVM_FLAVOR_AUTO, "gfx950", &ll, err, sizeof(err)); + if(lst != ROCKE_ERR_VALUE) + fail("LLVM lowering must reject a missing/unknown flag/fence scope or ordering", __LINE__); + char want[160]; + if(value) + snprintf(want, sizeof(want), "unknown %s '%s'", key, value); + else + snprintf(want, sizeof(want), "%s attribute missing or not a string", key); + if(!strstr(err, want)) + { + std::fprintf(stderr, "LLVM error '%s' lacks '%s'\n", err, want); + fail("LLVM lowering rejected a flag/fence attr with the wrong message", __LINE__); + } + std::free(ll); + rocke_ir_builder_free(&b); +} + +/* A flag op on a non-i32 pointer: the LLVM lowerer would GEP i32 while HIP + * indexes the pointee, so the builder and the verifier both refuse it. */ +void case_flag_ops_reject_non_i32_ptr() +{ + for(int wait = 0; wait < 2; ++wait) + { + rocke_ir_builder_t b; + rocke_ir_builder_init(&b, "flag_f16_ptr"); + const char* want = wait ? "global_flag_wait_eq expects a global i32 pointer, got " + "ptr" + : "global_flag_store expects a global i32 pointer, got " + "ptr"; + /* The builder raises (ckc::ValueError) or latches the status. */ + std::string got; + try + { + rocke_value_t* f = global_ptr_param(&b, "flags", rocke_f16()); + if(wait) + rocke_b_global_flag_wait_eq( + &b, f, rocke_b_const_i32(&b, 0), rocke_b_const_i32(&b, 1), NULL); + else + rocke_b_global_flag_store( + &b, f, rocke_b_const_i32(&b, 0), rocke_b_const_i32(&b, 1), NULL, NULL); + if(rocke_ir_builder_status(&b) == ROCKE_ERR_VALUE && rocke_ir_builder_error(&b)) + got = rocke_ir_builder_error(&b); + } + catch(const std::exception& e) + { + got = e.what(); + } + if(got != want) + { + std::fprintf(stderr, "builder error '%s', want '%s'\n", got.c_str(), want); + fail("flag-op builder must reject a non-i32 pointer", __LINE__); + } + rocke_ir_builder_free(&b); + } + + rocke_ir_builder_t b; + rocke_ir_builder_init(&b, "flag_f16_ptr_verify"); + rocke_value_t* operands[] = {global_ptr_param(&b, "flags", rocke_f16()), + rocke_b_const_i32(&b, 0), + rocke_b_const_i32(&b, 1)}; + const rocke_opcode_t kinds[] + = {ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ}; + for(rocke_opcode_t kind : kinds) + { + rocke_attr_map_t attrs; + rocke_attr_map_init(&attrs); + rocke_attr_set_str(&b, &attrs, "scope", "agent"); + if(kind == ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE) + rocke_attr_set_str(&b, &attrs, "ordering", "release"); + rocke_b_op(&b, kind, operands, 3, nullptr, 0, &attrs, nullptr, 0, nullptr, nullptr); + } + rocke_b_ret(&b); + rocke_diag_t* diags = nullptr; + size_t n = 0; + rocke_verify(rocke_ir_builder_kernel(&b), &diags, &n); + bool saw_store = false, saw_wait = false; + for(size_t i = 0; i < n; ++i) + { + const std::string m = diags[i].message ? diags[i].message : ""; + saw_store |= m.find("memref.global_flag_store: expected a global i32 pointer, got " + "ptr") + != std::string::npos; + saw_wait |= m.find("memref.global_flag_wait_eq: expected a global i32 pointer, got " + "ptr") + != std::string::npos; + } + rocke_diags_free(diags, n); + if(!saw_store || !saw_wait) + fail("verifier must reject a flag op on a non-i32 pointer", __LINE__); + rocke_ir_builder_free(&b); +} + +/* The verifier rejects what the builder would have refused (serialized IR), + * and the LLVM lowerer refuses a flag op it cannot unpack. */ +void case_flag_ops_verify_and_llvm_arity() +{ + rocke_ir_builder_t b; + rocke_ir_builder_init(&b, "flag_verify"); + rocke_value_t* f = global_ptr_param(&b, "flags", rocke_i32()); + rocke_value_t* two[] = {f, rocke_b_const_i32(&b, 0)}; + rocke_attr_map_t fence_attrs; + rocke_attr_map_init(&fence_attrs); + rocke_attr_set_str(&b, &fence_attrs, "scope", "agent"); + rocke_attr_set_str(&b, &fence_attrs, "ordering", "monotonic"); + rocke_op_t* fence = rocke_b_op(&b, + ROCKE_OP_MEMREF_FENCE, + nullptr, + 0, + nullptr, + 0, + &fence_attrs, + nullptr, + 0, + nullptr, + nullptr); + rocke_attr_map_t store_attrs; + rocke_attr_map_init(&store_attrs); + rocke_attr_set_str(&b, &store_attrs, "scope", "agent"); + rocke_attr_set_str(&b, &store_attrs, "ordering", "release"); + rocke_b_op(&b, + ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, + two, + 2, + nullptr, + 0, + &store_attrs, + nullptr, + 0, + nullptr, + nullptr); + rocke_b_ret(&b); + + rocke_diag_t* diags = nullptr; + size_t n = 0; + rocke_verify(rocke_ir_builder_kernel(&b), &diags, &n); + bool saw_ordering = false, saw_arity = false; + for(size_t i = 0; i < n; ++i) + { + const std::string m = diags[i].message ? diags[i].message : ""; + saw_ordering |= m.find("memref.fence: unknown ordering 'monotonic'") != std::string::npos; + saw_arity |= m.find("memref.global_flag_store: expected 3 operands / 0 results, got 2 / 0") + != std::string::npos; + } + rocke_diags_free(diags, n); + if(!saw_ordering) + fail("verifier must reject a fence ordering outside the builder's set", __LINE__); + if(!saw_arity) + fail("verifier must reject a two-operand flag store", __LINE__); + + /* Make the fence valid so the LLVM rejection is the flag store's arity. */ + rocke_attr_set_str(&b, &fence->attrs, "ordering", "acq_rel"); + char* ll = nullptr; + char err[ROCKE_ERR_MSG_CAP]; + err[0] = '\0'; + const rocke_status_t st = rocke_lower_kernel_to_llvm_ex( + rocke_ir_builder_kernel(&b), ROCKE_LLVM_FLAVOR_AUTO, "gfx950", &ll, err, sizeof(err)); + if(st != ROCKE_ERR_VALUE || !strstr(err, "memref.global_flag_store expects three operands")) + fail("LLVM lowering must reject a two-operand flag store", __LINE__); + std::free(ll); + rocke_ir_builder_free(&b); +} + +void case_flag_ops_lowerers_reject_bad_attrs() +{ + expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "ordering", "bogus"); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "scope", "bogus"); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "ordering", "bogus"); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "scope", "bogus"); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, "scope", "bogus"); + /* A load-only ordering is not a fence / store ordering. */ + expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "ordering", "monotonic"); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "ordering", "acquire"); + /* Missing attrs are rejected, not defaulted. */ + expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "ordering", nullptr); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "scope", nullptr); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "ordering", nullptr); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "scope", nullptr); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, "scope", nullptr); +} + /* ---- av.load / av.store (agent-scope 128-bit vector mem) ---- */ void case_av_load_b128() { @@ -1019,6 +1300,12 @@ const TestCase k_cases[] = { {"hip_zext_uses_unsigned_source_cast", case_hip_zext_uses_unsigned_source_cast}, {"gfx1250_scaled_wmma", case_gfx1250_scaled_wmma}, {"gfx1250_standalone_bridge", case_gfx1250_standalone_bridge}, + {"flag_ops", case_flag_ops}, + {"fence_system_scope", case_fence_system_scope}, + {"flag_ops_hip", case_flag_ops_hip}, + {"flag_ops_lowerers_reject_bad_attrs", case_flag_ops_lowerers_reject_bad_attrs}, + {"flag_ops_verify_and_llvm_arity", case_flag_ops_verify_and_llvm_arity}, + {"flag_ops_reject_non_i32_ptr", case_flag_ops_reject_non_i32_ptr}, {"opcode_names_are_aligned", case_opcode_names_are_aligned}, }; 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 386c3fc0cb9e..eea197985f50 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 @@ -922,6 +922,14 @@ "sha256": "5c0a33026f958a437bea9ad18b33d3e9dbcc9821923b2373e203f3edb1cb0222", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c16k32r3_sk_atomic_pers": { + "arch": "gfx942", + "bytes": 25156, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_sk_atomic_pers_N1H8W8C16_K32Y3X3_t64x32x16_w2x2_a16x16x16_mem_default_sk_skpers_skratomic", + "sha256": "1d70097054f9a71c691f4b89a87c5490e8b2371cc3c243cb7a725f8640744d20", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4": { "arch": "gfx942", "bytes": 22466, @@ -930,6 +938,22 @@ "sha256": "5bcf94f8ea7b81d9f15f2fd7edc5c7523831d876e472ba422b031cf0e4f9e21e", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_tree_pers": { + "arch": "gfx942", + "bytes": 44396, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_tree_pers_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk_skpers_skrtree", + "sha256": "4c3f744ae6f1ae5198f47461f06f90611d4d15fac5322b87147501133956c0c2", + "status": "ok" + }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_workspace": { + "arch": "gfx942", + "bytes": 25798, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_workspace_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk_skrworkspace", + "sha256": "21bee95d9ae4d04fde79971cb14764bc3207d0c63d590a45eb2e1e6340a01b0e", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", "bytes": 26896, @@ -946,6 +970,38 @@ "sha256": "3554714549ab49184e12b1c41aa45d1eb8aea9de8250f926d3674d9b3d8e4b32", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_atomic": { + "arch": "gfx950", + "bytes": 29815, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_atomic_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skratomic", + "sha256": "d9c46d394540a321408c425dc012480fd6a7688f52c697f0cb7c68f0f261e5bb", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_linear": { + "arch": "gfx950", + "bytes": 58937, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_linear_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk", + "sha256": "1819f9e2134239476b7315a2a059a7c2014004af3829fd5a764dbb06fea4b05d", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_tree": { + "arch": "gfx950", + "bytes": 60379, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_tree_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skrtree", + "sha256": "1d2674baa67db756ba52747582b5551d663aaf9e3d56b8069bd83665de9a5485", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_workspace_pers": { + "arch": "gfx950", + "bytes": 31327, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_workspace_pers_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers_skrworkspace", + "sha256": "9d25da3d386646386e6e4075eec6794d76c3eda4da66121976d59a15ee4f0cde", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", "bytes": 26145, @@ -970,6 +1026,14 @@ "sha256": "7ec4ccded2d060144039b3c5a644dcc8d4fd36c3f95c8a79384d205d687a5ba6", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c64k64r3_g4_sk_linear_pers": { + "arch": "gfx950", + "bytes": 61458, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_g4_sk_linear_pers_N1H8W8C64_K64Y3X3G4_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers", + "sha256": "7e23d8219721828777ad6b20909fc105514a4304d1cb638f0779dadafa35eb7f", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c64k64r3_g4_spk4": { "arch": "gfx950", "bytes": 38655, @@ -1428,8 +1492,8 @@ "gfx1201": 7, "gfx1250": 4, "gfx90a": 5, - "gfx942": 36, - "gfx950": 112 + "gfx942": 39, + "gfx950": 117 }, "expected_failures": 4, "families": { @@ -1443,7 +1507,7 @@ "conv_direct_depthwise_tiled": 3, "conv_direct_dgrad": 5, "conv_direct_nongrouped": 7, - "conv_wgrad": 13, + "conv_wgrad": 21, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, "deep_fused_conv": 9, @@ -1453,7 +1517,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 175 + "ok_cases": 183 } }, "llvm22": { @@ -2378,6 +2442,14 @@ "sha256": "a4d404a44b5f8a351ce2f2d24957c976c79d0ee13887ff15c7c7a47438575b82", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c16k32r3_sk_atomic_pers": { + "arch": "gfx942", + "bytes": 25289, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_sk_atomic_pers_N1H8W8C16_K32Y3X3_t64x32x16_w2x2_a16x16x16_mem_default_sk_skpers_skratomic", + "sha256": "f4fa2ad553aa4a1180c5a4a67d3d13cd86f0ee6671bfad4954dd55c93fe67e3d", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4": { "arch": "gfx942", "bytes": 22599, @@ -2386,6 +2458,22 @@ "sha256": "dbf4caa3716ae46dffa53f4ab0ec58b505d5196239c88d0e21216614d9141729", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_tree_pers": { + "arch": "gfx942", + "bytes": 44535, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_tree_pers_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk_skpers_skrtree", + "sha256": "2ef49bc86988a2fe9177d29325c801bb9ec0e7631a3255cc5a6d0fc2c86179be", + "status": "ok" + }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_workspace": { + "arch": "gfx942", + "bytes": 25931, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_workspace_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk_skrworkspace", + "sha256": "266d9829bcf2d91378ebbf7d2840fce771486b21d08c460fda3bbc1f206309e2", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", "bytes": 27035, @@ -2402,6 +2490,38 @@ "sha256": "46e8d22587b5a0f9bb3a0a5fe7d351933bd089e7d2c263535cc7d760631c2886", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_atomic": { + "arch": "gfx950", + "bytes": 29948, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_atomic_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skratomic", + "sha256": "ed4b63d61ea4edfd5c4e4c2a3aaef71fb586c7c0ebd5441ea975d8dfa516ba0d", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_linear": { + "arch": "gfx950", + "bytes": 59076, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_linear_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk", + "sha256": "a18cb4d993ab0e837f5e4cbda76848fce46e76a8393900e755b3cdaed05235f8", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_tree": { + "arch": "gfx950", + "bytes": 60518, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_tree_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skrtree", + "sha256": "efe578b01d1d931ed9ee46a7974a5167cb5032222f20718d28c8e546046d577a", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_workspace_pers": { + "arch": "gfx950", + "bytes": 31460, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_workspace_pers_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers_skrworkspace", + "sha256": "f2b8b993253d99f9d3b796085d5ac5c6ef3d15ae453ec0bd0aee0d6397591c28", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", "bytes": 26284, @@ -2426,6 +2546,14 @@ "sha256": "e24daa9ac08990da4b1d1ad4ccbeece698279dcda588cb28c785ee2d4e21cd38", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c64k64r3_g4_sk_linear_pers": { + "arch": "gfx950", + "bytes": 61597, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_g4_sk_linear_pers_N1H8W8C64_K64Y3X3G4_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers", + "sha256": "f978aa87c489a5989e562217a201cc15675aa1b8fd38a816846881a2697d55cb", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c64k64r3_g4_spk4": { "arch": "gfx950", "bytes": 38794, @@ -2884,8 +3012,8 @@ "gfx1201": 7, "gfx1250": 4, "gfx90a": 5, - "gfx942": 36, - "gfx950": 112 + "gfx942": 39, + "gfx950": 117 }, "expected_failures": 4, "families": { @@ -2899,7 +3027,7 @@ "conv_direct_depthwise_tiled": 3, "conv_direct_dgrad": 5, "conv_direct_nongrouped": 7, - "conv_wgrad": 13, + "conv_wgrad": 21, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, "deep_fused_conv": 9, @@ -2909,7 +3037,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 175 + "ok_cases": 183 } }, "llvm23": { @@ -3834,6 +3962,14 @@ "sha256": "24c106cf812be8655f0e9ed1a272346ecccc29dd2e893342eada4b6dbe9dd227", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c16k32r3_sk_atomic_pers": { + "arch": "gfx942", + "bytes": 25353, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_sk_atomic_pers_N1H8W8C16_K32Y3X3_t64x32x16_w2x2_a16x16x16_mem_default_sk_skpers_skratomic", + "sha256": "2863b1a0e57941df6e773f9e981c40f1843856237554d7615e92731db6ff06c4", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4": { "arch": "gfx942", "bytes": 22663, @@ -3842,6 +3978,22 @@ "sha256": "6b7dca0251498621500e6c6978cc778ec3a4a49e3304631cd1ab1939ce06b6c4", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_tree_pers": { + "arch": "gfx942", + "bytes": 44599, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_tree_pers_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk_skpers_skrtree", + "sha256": "2558f2ea4588e91e12e697b23a095805e6a83b20d9c01d9a7f78be69a7e688ce", + "status": "ok" + }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_workspace": { + "arch": "gfx942", + "bytes": 25995, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_workspace_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk_skrworkspace", + "sha256": "7c6ce1e064d7d92ebcfb0110334c89e674a828bb65bd784d566e5e9b58129721", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", "bytes": 27099, @@ -3858,6 +4010,38 @@ "sha256": "95493c1a74fc7fe192cbf74030f19bd0ecc57dfd28198ebdcb40eac381e134d5", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_atomic": { + "arch": "gfx950", + "bytes": 30012, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_atomic_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skratomic", + "sha256": "9fc28049b1eee41c32d58d6fe3fc2e4aa218e367d3ef90acaa97790dcb4d18a3", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_linear": { + "arch": "gfx950", + "bytes": 59140, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_linear_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk", + "sha256": "65640f53273c121fa0e1127e49cce824a2e7e4f7ec8a6ccb53725570d419b2e9", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_tree": { + "arch": "gfx950", + "bytes": 60582, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_tree_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skrtree", + "sha256": "e50b150b37785afabbc8c8f136c069bdd88ffbbff332d14dbe0f9db82b88c975", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_workspace_pers": { + "arch": "gfx950", + "bytes": 31524, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_workspace_pers_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers_skrworkspace", + "sha256": "3dbc6ebcdf5be7498ae4f73eeb9b99932a9cea7fe80df7077e4d629959f137dc", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", "bytes": 26348, @@ -3882,6 +4066,14 @@ "sha256": "822ab7ae23cca12ac3351b10bec5ee0cc5a18945c210f3f39e1c8f38b76252d0", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c64k64r3_g4_sk_linear_pers": { + "arch": "gfx950", + "bytes": 61661, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_g4_sk_linear_pers_N1H8W8C64_K64Y3X3G4_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers", + "sha256": "1b892d27711fae5606fa0bf1eca307a105ab3a9db98a9860d46e72db11937992", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c64k64r3_g4_spk4": { "arch": "gfx950", "bytes": 38858, @@ -4355,8 +4547,8 @@ "gfx1201": 7, "gfx1250": 8, "gfx90a": 5, - "gfx942": 36, - "gfx950": 112 + "gfx942": 39, + "gfx950": 117 }, "expected_failures": 0, "families": { @@ -4370,7 +4562,7 @@ "conv_direct_depthwise_tiled": 3, "conv_direct_dgrad": 5, "conv_direct_nongrouped": 7, - "conv_wgrad": 13, + "conv_wgrad": 21, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, "deep_fused_conv": 9, @@ -4380,7 +4572,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 179 + "ok_cases": 187 } } }, 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 3082488c96dc..5f3d4edf50dd 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 @@ -178,6 +178,9 @@ def build_conv_wgrad( split_k=1, dtype_d="fp16", two_stage=False, + streamk="off", + streamk_reduction="linear", + streamk_ctas=-1, ): def _build(): from kernels.common.conv_implicit_gemm_wgrad import ( @@ -207,6 +210,9 @@ def _build(): epilogue=epilogue, split_k=split_k, two_stage=two_stage, + streamk=streamk, + streamk_reduction=streamk_reduction, + streamk_ctas=streamk_ctas, ) return build_implicit_gemm_conv_wgrad(spec, arch=arch) @@ -1914,6 +1920,162 @@ def add(family, case_id, arch, build): epilogue="cshuffle", ), ) + # Stream-K wgrad: the linear and tree flag fixups, the atomic reduction, + # the persistent launch and the group fold into GEMM-M. MFMA-only. + add( + "conv_wgrad", + "conv_wgrad/gfx950/n1h8c16k32r3_sk_linear", + "gfx950", + build_conv_wgrad( + "irhash_wgrad_950_sk_linear", + "gfx950", + wgrad1, + wave_size=64, + wtm=32, + wtn=32, + wtk=16, + tile_m=64, + tile_n=64, + tile_k=16, + streamk="dp_sk", + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx950/n1h8c16k32r3_sk_atomic", + "gfx950", + build_conv_wgrad( + "irhash_wgrad_950_sk_atomic", + "gfx950", + wgrad1, + wave_size=64, + wtm=32, + wtn=32, + wtk=16, + tile_m=64, + tile_n=64, + tile_k=16, + dtype_d="fp32", + streamk="dp_sk", + streamk_reduction="atomic", + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_tree_pers", + "gfx942", + build_conv_wgrad( + "irhash_wgrad_942_g4_sk_tree_pers", + "gfx942", + wgrad_g4, + wave_size=64, + wtm=16, + wtn=16, + wtk=16, + tile_m=64, + tile_n=32, + tile_k=16, + streamk="persistent", + streamk_reduction="tree", + ), + ) + # The remaining {dp_sk, persistent} x {linear, tree, atomic, workspace} + # combinations, so every fixup is pinned under both launch modes. + add( + "conv_wgrad", + "conv_wgrad/gfx950/n1h8c16k32r3_sk_tree", + "gfx950", + build_conv_wgrad( + "irhash_wgrad_950_sk_tree", + "gfx950", + wgrad1, + wave_size=64, + wtm=32, + wtn=32, + wtk=16, + tile_m=64, + tile_n=64, + tile_k=16, + streamk="dp_sk", + streamk_reduction="tree", + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_workspace", + "gfx942", + build_conv_wgrad( + "irhash_wgrad_942_g4_sk_workspace", + "gfx942", + wgrad_g4, + wave_size=64, + wtm=16, + wtn=16, + wtk=16, + tile_m=64, + tile_n=32, + tile_k=16, + streamk="dp_sk", + streamk_reduction="workspace", + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx950/n1h8c64k64r3_g4_sk_linear_pers", + "gfx950", + build_conv_wgrad( + "irhash_wgrad_950_g4_sk_linear_pers", + "gfx950", + wgrad_g4, + wave_size=64, + wtm=32, + wtn=32, + wtk=16, + tile_m=64, + tile_n=64, + tile_k=16, + streamk="persistent", + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx942/n1h8c16k32r3_sk_atomic_pers", + "gfx942", + build_conv_wgrad( + "irhash_wgrad_942_sk_atomic_pers", + "gfx942", + wgrad1, + wave_size=64, + wtm=16, + wtn=16, + wtk=16, + tile_m=64, + tile_n=32, + tile_k=16, + dtype_d="fp32", + streamk="persistent", + streamk_reduction="atomic", + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx950/n1h8c16k32r3_sk_workspace_pers", + "gfx950", + build_conv_wgrad( + "irhash_wgrad_950_sk_workspace_pers", + "gfx950", + wgrad1, + wave_size=64, + wtm=32, + wtn=32, + wtk=16, + tile_m=64, + tile_n=64, + tile_k=16, + streamk="persistent", + streamk_reduction="workspace", + ), + ) # Grouped + cshuffle epilogue (MFMA): the LDS-staged store threads the # per-group k_out fold (group*kpg) and derives its store-vec from cpg. add( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp new file mode 100644 index 000000000000..1195cb0989a9 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp @@ -0,0 +1,180 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * tests/instances/streamk_iter_partition.cpp -- C parity for + * tests/instances/test_streamk_iter_partition.py. + * + * The iteration-balanced stream-K partition sizes the launch grid, which the + * byte-identity gate never sees, so a C/Python divergence here would surface + * as a silently wrong launch rather than an IR diff. The reference vectors are + * the Python test's, verbatim; the sweep checks the same invariants the fixup + * relies on. + */ +#include + +#include "rocke/helper_rocke.helpers.streamk.h" + +namespace +{ + +int g_failures = 0; + +#define CHECK(cond) \ + do \ + { \ + if(!(cond)) \ + { \ + fprintf(stderr, "FAIL: %s (%s:%d)\n", #cond, __FILE__, __LINE__); \ + ++g_failures; \ + } \ + } while(0) + +struct Vector +{ + int m, n, ipt, w; + bool persistent; + /* flags_bytes is for n_waves = 1; ipsc is clamped to 1 when all-DP. */ + int sk_tiles, sk_ctas, ipsc, extra, dp_tiles, grid, partners, rounds, flags_bytes; +}; + +const Vector k_vectors[] = { + {4, 4, 8, 16, false, 0, 0, 1, 0, 16, 16, 0, 0, 0}, + {2, 3, 10, 8, false, 6, 8, 7, 4, 0, 8, 2, 2, 256}, + {5, 7, 12, 16, false, 19, 16, 14, 4, 16, 32, 1, 1, 256}, + {1, 3, 2, 8, false, 0, 0, 1, 0, 3, 3, 0, 0, 0}, + {5, 7, 12, 16, true, 19, 16, 14, 4, 16, 16, 1, 1, 256}, + {3, 3, 64, 4, false, 5, 4, 80, 0, 4, 8, 1, 1, 256}, + {7, 5, 3, 16, false, 19, 16, 3, 9, 16, 32, 1, 1, 256}, +}; + +rocke_streamk_iter_partition_t make(int m, int n, int ipt, int w, bool persistent) +{ + rocke_status_t st = ROCKE_OK; + rocke_streamk_iter_partition_t p + = rocke_streamk_iter_partition_make(m, n, ipt, w, persistent, &st); + CHECK(st == ROCKE_OK); + return p; +} + +void check_vectors() +{ + for(const Vector& v : k_vectors) + { + const rocke_streamk_iter_partition_t p = make(v.m, v.n, v.ipt, v.w, v.persistent); + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&p); + CHECK(r.sk_tiles == v.sk_tiles); + CHECK(r.sk_ctas == v.sk_ctas); + CHECK(r.iters_per_sk_cta == v.ipsc); + CHECK(r.extra_iters == v.extra); + CHECK(r.dp_tiles == v.dp_tiles); + CHECK(r.grid_size == v.grid); + CHECK(r.max_linear_partners == v.partners); + CHECK(r.tree_rounds == v.rounds); + CHECK(rocke_streamk_flags_bytes(&r, 1, NULL) == (size_t)v.flags_bytes); + } + const rocke_streamk_iter_partition_t p = make(5, 7, 12, 16, false); + const int starts[] = {192, 207, 222, 237, 252}; + const int ends[] = {207, 222, 237, 252, 266}; + for(int c = 0; c < 5; ++c) + { + CHECK(rocke_streamk_start_iter(&p, c) == starts[c]); + CHECK(rocke_streamk_end_iter(&p, c) == ends[c]); + } +} + +void check_flags_bytes_per_wave() +{ + /* One i32 flag per (stream-K CTA, wave), padded to 256 bytes. */ + struct + { + int m, n, ipt, w, n_waves; + size_t want; + } const cases[] = { + {5, 7, 12, 16, 4, 256}, + {5, 7, 12, 16, 8, 512}, + {2, 3, 10, 8, 9, 512}, + {4, 4, 8, 16, 4, 0}, + }; + for(const auto& c : cases) + { + const rocke_streamk_iter_partition_t p = make(c.m, c.n, c.ipt, c.w, false); + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&p); + rocke_status_t st = ROCKE_ERR_VALUE; + CHECK(rocke_streamk_flags_bytes(&r, c.n_waves, &st) == c.want); + CHECK(st == ROCKE_OK); + } + const rocke_streamk_iter_partition_t p = make(5, 7, 12, 16, false); + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&p); + rocke_status_t st = ROCKE_OK; + CHECK(rocke_streamk_flags_bytes(&r, 0, &st) == 0); + CHECK(st == ROCKE_ERR_VALUE); +} + +void check_rejects_non_positive() +{ + rocke_status_t st = ROCKE_OK; + (void)rocke_streamk_iter_partition_make(0, 2, 4, 8, false, &st); + CHECK(st == ROCKE_ERR_VALUE); + st = ROCKE_OK; + (void)rocke_streamk_iter_partition_make(2, 2, 4, 0, false, &st); + CHECK(st == ROCKE_ERR_VALUE); +} + +void check_invariants(int m, int n, int ipt, int w) +{ + const rocke_streamk_iter_partition_t p = make(m, n, ipt, w, false); + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&p); + CHECK(r.dp_tiles + r.sk_tiles == r.num_tiles); + CHECK(r.iters_per_sk_cta >= 1); /* a kernel divisor even when all-DP */ + if(!r.sk_ctas) + return; + int cursor = r.total_dp_iters; + int prev_first_tile = -1; + int run = 0; /* contributors to the current tile */ + for(int c = 0; c < r.sk_ctas; ++c) + { + const int lo = rocke_streamk_start_iter(&p, c); + const int hi = rocke_streamk_end_iter(&p, c); + CHECK(lo == cursor); + CHECK(hi > lo); + cursor = hi; + const int first = lo / ipt; + const int last = (hi - 1) / ipt; + /* Contributors to `first` are this CTA plus the run carried from the + * previous CTA when it ended inside the same tile. */ + run = (first == prev_first_tile) ? run + 1 : 1; + CHECK(run - 1 <= r.max_linear_partners); + CHECK(run <= (1 << r.tree_rounds)); + if(last != first) + { + run = 1; + prev_first_tile = last; + } + else + prev_first_tile = first; + } + CHECK(cursor == r.total_dp_iters + r.total_sk_iters); +} + +} // namespace + +int main(void) +{ + check_vectors(); + check_flags_bytes_per_wave(); + check_rejects_non_positive(); + const int ipts[] = {1, 2, 3, 5, 8, 13}; + const int ws[] = {1, 2, 3, 4, 7, 8, 16}; + for(int m = 1; m < 7; ++m) + for(int n = 1; n < 5; ++n) + for(int ipt : ipts) + for(int w : ws) + check_invariants(m, n, ipt, w); + if(g_failures) + { + fprintf(stderr, "streamk_iter_partition: %d failure(s)\n", g_failures); + return 1; + } + printf("streamk_iter_partition: OK\n"); + return 0; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py new file mode 100644 index 000000000000..7471cfca79e6 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py @@ -0,0 +1,150 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Host-side math of the iteration-balanced stream-K partitioner. + +The partition drives the launch grid, which lives outside the IR, so the +byte-identity gate cannot see a disagreement here. The reference vectors +below are mirrored verbatim by ``streamk_iter_partition.cpp`` for the C +engine; the invariant sweep checks the properties the fixup relies on. +""" + +from __future__ import annotations + +import unittest + +from rocke.helpers.streamk import ( + StreamKIterPartition, + streamk_end_iter, + streamk_iter_partition, + streamk_start_iter, +) + +# (m_tiles, n_tiles, iters_per_tile, max_active_wgs, persistent) -> +# (sk_tiles, sk_ctas, iters_per_sk_cta, extra_iters, dp_tiles, grid_size, +# max_linear_partners, tree_rounds, flags_bytes(n_waves=1)) +# iters_per_sk_cta is clamped to 1 for an all-DP partition (a kernel divisor). +_VECTORS = ( + ((4, 4, 8, 16, False), (0, 0, 1, 0, 16, 16, 0, 0, 0)), # tiles % W == 0 + ((2, 3, 10, 8, False), (6, 8, 7, 4, 0, 8, 2, 2, 256)), # tiles < W + ((5, 7, 12, 16, False), (19, 16, 14, 4, 16, 32, 1, 1, 256)), # tiles > W + ((1, 3, 2, 8, False), (0, 0, 1, 0, 3, 3, 0, 0, 0)), # too little SK work + ((5, 7, 12, 16, True), (19, 16, 14, 4, 16, 16, 1, 1, 256)), # persistent + ((3, 3, 64, 4, False), (5, 4, 80, 0, 4, 8, 1, 1, 256)), + ((7, 5, 3, 16, False), (19, 16, 3, 9, 16, 32, 1, 1, 256)), +) + + +def _plan(p: StreamKIterPartition) -> tuple: + return ( + p.sk_tiles, + p.sk_ctas, + p.iters_per_sk_cta, + p.extra_iters, + p.dp_tiles, + p.grid_size, + p.max_linear_partners, + p.tree_rounds, + p.flags_bytes(1), + ) + + +def _make(m, n, ipt, w, persistent=False) -> StreamKIterPartition: + return streamk_iter_partition( + m_tiles=m, + n_tiles=n, + iters_per_tile=ipt, + max_active_wgs=w, + persistent=persistent, + ) + + +class TestStreamKIterPartitionVectors(unittest.TestCase): + def test_reference_vectors(self): + for args, want in _VECTORS: + with self.subTest(args=args): + self.assertEqual(_plan(_make(*args)), want) + + def test_start_end_iter_vectors(self): + p = _make(5, 7, 12, 16) + self.assertEqual( + [streamk_start_iter(p, c) for c in range(5)], [192, 207, 222, 237, 252] + ) + self.assertEqual( + [streamk_end_iter(p, c) for c in range(5)], [207, 222, 237, 252, 266] + ) + + def test_flags_bytes_per_wave(self): + # One i32 flag per (stream-K CTA, wave), padded to 256 bytes. + for args, n_waves, want in ( + ((5, 7, 12, 16), 4, 256), # 4 * 16 * 4 = 256 + ((5, 7, 12, 16), 8, 512), + ((2, 3, 10, 8), 9, 512), # 4 * 8 * 9 = 288 + ((4, 4, 8, 16), 4, 0), # all-DP: no flags + ): + with self.subTest(args=args, n_waves=n_waves): + self.assertEqual(_make(*args).flags_bytes(n_waves), want) + with self.assertRaisesRegex(ValueError, "n_waves must be > 0"): + _make(5, 7, 12, 16).flags_bytes(0) + + def test_rejects_non_positive_counts(self): + for kw in ("m_tiles", "n_tiles", "iters_per_tile", "max_active_wgs"): + args = { + "m_tiles": 2, + "n_tiles": 2, + "iters_per_tile": 4, + "max_active_wgs": 8, + } + args[kw] = 0 + with ( + self.subTest(field=kw), + self.assertRaisesRegex(ValueError, f"{kw} must be > 0"), + ): + streamk_iter_partition(**args) + + +class TestStreamKIterPartitionInvariants(unittest.TestCase): + """Exhaustive sweep of small shapes against the fixup's assumptions.""" + + def _check(self, p: StreamKIterPartition) -> None: + ipt = p.iters_per_tile + self.assertEqual(p.dp_tiles + p.sk_tiles, p.num_tiles) + # A kernel divisor even when all-DP. + self.assertGreaterEqual(p.iters_per_sk_cta, 1) + if not p.sk_ctas: + self.assertEqual(p.sk_tiles, 0) + return + # Every SK CTA owns >= 1 iteration and the ranges tile the SK span + # contiguously. + cursor = p.total_dp_iters + contributors = {} + for c in range(p.sk_ctas): + lo, hi = streamk_start_iter(p, c), streamk_end_iter(p, c) + self.assertEqual(lo, cursor) + self.assertGreater(hi, lo) + cursor = hi + tiles = range(lo // ipt, (hi - 1) // ipt + 1) + # A CTA is a non-owner in at most its first tile: one partial slot + # per SK CTA is enough. + non_owner = [t for t in tiles if lo > t * ipt] + self.assertLessEqual(len(non_owner), 1) + for t in tiles: + contributors.setdefault(t, []).append(c) + self.assertEqual(cursor, p.total_dp_iters + p.total_sk_iters) + for t, ctas in contributors.items(): + # Contributors to one tile are consecutive SK CTAs. + self.assertEqual(ctas, list(range(ctas[0], ctas[0] + len(ctas)))) + self.assertLessEqual(len(ctas) - 1, p.max_linear_partners) + self.assertLessEqual(len(ctas), 1 << p.tree_rounds) + + def test_sweep(self): + for m in range(1, 7): + for n in range(1, 5): + for ipt in (1, 2, 3, 5, 8, 13): + for w in (1, 2, 3, 4, 7, 8, 16): + with self.subTest(m=m, n=n, ipt=ipt, w=w): + self._check(_make(m, n, ipt, w)) + + +if __name__ == "__main__": + unittest.main() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp new file mode 100644 index 000000000000..74a5fd16e02d --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp @@ -0,0 +1,404 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * tests/instances/wgrad_streamk_host.cpp -- C parity for the stream-K wgrad + * host helpers and validator gates of instance_conv_implicit_gemm_wgrad.h. + * + * The launch grid, workspace layout and auto CTA pool these helpers compute + * never reach the emitted IR, so the byte-identity gate cannot see a C/Python + * divergence in them; it would surface as skipped work or an undersized + * buffer. The Python-vs-C pytest (library/tests/test_conv_wgrad_streamk.py) + * needs the rocke_engine binding, which CI does not build, so the vectors are + * pinned here as well. + * + * Every expected value below was produced by the Python implementation + * (kernels.common.conv_implicit_gemm_wgrad: wgrad_streamk_default_pool, + * wgrad_streamk_launch(...).pool, wgrad_streamk_partition, wgrad_streamk_grid, + * wgrad_streamk_workspace_layout, wgrad_streamk_workspace_nbytes and + * is_valid_wgrad_spec) for the same specs; regenerate them from there if the + * Python reference changes. + */ +#include +#include + +#include "rocke/instance_conv_implicit_gemm_wgrad.h" + +namespace +{ + +int g_failures = 0; + +#define CHECK(cond) \ + do \ + { \ + if(!(cond)) \ + { \ + fprintf(stderr, "FAIL: %s (%s:%d)\n", #cond, __FILE__, __LINE__); \ + ++g_failures; \ + } \ + } while(0) + +/* ConvProblem(N=2, Hi=12, Wi=12, C=32, K=96, Y=3, X=3, pH=1, pW=1): 10 tiles. */ +rocke_conv_problem_t dense() +{ + rocke_conv_problem_t p = rocke_conv_problem_default(2, 12, 12, 32, 96, 3, 3); + p.pH = 1; + p.pW = 1; + return p; +} + +/* ConvProblem(N=2, Hi=12, Wi=12, C=192, K=1056, Y=3, X=3, pH=1, pW=1): + * 17 x 27 = 459 tiles, more than either arch's default pool, so an auto-pool + * launch splits. */ +rocke_conv_problem_t splits() +{ + rocke_conv_problem_t p = rocke_conv_problem_default(2, 12, 12, 192, 1056, 3, 3); + p.pH = 1; + p.pW = 1; + return p; +} + +/* The library test's _spec(): 64x64x32 tiles, 2x2 warps, 32x32x16 atoms, + * dp_sk, linear, streamk_ctas=6. */ +rocke_implicit_gemm_conv_wgrad_spec_t spec(const rocke_conv_problem_t& p) +{ + rocke_implicit_gemm_conv_wgrad_spec_t s = rocke_implicit_gemm_conv_wgrad_spec_default(); + s.problem = p; + s.tile_m = 64; + s.tile_n = 64; + s.tile_k = 32; + s.warp_m = 2; + s.warp_n = 2; + s.warp_tile_m = 32; + s.warp_tile_n = 32; + s.warp_tile_k = 16; + s.streamk = "dp_sk"; + s.streamk_ctas = 6; + return s; +} + +void check_default_pool() +{ + static const struct + { + const char* arch; + int linear, atomic; + } k_pools[] = { + {"gfx950", 256, 1024}, + {"gfx942", 304, 1216}, + {"gfx1151", 64, 256}, + {"gfx90a", 256, 1024}, /* not in the CU table: 256-CU fallback */ + {"gfx9999", 256, 1024}, + {NULL, 256, 1024}, /* NULL arch is gfx950 */ + }; + for(const auto& e : k_pools) + { + CHECK(rocke_wgrad_conv_streamk_default_pool("linear", e.arch) == e.linear); + CHECK(rocke_wgrad_conv_streamk_default_pool("tree", e.arch) == e.linear); + CHECK(rocke_wgrad_conv_streamk_default_pool("atomic", e.arch) == e.atomic); + CHECK(rocke_wgrad_conv_streamk_default_pool("workspace", e.arch) == e.atomic); + } + CHECK(rocke_wgrad_conv_streamk_default_pool(NULL, "gfx942") == 304); /* linear */ +} + +void check_ctas() +{ + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(dense()); + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx942") == 6); + s.streamk_ctas = -1; + CHECK(rocke_wgrad_conv_streamk_ctas(&s, NULL) == 256); + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx950") == 256); + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx942") == 304); + s.streamk_reduction = "workspace"; + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx942") == 1216); + /* Python wgrad_streamk_launch raises for these; C reports no pool. */ + s.streamk_ctas = 0; + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx950") == 0); + s.streamk_ctas = -3; + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx950") == 0); + s.streamk_ctas = 6; + s.streamk = "off"; + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx950") == 0); + CHECK(rocke_wgrad_conv_streamk_ctas(NULL, "gfx950") == 0); +} + +struct PlanVector +{ + const char* arch; + const char* mode; + int ctas; + int sk_tiles, sk_ctas, dp_tiles, grid, ipsc, extra; + size_t flags_bytes, partials_bytes; +}; + +void check_plans() +{ + /* linear reduction on the 459-tile problem. */ + static const PlanVector k_plans[] = { + {"gfx950", "dp_sk", -1, 459, 256, 0, 256, 16, 35, 4096, 4194304}, + {"gfx950", "persistent", -1, 459, 256, 0, 256, 16, 35, 4096, 4194304}, + {"gfx942", "dp_sk", -1, 459, 304, 0, 304, 13, 179, 4864, 4980736}, + {"gfx942", "persistent", -1, 459, 304, 0, 304, 13, 179, 4864, 4980736}, + {"gfx950", "dp_sk", 7, 11, 7, 448, 455, 14, 1, 256, 114688}, + {"gfx942", "persistent", 7, 11, 7, 448, 7, 14, 1, 256, 114688}, + {NULL, "dp_sk", -1, 459, 256, 0, 256, 16, 35, 4096, 4194304}, + }; + for(const PlanVector& v : k_plans) + { + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(splits()); + s.streamk = v.mode; + s.streamk_ctas = v.ctas; + rocke_status_t st = ROCKE_OK; + const rocke_streamk_iter_partition_t part + = rocke_wgrad_conv_streamk_partition(&s, v.arch, &st); + CHECK(st == ROCKE_OK); + CHECK(part.m_tiles == 17); + CHECK(part.n_tiles == 27); + CHECK(part.iters_per_tile == 9); + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&part); + CHECK(r.sk_tiles == v.sk_tiles); + CHECK(r.sk_ctas == v.sk_ctas); + CHECK(r.dp_tiles == v.dp_tiles); + CHECK(r.grid_size == v.grid); + CHECK(r.iters_per_sk_cta == v.ipsc); + CHECK(r.extra_iters == v.extra); + + int grid[3] = {0, 0, 0}; + CHECK(rocke_wgrad_conv_streamk_grid(&s, v.arch, grid) == ROCKE_OK); + CHECK(grid[0] == v.grid && grid[1] == 1 && grid[2] == 1); + + size_t flags = 1, partials = 1; + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, v.arch, &flags, &partials) == ROCKE_OK); + CHECK(flags == v.flags_bytes); + CHECK(partials == v.partials_bytes); + CHECK(rocke_wgrad_conv_streamk_workspace_bytes(&s, v.arch) + == v.flags_bytes + v.partials_bytes); + } + + /* The other reductions on the same shape, auto pool. */ + static const struct + { + const char* arch; + const char* reduction; + const char* dtype_d; + int grid; + size_t workspace_bytes; + } k_reductions[] = { + {"gfx950", "tree", "fp16", 256, 4198400}, + {"gfx942", "tree", "fp16", 304, 4985600}, + {"gfx950", "atomic", "fp32", 1024, 0}, + {"gfx942", "atomic", "fp32", 1216, 0}, + {"gfx950", "workspace", "fp16", 1024, 58392576}, + {"gfx942", "workspace", "fp16", 1216, 58392576}, + }; + for(const auto& e : k_reductions) + { + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(splits()); + s.streamk_ctas = -1; + s.streamk_reduction = e.reduction; + s.dtype_d = e.dtype_d; + int grid[3] = {0, 0, 0}; + CHECK(rocke_wgrad_conv_streamk_grid(&s, e.arch, grid) == ROCKE_OK); + CHECK(grid[0] == e.grid); + size_t flags = 1, partials = 1; + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, e.arch, &flags, &partials) == ROCKE_OK); + const bool fixup = strcmp(e.reduction, "tree") == 0; + CHECK(fixup ? flags + partials == e.workspace_bytes : flags == 0 && partials == 0); + CHECK(rocke_wgrad_conv_streamk_workspace_bytes(&s, e.arch) == e.workspace_bytes); + CHECK(rocke_wgrad_conv_spec_is_deterministic(&s) == fixup); + } + + /* 10 tiles over a pool of 5: all data-parallel, no flags or slots. */ + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(dense()); + s.streamk_ctas = 5; + size_t flags = 1, partials = 1; + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, "gfx950", &flags, &partials) == ROCKE_OK); + CHECK(flags == 0 && partials == 0); + /* 10 tiles over 7: 256 flag bytes, 7 partial tiles. */ + s.streamk_ctas = 7; + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, "gfx950", &flags, &partials) == ROCKE_OK); + CHECK(flags == 256 && partials == 7 * 64 * 64 * 4); +} + +void check_invalid_inputs() +{ + rocke_status_t st = ROCKE_OK; + int grid[3] = {0, 0, 0}; + size_t flags = 0, partials = 0; + + rocke_streamk_iter_partition_t part = rocke_wgrad_conv_streamk_partition(NULL, NULL, &st); + CHECK(st == ROCKE_ERR_VALUE); + CHECK(part.m_tiles == 0 && part.max_active_wgs == 0); + CHECK(rocke_wgrad_conv_streamk_grid(NULL, NULL, grid) == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_layout(NULL, NULL, &flags, &partials) + == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_bytes(NULL, NULL) == 0); + + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(dense()); + CHECK(rocke_wgrad_conv_streamk_grid(&s, NULL, NULL) == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, NULL, NULL, &partials) == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, NULL, &flags, NULL) == ROCKE_ERR_VALUE); + + const int bad_ctas[] = {0, -2, -3}; + for(int c : bad_ctas) + { + s.streamk_ctas = c; + st = ROCKE_OK; + (void)rocke_wgrad_conv_streamk_partition(&s, "gfx950", &st); + CHECK(st == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_grid(&s, "gfx950", grid) == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, "gfx950", &flags, &partials) + == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_bytes(&s, "gfx950") == 0); + } + s.streamk_ctas = 6; + s.streamk = "off"; + st = ROCKE_OK; + (void)rocke_wgrad_conv_streamk_partition(&s, "gfx950", &st); + CHECK(st == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_bytes(&s, "gfx950") == 0); +} + +/* is_valid_wgrad_spec accept/reject reasons, verbatim from Python. */ +void check_validator() +{ + char why[1024]; + auto valid = [&](const rocke_implicit_gemm_conv_wgrad_spec_t& s, const char* arch) { + why[0] = '\0'; + return rocke_implicit_gemm_conv_wgrad_is_valid_spec(&s, arch, why, sizeof(why)); + }; + auto rejects = + [&](const rocke_implicit_gemm_conv_wgrad_spec_t& s, const char* arch, const char* reason) { + const bool ok = valid(s, arch); + if(ok || strcmp(why, reason) != 0) + { + fprintf(stderr, " want: %s\n got: %s\n", reason, ok ? "(accepted)" : why); + return false; + } + return true; + }; + + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(dense()); + CHECK(valid(s, "gfx950")); + CHECK(valid(s, NULL)); + s.streamk_ctas = -1; + CHECK(valid(s, "gfx950")); + s.streamk = "persistent"; + s.streamk_reduction = "tree"; + CHECK(valid(s, "gfx950")); + s = spec(dense()); + s.warp_tile_m = 16; + s.warp_tile_n = 16; + CHECK(valid(s, "gfx942")); + s = spec(dense()); + s.dtype_d = "fp32"; + s.streamk_reduction = "atomic"; + CHECK(valid(s, "gfx950")); + + s = spec(dense()); + s.streamk_ctas = 0; + CHECK(rejects(s, "gfx950", "streamk_ctas must be -1 (auto) or > 0 (got 0)")); + s.streamk_ctas = -3; + CHECK(rejects(s, "gfx950", "streamk_ctas must be -1 (auto) or > 0 (got -3)")); + + s = spec(dense()); + s.streamk = "sideways"; + CHECK(rejects(s, "gfx950", "streamk must be 'off', 'dp_sk' or 'persistent' (got 'sideways')")); + s = spec(dense()); + s.streamk_reduction = "ring"; + CHECK(rejects(s, + "gfx950", + "streamk_reduction must be 'linear', 'tree', 'atomic' or 'workspace' " + "(got 'ring')")); + s = spec(dense()); + s.split_k = 4; + CHECK(rejects(s, + "gfx950", + "streamk and split_k both partition K_wg; use split_k=1 with streamk " + "(got split_k=4)")); + s = spec(dense()); + s.async_dma = true; + CHECK(rejects(s, + "gfx950", + "streamk runs one plain K loop per work item; async_dma/unroll_k " + "pipeline the K loop across items")); + s = spec(dense()); + s.streamk_reduction = "atomic"; + CHECK(rejects(s, + "gfx950", + "streamk_reduction='atomic' needs dtype_d='fp32' (got 'fp16'); use " + "'linear', 'tree' or 'workspace'")); + s = spec(dense()); + CHECK(rejects(s, "gfx90a", "stream-K wgrad supports only gfx942, gfx950 (got gfx90a)")); + + /* Multi-cause rejects: the first gate in Python's order wins. */ + s = spec(dense()); + s.streamk_ctas = 0; + s.tile_m = 96; + CHECK(rejects(s, "gfx950", "tile_m not divisible by warp_m * warp_tile_m")); + s = spec(dense()); + s.warp_m = 8; + s.warp_n = 8; + s.tile_m = 256; + s.tile_n = 256; + s.streamk_ctas = 0; + CHECK(rejects(s, "gfx950", "block_size 4096 > 1024 (hardware cap) on gfx950")); + s = spec(dense()); + s.wave_size = 32; + s.streamk_ctas = 0; + CHECK(rejects(s, "gfx950", "spec wave_size 32 != gfx950 wave_size 64")); + s = spec(dense()); + CHECK(rejects(s, "gfx1250", "spec wave_size 64 != gfx1250 wave_size 32")); + s = spec(dense()); + s.has_vector_size_c = true; + s.vector_size_c = 4; + s.streamk_ctas = -5; + CHECK(rejects(s, "gfx950", "default epilogue is not supported with vector size c: 4")); + s = spec(dense()); + s.async_dma = true; + s.streamk = "off"; + CHECK(rejects(s, + "gfx950", + "wgrad async_dma requires lds_k_outer=True: the direct global->LDS load " + "needs a stride-1 reduction axis, which wgrad only has once the tile is " + "stored K-outer")); + s = spec(dense()); + s.tile_m = 96; + CHECK(rejects(s, + "gfx9999", + "\"unknown gfx target 'gfx9999'; known: ['gfx11-generic', 'gfx1151', " + "'gfx1201', 'gfx1250', 'gfx90a', 'gfx942', 'gfx950']. Add a row to " + "arch_specs.json.\"")); + + /* Depthwise ConvProblem(N=2, Hi=8, Wi=8, C=16, K=16, Y=3, X=3, pH=1, pW=1, + * groups=16): the shape group merging admits. */ + rocke_conv_problem_t dw = rocke_conv_problem_default(2, 8, 8, 16, 16, 3, 3); + dw.pH = 1; + dw.pW = 1; + dw.groups = 16; + s = spec(dw); + s.group_merge = 3; + CHECK(rejects(s, "gfx950", "group_merge must be 1 or one of (2, 4, 8, 16, 32, 64); got 3")); + s.group_merge = 2; + CHECK(rejects(s, "gfx950", "streamk folds conv groups into GEMM-M; use group_merge=1 (got 2)")); + s.tile_m = 96; + CHECK(rejects(s, "gfx950", "tile_m not divisible by warp_m * warp_tile_m")); +} + +} // namespace + +int main() +{ + check_default_pool(); + check_ctas(); + check_plans(); + check_invalid_inputs(); + check_validator(); + if(g_failures) + { + fprintf(stderr, "rocke_wgrad_streamk_host: %d failure(s)\n", g_failures); + return 1; + } + printf("rocke_wgrad_streamk_host: OK\n"); + return 0; +} 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 0eaa49e034f1..3b8d8dd7380b 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 @@ -21,6 +21,7 @@ from __future__ import annotations +import re import unittest from unittest.mock import patch, Mock @@ -3270,6 +3271,203 @@ def _lower(self, name, build, **kw): _assert_ir_assembles(self, ll, name) return ll + # ---- cross-workgroup fence / flag ops (stream-K fixup) ---- + @staticmethod + def _flag_kernel(b): + from rocke.core.ir import I32, PtrType + + return b.param("flags", PtrType(I32, "global")) + + def test_fence_renders_scope_and_ordering(self): + ll = self._lower("fence", lambda b: b.fence(scope="agent", ordering="release")) + self.assertIn(' fence syncscope("agent") release', ll) + + def test_fence_system_scope_has_no_syncscope_clause(self): + ll = self._lower( + "fencesys", lambda b: b.fence(scope="system", ordering="acquire") + ) + self.assertIn(" fence acquire", ll) + self.assertNotIn("syncscope", ll) + + def test_global_flag_store_is_scoped_atomic_store(self): + def build(b): + flags = self._flag_kernel(b) + b.global_flag_store(flags, b.const_i32(3), b.const_i32(1)) + + ll = self._lower("flagst", build) + self.assertRegex( + ll, + r"store atomic i32 1, ptr addrspace\(1\) %gep\.\d+ " + r'syncscope\("agent"\) release, align 4', + ) + + def test_global_flag_wait_eq_is_acquire_spin_loop(self): + def build(b): + flags = self._flag_kernel(b) + b.global_flag_wait_eq(flags, b.const_i32(2), b.const_i32(1)) + + ll = self._lower("flagwt", build) + self.assertRegex( + ll, + r"(?s)br label %(flag\.wait\.\d+)\n\1:\n" + r" (%flag\.\d+) = load atomic i32, ptr addrspace\(1\) %gep\.\d+ " + r'syncscope\("agent"\) acquire, align 4\n' + r" (%flag\.hit\.\d+) = icmp eq i32 \2, 1\n" + r" br i1 \3, label %flag\.done\.\d+, label %\1\n", + ) + + def test_flag_ops_lower_to_hip_atomics(self): + from rocke.core.lower_hip import lower_kernel_to_hip + + b = self._builder("flaghip") + flags = self._flag_kernel(b) + b.global_flag_store(flags, b.const_i32(0), b.const_i32(1)) + b.fence(scope="agent", ordering="acquire") + b.global_flag_wait_eq(flags, b.const_i32(0), b.const_i32(1)) + hip = lower_kernel_to_hip(b.kernel) + self.assertIn("__ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_AGENT);", hip) + self.assertIn('__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");', hip) + self.assertIn("while (__hip_atomic_load(", hip) + + def test_flag_ops_reject_bad_scope_ordering_and_type(self): + b = self._builder("flagbad") + flags = self._flag_kernel(b) + with self.assertRaisesRegex(ValueError, "unknown fence scope"): + b.fence(scope="device") + with self.assertRaisesRegex(ValueError, "unknown fence ordering"): + b.fence(ordering="monotonic") + with self.assertRaisesRegex(ValueError, "expects an i32 value"): + b.global_flag_store(flags, b.const_i32(0), b.const_f32(1.0)) + with self.assertRaisesRegex(ValueError, "unknown flag store ordering"): + b.global_flag_store( + flags, b.const_i32(0), b.const_i32(1), ordering="acquire" + ) + with self.assertRaisesRegex(ValueError, "unknown flag scope"): + b.global_flag_wait_eq(flags, b.const_i32(0), b.const_i32(1), scope="gpu") + + def test_flag_ops_reject_non_i32_pointer(self): + # LLVM GEPs the flag as i32 while HIP indexes the pointee type, so any + # other element type would make the two engines address different + # bytes. Builder and verifier both refuse it (C++ mirror: + # case_flag_ops_reject_non_i32_ptr). + from rocke.core.ir import F16, PtrType + from rocke.core.verify import ERROR, verify + + b = self._builder("flagf16") + half = b.param("flags", PtrType(F16, "global")) + for name, emit in ( + ("global_flag_store", b.global_flag_store), + ("global_flag_wait_eq", b.global_flag_wait_eq), + ): + with self.assertRaisesRegex( + ValueError, + f"^{name} expects a global i32 pointer, got ptr$", + ): + emit(half, b.const_i32(0), b.const_i32(1)) + + b = self._builder("flagf16verify") + flags = self._flag_kernel(b) + half = b.param("half", PtrType(F16, "global")) + b.global_flag_store(flags, b.const_i32(0), b.const_i32(1)) + b.global_flag_wait_eq(flags, b.const_i32(0), b.const_i32(1)) + b.ret() + for op in b.kernel.body.ops: + if op.name.startswith("memref.global_flag_"): + op.operands[0] = half + msgs = [d.message for d in verify(b.kernel) if d.severity == ERROR] + for name in ("memref.global_flag_store", "memref.global_flag_wait_eq"): + needle = f"{name}: expected a global i32 pointer, got ptr" + self.assertTrue(any(needle in m for m in msgs), (needle, msgs)) + + def test_flag_ops_lowerers_reject_bad_attrs(self): + # Python counterpart of the C++ expect_flag_attr_rejected: an op with a + # missing or unknown scope / ordering (serialized IR that skipped the + # builder and the verifier) is rejected by BOTH lowerers instead of + # defaulting or rendering invalid IR such as ``fence ... bogus``. + from rocke.core.backend import BackendError, _lower_via_cpp_engine + from rocke.core.ir_serialize import serialize + from rocke.core.lower_hip import lower_kernel_to_hip + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + valid = { + "memref.fence": {"scope": "agent", "ordering": "acq_rel"}, + "memref.global_flag_store": {"scope": "agent", "ordering": "release"}, + "memref.global_flag_wait_eq": {"scope": "agent"}, + } + cases = [ + ("memref.fence", "ordering", "bogus"), + ("memref.fence", "scope", "bogus"), + ("memref.fence", "ordering", "monotonic"), + ("memref.global_flag_store", "ordering", "bogus"), + ("memref.global_flag_store", "ordering", "acquire"), + ("memref.global_flag_store", "scope", "bogus"), + ("memref.global_flag_wait_eq", "scope", "bogus"), + ] + [(name, key, None) for name, attrs in valid.items() for key in attrs] + for name, key, value in cases: + with self.subTest(op=name, key=key, value=value): + b = self._builder("flag_bad_attr") + if name == "memref.fence": + b.fence() + else: + flags = self._flag_kernel(b) + emit = getattr(b, name.split(".", 1)[1]) + emit(flags, b.const_i32(0), b.const_i32(1)) + b.ret() + (op,) = [o for o in b.kernel.body.ops if o.name == name] + op.attrs.clear() + op.attrs.update(valid[name]) + if value is None: + del op.attrs[key] + want = f"{name}: {key} attribute missing or not a string" + else: + op.attrs[key] = value + want = f"{name}: unknown {key} '{value}'" + with self.assertRaisesRegex(ValueError, re.escape(want)): + lower_kernel_to_hip(b.kernel) + with self.assertRaisesRegex(ValueError, re.escape(want)): + _lower_kernel_to_llvm_python(b.kernel) + try: + _lower_via_cpp_engine(serialize(b.kernel), "gfx950", "llvm20") + except BackendError: + continue # engine not built: C++ is covered by its ctest + except Exception as e: # noqa: BLE001 -- engine error type + self.assertIn(want, str(e)) + else: + self.fail(f"C++ engine accepted {name} with {key}={value!r}") + + def test_verifier_rejects_malformed_flag_ops(self): + # Serialized IR can carry flag ops the builder would have refused; the + # verifier has to catch them before a lowerer unpacks the operands. + from rocke.core.ir import F32 + from rocke.core.verify import ERROR, verify + + b = self._builder("flagverify") + flags = self._flag_kernel(b) + b.fence(scope="agent", ordering="release") + b.global_flag_store(flags, b.const_i32(0), b.const_i32(1)) + b.global_flag_wait_eq(flags, b.const_i32(0), b.const_i32(1)) + b.ret() + ops = {op.name: op for op in b.kernel.body.ops} + self.assertEqual( + [d for d in verify(b.kernel) if d.severity == ERROR], [], "well-formed" + ) + ops["memref.fence"].attrs["ordering"] = "monotonic" + ops["memref.global_flag_store"].attrs["scope"] = "gpu" + wait = ops["memref.global_flag_wait_eq"] + wait.operands[2] = b.const_f32(1.0) + msgs = [d.message for d in verify(b.kernel) if d.severity == ERROR] + for needle in ( + "memref.fence: unknown ordering 'monotonic'", + "memref.global_flag_store: unknown scope 'gpu'", + f"memref.global_flag_wait_eq: value must be i32, got {F32.name}", + ): + self.assertTrue(any(needle in m for m in msgs), (needle, msgs)) + del wait.operands[2] + msgs = [d.message for d in verify(b.kernel) if d.severity == ERROR] + self.assertTrue( + any("expected 3 operands / 0 results, got 2 / 0" in m for m in msgs), msgs + ) + # ---- ds_swizzle (raw offset + XOR-butterfly encoding) ---- def test_ds_swizzle_passes_raw_offset_immediate(self): ll = self._lower("dssw", lambda b: b.ds_swizzle(b.const_i32(1), 0x041F))