From ffdc64cef0cd60e6a2633cadec85fb57abeb9729 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Tue, 22 Sep 2026 12:07:13 +0000 Subject: [PATCH 1/5] rocke(feat): direct conv bf16 support --- .../common/benchmark_direct_conv.py | 92 +-- .../kernels/common/conv_direct_grouped.py | 533 +++++++++++++----- .../rocke/instance_conv_direct_grouped.h | 42 +- .../instance_conv_direct_grouped_internal.h | 3 + ...direct_grouped_build_16c_loop_closures.cpp | 70 ++- .../conv_direct_grouped_build_16c_phases.cpp | 53 +- .../common/conv_direct_grouped_build_32c.cpp | 65 ++- .../common/conv_direct_grouped_build_8c.cpp | 72 ++- .../conv_direct_grouped_spec_and_validate.cpp | 224 +++++++- 9 files changed, 871 insertions(+), 283 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py index acd7673d0d6a..71deebceb802 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py @@ -80,10 +80,10 @@ class DepthwiseResult: # --------------------------------------------------------------------------- _MIOPEN_DTYPE_MAP = { - "conv": "fp32", + "conv": "fp32", # rejected — fp32 not supported "convfp16": "fp16", "convbfp16": "bf16", - "convint8": "fp16", # int8 not supported; fall back to fp16 and warn + "convint8": "int8", # rejected — int8 not supported } @@ -117,10 +117,14 @@ def parse_miopen_cmd_direct(cmd: str): f"(expected one of: {list(_MIOPEN_DTYPE_MAP)})" ) dtype = _MIOPEN_DTYPE_MAP[driver_kw] + if dtype == "fp32": + raise ValueError( + f"fp32 ({driver_kw!r}) is not supported by this benchmark; " + f"use convfp16 or convbfp16" + ) if driver_kw == "convint8": - print( - "[warn] convint8 is not supported by this benchmark; treating as fp16", - file=sys.stderr, + raise ValueError( + "convint8 is not supported by this benchmark; " "use convfp16 or convbfp16" ) sub = argparse.ArgumentParser(add_help=False) @@ -200,6 +204,7 @@ def parse_miopen_cmd_direct(cmd: str): KW=miopen_args.X, PAD=miopen_args.pH, stride=sH, + dtype=dtype if dtype in ("fp16", "bf16") else "fp16", ) return problem, dtype, miopen_args.forw @@ -349,12 +354,17 @@ def __init__(self, p): def _print_results( - results: List[Result], top_n_arg: int, arch: str, p, show_verify: bool + results: List[Result], + top_n_arg: int, + arch: str, + p, + show_verify: bool, + dtype: str = "fp16", ): top_n = min(top_n_arg, len(results)) width = 100 if show_verify else 88 print(f"\n{'='*width}") - print(f"Top {top_n} configurations for {arch} fp16 {p.short()}") + print(f"Top {top_n} configurations for {arch} {dtype} {p.short()}") print(f"{'='*width}") hdr = ( f"{'rank':>4} {'TFLOPS':>7} {'ms':>8} {'GBps':>7} {'verify':>6} config" @@ -380,12 +390,17 @@ def _print_results( def _print_depthwise_results( - results: "List[DepthwiseResult]", top_n_arg: int, arch: str, p, show_verify: bool + results: "List[DepthwiseResult]", + top_n_arg: int, + arch: str, + p, + show_verify: bool, + dtype: str = "fp16", ): top_n = min(top_n_arg, len(results)) width = 96 if show_verify else 84 print(f"\n{'='*width}") - print(f"Top {top_n} depthwise configurations for {arch} fp16 {p.short()}") + print(f"Top {top_n} depthwise configurations for {arch} {dtype} {p.short()}") print(f"{'='*width}") hdr = ( f"{'rank':>4} {'TFLOPS':>7} {'ms':>8} {'GBps':>7} {'verify':>6} config" @@ -419,6 +434,7 @@ def _run_depthwise_sweep( *, args, problem, + dtype: str = "fp16", arch: str, compile_kernel, jobs: int, @@ -451,13 +467,14 @@ def _run_depthwise_sweep( _use_spatial = p.groups < _wave_size torch.manual_seed(42) - A_t = torch.empty(p.N, p.H, p.W, p.total_c, dtype=torch.float16).uniform_(-1.0, 1.0) - B_t = torch.empty(p.total_k, p.KH, p.KW, 1, dtype=torch.float16).uniform_(-1.0, 1.0) - D_t = torch.empty(p.N, p.Ho, p.Wo, p.total_k, dtype=torch.float16) + _torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16 + A_t = torch.empty(p.N, p.H, p.W, p.total_c, dtype=_torch_dtype).uniform_(-1.0, 1.0) + B_t = torch.empty(p.total_k, p.KH, p.KW, 1, dtype=_torch_dtype).uniform_(-1.0, 1.0) + D_t = torch.empty(p.N, p.Ho, p.Wo, p.total_k, dtype=_torch_dtype) bytes_xfer = float(A_t.nbytes + B_t.nbytes + D_t.nbytes) flop = float(p.flops) - sig = conv_args_signature("fp16") + sig = conv_args_signature(dtype) # For the spatial layout block_w is derived from block_waves internally, # so sweeping block_w would produce duplicate kernels; use a dummy value. @@ -476,7 +493,7 @@ def _run_depthwise_sweep( ) print( - f"Sweeping {len(combos)} depthwise combinations for {arch} fp16 {p.short()} ...", + f"Sweeping {len(combos)} depthwise combinations for {arch} {dtype} {p.short()} ...", flush=True, ) @@ -637,7 +654,7 @@ def _run_depthwise_sweep( return 1, [] results.sort(key=lambda r: r.tflops, reverse=True) - _print_depthwise_results(results, args.top, arch, p, args.verify) + _print_depthwise_results(results, args.top, arch, p, args.verify, dtype=dtype) return 0, results @@ -645,6 +662,7 @@ def _run_sweep( *, args, problem, + dtype: str = "fp16", arch: str, compile_kernel, jobs: int, @@ -668,15 +686,16 @@ def _run_sweep( p = problem torch.manual_seed(42) - A_t = torch.empty(p.N, p.H, p.W, p.total_c, dtype=torch.float16).uniform_(-1.0, 1.0) - B_t = torch.empty(p.total_k, p.KH, p.KW, p.cpg, dtype=torch.float16).uniform_( + _torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16 + A_t = torch.empty(p.N, p.H, p.W, p.total_c, dtype=_torch_dtype).uniform_(-1.0, 1.0) + B_t = torch.empty(p.total_k, p.KH, p.KW, p.cpg, dtype=_torch_dtype).uniform_( -1.0, 1.0 ) - D_t = torch.empty(p.N, p.Ho, p.Wo, p.total_k, dtype=torch.float16) + D_t = torch.empty(p.N, p.Ho, p.Wo, p.total_k, dtype=_torch_dtype) bytes_xfer = float(A_t.nbytes + B_t.nbytes + D_t.nbytes) flop = float(p.flops) - sig = conv_args_signature("fp16") + sig = conv_args_signature(dtype) combos = list(itertools.product(_BLOCK_Q, _BLOCK_GROUPS, _DOUBLE_BUFFER)) @@ -690,7 +709,7 @@ def _run_sweep( ) print( - f"Sweeping {len(combos)} combinations for {arch} fp16 {p.short()} " + f"Sweeping {len(combos)} combinations for {arch} {dtype} {p.short()} " f"(cpg={p.cpg}) ...", flush=True, ) @@ -840,7 +859,7 @@ def _run_sweep( return 1, [] results.sort(key=lambda r: r.tflops, reverse=True) - _print_results(results, args.top, arch, p, args.verify) + _print_results(results, args.top, arch, p, args.verify, dtype=dtype) return 0, results @@ -853,6 +872,7 @@ def _run_dgrad_sweep( *, args, problem, + dtype: str = "fp16", arch: str, compile_kernel, jobs: int, @@ -886,17 +906,18 @@ def _run_dgrad_sweep( p = problem torch.manual_seed(42) - dY_t = torch.empty(p.N, p.Ho, p.Wo, p.total_k, dtype=torch.float16).uniform_( + _torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16 + dY_t = torch.empty(p.N, p.Ho, p.Wo, p.total_k, dtype=_torch_dtype).uniform_( -1.0, 1.0 ) - W_t = torch.empty(p.total_k, p.KH, p.KW, p.cpg, dtype=torch.float16).uniform_( + W_t = torch.empty(p.total_k, p.KH, p.KW, p.cpg, dtype=_torch_dtype).uniform_( -1.0, 1.0 ) - dX_t = torch.empty(p.N, p.H, p.W, p.total_c, dtype=torch.float16) + dX_t = torch.empty(p.N, p.H, p.W, p.total_c, dtype=_torch_dtype) bytes_xfer = float(dY_t.nbytes + W_t.nbytes + dX_t.nbytes) flop = float(p.flops) - sig = conv_args_signature("fp16") + sig = conv_args_signature(dtype) is_depthwise = p.cpg == 1 @@ -910,7 +931,7 @@ def _run_dgrad_sweep( combos_dw = list(itertools.product(_DW_BLOCK_W, _DW_BLOCK_WAVES)) print( - f"Sweeping {len(combos_dw)} depthwise dgrad combinations for {arch} fp16 " + f"Sweeping {len(combos_dw)} depthwise dgrad combinations for {arch} {dtype} " f"{p.short()} (stride={p.stride}) ...", flush=True, ) @@ -956,7 +977,7 @@ def _run_dgrad_sweep( valid_bgs = [bg for bg in _BLOCK_GROUPS if p.groups % bg == 0] combos = list(itertools.product(_BLOCK_Q, valid_bgs)) print( - f"Sweeping {len(combos)} MFMA dgrad combinations for {arch} fp16 {p.short()} " + f"Sweeping {len(combos)} MFMA dgrad combinations for {arch} {dtype} {p.short()} " f"(cpg={p.cpg}, kpg={p.kpg}) ...", flush=True, ) @@ -1046,7 +1067,7 @@ def _run_dgrad_sweep( _DGRAD_BLOCK_Q = (4, 8, 16, 32) combos = list(itertools.product(_DGRAD_BLOCK_Q, valid_bgs)) print( - f"Sweeping {len(combos)} scalar-FMA dgrad combinations for {arch} fp16 {p.short()} " + f"Sweeping {len(combos)} scalar-FMA dgrad combinations for {arch} {dtype} {p.short()} " f"(cpg={p.cpg}, kpg={p.kpg}, stride={p.stride}) ...", flush=True, ) @@ -1552,6 +1573,12 @@ def main() -> int: default=1, help="number of conv groups; C and K must each be divisible by groups (default: 1)", ) + conv.add_argument( + "--dtype", + choices=("fp16", "bf16"), + default="fp16", + help="I/O data type when using shape flags (default: fp16; ignored when using --miopen-cmd/--miopen-file)", + ) args = parser.parse_args() @@ -1663,8 +1690,9 @@ def _u8(t): KW=args.X, PAD=args.pH, stride=args.sH, + dtype=args.dtype, ) - cases = [(problem, "fp16", args.direction)] + cases = [(problem, args.dtype, args.direction)] _common = dict( args=args, @@ -1691,11 +1719,11 @@ def _u8(t): cpg = problem.cpg if direction == "dgrad": - rc, _ = _run_dgrad_sweep(problem=problem, **_common) + rc, _ = _run_dgrad_sweep(problem=problem, dtype=dtype, **_common) elif cpg == 1: - rc, _ = _run_depthwise_sweep(problem=problem, **_common) + rc, _ = _run_depthwise_sweep(problem=problem, dtype=dtype, **_common) else: - rc, _ = _run_sweep(problem=problem, **_common) + rc, _ = _run_sweep(problem=problem, dtype=dtype, **_common) all_rc = all_rc or rc return all_rc diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py index 3eaaee4dfaae..13f14f66c0cb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py @@ -72,6 +72,7 @@ from typing import List, Tuple from rocke.core.ir import ( + BF16, F16, F32, I32, @@ -80,6 +81,54 @@ PtrType, Value, ) + + +def _io_type(dtype: str): + """Return the IR type for a given ``dtype`` string (``"fp16"`` or ``"bf16"``).""" + if dtype == "bf16": + return BF16 + if dtype == "fp16": + return F16 + raise ValueError( + f"unsupported direct_conv dtype: {dtype!r}; expected 'fp16' or 'bf16'" + ) + + +def _buf_load_vN( + b: IRBuilder, dtype: str, rsrc: Value, voff: Value, soff: Value, dwords: int +) -> Value: + """Dtype-dispatch for vectorised buffer load. + + ``dwords`` matches the ``dwords`` parameter of ``buffer_load_vN_f16`` / + ``buffer_load_vN_bf16``: each dword holds two 16-bit elements (dwords=1 → + 2 elements, dwords=2 → 4 elements, dwords=4 → 8 elements). + Uses the type-specific op names to stay byte-identical with the C++ engine. + """ + if dtype == "bf16": + return b.buffer_load_vN_bf16(rsrc, voff, soff, dwords) + return b.buffer_load_vN_f16(rsrc, voff, soff, dwords) + + +def _buf_store_vN( + b: IRBuilder, + dtype: str, + rsrc: Value, + voff: Value, + soff: Value, + val: Value, + dwords: int, +) -> None: + """Dtype-dispatch for vectorised buffer store. + + ``dwords`` matches the ``dwords`` parameter of ``buffer_store_vN_f16`` / + ``buffer_store_vN_bf16``: each dword holds two 16-bit elements. + """ + if dtype == "bf16": + b.buffer_store_vN_bf16(rsrc, voff, soff, val, dwords) + else: + b.buffer_store_vN_f16(rsrc, voff, soff, val, dwords) + + from rocke.helpers.transforms import TensorDescriptor, embed, unmerge_magic @@ -88,9 +137,9 @@ class DirectConvProblem: """The grouped direct-conv shape parameters. Layouts: - A: NHWC fp16, `[N, H, W, groups*cpg]` - B: KRSC fp16, `[groups*kpg, KH, KW, cpg]` - D: NHWK fp16, `[N, H, W, groups*kpg]` + A: NHWC, `[N, H, W, groups*cpg]` + B: KRSC, `[groups*kpg, KH, KW, cpg]` + D: NHWK, `[N, H, W, groups*kpg]` """ N: int @@ -103,6 +152,7 @@ class DirectConvProblem: KW: int = 3 PAD: int = 1 stride: int = 1 + dtype: str = "fp16" # "fp16" or "bf16" @property def total_c(self) -> int: @@ -190,11 +240,13 @@ def kernel_name(self) -> str: f"bq{self.block_q}", f"bg{self.block_groups}", "db" if self.double_buffer else "sb", - flags={"k32": self.fold_k32}, + flags={"k32": self.fold_k32, "bf16": p.dtype == "bf16"}, ) def validate(self) -> None: p = self.problem + if p.dtype not in ("fp16", "bf16"): + raise ValueError(f"DirectConv16cSpec: unsupported dtype {p.dtype!r}") if p.cpg != 16 or p.kpg != 16: raise ValueError( f"DirectConv16cSpec expects cpg=kpg=16 (got {p.cpg}, {p.kpg})" @@ -230,6 +282,8 @@ def is_valid_spec_16c( except KeyError as e: return False, str(e) p = spec.problem + if p.dtype not in ("fp16", "bf16"): + return False, f"unsupported dtype {p.dtype!r}; expected 'fp16' or 'bf16'" if p.stride != 1: return False, f"stride > 1 is not supported (got {p.stride})" if p.cpg != 16 or p.kpg != 16: @@ -238,15 +292,16 @@ def is_valid_spec_16c( return False, ( f"groups {p.groups} not divisible by block_groups {spec.block_groups}" ) + ab_dtype = "bf16" if p.dtype == "bf16" else "f16" if not target.mma.has_shape( - a_dtype="f16", b_dtype="f16", c_dtype="fp32", m=16, n=16, k=16 + a_dtype=ab_dtype, b_dtype=ab_dtype, c_dtype="fp32", m=16, n=16, k=16 ): - return False, f"missing 16x16x16 f16 MFMA atom on {arch}" + return False, f"missing 16x16x16 {ab_dtype} MFMA atom on {arch}" if spec.fold_k32 and not target.mma.has_shape( - a_dtype="f16", b_dtype="f16", c_dtype="fp32", m=16, n=16, k=32 + a_dtype=ab_dtype, b_dtype=ab_dtype, c_dtype="fp32", m=16, n=16, k=32 ): return False, ( - f"fold_k32=True needs the 16x16x32 f16 MFMA atom, absent on " + f"fold_k32=True needs the 16x16x32 {ab_dtype} MFMA atom, absent on " f"{arch}; use fold_k32=False for a {arch}-capable kernel" ) return True, "ok" @@ -273,6 +328,7 @@ def build_direct_conv_16c( if not ok: raise ValueError(f"invalid direct_conv_16c spec for {arch}: {why}") p = spec.problem + io_type = _io_type(p.dtype) BLOCK_Q = spec.block_q BLOCK_GROUPS = spec.block_groups WAVE = spec.wave_size @@ -292,9 +348,9 @@ def build_direct_conv_16c( b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = THREADS - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - Bp = b.param("B", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + Bp = b.param("B", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) B_bytes = b.param("B_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -355,9 +411,9 @@ def build_direct_conv_16c( # halves so the OOB-zeroed writes land in the slack region of the # allocation and never alias a valid chunk. lds_total_fp16 = PASSES * THREADS * LOAD_VEC - A_smem = b.smem_alloc(F16, [1, lds_total_fp16], name_hint="lds_a") + A_smem = b.smem_alloc(io_type, [1, lds_total_fp16], name_hint="lds_a") B_smem = ( - b.smem_alloc(F16, [1, lds_total_fp16], name_hint="lds_b") + b.smem_alloc(io_type, [1, lds_total_fp16], name_hint="lds_b") if spec.double_buffer else A_smem ) @@ -369,7 +425,7 @@ def build_direct_conv_16c( c_half_bytes = b.const_i32(2) oob_sentinel = b.const_i32((1 << 31) - 1) - fp16x4_zero = b.zero_vec_f16(4) + fp16x4_zero = b.zero_vec(io_type, 4) zero_acc = b.zero_vec_f32(4) # ---- weight loads (constant across H-loop) ---- @@ -396,7 +452,7 @@ def build_direct_conv_16c( # c4 in {2, 3}) are zero-padded, so its accumulator chain stays the # same width as the S=0/1 atom (see the MFMA comment below). lane_in_lo_half = b.cmp_lt(c4, b.const_i32(2)) - fp16x8_zero = b.zero_vec_f16(8) + fp16x8_zero = b.zero_vec(io_type, 8) if spec.fold_k32: for r_const in range(p.KH): r_i = b.const_i32(r_const) @@ -410,7 +466,7 @@ def build_direct_conv_16c( c=ch_lane_k32, ) weights_k32.append( - b.buffer_load_vN_f16(b_rsrc, b.mul(w_off_k32, c_half_bytes), c0, 4) + _buf_load_vN(b, p.dtype, b_rsrc, b.mul(w_off_k32, c_half_bytes), c0, 4) ) # Residual S=2 promoted to a zero-padded K=32 atom. The low # half (c4 in {0,1}) carries B[k_out, r, 2, 0:8] / [8:16]; the @@ -422,7 +478,9 @@ def build_direct_conv_16c( s=b.const_i32(2), c=ch_lane_k32, ) - w_s2 = b.buffer_load_vN_f16(b_rsrc, b.mul(w_off_s2, c_half_bytes), c0, 4) + w_s2 = _buf_load_vN( + b, p.dtype, b_rsrc, b.mul(w_off_s2, c_half_bytes), c0, 4 + ) weights_s2_k32.append(b.select(lane_in_lo_half, w_s2, fp16x8_zero)) else: for r_const in range(p.KH): @@ -437,7 +495,7 @@ def build_direct_conv_16c( c=ch_lane_k16, ) weights.append( - b.buffer_load_vN_f16(b_rsrc, b.mul(w_off, c_half_bytes), c0, 2) + _buf_load_vN(b, p.dtype, b_rsrc, b.mul(w_off, c_half_bytes), c0, 2) ) # ---- LDS load helper ---- @@ -579,7 +637,7 @@ def issue_dram_load(y_iter_val: Value): valid = b.land(addr_valid, cm["in_bounds"]) a_off_bytes = b.mul(a_off_elems, c_half_bytes) safe_off = b.select(valid, a_off_bytes, oob_sentinel) - a_vec = b.buffer_load_vN_f16(a_rsrc, safe_off, c0, 2) + a_vec = _buf_load_vN(b, p.dtype, a_rsrc, safe_off, c0, 2) a_vec = b.select(valid, a_vec, fp16x4_zero) # LDS index in halves: chunk_idx * 4. Allocation is 2D # `[1, ROW]` so we pass (row=0, col=lds_idx). @@ -589,7 +647,7 @@ def issue_dram_load(y_iter_val: Value): def store_to_lds(loads, lds: Value) -> None: for a_vec, lds_idx in loads: - b.smem_store_vN_f16(lds, [c0, lds_idx], a_vec, 4) + b.smem_store_vN(lds, [c0, lds_idx], a_vec, 4) q_subtiles = BLOCK_Q // 16 @@ -613,7 +671,7 @@ def lds_read_input(q_subtile: int, s_const: int, lds: Value) -> Value: ), b.mul(c4, b.const_i32(4)), ) - return b.smem_load_vN_f16(lds, c0, lds_idx, n=4) + return b.smem_load_vN(lds, c0, lds_idx, dtype=io_type, n=4) def lds_read_input_k32(q_subtile: int, lds: Value) -> Value: """Per-lane <8 x half> read for the folded K=32 MFMA. @@ -633,7 +691,7 @@ def lds_read_input_k32(q_subtile: int, lds: Value) -> Value: ), ch_lane_k32, ) - return b.smem_load_vN_f16(lds, c0, lds_idx, n=8) + return b.smem_load_vN(lds, c0, lds_idx, dtype=io_type, n=8) def lds_read_input_s2_k32(q_subtile: int, lds: Value) -> Value: """Per-lane <8 x half> input read for the S=2 residual, promoted to @@ -657,7 +715,7 @@ def lds_read_input_s2_k32(q_subtile: int, lds: Value) -> Value: ), ch_lane_k32, ) - vec = b.smem_load_vN_f16(lds, c0, lds_idx, n=8) + vec = b.smem_load_vN(lds, c0, lds_idx, dtype=io_type, n=8) return b.select(lane_in_lo_half, vec, fp16x8_zero) # ---- prologue: prefetch row 0 (= -PAD..-PAD+1 = -1) into A_smem ---- @@ -752,19 +810,32 @@ def lds_read_input_s2_k32(q_subtile: int, lds: Value) -> Value: # Migrating the C-accumulator readout + A/B K-pack to # c_layout().coord would delete this whole hazard class at # the source; tracked as a follow-up. - acc_in = b.mfma_f32_16x16x32_f16( - weights_k32[r_const], input_k32, acc_in - ) - acc_in = b.mfma_f32_16x16x32_f16( - weights_s2_k32[r_const], input_s2, acc_in - ) + if p.dtype == "bf16": + acc_in = b.mfma_f32_16x16x32_bf16( + weights_k32[r_const], input_k32, acc_in + ) + acc_in = b.mfma_f32_16x16x32_bf16( + weights_s2_k32[r_const], input_s2, acc_in + ) + else: + acc_in = b.mfma_f32_16x16x32_f16( + weights_k32[r_const], input_k32, acc_in + ) + acc_in = b.mfma_f32_16x16x32_f16( + weights_s2_k32[r_const], input_s2, acc_in + ) else: inputs = inputs_by_q[qt] for s_const in range(p.KW): w_idx = r_const * p.KW + s_const - acc_in = b.mfma_f32_16x16x16_f16( - weights[w_idx], inputs[s_const], acc_in - ) + if p.dtype == "bf16": + acc_in = b.mfma_f32_16x16x16_bf16( + weights[w_idx], inputs[s_const], acc_in + ) + else: + acc_in = b.mfma_f32_16x16x16_f16( + weights[w_idx], inputs[s_const], acc_in + ) accs[p_idx] = acc_in if loads_next is not None: @@ -828,8 +899,12 @@ def lds_read_input_s2_k32(q_subtile: int, lds: Value) -> Value: # k_out = g*kpg + c4*4 + [0..3]. Store them as one # 64-bit vector instead of four scalar buffer_store_short # ops. - acc_h = b.vec_trunc_f32_to_f16(acc_to_flush) - b.buffer_store_vN_f16(d_rsrc, safe_d_off, c0, acc_h, 2) + if p.dtype == "bf16": + acc_h = b.vec_trunc_f32_to_bf16(acc_to_flush) + b.buffer_store_vN_bf16(d_rsrc, safe_d_off, c0, acc_h, 2) + else: + acc_h = b.vec_trunc_f32_to_f16(acc_to_flush) + b.buffer_store_vN_f16(d_rsrc, safe_d_off, c0, acc_h, 2) # Unconditional slot reset — kills early-iter leaks before they # pollute a later output row. for qt in range(q_subtiles): @@ -867,6 +942,11 @@ def kernel_name(self) -> str: def validate(self) -> None: p = self.problem + if p.dtype != "fp16": + raise ValueError( + f"DirectConv4cSpec: bf16 is not supported — the mfma_f32_4x4x4 atom " + f"is fp16-only on CDNA; use fp16 dtype or a different cpg variant" + ) if p.cpg != 4 or p.kpg != 4: raise ValueError( f"DirectConv4cSpec expects cpg=kpg=4 (got {p.cpg}, {p.kpg})" @@ -892,6 +972,7 @@ def is_valid_spec_4c(spec: DirectConv4cSpec, arch: str = "gfx950") -> Tuple[bool is deliberately not gated through the MMA catalog ``has_shape`` check because the catalog lists only the warp-tile (16x16 / 32x32) shapes, while comgr selects the 4x4x4 intrinsic directly on both targets. + bf16 is not supported — no 4x4x4 bf16 MFMA atom exists on CDNA. """ from rocke.core.arch import ArchTarget @@ -900,6 +981,10 @@ def is_valid_spec_4c(spec: DirectConv4cSpec, arch: str = "gfx950") -> Tuple[bool except KeyError as e: return False, str(e) p = spec.problem + if p.dtype != "fp16": + return False, ( + f"DirectConv4cSpec: bf16 not supported — no mfma_f32_4x4x4_bf16 atom on CDNA" + ) if p.stride != 1: return False, f"stride > 1 is not supported (got {p.stride})" if p.cpg != 4 or p.kpg != 4: @@ -1172,10 +1257,13 @@ def kernel_name(self) -> str: f"bq{self.block_q}", f"bg{self.block_groups}", "db" if self.double_buffer else "sb", + flags={"bf16": p.dtype == "bf16"}, ) def validate(self) -> None: p = self.problem + if p.dtype not in ("fp16", "bf16"): + raise ValueError(f"DirectConv8cSpec: unsupported dtype {p.dtype!r}") if p.cpg != 8 or p.kpg != 8: raise ValueError( f"DirectConv8cSpec expects cpg=kpg=8 (got {p.cpg}, {p.kpg})" @@ -1192,15 +1280,18 @@ def is_valid_spec_8c(spec: DirectConv8cSpec, arch: str = "gfx950") -> Tuple[bool """Return ``(ok, reason)`` for an 8c spec on ``arch``. The 8c kernel folds two S-positions into the K=16 dimension of - ``mfma_f32_16x16x16_f16``, which is present on both gfx942 and gfx950. + ``mfma_f32_16x16x16_f16`` (or the bf16 counterpart), which is present on + both gfx942 and gfx950. """ from rocke.core.arch import ArchTarget try: - ArchTarget.from_gfx(arch) + target = ArchTarget.from_gfx(arch) except KeyError as e: return False, str(e) p = spec.problem + if p.dtype not in ("fp16", "bf16"): + return False, f"unsupported dtype {p.dtype!r}; expected 'fp16' or 'bf16'" if p.stride != 1: return False, f"stride > 1 is not supported (got {p.stride})" if p.cpg != 8 or p.kpg != 8: @@ -1212,6 +1303,11 @@ def is_valid_spec_8c(spec: DirectConv8cSpec, arch: str = "gfx950") -> Tuple[bool ) if spec.block_q % 16 != 0: return False, "DirectConv8cSpec block_q must be a multiple of 16" + ab_dtype = "bf16" if p.dtype == "bf16" else "f16" + if not target.mma.has_shape( + a_dtype=ab_dtype, b_dtype=ab_dtype, c_dtype="fp32", m=16, n=16, k=16 + ): + return False, f"missing 16x16x16 {ab_dtype} MFMA atom on {arch}" return True, "ok" @@ -1244,6 +1340,7 @@ def build_direct_conv_8c(spec: DirectConv8cSpec, arch: str = "gfx950") -> Kernel if not ok: raise ValueError(f"invalid direct_conv_8c spec for {arch}: {why}") p = spec.problem + io_type = _io_type(p.dtype) BLOCK_Q = spec.block_q BLOCK_GROUPS = spec.block_groups @@ -1261,9 +1358,9 @@ def build_direct_conv_8c(spec: DirectConv8cSpec, arch: str = "gfx950") -> Kernel b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = THREADS - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - Bp = b.param("B", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + Bp = b.param("B", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) B_bytes = b.param("B_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -1308,9 +1405,9 @@ def build_direct_conv_8c(spec: DirectConv8cSpec, arch: str = "gfx950") -> Kernel # LDS allocation (same over-allocation scheme as 16c to absorb OOB writes). lds_total_fp16 = PASSES * THREADS * LOAD_VEC - A_smem = b.smem_alloc(F16, [1, lds_total_fp16], name_hint="lds_a") + A_smem = b.smem_alloc(io_type, [1, lds_total_fp16], name_hint="lds_a") B_smem = ( - b.smem_alloc(F16, [1, lds_total_fp16], name_hint="lds_b") + b.smem_alloc(io_type, [1, lds_total_fp16], name_hint="lds_b") if spec.double_buffer else A_smem ) @@ -1319,7 +1416,7 @@ def build_direct_conv_8c(spec: DirectConv8cSpec, arch: str = "gfx950") -> Kernel b_rsrc = b.buffer_rsrc(Bp, B_bytes) d_rsrc = b.buffer_rsrc(D, D_bytes) - fp16x4_zero = b.zero_vec_f16(4) + fp16x4_zero = b.zero_vec(io_type, 4) zero_acc = b.zero_vec_f32(4) # Weight loads (constant across the H-loop). @@ -1348,13 +1445,13 @@ def build_direct_conv_8c(spec: DirectConv8cSpec, arch: str = "gfx950") -> Kernel # Main atom: s_lane selects s=0 (c4 ∈ {0,1}) or s=1 (c4 ∈ {2,3}). w_off_main, _ = b_desc.offset(b, k_out=k_out_val, r=r_i, s=s_lane, c=ch_lane) weights_main.append( - b.buffer_load_vN_f16(b_rsrc, b.mul(w_off_main, c_half_bytes), c0, 2) + _buf_load_vN(b, p.dtype, b_rsrc, b.mul(w_off_main, c_half_bytes), c0, 2) ) # Residual s=2: valid only for c4 ∈ {0,1} (lower half of K). w_off_s2, _ = b_desc.offset( b, k_out=k_out_val, r=r_i, s=b.const_i32(2), c=ch_lane ) - w_s2 = b.buffer_load_vN_f16(b_rsrc, b.mul(w_off_s2, c_half_bytes), c0, 2) + w_s2 = _buf_load_vN(b, p.dtype, b_rsrc, b.mul(w_off_s2, c_half_bytes), c0, 2) weights_s2.append(b.select(lane_in_lo_half, w_s2, fp16x4_zero)) # LDS loader (same chunk-decomposition algebra as 16c but with cpg=8). @@ -1434,7 +1531,7 @@ def issue_dram_load(y_iter_val): valid = b.land(addr_valid, cm["in_bounds"]) a_off_bytes = b.mul(a_off_elems, c_half_bytes) safe_off = b.select(valid, a_off_bytes, oob_sentinel) - a_vec = b.buffer_load_vN_f16(a_rsrc, safe_off, c0, 2) + a_vec = _buf_load_vN(b, p.dtype, a_rsrc, safe_off, c0, 2) a_vec = b.select(valid, a_vec, fp16x4_zero) lds_idx = b.mul(cm["chunk_idx"], b.const_i32(4)) out.append((a_vec, lds_idx)) @@ -1442,7 +1539,7 @@ def issue_dram_load(y_iter_val): def store_to_lds(loads, lds): for a_vec, lds_idx in loads: - b.smem_store_vN_f16(lds, [c0, lds_idx], a_vec, 4) + b.smem_store_vN(lds, [c0, lds_idx], a_vec, 4) q_subtiles = BLOCK_Q // 16 @@ -1467,7 +1564,7 @@ def lds_read_input_main(q_subtile: int, lds) -> Value: ), b.mul(ch_block_idx, b.const_i32(4)), ) - return b.smem_load_vN_f16(lds, c0, lds_idx, n=4) + return b.smem_load_vN(lds, c0, lds_idx, dtype=io_type, n=4) def lds_read_input_s2(q_subtile: int, lds) -> Value: """Per-lane <4 x half> read from LDS for the s=2 residual. @@ -1489,7 +1586,7 @@ def lds_read_input_s2(q_subtile: int, lds) -> Value: ), b.mul(ch_block_idx, b.const_i32(4)), ) - vec = b.smem_load_vN_f16(lds, c0, lds_idx, n=4) + vec = b.smem_load_vN(lds, c0, lds_idx, dtype=io_type, n=4) return b.select(lane_in_lo_half, vec, fp16x4_zero) # Prologue: prefetch row 0 into A_smem. @@ -1526,11 +1623,22 @@ def lds_read_input_s2(q_subtile: int, lds) -> Value: p_idx = (y - r_const) % p.KH acc_in = accs[p_idx] # Main atom: s=0 and s=1 folded into K=16. - acc_in = b.mfma_f32_16x16x16_f16( - weights_main[r_const], inp_main, acc_in - ) - # Residual atom: s=2, zero-padded to K=16 (upper lanes carry zeros). - acc_in = b.mfma_f32_16x16x16_f16(weights_s2[r_const], inp_s2, acc_in) + if p.dtype == "bf16": + acc_in = b.mfma_f32_16x16x16_bf16( + weights_main[r_const], inp_main, acc_in + ) + # Residual atom: s=2, zero-padded to K=16 (upper lanes carry zeros). + acc_in = b.mfma_f32_16x16x16_bf16( + weights_s2[r_const], inp_s2, acc_in + ) + else: + acc_in = b.mfma_f32_16x16x16_f16( + weights_main[r_const], inp_main, acc_in + ) + # Residual atom: s=2, zero-padded to K=16 (upper lanes carry zeros). + acc_in = b.mfma_f32_16x16x16_f16( + weights_s2[r_const], inp_s2, acc_in + ) accs[p_idx] = acc_in if loads_next is not None: @@ -1563,8 +1671,12 @@ def lds_read_input_s2(q_subtile: int, lds) -> Value: d_base_bytes = b.mul(d_base, c_half_bytes) store_valid = b.land(out_q_valid, c4_valid) safe_d_off = b.select(store_valid, d_base_bytes, oob_sentinel) - acc_h = b.vec_trunc_f32_to_f16(acc_to_flush) - b.buffer_store_vN_f16(d_rsrc, safe_d_off, c0, acc_h, 2) + if p.dtype == "bf16": + acc_h = b.vec_trunc_f32_to_bf16(acc_to_flush) + b.buffer_store_vN_bf16(d_rsrc, safe_d_off, c0, acc_h, 2) + else: + acc_h = b.vec_trunc_f32_to_f16(acc_to_flush) + b.buffer_store_vN_f16(d_rsrc, safe_d_off, c0, acc_h, 2) for qt in range(q_subtiles): acc_tiles[qt][P_FLUSH] = zero_acc @@ -1627,10 +1739,13 @@ def kernel_name(self) -> str: f"bq{self.block_q}", f"bg{self.block_groups}", "db" if self.double_buffer else "sb", + flags={"bf16": p.dtype == "bf16"}, ) def validate(self) -> None: p = self.problem + if p.dtype not in ("fp16", "bf16"): + raise ValueError(f"DirectConv32cSpec: unsupported dtype {p.dtype!r}") if p.cpg != 32 or p.kpg != 32: raise ValueError( f"DirectConv32cSpec expects cpg=kpg=32 (got {p.cpg}, {p.kpg})" @@ -1648,8 +1763,8 @@ def is_valid_spec_32c( ) -> Tuple[bool, str]: """Return ``(ok, reason)`` for a 32c spec on ``arch``. - The 32c kernel uses ``mfma_f32_32x32x8_f16``, which is present in the - rocke MMA catalog for both gfx942 and gfx950. + The 32c kernel uses ``mfma_f32_32x32x8_f16`` or ``mfma_f32_32x32x8_bf16``, + which are present in the rocke MMA catalog for both gfx942 and gfx950. """ from rocke.core.arch import ArchTarget @@ -1658,6 +1773,8 @@ def is_valid_spec_32c( except KeyError as e: return False, str(e) p = spec.problem + if p.dtype not in ("fp16", "bf16"): + return False, f"unsupported dtype {p.dtype!r}; expected 'fp16' or 'bf16'" if p.stride != 1: return False, f"stride > 1 is not supported (got {p.stride})" if p.cpg != 32 or p.kpg != 32: @@ -1669,6 +1786,11 @@ def is_valid_spec_32c( ) if spec.block_q % 32 != 0: return False, "DirectConv32cSpec block_q must be a multiple of 32" + ab_dtype = "bf16" if p.dtype == "bf16" else "f16" + if not target.mma.has_shape( + a_dtype=ab_dtype, b_dtype=ab_dtype, c_dtype="fp32", m=32, n=32, k=8 + ): + return False, f"missing 32x32x8 {ab_dtype} MFMA atom on {arch}" return True, "ok" @@ -1696,6 +1818,7 @@ def build_direct_conv_32c(spec: DirectConv32cSpec, arch: str = "gfx950") -> Kern if not ok: raise ValueError(f"invalid direct_conv_32c spec for {arch}: {why}") p = spec.problem + io_type = _io_type(p.dtype) BLOCK_Q = spec.block_q BLOCK_GROUPS = spec.block_groups @@ -1713,9 +1836,9 @@ def build_direct_conv_32c(spec: DirectConv32cSpec, arch: str = "gfx950") -> Kern b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = THREADS - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - Bp = b.param("B", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + Bp = b.param("B", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) B_bytes = b.param("B_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -1734,7 +1857,7 @@ def build_direct_conv_32c(spec: DirectConv32cSpec, arch: str = "gfx950") -> Kern tid = b.thread_id_x() wave_id = b.div(tid, c_wave) lane = b.mod(tid, c_wave) - # Lane decomposition for mfma_f32_32x32x8_f16 (wave64): + # Lane decomposition for mfma_f32_32x32x8_{f16,bf16} (wave64): # q_in_lane = lane % 32 → M row (k_out within group) and N col (output W) # k_blk = lane // 32 → K-block (0 → ch=0..3, 1 → ch=4..7 within each atom) q_in_lane = b.mod(lane, b.const_i32(32)) @@ -1748,9 +1871,9 @@ def build_direct_conv_32c(spec: DirectConv32cSpec, arch: str = "gfx950") -> Kern q_tile_start = b.mul(bx, c_BQ) lds_total_fp16 = PASSES * THREADS * LOAD_VEC - A_smem = b.smem_alloc(F16, [1, lds_total_fp16], name_hint="lds_a") + A_smem = b.smem_alloc(io_type, [1, lds_total_fp16], name_hint="lds_a") B_smem = ( - b.smem_alloc(F16, [1, lds_total_fp16], name_hint="lds_b") + b.smem_alloc(io_type, [1, lds_total_fp16], name_hint="lds_b") if spec.double_buffer else A_smem ) @@ -1759,10 +1882,10 @@ def build_direct_conv_32c(spec: DirectConv32cSpec, arch: str = "gfx950") -> Kern b_rsrc = b.buffer_rsrc(Bp, B_bytes) d_rsrc = b.buffer_rsrc(D, D_bytes) - fp16x4_zero = b.zero_vec_f16(4) + fp16x4_zero = b.zero_vec(io_type, 4) zero_acc = b.zero_vec_f32(16) - # Weight loads: per (r, s, atom_idx), each lane loads 4 f16 at + # Weight loads: per (r, s, atom_idx), each lane loads 4 elements at # weight[k_out_val, r, s, ch_start + k_blk*4 .. ch_start + k_blk*4 + 3] # where ch_start = atom_idx * 8, k_out_val = g*kpg + q_in_lane. b_desc = TensorDescriptor.naive( @@ -1789,7 +1912,7 @@ def build_direct_conv_32c(spec: DirectConv32cSpec, arch: str = "gfx950") -> Kern c=ch_off, ) weights_rs.append( - b.buffer_load_vN_f16(b_rsrc, b.mul(w_off, c_half_bytes), c0, 2) + _buf_load_vN(b, p.dtype, b_rsrc, b.mul(w_off, c_half_bytes), c0, 2) ) weights_r.append(weights_rs) weights.append(weights_r) @@ -1869,7 +1992,7 @@ def issue_dram_load(y_iter_val): valid = b.land(addr_valid, cm["in_bounds"]) a_off_bytes = b.mul(a_off_elems, c_half_bytes) safe_off = b.select(valid, a_off_bytes, oob_sentinel) - a_vec = b.buffer_load_vN_f16(a_rsrc, safe_off, c0, 2) + a_vec = _buf_load_vN(b, p.dtype, a_rsrc, safe_off, c0, 2) a_vec = b.select(valid, a_vec, fp16x4_zero) lds_idx = b.mul(cm["chunk_idx"], b.const_i32(4)) out.append((a_vec, lds_idx)) @@ -1877,10 +2000,10 @@ def issue_dram_load(y_iter_val): def store_to_lds(loads, lds): for a_vec, lds_idx in loads: - b.smem_store_vN_f16(lds, [c0, lds_idx], a_vec, 4) + b.smem_store_vN(lds, [c0, lds_idx], a_vec, 4) def lds_read_input(q_subtile: int, s_const: int, atom_idx: int, lds) -> Value: - """Per-lane <4 x half> read from LDS for one K-atom of the 32c kernel. + """Per-lane <4 x half/bfloat> read from LDS for one K-atom of the 32c kernel. Lane ``q_in_lane`` owns output column ``q_subtile*32 + q_in_lane``. LDS offset: ``(q_in_lane + q_subtile*32) * stride + s_const``. @@ -1901,7 +2024,7 @@ def lds_read_input(q_subtile: int, s_const: int, atom_idx: int, lds) -> Value: ), ch_off, ) - return b.smem_load_vN_f16(lds, c0, lds_idx, n=4) + return b.smem_load_vN(lds, c0, lds_idx, dtype=io_type, n=4) q_subtiles = BLOCK_Q // 32 @@ -1942,11 +2065,18 @@ def lds_read_input(q_subtile: int, s_const: int, atom_idx: int, lds) -> Value: acc_in = accs[p_idx] for s_const in range(p.KW): for atom_idx in range(4): - acc_in = b.mfma_f32_32x32x8_f16( - weights[r_const][s_const][atom_idx], - inputs_by_q[qt][s_const][atom_idx], - acc_in, - ) + if p.dtype == "bf16": + acc_in = b.mfma_f32_32x32x8_bf16( + weights[r_const][s_const][atom_idx], + inputs_by_q[qt][s_const][atom_idx], + acc_in, + ) + else: + acc_in = b.mfma_f32_32x32x8_f16( + weights[r_const][s_const][atom_idx], + inputs_by_q[qt][s_const][atom_idx], + acc_in, + ) accs[p_idx] = acc_in if loads_next is not None: @@ -1986,8 +2116,12 @@ def lds_read_input(q_subtile: int, s_const: int, atom_idx: int, lds) -> Value: e0 = b.vec_extract(acc_to_flush, slot0) e1 = b.vec_extract(acc_to_flush, slot1) pair_acc = b.vec_pack([e0, e1], F32) - pair_h = b.vec_trunc_f32_to_f16(pair_acc) - b.buffer_store_vN_f16(d_rsrc, safe_d_off, c0, pair_h, 1) + if p.dtype == "bf16": + pair_h = b.vec_trunc_f32_to_bf16(pair_acc) + b.buffer_store_vN_bf16(d_rsrc, safe_d_off, c0, pair_h, 1) + else: + pair_h = b.vec_trunc_f32_to_f16(pair_acc) + b.buffer_store_vN_f16(d_rsrc, safe_d_off, c0, pair_h, 1) for qt in range(q_subtiles): acc_tiles[qt][P_FLUSH] = zero_acc @@ -2055,6 +2189,7 @@ def kernel_name(self) -> str: wk_flag = f"wk{self.waves_k}" if self.waves_k > 1 else "" rk_flag = "rk" if self.runtime_k_loop else "" k32_flag = "k32" if self.fold_k32 else "" + bf16_flag = "bf16" if p.dtype == "bf16" else "" return kernel_name_join( self.name, p.short(), @@ -2066,10 +2201,13 @@ def kernel_name(self) -> str: wk_flag, rk_flag, k32_flag, + bf16_flag, ) def validate(self) -> None: p = self.problem + if p.dtype not in ("fp16", "bf16"): + raise ValueError(f"DirectConvSpec: unsupported dtype {p.dtype!r}") if p.cpg % 4 != 0 or p.cpg < 4: raise ValueError( f"DirectConvSpec requires cpg to be a positive multiple of 4 " @@ -2125,6 +2263,8 @@ def is_valid_spec(spec: "DirectConvSpec", arch: str = "gfx950") -> Tuple[bool, s return False, str(e) p = spec.problem + if p.dtype not in ("fp16", "bf16"): + return False, f"unsupported dtype {p.dtype!r}; expected 'fp16' or 'bf16'" if p.stride != 1: return False, f"stride > 1 is not supported (got {p.stride})" if p.cpg % 4 != 0 or p.cpg < 4: @@ -2138,10 +2278,15 @@ def is_valid_spec(spec: "DirectConvSpec", arch: str = "gfx950") -> Tuple[bool, s ) if spec.block_q % 16 != 0: return False, "block_q must be a multiple of 16" + ab_dtype = "bf16" if p.dtype == "bf16" else "f16" if not target.mma.has_shape( - a_dtype="f16", b_dtype="f16", c_dtype="fp32", m=16, n=16, k=16 + a_dtype=ab_dtype, b_dtype=ab_dtype, c_dtype="fp32", m=16, n=16, k=16 + ): + return False, f"missing mfma_f32_16x16x16_{ab_dtype} on {arch}" + if spec.fold_k32 and not target.mma.has_shape( + a_dtype=ab_dtype, b_dtype=ab_dtype, c_dtype="fp32", m=16, n=16, k=32 ): - return False, f"missing mfma_f32_16x16x16_f16 on {arch}" + return False, f"fold_k32 requires mfma_f32_16x16x32_{ab_dtype} on {arch}" return True, "ok" @@ -2174,6 +2319,7 @@ def build_direct_conv(spec: "DirectConvSpec", arch: str = "gfx950") -> KernelDef raise ValueError(f"invalid DirectConvSpec for {arch}: {why}") p = spec.problem + io_type = _io_type(p.dtype) BLOCK_Q = spec.block_q BLOCK_GROUPS = spec.block_groups @@ -2217,9 +2363,9 @@ def build_direct_conv(spec: "DirectConvSpec", arch: str = "gfx950") -> KernelDef b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = THREADS - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - Bp = b.param("B", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + Bp = b.param("B", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) B_bytes = b.param("B_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -2239,8 +2385,8 @@ def build_direct_conv(spec: "DirectConvSpec", arch: str = "gfx950") -> KernelDef c_half_bytes = b.const_i32(2) oob_sentinel = b.const_i32((1 << 31) - 1) - fp16x4_zero = b.zero_vec_f16(4) - zero_acc = b.zero_vec_f32(4) # mfma_f32_16x16x16_f16: 4 f32 per lane + fp16x4_zero = b.zero_vec(io_type, 4) + zero_acc = b.zero_vec_f32(4) # mfma_f32_16x16x16_{f16,bf16}: 4 f32 per lane tid = b.thread_id_x() # Wave decomposition: wave_id encodes (group, waves_q, waves_k) as: @@ -2326,9 +2472,9 @@ def build_direct_conv(spec: "DirectConvSpec", arch: str = "gfx950") -> KernelDef # LDS loader uses the full block Q range (all waves cooperate on same LDS row). q_tile_start_lds = block_q_start - A_smem = b.smem_alloc(F16, [1, lds_total_fp16], name_hint="lds_a") + A_smem = b.smem_alloc(io_type, [1, lds_total_fp16], name_hint="lds_a") B_smem = ( - b.smem_alloc(F16, [1, lds_total_fp16], name_hint="lds_b") + b.smem_alloc(io_type, [1, lds_total_fp16], name_hint="lds_b") if spec.double_buffer else A_smem ) @@ -2442,10 +2588,10 @@ def issue_dram_load(y_iter_val: Value, g_tile_val=None): ) valid = b.land(addr_valid, cm["in_bounds"]) safe_off = b.select(valid, b.mul(a_off, c_half_bytes), oob_sentinel) - # LOAD_VEC=4 → dwordx2 (4 halves); LOAD_VEC=8 → dwordx4 (8 halves). - _n_dwords = LOAD_VEC // 2 # dwords = halves / 2 - a_vec = b.buffer_load_vN_f16(a_rsrc, safe_off, c0, _n_dwords) - _zero_vec = b.zero_vec_f16(LOAD_VEC) + # LOAD_VEC=4 → dwordx2 (4 elements); LOAD_VEC=8 → dwordx4 (8 elements). + _n_dwords = LOAD_VEC // 2 # dwords = elements / 2 + a_vec = _buf_load_vN(b, p.dtype, a_rsrc, safe_off, c0, _n_dwords) + _zero_vec = b.zero_vec(io_type, LOAD_VEC) a_vec = b.select(valid, a_vec, _zero_vec) lds_idx = b.mul(cm["chunk_idx"], b.const_i32(LOAD_VEC)) out.append((a_vec, lds_idx)) @@ -2453,7 +2599,7 @@ def issue_dram_load(y_iter_val: Value, g_tile_val=None): def store_to_lds(loads, lds) -> None: for a_vec, lds_idx in loads: - b.smem_store_vN_f16(lds, [c0, lds_idx], a_vec, LOAD_VEC) + b.smem_store_vN(lds, [c0, lds_idx], a_vec, LOAD_VEC) # ---- Persistent cell loop (when persistent_grid=True) ---- # Each of 256 blocks iterates over its assigned cells: (n, h_tile, q_tile). @@ -2570,8 +2716,13 @@ def store_to_lds(loads, lds) -> None: b.mul(block_idx, c_block_sz), b.mul(lane_id_pw, b.const_i32(LOAD_VEC)), ) - w_frag_pw = b.buffer_load_vN_f16( - b_rsrc, b.mul(elem_off, c_half_bytes), c0, _pw_n_dwords + w_frag_pw = _buf_load_vN( + b, + p.dtype, + b_rsrc, + b.mul(elem_off, c_half_bytes), + c0, + _pw_n_dwords, ) preloaded_w[(r_const, s_const, local_atom, m)] = w_frag_pw @@ -2645,12 +2796,12 @@ def store_to_lds(loads, lds) -> None: ch_off, ) # For fold_k32: load 8 halves from LDS (K=32 atom). - x_frag_pw = b.smem_load_vN_f16( - cur, c0, lds_idx_pw, n=LOAD_VEC + x_frag_pw = b.smem_load_vN( + cur, c0, lds_idx_pw, dtype=io_type, n=LOAD_VEC ) if p.cpg % K_ATOM_SIZE != 0: c4_oob_pw = b.cmp_ge(ch_off, b.const_i32(p.cpg)) - _zero_pw = b.zero_vec_f16(LOAD_VEC) + _zero_pw = b.zero_vec(io_type, LOAD_VEC) x_frag_pw = b.select(c4_oob_pw, _zero_pw, x_frag_pw) for m in range(N_M_TILES): @@ -2660,13 +2811,39 @@ def store_to_lds(loads, lds) -> None: (r_const, s_const, local_atom, m) ] if FOLD_K32: - acc_tiles[qt][m][p_idx] = b.mfma_f32_16x16x32_f16( - w_frag_pw, x_frag_pw, acc_tiles[qt][m][p_idx] - ) + if p.dtype == "bf16": + acc_tiles[qt][m][p_idx] = ( + b.mfma_f32_16x16x32_bf16( + w_frag_pw, + x_frag_pw, + acc_tiles[qt][m][p_idx], + ) + ) + else: + acc_tiles[qt][m][p_idx] = ( + b.mfma_f32_16x16x32_f16( + w_frag_pw, + x_frag_pw, + acc_tiles[qt][m][p_idx], + ) + ) else: - acc_tiles[qt][m][p_idx] = b.mfma_f32_16x16x16_f16( - w_frag_pw, x_frag_pw, acc_tiles[qt][m][p_idx] - ) + if p.dtype == "bf16": + acc_tiles[qt][m][p_idx] = ( + b.mfma_f32_16x16x16_bf16( + w_frag_pw, + x_frag_pw, + acc_tiles[qt][m][p_idx], + ) + ) + else: + acc_tiles[qt][m][p_idx] = ( + b.mfma_f32_16x16x16_f16( + w_frag_pw, + x_frag_pw, + acc_tiles[qt][m][p_idx], + ) + ) continue # skip the runtime s_loop/atom_loop below # ---- Runtime K-atom loop path (runtime_k_loop=True) ---- @@ -2722,12 +2899,12 @@ def store_to_lds(loads, lds) -> None: ), ch_off_rk, ) - x_frag_rk = b.smem_load_vN_f16( - cur, c0, lds_idx_rk, n=LOAD_VEC + x_frag_rk = b.smem_load_vN( + cur, c0, lds_idx_rk, dtype=io_type, n=LOAD_VEC ) if p.cpg % K_ATOM_SIZE != 0: c4_oob_rk = b.cmp_ge(ch_off_rk, b.const_i32(p.cpg)) - _zero_rk = b.zero_vec_f16(LOAD_VEC) + _zero_rk = b.zero_vec(io_type, LOAD_VEC) x_frag_rk = b.select(c4_oob_rk, _zero_rk, x_frag_rk) for m in range(N_M_TILES): if m * 16 >= p.kpg: @@ -2759,20 +2936,33 @@ def store_to_lds(loads, lds) -> None: b.const_i32(LOAD_VEC), ), ) - w_frag_rk = b.buffer_load_vN_f16( + w_frag_rk = _buf_load_vN( + b, + p.dtype, b_rsrc, b.mul(elem_off_rk, c_half_bytes), c0, + io_type, LOAD_VEC // 2, ) if FOLD_K32: - new_k_accs_rk[m] = b.mfma_f32_16x16x32_f16( - w_frag_rk, x_frag_rk, new_k_accs_rk[m] - ) + if p.dtype == "bf16": + new_k_accs_rk[m] = b.mfma_f32_16x16x32_bf16( + w_frag_rk, x_frag_rk, new_k_accs_rk[m] + ) + else: + new_k_accs_rk[m] = b.mfma_f32_16x16x32_f16( + w_frag_rk, x_frag_rk, new_k_accs_rk[m] + ) else: - new_k_accs_rk[m] = b.mfma_f32_16x16x16_f16( - w_frag_rk, x_frag_rk, new_k_accs_rk[m] - ) + if p.dtype == "bf16": + new_k_accs_rk[m] = b.mfma_f32_16x16x16_bf16( + w_frag_rk, x_frag_rk, new_k_accs_rk[m] + ) + else: + new_k_accs_rk[m] = b.mfma_f32_16x16x16_f16( + w_frag_rk, x_frag_rk, new_k_accs_rk[m] + ) b.scf_yield(*new_k_accs_rk) for m in range(N_M_TILES): acc_tiles[qt][m][p_idx] = k_loop_rk.results[m] @@ -2846,11 +3036,13 @@ def store_to_lds(loads, lds) -> None: ), ch_off, ) - x_frag = b.smem_load_vN_f16(cur, c0, lds_idx, n=LOAD_VEC) + x_frag = b.smem_load_vN( + cur, c0, lds_idx, dtype=io_type, n=LOAD_VEC + ) if p.cpg % K_ATOM_SIZE != 0: c4_oob = b.cmp_ge(ch_off, b.const_i32(p.cpg)) - _zero_std = b.zero_vec_f16(LOAD_VEC) + _zero_std = b.zero_vec(io_type, LOAD_VEC) x_frag = b.select(c4_oob, _zero_std, x_frag) new_atom_accs = [] @@ -2866,20 +3058,35 @@ def store_to_lds(loads, lds) -> None: s=s_iv, c=ch_off, ) - w_frag = b.buffer_load_vN_f16( - b_rsrc, b.mul(w_off, c_half_bytes), c0, LOAD_VEC // 2 + w_frag = _buf_load_vN( + b, + p.dtype, + b_rsrc, + b.mul(w_off, c_half_bytes), + c0, + LOAD_VEC // 2, ) if p.cpg % K_ATOM_SIZE != 0: w_frag = b.select(c4_oob, _zero_std, w_frag) if FOLD_K32: - new_acc = b.mfma_f32_16x16x32_f16( - w_frag, x_frag, atom_accs[m] - ) + if p.dtype == "bf16": + new_acc = b.mfma_f32_16x16x32_bf16( + w_frag, x_frag, atom_accs[m] + ) + else: + new_acc = b.mfma_f32_16x16x32_f16( + w_frag, x_frag, atom_accs[m] + ) else: - new_acc = b.mfma_f32_16x16x16_f16( - w_frag, x_frag, atom_accs[m] - ) + if p.dtype == "bf16": + new_acc = b.mfma_f32_16x16x16_bf16( + w_frag, x_frag, atom_accs[m] + ) + else: + new_acc = b.mfma_f32_16x16x16_f16( + w_frag, x_frag, atom_accs[m] + ) new_atom_accs.append(new_acc) b.scf_yield(*new_atom_accs) @@ -2994,18 +3201,26 @@ def store_to_lds(loads, lds) -> None: s = b.fadd(s, b.vec_extract(rows_f32[wk], slot)) sum_slots.append(s) partial = b.vec_pack(sum_slots, F32) - acc_h = b.vec_trunc_f32_to_f16(partial) safe_d = b.select( store_ok, b.mul(d_base, c_half_bytes), oob_sentinel ) - b.buffer_store_vN_f16(d_rsrc, safe_d, c0, acc_h, 2) + if p.dtype == "bf16": + acc_h = b.vec_trunc_f32_to_bf16(partial) + b.buffer_store_vN_bf16(d_rsrc, safe_d, c0, acc_h, 2) + else: + acc_h = b.vec_trunc_f32_to_f16(partial) + b.buffer_store_vN_f16(d_rsrc, safe_d, c0, acc_h, 2) b.sync() # allow red_lds reuse by next flush else: - acc_h = b.vec_trunc_f32_to_f16(acc_to_flush) safe_d = b.select( store_ok, b.mul(d_base, c_half_bytes), oob_sentinel ) - b.buffer_store_vN_f16(d_rsrc, safe_d, c0, acc_h, 2) + if p.dtype == "bf16": + acc_h = b.vec_trunc_f32_to_bf16(acc_to_flush) + b.buffer_store_vN_bf16(d_rsrc, safe_d, c0, acc_h, 2) + else: + acc_h = b.vec_trunc_f32_to_f16(acc_to_flush) + b.buffer_store_vN_f16(d_rsrc, safe_d, c0, acc_h, 2) for qt in range(q_subtiles): for m in range(N_M_TILES): @@ -3071,13 +3286,14 @@ def build_direct_transpose_weights_dgrad( Block: (64, 1, 1) """ p = spec.problem + io_type = _io_type(p.dtype) BLOCK = 64 b = IRBuilder(f"direct_transpose_weights_dgrad_{p.short()}") b.kernel.attrs["max_workgroup_size"] = BLOCK - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -3127,7 +3343,10 @@ def build_direct_transpose_weights_dgrad( src_safe = b.select(valid, b.mul(src_off, c_half_bytes), oob_sentinel) a_rsrc = b.buffer_rsrc(A, A_bytes) d_rsrc = b.buffer_rsrc(D, D_bytes) - val = b.buffer_load_f16(a_rsrc, src_safe, c0) + if p.dtype == "bf16": + val = b.buffer_load_bf16(a_rsrc, src_safe, c0) + else: + val = b.buffer_load_f16(a_rsrc, src_safe, c0) # Destination: W_T[c_abs, r', s', k_in_g] dst_desc = TensorDescriptor.naive( @@ -3135,7 +3354,10 @@ def build_direct_transpose_weights_dgrad( ) dst_off, _ = dst_desc.offset(b, c=c_abs, r=r_prime, s=s_prime, k=k_in_g) dst_safe = b.select(valid, b.mul(dst_off, c_half_bytes), oob_sentinel) - b.buffer_store_f16(d_rsrc, dst_safe, c0, val) + if p.dtype == "bf16": + b.buffer_store_bf16(d_rsrc, dst_safe, c0, val) + else: + b.buffer_store_f16(d_rsrc, dst_safe, c0, val) return b.kernel @@ -3214,8 +3436,9 @@ def build_direct_reorganize_weights( b = IRBuilder(f"direct_reorg_wt{'32' if fold_k32 else ''}_{p.short()}") b.kernel.attrs["max_workgroup_size"] = WAVE - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + io_type = _io_type(p.dtype) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -3280,8 +3503,8 @@ def build_direct_reorganize_weights( b, k_new=k_new_val, r=r_prime, s=s_prime, c_new=c_new_base ) safe_src = b.select(src_ok, b.mul(src_off, c_half_bytes), oob_sentinel) - # Load ELEMS_PER_LANE consecutive halves: n_dwords = ELEMS_PER_LANE // 2 - val = b.buffer_load_vN_f16(a_rsrc, safe_src, c0, ELEMS_PER_LANE // 2) + # Load ELEMS_PER_LANE consecutive elements: n_dwords = ELEMS_PER_LANE // 2 + val = _buf_load_vN(b, p.dtype, a_rsrc, safe_src, c0, ELEMS_PER_LANE // 2) # Destination: stride ELEMS_PER_LANE between consecutive lanes → coalesced. dst_off = b.add( @@ -3289,7 +3512,10 @@ def build_direct_reorganize_weights( b.mul(tid, b.const_i32(ELEMS_PER_LANE)), ) safe_dst = b.select(src_ok, b.mul(dst_off, c_half_bytes), oob_sentinel) - b.buffer_store_vN_f16(d_rsrc, safe_dst, c0, val, ELEMS_PER_LANE // 2) + if p.dtype == "bf16": + b.buffer_store_vN_bf16(d_rsrc, safe_dst, c0, val, ELEMS_PER_LANE // 2) + else: + b.buffer_store_vN_f16(d_rsrc, safe_dst, c0, val, ELEMS_PER_LANE // 2) return b.kernel @@ -3320,9 +3546,10 @@ def build_direct_coalesced_weights_dgrad( b = IRBuilder(f"direct_coa_wt_dgrad_{p.short()}") b.kernel.attrs["max_workgroup_size"] = WAVE + io_type = _io_type(p.dtype) # A = W source, D = W_coalesced destination - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -3388,7 +3615,7 @@ def build_direct_coalesced_weights_dgrad( ) src_off, _ = src_desc.offset(b, k=k_src, r=r_flip, s=s_flip, c=c_src) safe_src = b.select(k_ok, b.mul(src_off, c_half_bytes), oob_sentinel) - val = b.buffer_load_vN_f16(a_rsrc, safe_src, c0, 2) + val = _buf_load_vN(b, p.dtype, a_rsrc, safe_src, c0, 2) # Destination: W_coa[bx * WAVE * 4 + tid * 4] (coalesced: tid*4 stride) dst_off = b.add( @@ -3396,7 +3623,10 @@ def build_direct_coalesced_weights_dgrad( b.mul(tid, b.const_i32(4)), ) safe_dst = b.select(k_ok, b.mul(dst_off, c_half_bytes), oob_sentinel) - b.buffer_store_vN_f16(d_rsrc, safe_dst, c0, val, 2) + if p.dtype == "bf16": + b.buffer_store_vN_bf16(d_rsrc, safe_dst, c0, val, 2) + else: + b.buffer_store_vN_f16(d_rsrc, safe_dst, c0, val, 2) return b.kernel @@ -3557,10 +3787,13 @@ def kernel_name(self) -> str: p.short(), f"bq{self.block_q}", f"bg{self.block_groups}", + flags={"bf16": p.dtype == "bf16"}, ) def validate(self) -> None: p = self.problem + if p.dtype not in ("fp16", "bf16"): + raise ValueError(f"DirectConvDgradSpec: unsupported dtype {p.dtype!r}") if p.cpg < 1: raise ValueError(f"DirectConvDgradSpec requires cpg >= 1 (got {p.cpg})") if p.kpg < 1: @@ -3630,6 +3863,7 @@ def build_direct_conv_dgrad( raise ValueError(f"invalid DirectConvDgradSpec for {arch}: {why}") p = spec.problem + io_type = _io_type(p.dtype) BLOCK_W = spec.block_q # reuse block_q field as input-W tile BLOCK_WAVES = spec.block_groups # reuse block_groups as wave count per block WAVE = spec.wave_size @@ -3643,9 +3877,9 @@ def build_direct_conv_dgrad( b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = THREADS - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - Bp = b.param("B", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + Bp = b.param("B", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) B_bytes = b.param("B_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -3794,13 +4028,19 @@ def build_direct_conv_dgrad( k_byte_off = b.mul(k_iv, k_stride_bytes) w_byte = b.add(w_off0_bytes, k_byte_off) safe_w = b.select(tap_valid, w_byte, oob_sentinel) - w_h = b.buffer_load_f16(b_rsrc, safe_w, c0) + if p.dtype == "bf16": + w_h = b.buffer_load_bf16(b_rsrc, safe_w, c0) + else: + w_h = b.buffer_load_f16(b_rsrc, safe_w, c0) w_f32 = b.select(tap_valid, b.cast_to_f32(w_h), zero_f32) # dY[n, ho, wo, k_base + k_iv]: k is contiguous in NHWK. dy_byte = b.add(dy_off0_bytes, b.mul(k_iv, c_half_bytes)) safe_dy = b.select(tap_valid, dy_byte, oob_sentinel) - dy_h = b.buffer_load_f16(a_rsrc, safe_dy, c0) + if p.dtype == "bf16": + dy_h = b.buffer_load_bf16(a_rsrc, safe_dy, c0) + else: + dy_h = b.buffer_load_f16(a_rsrc, safe_dy, c0) dy_f32 = b.select(tap_valid, b.cast_to_f32(dy_h), zero_f32) new_acc = b.fma(w_f32, dy_f32, acc_k) @@ -3813,7 +4053,10 @@ def build_direct_conv_dgrad( safe_d = b.select( b.land(c_in_ok, wi_ok), b.mul(d_off, c_half_bytes), oob_sentinel ) - b.buffer_store_f16(d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc)) + if p.dtype == "bf16": + b.buffer_store_bf16(d_rsrc, safe_d, c0, b.trunc_f32_to_bf16(acc)) + else: + b.buffer_store_f16(d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc)) b.scf_yield(dummy_in) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h index 207da1632d9d..96a427176d3a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h @@ -76,11 +76,12 @@ struct rocke_arena; /* fwd (rocke/arena.h) */ * class DirectConvProblem: * N, H, W, groups, cpg, kpg # required * KH=3, KW=3, PAD=1, stride=1 + * dtype="fp16" # "fp16" or "bf16" * * Layouts: - * A: NHWC fp16, [N, H, W, groups*cpg] - * B: KRSC fp16, [groups*kpg, KH, KW, cpg] - * D: NHWK fp16, [N, H, W, groups*kpg] + * A: NHWC, [N, H, W, groups*cpg] + * B: KRSC, [groups*kpg, KH, KW, cpg] + * D: NHWK, [N, H, W, groups*kpg] * ===================================================================== */ typedef struct rocke_direct_conv_problem { @@ -94,6 +95,7 @@ typedef struct rocke_direct_conv_problem int KW; /* default 3 */ int PAD; /* default 1 */ int stride; /* default 1 */ + const char* dtype; /* "fp16" or "bf16"; default "fp16" */ } rocke_direct_conv_problem_t; /* DirectConvProblem with dataclass defaults (KH=KW=3, PAD=1, stride=1) and the @@ -148,25 +150,27 @@ int rocke_direct_conv_16c_n_acc_slots(const rocke_direct_conv_16c_spec_t* spec); /* kernel_name(): * kernel_name_join(name, problem.short(), f"bq{block_q}", f"bg{block_groups}", - * "db" if double_buffer else "sb", flags={"k32": fold_k32}) + * "db" if double_buffer else "sb", + * flags={"k32": fold_k32, "bf16": problem.dtype=="bf16"}) * Writes NUL-terminated into out (capacity out_cap). */ rocke_status_t rocke_direct_conv_16c_kernel_name(const rocke_direct_conv_16c_spec_t* spec, char* out, size_t out_cap); /* validate(): the hard assertions of DirectConv16cSpec.validate (cpg==kpg==16, - * groups % block_groups == 0). On a violated invariant returns ROCKE_ERR_VALUE and - * (if reason non-NULL, cap reason_cap) writes the message; else ROCKE_OK. */ + * groups % block_groups == 0, dtype in {"fp16","bf16"}). On a violated invariant + * returns ROCKE_ERR_VALUE and (if reason non-NULL, cap reason_cap) writes the + * message; else ROCKE_OK. */ rocke_status_t rocke_direct_conv_16c_validate(const rocke_direct_conv_16c_spec_t* spec, char* reason, size_t reason_cap); /* is_valid_spec_16c(spec, arch) -> (ok, reason). `arch` NULL => "gfx950". - * Checks: ArchTarget.from_gfx(arch) resolves; cpg==kpg==16; groups % block_groups - * == 0; the 16x16x16 f16 MFMA atom present on arch; and when fold_k32 the - * 16x16x32 f16 atom present on arch (absent on gfx942 -> clean reject). On reject - * writes the reason (if non-NULL) and returns false; on accept writes "ok", - * returns true. */ + * Checks: ArchTarget.from_gfx(arch) resolves; dtype in {"fp16","bf16"}; + * cpg==kpg==16; groups % block_groups == 0; the 16x16x16 {f16,bf16} MFMA atom + * present on arch; and when fold_k32 the 16x16x32 {f16,bf16} atom present on arch + * (absent on gfx942 -> clean reject). On reject writes the reason (if non-NULL) + * and returns false; on accept writes "ok", returns true. */ bool rocke_direct_conv_16c_is_valid_spec(const rocke_direct_conv_16c_spec_t* spec, const char* arch, char* reason, @@ -510,19 +514,27 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_dgrad_new( rocke_ir_builder_t* b, const rocke_direct_depthwise_dgrad_spec_t* spec, const char* arch); /* ===================================================================== * - * SIGNATURE (manifest) -- both kernels share the 6-entry ABI: - * ptr A:f16, ptr B:f16, ptr D:f16, scalar A_bytes:i32, B_bytes:i32, - * D_bytes:i32. + * SIGNATURE (manifest) -- all kernels share the 6-entry ABI: + * ptr A:{dtype}, ptr B:{dtype}, ptr D:{dtype}, scalar A_bytes:i32, + * B_bytes:i32, D_bytes:i32. * ===================================================================== */ /* Writes the 6 manifest entries into out[] (capacity out_cap) and sets * *out_count = 6. Strings live in `arena`. Returns ROCKE_OK or ROCKE_ERR_VALUE - * (NULL args / out_cap < 6). One signature serves both 16c and 4c. */ + * (NULL args / out_cap < 6). One signature serves both 16c and 4c (fp16 only). */ rocke_status_t rocke_direct_conv_signature(struct rocke_arena* arena, struct rocke_sig_entry* out, size_t out_cap, size_t* out_count); +/* Dtype-aware variant. `dtype` is "f16" (alias "fp16") or "bf16". + * NULL `dtype` defaults to "f16". */ +rocke_status_t rocke_direct_conv_signature_for_dtype(struct rocke_arena* arena, + const char* dtype, + struct rocke_sig_entry* out, + size_t out_cap, + size_t* out_count); + /* ===================================================================== * * CONVENIENCE: build -> lower to LLVM .ll text. * `arch` NULL => "gfx950". On ROCKE_OK *out_ll receives a malloc'd diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h index bdef6173fa6b..08b60dbe5085 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h @@ -89,6 +89,7 @@ typedef struct rocke_dconv_16c_ctx const rocke_direct_conv_16c_spec_t* spec; /* the DirectConv16cSpec */ const char* arch; /* NULL-normalised "gfx950" */ rocke_direct_conv_problem_t p; /* spec->problem (by value) */ + const rocke_type_t* io_type; /* rocke_f16() or rocke_bf16() depending on p.dtype */ /* ---- block-geometry scalars (Python all-caps locals) -- */ int BLOCK_Q; /* spec.block_q */ @@ -386,6 +387,7 @@ typedef struct rocke_dconv_8c_ctx const rocke_direct_conv_8c_spec_t* spec; const char* arch; rocke_direct_conv_problem_t p; + const rocke_type_t* io_type; /* rocke_f16() or rocke_bf16() */ int BLOCK_Q; int BLOCK_GROUPS; @@ -485,6 +487,7 @@ typedef struct rocke_dconv_32c_ctx const rocke_direct_conv_32c_spec_t* spec; const char* arch; rocke_direct_conv_problem_t p; + const rocke_type_t* io_type; /* rocke_f16() or rocke_bf16() */ int BLOCK_Q; int BLOCK_GROUPS; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_loop_closures.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_loop_closures.cpp index ff8b0bf83324..76bb13f4915f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_loop_closures.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_loop_closures.cpp @@ -23,6 +23,8 @@ */ #include "rocke/instance_conv_direct_grouped_internal.h" +#include + #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/ir.h" @@ -98,8 +100,11 @@ int rocke_dconv16c_issue_dram_load(rocke_dconv_16c_ctx_t* ctx, a_off_bytes = rocke_b_mul(b, a_off_elems, ctx->c_half_bytes); /* safe_off = b.select(valid, a_off_bytes, oob_sentinel) */ safe_off = rocke_b_select(b, valid, a_off_bytes, ctx->oob_sentinel); - /* a_vec = b.buffer_load_vN_f16(a_rsrc, safe_off, c0, 2) */ - a_vec = rocke_b_buffer_load_vN_f16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); + /* a_vec = _buf_load_vN(a_rsrc, safe_off, c0, 2) -- dtype-dispatched */ + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + a_vec = rocke_b_buffer_load_vN_bf16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); + else + a_vec = rocke_b_buffer_load_vN_f16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); /* a_vec = b.select(valid, a_vec, fp16x4_zero) */ a_vec = rocke_b_select(b, valid, a_vec, ctx->fp16x4_zero); /* lds_idx = b.mul(cm["chunk_idx"], b.const_i32(4)) */ @@ -127,13 +132,13 @@ void rocke_dconv16c_store_to_lds(rocke_dconv_16c_ctx_t* ctx, int i; /* for a_vec, lds_idx in loads: - * b.smem_store_vN_f16(lds, [c0, lds_idx], a_vec, 4) */ + * b.smem_store_vN(lds, [c0, lds_idx], a_vec, 4) -- generic, dtype from value */ for(i = 0; i < n; ++i) { rocke_value_t* indices[2]; indices[0] = ctx->c0; indices[1] = lds_idx[i]; - rocke_b_smem_store_vN_f16(b, lds, indices, 2, vecs[i], 4); + rocke_b_smem_store_vN(b, lds, indices, 2, vecs[i], 4); } } @@ -176,10 +181,10 @@ rocke_value_t* rocke_dconv16c_lds_read_input(rocke_dconv_16c_ctx_t* ctx, rocke_value_t* mul_c4 = rocke_b_mul(b, ctx->c4, rocke_b_const_i32(b, 4)); lds_idx = rocke_b_add(b, inner, mul_c4); } - /* return b.smem_load_vN_f16(lds, c0, lds_idx, n=4) */ + /* return b.smem_load_vN(lds, c0, lds_idx, dtype=io_type, n=4) */ indices[0] = ctx->c0; indices[1] = lds_idx; - return rocke_b_smem_load_vN_f16(b, lds, indices, 2, 4); + return rocke_b_smem_load_vN(b, lds, indices, 2, ctx->io_type, 4); } /* ===================================================================== * @@ -214,10 +219,10 @@ rocke_value_t* rocke_value_t* inner = rocke_b_add(b, mul_wlds, mul_wave); lds_idx = rocke_b_add(b, inner, ctx->ch_lane_k32); } - /* return b.smem_load_vN_f16(lds, c0, lds_idx, n=8) */ + /* return b.smem_load_vN(lds, c0, lds_idx, dtype=io_type, n=8) */ indices[0] = ctx->c0; indices[1] = lds_idx; - return rocke_b_smem_load_vN_f16(b, lds, indices, 2, 8); + return rocke_b_smem_load_vN(b, lds, indices, 2, ctx->io_type, 8); } /* ===================================================================== * @@ -257,10 +262,10 @@ rocke_value_t* rocke_dconv16c_lds_read_input_s2_k32(rocke_dconv_16c_ctx_t* ctx, rocke_value_t* inner = rocke_b_add(b, mul_wlds, mul_wave); lds_idx = rocke_b_add(b, inner, ctx->ch_lane_k32); } - /* vec = b.smem_load_vN_f16(lds, c0, lds_idx, n=8) */ + /* vec = b.smem_load_vN(lds, c0, lds_idx, dtype=io_type, n=8) */ indices[0] = ctx->c0; indices[1] = lds_idx; - vec = rocke_b_smem_load_vN_f16(b, lds, indices, 2, 8); + vec = rocke_b_smem_load_vN(b, lds, indices, 2, ctx->io_type, 8); /* return b.select(lane_in_lo_half, vec, fp16x8_zero) */ return rocke_b_select(b, ctx->lane_in_lo_half, vec, ctx->fp16x8_zero); } @@ -415,12 +420,22 @@ rocke_kernel_def_t* rocke_dconv16c_stream_h_loop(rocke_dconv_16c_ctx_t* ctx) * hazard that BOTH comgr and hipcc miscompile in this * fully-unrolled kernel, corrupting H-edge output rows in a * shape-dependent way. Op order matches Python: - * acc_in = mfma_f32_16x16x32_f16(weights_k32[r], in_k32, acc_in) - * acc_in = mfma_f32_16x16x32_f16(weights_s2_k32[r], in_s2, acc_in) */ - acc_in = rocke_b_mfma_f32_16x16x32_f16( - b, ctx->weights_k32[r_const], in_k32[qt], acc_in); - acc_in = rocke_b_mfma_f32_16x16x32_f16( - b, ctx->weights_s2_k32[r_const], in_s2[qt], acc_in); + * acc_in = mfma_f32_16x16x32_{f16,bf16}(weights_k32[r], in_k32, acc_in) + * acc_in = mfma_f32_16x16x32_{f16,bf16}(weights_s2_k32[r], in_s2, acc_in) */ + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + { + acc_in = rocke_b_mfma_f32_16x16x32_bf16( + b, ctx->weights_k32[r_const], in_k32[qt], acc_in); + acc_in = rocke_b_mfma_f32_16x16x32_bf16( + b, ctx->weights_s2_k32[r_const], in_s2[qt], acc_in); + } + else + { + acc_in = rocke_b_mfma_f32_16x16x32_f16( + b, ctx->weights_k32[r_const], in_k32[qt], acc_in); + acc_in = rocke_b_mfma_f32_16x16x32_f16( + b, ctx->weights_s2_k32[r_const], in_s2[qt], acc_in); + } } else { @@ -429,8 +444,12 @@ rocke_kernel_def_t* rocke_dconv16c_stream_h_loop(rocke_dconv_16c_ctx_t* ctx) { /* w_idx = r_const * KW + s_const */ int w_idx = r_const * KW + s_const; - acc_in = rocke_b_mfma_f32_16x16x16_f16( - b, ctx->weights[w_idx], in_s[qt][s_const], acc_in); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + acc_in = rocke_b_mfma_f32_16x16x16_bf16( + b, ctx->weights[w_idx], in_s[qt][s_const], acc_in); + else + acc_in = rocke_b_mfma_f32_16x16x16_f16( + b, ctx->weights[w_idx], in_s[qt][s_const], acc_in); } } /* accs[p_idx] = acc_in */ @@ -525,10 +544,17 @@ rocke_kernel_def_t* rocke_dconv16c_stream_h_loop(rocke_dconv_16c_ctx_t* ctx) d_base_bytes = rocke_b_mul(b, d_base, ctx->c_half_bytes); /* safe_d_off = b.select(out_q_valid, d_base_bytes, oob_sentinel) */ safe_d_off = rocke_b_select(b, out_q_valid, d_base_bytes, ctx->oob_sentinel); - /* acc_h = b.vec_trunc_f32_to_f16(acc_to_flush) */ - acc_h = rocke_b_vec_trunc_f32_to_f16(b, acc_to_flush); - /* b.buffer_store_vN_f16(d_rsrc, safe_d_off, c0, acc_h, 2) */ - rocke_b_buffer_store_vN_f16(b, ctx->d_rsrc, safe_d_off, ctx->c0, acc_h, 2); + /* acc_h = vec_trunc_f32_to_{f16,bf16}(acc_to_flush) */ + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + { + acc_h = rocke_b_vec_trunc_f32_to_bf16(b, acc_to_flush); + rocke_b_buffer_store_vN_bf16(b, ctx->d_rsrc, safe_d_off, ctx->c0, acc_h, 2); + } + else + { + acc_h = rocke_b_vec_trunc_f32_to_f16(b, acc_to_flush); + rocke_b_buffer_store_vN_f16(b, ctx->d_rsrc, safe_d_off, ctx->c0, acc_h, 2); + } } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_phases.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_phases.cpp index f580c63a4a3f..bd3bcd4fffc4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_phases.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_phases.cpp @@ -18,7 +18,9 @@ #include #include #include +#include +#include "rocke/helper_rocke.helpers.io.h" #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/instance_conv_direct_grouped.h" #include "rocke/instance_conv_direct_grouped_internal.h" @@ -61,6 +63,13 @@ bool rocke_dconv16c_prologue(rocke_dconv_16c_ctx_t* ctx) * mirror the Python local binding explicitly). */ ctx->p = spec->problem; + /* io_type = _io_type(p.dtype): f16 or bf16 IR type. */ + ctx->io_type = rocke_b_io_ir_type(b, ctx->p.dtype ? ctx->p.dtype : "fp16"); + if(ctx->io_type == NULL) + { + return false; /* builder sticky error already set by rocke_b_io_ir_type */ + } + /* ---- block-geometry scalars ---- */ ctx->BLOCK_Q = spec->block_q; ctx->BLOCK_GROUPS = spec->block_groups; @@ -88,12 +97,12 @@ bool rocke_dconv16c_prologue(rocke_dconv_16c_ctx_t* ctx) /* ---- params ---- */ { - const rocke_type_t* f16ptr = rocke_ptr_type(b, rocke_f16(), "global"); + const rocke_type_t* io_ptr = rocke_ptr_type(b, ctx->io_type, "global"); rocke_param_opts_t ro; rocke_param_opts_t wo; rocke_param_opts_t none; - /* A = b.param("A", PtrType(F16,"global"), noalias=True, readonly=True, align=16) */ + /* A = b.param("A", PtrType(io_type,"global"), noalias=True, readonly=True, align=16) */ ro = (rocke_param_opts_t){0}; ro.noalias = true; ro.noalias_set = true; @@ -101,10 +110,10 @@ bool rocke_dconv16c_prologue(rocke_dconv_16c_ctx_t* ctx) ro.readonly_set = true; ro.align = 16; ro.align_set = true; - ctx->A = rocke_b_param(b, "A", f16ptr, &ro); - ctx->Bp = rocke_b_param(b, "B", f16ptr, &ro); + ctx->A = rocke_b_param(b, "A", io_ptr, &ro); + ctx->Bp = rocke_b_param(b, "B", io_ptr, &ro); - /* D = b.param("D", PtrType(F16,"global"), noalias=True, writeonly=True, align=16) */ + /* D = b.param("D", PtrType(io_type,"global"), noalias=True, writeonly=True, align=16) */ wo = (rocke_param_opts_t){0}; wo.noalias = true; wo.noalias_set = true; @@ -112,7 +121,7 @@ bool rocke_dconv16c_prologue(rocke_dconv_16c_ctx_t* ctx) wo.writeonly_set = true; wo.align = 16; wo.align_set = true; - ctx->D = rocke_b_param(b, "D", f16ptr, &wo); + ctx->D = rocke_b_param(b, "D", io_ptr, &wo); none = (rocke_param_opts_t){0}; ctx->A_bytes = rocke_b_param(b, "A_bytes", rocke_i32(), &none); @@ -165,9 +174,9 @@ bool rocke_dconv16c_prologue(rocke_dconv_16c_ctx_t* ctx) int shape[2]; shape[0] = 1; shape[1] = ctx->lds_total_fp16; - ctx->A_smem = rocke_b_smem_alloc(b, rocke_f16(), shape, 2, "lds_a"); + ctx->A_smem = rocke_b_smem_alloc(b, ctx->io_type, shape, 2, "lds_a"); if(spec->double_buffer) - ctx->B_smem = rocke_b_smem_alloc(b, rocke_f16(), shape, 2, "lds_b"); + ctx->B_smem = rocke_b_smem_alloc(b, ctx->io_type, shape, 2, "lds_b"); else ctx->B_smem = ctx->A_smem; } @@ -179,7 +188,7 @@ bool rocke_dconv16c_prologue(rocke_dconv_16c_ctx_t* ctx) ctx->c_half_bytes = rocke_b_const_i32(b, 2); ctx->oob_sentinel = rocke_b_const_i32(b, ((int64_t)1 << 31) - 1); - ctx->fp16x4_zero = rocke_b_zero_vec_f16(b, 4); + ctx->fp16x4_zero = rocke_b_zero_vec(b, ctx->io_type, 4); ctx->zero_acc = rocke_b_zero_vec_f32(b, 4); return rocke_ir_builder_ok(b); @@ -221,7 +230,7 @@ void rocke_dconv16c_load_weights(rocke_dconv_16c_ctx_t* ctx) * fp16x8_zero = b.zero_vec_f16(8) * (emitted here, matching the Python SSA order, before the per-r loop). */ ctx->lane_in_lo_half = rocke_b_cmp_lt(b, ctx->c4, rocke_b_const_i32(b, 2)); - ctx->fp16x8_zero = rocke_b_zero_vec_f16(b, 8); + ctx->fp16x8_zero = rocke_b_zero_vec(b, ctx->io_type, 8); if(ctx->spec->fold_k32) { @@ -246,8 +255,12 @@ void rocke_dconv16c_load_weights(rocke_dconv_16c_ctx_t* ctx) rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off_k32, &valid); } - ctx->weights_k32[r_const] = rocke_b_buffer_load_vN_f16( - b, ctx->b_rsrc, rocke_b_mul(b, w_off_k32, ctx->c_half_bytes), ctx->c0, 4); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + ctx->weights_k32[r_const] = rocke_b_buffer_load_vN_bf16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off_k32, ctx->c_half_bytes), ctx->c0, 4); + else + ctx->weights_k32[r_const] = rocke_b_buffer_load_vN_f16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off_k32, ctx->c_half_bytes), ctx->c0, 4); /* Residual S=2 promoted to a zero-padded K=32 atom: low half * (c4 in {0,1}) carries B[k_out,r,2,0:8]/[8:16]; high half zero. */ @@ -261,8 +274,12 @@ void rocke_dconv16c_load_weights(rocke_dconv_16c_ctx_t* ctx) rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off_s2, &valid); } - w_s2 = rocke_b_buffer_load_vN_f16( - b, ctx->b_rsrc, rocke_b_mul(b, w_off_s2, ctx->c_half_bytes), ctx->c0, 4); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + w_s2 = rocke_b_buffer_load_vN_bf16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off_s2, ctx->c_half_bytes), ctx->c0, 4); + else + w_s2 = rocke_b_buffer_load_vN_f16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off_s2, ctx->c_half_bytes), ctx->c0, 4); ctx->weights_s2_k32[r_const] = rocke_b_select(b, ctx->lane_in_lo_half, w_s2, ctx->fp16x8_zero); } @@ -287,8 +304,12 @@ void rocke_dconv16c_load_weights(rocke_dconv_16c_ctx_t* ctx) in_values[3] = ctx->ch_lane_k16; rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off, &valid); - ctx->weights[ctx->n_weights++] = rocke_b_buffer_load_vN_f16( - b, ctx->b_rsrc, rocke_b_mul(b, w_off, ctx->c_half_bytes), ctx->c0, 2); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + ctx->weights[ctx->n_weights++] = rocke_b_buffer_load_vN_bf16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off, ctx->c_half_bytes), ctx->c0, 2); + else + ctx->weights[ctx->n_weights++] = rocke_b_buffer_load_vN_f16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off, ctx->c_half_bytes), ctx->c0, 2); } } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_32c.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_32c.cpp index 19f0a6a572b8..b8e93f116f65 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_32c.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_32c.cpp @@ -18,7 +18,9 @@ #include #include #include +#include +#include "rocke/helper_rocke.helpers.io.h" #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/instance_conv_direct_grouped.h" #include "rocke/instance_conv_direct_grouped_internal.h" @@ -52,6 +54,13 @@ bool rocke_dconv32c_prologue(rocke_dconv_32c_ctx_t* ctx) ctx->p = spec->problem; + /* io_type = _io_type(p.dtype): f16 or bf16 IR type. */ + ctx->io_type = rocke_b_io_ir_type(b, ctx->p.dtype ? ctx->p.dtype : "fp16"); + if(ctx->io_type == NULL) + { + return false; /* builder sticky error already set by rocke_b_io_ir_type */ + } + ctx->BLOCK_Q = spec->block_q; ctx->BLOCK_GROUPS = spec->block_groups; ctx->WAVE = spec->wave_size; @@ -76,7 +85,7 @@ bool rocke_dconv32c_prologue(rocke_dconv_32c_ctx_t* ctx) rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", ctx->THREADS); { - const rocke_type_t* f16ptr = rocke_ptr_type(b, rocke_f16(), "global"); + const rocke_type_t* f16ptr = rocke_ptr_type(b, ctx->io_type, "global"); rocke_param_opts_t ro; rocke_param_opts_t wo; rocke_param_opts_t none; @@ -136,10 +145,10 @@ bool rocke_dconv32c_prologue(rocke_dconv_32c_ctx_t* ctx) int shape[2]; shape[0] = 1; shape[1] = ctx->lds_total_fp16; - ctx->A_smem = rocke_b_smem_alloc(b, rocke_f16(), shape, 2, "lds_a"); + ctx->A_smem = rocke_b_smem_alloc(b, ctx->io_type, shape, 2, "lds_a"); if(spec->double_buffer) { - ctx->B_smem = rocke_b_smem_alloc(b, rocke_f16(), shape, 2, "lds_b"); + ctx->B_smem = rocke_b_smem_alloc(b, ctx->io_type, shape, 2, "lds_b"); } else { @@ -151,7 +160,7 @@ bool rocke_dconv32c_prologue(rocke_dconv_32c_ctx_t* ctx) ctx->b_rsrc = rocke_b_buffer_rsrc(b, ctx->Bp, ctx->B_bytes); ctx->d_rsrc = rocke_b_buffer_rsrc(b, ctx->D, ctx->D_bytes); - ctx->fp16x4_zero = rocke_b_zero_vec_f16(b, 4); + ctx->fp16x4_zero = rocke_b_zero_vec(b, ctx->io_type, 4); /* zero_acc = b.zero_vec_f32(16) — 16 f32 elements (32*32/64) */ ctx->zero_acc = rocke_b_zero_vec_f32(b, 16); @@ -212,8 +221,12 @@ void rocke_dconv32c_load_weights(rocke_dconv_32c_ctx_t* ctx) in_values[3] = ch_off; rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off, &valid); - ctx->weights[r_const][s_const][atom_idx] = rocke_b_buffer_load_vN_f16( - b, ctx->b_rsrc, rocke_b_mul(b, w_off, ctx->c_half_bytes), ctx->c0, 2); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + ctx->weights[r_const][s_const][atom_idx] = rocke_b_buffer_load_vN_bf16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off, ctx->c_half_bytes), ctx->c0, 2); + else + ctx->weights[r_const][s_const][atom_idx] = rocke_b_buffer_load_vN_f16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off, ctx->c_half_bytes), ctx->c0, 2); } } } @@ -404,7 +417,10 @@ static int rocke_dconv32c_issue_dram_load(rocke_dconv_32c_ctx_t* ctx, valid = rocke_b_land(b, addr_valid, ctx->chunk_meta[i].in_bounds); a_off_bytes = rocke_b_mul(b, a_off_elems, ctx->c_half_bytes); safe_off = rocke_b_select(b, valid, a_off_bytes, ctx->oob_sentinel); - a_vec = rocke_b_buffer_load_vN_f16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + a_vec = rocke_b_buffer_load_vN_bf16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); + else + a_vec = rocke_b_buffer_load_vN_f16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); a_vec = rocke_b_select(b, valid, a_vec, ctx->fp16x4_zero); lds_idx = rocke_b_mul(b, ctx->chunk_meta[i].chunk_idx, rocke_b_const_i32(b, 4)); @@ -428,7 +444,7 @@ static void rocke_dconv32c_store_to_lds(rocke_dconv_32c_ctx_t* ctx, rocke_value_t* indices[2]; indices[0] = ctx->c0; indices[1] = lds_idx[i]; - rocke_b_smem_store_vN_f16(b, lds, indices, 2, vecs[i], 4); + rocke_b_smem_store_vN(b, lds, indices, 2, vecs[i], 4); } } @@ -476,7 +492,7 @@ static rocke_value_t* rocke_dconv32c_lds_read_input( } indices[0] = ctx->c0; indices[1] = lds_idx; - return rocke_b_smem_load_vN_f16(b, lds, indices, 2, 4); + return rocke_b_smem_load_vN(b, lds, indices, 2, ctx->io_type, 4); } /* ===================================================================== * @@ -588,11 +604,18 @@ rocke_kernel_def_t* rocke_dconv32c_stream_h_loop(rocke_dconv_32c_ctx_t* ctx) { for(atom_idx = 0; atom_idx < 4; ++atom_idx) { - acc_in - = rocke_b_mfma_f32_32x32x8_f16(b, - ctx->weights[r_const][s_const][atom_idx], - inputs_by_q[qt][s_const][atom_idx], - acc_in); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + acc_in = rocke_b_mfma_f32_32x32x8_bf16( + b, + ctx->weights[r_const][s_const][atom_idx], + inputs_by_q[qt][s_const][atom_idx], + acc_in); + else + acc_in = rocke_b_mfma_f32_32x32x8_f16( + b, + ctx->weights[r_const][s_const][atom_idx], + inputs_by_q[qt][s_const][atom_idx], + acc_in); } } ctx->acc_tiles[qt][p_idx] = acc_in; @@ -698,8 +721,18 @@ rocke_kernel_def_t* rocke_dconv32c_stream_h_loop(rocke_dconv_32c_ctx_t* ctx) pair_elems[0] = e0; pair_elems[1] = e1; pair_f32 = rocke_b_vec_pack(b, pair_elems, 2, rocke_f32()); - pair_h = rocke_b_vec_trunc_f32_to_f16(b, pair_f32); - rocke_b_buffer_store_vN_f16(b, ctx->d_rsrc, safe_d_off, ctx->c0, pair_h, 1); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + { + pair_h = rocke_b_vec_trunc_f32_to_bf16(b, pair_f32); + rocke_b_buffer_store_vN_bf16( + b, ctx->d_rsrc, safe_d_off, ctx->c0, pair_h, 1); + } + else + { + pair_h = rocke_b_vec_trunc_f32_to_f16(b, pair_f32); + rocke_b_buffer_store_vN_f16( + b, ctx->d_rsrc, safe_d_off, ctx->c0, pair_h, 1); + } } } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_8c.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_8c.cpp index 784551e50d01..d509478241ac 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_8c.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_8c.cpp @@ -20,7 +20,9 @@ #include #include #include +#include +#include "rocke/helper_rocke.helpers.io.h" #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/instance_conv_direct_grouped.h" #include "rocke/instance_conv_direct_grouped_internal.h" @@ -54,6 +56,13 @@ bool rocke_dconv8c_prologue(rocke_dconv_8c_ctx_t* ctx) ctx->p = spec->problem; + /* io_type = _io_type(p.dtype): f16 or bf16 IR type. */ + ctx->io_type = rocke_b_io_ir_type(b, ctx->p.dtype ? ctx->p.dtype : "fp16"); + if(ctx->io_type == NULL) + { + return false; /* builder sticky error already set by rocke_b_io_ir_type */ + } + /* ---- block geometry ---- */ ctx->BLOCK_Q = spec->block_q; ctx->BLOCK_GROUPS = spec->block_groups; @@ -79,7 +88,7 @@ bool rocke_dconv8c_prologue(rocke_dconv_8c_ctx_t* ctx) /* ---- params ---- */ { - const rocke_type_t* f16ptr = rocke_ptr_type(b, rocke_f16(), "global"); + const rocke_type_t* f16ptr = rocke_ptr_type(b, ctx->io_type, "global"); rocke_param_opts_t ro; rocke_param_opts_t wo; rocke_param_opts_t none; @@ -153,10 +162,10 @@ bool rocke_dconv8c_prologue(rocke_dconv_8c_ctx_t* ctx) int shape[2]; shape[0] = 1; shape[1] = ctx->lds_total_fp16; - ctx->A_smem = rocke_b_smem_alloc(b, rocke_f16(), shape, 2, "lds_a"); + ctx->A_smem = rocke_b_smem_alloc(b, ctx->io_type, shape, 2, "lds_a"); if(spec->double_buffer) { - ctx->B_smem = rocke_b_smem_alloc(b, rocke_f16(), shape, 2, "lds_b"); + ctx->B_smem = rocke_b_smem_alloc(b, ctx->io_type, shape, 2, "lds_b"); } else { @@ -168,7 +177,7 @@ bool rocke_dconv8c_prologue(rocke_dconv_8c_ctx_t* ctx) ctx->b_rsrc = rocke_b_buffer_rsrc(b, ctx->Bp, ctx->B_bytes); ctx->d_rsrc = rocke_b_buffer_rsrc(b, ctx->D, ctx->D_bytes); - ctx->fp16x4_zero = rocke_b_zero_vec_f16(b, 4); + ctx->fp16x4_zero = rocke_b_zero_vec(b, ctx->io_type, 4); ctx->zero_acc = rocke_b_zero_vec_f32(b, 4); return rocke_ir_builder_ok(b); @@ -225,8 +234,12 @@ void rocke_dconv8c_load_weights(rocke_dconv_8c_ctx_t* ctx) rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off_main, &valid); } - ctx->weights_main[r_const] = rocke_b_buffer_load_vN_f16( - b, ctx->b_rsrc, rocke_b_mul(b, w_off_main, ctx->c_half_bytes), ctx->c0, 2); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + ctx->weights_main[r_const] = rocke_b_buffer_load_vN_bf16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off_main, ctx->c_half_bytes), ctx->c0, 2); + else + ctx->weights_main[r_const] = rocke_b_buffer_load_vN_f16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off_main, ctx->c_half_bytes), ctx->c0, 2); /* Residual s=2 atom: valid only for c4∈{0,1}; zeroed for c4∈{2,3}. */ { @@ -239,8 +252,12 @@ void rocke_dconv8c_load_weights(rocke_dconv_8c_ctx_t* ctx) rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off_s2, &valid); } - w_s2 = rocke_b_buffer_load_vN_f16( - b, ctx->b_rsrc, rocke_b_mul(b, w_off_s2, ctx->c_half_bytes), ctx->c0, 2); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + w_s2 = rocke_b_buffer_load_vN_bf16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off_s2, ctx->c_half_bytes), ctx->c0, 2); + else + w_s2 = rocke_b_buffer_load_vN_f16( + b, ctx->b_rsrc, rocke_b_mul(b, w_off_s2, ctx->c_half_bytes), ctx->c0, 2); ctx->weights_s2[r_const] = rocke_b_select(b, ctx->lane_in_lo_half, w_s2, ctx->fp16x4_zero); } ctx->n_weights = ctx->p.KH; @@ -438,7 +455,10 @@ static int rocke_dconv8c_issue_dram_load(rocke_dconv_8c_ctx_t* ctx, valid = rocke_b_land(b, addr_valid, ctx->chunk_meta[i].in_bounds); a_off_bytes = rocke_b_mul(b, a_off_elems, ctx->c_half_bytes); safe_off = rocke_b_select(b, valid, a_off_bytes, ctx->oob_sentinel); - a_vec = rocke_b_buffer_load_vN_f16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + a_vec = rocke_b_buffer_load_vN_bf16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); + else + a_vec = rocke_b_buffer_load_vN_f16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); a_vec = rocke_b_select(b, valid, a_vec, ctx->fp16x4_zero); lds_idx = rocke_b_mul(b, ctx->chunk_meta[i].chunk_idx, rocke_b_const_i32(b, 4)); @@ -462,7 +482,7 @@ static void rocke_dconv8c_store_to_lds(rocke_dconv_8c_ctx_t* ctx, rocke_value_t* indices[2]; indices[0] = ctx->c0; indices[1] = lds_idx[i]; - rocke_b_smem_store_vN_f16(b, lds, indices, 2, vecs[i], 4); + rocke_b_smem_store_vN(b, lds, indices, 2, vecs[i], 4); } } @@ -500,7 +520,7 @@ static rocke_value_t* } indices[0] = ctx->c0; indices[1] = lds_idx; - return rocke_b_smem_load_vN_f16(b, lds, indices, 2, 4); + return rocke_b_smem_load_vN(b, lds, indices, 2, ctx->io_type, 4); } /* lds_read_input_s2(q_subtile, lds) -- per-lane <4 x half> read for the s=2 residual. @@ -536,7 +556,7 @@ static rocke_value_t* } indices[0] = ctx->c0; indices[1] = lds_idx; - vec = rocke_b_smem_load_vN_f16(b, lds, indices, 2, 4); + vec = rocke_b_smem_load_vN(b, lds, indices, 2, ctx->io_type, 4); return rocke_b_select(b, ctx->lane_in_lo_half, vec, ctx->fp16x4_zero); } @@ -642,11 +662,19 @@ rocke_kernel_def_t* rocke_dconv8c_stream_h_loop(rocke_dconv_8c_ctx_t* ctx) rocke_value_t* acc_in = ctx->acc_tiles[qt][p_idx]; /* acc_in = mfma_16x16x16(weights_main[r], in_main, acc_in) */ - acc_in = rocke_b_mfma_f32_16x16x16_f16( - b, ctx->weights_main[r_const], in_main[qt], acc_in); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + acc_in = rocke_b_mfma_f32_16x16x16_bf16( + b, ctx->weights_main[r_const], in_main[qt], acc_in); + else + acc_in = rocke_b_mfma_f32_16x16x16_f16( + b, ctx->weights_main[r_const], in_main[qt], acc_in); /* acc_in = mfma_16x16x16(weights_s2[r], in_s2, acc_in) */ - acc_in - = rocke_b_mfma_f32_16x16x16_f16(b, ctx->weights_s2[r_const], in_s2[qt], acc_in); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + acc_in = rocke_b_mfma_f32_16x16x16_bf16( + b, ctx->weights_s2[r_const], in_s2[qt], acc_in); + else + acc_in = rocke_b_mfma_f32_16x16x16_f16( + b, ctx->weights_s2[r_const], in_s2[qt], acc_in); ctx->acc_tiles[qt][p_idx] = acc_in; } } @@ -722,8 +750,16 @@ rocke_kernel_def_t* rocke_dconv8c_stream_h_loop(rocke_dconv_8c_ctx_t* ctx) /* store_valid = b.land(out_q_valid, c4_valid) */ store_valid = rocke_b_land(b, out_q_valid, c4_valid); safe_d_off = rocke_b_select(b, store_valid, d_base_bytes, ctx->oob_sentinel); - acc_h = rocke_b_vec_trunc_f32_to_f16(b, acc_to_flush); - rocke_b_buffer_store_vN_f16(b, ctx->d_rsrc, safe_d_off, ctx->c0, acc_h, 2); + if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + { + acc_h = rocke_b_vec_trunc_f32_to_bf16(b, acc_to_flush); + rocke_b_buffer_store_vN_bf16(b, ctx->d_rsrc, safe_d_off, ctx->c0, acc_h, 2); + } + else + { + acc_h = rocke_b_vec_trunc_f32_to_f16(b, acc_to_flush); + rocke_b_buffer_store_vN_f16(b, ctx->d_rsrc, safe_d_off, ctx->c0, acc_h, 2); + } } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp index 3a3ee1d936ee..6127001534bd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp @@ -116,6 +116,7 @@ rocke_direct_conv_problem_t rocke_direct_conv_problem_default(void) p.KW = 3; p.PAD = 1; p.stride = 1; + p.dtype = "fp16"; return p; } @@ -200,9 +201,10 @@ rocke_status_t rocke_direct_conv_16c_kernel_name(const rocke_direct_conv_16c_spe char bg_buf[32]; const char* db_part; const char* parts[4]; - const char* flag_names[1]; - int flag_on[1]; + const char* flag_names[2]; + int flag_on[2]; rocke_status_t st; + const char* dtype; if(spec == NULL || out == NULL) { @@ -220,16 +222,19 @@ rocke_status_t rocke_direct_conv_16c_kernel_name(const rocke_direct_conv_16c_spe db_part = spec->double_buffer ? "db" : "sb"; /* kernel_name_join(name, short, "bq..", "bg..", "db"/"sb", - * flags={"k32": fold_k32}) */ + * flags={"k32": fold_k32, "bf16": dtype=="bf16"}) */ parts[0] = short_buf; parts[1] = bq_buf; parts[2] = bg_buf; parts[3] = db_part; + dtype = spec->problem.dtype ? spec->problem.dtype : "fp16"; flag_names[0] = "k32"; flag_on[0] = spec->fold_k32 ? 1 : 0; + flag_names[1] = "bf16"; + flag_on[1] = (strcmp(dtype, "bf16") == 0) ? 1 : 0; - return rocke_kernel_name_join(spec->name, parts, 4, flag_names, flag_on, 1, out, out_cap, NULL); + return rocke_kernel_name_join(spec->name, parts, 4, flag_names, flag_on, 2, out, out_cap, NULL); } rocke_status_t rocke_direct_conv_16c_validate(const rocke_direct_conv_16c_spec_t* spec, @@ -242,6 +247,18 @@ rocke_status_t rocke_direct_conv_16c_validate(const rocke_direct_conv_16c_spec_t return ROCKE_ERR_VALUE; } p = &spec->problem; + /* if p.dtype not in ("fp16", "bf16"): raise ValueError(...) */ + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) + { + if(reason != NULL && reason_cap > 0) + { + snprintf(reason, reason_cap, "DirectConv16cSpec: unsupported dtype '%s'", dt); + } + return ROCKE_ERR_VALUE; + } + } /* if p.cpg != 16 or p.kpg != 16: raise ValueError(...) */ if(p->cpg != 16 || p->kpg != 16) { @@ -309,6 +326,13 @@ bool rocke_direct_conv_16c_is_valid_spec(const rocke_direct_conv_16c_spec_t* spe } p = &spec->problem; + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) + { + CK_DCONV16C_REJECT("unsupported dtype '%s'; expected 'fp16' or 'bf16'", dt); + } + } /* if p.cpg != 16 or p.kpg != 16: return False, ... */ if(p->cpg != 16 || p->kpg != 16) { @@ -322,18 +346,23 @@ bool rocke_direct_conv_16c_is_valid_spec(const rocke_direct_conv_16c_spec_t* spe } mma = rocke_archtarget_mma(target); - /* if not target.mma.has_shape(f16,f16,fp32, 16,16,16): return False, ... */ - if(!rocke_mma_catalog_has_shape(mma, "mma", "f16", "f16", "fp32", 16, 16, 16)) - { - CK_DCONV16C_REJECT("missing 16x16x16 f16 MFMA atom on %s", arch); - } - /* if spec.fold_k32 and not target.mma.has_shape(f16,f16,fp32, 16,16,32): ... */ - if(spec->fold_k32 && !rocke_mma_catalog_has_shape(mma, "mma", "f16", "f16", "fp32", 16, 16, 32)) { - CK_DCONV16C_REJECT("fold_k32=True needs the 16x16x32 f16 MFMA atom, absent on %s; use " - "fold_k32=False for a %s-capable kernel", - arch, - arch); + /* ab_dtype: "f16" or "bf16" based on problem.dtype */ + const char* ab = (p->dtype && strcmp(p->dtype, "bf16") == 0) ? "bf16" : "f16"; + /* if not target.mma.has_shape(ab,ab,fp32, 16,16,16): return False, ... */ + if(!rocke_mma_catalog_has_shape(mma, "mma", ab, ab, "fp32", 16, 16, 16)) + { + CK_DCONV16C_REJECT("missing 16x16x16 %s MFMA atom on %s", ab, arch); + } + /* if spec.fold_k32 and not target.mma.has_shape(ab,ab,fp32, 16,16,32): ... */ + if(spec->fold_k32 && !rocke_mma_catalog_has_shape(mma, "mma", ab, ab, "fp32", 16, 16, 32)) + { + CK_DCONV16C_REJECT("fold_k32=True needs the 16x16x32 %s MFMA atom, absent on %s; use " + "fold_k32=False for a %s-capable kernel", + ab, + arch, + arch); + } } if(reason != NULL && reason_cap > 0) @@ -585,6 +614,71 @@ rocke_status_t rocke_direct_conv_signature(struct rocke_arena* arena, return ROCKE_OK; } +rocke_status_t rocke_direct_conv_signature_for_dtype(struct rocke_arena* arena, + const char* dtype, + struct rocke_sig_entry* out, + size_t out_cap, + size_t* out_count) +{ + rocke_status_t st; + /* Resolve dtype: "fp16" -> "f16", "bf16" -> "bf16", NULL -> "f16". */ + const char* dt; + + if(arena == NULL || out == NULL || out_cap < 6) + { + return ROCKE_ERR_VALUE; + } + if(dtype == NULL || strcmp(dtype, "fp16") == 0 || strcmp(dtype, "f16") == 0) + { + dt = "f16"; + } + else if(strcmp(dtype, "bf16") == 0) + { + dt = "bf16"; + } + else + { + return ROCKE_ERR_VALUE; /* unsupported dtype */ + } + + st = rocke_sig_param(arena, "A", dt, NULL, &out[0]); + if(st != ROCKE_OK) + { + return st; + } + st = rocke_sig_param(arena, "B", dt, NULL, &out[1]); + if(st != ROCKE_OK) + { + return st; + } + st = rocke_sig_param(arena, "D", dt, NULL, &out[2]); + if(st != ROCKE_OK) + { + return st; + } + st = rocke_sig_scalar(arena, "A_bytes", "i32", &out[3]); + if(st != ROCKE_OK) + { + return st; + } + st = rocke_sig_scalar(arena, "B_bytes", "i32", &out[4]); + if(st != ROCKE_OK) + { + return st; + } + st = rocke_sig_scalar(arena, "D_bytes", "i32", &out[5]); + if(st != ROCKE_OK) + { + return st; + } + + if(out_count != NULL) + { + *out_count = 6; + } + return ROCKE_OK; +} + /* ===================================================================== * * DirectConv8cSpec (cpg = kpg = 8) * ===================================================================== */ @@ -630,7 +724,15 @@ rocke_status_t rocke_direct_conv_8c_kernel_name(const rocke_direct_conv_8c_spec_ parts[1] = bq_buf; parts[2] = bg_buf; parts[3] = spec->double_buffer ? "db" : "sb"; - return rocke_kernel_name_join(spec->name, parts, 4, NULL, NULL, 0, out, out_cap, NULL); + { + /* flags={"bf16": dtype=="bf16"} */ + const char* flag_names8c[1] = {"bf16"}; + const char* dt8c = spec->problem.dtype ? spec->problem.dtype : "fp16"; + int flag_on8c[1]; + flag_on8c[0] = (strcmp(dt8c, "bf16") == 0) ? 1 : 0; + return rocke_kernel_name_join( + spec->name, parts, 4, flag_names8c, flag_on8c, 1, out, out_cap, NULL); + } } rocke_status_t rocke_direct_conv_8c_validate(const rocke_direct_conv_8c_spec_t* spec, @@ -643,6 +745,17 @@ rocke_status_t rocke_direct_conv_8c_validate(const rocke_direct_conv_8c_spec_t* return ROCKE_ERR_VALUE; } p = &spec->problem; + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) + { + if(reason && reason_cap > 0) + { + snprintf(reason, reason_cap, "DirectConv8cSpec: unsupported dtype '%s'", dt); + } + return ROCKE_ERR_VALUE; + } + } if(p->cpg != 8 || p->kpg != 8) { if(reason && reason_cap > 0) @@ -689,6 +802,10 @@ bool rocke_direct_conv_8c_is_valid_spec(const rocke_direct_conv_8c_spec_t* spec, size_t reason_cap) { const rocke_direct_conv_problem_t* p; + const rocke_archtarget_t* target; + const rocke_arch_mma_catalog_t* mma; + const char* ab; + if(spec == NULL) { if(reason && reason_cap > 0) @@ -701,12 +818,25 @@ bool rocke_direct_conv_8c_is_valid_spec(const rocke_direct_conv_8c_spec_t* spec, { arch = "gfx950"; } - if(rocke_archtarget_from_gfx(arch) == NULL) + target = rocke_archtarget_from_gfx(arch); + if(target == NULL) { rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); return false; } p = &spec->problem; + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) + { + if(reason && reason_cap > 0) + { + snprintf( + reason, reason_cap, "unsupported dtype '%s'; expected 'fp16' or 'bf16'", dt); + } + return false; + } + } if(p->cpg != 8 || p->kpg != 8) { if(reason && reason_cap > 0) @@ -739,6 +869,16 @@ bool rocke_direct_conv_8c_is_valid_spec(const rocke_direct_conv_8c_spec_t* spec, } return false; } + mma = rocke_archtarget_mma(target); + ab = (p->dtype && strcmp(p->dtype, "bf16") == 0) ? "bf16" : "f16"; + if(!rocke_mma_catalog_has_shape(mma, "mma", ab, ab, "fp32", 16, 16, 16)) + { + if(reason && reason_cap > 0) + { + snprintf(reason, reason_cap, "missing 16x16x16 %s MFMA atom on %s", ab, arch); + } + return false; + } if(reason && reason_cap > 0) { strncpy(reason, "ok", reason_cap); @@ -792,7 +932,15 @@ rocke_status_t rocke_direct_conv_32c_kernel_name(const rocke_direct_conv_32c_spe parts[1] = bq_buf; parts[2] = bg_buf; parts[3] = spec->double_buffer ? "db" : "sb"; - return rocke_kernel_name_join(spec->name, parts, 4, NULL, NULL, 0, out, out_cap, NULL); + { + /* flags={"bf16": dtype=="bf16"} */ + const char* flag_names32c[1] = {"bf16"}; + const char* dt32c = spec->problem.dtype ? spec->problem.dtype : "fp16"; + int flag_on32c[1]; + flag_on32c[0] = (strcmp(dt32c, "bf16") == 0) ? 1 : 0; + return rocke_kernel_name_join( + spec->name, parts, 4, flag_names32c, flag_on32c, 1, out, out_cap, NULL); + } } rocke_status_t rocke_direct_conv_32c_validate(const rocke_direct_conv_32c_spec_t* spec, @@ -805,6 +953,17 @@ rocke_status_t rocke_direct_conv_32c_validate(const rocke_direct_conv_32c_spec_t return ROCKE_ERR_VALUE; } p = &spec->problem; + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) + { + if(reason && reason_cap > 0) + { + snprintf(reason, reason_cap, "DirectConv32cSpec: unsupported dtype '%s'", dt); + } + return ROCKE_ERR_VALUE; + } + } if(p->cpg != 32 || p->kpg != 32) { if(reason && reason_cap > 0) @@ -851,6 +1010,10 @@ bool rocke_direct_conv_32c_is_valid_spec(const rocke_direct_conv_32c_spec_t* spe size_t reason_cap) { const rocke_direct_conv_problem_t* p; + const rocke_archtarget_t* target; + const rocke_arch_mma_catalog_t* mma; + const char* ab; + if(spec == NULL) { if(reason && reason_cap > 0) @@ -863,12 +1026,25 @@ bool rocke_direct_conv_32c_is_valid_spec(const rocke_direct_conv_32c_spec_t* spe { arch = "gfx950"; } - if(rocke_archtarget_from_gfx(arch) == NULL) + target = rocke_archtarget_from_gfx(arch); + if(target == NULL) { rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); return false; } p = &spec->problem; + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) + { + if(reason && reason_cap > 0) + { + snprintf( + reason, reason_cap, "unsupported dtype '%s'; expected 'fp16' or 'bf16'", dt); + } + return false; + } + } if(p->cpg != 32 || p->kpg != 32) { if(reason && reason_cap > 0) @@ -901,6 +1077,16 @@ bool rocke_direct_conv_32c_is_valid_spec(const rocke_direct_conv_32c_spec_t* spe } return false; } + mma = rocke_archtarget_mma(target); + ab = (p->dtype && strcmp(p->dtype, "bf16") == 0) ? "bf16" : "f16"; + if(!rocke_mma_catalog_has_shape(mma, "mma", ab, ab, "fp32", 32, 32, 8)) + { + if(reason && reason_cap > 0) + { + snprintf(reason, reason_cap, "missing 32x32x8 %s MFMA atom on %s", ab, arch); + } + return false; + } if(reason && reason_cap > 0) { strncpy(reason, "ok", reason_cap); From 51bdfad369a083eae2351f0266bf5d8a4e70860b Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Tue, 22 Sep 2026 21:52:12 +0000 Subject: [PATCH 2/5] fixes --- .../kernels/common/conv_direct_grouped.py | 1 - .../tests/test_direct_conv_correctness.py | 108 ++++++++++++++---- 2 files changed, 88 insertions(+), 21 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py index 13f14f66c0cb..98d1d48af442 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py @@ -2942,7 +2942,6 @@ def store_to_lds(loads, lds) -> None: b_rsrc, b.mul(elem_off_rk, c_half_bytes), c0, - io_type, LOAD_VEC // 2, ) if FOLD_K32: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py index 856e1359de02..829a4e30cb5c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py @@ -57,6 +57,7 @@ def _skip_reason() -> str: _SKIP_REASON = _skip_reason() _TOL = 5e-2 +_TOL_BF16 = 1e-1 # bf16 has ~3× less mantissa precision than f16 # --------------------------------------------------------------------------- @@ -150,7 +151,7 @@ def _conv_ref_grouped(A_t, B_t, p) -> "torch.Tensor": return out_nchw.permute(0, 2, 3, 1).contiguous().cuda() -def _run_grouped_one(arch: str, shape: _Shape) -> Tuple[bool, str]: +def _run_grouped_one(arch: str, shape: _Shape, dtype: str = "fp16") -> Tuple[bool, str]: """Build, compile, launch, and verify one grouped direct-conv kernel. Uses the generic ``DirectConvSpec`` dispatcher which selects the right @@ -184,6 +185,7 @@ def _run_grouped_one(arch: str, shape: _Shape) -> Tuple[bool, str]: KW=shape.KW, PAD=shape.PAD, stride=shape.stride, + dtype=dtype, ) spec = DirectConvSpec( @@ -208,11 +210,10 @@ def _run_grouped_one(arch: str, shape: _Shape) -> Tuple[bool, str]: torch.manual_seed(0) total_c = shape.groups * shape.cpg total_k = shape.groups * shape.cpg - A_t = torch.empty(p.N, p.H, p.W, total_c, dtype=torch.float16).uniform_(-1.0, 1.0) - B_t = torch.empty(total_k, p.KH, p.KW, shape.cpg, dtype=torch.float16).uniform_( - -1.0, 1.0 - ) - D_t = torch.empty(p.N, p.Ho, p.Wo, total_k, dtype=torch.float16) + _td = torch.bfloat16 if dtype == "bf16" else torch.float16 + A_t = torch.empty(p.N, p.H, p.W, total_c, dtype=_td).uniform_(-1.0, 1.0) + B_t = torch.empty(total_k, p.KH, p.KW, shape.cpg, dtype=_td).uniform_(-1.0, 1.0) + D_t = torch.empty(p.N, p.Ho, p.Wo, total_k, dtype=_td) ref = _conv_ref_grouped(A_t, B_t, p) @@ -224,7 +225,7 @@ def _run_grouped_one(arch: str, shape: _Shape) -> Tuple[bool, str]: rt.memcpy_h2d(B_dev, _u8(B_t), B_t.nbytes) rt.memset(D_dev, 0, D_t.nbytes) - sig = conv_args_signature("fp16") + sig = conv_args_signature(dtype) try: launcher = KernelLauncher( hsaco=artifact.hsaco, @@ -264,11 +265,12 @@ def _run_grouped_one(arch: str, shape: _Shape) -> Tuple[bool, str]: abs_diff = (out_f32 - ref_f32).abs() ref_scale = ref_f32.abs().max().clamp(min=1.0) rel_err = float(abs_diff.max() / ref_scale) - passed = rel_err < _TOL + tol = _TOL_BF16 if dtype == "bf16" else _TOL + passed = rel_err < tol if not passed: - return False, f"rel_err={rel_err:.3e} > tol={_TOL:.1e}" + return False, f"rel_err={rel_err:.3e} > tol={tol:.1e}" print( - f" PASS {shape.id} {arch} rel_err={rel_err:.2e}", + f" PASS {shape.id} {arch} {dtype} rel_err={rel_err:.2e}", flush=True, ) return True, "" @@ -621,7 +623,7 @@ def test_depthwise_spatial(self): ] -def _run_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: +def _run_dgrad_one(arch: str, shape: _Shape, dtype: str = "fp16") -> Tuple[bool, str]: """Build, compile, launch, and verify the direct dgrad kernel. Returns ``(passed, reason)``. @@ -653,6 +655,7 @@ def _run_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: KW=shape.KW, PAD=shape.PAD, stride=shape.stride, + dtype=dtype, ) spec_kwargs = {"problem": p, "name": f"test_dgrad_{shape.id}"} if shape.block_groups > 0: @@ -676,14 +679,13 @@ def _run_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: torch.manual_seed(42) total_c = shape.groups * shape.cpg total_k = shape.groups * kpg + _td = torch.bfloat16 if dtype == "bf16" else torch.float16 # dY: output gradient [N, Ho, Wo, K] - dY = torch.empty(p.N, p.Ho, p.Wo, total_k, dtype=torch.float16).uniform_(-0.5, 0.5) + dY = torch.empty(p.N, p.Ho, p.Wo, total_k, dtype=_td).uniform_(-0.5, 0.5) # W: weights [K, KH, KW, cpg] - W = torch.empty(total_k, p.KH, p.KW, shape.cpg, dtype=torch.float16).uniform_( - -0.5, 0.5 - ) - dX = torch.zeros(p.N, p.H, p.W, total_c, dtype=torch.float16) + W = torch.empty(total_k, p.KH, p.KW, shape.cpg, dtype=_td).uniform_(-0.5, 0.5) + dX = torch.zeros(p.N, p.H, p.W, total_c, dtype=_td) # Reference: dX = conv_transpose2d(dY, W) # output_padding recovers the exact input H, W (matters when stride > 1). @@ -709,7 +711,7 @@ def _run_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: rt.memcpy_h2d(W_dev, _u8(W), W.nbytes) rt.memset(dX_dev, 0, dX.nbytes) - sig = conv_args_signature("fp16") + sig = conv_args_signature(dtype) try: launcher = KernelLauncher( hsaco=artifact.hsaco, @@ -751,10 +753,11 @@ def _run_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: abs_diff = (out_f32 - ref_f32).abs() ref_scale = ref_f32.abs().max().clamp(min=1.0) rel_err = float(abs_diff.max() / ref_scale) - passed = rel_err < _TOL + tol = _TOL_BF16 if dtype == "bf16" else _TOL + passed = rel_err < tol if not passed: - return False, f"rel_err={rel_err:.3e} > tol={_TOL:.1e}" - print(f" PASS {shape.id} {arch} rel_err={rel_err:.2e}", flush=True) + return False, f"rel_err={rel_err:.3e} > tol={tol:.1e}" + print(f" PASS {shape.id} {arch} {dtype} rel_err={rel_err:.2e}", flush=True) return True, "" @@ -900,5 +903,70 @@ def test_dw_dgrad(self): self._run_dw_dgrad(s) +# --------------------------------------------------------------------------- +# bf16 correctness tests +# bf16 is supported by cpg=8, cpg=16, cpg=32 (not cpg=4 — no 4x4x4 bf16 atom) +# and by the scalar dgrad path. gfx950 is required for the 16x16x32 fold_k32 +# atom; 16x16x16 bf16 (non-fold path) works on both gfx942 and gfx950. +# --------------------------------------------------------------------------- + +# Subset of _SHAPES with cpg values that support bf16. +_BF16_FWD_SHAPES: List[_Shape] = [s for s in _SHAPES if s.cpg in (8, 16, 32)] + +# Dgrad shapes that support bf16 (scalar FMA dgrad handles all cpg/kpg). +_BF16_DGRAD_SHAPES: List[_Shape] = list(_DGRAD_SHAPES) + + +@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") +class TestDirectConvBf16Correctness(unittest.TestCase): + """Correctness tests for direct conv with bf16 I/O tensors. + + Uses the same harness as ``TestDirectConvCorrectness`` but with + ``dtype="bf16"`` and a looser tolerance (``_TOL_BF16``). cpg=4 is + excluded because there is no ``mfma_f32_4x4x4_bf16`` atom on CDNA. + """ + + def _run_fwd(self, shape: _Shape) -> None: + passed, reason = _run_grouped_one(GPU_ARCH, shape, dtype="bf16") + if reason.startswith("skip"): + self.skipTest(reason) + self.assertTrue( + passed, + f"FAIL bf16 fwd {shape.id} on {GPU_ARCH}: {reason}", + ) + + def test_bf16_cpg8(self): + for s in _BF16_FWD_SHAPES: + if s.cpg == 8: + with self.subTest(shape=s.id): + self._run_fwd(s) + + def test_bf16_cpg16(self): + for s in _BF16_FWD_SHAPES: + if s.cpg == 16: + with self.subTest(shape=s.id): + self._run_fwd(s) + + def test_bf16_cpg32(self): + for s in _BF16_FWD_SHAPES: + if s.cpg == 32: + with self.subTest(shape=s.id): + self._run_fwd(s) + + def _run_dgrad(self, shape: _Shape) -> None: + passed, reason = _run_dgrad_one(GPU_ARCH, shape, dtype="bf16") + if reason.startswith("skip"): + self.skipTest(reason) + self.assertTrue( + passed, + f"FAIL bf16 dgrad {shape.id} on {GPU_ARCH}: {reason}", + ) + + def test_bf16_dgrad(self): + for s in _BF16_DGRAD_SHAPES: + with self.subTest(shape=s.id): + self._run_dgrad(s) + + if __name__ == "__main__": unittest.main(verbosity=2) From a1c3bc8931fce7ea391baafba3b94f78d62e3d0f Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Wed, 23 Sep 2026 10:30:00 +0000 Subject: [PATCH 3/5] fixes --- .../kernels/common/conv_direct_grouped.py | 286 ++++---- .../tests/parity/conv_direct_grouped_emit.c | 33 + .../tests/parity/conv_direct_grouped_emit.py | 40 ++ .../tests/test_direct_conv_correctness.py | 18 +- .../library/tests/test_dtype_dispatch.py | 645 ++++++++++++++++++ .../rocke/instance_conv_direct_grouped.h | 10 +- .../instance_conv_direct_grouped_internal.h | 16 +- ...direct_grouped_build_16c_loop_closures.cpp | 14 +- .../conv_direct_grouped_build_16c_phases.cpp | 13 +- .../common/conv_direct_grouped_build_32c.cpp | 27 +- ...irect_grouped_build_4c_phases_and_loop.cpp | 6 +- .../common/conv_direct_grouped_build_8c.cpp | 37 +- .../conv_direct_grouped_spec_and_validate.cpp | 61 +- 13 files changed, 982 insertions(+), 224 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_dtype_dispatch.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py index 98d1d48af442..1e2d1955c03d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py @@ -81,6 +81,7 @@ PtrType, Value, ) +from rocke.helpers.transforms import TensorDescriptor, embed, unmerge_magic def _io_type(dtype: str): @@ -100,8 +101,8 @@ def _buf_load_vN( """Dtype-dispatch for vectorised buffer load. ``dwords`` matches the ``dwords`` parameter of ``buffer_load_vN_f16`` / - ``buffer_load_vN_bf16``: each dword holds two 16-bit elements (dwords=1 → - 2 elements, dwords=2 → 4 elements, dwords=4 → 8 elements). + ``buffer_load_vN_bf16``: each dword holds two 16-bit elements (dwords=1 -> + 2 elements, dwords=2 -> 4 elements, dwords=4 -> 8 elements). Uses the type-specific op names to stay byte-identical with the C++ engine. """ if dtype == "bf16": @@ -129,7 +130,31 @@ def _buf_store_vN( b.buffer_store_vN_f16(rsrc, voff, soff, val, dwords) -from rocke.helpers.transforms import TensorDescriptor, embed, unmerge_magic +def _trunc_f32(b: IRBuilder, dtype: str, val: Value) -> Value: + """Truncate a vector of f32 accumulators to the output dtype.""" + if dtype == "bf16": + return b.vec_trunc_f32_to_bf16(val) + return b.vec_trunc_f32_to_f16(val) + + +def _mfma( + b: IRBuilder, + dtype: str, + shape: str, + a: Value, + b_val: Value, + acc: Value, +) -> Value: + """Dtype-dispatch for a single MFMA call. + + ``shape`` is the size suffix without the dtype, e.g. ``"16x16x16"``, + ``"16x16x32"``, or ``"32x32x8"``. + """ + if dtype == "bf16": + fn = getattr(b, f"mfma_f32_{shape}_bf16") + else: + fn = getattr(b, f"mfma_f32_{shape}_f16") + return fn(a, b_val, acc) @dataclass(frozen=True) @@ -810,32 +835,29 @@ def lds_read_input_s2_k32(q_subtile: int, lds: Value) -> Value: # Migrating the C-accumulator readout + A/B K-pack to # c_layout().coord would delete this whole hazard class at # the source; tracked as a follow-up. - if p.dtype == "bf16": - acc_in = b.mfma_f32_16x16x32_bf16( - weights_k32[r_const], input_k32, acc_in - ) - acc_in = b.mfma_f32_16x16x32_bf16( - weights_s2_k32[r_const], input_s2, acc_in - ) - else: - acc_in = b.mfma_f32_16x16x32_f16( - weights_k32[r_const], input_k32, acc_in - ) - acc_in = b.mfma_f32_16x16x32_f16( - weights_s2_k32[r_const], input_s2, acc_in - ) + acc_in = _mfma( + b, p.dtype, "16x16x32", weights_k32[r_const], input_k32, acc_in + ) + acc_in = _mfma( + b, + p.dtype, + "16x16x32", + weights_s2_k32[r_const], + input_s2, + acc_in, + ) else: inputs = inputs_by_q[qt] for s_const in range(p.KW): w_idx = r_const * p.KW + s_const - if p.dtype == "bf16": - acc_in = b.mfma_f32_16x16x16_bf16( - weights[w_idx], inputs[s_const], acc_in - ) - else: - acc_in = b.mfma_f32_16x16x16_f16( - weights[w_idx], inputs[s_const], acc_in - ) + acc_in = _mfma( + b, + p.dtype, + "16x16x16", + weights[w_idx], + inputs[s_const], + acc_in, + ) accs[p_idx] = acc_in if loads_next is not None: @@ -899,13 +921,9 @@ def lds_read_input_s2_k32(q_subtile: int, lds: Value) -> Value: # k_out = g*kpg + c4*4 + [0..3]. Store them as one # 64-bit vector instead of four scalar buffer_store_short # ops. - if p.dtype == "bf16": - acc_h = b.vec_trunc_f32_to_bf16(acc_to_flush) - b.buffer_store_vN_bf16(d_rsrc, safe_d_off, c0, acc_h, 2) - else: - acc_h = b.vec_trunc_f32_to_f16(acc_to_flush) - b.buffer_store_vN_f16(d_rsrc, safe_d_off, c0, acc_h, 2) - # Unconditional slot reset — kills early-iter leaks before they + acc_h = _trunc_f32(b, p.dtype, acc_to_flush) + _buf_store_vN(b, p.dtype, d_rsrc, safe_d_off, c0, acc_h, 2) + # Unconditional slot reset - kills early-iter leaks before they # pollute a later output row. for qt in range(q_subtiles): acc_tiles[qt][P_FLUSH] = zero_acc @@ -944,7 +962,7 @@ def validate(self) -> None: p = self.problem if p.dtype != "fp16": raise ValueError( - f"DirectConv4cSpec: bf16 is not supported — the mfma_f32_4x4x4 atom " + f"DirectConv4cSpec: bf16 is not supported - the mfma_f32_4x4x4 atom " f"is fp16-only on CDNA; use fp16 dtype or a different cpg variant" ) if p.cpg != 4 or p.kpg != 4: @@ -983,7 +1001,7 @@ def is_valid_spec_4c(spec: DirectConv4cSpec, arch: str = "gfx950") -> Tuple[bool p = spec.problem if p.dtype != "fp16": return False, ( - f"DirectConv4cSpec: bf16 not supported — no mfma_f32_4x4x4_bf16 atom on CDNA" + f"DirectConv4cSpec: bf16 not supported - no mfma_f32_4x4x4_bf16 atom on CDNA" ) if p.stride != 1: return False, f"stride > 1 is not supported (got {p.stride})" @@ -1191,8 +1209,8 @@ def build_direct_conv_4c(spec: DirectConv4cSpec, *, arch: str = "gfx950") -> Ker safe_d = b.select(out_q_ok, b.mul(d_base, c_half_bytes), oob_sentinel) # MFMA 4x4x4 wave64 per-lane output layout: # acc[i] -> D[n, ho_row, out_q, g*kpg + i] for i in 0..3 - acc_h = b.vec_trunc_f32_to_f16(acc) - b.buffer_store_vN_f16(d_rsrc, safe_d, c0, acc_h, 2) + acc_h = _trunc_f32(b, p.dtype, acc) + _buf_store_vN(b, p.dtype, d_rsrc, safe_d, c0, acc_h, 2) for qt in range(q_tiles_per_wave): acc_tiles[qt][P_FLUSH] = zero_acc @@ -1623,22 +1641,13 @@ def lds_read_input_s2(q_subtile: int, lds) -> Value: p_idx = (y - r_const) % p.KH acc_in = accs[p_idx] # Main atom: s=0 and s=1 folded into K=16. - if p.dtype == "bf16": - acc_in = b.mfma_f32_16x16x16_bf16( - weights_main[r_const], inp_main, acc_in - ) - # Residual atom: s=2, zero-padded to K=16 (upper lanes carry zeros). - acc_in = b.mfma_f32_16x16x16_bf16( - weights_s2[r_const], inp_s2, acc_in - ) - else: - acc_in = b.mfma_f32_16x16x16_f16( - weights_main[r_const], inp_main, acc_in - ) - # Residual atom: s=2, zero-padded to K=16 (upper lanes carry zeros). - acc_in = b.mfma_f32_16x16x16_f16( - weights_s2[r_const], inp_s2, acc_in - ) + acc_in = _mfma( + b, p.dtype, "16x16x16", weights_main[r_const], inp_main, acc_in + ) + # Residual atom: s=2, zero-padded to K=16 (upper lanes carry zeros). + acc_in = _mfma( + b, p.dtype, "16x16x16", weights_s2[r_const], inp_s2, acc_in + ) accs[p_idx] = acc_in if loads_next is not None: @@ -1671,12 +1680,8 @@ def lds_read_input_s2(q_subtile: int, lds) -> Value: d_base_bytes = b.mul(d_base, c_half_bytes) store_valid = b.land(out_q_valid, c4_valid) safe_d_off = b.select(store_valid, d_base_bytes, oob_sentinel) - if p.dtype == "bf16": - acc_h = b.vec_trunc_f32_to_bf16(acc_to_flush) - b.buffer_store_vN_bf16(d_rsrc, safe_d_off, c0, acc_h, 2) - else: - acc_h = b.vec_trunc_f32_to_f16(acc_to_flush) - b.buffer_store_vN_f16(d_rsrc, safe_d_off, c0, acc_h, 2) + acc_h = _trunc_f32(b, p.dtype, acc_to_flush) + _buf_store_vN(b, p.dtype, d_rsrc, safe_d_off, c0, acc_h, 2) for qt in range(q_subtiles): acc_tiles[qt][P_FLUSH] = zero_acc @@ -2065,18 +2070,14 @@ def lds_read_input(q_subtile: int, s_const: int, atom_idx: int, lds) -> Value: acc_in = accs[p_idx] for s_const in range(p.KW): for atom_idx in range(4): - if p.dtype == "bf16": - acc_in = b.mfma_f32_32x32x8_bf16( - weights[r_const][s_const][atom_idx], - inputs_by_q[qt][s_const][atom_idx], - acc_in, - ) - else: - acc_in = b.mfma_f32_32x32x8_f16( - weights[r_const][s_const][atom_idx], - inputs_by_q[qt][s_const][atom_idx], - acc_in, - ) + acc_in = _mfma( + b, + p.dtype, + "32x32x8", + weights[r_const][s_const][atom_idx], + inputs_by_q[qt][s_const][atom_idx], + acc_in, + ) accs[p_idx] = acc_in if loads_next is not None: @@ -2116,12 +2117,8 @@ def lds_read_input(q_subtile: int, s_const: int, atom_idx: int, lds) -> Value: e0 = b.vec_extract(acc_to_flush, slot0) e1 = b.vec_extract(acc_to_flush, slot1) pair_acc = b.vec_pack([e0, e1], F32) - if p.dtype == "bf16": - pair_h = b.vec_trunc_f32_to_bf16(pair_acc) - b.buffer_store_vN_bf16(d_rsrc, safe_d_off, c0, pair_h, 1) - else: - pair_h = b.vec_trunc_f32_to_f16(pair_acc) - b.buffer_store_vN_f16(d_rsrc, safe_d_off, c0, pair_h, 1) + pair_h = _trunc_f32(b, p.dtype, pair_acc) + _buf_store_vN(b, p.dtype, d_rsrc, safe_d_off, c0, pair_h, 1) for qt in range(q_subtiles): acc_tiles[qt][P_FLUSH] = zero_acc @@ -2810,40 +2807,15 @@ def store_to_lds(loads, lds) -> None: w_frag_pw = preloaded_w[ (r_const, s_const, local_atom, m) ] - if FOLD_K32: - if p.dtype == "bf16": - acc_tiles[qt][m][p_idx] = ( - b.mfma_f32_16x16x32_bf16( - w_frag_pw, - x_frag_pw, - acc_tiles[qt][m][p_idx], - ) - ) - else: - acc_tiles[qt][m][p_idx] = ( - b.mfma_f32_16x16x32_f16( - w_frag_pw, - x_frag_pw, - acc_tiles[qt][m][p_idx], - ) - ) - else: - if p.dtype == "bf16": - acc_tiles[qt][m][p_idx] = ( - b.mfma_f32_16x16x16_bf16( - w_frag_pw, - x_frag_pw, - acc_tiles[qt][m][p_idx], - ) - ) - else: - acc_tiles[qt][m][p_idx] = ( - b.mfma_f32_16x16x16_f16( - w_frag_pw, - x_frag_pw, - acc_tiles[qt][m][p_idx], - ) - ) + mfma_shape = "16x16x32" if FOLD_K32 else "16x16x16" + acc_tiles[qt][m][p_idx] = _mfma( + b, + p.dtype, + mfma_shape, + w_frag_pw, + x_frag_pw, + acc_tiles[qt][m][p_idx], + ) continue # skip the runtime s_loop/atom_loop below # ---- Runtime K-atom loop path (runtime_k_loop=True) ---- @@ -2944,24 +2916,15 @@ def store_to_lds(loads, lds) -> None: c0, LOAD_VEC // 2, ) - if FOLD_K32: - if p.dtype == "bf16": - new_k_accs_rk[m] = b.mfma_f32_16x16x32_bf16( - w_frag_rk, x_frag_rk, new_k_accs_rk[m] - ) - else: - new_k_accs_rk[m] = b.mfma_f32_16x16x32_f16( - w_frag_rk, x_frag_rk, new_k_accs_rk[m] - ) - else: - if p.dtype == "bf16": - new_k_accs_rk[m] = b.mfma_f32_16x16x16_bf16( - w_frag_rk, x_frag_rk, new_k_accs_rk[m] - ) - else: - new_k_accs_rk[m] = b.mfma_f32_16x16x16_f16( - w_frag_rk, x_frag_rk, new_k_accs_rk[m] - ) + mfma_shape = "16x16x32" if FOLD_K32 else "16x16x16" + new_k_accs_rk[m] = _mfma( + b, + p.dtype, + mfma_shape, + w_frag_rk, + x_frag_rk, + new_k_accs_rk[m], + ) b.scf_yield(*new_k_accs_rk) for m in range(N_M_TILES): acc_tiles[qt][m][p_idx] = k_loop_rk.results[m] @@ -3068,24 +3031,10 @@ def store_to_lds(loads, lds) -> None: if p.cpg % K_ATOM_SIZE != 0: w_frag = b.select(c4_oob, _zero_std, w_frag) - if FOLD_K32: - if p.dtype == "bf16": - new_acc = b.mfma_f32_16x16x32_bf16( - w_frag, x_frag, atom_accs[m] - ) - else: - new_acc = b.mfma_f32_16x16x32_f16( - w_frag, x_frag, atom_accs[m] - ) - else: - if p.dtype == "bf16": - new_acc = b.mfma_f32_16x16x16_bf16( - w_frag, x_frag, atom_accs[m] - ) - else: - new_acc = b.mfma_f32_16x16x16_f16( - w_frag, x_frag, atom_accs[m] - ) + mfma_shape = "16x16x32" if FOLD_K32 else "16x16x16" + new_acc = _mfma( + b, p.dtype, mfma_shape, w_frag, x_frag, atom_accs[m] + ) new_atom_accs.append(new_acc) b.scf_yield(*new_atom_accs) @@ -3203,23 +3152,15 @@ def store_to_lds(loads, lds) -> None: safe_d = b.select( store_ok, b.mul(d_base, c_half_bytes), oob_sentinel ) - if p.dtype == "bf16": - acc_h = b.vec_trunc_f32_to_bf16(partial) - b.buffer_store_vN_bf16(d_rsrc, safe_d, c0, acc_h, 2) - else: - acc_h = b.vec_trunc_f32_to_f16(partial) - b.buffer_store_vN_f16(d_rsrc, safe_d, c0, acc_h, 2) + acc_h = _trunc_f32(b, p.dtype, partial) + _buf_store_vN(b, p.dtype, d_rsrc, safe_d, c0, acc_h, 2) b.sync() # allow red_lds reuse by next flush else: safe_d = b.select( store_ok, b.mul(d_base, c_half_bytes), oob_sentinel ) - if p.dtype == "bf16": - acc_h = b.vec_trunc_f32_to_bf16(acc_to_flush) - b.buffer_store_vN_bf16(d_rsrc, safe_d, c0, acc_h, 2) - else: - acc_h = b.vec_trunc_f32_to_f16(acc_to_flush) - b.buffer_store_vN_f16(d_rsrc, safe_d, c0, acc_h, 2) + acc_h = _trunc_f32(b, p.dtype, acc_to_flush) + _buf_store_vN(b, p.dtype, d_rsrc, safe_d, c0, acc_h, 2) for qt in range(q_subtiles): for m in range(N_M_TILES): @@ -3511,10 +3452,7 @@ def build_direct_reorganize_weights( b.mul(tid, b.const_i32(ELEMS_PER_LANE)), ) safe_dst = b.select(src_ok, b.mul(dst_off, c_half_bytes), oob_sentinel) - if p.dtype == "bf16": - b.buffer_store_vN_bf16(d_rsrc, safe_dst, c0, val, ELEMS_PER_LANE // 2) - else: - b.buffer_store_vN_f16(d_rsrc, safe_dst, c0, val, ELEMS_PER_LANE // 2) + _buf_store_vN(b, p.dtype, d_rsrc, safe_dst, c0, val, ELEMS_PER_LANE // 2) return b.kernel @@ -3622,10 +3560,7 @@ def build_direct_coalesced_weights_dgrad( b.mul(tid, b.const_i32(4)), ) safe_dst = b.select(k_ok, b.mul(dst_off, c_half_bytes), oob_sentinel) - if p.dtype == "bf16": - b.buffer_store_vN_bf16(d_rsrc, safe_dst, c0, val, 2) - else: - b.buffer_store_vN_f16(d_rsrc, safe_dst, c0, val, 2) + _buf_store_vN(b, p.dtype, d_rsrc, safe_dst, c0, val, 2) return b.kernel @@ -3669,6 +3604,7 @@ def build_direct_mfma_dgrad( KW=orig_p.KW, PAD=orig_p.KH - 1 - orig_p.PAD, # undo the PAD swap stride=1, + dtype=orig_p.dtype, ) transpose_kernel = build_direct_transpose_weights_dgrad( DirectTransposeWeightsDgradSpec(problem=orig_problem), arch=arch @@ -3721,6 +3657,7 @@ def make_dgrad_fprop_spec( KW=p.KW, PAD=pad_new, stride=1, + dtype=p.dtype, ) return DirectConvSpec( problem=transposed_problem, @@ -4117,6 +4054,10 @@ def kernel_name(self) -> str: def validate(self) -> None: p = self.problem + if p.dtype != "fp16": + raise ValueError( + f"DirectDepthwiseDgradSpec: bf16 not supported; depthwise kernels are fp16-only" + ) if p.cpg != 1 or p.kpg != 1: raise ValueError( f"DirectDepthwiseDgradSpec requires cpg=kpg=1 (got {p.cpg}, {p.kpg})" @@ -4368,6 +4309,10 @@ def kernel_name(self) -> str: def validate(self) -> None: p = self.problem + if p.dtype != "fp16": + raise ValueError( + f"DirectDepthwiseDgradStreamSpec: bf16 not supported; depthwise kernels are fp16-only" + ) if p.cpg != 1 or p.kpg != 1: raise ValueError( f"DirectDepthwiseDgradStreamSpec requires cpg=kpg=1 (got {p.cpg}, {p.kpg})" @@ -4639,11 +4584,14 @@ def kernel_name(self) -> str: def validate(self) -> None: p = self.problem + if p.dtype != "fp16": + raise ValueError( + f"DirectDepthwiseSpec: bf16 not supported; depthwise kernels are fp16-only" + ) if p.cpg != 1 or p.kpg != 1: raise ValueError( f"DirectDepthwiseSpec requires cpg=kpg=1 (got cpg={p.cpg}, kpg={p.kpg})" ) - pass def is_valid_depthwise_spec( @@ -5001,6 +4949,10 @@ def kernel_name(self) -> str: def validate(self) -> None: p = self.problem + if p.dtype != "fp16": + raise ValueError( + f"DirectDepthwiseSpatialSpec: bf16 not supported; depthwise kernels are fp16-only" + ) if p.cpg != 1 or p.kpg != 1: raise ValueError( f"DirectDepthwiseSpatialSpec requires cpg=kpg=1 " diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c index 6aba98f09073..8bcd48da3d0a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c @@ -305,6 +305,39 @@ static int make_cfg(int idx, *kind = KIND_DW_DGRAD; *arch = "gfx950"; return 0; + case 17: + /* 16c bf16: exercises bf16 I/O, bf16 load/store taps, mfma_f32_16x16x16_bf16 */ + p.N = 32; + p.H = 200; + p.W = 200; + p.groups = 16; + p.cpg = 16; + p.kpg = 16; + p.dtype = "bf16"; + *s16 = rocke_direct_conv_16c_spec_default(); + s16->problem = p; + s16->block_groups = 8; + s16->fold_k32 = false; + *kind = KIND_16C; + *arch = "gfx950"; + return 0; + case 18: + /* 8c bf16: exercises bf16 I/O, bf16 load/store taps, mfma_f32_16x16x16_bf16 */ + p.N = 32; + p.H = 200; + p.W = 200; + p.groups = 16; + p.cpg = 8; + p.kpg = 8; + p.dtype = "bf16"; + *s8 = rocke_direct_conv_8c_spec_default(); + s8->problem = p; + s8->block_q = 16; + s8->block_groups = 8; + s8->double_buffer = true; + *kind = KIND_8C; + *arch = "gfx950"; + return 0; default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py index c448409b183b..11f6e66351ee 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py @@ -195,6 +195,46 @@ def _spec(idx: int): DirectDepthwiseDgradSpec(problem=p, block_w=8, block_waves=1), "gfx950", ) + if idx == 17: + # 16c bf16: exercises bf16 I/O, bf16 load/store taps, mfma_f32_16x16x16_bf16 + p = DirectConvProblem( + N=32, + H=200, + W=200, + groups=16, + cpg=16, + kpg=16, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + ) + return ( + "16c", + DirectConv16cSpec(problem=p, block_groups=8, fold_k32=False), + "gfx950", + ) + if idx == 18: + # 8c bf16: exercises bf16 I/O, bf16 load/store taps, mfma_f32_16x16x16_bf16 + p = DirectConvProblem( + N=32, + H=200, + W=200, + groups=16, + cpg=8, + kpg=8, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + ) + return ( + "8c", + DirectConv8cSpec(problem=p, block_q=16, block_groups=8, double_buffer=True), + "gfx950", + ) raise SystemExit(f"unknown config index {idx}") diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py index 829a4e30cb5c..5616edb34280 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py @@ -57,7 +57,7 @@ def _skip_reason() -> str: _SKIP_REASON = _skip_reason() _TOL = 5e-2 -_TOL_BF16 = 1e-1 # bf16 has ~3× less mantissa precision than f16 +_TOL_BF16 = 1e-1 # bf16 has 3 fewer mantissa bits than fp16 (~8x coarser precision) # --------------------------------------------------------------------------- @@ -968,5 +968,21 @@ def test_bf16_dgrad(self): self._run_dgrad(s) +class TestDirectConvValidation(unittest.TestCase): + """Validation-only tests that do not require a GPU.""" + + def test_cpg4_bf16_rejected(self): + """cpg=4 + bf16 must raise ValueError (no mfma_f32_4x4x4_bf16 on CDNA).""" + from kernels.common.conv_direct_grouped import ( + DirectConv4cSpec, + DirectConvProblem, + ) + + p = DirectConvProblem(N=1, H=8, W=8, groups=16, cpg=4, kpg=4, dtype="bf16") + spec = DirectConv4cSpec(problem=p) + with self.assertRaises(ValueError): + spec.validate() + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_dtype_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_dtype_dispatch.py new file mode 100644 index 000000000000..c08f007a883a --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_dtype_dispatch.py @@ -0,0 +1,645 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Dtype-dispatch sanity tests for direct-conv and implicit-gemm families. + +Two test tiers: + + Level 1 — IR smoke (no GPU required). + Lowers each (family × direction × dtype) combination to LLVM-IR via the + Python engine and asserts that the parameter pointer element type matches + the requested dtype. Specifically: + - bf16 specs → the function signature must contain ``bfloat addrspace`` + - fp16 specs → must NOT contain ``bfloat addrspace`` + + This would have caught the pre-fix bug where ``make_dgrad_fprop_spec`` + silently dropped ``dtype`` from ``transposed_problem``, so a bf16 dgrad + spec emitted fp16 IR for the underlying fprop kernel. + + Level 2 — Execution round-trip (GPU + torch required). + Runs a minimal kernel (N=1, tiny spatial) with inputs crafted to + distinguish fp16 from bf16 at the bit level: + + 1.0 as bf16 → bits 0x3F80 (interpreted as fp16 → value ≈ 1.5) + 1.0 as fp16 → bits 0x3C00 (interpreted as bf16 → value ≈ 1e-4) + + A kernel that internally reinterprets bf16 bits as fp16 (or vice-versa) + computes a convolution with the *wrong* inputs and the result diverges + from the float32 reference by far more than the dtype tolerance. + +Run without GPU (IR only): + PYTHONPATH=rocke/platform/python:rocke/library pytest \\ + rocke/library/tests/test_dtype_dispatch.py -k "ir" + +Run with GPU (all tests): + PYTHONPATH=rocke/platform/python:rocke/library -m pytest \\ + rocke/library/tests/test_dtype_dispatch.py +""" +from __future__ import annotations + +import ctypes +import importlib.util +import unittest +from dataclasses import dataclass, field +from typing import Tuple + +from rocke.runtime.hip_module import get_device_arch + +_HAS_TORCH = importlib.util.find_spec("torch") is not None + +if _HAS_TORCH: + import torch + + torch.cuda.is_available() + +GPU_ARCH = get_device_arch(0) +_IS_MFMA = GPU_ARCH in ("gfx942", "gfx950") + + +def _skip_reason_gpu() -> str: + if not GPU_ARCH: + return "no ROCm GPU detected" + if not _HAS_TORCH: + return "torch not importable" + if not _IS_MFMA: + return f"unsupported arch {GPU_ARCH!r} (need gfx942 or gfx950)" + return "" + + +_SKIP_GPU = _skip_reason_gpu() + +_ARCH = GPU_ARCH if GPU_ARCH in ("gfx942", "gfx950") else "gfx950" + +# Tolerances (relative error against float32 reference). +_TOL_FP16 = 5e-2 +_TOL_BF16 = 1e-1 + + +# --------------------------------------------------------------------------- +# Shared lower helper +# --------------------------------------------------------------------------- + + +def _lower(kernel, arch: str) -> str: + try: + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python as _f + except ImportError: + from rocke import lower_kernel_to_llvm as _f + return _f(kernel, arch=arch) + + +# --------------------------------------------------------------------------- +# Level 1 helpers +# --------------------------------------------------------------------------- + + +def _assert_ir_dtype(ir: str, dtype: str, kernel_label: str) -> None: + """Assert that the LLVM-IR load/store element type matches *dtype*. + + Modern LLVM uses opaque pointers (``ptr addrspace(N)``), so the element + type appears in the load/store instruction, not in the function signature. + We check for ``load `` / ``load `` in the IR body. + """ + if dtype == "bf16": + # The IR must reference 'bfloat' for load/store/MFMA element type. + # Scalar dgrad uses individual bfloat loads; vector kernels use + # loads. Either form proves dtype reached the builder. + assert "bfloat" in ir, ( + f"{kernel_label}: expected 'bfloat' in IR for bf16 spec " + f"but not found — dtype was likely not forwarded to the kernel builder" + ) + else: + assert "bfloat" not in ir, f"{kernel_label}: found 'bfloat' in IR for fp16 spec" + + +# =========================================================================== +# Level 1 — IR dtype checks (no GPU) +# =========================================================================== + + +class TestDirectConvIRDtype(unittest.TestCase): + """IR-level dtype checks for the direct-conv family (no GPU required). + + Builds each (variant × direction × dtype) combination and verifies that + the emitted LLVM-IR parameter pointers use the correct element type. + """ + + # ---- forward ----------------------------------------------------------- + + def _check_fwd_ir(self, cpg: int, dtype: str) -> None: + from kernels.common.conv_direct_grouped import ( + DirectConvProblem, + DirectConvSpec, + build_direct_conv, + ) + + # groups=8 satisfies default block_groups=8 for 8c/16c/32c. + groups = 8 + p = DirectConvProblem( + N=1, + H=8, + W=8, + groups=groups, + cpg=cpg, + kpg=cpg, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype=dtype, + ) + spec = DirectConvSpec(problem=p, name=f"ir_fwd_{cpg}c_{dtype}") + kernel = build_direct_conv(spec, arch=_ARCH) + ir = _lower(kernel, arch=_ARCH) + _assert_ir_dtype(ir, dtype, f"direct_fwd cpg={cpg} {dtype}") + + def test_fwd_8c_fp16(self): + self._check_fwd_ir(8, "fp16") + + def test_fwd_8c_bf16(self): + self._check_fwd_ir(8, "bf16") + + def test_fwd_16c_fp16(self): + self._check_fwd_ir(16, "fp16") + + def test_fwd_16c_bf16(self): + self._check_fwd_ir(16, "bf16") + + def test_fwd_32c_fp16(self): + self._check_fwd_ir(32, "fp16") + + def test_fwd_32c_bf16(self): + self._check_fwd_ir(32, "bf16") + + # ---- dgrad (scalar-FMA path, cpg=kpg=1 → 16c-as-dgrad) --------------- + # The scalar-FMA dgrad (build_direct_conv_dgrad) uses _io_type(p.dtype) + # directly. The MFMA dgrad path (cpg>1 → make_dgrad_fprop_spec) is the + # one that previously dropped dtype; it maps to DirectConvDgradSpec. + + def _check_dgrad_scalar_ir(self, dtype: str) -> None: + from kernels.common.conv_direct_grouped import ( + DirectConvDgradSpec, + DirectConvProblem, + build_direct_conv_dgrad, + ) + + # cpg=kpg=1 → scalar FMA dgrad (single kernel, no weight-transpose) + p = DirectConvProblem( + N=1, + H=8, + W=8, + groups=8, + cpg=1, + kpg=1, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype=dtype, + ) + spec = DirectConvDgradSpec( + problem=p, name=f"ir_dgrad_dw_{dtype}", block_q=8, block_groups=8 + ) + kernel = build_direct_conv_dgrad(spec, arch=_ARCH) + ir = _lower(kernel, arch=_ARCH) + _assert_ir_dtype(ir, dtype, f"direct_dgrad dw {dtype}") + + def test_dgrad_scalar_fp16(self): + self._check_dgrad_scalar_ir("fp16") + + def test_dgrad_scalar_bf16(self): + self._check_dgrad_scalar_ir("bf16") + + # ---- dgrad MFMA path (cpg>1 → transposed fprop) ----------------------- + # THIS test would have caught the pre-fix bug: make_dgrad_fprop_spec + # dropped dtype, so the transposed fprop kernel always emitted fp16 IR. + + def _check_dgrad_mfma_ir(self, cpg: int, dtype: str) -> None: + from kernels.common.conv_direct_grouped import ( + DirectConvProblem, + make_dgrad_fprop_spec, + build_direct_conv, + ) + + p = DirectConvProblem( + N=1, + H=8, + W=8, + groups=8, + cpg=cpg, + kpg=cpg, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype=dtype, + ) + spec = make_dgrad_fprop_spec(p, block_q=16, block_groups=8 if cpg <= 16 else 4) + kernel = build_direct_conv(spec, arch=_ARCH) + ir = _lower(kernel, arch=_ARCH) + _assert_ir_dtype(ir, dtype, f"direct_dgrad mfma cpg={cpg} {dtype}") + + def test_dgrad_mfma_16c_fp16(self): + self._check_dgrad_mfma_ir(16, "fp16") + + def test_dgrad_mfma_16c_bf16(self): + # Regression test: before the fix this emitted 'half addrspace' even + # though dtype="bf16" because make_dgrad_fprop_spec dropped dtype. + self._check_dgrad_mfma_ir(16, "bf16") + + def test_dgrad_mfma_32c_fp16(self): + self._check_dgrad_mfma_ir(32, "fp16") + + def test_dgrad_mfma_32c_bf16(self): + self._check_dgrad_mfma_ir(32, "bf16") + + +class TestImplicitGemmIRDtype(unittest.TestCase): + """IR-level dtype checks for the implicit-gemm family (no GPU required).""" + + def _make_problem(self): + from kernels.common._conv_implicit_gemm_common import ConvProblem + + # Tiny shape that satisfies tile alignment constraints. + # C=K=64 so default_vector_sizes returns vec_c=8, but default epilogue + # does not support vec_c > 1, so we override vector_size_c=1 below. + return ConvProblem(N=1, Hi=8, Wi=8, C=64, K=64, Y=3, X=3, pH=1, pW=1) + + def _check_fwd_ir(self, dtype: str) -> None: + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + from kernels.common.conv_implicit_gemm import ( + ImplicitGemmConvSpec, + build_implicit_gemm_conv, + ) + + p = self._make_problem() + data = ConvDataSpec(dtype_a=dtype, dtype_b=dtype, dtype_d=dtype) + spec = ImplicitGemmConvSpec( + problem=p, + name=f"ir_igemm_fwd_{dtype}", + data=data, + tile_m=64, + tile_n=64, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=32, + warp_tile_n=32, + warp_tile_k=16, + # vec_c > 1 is incompatible with epilogue="default"; force 1. + vector_size_c=1, + ) + kernel = build_implicit_gemm_conv(spec, arch=_ARCH) + ir = _lower(kernel, arch=_ARCH) + _assert_ir_dtype(ir, dtype, f"igemm_fwd {dtype}") + + def test_fwd_fp16(self): + self._check_fwd_ir("fp16") + + def test_fwd_bf16(self): + self._check_fwd_ir("bf16") + + def _check_wgrad_ir(self, dtype: str) -> None: + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + from kernels.common.conv_implicit_gemm_wgrad import ( + WgradConvSpec, + build_implicit_gemm_conv_wgrad, + ) + + p = self._make_problem() + data = ConvDataSpec(dtype_a=dtype, dtype_b=dtype, dtype_d=dtype) + spec = WgradConvSpec( + problem=p, + name=f"ir_igemm_wgrad_{dtype}", + data=data, + tile_m=64, + tile_n=64, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=32, + warp_tile_n=32, + warp_tile_k=16, + ) + kernel = build_implicit_gemm_conv_wgrad(spec, arch=_ARCH) + ir = _lower(kernel, arch=_ARCH) + _assert_ir_dtype(ir, dtype, f"igemm_wgrad {dtype}") + + def test_wgrad_fp16(self): + self._check_wgrad_ir("fp16") + + def test_wgrad_bf16(self): + self._check_wgrad_ir("bf16") + + def _check_dgrad_ir(self, dtype: str) -> None: + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + from kernels.common.conv_implicit_gemm_dgrad import ( + DgradConvSpec, + build_implicit_gemm_conv_dgrad, + ) + + p = self._make_problem() + data = ConvDataSpec(dtype_a=dtype, dtype_b=dtype, dtype_d=dtype) + spec = DgradConvSpec( + problem=p, + name=f"ir_igemm_dgrad_{dtype}", + data=data, + tile_m=64, + tile_n=64, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=32, + warp_tile_n=32, + warp_tile_k=16, + ) + kernel = build_implicit_gemm_conv_dgrad(spec, arch=_ARCH) + ir = _lower(kernel, arch=_ARCH) + _assert_ir_dtype(ir, dtype, f"igemm_dgrad {dtype}") + + def test_dgrad_fp16(self): + self._check_dgrad_ir("fp16") + + def test_dgrad_bf16(self): + self._check_dgrad_ir("bf16") + + +# =========================================================================== +# Level 2 — GPU execution with dtype-distinguishing inputs +# =========================================================================== +# +# Key insight: 1.0 represented as bf16 has bits 0x3F80. +# Those same bits as fp16 decode to ≈ 1.5 (sign=0, exp=15→0, frac=0.5). +# So a kernel that silently reinterprets bf16 data as fp16 computes a +# convolution with inputs ≈ 1.5 instead of 1.0, producing wrong results. +# The relative error against the correct float32 reference is >> tolerance. + + +def _u8(t): + import torch # noqa: F401 + + return (ctypes.c_uint8 * t.nbytes).from_address(t.data_ptr()) + + +def _conv_ref_f32(A_t, B_t, p_stride: int, p_pad: int, groups: int): + """Return NHWK float32 reference computed by torch.nn.functional.conv2d.""" + import torch + import torch.nn.functional as F + + A_f = A_t.float() + B_f = B_t.float() + A_nchw = A_f.permute(0, 3, 1, 2) + B_nchw = B_f.permute(0, 3, 1, 2) + out = F.conv2d(A_nchw, B_nchw, padding=p_pad, stride=p_stride, groups=groups) + return out.permute(0, 2, 3, 1).contiguous().cuda() + + +def _run_direct_fwd(arch: str, cpg: int, dtype: str) -> Tuple[bool, str]: + import torch + from rocke import compile_kernel + from rocke.helpers.manifest import conv_args_signature + from kernels.common.conv_direct_grouped import ( + DirectConvProblem, + DirectConvSpec, + build_direct_conv, + is_valid_spec, + ) + from rocke.runtime import synchronize_and_release + from rocke.runtime.hip_module import HipError, Runtime + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + + groups = 8 + N, H, W = 1, 8, 8 + p = DirectConvProblem( + N=N, + H=H, + W=W, + groups=groups, + cpg=cpg, + kpg=cpg, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype=dtype, + ) + spec = DirectConvSpec( + problem=p, name=f"exec_fwd_{cpg}c_{dtype}", block_groups=groups + ) + ok, reason = is_valid_spec(spec, arch=arch) + if not ok: + return False, f"skip {reason}" + + kernel = build_direct_conv(spec, arch=arch) + artifact = compile_kernel(kernel, arch=arch) + + td = torch.bfloat16 if dtype == "bf16" else torch.float16 + torch.manual_seed(0) + A = torch.ones(N, H, W, groups * cpg, dtype=td).cuda() + B = torch.ones(groups * cpg, 3, 3, cpg, dtype=td).cuda() + D = torch.zeros(N, p.Ho, p.Wo, groups * cpg, dtype=td).cuda() + + ref = _conv_ref_f32(A.cpu(), B.cpu(), p.stride, p.PAD, groups) + + rt = Runtime() + A_dev = rt.alloc(A.nbytes) + rt.memcpy_h2d(A_dev, _u8(A), A.nbytes) + B_dev = rt.alloc(B.nbytes) + rt.memcpy_h2d(B_dev, _u8(B), B.nbytes) + D_dev = rt.alloc(D.nbytes) + rt.memset(D_dev, 0, D.nbytes) + + sig = conv_args_signature(dtype) + try: + launcher = KernelLauncher( + hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, signature=sig + ) + except HipError as e: + for dev in (A_dev, B_dev, D_dev): + rt.free(dev) + return False, f"load failed: {e}" + + q_tiles = (p.Wo + spec.block_q - 1) // spec.block_q + g_tiles = p.groups // spec.block_groups + launcher( + { + "A": A_dev, + "B": B_dev, + "D": D_dev, + "A_bytes": A.nbytes, + "B_bytes": B.nbytes, + "D_bytes": D.nbytes, + }, + config=LaunchConfig( + grid=(q_tiles, g_tiles, N), block=(spec.threads_per_block, 1, 1), fence=True + ), + ) + + D_cpu = torch.empty_like(D.cpu()) + rt.memcpy_d2h(_u8(D_cpu), D_dev, D.nbytes) + for dev in (A_dev, B_dev, D_dev): + rt.free(dev) + synchronize_and_release(0) + + rel = float( + (D_cpu.float() - ref.float().cpu()).abs().max() + / ref.float().cpu().abs().max().clamp(min=1.0) + ) + tol = _TOL_BF16 if dtype == "bf16" else _TOL_FP16 + return (rel < tol), f"rel_err={rel:.3e} tol={tol:.1e}" + + +def _run_direct_dgrad(arch: str, cpg: int, dtype: str) -> Tuple[bool, str]: + import torch + from rocke import compile_kernel + from rocke.helpers.manifest import conv_args_signature + from kernels.common.conv_direct_grouped import ( + DirectConvDgradSpec, + DirectConvProblem, + build_direct_conv_dgrad, + is_valid_dgrad_spec, + ) + from rocke.runtime import synchronize_and_release + from rocke.runtime.hip_module import HipError, Runtime + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + + groups = 8 + N, H, W = 1, 8, 8 + kpg = cpg + p = DirectConvProblem( + N=N, + H=H, + W=W, + groups=groups, + cpg=cpg, + kpg=kpg, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype=dtype, + ) + spec = DirectConvDgradSpec( + problem=p, name=f"exec_dgrad_{cpg}c_{dtype}", block_q=16, block_groups=groups + ) + ok, reason = is_valid_dgrad_spec(spec, arch=arch) + if not ok: + return False, f"skip {reason}" + + kernel = build_direct_conv_dgrad(spec, arch=arch) + artifact = compile_kernel(kernel, arch=arch) + + td = torch.bfloat16 if dtype == "bf16" else torch.float16 + torch.manual_seed(42) + total_c = groups * cpg + total_k = groups * kpg + dY = torch.ones(N, p.Ho, p.Wo, total_k, dtype=td).cuda() + Wt = torch.ones(total_k, 3, 3, cpg, dtype=td).cuda() + dX = torch.zeros(N, H, W, total_c, dtype=td).cuda() + + dY_nchw = dY.float().permute(0, 3, 1, 2) + W_nchw = Wt.float().permute(0, 3, 1, 2) + h_base = (p.Ho - 1) * p.stride - 2 * p.PAD + p.KH + w_base = (p.Wo - 1) * p.stride - 2 * p.PAD + p.KW + ref_nchw = torch.nn.functional.conv_transpose2d( + dY_nchw, + W_nchw, + padding=p.PAD, + stride=p.stride, + groups=groups, + output_padding=(H - h_base, W - w_base), + ) + ref = ref_nchw.permute(0, 2, 3, 1).contiguous() + + rt = Runtime() + dY_dev = rt.alloc(dY.nbytes) + rt.memcpy_h2d(dY_dev, _u8(dY), dY.nbytes) + W_dev = rt.alloc(Wt.nbytes) + rt.memcpy_h2d(W_dev, _u8(Wt), Wt.nbytes) + dX_dev = rt.alloc(dX.nbytes) + rt.memset(dX_dev, 0, dX.nbytes) + + sig = conv_args_signature(dtype) + try: + launcher = KernelLauncher( + hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, signature=sig + ) + except HipError as e: + for dev in (dY_dev, W_dev, dX_dev): + rt.free(dev) + return False, f"load failed: {e}" + + q_tiles = (W + spec.block_q - 1) // spec.block_q + g_tiles = (total_c + spec.threads_per_block - 1) // spec.threads_per_block + launcher( + { + "A": dY_dev, + "B": W_dev, + "D": dX_dev, + "A_bytes": dY.nbytes, + "B_bytes": Wt.nbytes, + "D_bytes": dX.nbytes, + }, + config=LaunchConfig( + grid=(q_tiles, g_tiles, N * H), + block=(spec.threads_per_block, 1, 1), + fence=True, + ), + ) + + dX_cpu = torch.empty_like(dX.cpu()) + rt.memcpy_d2h(_u8(dX_cpu), dX_dev, dX.nbytes) + for dev in (dY_dev, W_dev, dX_dev): + rt.free(dev) + synchronize_and_release(0) + + scale = ref.float().abs().max().clamp(min=1.0) + rel = float((dX_cpu.float() - ref.float().cpu()).abs().max() / scale) + tol = _TOL_BF16 if dtype == "bf16" else _TOL_FP16 + return (rel < tol), f"rel_err={rel:.3e} tol={tol:.1e}" + + +@unittest.skipUnless(not _SKIP_GPU, _SKIP_GPU or "no GPU") +class TestDirectConvExecDtype(unittest.TestCase): + """GPU execution tests: verify dtype-correct output for direct-conv. + + Uses all-ones inputs so the expected result is known analytically + (sum of kernel-window elements) and distinguishes fp16 vs bf16 because + their bit patterns for the same float value differ: + 1.0_bf16 = 0x3F80 → read as fp16 → ≈ 1.5 + 1.0_fp16 = 0x3C00 → read as bf16 → ≈ 9.77e-4 + A kernel with wrong dtype reinterpretation fails the relative-error gate. + """ + + def _fwd(self, cpg: int, dtype: str) -> None: + ok, msg = _run_direct_fwd(_ARCH, cpg, dtype) + if msg.startswith("skip"): + self.skipTest(msg) + self.assertTrue(ok, f"FAIL direct_fwd cpg={cpg} {dtype}: {msg}") + + def _dgrad(self, cpg: int, dtype: str) -> None: + ok, msg = _run_direct_dgrad(_ARCH, cpg, dtype) + if msg.startswith("skip"): + self.skipTest(msg) + self.assertTrue(ok, f"FAIL direct_dgrad cpg={cpg} {dtype}: {msg}") + + # -- forward fp16 + def test_exec_fwd_16c_fp16(self): + self._fwd(16, "fp16") + + def test_exec_fwd_8c_fp16(self): + self._fwd(8, "fp16") + + # -- forward bf16 + def test_exec_fwd_16c_bf16(self): + self._fwd(16, "bf16") + + def test_exec_fwd_8c_bf16(self): + self._fwd(8, "bf16") + + # -- dgrad fp16 + def test_exec_dgrad_16c_fp16(self): + self._dgrad(16, "fp16") + + # -- dgrad bf16 (regression: make_dgrad_fprop_spec used to drop dtype) + def test_exec_dgrad_16c_bf16(self): + self._dgrad(16, "bf16") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h index 96a427176d3a..8898ff929c31 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h @@ -95,11 +95,17 @@ typedef struct rocke_direct_conv_problem int KW; /* default 3 */ int PAD; /* default 1 */ int stride; /* default 1 */ - const char* dtype; /* "fp16" or "bf16"; default "fp16" */ + const char* dtype; /* "fp16" or "bf16"; NULL is treated as "fp16" by all + * build functions. Always set this field explicitly or + * use rocke_direct_conv_problem_default() which sets it + * to "fp16". Zero-initialising the struct leaves dtype + * NULL, which silently selects fp16 and will silently + * drop a bf16 request. */ } rocke_direct_conv_problem_t; /* DirectConvProblem with dataclass defaults (KH=KW=3, PAD=1, stride=1) and the - * six required dims zeroed. Caller fills N,H,W,groups,cpg,kpg. */ + * six required dims zeroed. dtype is initialised to "fp16". + * Caller fills N,H,W,groups,cpg,kpg; override dtype for bf16. */ rocke_direct_conv_problem_t rocke_direct_conv_problem_default(void); /* @property total_c -> groups * cpg. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h index 08b60dbe5085..e9624c61cf96 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h @@ -101,7 +101,7 @@ typedef struct rocke_dconv_16c_ctx int LOAD_VEC; /* 4 */ int NUM_VEC4; /* LDS_ROW_FP16 / LOAD_VEC */ int PASSES; /* ceil(NUM_VEC4 / THREADS) */ - int lds_total_fp16; /* PASSES * THREADS * LOAD_VEC */ + int lds_total_elems; /* PASSES * THREADS * LOAD_VEC */ int q_subtiles; /* BLOCK_Q // 16 */ int n_iters; /* H + KH - 1 */ @@ -125,7 +125,7 @@ typedef struct rocke_dconv_16c_ctx rocke_value_t* c_BG_cpg; /* const_i32(BLOCK_GROUPS * cpg) */ rocke_value_t* c_half_bytes; /* const_i32(2) */ rocke_value_t* oob_sentinel; /* const_i32((1<<31)-1) */ - rocke_value_t* fp16x4_zero; /* zero_vec_f16(4) */ + rocke_value_t* io_vec4_zero; /* zero_vec_f16(4) */ rocke_value_t* zero_acc; /* zero_vec_f32(4) */ /* ---- thread / wave / lane decode (SSA) -- */ @@ -147,7 +147,7 @@ typedef struct rocke_dconv_16c_ctx rocke_value_t* q_tile_start; /* bx * BLOCK_Q */ /* ---- LDS ping-pong buffers + buffer rsrcs -- */ - rocke_value_t* A_smem; /* smem_alloc lds_a [1, lds_total_fp16] */ + rocke_value_t* A_smem; /* smem_alloc lds_a [1, lds_total_elems] */ rocke_value_t* B_smem; /* smem_alloc lds_b (or == A_smem) */ rocke_value_t* a_rsrc; /* buffer_rsrc(A, A_bytes) */ rocke_value_t* b_rsrc; /* buffer_rsrc(Bp, B_bytes) */ @@ -225,7 +225,7 @@ typedef struct rocke_dconv_4c_ctx rocke_value_t* c_kpg; /* const_i32(kpg) */ rocke_value_t* c_half_bytes; /* const_i32(2) */ rocke_value_t* oob_sentinel; /* const_i32((1<<31)-1) */ - rocke_value_t* fp16x4_zero; /* zero_vec_f16(4) */ + rocke_value_t* io_vec4_zero; /* zero_vec_f16(4) */ rocke_value_t* zero_acc; /* zero_vec_f32(4) */ /* ---- thread / wave / lane decode (SSA) -- */ @@ -398,7 +398,7 @@ typedef struct rocke_dconv_8c_ctx int LOAD_VEC; int NUM_VEC4; int PASSES; - int lds_total_fp16; + int lds_total_elems; int q_subtiles; /* BLOCK_Q / 16 */ int n_iters; @@ -420,7 +420,7 @@ typedef struct rocke_dconv_8c_ctx rocke_value_t* c_BG_cpg; rocke_value_t* c_half_bytes; rocke_value_t* oob_sentinel; - rocke_value_t* fp16x4_zero; + rocke_value_t* io_vec4_zero; rocke_value_t* zero_acc; /* <4 x float> */ rocke_value_t* tid; @@ -499,7 +499,7 @@ typedef struct rocke_dconv_32c_ctx int NUM_VEC4; int PASSES; int N_CH_BLOCKS; /* cpg / 4 = 8 for cpg=32 */ - int lds_total_fp16; + int lds_total_elems; int q_subtiles; /* BLOCK_Q / 32 */ int n_iters; @@ -521,7 +521,7 @@ typedef struct rocke_dconv_32c_ctx rocke_value_t* c_BG_cpg; rocke_value_t* c_half_bytes; rocke_value_t* oob_sentinel; - rocke_value_t* fp16x4_zero; + rocke_value_t* io_vec4_zero; rocke_value_t* zero_acc; /* <16 x float> */ rocke_value_t* tid; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_loop_closures.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_loop_closures.cpp index 76bb13f4915f..f8b7fc31dbd8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_loop_closures.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_loop_closures.cpp @@ -42,6 +42,7 @@ int rocke_dconv16c_issue_dram_load(rocke_dconv_16c_ctx_t* ctx, int out_cap) { rocke_ir_builder_t* b = ctx->b; + const int is_bf16 = ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0; int count = 0; int i; @@ -101,12 +102,12 @@ int rocke_dconv16c_issue_dram_load(rocke_dconv_16c_ctx_t* ctx, /* safe_off = b.select(valid, a_off_bytes, oob_sentinel) */ safe_off = rocke_b_select(b, valid, a_off_bytes, ctx->oob_sentinel); /* a_vec = _buf_load_vN(a_rsrc, safe_off, c0, 2) -- dtype-dispatched */ - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) a_vec = rocke_b_buffer_load_vN_bf16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); else a_vec = rocke_b_buffer_load_vN_f16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); - /* a_vec = b.select(valid, a_vec, fp16x4_zero) */ - a_vec = rocke_b_select(b, valid, a_vec, ctx->fp16x4_zero); + /* a_vec = b.select(valid, a_vec, io_vec4_zero) */ + a_vec = rocke_b_select(b, valid, a_vec, ctx->io_vec4_zero); /* lds_idx = b.mul(cm["chunk_idx"], b.const_i32(4)) */ lds_idx = rocke_b_mul(b, ctx->chunk_meta[i].chunk_idx, rocke_b_const_i32(b, 4)); @@ -305,6 +306,7 @@ void rocke_dconv16c_prologue_prefetch(rocke_dconv_16c_ctx_t* ctx) rocke_kernel_def_t* rocke_dconv16c_stream_h_loop(rocke_dconv_16c_ctx_t* ctx) { rocke_ir_builder_t* b = ctx->b; + const int is_bf16 = ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0; const rocke_direct_conv_problem_t* p = &ctx->p; int KH = p->KH; int KW = p->KW; @@ -422,7 +424,7 @@ rocke_kernel_def_t* rocke_dconv16c_stream_h_loop(rocke_dconv_16c_ctx_t* ctx) * shape-dependent way. Op order matches Python: * acc_in = mfma_f32_16x16x32_{f16,bf16}(weights_k32[r], in_k32, acc_in) * acc_in = mfma_f32_16x16x32_{f16,bf16}(weights_s2_k32[r], in_s2, acc_in) */ - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) { acc_in = rocke_b_mfma_f32_16x16x32_bf16( b, ctx->weights_k32[r_const], in_k32[qt], acc_in); @@ -444,7 +446,7 @@ rocke_kernel_def_t* rocke_dconv16c_stream_h_loop(rocke_dconv_16c_ctx_t* ctx) { /* w_idx = r_const * KW + s_const */ int w_idx = r_const * KW + s_const; - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) acc_in = rocke_b_mfma_f32_16x16x16_bf16( b, ctx->weights[w_idx], in_s[qt][s_const], acc_in); else @@ -545,7 +547,7 @@ rocke_kernel_def_t* rocke_dconv16c_stream_h_loop(rocke_dconv_16c_ctx_t* ctx) /* safe_d_off = b.select(out_q_valid, d_base_bytes, oob_sentinel) */ safe_d_off = rocke_b_select(b, out_q_valid, d_base_bytes, ctx->oob_sentinel); /* acc_h = vec_trunc_f32_to_{f16,bf16}(acc_to_flush) */ - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) { acc_h = rocke_b_vec_trunc_f32_to_bf16(b, acc_to_flush); rocke_b_buffer_store_vN_bf16(b, ctx->d_rsrc, safe_d_off, ctx->c0, acc_h, 2); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_phases.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_phases.cpp index bd3bcd4fffc4..45cede74c216 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_phases.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_16c_phases.cpp @@ -169,11 +169,11 @@ bool rocke_dconv16c_prologue(rocke_dconv_16c_ctx_t* ctx) ctx->q_tile_start = rocke_b_mul(b, ctx->bx, ctx->c_BQ); /* ---- LDS ping-pong buffers ---- */ - ctx->lds_total_fp16 = ctx->PASSES * ctx->THREADS * ctx->LOAD_VEC; + ctx->lds_total_elems = ctx->PASSES * ctx->THREADS * ctx->LOAD_VEC; { int shape[2]; shape[0] = 1; - shape[1] = ctx->lds_total_fp16; + shape[1] = ctx->lds_total_elems; ctx->A_smem = rocke_b_smem_alloc(b, ctx->io_type, shape, 2, "lds_a"); if(spec->double_buffer) ctx->B_smem = rocke_b_smem_alloc(b, ctx->io_type, shape, 2, "lds_b"); @@ -188,7 +188,7 @@ bool rocke_dconv16c_prologue(rocke_dconv_16c_ctx_t* ctx) ctx->c_half_bytes = rocke_b_const_i32(b, 2); ctx->oob_sentinel = rocke_b_const_i32(b, ((int64_t)1 << 31) - 1); - ctx->fp16x4_zero = rocke_b_zero_vec(b, ctx->io_type, 4); + ctx->io_vec4_zero = rocke_b_zero_vec(b, ctx->io_type, 4); ctx->zero_acc = rocke_b_zero_vec_f32(b, 4); return rocke_ir_builder_ok(b); @@ -206,6 +206,7 @@ bool rocke_dconv16c_prologue(rocke_dconv_16c_ctx_t* ctx) void rocke_dconv16c_load_weights(rocke_dconv_16c_ctx_t* ctx) { rocke_ir_builder_t* b = ctx->b; + const int is_bf16 = ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0; int total_k = rocke_direct_conv_problem_total_k(&ctx->p); /* TensorDescriptor.naive("B", lengths=[total_k, KH, KW, cpg], @@ -255,7 +256,7 @@ void rocke_dconv16c_load_weights(rocke_dconv_16c_ctx_t* ctx) rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off_k32, &valid); } - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) ctx->weights_k32[r_const] = rocke_b_buffer_load_vN_bf16( b, ctx->b_rsrc, rocke_b_mul(b, w_off_k32, ctx->c_half_bytes), ctx->c0, 4); else @@ -274,7 +275,7 @@ void rocke_dconv16c_load_weights(rocke_dconv_16c_ctx_t* ctx) rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off_s2, &valid); } - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) w_s2 = rocke_b_buffer_load_vN_bf16( b, ctx->b_rsrc, rocke_b_mul(b, w_off_s2, ctx->c_half_bytes), ctx->c0, 4); else @@ -304,7 +305,7 @@ void rocke_dconv16c_load_weights(rocke_dconv_16c_ctx_t* ctx) in_values[3] = ctx->ch_lane_k16; rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off, &valid); - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) ctx->weights[ctx->n_weights++] = rocke_b_buffer_load_vN_bf16( b, ctx->b_rsrc, rocke_b_mul(b, w_off, ctx->c_half_bytes), ctx->c0, 2); else diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_32c.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_32c.cpp index b8e93f116f65..aed39cdeca36 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_32c.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_32c.cpp @@ -85,7 +85,7 @@ bool rocke_dconv32c_prologue(rocke_dconv_32c_ctx_t* ctx) rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", ctx->THREADS); { - const rocke_type_t* f16ptr = rocke_ptr_type(b, ctx->io_type, "global"); + const rocke_type_t* io_ptr = rocke_ptr_type(b, ctx->io_type, "global"); rocke_param_opts_t ro; rocke_param_opts_t wo; rocke_param_opts_t none; @@ -97,8 +97,8 @@ bool rocke_dconv32c_prologue(rocke_dconv_32c_ctx_t* ctx) ro.readonly_set = true; ro.align = 16; ro.align_set = true; - ctx->A = rocke_b_param(b, "A", f16ptr, &ro); - ctx->Bp = rocke_b_param(b, "B", f16ptr, &ro); + ctx->A = rocke_b_param(b, "A", io_ptr, &ro); + ctx->Bp = rocke_b_param(b, "B", io_ptr, &ro); wo = (rocke_param_opts_t){0}; wo.noalias = true; @@ -107,7 +107,7 @@ bool rocke_dconv32c_prologue(rocke_dconv_32c_ctx_t* ctx) wo.writeonly_set = true; wo.align = 16; wo.align_set = true; - ctx->D = rocke_b_param(b, "D", f16ptr, &wo); + ctx->D = rocke_b_param(b, "D", io_ptr, &wo); none = (rocke_param_opts_t){0}; ctx->A_bytes = rocke_b_param(b, "A_bytes", rocke_i32(), &none); @@ -140,11 +140,11 @@ bool rocke_dconv32c_prologue(rocke_dconv_32c_ctx_t* ctx) ctx->g = rocke_b_add(b, rocke_b_mul(b, ctx->g_tile, ctx->c_BG), ctx->wave_id); ctx->q_tile_start = rocke_b_mul(b, ctx->bx, ctx->c_BQ); - ctx->lds_total_fp16 = ctx->PASSES * ctx->THREADS * ctx->LOAD_VEC; + ctx->lds_total_elems = ctx->PASSES * ctx->THREADS * ctx->LOAD_VEC; { int shape[2]; shape[0] = 1; - shape[1] = ctx->lds_total_fp16; + shape[1] = ctx->lds_total_elems; ctx->A_smem = rocke_b_smem_alloc(b, ctx->io_type, shape, 2, "lds_a"); if(spec->double_buffer) { @@ -160,7 +160,7 @@ bool rocke_dconv32c_prologue(rocke_dconv_32c_ctx_t* ctx) ctx->b_rsrc = rocke_b_buffer_rsrc(b, ctx->Bp, ctx->B_bytes); ctx->d_rsrc = rocke_b_buffer_rsrc(b, ctx->D, ctx->D_bytes); - ctx->fp16x4_zero = rocke_b_zero_vec(b, ctx->io_type, 4); + ctx->io_vec4_zero = rocke_b_zero_vec(b, ctx->io_type, 4); /* zero_acc = b.zero_vec_f32(16) — 16 f32 elements (32*32/64) */ ctx->zero_acc = rocke_b_zero_vec_f32(b, 16); @@ -177,6 +177,7 @@ bool rocke_dconv32c_prologue(rocke_dconv_32c_ctx_t* ctx) void rocke_dconv32c_load_weights(rocke_dconv_32c_ctx_t* ctx) { rocke_ir_builder_t* b = ctx->b; + const int is_bf16 = ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0; int total_k = rocke_direct_conv_problem_total_k(&ctx->p); int r_const, s_const, atom_idx; @@ -221,7 +222,7 @@ void rocke_dconv32c_load_weights(rocke_dconv_32c_ctx_t* ctx) in_values[3] = ch_off; rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off, &valid); - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) ctx->weights[r_const][s_const][atom_idx] = rocke_b_buffer_load_vN_bf16( b, ctx->b_rsrc, rocke_b_mul(b, w_off, ctx->c_half_bytes), ctx->c0, 2); else @@ -370,6 +371,7 @@ static int rocke_dconv32c_issue_dram_load(rocke_dconv_32c_ctx_t* ctx, int out_cap) { rocke_ir_builder_t* b = ctx->b; + const int is_bf16 = ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0; int count = 0; int i; @@ -417,11 +419,11 @@ static int rocke_dconv32c_issue_dram_load(rocke_dconv_32c_ctx_t* ctx, valid = rocke_b_land(b, addr_valid, ctx->chunk_meta[i].in_bounds); a_off_bytes = rocke_b_mul(b, a_off_elems, ctx->c_half_bytes); safe_off = rocke_b_select(b, valid, a_off_bytes, ctx->oob_sentinel); - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) a_vec = rocke_b_buffer_load_vN_bf16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); else a_vec = rocke_b_buffer_load_vN_f16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); - a_vec = rocke_b_select(b, valid, a_vec, ctx->fp16x4_zero); + a_vec = rocke_b_select(b, valid, a_vec, ctx->io_vec4_zero); lds_idx = rocke_b_mul(b, ctx->chunk_meta[i].chunk_idx, rocke_b_const_i32(b, 4)); out_vecs[count] = a_vec; @@ -519,6 +521,7 @@ void rocke_dconv32c_prologue_prefetch(rocke_dconv_32c_ctx_t* ctx) rocke_kernel_def_t* rocke_dconv32c_stream_h_loop(rocke_dconv_32c_ctx_t* ctx) { rocke_ir_builder_t* b = ctx->b; + const int is_bf16 = ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0; const rocke_direct_conv_problem_t* p = &ctx->p; int KH = p->KH; int KW = p->KW; @@ -604,7 +607,7 @@ rocke_kernel_def_t* rocke_dconv32c_stream_h_loop(rocke_dconv_32c_ctx_t* ctx) { for(atom_idx = 0; atom_idx < 4; ++atom_idx) { - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) acc_in = rocke_b_mfma_f32_32x32x8_bf16( b, ctx->weights[r_const][s_const][atom_idx], @@ -721,7 +724,7 @@ rocke_kernel_def_t* rocke_dconv32c_stream_h_loop(rocke_dconv_32c_ctx_t* ctx) pair_elems[0] = e0; pair_elems[1] = e1; pair_f32 = rocke_b_vec_pack(b, pair_elems, 2, rocke_f32()); - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) { pair_h = rocke_b_vec_trunc_f32_to_bf16(b, pair_f32); rocke_b_buffer_store_vN_bf16( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_4c_phases_and_loop.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_4c_phases_and_loop.cpp index 88847ab3b69a..666ac4fbd281 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_4c_phases_and_loop.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_4c_phases_and_loop.cpp @@ -145,7 +145,7 @@ bool rocke_dconv4c_prologue(rocke_dconv_4c_ctx_t* ctx) ctx->a_rsrc = rocke_b_buffer_rsrc(b, ctx->A, ctx->A_bytes); ctx->b_rsrc = rocke_b_buffer_rsrc(b, ctx->Bp, ctx->B_bytes); ctx->d_rsrc = rocke_b_buffer_rsrc(b, ctx->D, ctx->D_bytes); - ctx->fp16x4_zero = rocke_b_zero_vec_f16(b, 4); + ctx->io_vec4_zero = rocke_b_zero_vec_f16(b, 4); ctx->zero_acc = rocke_b_zero_vec_f32(b, 4); return rocke_ir_builder_ok(b); @@ -338,8 +338,8 @@ rocke_kernel_def_t* rocke_dconv4c_stream_h_loop(rocke_dconv_4c_ctx_t* ctx) b, valid, rocke_b_mul(b, a_off, ctx->c_half_bytes), ctx->oob_sentinel); /* Line 985: vec = buffer_load_vN_f16(a_rsrc, safe_a, c0, 2). */ vec = rocke_b_buffer_load_vN_f16(b, ctx->a_rsrc, safe_a, ctx->c0, 2); - /* Line 986: vec = select(valid, vec, fp16x4_zero). */ - vec = rocke_b_select(b, valid, vec, ctx->fp16x4_zero); + /* Line 986: vec = select(valid, vec, io_vec4_zero). */ + vec = rocke_b_select(b, valid, vec, ctx->io_vec4_zero); inputs_by_qtile[qt][s_idx] = vec; } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_8c.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_8c.cpp index d509478241ac..6c806fa0e19e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_8c.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_8c.cpp @@ -88,7 +88,7 @@ bool rocke_dconv8c_prologue(rocke_dconv_8c_ctx_t* ctx) /* ---- params ---- */ { - const rocke_type_t* f16ptr = rocke_ptr_type(b, ctx->io_type, "global"); + const rocke_type_t* io_ptr = rocke_ptr_type(b, ctx->io_type, "global"); rocke_param_opts_t ro; rocke_param_opts_t wo; rocke_param_opts_t none; @@ -100,8 +100,8 @@ bool rocke_dconv8c_prologue(rocke_dconv_8c_ctx_t* ctx) ro.readonly_set = true; ro.align = 16; ro.align_set = true; - ctx->A = rocke_b_param(b, "A", f16ptr, &ro); - ctx->Bp = rocke_b_param(b, "B", f16ptr, &ro); + ctx->A = rocke_b_param(b, "A", io_ptr, &ro); + ctx->Bp = rocke_b_param(b, "B", io_ptr, &ro); wo = (rocke_param_opts_t){0}; wo.noalias = true; @@ -110,7 +110,7 @@ bool rocke_dconv8c_prologue(rocke_dconv_8c_ctx_t* ctx) wo.writeonly_set = true; wo.align = 16; wo.align_set = true; - ctx->D = rocke_b_param(b, "D", f16ptr, &wo); + ctx->D = rocke_b_param(b, "D", io_ptr, &wo); none = (rocke_param_opts_t){0}; ctx->A_bytes = rocke_b_param(b, "A_bytes", rocke_i32(), &none); @@ -157,11 +157,11 @@ bool rocke_dconv8c_prologue(rocke_dconv_8c_ctx_t* ctx) ctx->q_tile_start = rocke_b_mul(b, ctx->bx, ctx->c_BQ); /* ---- LDS (Python lines 1305-1315) ---- */ - ctx->lds_total_fp16 = ctx->PASSES * ctx->THREADS * ctx->LOAD_VEC; + ctx->lds_total_elems = ctx->PASSES * ctx->THREADS * ctx->LOAD_VEC; { int shape[2]; shape[0] = 1; - shape[1] = ctx->lds_total_fp16; + shape[1] = ctx->lds_total_elems; ctx->A_smem = rocke_b_smem_alloc(b, ctx->io_type, shape, 2, "lds_a"); if(spec->double_buffer) { @@ -177,7 +177,7 @@ bool rocke_dconv8c_prologue(rocke_dconv_8c_ctx_t* ctx) ctx->b_rsrc = rocke_b_buffer_rsrc(b, ctx->Bp, ctx->B_bytes); ctx->d_rsrc = rocke_b_buffer_rsrc(b, ctx->D, ctx->D_bytes); - ctx->fp16x4_zero = rocke_b_zero_vec(b, ctx->io_type, 4); + ctx->io_vec4_zero = rocke_b_zero_vec(b, ctx->io_type, 4); ctx->zero_acc = rocke_b_zero_vec_f32(b, 4); return rocke_ir_builder_ok(b); @@ -194,6 +194,7 @@ bool rocke_dconv8c_prologue(rocke_dconv_8c_ctx_t* ctx) void rocke_dconv8c_load_weights(rocke_dconv_8c_ctx_t* ctx) { rocke_ir_builder_t* b = ctx->b; + const int is_bf16 = ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0; int total_k = rocke_direct_conv_problem_total_k(&ctx->p); int r_const; @@ -234,7 +235,7 @@ void rocke_dconv8c_load_weights(rocke_dconv_8c_ctx_t* ctx) rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off_main, &valid); } - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) ctx->weights_main[r_const] = rocke_b_buffer_load_vN_bf16( b, ctx->b_rsrc, rocke_b_mul(b, w_off_main, ctx->c_half_bytes), ctx->c0, 2); else @@ -252,13 +253,13 @@ void rocke_dconv8c_load_weights(rocke_dconv_8c_ctx_t* ctx) rocke_transforms_descriptor_offset( b, ctx->b_desc, in_names, in_values, 4, &w_off_s2, &valid); } - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) w_s2 = rocke_b_buffer_load_vN_bf16( b, ctx->b_rsrc, rocke_b_mul(b, w_off_s2, ctx->c_half_bytes), ctx->c0, 2); else w_s2 = rocke_b_buffer_load_vN_f16( b, ctx->b_rsrc, rocke_b_mul(b, w_off_s2, ctx->c_half_bytes), ctx->c0, 2); - ctx->weights_s2[r_const] = rocke_b_select(b, ctx->lane_in_lo_half, w_s2, ctx->fp16x4_zero); + ctx->weights_s2[r_const] = rocke_b_select(b, ctx->lane_in_lo_half, w_s2, ctx->io_vec4_zero); } ctx->n_weights = ctx->p.KH; } @@ -407,6 +408,7 @@ static int rocke_dconv8c_issue_dram_load(rocke_dconv_8c_ctx_t* ctx, int out_cap) { rocke_ir_builder_t* b = ctx->b; + const int is_bf16 = ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0; int count = 0; int i; @@ -455,11 +457,11 @@ static int rocke_dconv8c_issue_dram_load(rocke_dconv_8c_ctx_t* ctx, valid = rocke_b_land(b, addr_valid, ctx->chunk_meta[i].in_bounds); a_off_bytes = rocke_b_mul(b, a_off_elems, ctx->c_half_bytes); safe_off = rocke_b_select(b, valid, a_off_bytes, ctx->oob_sentinel); - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) a_vec = rocke_b_buffer_load_vN_bf16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); else a_vec = rocke_b_buffer_load_vN_f16(b, ctx->a_rsrc, safe_off, ctx->c0, 2); - a_vec = rocke_b_select(b, valid, a_vec, ctx->fp16x4_zero); + a_vec = rocke_b_select(b, valid, a_vec, ctx->io_vec4_zero); lds_idx = rocke_b_mul(b, ctx->chunk_meta[i].chunk_idx, rocke_b_const_i32(b, 4)); out_vecs[count] = a_vec; @@ -524,7 +526,7 @@ static rocke_value_t* } /* lds_read_input_s2(q_subtile, lds) -- per-lane <4 x half> read for the s=2 residual. - * Valid only for c4∈{0,1}; zeroed via select(lane_in_lo_half, vec, fp16x4_zero). */ + * Valid only for c4∈{0,1}; zeroed via select(lane_in_lo_half, vec, io_vec4_zero). */ static rocke_value_t* rocke_dconv8c_lds_read_s2(rocke_dconv_8c_ctx_t* ctx, int q_subtile, rocke_value_t* lds) { @@ -557,7 +559,7 @@ static rocke_value_t* indices[0] = ctx->c0; indices[1] = lds_idx; vec = rocke_b_smem_load_vN(b, lds, indices, 2, ctx->io_type, 4); - return rocke_b_select(b, ctx->lane_in_lo_half, vec, ctx->fp16x4_zero); + return rocke_b_select(b, ctx->lane_in_lo_half, vec, ctx->io_vec4_zero); } /* ===================================================================== * @@ -584,6 +586,7 @@ void rocke_dconv8c_prologue_prefetch(rocke_dconv_8c_ctx_t* ctx) rocke_kernel_def_t* rocke_dconv8c_stream_h_loop(rocke_dconv_8c_ctx_t* ctx) { rocke_ir_builder_t* b = ctx->b; + const int is_bf16 = ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0; const rocke_direct_conv_problem_t* p = &ctx->p; int KH = p->KH; int y; @@ -662,14 +665,14 @@ rocke_kernel_def_t* rocke_dconv8c_stream_h_loop(rocke_dconv_8c_ctx_t* ctx) rocke_value_t* acc_in = ctx->acc_tiles[qt][p_idx]; /* acc_in = mfma_16x16x16(weights_main[r], in_main, acc_in) */ - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) acc_in = rocke_b_mfma_f32_16x16x16_bf16( b, ctx->weights_main[r_const], in_main[qt], acc_in); else acc_in = rocke_b_mfma_f32_16x16x16_f16( b, ctx->weights_main[r_const], in_main[qt], acc_in); /* acc_in = mfma_16x16x16(weights_s2[r], in_s2, acc_in) */ - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) acc_in = rocke_b_mfma_f32_16x16x16_bf16( b, ctx->weights_s2[r_const], in_s2[qt], acc_in); else @@ -750,7 +753,7 @@ rocke_kernel_def_t* rocke_dconv8c_stream_h_loop(rocke_dconv_8c_ctx_t* ctx) /* store_valid = b.land(out_q_valid, c4_valid) */ store_valid = rocke_b_land(b, out_q_valid, c4_valid); safe_d_off = rocke_b_select(b, store_valid, d_base_bytes, ctx->oob_sentinel); - if(ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) + if(is_bf16) { acc_h = rocke_b_vec_trunc_f32_to_bf16(b, acc_to_flush); rocke_b_buffer_store_vN_bf16(b, ctx->d_rsrc, safe_d_off, ctx->c0, acc_h, 2); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp index 6127001534bd..af9cdce8d86f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp @@ -438,6 +438,21 @@ rocke_status_t rocke_direct_conv_4c_validate(const rocke_direct_conv_4c_spec_t* return ROCKE_ERR_VALUE; } p = &spec->problem; + /* if p.dtype != "fp16": raise ValueError(...) — no mfma_f32_4x4x4_bf16 atom on CDNA */ + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0) + { + if(reason != NULL && reason_cap > 0) + { + snprintf(reason, + reason_cap, + "DirectConv4cSpec: bf16 is not supported - the mfma_f32_4x4x4 atom " + "is fp16-only on CDNA; use 8c/16c/32c for bf16"); + } + return ROCKE_ERR_VALUE; + } + } /* if p.cpg != 4 or p.kpg != 4: raise ValueError(...) */ if(p->cpg != 4 || p->kpg != 4) { @@ -522,6 +537,15 @@ bool rocke_direct_conv_4c_is_valid_spec(const rocke_direct_conv_4c_spec_t* spec, } p = &spec->problem; + /* if p.dtype != "fp16": return False, ... — no mfma_f32_4x4x4_bf16 atom on CDNA */ + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0) + { + CK_DCONV4C_REJECT("DirectConv4cSpec: bf16 is not supported - the mfma_f32_4x4x4 atom " + "is fp16-only on CDNA; use 8c/16c/32c for bf16"); + } + } /* if p.cpg != 4 or p.kpg != 4: return False, ... */ if(p->cpg != 4 || p->kpg != 4) { @@ -1387,13 +1411,21 @@ rocke_status_t rocke_direct_conv_dgrad_kernel_name(const rocke_direct_conv_dgrad { return ROCKE_ERR_VALUE; } - /* kernel_name_join(name, p.short(), f"bq{block_q}", f"bg{block_groups}") */ + /* kernel_name_join(name, p.short(), f"bq{block_q}", f"bg{block_groups}", + * flags={"bf16": dtype=="bf16"}) */ snprintf(bq_buf, sizeof(bq_buf), "bq%d", spec->block_q); snprintf(bg_buf, sizeof(bg_buf), "bg%d", spec->block_groups); parts[0] = prob_short; parts[1] = bq_buf; parts[2] = bg_buf; - return rocke_kernel_name_join(spec->name, parts, 3, NULL, NULL, 0, out, out_cap, NULL); + { + const char* flag_names[1] = {"bf16"}; + int flag_on[1]; + const char* dt = spec->problem.dtype ? spec->problem.dtype : "fp16"; + flag_on[0] = (strcmp(dt, "bf16") == 0) ? 1 : 0; + return rocke_kernel_name_join( + spec->name, parts, 3, flag_names, flag_on, 1, out, out_cap, NULL); + } } rocke_status_t rocke_direct_conv_dgrad_validate(const rocke_direct_conv_dgrad_spec_t* spec, @@ -1406,6 +1438,18 @@ rocke_status_t rocke_direct_conv_dgrad_validate(const rocke_direct_conv_dgrad_sp return ROCKE_ERR_VALUE; } p = &spec->problem; + /* if p.dtype not in ("fp16", "bf16"): raise ValueError(...) */ + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) + { + if(reason && reason_cap > 0) + { + snprintf(reason, reason_cap, "DirectConvDgradSpec: unsupported dtype '%s'", dt); + } + return ROCKE_ERR_VALUE; + } + } if(p->cpg < 1) { if(reason && reason_cap > 0) @@ -1461,6 +1505,19 @@ bool rocke_direct_conv_dgrad_is_valid_spec(const rocke_direct_conv_dgrad_spec_t* return false; } p = &spec->problem; + /* if p.dtype not in ("fp16", "bf16"): return False, ... */ + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) + { + if(reason && reason_cap > 0) + { + snprintf( + reason, reason_cap, "unsupported dtype '%s'; expected 'fp16' or 'bf16'", dt); + } + return false; + } + } if(p->cpg < 1) { if(reason && reason_cap > 0) From 15d13c324ce4140e2a157b664bc5cbaa1f206ea3 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Wed, 23 Sep 2026 13:20:31 +0000 Subject: [PATCH 4/5] fixes --- .../conv_direct_grouped_build_dgrad.cpp | 50 +++++++++++++------ .../conv_direct_grouped_spec_and_validate.cpp | 19 +++++-- 2 files changed, 50 insertions(+), 19 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_dgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_dgrad.cpp index a344b2345ca4..21674cbaaf57 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_dgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_dgrad.cpp @@ -155,8 +155,10 @@ rocke_kernel_def_t* rocke_build_direct_conv_dgrad(rocke_ir_builder_t* b, rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", THREADS); - /* Params */ - const rocke_type_t* f16ptr = rocke_ptr_type(b, rocke_f16(), "global"); + /* Params — io_type = _io_type(p.dtype): f16 or bf16 IR type. */ + const rocke_type_t* io_type = rocke_b_io_ir_type(b, p->dtype ? p->dtype : "fp16"); + const rocke_type_t* ioptr = rocke_ptr_type(b, io_type, "global"); + const int is_bf16 = (p->dtype && strcmp(p->dtype, "bf16") == 0); rocke_param_opts_t ro = {0}; ro.noalias = true; ro.noalias_set = true; @@ -173,9 +175,9 @@ rocke_kernel_def_t* rocke_build_direct_conv_dgrad(rocke_ir_builder_t* b, wo_opts.align_set = true; rocke_param_opts_t none = {0}; - rocke_value_t* A = rocke_b_param(b, "A", f16ptr, &ro); - rocke_value_t* Bp = rocke_b_param(b, "B", f16ptr, &ro); - rocke_value_t* D = rocke_b_param(b, "D", f16ptr, &wo_opts); + rocke_value_t* A = rocke_b_param(b, "A", ioptr, &ro); + rocke_value_t* Bp = rocke_b_param(b, "B", ioptr, &ro); + rocke_value_t* D = rocke_b_param(b, "D", ioptr, &wo_opts); rocke_value_t* A_bytes = rocke_b_param(b, "A_bytes", rocke_i32(), &none); rocke_value_t* B_bytes = rocke_b_param(b, "B_bytes", rocke_i32(), &none); rocke_value_t* D_bytes = rocke_b_param(b, "D_bytes", rocke_i32(), &none); @@ -365,7 +367,8 @@ rocke_kernel_def_t* rocke_build_direct_conv_dgrad(rocke_ir_builder_t* b, rocke_value_t* k_byte_off = rocke_b_mul(b, k_iv, k_stride_bytes); rocke_value_t* w_byte = rocke_b_add(b, w_off0_bytes, k_byte_off); rocke_value_t* safe_w = rocke_b_select(b, tap_valid, w_byte, oob_sentinel); - rocke_value_t* w_h = rocke_b_buffer_load_f16(b, b_rsrc, safe_w, c0); + rocke_value_t* w_h = is_bf16 ? rocke_b_buffer_load_bf16(b, b_rsrc, safe_w, c0) + : rocke_b_buffer_load_f16(b, b_rsrc, safe_w, c0); rocke_value_t* w_cast = rocke_b_cast_to_f32(b, w_h); rocke_value_t* w_f32 = rocke_b_select(b, tap_valid, w_cast, zero_f32); @@ -373,7 +376,8 @@ rocke_kernel_def_t* rocke_build_direct_conv_dgrad(rocke_ir_builder_t* b, rocke_value_t* dy_byte = rocke_b_add(b, dy_off0_bytes, rocke_b_mul(b, k_iv, c_half_bytes)); rocke_value_t* safe_dy = rocke_b_select(b, tap_valid, dy_byte, oob_sentinel); - rocke_value_t* dy_h = rocke_b_buffer_load_f16(b, a_rsrc, safe_dy, c0); + rocke_value_t* dy_h = is_bf16 ? rocke_b_buffer_load_bf16(b, a_rsrc, safe_dy, c0) + : rocke_b_buffer_load_f16(b, a_rsrc, safe_dy, c0); rocke_value_t* dy_cast = rocke_b_cast_to_f32(b, dy_h); rocke_value_t* dy_f32 = rocke_b_select(b, tap_valid, dy_cast, zero_f32); @@ -400,7 +404,12 @@ rocke_kernel_def_t* rocke_build_direct_conv_dgrad(rocke_ir_builder_t* b, rocke_value_t* store_guard = rocke_b_land(b, c_in_ok, wi_ok); rocke_value_t* d_bytes = rocke_b_mul(b, d_off, c_half_bytes); rocke_value_t* safe_d = rocke_b_select(b, store_guard, d_bytes, oob_sentinel); - rocke_b_buffer_store_f16(b, d_rsrc, safe_d, c0, rocke_b_trunc_f32_to_f16(b, acc)); + rocke_value_t* acc_h + = is_bf16 ? rocke_b_trunc_f32_to_bf16(b, acc) : rocke_b_trunc_f32_to_f16(b, acc); + if(is_bf16) + rocke_b_buffer_store_bf16(b, d_rsrc, safe_d, c0, acc_h); + else + rocke_b_buffer_store_f16(b, d_rsrc, safe_d, c0, acc_h); } rocke_value_t* hi_yield[1]; @@ -448,8 +457,10 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_dgrad( rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", THREADS); - /* Params */ - const rocke_type_t* f16ptr = rocke_ptr_type(b, rocke_f16(), "global"); + /* Params — io_type = _io_type(p.dtype): f16 or bf16 IR type. */ + const rocke_type_t* io_type = rocke_b_io_ir_type(b, p->dtype ? p->dtype : "fp16"); + const rocke_type_t* ioptr = rocke_ptr_type(b, io_type, "global"); + const int is_bf16 = (p->dtype && strcmp(p->dtype, "bf16") == 0); rocke_param_opts_t ro = {0}; ro.noalias = true; ro.noalias_set = true; @@ -466,9 +477,9 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_dgrad( wo_opts.align_set = true; rocke_param_opts_t none = {0}; - rocke_value_t* A = rocke_b_param(b, "A", f16ptr, &ro); - rocke_value_t* Bp = rocke_b_param(b, "B", f16ptr, &ro); - rocke_value_t* D = rocke_b_param(b, "D", f16ptr, &wo_opts); + rocke_value_t* A = rocke_b_param(b, "A", ioptr, &ro); + rocke_value_t* Bp = rocke_b_param(b, "B", ioptr, &ro); + rocke_value_t* D = rocke_b_param(b, "D", ioptr, &wo_opts); rocke_value_t* A_bytes = rocke_b_param(b, "A_bytes", rocke_i32(), &none); rocke_value_t* B_bytes = rocke_b_param(b, "B_bytes", rocke_i32(), &none); rocke_value_t* D_bytes = rocke_b_param(b, "D_bytes", rocke_i32(), &none); @@ -568,7 +579,8 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_dgrad( /* Python: select(ch_in_range, mul(w_off, c_half_bytes), oob_sentinel) */ rocke_value_t* w_bytes = rocke_b_mul(b, w_off, c_half_bytes); rocke_value_t* safe_w = rocke_b_select(b, ch_in_range, w_bytes, oob_sentinel); - rocke_value_t* w_h = rocke_b_buffer_load_f16(b, b_rsrc, safe_w, c0); + rocke_value_t* w_h = is_bf16 ? rocke_b_buffer_load_bf16(b, b_rsrc, safe_w, c0) + : rocke_b_buffer_load_f16(b, b_rsrc, safe_w, c0); rocke_value_t* w_cast = rocke_b_cast_to_f32(b, w_h); weights_f32[r_const * p->KW + s_const] = rocke_b_select(b, ch_in_range, w_cast, zero_f32); @@ -638,7 +650,8 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_dgrad( /* Python: select(valid, mul(dy_off, c_half_bytes), oob_sentinel) */ rocke_value_t* dy_bytes = rocke_b_mul(b, dy_off, c_half_bytes); rocke_value_t* safe_dy = rocke_b_select(b, valid, dy_bytes, oob_sentinel); - rocke_value_t* dy_h = rocke_b_buffer_load_f16(b, a_rsrc, safe_dy, c0); + rocke_value_t* dy_h = is_bf16 ? rocke_b_buffer_load_bf16(b, a_rsrc, safe_dy, c0) + : rocke_b_buffer_load_f16(b, a_rsrc, safe_dy, c0); rocke_value_t* dy_cast = rocke_b_cast_to_f32(b, dy_h); rocke_value_t* dy_f32 = rocke_b_select(b, valid, dy_cast, zero_f32); acc = rocke_b_fma(b, weights_f32[r_const * p->KW + s_const], dy_f32, acc); @@ -658,7 +671,12 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_dgrad( rocke_value_t* store_guard = rocke_b_land(b, ch_in_range, wi_ok); rocke_value_t* d_bytes = rocke_b_mul(b, d_off, c_half_bytes); rocke_value_t* safe_d = rocke_b_select(b, store_guard, d_bytes, oob_sentinel); - rocke_b_buffer_store_f16(b, d_rsrc, safe_d, c0, rocke_b_trunc_f32_to_f16(b, acc)); + rocke_value_t* acc_h + = is_bf16 ? rocke_b_trunc_f32_to_bf16(b, acc) : rocke_b_trunc_f32_to_f16(b, acc); + if(is_bf16) + rocke_b_buffer_store_bf16(b, d_rsrc, safe_d, c0, acc_h); + else + rocke_b_buffer_store_f16(b, d_rsrc, safe_d, c0, acc_h); } rocke_value_t* hi_yield[1]; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp index af9cdce8d86f..503eb1ec47b6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp @@ -448,7 +448,7 @@ rocke_status_t rocke_direct_conv_4c_validate(const rocke_direct_conv_4c_spec_t* snprintf(reason, reason_cap, "DirectConv4cSpec: bf16 is not supported - the mfma_f32_4x4x4 atom " - "is fp16-only on CDNA; use 8c/16c/32c for bf16"); + "is fp16-only; use fp16 dtype or a different cpg variant"); } return ROCKE_ERR_VALUE; } @@ -542,8 +542,8 @@ bool rocke_direct_conv_4c_is_valid_spec(const rocke_direct_conv_4c_spec_t* spec, const char* dt = p->dtype ? p->dtype : "fp16"; if(strcmp(dt, "fp16") != 0) { - CK_DCONV4C_REJECT("DirectConv4cSpec: bf16 is not supported - the mfma_f32_4x4x4 atom " - "is fp16-only on CDNA; use 8c/16c/32c for bf16"); + CK_DCONV4C_REJECT("DirectConv4cSpec: bf16 not supported - " + "no mfma_f32_4x4x4_bf16 atom"); } } /* if p.cpg != 4 or p.kpg != 4: return False, ... */ @@ -1613,6 +1613,19 @@ rocke_status_t rocke_direct_depthwise_dgrad_validate( return ROCKE_ERR_VALUE; } p = &spec->problem; + /* if p.dtype != "fp16": raise ValueError(...) — depthwise kernels are fp16-only */ + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0) + { + if(reason && reason_cap > 0) + snprintf(reason, + reason_cap, + "DirectDepthwiseDgradSpec: bf16 not supported; " + "depthwise kernels are fp16-only"); + return ROCKE_ERR_VALUE; + } + } if(p->cpg != 1 || p->kpg != 1) { if(reason && reason_cap > 0) From 07480119bca6e1bcc018bb74b4262529a743c5b2 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Wed, 23 Sep 2026 19:41:12 +0000 Subject: [PATCH 5/5] fixes --- .../kernels/common/conv_direct_grouped.py | 188 ++++++++++++++---- .../tests/parity/conv_direct_grouped_emit.c | 94 +++++++++ .../tests/parity/conv_direct_grouped_emit.py | 120 +++++++++++ .../rocke/instance_conv_direct_grouped.h | 2 + .../instance_conv_direct_grouped_internal.h | 1 + .../conv_direct_grouped_build_depthwise.cpp | 38 +++- ...direct_grouped_build_depthwise_spatial.cpp | 37 +++- .../conv_direct_grouped_build_dgrad.cpp | 2 + .../conv_direct_grouped_spec_and_validate.cpp | 96 ++++++++- 9 files changed, 503 insertions(+), 75 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py index 1e2d1955c03d..f2835cb33cbd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py @@ -4050,13 +4050,14 @@ def kernel_name(self) -> str: p.short(), f"bw{self.block_w}", f"bw{self.block_waves}wv", + flags={"bf16": p.dtype == "bf16"}, ) def validate(self) -> None: p = self.problem - if p.dtype != "fp16": + if p.dtype not in ("fp16", "bf16"): raise ValueError( - f"DirectDepthwiseDgradSpec: bf16 not supported; depthwise kernels are fp16-only" + f"DirectDepthwiseDgradSpec: unsupported dtype {p.dtype!r}; expected fp16 or bf16" ) if p.cpg != 1 or p.kpg != 1: raise ValueError( @@ -4075,6 +4076,11 @@ def is_valid_depthwise_dgrad_spec( except KeyError as e: return False, str(e) p = spec.problem + if p.dtype not in ("fp16", "bf16"): + return ( + False, + f"DirectDepthwiseDgradSpec: unsupported dtype {p.dtype!r}; expected fp16 or bf16", + ) if p.cpg != 1 or p.kpg != 1: return False, f"requires cpg=kpg=1 (got {p.cpg}, {p.kpg})" return True, "ok" @@ -4119,9 +4125,10 @@ def build_direct_depthwise_dgrad( b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = THREADS - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - Bp = b.param("B", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + io_type = _io_type(p.dtype) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + Bp = b.param("B", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) B_bytes = b.param("B_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -4171,7 +4178,11 @@ def build_direct_depthwise_dgrad( b, k=ch, r=b.const_i32(r_const), s=b.const_i32(s_const), c=c0 ) safe_w = b.select(ch_in_range, b.mul(w_off, c_half_bytes), oob_sentinel) - w_h = b.buffer_load_f16(b_rsrc, safe_w, c0) + w_h = ( + b.buffer_load_bf16(b_rsrc, safe_w, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(b_rsrc, safe_w, c0) + ) row.append(b.select(ch_in_range, b.cast_to_f32(w_h), zero_f32)) weights_f32.append(row) @@ -4234,7 +4245,11 @@ def build_direct_depthwise_dgrad( dy_off, _ = dy_desc.offset(b, n=n, ho=ho, wo=wo, ch=ch) safe_dy = b.select(valid, b.mul(dy_off, c_half_bytes), oob_sentinel) - dy_h = b.buffer_load_f16(a_rsrc, safe_dy, c0) + dy_h = ( + b.buffer_load_bf16(a_rsrc, safe_dy, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(a_rsrc, safe_dy, c0) + ) dy_f32 = b.select(valid, b.cast_to_f32(dy_h), zero_f32) acc = b.fma(weights_f32[r_const][s_const], dy_f32, acc) @@ -4243,7 +4258,10 @@ def build_direct_depthwise_dgrad( safe_d = b.select( b.land(ch_in_range, wi_ok), b.mul(d_off, c_half_bytes), oob_sentinel ) - b.buffer_store_f16(d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc)) + if p.dtype == "bf16": + b.buffer_store_bf16(d_rsrc, safe_d, c0, b.trunc_f32_to_bf16(acc)) + else: + b.buffer_store_f16(d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc)) b.scf_yield(dummy_in) @@ -4305,13 +4323,14 @@ def kernel_name(self) -> str: p.short(), f"bw{self.block_w}", f"bw{self.block_waves}wv", + flags={"bf16": p.dtype == "bf16"}, ) def validate(self) -> None: p = self.problem - if p.dtype != "fp16": + if p.dtype not in ("fp16", "bf16"): raise ValueError( - f"DirectDepthwiseDgradStreamSpec: bf16 not supported; depthwise kernels are fp16-only" + f"DirectDepthwiseDgradStreamSpec: unsupported dtype {p.dtype!r}; expected fp16 or bf16" ) if p.cpg != 1 or p.kpg != 1: raise ValueError( @@ -4330,6 +4349,11 @@ def is_valid_depthwise_dgrad_stream_spec( except KeyError as e: return False, str(e) p = spec.problem + if p.dtype not in ("fp16", "bf16"): + return ( + False, + f"DirectDepthwiseDgradStreamSpec: unsupported dtype {p.dtype!r}; expected fp16 or bf16", + ) if p.cpg != 1 or p.kpg != 1: return False, f"requires cpg=kpg=1 (got {p.cpg}, {p.kpg})" return True, "ok" @@ -4373,9 +4397,10 @@ def build_direct_depthwise_dgrad_streaming( b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = THREADS - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - Bp = b.param("B", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + io_type = _io_type(p.dtype) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + Bp = b.param("B", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) B_bytes = b.param("B_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -4427,7 +4452,11 @@ def build_direct_depthwise_dgrad_streaming( b, k=ch, r=b.const_i32(r_const), s=b.const_i32(s_const), c=c0 ) safe_w = b.select(ch_ok, b.mul(w_off, c_half_bytes), oob_sentinel) - w_h = b.buffer_load_f16(b_rsrc, safe_w, c0) + w_h = ( + b.buffer_load_bf16(b_rsrc, safe_w, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(b_rsrc, safe_w, c0) + ) row.append(b.select(ch_ok, b.cast_to_f32(w_h), zero_f32)) weights_f32.append(row) @@ -4492,7 +4521,11 @@ def build_direct_depthwise_dgrad_streaming( safe_dy = b.select( tap_valid, b.mul(dy_off, c_half_bytes), oob_sentinel ) - dy_h = b.buffer_load_f16(a_rsrc, safe_dy, c0) + dy_h = ( + b.buffer_load_bf16(a_rsrc, safe_dy, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(a_rsrc, safe_dy, c0) + ) dy_f32 = b.select(tap_valid, b.cast_to_f32(dy_h), zero_f32) acc_slots[slot][j] = b.fma( @@ -4519,9 +4552,14 @@ def build_direct_depthwise_dgrad_streaming( safe_d = b.select( b.land(ch_ok, wi_ok), b.mul(d_off, c_half_bytes), oob_sentinel ) - b.buffer_store_f16( - d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc_slots[slot][j]) - ) + if p.dtype == "bf16": + b.buffer_store_bf16( + d_rsrc, safe_d, c0, b.trunc_f32_to_bf16(acc_slots[slot][j]) + ) + else: + b.buffer_store_f16( + d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc_slots[slot][j]) + ) # Reset slot for future use. acc_slots[slot][j] = zero_f32 @@ -4580,13 +4618,14 @@ def kernel_name(self) -> str: p.short(), f"bw{self.block_w}", f"bw{self.block_waves}wv", + flags={"bf16": p.dtype == "bf16"}, ) def validate(self) -> None: p = self.problem - if p.dtype != "fp16": + if p.dtype not in ("fp16", "bf16"): raise ValueError( - f"DirectDepthwiseSpec: bf16 not supported; depthwise kernels are fp16-only" + f"DirectDepthwiseSpec: unsupported dtype {p.dtype!r}; expected fp16 or bf16" ) if p.cpg != 1 or p.kpg != 1: raise ValueError( @@ -4610,6 +4649,11 @@ def is_valid_depthwise_spec( return False, str(e) p = spec.problem + if p.dtype not in ("fp16", "bf16"): + return ( + False, + f"DirectDepthwiseSpec: unsupported dtype {p.dtype!r}; expected fp16 or bf16", + ) if p.cpg != 1 or p.kpg != 1: return False, f"cpg and kpg must both be 1 (got cpg={p.cpg}, kpg={p.kpg})" return True, "ok" @@ -4660,9 +4704,10 @@ def build_direct_depthwise( b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = THREADS - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - Bp = b.param("B", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + io_type = _io_type(p.dtype) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + Bp = b.param("B", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) B_bytes = b.param("B_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -4749,7 +4794,11 @@ def build_direct_depthwise( c=c0, ) safe_w_off = b.select(ch_in_range, b.mul(w_off, c_half_bytes), oob_sentinel) - w_h = b.buffer_load_f16(b_rsrc, safe_w_off, c0) + w_h = ( + b.buffer_load_bf16(b_rsrc, safe_w_off, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(b_rsrc, safe_w_off, c0) + ) w_f32 = b.select(ch_in_range, b.cast_to_f32(w_h), zero_f32) row.append(w_f32) weights_f32.append(row) @@ -4784,7 +4833,11 @@ def build_direct_depthwise( safe_off = b.select( load_ok, b.mul(a_off, c_half_bytes), oob_sentinel ) - a_h = b.buffer_load_f16(a_rsrc, safe_off, c0) + a_h = ( + b.buffer_load_bf16(a_rsrc, safe_off, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(a_rsrc, safe_off, c0) + ) a_f32 = b.select(load_ok, b.cast_to_f32(a_h), zero_f32) for r_const in range(p.KH): p_idx = (y - r_const + p.KH) % p.KH @@ -4807,9 +4860,14 @@ def build_direct_depthwise( safe_d = b.select( out_q_ok, b.mul(d_off, c_half_bytes), oob_sentinel ) - b.buffer_store_f16( - d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc[w_out][P_FLUSH]) - ) + if p.dtype == "bf16": + b.buffer_store_bf16( + d_rsrc, safe_d, c0, b.trunc_f32_to_bf16(acc[w_out][P_FLUSH]) + ) + else: + b.buffer_store_f16( + d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc[w_out][P_FLUSH]) + ) for w_out in range(BLOCK_W): acc[w_out][P_FLUSH] = zero_f32 @@ -4854,7 +4912,11 @@ def build_direct_depthwise( safe_off = b.select( ok, b.mul(a_off, c_half_bytes), oob_sentinel ) - a_h = b.buffer_load_f16(a_rsrc, safe_off, c0) + a_h = ( + b.buffer_load_bf16(a_rsrc, safe_off, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(a_rsrc, safe_off, c0) + ) a_f32 = b.select(ok, b.cast_to_f32(a_h), zero_f32) for r_const in range(p.KH): p_idx = (j - r_const + p.KH) % p.KH # STATIC slot @@ -4888,7 +4950,14 @@ def build_direct_depthwise( safe_d = b.select( store_ok, b.mul(d_off, c_half_bytes), oob_sentinel ) - b.buffer_store_f16(d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc_val)) + if p.dtype == "bf16": + b.buffer_store_bf16( + d_rsrc, safe_d, c0, b.trunc_f32_to_bf16(acc_val) + ) + else: + b.buffer_store_f16( + d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc_val) + ) for w_out in range(BLOCK_W): new_accs[P_FLUSH_j * BLOCK_W + w_out] = zero_f32 @@ -4945,13 +5014,18 @@ def kernel_name(self) -> str: from rocke.helpers.spec import kernel_name_join p = self.problem - return kernel_name_join(self.name, p.short(), f"bwv{self.block_waves}") + return kernel_name_join( + self.name, + p.short(), + f"bwv{self.block_waves}", + flags={"bf16": p.dtype == "bf16"}, + ) def validate(self) -> None: p = self.problem - if p.dtype != "fp16": + if p.dtype not in ("fp16", "bf16"): raise ValueError( - f"DirectDepthwiseSpatialSpec: bf16 not supported; depthwise kernels are fp16-only" + f"DirectDepthwiseSpatialSpec: unsupported dtype {p.dtype!r}; expected fp16 or bf16" ) if p.cpg != 1 or p.kpg != 1: raise ValueError( @@ -4981,6 +5055,11 @@ def is_valid_depthwise_spatial_spec( return False, str(e) p = spec.problem + if p.dtype not in ("fp16", "bf16"): + return ( + False, + f"DirectDepthwiseSpatialSpec: unsupported dtype {p.dtype!r}; expected fp16 or bf16", + ) if p.cpg != 1 or p.kpg != 1: return False, f"cpg and kpg must both be 1 (got cpg={p.cpg}, kpg={p.kpg})" if p.groups > spec.wave_size: @@ -5020,9 +5099,10 @@ def build_direct_depthwise_spatial( b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = THREADS - A = b.param("A", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - Bp = b.param("B", PtrType(F16, "global"), noalias=True, readonly=True, align=16) - D = b.param("D", PtrType(F16, "global"), noalias=True, writeonly=True, align=16) + io_type = _io_type(p.dtype) + A = b.param("A", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + Bp = b.param("B", PtrType(io_type, "global"), noalias=True, readonly=True, align=16) + D = b.param("D", PtrType(io_type, "global"), noalias=True, writeonly=True, align=16) A_bytes = b.param("A_bytes", I32) B_bytes = b.param("B_bytes", I32) D_bytes = b.param("D_bytes", I32) @@ -5088,7 +5168,11 @@ def build_direct_depthwise_spatial( b, k=ch, r=b.const_i32(r_const), s=b.const_i32(s_const), c=c0 ) safe_w = b.select(w_valid, b.mul(w_off, c_half_bytes), oob_sentinel) - w_h = b.buffer_load_f16(b_rsrc, safe_w, c0) + w_h = ( + b.buffer_load_bf16(b_rsrc, safe_w, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(b_rsrc, safe_w, c0) + ) row.append(b.select(w_valid, b.cast_to_f32(w_h), zero_f32)) weights_f32.append(row) @@ -5103,7 +5187,11 @@ def build_direct_depthwise_spatial( ) ok = b.land(valid, q_ok) safe_off = b.select(ok, b.mul(a_off, c_half_bytes), oob_sentinel) - a_h = b.buffer_load_f16(a_rsrc, safe_off, c0) + a_h = ( + b.buffer_load_bf16(a_rsrc, safe_off, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(a_rsrc, safe_off, c0) + ) a_f32 = b.select(ok, b.cast_to_f32(a_h), zero_f32) for r_const in range(p.KH): p_idx = (y - r_const + p.KH) % p.KH @@ -5118,9 +5206,14 @@ def build_direct_depthwise_spatial( b, n=n, h=b.const_i32(ho_row), w=q_out, k=ch ) safe_d = b.select(q_ok, b.mul(d_off, c_half_bytes), oob_sentinel) - b.buffer_store_f16( - d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc[P_FLUSH]) - ) + if p.dtype == "bf16": + b.buffer_store_bf16( + d_rsrc, safe_d, c0, b.trunc_f32_to_bf16(acc[P_FLUSH]) + ) + else: + b.buffer_store_f16( + d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc[P_FLUSH]) + ) acc[P_FLUSH] = zero_f32 else: @@ -5151,7 +5244,11 @@ def build_direct_depthwise_spatial( ) ok = b.land(b.land(valid, j_valid), q_ok) safe_off = b.select(ok, b.mul(a_off, c_half_bytes), oob_sentinel) - a_h = b.buffer_load_f16(a_rsrc, safe_off, c0) + a_h = ( + b.buffer_load_bf16(a_rsrc, safe_off, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(a_rsrc, safe_off, c0) + ) a_f32 = b.select(ok, b.cast_to_f32(a_h), zero_f32) for r_const in range(p.KH): p_idx = (j - r_const + p.KH) % p.KH # STATIC @@ -5179,7 +5276,12 @@ def build_direct_depthwise_spatial( acc_val = new_accs[P_FLUSH_j] # STATIC index d_off, _ = d_desc.offset(b, n=n, h=ho_row_j, w=q_out, k=ch) safe_d = b.select(store_ok, b.mul(d_off, c_half_bytes), oob_sentinel) - b.buffer_store_f16(d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc_val)) + if p.dtype == "bf16": + b.buffer_store_bf16( + d_rsrc, safe_d, c0, b.trunc_f32_to_bf16(acc_val) + ) + else: + b.buffer_store_f16(d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc_val)) new_accs[P_FLUSH_j] = zero_f32 # unconditional static reset diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c index 8bcd48da3d0a..0acbf56925e2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c @@ -338,6 +338,100 @@ static int make_cfg(int idx, *kind = KIND_8C; *arch = "gfx950"; return 0; + case 19: + /* dgrad bf16: exercises bf16 I/O on the scalar-FMA grouped dgrad path */ + p.N = 2; + p.H = 8; + p.W = 8; + p.groups = 8; + p.cpg = 16; + p.kpg = 16; + p.dtype = "bf16"; + *sdgrad = rocke_direct_conv_dgrad_spec_default(); + sdgrad->problem = p; + sdgrad->block_q = 16; + sdgrad->block_groups = 8; + *kind = KIND_DGRAD; + *arch = "gfx950"; + return 0; + case 20: + /* 16c bf16 with fold_k32=True: pins the non-default fold_k32 path under bf16 */ + p.N = 32; + p.H = 200; + p.W = 200; + p.groups = 16; + p.cpg = 16; + p.kpg = 16; + p.dtype = "bf16"; + *s16 = rocke_direct_conv_16c_spec_default(); + s16->problem = p; + s16->block_groups = 8; + s16->fold_k32 = true; + *kind = KIND_16C; + *arch = "gfx950"; + return 0; + case 21: + /* 32c bf16: exercises bf16 I/O on the 32c MFMA path */ + p.N = 32; + p.H = 200; + p.W = 200; + p.groups = 32; + p.cpg = 32; + p.kpg = 32; + p.dtype = "bf16"; + *s32 = rocke_direct_conv_32c_spec_default(); + s32->problem = p; + s32->block_groups = 8; + *kind = KIND_32C; + *arch = "gfx950"; + return 0; + case 22: + /* depthwise forward bf16: exercises bf16 I/O on the scalar-FMA depthwise path */ + p.N = 2; + p.H = 14; + p.W = 14; + p.groups = 64; + p.cpg = 1; + p.kpg = 1; + p.dtype = "bf16"; + *sdw = rocke_direct_depthwise_spec_default(); + sdw->problem = p; + sdw->block_w = 8; + sdw->block_waves = 1; + *kind = KIND_DW; + *arch = "gfx950"; + return 0; + case 23: + /* depthwise spatial bf16: exercises bf16 I/O on the small-group spatial path */ + p.N = 2; + p.H = 14; + p.W = 14; + p.groups = 16; + p.cpg = 1; + p.kpg = 1; + p.dtype = "bf16"; + *ssp = rocke_direct_depthwise_spatial_spec_default(); + ssp->problem = p; + ssp->block_waves = 1; + *kind = KIND_SPATIAL; + *arch = "gfx950"; + return 0; + case 24: + /* depthwise dgrad bf16: exercises bf16 I/O on the scalar-FMA depthwise dgrad path */ + p.N = 2; + p.H = 14; + p.W = 14; + p.groups = 64; + p.cpg = 1; + p.kpg = 1; + p.dtype = "bf16"; + *sdw_dgrad = rocke_direct_depthwise_dgrad_spec_default(); + sdw_dgrad->problem = p; + sdw_dgrad->block_w = 8; + sdw_dgrad->block_waves = 1; + *kind = KIND_DW_DGRAD; + *arch = "gfx950"; + return 0; default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py index 11f6e66351ee..e348579161bc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py @@ -235,6 +235,126 @@ def _spec(idx: int): DirectConv8cSpec(problem=p, block_q=16, block_groups=8, double_buffer=True), "gfx950", ) + if idx == 19: + # dgrad bf16: exercises bf16 I/O on the scalar-FMA grouped dgrad path + p = DirectConvProblem( + N=2, + H=8, + W=8, + groups=8, + cpg=16, + kpg=16, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + ) + return ( + "dgrad", + DirectConvDgradSpec(problem=p, block_q=16, block_groups=8), + "gfx950", + ) + if idx == 20: + # 16c bf16 with fold_k32=True: pins the non-default fold_k32 path under bf16 + p = DirectConvProblem( + N=32, + H=200, + W=200, + groups=16, + cpg=16, + kpg=16, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + ) + return ( + "16c", + DirectConv16cSpec(problem=p, block_groups=8, fold_k32=True), + "gfx950", + ) + if idx == 21: + # 32c bf16: exercises bf16 I/O on the 32c MFMA path + p = DirectConvProblem( + N=32, + H=200, + W=200, + groups=32, + cpg=32, + kpg=32, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + ) + return ( + "32c", + DirectConv32cSpec(problem=p, block_groups=8), + "gfx950", + ) + if idx == 22: + # depthwise forward bf16: exercises bf16 I/O on the scalar-FMA depthwise path + p = DirectConvProblem( + N=2, + H=14, + W=14, + groups=64, + cpg=1, + kpg=1, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + ) + return ( + "dw", + DirectDepthwiseSpec(problem=p, block_w=8, block_waves=1), + "gfx950", + ) + if idx == 23: + # depthwise spatial bf16: exercises bf16 I/O on the small-group spatial path + p = DirectConvProblem( + N=2, + H=14, + W=14, + groups=16, + cpg=1, + kpg=1, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + ) + return ( + "spatial", + DirectDepthwiseSpatialSpec(problem=p, block_waves=1), + "gfx950", + ) + if idx == 24: + # depthwise dgrad bf16: exercises bf16 I/O on the scalar-FMA depthwise dgrad path + p = DirectConvProblem( + N=2, + H=14, + W=14, + groups=64, + cpg=1, + kpg=1, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + ) + return ( + "dw_dgrad", + DirectDepthwiseDgradSpec(problem=p, block_w=8, block_waves=1), + "gfx950", + ) raise SystemExit(f"unknown config index {idx}") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h index 8898ff929c31..b56a0cae28ff 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h @@ -366,6 +366,8 @@ bool rocke_direct_depthwise_spatial_is_valid_spec(const rocke_direct_depthwise_s const char* arch, char* reason, size_t reason_cap); +rocke_status_t rocke_direct_depthwise_spatial_validate( + const rocke_direct_depthwise_spatial_spec_t* spec, char* reason, size_t reason_cap); /* ===================================================================== * * DirectConvDgradSpec (grouped dgrad: scalar FMA, any cpg/kpg, stride>=1) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h index e9624c61cf96..095f26b99816 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h @@ -596,6 +596,7 @@ typedef struct rocke_dconv_dw_ctx int Ho; /* output height: (H + 2*PAD - KH) / stride + 1 */ int Wo; /* output width: (W + 2*PAD - KW) / stride + 1 */ int c_stride_dw; /* p.stride (kept as int for flush-loop modulo) */ + int is_bf16; /* 1 when p.dtype == "bf16", 0 otherwise */ rocke_value_t* A; rocke_value_t* Bp; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp index f1fe21618ad6..fab8a19cbcb2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp @@ -28,7 +28,9 @@ #include #include #include +#include +#include "rocke/helper_rocke.helpers.io.h" #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/instance_conv_direct_grouped.h" #include "rocke/instance_conv_direct_grouped_internal.h" @@ -74,8 +76,10 @@ bool rocke_dconv_dw_prologue(rocke_dconv_dw_ctx_t* ctx) rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", ctx->THREADS); + ctx->is_bf16 = (ctx->p.dtype && strcmp(ctx->p.dtype, "bf16") == 0) ? 1 : 0; { - const rocke_type_t* f16ptr = rocke_ptr_type(b, rocke_f16(), "global"); + const rocke_type_t* io_type = rocke_b_io_ir_type(b, ctx->p.dtype ? ctx->p.dtype : "fp16"); + const rocke_type_t* ioptr = rocke_ptr_type(b, io_type, "global"); rocke_param_opts_t ro; rocke_param_opts_t wo; rocke_param_opts_t none; @@ -87,8 +91,8 @@ bool rocke_dconv_dw_prologue(rocke_dconv_dw_ctx_t* ctx) ro.readonly_set = true; ro.align = 16; ro.align_set = true; - ctx->A = rocke_b_param(b, "A", f16ptr, &ro); - ctx->Bp = rocke_b_param(b, "B", f16ptr, &ro); + ctx->A = rocke_b_param(b, "A", ioptr, &ro); + ctx->Bp = rocke_b_param(b, "B", ioptr, &ro); wo = (rocke_param_opts_t){0}; wo.noalias = true; @@ -97,7 +101,7 @@ bool rocke_dconv_dw_prologue(rocke_dconv_dw_ctx_t* ctx) wo.writeonly_set = true; wo.align = 16; wo.align_set = true; - ctx->D = rocke_b_param(b, "D", f16ptr, &wo); + ctx->D = rocke_b_param(b, "D", ioptr, &wo); none = (rocke_param_opts_t){0}; ctx->A_bytes = rocke_b_param(b, "A_bytes", rocke_i32(), &none); @@ -253,7 +257,8 @@ void rocke_dconv_dw_load_weights(rocke_dconv_dw_ctx_t* ctx) ctx->ch_in_range, rocke_b_mul(b, w_off, ctx->c_half_bytes), ctx->oob_sentinel); - w_h = rocke_b_buffer_load_f16(b, ctx->b_rsrc, safe_w, ctx->c0); + w_h = ctx->is_bf16 ? rocke_b_buffer_load_bf16(b, ctx->b_rsrc, safe_w, ctx->c0) + : rocke_b_buffer_load_f16(b, ctx->b_rsrc, safe_w, ctx->c0); ctx->weights_f32[r_const][s_const] = rocke_b_select( b, ctx->ch_in_range, rocke_b_cast_to_f32(b, w_h), ctx->zero_f32); } @@ -333,7 +338,8 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) load_ok = rocke_b_land(b, valid, ctx->ch_in_range); safe_off = rocke_b_select( b, load_ok, rocke_b_mul(b, a_off, ctx->c_half_bytes), ctx->oob_sentinel); - a_h = rocke_b_buffer_load_f16(b, ctx->a_rsrc, safe_off, ctx->c0); + a_h = ctx->is_bf16 ? rocke_b_buffer_load_bf16(b, ctx->a_rsrc, safe_off, ctx->c0) + : rocke_b_buffer_load_f16(b, ctx->a_rsrc, safe_off, ctx->c0); a_f32 = rocke_b_select(b, load_ok, rocke_b_cast_to_f32(b, a_h), ctx->zero_f32); for(r_const = 0; r_const < KH; ++r_const) @@ -379,8 +385,13 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) out_q_ok, rocke_b_mul(b, d_off, ctx->c_half_bytes), ctx->oob_sentinel); - acc_h = rocke_b_trunc_f32_to_f16(b, ctx->acc[w_out][P_FLUSH]); - rocke_b_buffer_store_f16(b, ctx->d_rsrc, safe_d, ctx->c0, acc_h); + acc_h = ctx->is_bf16 + ? rocke_b_trunc_f32_to_bf16(b, ctx->acc[w_out][P_FLUSH]) + : rocke_b_trunc_f32_to_f16(b, ctx->acc[w_out][P_FLUSH]); + if(ctx->is_bf16) + rocke_b_buffer_store_bf16(b, ctx->d_rsrc, safe_d, ctx->c0, acc_h); + else + rocke_b_buffer_store_f16(b, ctx->d_rsrc, safe_d, ctx->c0, acc_h); } } } @@ -500,7 +511,8 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) ok = rocke_b_land(b, rocke_b_land(b, valid, j_valid), ctx->ch_in_range); safe_off = rocke_b_select( b, ok, rocke_b_mul(b, a_off, ctx->c_half_bytes), ctx->oob_sentinel); - a_h = rocke_b_buffer_load_f16(b, ctx->a_rsrc, safe_off, ctx->c0); + a_h = ctx->is_bf16 ? rocke_b_buffer_load_bf16(b, ctx->a_rsrc, safe_off, ctx->c0) + : rocke_b_buffer_load_f16(b, ctx->a_rsrc, safe_off, ctx->c0); a_f32 = rocke_b_select(b, ok, rocke_b_cast_to_f32(b, a_h), ctx->zero_f32); for(r_const = 0; r_const < KH; ++r_const) @@ -564,8 +576,12 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) safe_d = rocke_b_select( b, store_ok, rocke_b_mul(b, d_off, ctx->c_half_bytes), ctx->oob_sentinel); - rocke_b_buffer_store_f16( - b, ctx->d_rsrc, safe_d, ctx->c0, rocke_b_trunc_f32_to_f16(b, acc_val)); + if(ctx->is_bf16) + rocke_b_buffer_store_bf16( + b, ctx->d_rsrc, safe_d, ctx->c0, rocke_b_trunc_f32_to_bf16(b, acc_val)); + else + rocke_b_buffer_store_f16( + b, ctx->d_rsrc, safe_d, ctx->c0, rocke_b_trunc_f32_to_f16(b, acc_val)); /* Unconditional static reset */ new_accs[P_FLUSH_j * BLOCK_W + w_out] = ctx->zero_f32; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp index b1cad4e615c5..98ac0678b07d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp @@ -42,6 +42,7 @@ #include #include +#include "rocke/helper_rocke.helpers.io.h" #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/instance_conv_direct_grouped.h" #include "rocke/instance_conv_direct_grouped_internal.h" @@ -79,6 +80,7 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( int n_iters; int _use_unroll; int total_c, total_k; + int is_bf16; rocke_value_t* A; rocke_value_t* Bp; @@ -152,12 +154,14 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( n_iters = p->H + p->KH - 1; _use_unroll = ((long)n_iters * p->KH * p->KW) <= 20000; + is_bf16 = (p->dtype && strcmp(p->dtype, "bf16") == 0) ? 1 : 0; rocke_attr_set_int(bld, &bld->kernel->attrs, "max_workgroup_size", THREADS); /* ---- parameters ---- */ { - const rocke_type_t* f16ptr = rocke_ptr_type(bld, rocke_f16(), "global"); + const rocke_type_t* io_type = rocke_b_io_ir_type(bld, p->dtype ? p->dtype : "fp16"); + const rocke_type_t* ioptr = rocke_ptr_type(bld, io_type, "global"); rocke_param_opts_t ro, wo_opt, none; ro = (rocke_param_opts_t){0}; @@ -165,15 +169,15 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( ro.readonly = ro.readonly_set = true; ro.align = 16; ro.align_set = true; - A = rocke_b_param(bld, "A", f16ptr, &ro); - Bp = rocke_b_param(bld, "B", f16ptr, &ro); + A = rocke_b_param(bld, "A", ioptr, &ro); + Bp = rocke_b_param(bld, "B", ioptr, &ro); wo_opt = (rocke_param_opts_t){0}; wo_opt.noalias = wo_opt.noalias_set = true; wo_opt.writeonly = wo_opt.writeonly_set = true; wo_opt.align = 16; wo_opt.align_set = true; - D = rocke_b_param(bld, "D", f16ptr, &wo_opt); + D = rocke_b_param(bld, "D", ioptr, &wo_opt); none = (rocke_param_opts_t){0}; A_bytes = rocke_b_param(bld, "A_bytes", rocke_i32(), &none); @@ -296,7 +300,8 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( safe_w = rocke_b_select( bld, w_valid, rocke_b_mul(bld, w_off, c_half_bytes), oob_sentinel); - w_h = rocke_b_buffer_load_f16(bld, b_rsrc, safe_w, c0); + w_h = is_bf16 ? rocke_b_buffer_load_bf16(bld, b_rsrc, safe_w, c0) + : rocke_b_buffer_load_f16(bld, b_rsrc, safe_w, c0); weights_f32[r_const][s_const] = rocke_b_select(bld, w_valid, rocke_b_cast_to_f32(bld, w_h), zero_f32); } @@ -349,7 +354,8 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( ok = rocke_b_land(bld, valid, q_ok); safe_off = rocke_b_select(bld, ok, rocke_b_mul(bld, a_off, c_half_bytes), oob_sentinel); - a_h = rocke_b_buffer_load_f16(bld, a_rsrc, safe_off, c0); + a_h = is_bf16 ? rocke_b_buffer_load_bf16(bld, a_rsrc, safe_off, c0) + : rocke_b_buffer_load_f16(bld, a_rsrc, safe_off, c0); a_f32 = rocke_b_select(bld, ok, rocke_b_cast_to_f32(bld, a_h), zero_f32); for(r_const = 0; r_const < p->KH; ++r_const) @@ -383,8 +389,12 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( safe_d = rocke_b_select( bld, q_ok, rocke_b_mul(bld, d_off, c_half_bytes), oob_sentinel); - rocke_b_buffer_store_f16( - bld, d_rsrc, safe_d, c0, rocke_b_trunc_f32_to_f16(bld, acc[P_FLUSH])); + if(is_bf16) + rocke_b_buffer_store_bf16( + bld, d_rsrc, safe_d, c0, rocke_b_trunc_f32_to_bf16(bld, acc[P_FLUSH])); + else + rocke_b_buffer_store_f16( + bld, d_rsrc, safe_d, c0, rocke_b_trunc_f32_to_f16(bld, acc[P_FLUSH])); } } @@ -490,7 +500,8 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( ok = rocke_b_land(bld, rocke_b_land(bld, valid, j_valid), q_ok); safe_off = rocke_b_select(bld, ok, rocke_b_mul(bld, a_off, c_half_bytes), oob_sentinel); - a_h = rocke_b_buffer_load_f16(bld, a_rsrc, safe_off, c0); + a_h = is_bf16 ? rocke_b_buffer_load_bf16(bld, a_rsrc, safe_off, c0) + : rocke_b_buffer_load_f16(bld, a_rsrc, safe_off, c0); a_f32 = rocke_b_select(bld, ok, rocke_b_cast_to_f32(bld, a_h), zero_f32); for(r_const = 0; r_const < KH; ++r_const) @@ -542,8 +553,12 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( safe_d = rocke_b_select( bld, store_ok, rocke_b_mul(bld, d_off, c_half_bytes), oob_sentinel); - rocke_b_buffer_store_f16( - bld, d_rsrc, safe_d, c0, rocke_b_trunc_f32_to_f16(bld, acc_val)); + if(is_bf16) + rocke_b_buffer_store_bf16( + bld, d_rsrc, safe_d, c0, rocke_b_trunc_f32_to_bf16(bld, acc_val)); + else + rocke_b_buffer_store_f16( + bld, d_rsrc, safe_d, c0, rocke_b_trunc_f32_to_f16(bld, acc_val)); /* Unconditional static reset */ new_accs[P_FLUSH_j] = zero_f32; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_dgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_dgrad.cpp index 21674cbaaf57..93cc4140059f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_dgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_dgrad.cpp @@ -16,7 +16,9 @@ #include #include #include +#include +#include "rocke/helper_rocke.helpers.io.h" #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/instance_conv_direct_grouped.h" #include "rocke/instance_conv_direct_grouped_internal.h" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp index 503eb1ec47b6..106ad3d87621 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp @@ -1162,13 +1162,21 @@ rocke_status_t rocke_direct_depthwise_kernel_name(const rocke_direct_depthwise_s { return ROCKE_ERR_VALUE; } - /* kernel_name_join(name, p.short(), f"bw{block_w}", f"bw{block_waves}wv") */ + /* kernel_name_join(name, p.short(), f"bw{block_w}", f"bw{block_waves}wv", + * flags={"bf16": dtype=="bf16"}) */ snprintf(bw_buf, sizeof(bw_buf), "bw%d", spec->block_w); snprintf(bwv_buf, sizeof(bwv_buf), "bw%dwv", spec->block_waves); parts[0] = prob_short; parts[1] = bw_buf; parts[2] = bwv_buf; - return rocke_kernel_name_join(spec->name, parts, 3, NULL, NULL, 0, out, out_cap, NULL); + { + const char* flag_names[1] = {"bf16"}; + int flag_on[1]; + const char* dt = spec->problem.dtype ? spec->problem.dtype : "fp16"; + flag_on[0] = (strcmp(dt, "bf16") == 0) ? 1 : 0; + return rocke_kernel_name_join( + spec->name, parts, 3, flag_names, flag_on, 1, out, out_cap, NULL); + } } rocke_status_t rocke_direct_depthwise_validate(const rocke_direct_depthwise_spec_t* spec, @@ -1182,6 +1190,19 @@ rocke_status_t rocke_direct_depthwise_validate(const rocke_direct_depthwise_spec return ROCKE_ERR_VALUE; } p = &spec->problem; + /* if p.dtype not in ("fp16", "bf16"): raise ValueError(...) */ + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) + { + if(reason && reason_cap > 0) + snprintf(reason, + reason_cap, + "DirectDepthwiseSpec: unsupported dtype '%s'; expected fp16 or bf16", + dt); + return ROCKE_ERR_VALUE; + } + } if(p->cpg != 1 || p->kpg != 1) { if(reason && reason_cap > 0) @@ -1302,11 +1323,18 @@ rocke_status_t rocke_direct_depthwise_spatial_kernel_name( { return ROCKE_ERR_VALUE; } - /* kernel_name_join(name, p.short(), f"bwv{block_waves}") */ + /* kernel_name_join(name, p.short(), f"bwv{block_waves}", flags={"bf16": dtype=="bf16"}) */ snprintf(bwv_buf, sizeof(bwv_buf), "bwv%d", spec->block_waves); parts[0] = prob_short; parts[1] = bwv_buf; - return rocke_kernel_name_join(spec->name, parts, 2, NULL, NULL, 0, out, out_cap, NULL); + { + const char* flag_names[1] = {"bf16"}; + int flag_on[1]; + const char* dt = spec->problem.dtype ? spec->problem.dtype : "fp16"; + flag_on[0] = (strcmp(dt, "bf16") == 0) ? 1 : 0; + return rocke_kernel_name_join( + spec->name, parts, 2, flag_names, flag_on, 1, out, out_cap, NULL); + } } bool rocke_direct_depthwise_spatial_is_valid_spec(const rocke_direct_depthwise_spatial_spec_t* spec, @@ -1374,6 +1402,45 @@ bool rocke_direct_depthwise_spatial_is_valid_spec(const rocke_direct_depthwise_s return true; } +rocke_status_t rocke_direct_depthwise_spatial_validate( + const rocke_direct_depthwise_spatial_spec_t* spec, char* reason, size_t reason_cap) +{ + const rocke_direct_conv_problem_t* p; + if(spec == NULL) + return ROCKE_ERR_VALUE; + p = &spec->problem; + /* if p.dtype not in ("fp16", "bf16"): raise ValueError(...) */ + { + const char* dt = p->dtype ? p->dtype : "fp16"; + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) + { + if(reason && reason_cap > 0) + snprintf(reason, + reason_cap, + "DirectDepthwiseSpatialSpec: unsupported dtype '%s'; " + "expected fp16 or bf16", + dt); + return ROCKE_ERR_VALUE; + } + } + if(p->cpg != 1 || p->kpg != 1) + { + if(reason && reason_cap > 0) + snprintf(reason, + reason_cap, + "DirectDepthwiseSpatialSpec requires cpg=kpg=1 (got cpg=%d, kpg=%d)", + p->cpg, + p->kpg); + return ROCKE_ERR_VALUE; + } + if(reason && reason_cap > 0) + { + strncpy(reason, "ok", reason_cap); + reason[reason_cap - 1] = '\0'; + } + return ROCKE_OK; +} + /* ===================================================================== * * DirectConvDgradSpec (grouped dgrad — scalar FMA) * ===================================================================== */ @@ -1595,13 +1662,21 @@ rocke_status_t rocke_direct_depthwise_dgrad_kernel_name( { return ROCKE_ERR_VALUE; } - /* kernel_name_join(name, p.short(), f"bw{block_w}", f"bw{block_waves}wv") */ + /* kernel_name_join(name, p.short(), f"bw{block_w}", f"bw{block_waves}wv", + * flags={"bf16": dtype=="bf16"}) */ snprintf(bw_buf, sizeof(bw_buf), "bw%d", spec->block_w); snprintf(bwv_buf, sizeof(bwv_buf), "bw%dwv", spec->block_waves); parts[0] = prob_short; parts[1] = bw_buf; parts[2] = bwv_buf; - return rocke_kernel_name_join(spec->name, parts, 3, NULL, NULL, 0, out, out_cap, NULL); + { + const char* flag_names[1] = {"bf16"}; + int flag_on[1]; + const char* dt = spec->problem.dtype ? spec->problem.dtype : "fp16"; + flag_on[0] = (strcmp(dt, "bf16") == 0) ? 1 : 0; + return rocke_kernel_name_join( + spec->name, parts, 3, flag_names, flag_on, 1, out, out_cap, NULL); + } } rocke_status_t rocke_direct_depthwise_dgrad_validate( @@ -1613,16 +1688,17 @@ rocke_status_t rocke_direct_depthwise_dgrad_validate( return ROCKE_ERR_VALUE; } p = &spec->problem; - /* if p.dtype != "fp16": raise ValueError(...) — depthwise kernels are fp16-only */ + /* if p.dtype not in ("fp16", "bf16"): raise ValueError(...) */ { const char* dt = p->dtype ? p->dtype : "fp16"; - if(strcmp(dt, "fp16") != 0) + if(strcmp(dt, "fp16") != 0 && strcmp(dt, "bf16") != 0) { if(reason && reason_cap > 0) snprintf(reason, reason_cap, - "DirectDepthwiseDgradSpec: bf16 not supported; " - "depthwise kernels are fp16-only"); + "DirectDepthwiseDgradSpec: unsupported dtype '%s'; " + "expected fp16 or bf16", + dt); return ROCKE_ERR_VALUE; } }