diff --git a/kernels/conv/conv3d_implicit.py b/kernels/conv/conv3d_implicit.py new file mode 100644 index 000000000..1bb9abbb0 --- /dev/null +++ b/kernels/conv/conv3d_implicit.py @@ -0,0 +1,765 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Double-buffered implicit-GEMM conv3d (BF16). + +x: (N, C, D, H, W) bf16 NCDHW, weight: (K, C, T, R, S) bf16 KCTRS. +Returns (N, K, Do, Ho, Wo) bf16. Supports stride, padding, bias, and split-K. +""" + +import functools +import os +import weakref + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, const_expr, range_constexpr, rocdl +from flydsl.expr.typing import T +from kernels.common import buffer_ops +from kernels.common.mem_ops import buffer_atomic_add + +TILE_K = 32 +STAGES = 2 +WARP_SIZE = 64 + +MFMA_M = 16 +MFMA_N = 16 +MFMA_A_VALUES = 8 +MFMA_B_VALUES = 8 +MFMA_C_VALUES = 4 + +LDG_VEC = 8 + +BF16_BYTES = 2 + +DEFAULT_TILE = (128, 128, 2, 4) + + +def _autotune_enabled(): + return os.environ.get("FLYDSL_CONV3D_AUTOTUNE", "0").lower() in ("1", "true", "yes") + + +_WEIGHT_CACHE = {} + + +def _prep_weight(w, k, kt, kh, kw, c): + key = id(w) + ent = _WEIGHT_CACHE.get(key) + if ent is not None and ent[0]() is w: + return ent[1] + wk = w.permute(0, 2, 3, 4, 1).contiguous().reshape(k, kt * kh * kw * c) + _WEIGHT_CACHE[key] = (weakref.ref(w), wk) + return wk + + +TR_TILE = 64 +TR_VEC = 8 +TR_THREADS = 256 +_TR_VPL = TR_TILE // TR_VEC +_TR_ITERS = (TR_TILE * TR_TILE) // (TR_VEC * TR_THREADS) +_TR_PAD = 8 +_TR_LDS_S = TR_TILE + _TR_PAD + + +@functools.lru_cache(maxsize=64) +def compile_transpose_ncdhw_ndhwc(n, c, s): + """Transpose flat (N, C, S) -> (N, S, C) (S == T*H*W). Requires c%8==0, s%8==0.""" + grid_s = (s + TR_TILE - 1) // TR_TILE + grid_c = (c + TR_TILE - 1) // TR_TILE + elem_ty = fx.BFloat16 + BIG = (n * c * s) > 0x7FFFFFFF + + @flyc.kernel(known_block_size=[TR_THREADS, 1, 1]) + def transpose_kernel(out: fx.Tensor, inp: fx.Tensor): + in_rsrc = buffer_ops.create_buffer_resource(inp) + out_rsrc = buffer_ops.create_buffer_resource(out) + lds_alloc = fx.SharedAllocator(static=False) + lds = lds_alloc.allocate(fx.Array[elem_ty, TR_TILE * _TR_LDS_S, 16]).peek() + + class BF16Ty: + ir_type = elem_ty.ir_type + + tid = fx.thread_idx.x + s0 = fx.block_idx.x * TR_TILE + c0 = fx.block_idx.y * TR_TILE + nb = fx.block_idx.z + if const_expr(BIG): + in_base_elem = fx.Index(nb) * fx.Index(c) * fx.Index(s) + fx.Index(c0) * fx.Index(s) + fx.Index(s0) + in_addr = fx.Int64(buffer_ops.extract_base_index(inp)) + fx.Int64(in_base_elem) * fx.Int64(2) + in_rsrc = buffer_ops.create_buffer_resource_from_addr(in_addr) + out_base_elem = fx.Index(nb) * fx.Index(s) * fx.Index(c) + fx.Index(s0) * fx.Index(c) + fx.Index(c0) + out_addr = fx.Int64(buffer_ops.extract_base_index(out)) + fx.Int64(out_base_elem) * fx.Int64(2) + out_rsrc = buffer_ops.create_buffer_resource_from_addr(out_addr) + else: + in_base = nb * c * s + out_base = nb * s * c + + def lds_store_vec8(elem_offset, value): + base = fx.Int64(fx.ptrtoint(lds.ptr)) + fx.Int64(elem_offset * 2) + ptr = buffer_ops.create_llvm_ptr(base, address_space=3) + llvm.StoreOp(value, ptr, alignment=16) + + def lds_load_scalar(elem_offset): + u8 = fx.recast_iter(fx.Uint8, lds.ptr) + return fx.ptr_load(u8 + fx.Int32(elem_offset * 2), result_type=BF16Ty) + + # Read: coalesced vec8 along contiguous S -> LDS[c_local][s_local]. + for i in range_constexpr(_TR_ITERS): + lin = tid + i * TR_THREADS + rc = lin // _TR_VPL + sv = (lin % _TR_VPL) * TR_VEC + cc = c0 + rc + ss = s0 + sv + valid = (cc < c) & (ss < s) + if const_expr(BIG): + g = fx.Int32(rc * s + sv) + else: + g = fx.Int32(in_base + cc * s + ss) + safe = arith.select(valid, g, fx.Int32(0)) + v = buffer_ops.buffer_load(in_rsrc, safe, vec_width=TR_VEC, dtype=elem_ty) + lds_store_vec8(rc * _TR_LDS_S + sv, v) + + llvm.InlineAsmOp(None, [], "s_waitcnt lgkmcnt(0)\n\ts_barrier", "", has_side_effects=True) + + for i in range_constexpr(_TR_ITERS): + lin = tid + i * TR_THREADS + rs = lin // _TR_VPL + cv = (lin % _TR_VPL) * TR_VEC + ss = s0 + rs + cc = c0 + cv + scalars = [lds_load_scalar((cv + j) * _TR_LDS_S + rs) for j in range_constexpr(TR_VEC)] + vv = fx.Vector.from_elements(scalars, dtype=elem_ty) + valid = (ss < s) & (cc < c) + if valid: + if const_expr(BIG): + go = fx.Int32(rs * c + cv) + else: + go = fx.Int32(out_base + ss * c + cc) + buffer_ops.buffer_store(vv, out_rsrc, go) + + @flyc.jit + def launch_transpose(out: fx.Tensor, inp: fx.Tensor, stream: fx.Stream = fx.Stream(None)): + transpose_kernel(out, inp).launch( + grid=(grid_s, grid_c, n), + block=(TR_THREADS, 1, 1), + stream=stream, + ) + + return launch_transpose + + +def _ncdhw_to_ndhwc(x, stream): + """Fast NCDHW->NDHWC via the tiled transpose kernel; falls back to torch.""" + n, c, t, h, w = x.shape + s = t * h * w + if not (x.is_contiguous() and x.dtype == torch.bfloat16 and c % 8 == 0 and s % 8 == 0): + return x.permute(0, 2, 3, 4, 1).contiguous() + out = torch.empty((n, t, h, w, c), device=x.device, dtype=x.dtype) + exe = compile_transpose_ncdhw_ndhwc(n, c, s) + exe(out, x, torch.cuda.current_stream() if stream is None else stream) + return out + + +@functools.lru_cache(maxsize=256) +def compile_conv3d_implicit( + n, c, d, h, w, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias=False, splitk=1, tile=DEFAULT_TILE, wgm=1 +): + TILE_M, TILE_N, WAVE_M, WAVE_N = tile + BLOCK_THREADS = WAVE_M * WAVE_N * WARP_SIZE + # Per-wave MFMA grid (flat acc[mi * MI_N + ni]); WARP_M/N is the per-wave tile span. + MI_M = TILE_M // WAVE_M // MFMA_M + MI_N = TILE_N // WAVE_N // MFMA_N + N_ACC = MI_M * MI_N + WARP_M = MI_M * MFMA_M + WARP_N = MI_N * MFMA_N + BLOCK_VECS = LDG_VEC * BLOCK_THREADS + LDG_A_COUNT = TILE_M * TILE_K // BLOCK_VECS + LDG_B_COUNT = TILE_N * TILE_K // BLOCK_VECS + + assert TILE_K == 32 + assert TILE_M % (WAVE_M * MFMA_M) == 0, f"TILE_M={TILE_M} not divisible by WAVE_M*16" + assert TILE_N % (WAVE_N * MFMA_N) == 0, f"TILE_N={TILE_N} not divisible by WAVE_N*16" + assert (TILE_M * TILE_K) % BLOCK_VECS == 0, f"A tile {TILE_M}x{TILE_K} not a multiple of {BLOCK_VECS} vecs" + assert (TILE_N * TILE_K) % BLOCK_VECS == 0, f"B tile {TILE_N}x{TILE_K} not a multiple of {BLOCK_VECS} vecs" + assert LDG_A_COUNT >= 1 and LDG_B_COUNT >= 1 + assert c % LDG_VEC == 0 + assert BLOCK_THREADS <= 1024, f"BLOCK_THREADS={BLOCK_THREADS} exceeds 1024" + + do = (d + 2 * pt - kt) // st + 1 + ho = (h + 2 * ph - kh) // sh + 1 + wo = (w + 2 * pw - kw) // sw + 1 + dhw = do * ho * wo + hw_o = ho * wo + npq = n * dhw + crs = c * kt * kh * kw + k_tiles = (crs + TILE_K - 1) // TILE_K + + BIG_IN = (n * c * d * h * w) > 0x7FFFFFFF + BIG_OUT = (n * k * do * ho * wo * BF16_BYTES) > 0x7FFFFFFF + + X_BYTES = n * c * d * h * w * BF16_BYTES + W_BYTES = k * c * kt * kh * kw * BF16_BYTES + OOB_SENTINEL_ELEM = 0x7FFFFF80 # *2 = 0xFFFFFF00 bytes (~4.2950 GB), just under 2^32 + OOB_SENTINEL_BYTES = OOB_SENTINEL_ELEM * BF16_BYTES + BIG_IN_NR = 0x80000000 # 2 GB num_records for the rebased BIG_IN resource + assert W_BYTES < OOB_SENTINEL_BYTES, f"weight {W_BYTES}B exceeds limit {OOB_SENTINEL_BYTES}B" + assert X_BYTES < OOB_SENTINEL_BYTES or BIG_IN, f"input {X_BYTES}B exceeds limit" + BIG_IN_N1 = BIG_IN and n == 1 + BIG_IN_NM = BIG_IN and n > 1 + X_SAMPLE_BYTES = c * d * h * w * BF16_BYTES + + n_tail = k % TILE_N != 0 + grid_n = (k + TILE_N - 1) // TILE_N + + splitk = max(1, min(splitk, k_tiles)) + tiles_per_split = k_tiles // splitk + use_splitk = splitk > 1 + + # Software-pipeline depth. 4 stages is optimal across all shapes on gfx950 -- + # even short-K, memory-bound 3x1x1 depends more (not less) on deep prefetch to + # hide DMA latency; a shallower pipeline measured slower (2/3/4-stage A/B). + PIPE_STAGES = 4 + + LDS_A_SIZE = PIPE_STAGES * TILE_M * TILE_K + LDS_B_SIZE = PIPE_STAGES * TILE_N * TILE_K + + grid_m = (npq + TILE_M - 1) // TILE_M + WGM = max(1, int(wgm)) + elem_ty = fx.BFloat16 + mfma_fn = rocdl.mfma_f32_16x16x32_bf16 + temporal_only_fast = ( + kh == 1 + and kw == 1 + and st == 1 + and sh == 1 + and sw == 1 + and ph == 0 + and pw == 0 + and do == d + and ho == h + and wo == w + ) + + @flyc.kernel(known_block_size=[BLOCK_THREADS, 1, 1]) + def conv3d_implicit_kernel(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias: fx.Tensor): + w_rsrc = buffer_ops.create_buffer_resource(weight, num_records_bytes=W_BYTES) + if const_expr(not BIG_IN): + x_rsrc = buffer_ops.create_buffer_resource(x, num_records_bytes=X_BYTES) + y_rsrc = buffer_ops.create_buffer_resource(y) + if const_expr(has_bias): + bias_rsrc = buffer_ops.create_buffer_resource(bias) + + lds_alloc = fx.SharedAllocator(static=False) + a_lds = lds_alloc.allocate(fx.Array[elem_ty, LDS_A_SIZE, 16]).peek() + b_lds = lds_alloc.allocate(fx.Array[elem_ty, LDS_B_SIZE, 16]).peek() + + tid = fx.thread_idx.x + if const_expr(WGM > 1): + pid = fx.Index(fx.block_idx.x) + fx.Index(fx.block_idx.y) * fx.Index(grid_m) + blocks_per_group = fx.Index(WGM * grid_n) + group_id = pid // blocks_per_group + first_m = group_id * fx.Index(WGM) + group_rows = fx.Index(grid_m) - first_m + group_rows = fx.Index(arith.select(group_rows < fx.Index(WGM), group_rows, fx.Index(WGM))) + local = pid % blocks_per_group + m_offset = fx.Index(first_m + (local % group_rows)) * TILE_M + n_offset = fx.Index(local // group_rows) * TILE_N + else: + m_offset = fx.block_idx.x * TILE_M + n_offset = fx.block_idx.y * TILE_N + if const_expr(use_splitk): + k_off = fx.block_idx.z * (tiles_per_split * TILE_K) + else: + k_off = 0 + + if const_expr(BIG_IN_N1): + nbase = m_offset // dhw + ot_base0 = (m_offset % dhw) // hw_o + base_t = ot_base0 - fx.Index(pt) + base_t = arith.select(base_t < fx.Index(0), fx.Index(0), base_t) + x_base_elem = ((nbase * fx.Index(d) + base_t) * fx.Index(h) + fx.Index(0)) * fx.Index(w) * fx.Index(c) + x_addr = fx.Int64(buffer_ops.extract_base_index(x)) + fx.Int64(x_base_elem) * fx.Int64(2) + x_rsrc = buffer_ops.create_buffer_resource_from_addr(x_addr, num_records_bytes=BIG_IN_NR) + if const_expr(BIG_IN_NM): + x_base_addr = fx.Int64(buffer_ops.extract_base_index(x)) + + wid = tid // WARP_SIZE + lane = tid % WARP_SIZE + wave_m = wid // WAVE_N + wave_n = wid % WAVE_N + + lane_m = lane % MFMA_M + lane_n = lane % MFMA_N + lane_k_a = lane // MFMA_M * MFMA_A_VALUES + lane_k_b = lane // MFMA_N * MFMA_B_VALUES + c_m_vec = lane // MFMA_N * MFMA_C_VALUES + c_n = lane % MFMA_N + + Vec = fx.Vector + + class Vec8Ty: + ir_type = Vec.make_type(8, elem_ty) + + acc0 = Vec.filled(MFMA_C_VALUES, 0.0, fx.Float32) + acc = [acc0 for _ in range_constexpr(N_ACC)] + + def barrier(vmcnt=0, lgkmcnt=None): + waits = [] + if vmcnt is not None: + waits.append(f"vmcnt({vmcnt})") + if lgkmcnt is not None: + waits.append(f"lgkmcnt({lgkmcnt})") + pre = ("s_waitcnt " + " ".join(waits) + "\n\t") if waits else "" + llvm.InlineAsmOp(None, [], f"{pre}s_barrier", "", has_side_effects=True) + + def lds_load_vec8(lds_array, elem_offset): + u8_ptr = fx.recast_iter(fx.Uint8, lds_array.ptr) + return fx.ptr_load(u8_ptr + fx.Int32(elem_offset * 2), result_type=Vec8Ty) + + def a_lds_off(stage, row, col): + return (fx.Index(stage) * TILE_M + row) * TILE_K + col + + def b_lds_off(stage, row, col): + return (fx.Index(stage) * TILE_N + row) * TILE_K + col + + def in_range(v, hi): + return (v >= 0) & (v < fx.Index(hi)) + + # ---- Per-thread row decomposition (loop-invariant across K) ---- + _row_dec = [] # per-i tuple of precomputed row terms + for i in range_constexpr(LDG_A_COUNT): + linear = (tid + i * BLOCK_THREADS) * LDG_VEC + local_m = linear // TILE_K + local_k = linear % TILE_K + row = m_offset + local_m + row_valid = row < fx.Index(npq) + if const_expr(temporal_only_fast): + out_t = (row // hw_o) % d + _row_dec.append((local_k, row, row_valid, out_t)) + else: + n_idx = row // dhw + rem = row % dhw + ot = rem // hw_o + rem2 = rem % hw_o + oh = rem2 // wo + ow = rem2 % wo + in_t0 = ot * st - pt + in_h0 = oh * sh - ph + in_w0 = ow * sw - pw + if const_expr(BIG_IN_N1): + di = n_idx - nbase + _row_dec.append((local_k, row_valid, di, in_t0, in_h0, in_w0)) + elif const_expr(BIG_IN_NM): + _row_dec.append((local_k, row_valid, n_idx, in_t0, in_h0, in_w0)) + else: + _row_dec.append((local_k, row_valid, n_idx, in_t0, in_h0, in_w0)) + + SCALAR_K = c % TILE_K == 0 + + # ---- 3D im2col address math ---- + def _a_addr(i, kbase_i, cc_base, ckk_base): + dec = _row_dec[i] + local_k = dec[0] + k_abs = kbase_i + fx.Index(local_k) + if const_expr(SCALAR_K): + cc = cc_base + fx.Index(local_k) + else: + cc = k_abs % c + k_valid = k_abs < fx.Index(crs) + if const_expr(temporal_only_fast): + _, row, row_valid, out_t = dec + kt_i = ckk_base if const_expr(SCALAR_K) else k_abs // c + temporal_delta = kt_i - pt + in_t = out_t + temporal_delta + valid = row_valid & k_valid & in_range(in_t, d) + if const_expr(BIG_IN_N1): + g_off = ((row + temporal_delta * hw_o) - (fx.Index(nbase) * dhw + base_t * hw_o)) * c + cc + else: + g_off = (row + temporal_delta * hw_o) * c + cc + else: + ckk = ckk_base if const_expr(SCALAR_K) else k_abs // c + kw_i = ckk % kw + ckk2 = ckk // kw + kh_i = ckk2 % kh + kt_i = ckk2 // kh + if const_expr(BIG_IN_N1): + _, row_valid, di, in_t0, in_h0, in_w0 = dec + in_t = in_t0 + kt_i + in_h = in_h0 + kh_i + in_w = in_w0 + kw_i + valid = row_valid & k_valid & in_range(in_t, d) & in_range(in_h, h) & in_range(in_w, w) + g_off = (((di * d + (in_t - base_t)) * h + in_h) * w + in_w) * c + cc + elif const_expr(BIG_IN_NM): + _, row_valid, n_idx, in_t0, in_h0, in_w0 = dec + in_t = in_t0 + kt_i + in_h = in_h0 + kh_i + in_w = in_w0 + kw_i + valid = row_valid & k_valid & in_range(in_t, d) & in_range(in_h, h) & in_range(in_w, w) + g_off = ((in_t * h + in_h) * w + in_w) * c + cc + return fx.Int32(g_off), valid, n_idx + else: + _, row_valid, n_idx, in_t0, in_h0, in_w0 = dec + in_t = in_t0 + kt_i + in_h = in_h0 + kh_i + in_w = in_w0 + kw_i + valid = row_valid & k_valid & in_range(in_t, d) & in_range(in_h, h) & in_range(in_w, w) + g_off = (((n_idx * d + in_t) * h + in_h) * w + in_w) * c + cc + return fx.Int32(g_off), valid + + def _b_addr(i, k_base): + linear = (tid + i * BLOCK_THREADS) * LDG_VEC + local_n = linear // TILE_K + local_k = linear % TILE_K + col = n_offset + fx.Index(local_n) + g_off = fx.Int32(col * crs + (fx.Index(k_base) + fx.Index(local_k))) + col_valid = (col < fx.Index(k)) if const_expr(n_tail) else None + return g_off, col_valid + + # ---- global -> LDS DMA copy, masking via OOB routing ---- + DMA_BYTES = LDG_VEC * BF16_BYTES # 16 + OOB_ELEM = fx.Int32(OOB_SENTINEL_ELEM) + + def _lds_dma_ptr(lds_array, stage_tile, i): + off_elems = fx.Index(stage_tile) + (fx.Index(tid) + fx.Index(i * BLOCK_THREADS)) * fx.Index(LDG_VEC) + base_bytes = off_elems * fx.Index(BF16_BYTES) + addr = fx.Int64(fx.ptrtoint(lds_array.ptr)) + fx.Int64(base_bytes) + addr = rocdl.readfirstlane(T.i64, arith.index_cast(T.i64, addr.ir_value())) + return llvm.inttoptr(ir.Type.parse("!llvm.ptr<3>"), addr) + + def _dma_to_lds(rsrc, lds_ptr, voff_elem): + voff_b = (voff_elem * fx.Int32(BF16_BYTES)).ir_value() + rocdl.raw_ptr_buffer_load_lds( + rsrc, + lds_ptr, + arith.constant(DMA_BYTES, type=T.i32), + voff_b, + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + ) + + def _load_a(stage, k_base): + kbase_i = fx.Index(k_base) + cc_base = ckk_base = None + if const_expr(SCALAR_K): + cc_base = kbase_i % c + ckk_base = kbase_i // c + stage_tile = fx.Index(stage) * TILE_M * TILE_K + for i in range_constexpr(LDG_A_COUNT): + if const_expr(BIG_IN_NM): + addr_ret = _a_addr(i, kbase_i, cc_base, ckk_base) + g_off_i, valid, n_idx_i = addr_ret + sample_addr = x_base_addr + fx.Int64(n_idx_i) * fx.Int64(X_SAMPLE_BYTES) + x_rsrc_i = buffer_ops.create_buffer_resource_from_addr(sample_addr, num_records_bytes=BIG_IN_NR) + voff = fx.Int32(arith.select(valid, g_off_i, OOB_ELEM)) + _dma_to_lds(x_rsrc_i, _lds_dma_ptr(a_lds, stage_tile, i), voff) + else: + g_off_i, valid = _a_addr(i, kbase_i, cc_base, ckk_base) + voff = fx.Int32(arith.select(valid, g_off_i, OOB_ELEM)) + _dma_to_lds(x_rsrc, _lds_dma_ptr(a_lds, stage_tile, i), voff) + + def _load_b(stage, k_base): + stage_tile = fx.Index(stage) * TILE_N * TILE_K + for i in range_constexpr(LDG_B_COUNT): + g_off, col_valid = _b_addr(i, k_base) + if const_expr(n_tail): + voff = fx.Int32(arith.select(col_valid, g_off, OOB_ELEM)) + else: + voff = g_off + _dma_to_lds(w_rsrc, _lds_dma_ptr(b_lds, stage_tile, i), voff) + + # ---- single-vec ds_read (LDS -> register), indexed by per-wave MFMA row ---- + def read_a_vec(stage, mi): + a_row = wave_m * WARP_M + mi * MFMA_M + lane_m + return lds_load_vec8(a_lds, a_lds_off(stage, fx.Index(a_row), fx.Index(lane_k_a))) + + def read_b_vec(stage, ni): + b_row = wave_n * WARP_N + ni * MFMA_N + lane_n + return lds_load_vec8(b_lds, b_lds_off(stage, fx.Index(b_row), fx.Index(lane_k_b))) + + def mfma_one(a_frag, b_frag, c_frag): + out = mfma_fn( + T.vec(MFMA_C_VALUES, T.f32), + [a_frag, b_frag, c_frag, 0, 0, 0], + ) + rocdl.sched_mfma(1) + return out + + def read_a_frags(stage): + frags = [read_a_vec(stage, mi) for mi in range_constexpr(MI_M)] + rocdl.sched_dsrd(MI_M) + return frags + + def read_b_frags(stage): + frags = [read_b_vec(stage, ni) for ni in range_constexpr(MI_N)] + rocdl.sched_dsrd(MI_N) + return frags + + def do_compute(acc_values, a_frag_values, b_frag_values): + rocdl.s_setprio(1) + for mi in range_constexpr(MI_M): + for ni in range_constexpr(MI_N): + idx = mi * MI_N + ni + acc_values[idx] = mfma_one(a_frag_values[mi], b_frag_values[ni], acc_values[idx]) + rocdl.s_setprio(0) + return acc_values + + # global->LDS software pipeline + # ---- prologue: fill the pipeline with the first PREFETCH tiles' DMAs ---- + PREFETCH = PIPE_STAGES - 1 + for s in range_constexpr(PREFETCH): + if const_expr(s < tiles_per_split): + _load_a(s, k_off + s * TILE_K) + _load_b(s, k_off + s * TILE_K) + LDG_PER_TILE = LDG_A_COUNT + LDG_B_COUNT + + # ---- main loop: wait oldest tile, read frags, launch tile PREFETCH ahead, compute ---- + for kt_idx in range_constexpr(tiles_per_split): + cur = kt_idx % PIPE_STAGES + inflight_tiles = min(PREFETCH - 1, tiles_per_split - 1 - kt_idx) + barrier(vmcnt=inflight_tiles * LDG_PER_TILE, lgkmcnt=0) + a_frags = read_a_frags(cur) + b_frags = read_b_frags(cur) + nxt = kt_idx + PREFETCH + if const_expr(nxt < tiles_per_split): + _load_a(nxt % PIPE_STAGES, k_off + nxt * TILE_K) + _load_b(nxt % PIPE_STAGES, k_off + nxt * TILE_K) + rocdl.sched_vmem(LDG_A_COUNT + LDG_B_COUNT) + acc = do_compute(acc, a_frags, b_frags) + + _row_chk = npq % TILE_M != 0 + _need_chk = _row_chk or n_tail + _vec_store = (n == 1) and (not use_splitk) and (dhw % MFMA_C_VALUES == 0) and (not BIG_OUT) + + if const_expr(BIG_OUT): + y_elem_base = fx.Int64(buffer_ops.extract_base_index(y)) + + def _big_store(off_nk_i64, value): + addr = y_elem_base + off_nk_i64 * fx.Int64(BF16_BYTES) + ptr = buffer_ops.create_llvm_ptr(addr, address_space=1) + llvm.StoreOp(value.ir_value() if hasattr(value, "ir_value") else value, ptr, alignment=2) + + def _valid_raw(row, col): + if const_expr(_row_chk and n_tail): + return arith.andi(row < fx.Index(npq), col < fx.Index(k)) + if const_expr(_row_chk): + v = row < fx.Index(npq) + return arith.andi(v, v) + v = col < fx.Index(k) + return arith.andi(v, v) + + def store_acc(): + for mi in range_constexpr(MI_M): + row_base = m_offset + wave_m * WARP_M + mi * MFMA_M + c_m_vec + for ni in range_constexpr(MI_N): + col = n_offset + fx.Index(wave_n * WARP_N + ni * MFMA_N + c_n) + a = Vec(acc[mi * MI_N + ni]) + if const_expr(has_bias and not use_splitk): + col_i = fx.Int32(col) + if const_expr(n_tail): + col_i = arith.select(col < fx.Index(k), col_i, fx.Int32(0)) + bias_val = fx.Float32(buffer_ops.buffer_load(bias_rsrc, col_i, vec_width=1, dtype=fx.Float32)) + + if const_expr(_vec_store): + row0 = fx.Index(row_base) + off_nk0 = col * dhw + row0 + + def _emit_vec(): + vals = [] + for i in range_constexpr(MFMA_C_VALUES): + cval = (a[i] + bias_val) if const_expr(has_bias) else a[i] + vals.append(cval.to(elem_ty)) + v4 = fx.Vector.from_elements(vals, dtype=elem_ty) + buffer_ops.buffer_store(v4, y_rsrc, off_nk0) + + if const_expr(_need_chk): + if _valid_raw(row0, col): + _emit_vec() + else: + _emit_vec() + continue + + for i in range_constexpr(MFMA_C_VALUES): + row = fx.Index(row_base + i) + off_sk = row * k + col + + if const_expr(n == 1): + off_nk = col * dhw + row + else: + n_idx = row // dhw + sp = row % dhw + off_nk = n_idx * (k * dhw) + col * dhw + sp + + def _emit(): + if const_expr(use_splitk): + off_b = fx.Int32(off_sk * 4) + z0 = fx.Int32(0) + buffer_atomic_add(a[i], y_rsrc, off_b, z0, z0) + else: + cval = (a[i] + bias_val).to(elem_ty) if const_expr(has_bias) else a[i].to(elem_ty) + if const_expr(BIG_OUT): + _big_store(fx.Int64(off_nk), cval) + else: + buffer_ops.buffer_store(cval, y_rsrc, off_nk) + + if const_expr(_need_chk): + if _valid_raw(row, col): + _emit() + else: + _emit() + + store_acc() + + @flyc.jit + def launch(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias: fx.Tensor, stream: fx.Stream = fx.Stream(None)): + conv3d_implicit_kernel(y, x, weight, bias).launch( + grid=(grid_m, grid_n, splitk), block=(BLOCK_THREADS, 1, 1), stream=stream + ) + + return launch + + +def _resolve_splitk(splitk, npq, crs, k, device, tile=DEFAULT_TILE): + k_tiles = (crs + TILE_K - 1) // TILE_K + if splitk is None: + tile_m, tile_n = tile[0], tile[1] + base = ((npq + tile_m - 1) // tile_m) * ((k + tile_n - 1) // tile_n) + if ( + npq < 4096 + or k_tiles < 16 + or k % tile_n != 0 + or npq % tile_m != 0 + or crs % TILE_K != 0 + or npq * k * 4 > 0x7FFFFFFF + ): + sk = 1 + else: + try: + num_cu = torch.cuda.get_device_properties(device).multi_processor_count + except Exception: + num_cu = 256 + if base >= (3 * num_cu) // 4: + sk = 1 + else: + sk = min(4, max(1, num_cu // base), k_tiles) + else: + sk = max(1, splitk) + while sk > 1 and k_tiles % sk != 0: + sk -= 1 + return sk + + +def _conv3d_impl(x, weight, bias=None, stride=1, padding=0, splitk=None, stream=None, tile=None, autotune=None): + n, c, d, h, w = x.shape + k, wc, kt, kh, kw = weight.shape + assert c == wc + assert x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16 + st, sh, sw = (stride, stride, stride) if isinstance(stride, int) else stride + pt, ph, pw = (padding, padding, padding) if isinstance(padding, int) else padding + + # 1x1x1 fast path: y[n,k,dhw] = sum_c weight[k,c] * x[n,c,dhw] — pure channel GEMM. + if kt == 1 and kh == 1 and kw == 1 and st == 1 and sh == 1 and sw == 1 and pt == 0 and ph == 0 and pw == 0: + wm = weight.reshape(k, c) + if n == 1: + y = torch.matmul(wm, x.reshape(c, d * h * w)).reshape(n, k, d, h, w) + else: + y = torch.matmul(wm, x.reshape(n, c, d * h * w)).reshape(n, k, d, h, w) + if bias is not None: + y = y + bias.to(y.dtype).view(1, k, 1, 1, 1) + return y + + do = (d + 2 * pt - kt) // st + 1 + ho = (h + 2 * ph - kh) // sh + 1 + wo = (w + 2 * pw - kw) // sw + 1 + npq = n * do * ho * wo + crs = c * kt * kh * kw + + launch_stream = torch.cuda.current_stream() if stream is None else stream + has_bias = bias is not None + bias_arg = bias.to(torch.float32).contiguous() if has_bias else torch.empty(1, device=x.device, dtype=torch.float32) + + x_ndhwc = _ncdhw_to_ndhwc(x, stream) + w_packed = _prep_weight(weight, k, kt, kh, kw, c) + + shape = (n, c, d, h, w, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias) + + def _run(the_tile, the_wgm=1): + sk = _resolve_splitk(splitk, npq, crs, k, x.device, the_tile) + if sk > 1: + y = torch.zeros((npq, k), device=x.device, dtype=torch.float32) + else: + y = torch.empty((n, k, do, ho, wo), device=x.device, dtype=torch.bfloat16) + exe = compile_conv3d_implicit( + n, c, d, h, w, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias, sk, the_tile, the_wgm + ) + exe(y, x_ndhwc, w_packed, bias_arg, launch_stream) + return y, sk + + if tile is not None: + chosen_tile = tuple(tile) + chosen_wgm = 1 + elif autotune or (autotune is None and _autotune_enabled()): + from kernels.conv.conv3d_implicit_autotune import BF16_CANDIDATES, WGM_VALUES, autotune_conv3d + + candidates = [(t, w) for t in BF16_CANDIDATES for w in WGM_VALUES] + best = autotune_conv3d("bf16", shape, "bf16", candidates, x.device, lambda tw: _run(tw[0], tw[1])[0]) + chosen_tile, chosen_wgm = best + else: + chosen_tile = DEFAULT_TILE + chosen_wgm = 1 + + y, sk = _run(chosen_tile, chosen_wgm) + if sk > 1: + if has_bias: + y = y + bias_arg.view(1, k) + y = y.to(torch.bfloat16) + return y.view(n, do, ho, wo, k).permute(0, 4, 1, 2, 3) + return y + + +def _conv2d_impl(x, weight, bias=None, stride=1, padding=0, **kwargs): + assert x.dim() == 4 and weight.dim() == 4, "conv2d expects (N,C,H,W) / (K,C,R,S)" + sh, sw = (stride, stride) if isinstance(stride, int) else stride + ph, pw = (padding, padding) if isinstance(padding, int) else padding + n, c, h, w = x.shape + k, wc, r, s = weight.shape + x5 = x.reshape(n, c, 1, h, w) + w5 = weight.reshape(k, wc, 1, r, s) + y5 = _conv3d_impl(x5, w5, bias=bias, stride=(1, sh, sw), padding=(0, ph, pw), **kwargs) + return y5.reshape(y5.shape[0], y5.shape[1], y5.shape[3], y5.shape[4]) + + +def _conv1d_impl(x, weight, bias=None, stride=1, padding=0, **kwargs): + assert x.dim() == 3 and weight.dim() == 3, "conv1d expects (N,C,W) / (K,C,S)" + sw = stride if isinstance(stride, int) else stride[0] + pw = padding if isinstance(padding, int) else padding[0] + n, c, w = x.shape + k, wc, s = weight.shape + x5 = x.reshape(n, c, 1, 1, w) + w5 = weight.reshape(k, wc, 1, 1, s) + y5 = _conv3d_impl(x5, w5, bias=bias, stride=(1, 1, sw), padding=(0, 0, pw), **kwargs) + return y5.reshape(y5.shape[0], y5.shape[1], y5.shape[4]) + + +def conv3d_implicit(x, weight, bias=None, stride=1, padding=0, **kwargs): + """Main implicit-GEMM conv entry; dispatches 1D/2D/3D by filter rank. + + Rank is taken from the filter (weight.dim() - 2): 3 -> 3D (N,C,D,H,W)/(K,C,T,R,S), + 2 -> 2D (N,C,H,W)/(K,C,R,S), 1 -> 1D (N,C,W)/(K,C,S); x and weight must match. + True 3D calls run the implementation directly; 2D/1D reshape to the degenerate + 5D case. stride/padding/bias and extra kwargs (splitk, tile, autotune, stream) + forward to the chosen path. + """ + assert x.dim() == weight.dim(), f"x rank {x.dim()} != weight rank {weight.dim()}" + spatial_rank = weight.dim() - 2 + if spatial_rank == 3: + return _conv3d_impl(x, weight, bias=bias, stride=stride, padding=padding, **kwargs) + if spatial_rank == 2: + return _conv2d_impl(x, weight, bias=bias, stride=stride, padding=padding, **kwargs) + if spatial_rank == 1: + return _conv1d_impl(x, weight, bias=bias, stride=stride, padding=padding, **kwargs) + raise ValueError(f"conv3d_implicit supports 1D/2D/3D; got filter rank {weight.dim()}") diff --git a/kernels/conv/conv3d_implicit_8wave.py b/kernels/conv/conv3d_implicit_8wave.py deleted file mode 100644 index 02d5c052a..000000000 --- a/kernels/conv/conv3d_implicit_8wave.py +++ /dev/null @@ -1,597 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""8-wave double-buffered implicit-GEMM conv3d (BF16). - -x: (N, C, D, H, W) bf16 NCDHW, weight: (K, C, T, R, S) bf16 KCTRS. -Returns (N, K, Do, Ho, Wo) bf16. Supports stride, padding, bias, and split-K. -""" - -import functools -import weakref - -import torch - -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl._mlir.dialects import llvm -from flydsl.expr import arith, const_expr, range_constexpr, rocdl -from flydsl.expr.typing import T -from kernels.common import buffer_ops -from kernels.common.mem_ops import buffer_atomic_add - -TILE_M = 128 -TILE_N = 128 -TILE_K = 32 -STAGES = 2 - -WAVE_M = 2 -WAVE_N = 4 -WARP_SIZE = 64 -BLOCK_THREADS = WAVE_M * WAVE_N * WARP_SIZE # 512 - -MFMA_M = 16 -MFMA_N = 16 -MFMA_A_VALUES = 8 -MFMA_B_VALUES = 8 -MFMA_C_VALUES = 4 - -HALF_M = TILE_M // 2 -HALF_N = TILE_N // 2 -QM_STEPS = HALF_M // WAVE_M // MFMA_M # 2 -QN_STEPS = HALF_N // WAVE_N // MFMA_N # 1 -N_SUB = QM_STEPS * QN_STEPS - -# The main loop below is handwritten for this exact 8-wave shape. -assert QM_STEPS == 2 and QN_STEPS == 1 - -LDG_VEC = 8 -BLOCK_VECS = LDG_VEC * BLOCK_THREADS -LDG_A_COUNT = TILE_M * TILE_K // BLOCK_VECS -LDS_A_SIZE = STAGES * TILE_M * TILE_K -LDS_B_SIZE = STAGES * TILE_N * TILE_K - - -_WEIGHT_CACHE = {} - - -def _prep_weight(w, k, kt, kh, kw, c): - key = id(w) - ent = _WEIGHT_CACHE.get(key) - if ent is not None and ent[0]() is w: - return ent[1] - wk = w.permute(0, 2, 3, 4, 1).contiguous().reshape(k, kt * kh * kw * c) - _WEIGHT_CACHE[key] = (weakref.ref(w), wk) - return wk - - -TR_TILE = 64 -TR_VEC = 8 -TR_THREADS = 256 -_TR_VPL = TR_TILE // TR_VEC -_TR_ITERS = (TR_TILE * TR_TILE) // (TR_VEC * TR_THREADS) -_TR_PAD = 8 -_TR_LDS_S = TR_TILE + _TR_PAD - - -@functools.lru_cache(maxsize=64) -def compile_transpose_ncdhw_ndhwc(n, c, s): - """Transpose flat (N, C, S) -> (N, S, C) (S == T*H*W). Requires c%8==0, s%8==0.""" - grid_s = (s + TR_TILE - 1) // TR_TILE - grid_c = (c + TR_TILE - 1) // TR_TILE - elem_ty = fx.BFloat16 - - @flyc.kernel(known_block_size=[TR_THREADS, 1, 1]) - def transpose_kernel(out: fx.Tensor, inp: fx.Tensor): - in_rsrc = buffer_ops.create_buffer_resource(inp) - out_rsrc = buffer_ops.create_buffer_resource(out) - lds_alloc = fx.SharedAllocator(static=False) - lds = lds_alloc.allocate(fx.Array[elem_ty, TR_TILE * _TR_LDS_S, 16]).peek() - - Vec = fx.Vector - - class Vec8Ty: - ir_type = Vec.make_type(TR_VEC, elem_ty) - - class BF16Ty: - ir_type = elem_ty.ir_type - - tid = fx.thread_idx.x - s0 = fx.block_idx.x * TR_TILE - c0 = fx.block_idx.y * TR_TILE - nb = fx.block_idx.z - in_base = nb * c * s - out_base = nb * s * c - - def lds_store_vec8(elem_offset, value): - base = fx.Int64(fx.ptrtoint(lds.ptr)) + fx.Int64(elem_offset * 2) - ptr = buffer_ops.create_llvm_ptr(base, address_space=3) - llvm.StoreOp(value, ptr, alignment=16) - - def lds_load_scalar(elem_offset): - u8 = fx.recast_iter(fx.Uint8, lds.ptr) - return fx.ptr_load(u8 + fx.Int32(elem_offset * 2), result_type=BF16Ty) - - # Read: coalesced vec8 along contiguous S -> LDS[c_local][s_local]. - for i in range_constexpr(_TR_ITERS): - lin = tid + i * TR_THREADS - rc = lin // _TR_VPL - sv = (lin % _TR_VPL) * TR_VEC - cc = c0 + rc - ss = s0 + sv - valid = (cc < c) & (ss < s) - g = fx.Int32(in_base + cc * s + ss) - safe = arith.select(valid, g, fx.Int32(0)) - v = buffer_ops.buffer_load(in_rsrc, safe, vec_width=TR_VEC, dtype=elem_ty) - lds_store_vec8(rc * _TR_LDS_S + sv, v) - - llvm.InlineAsmOp(None, [], "s_waitcnt lgkmcnt(0)\n\ts_barrier", "", has_side_effects=True) - - for i in range_constexpr(_TR_ITERS): - lin = tid + i * TR_THREADS - rs = lin // _TR_VPL - cv = (lin % _TR_VPL) * TR_VEC - ss = s0 + rs - cc = c0 + cv - scalars = [lds_load_scalar((cv + j) * _TR_LDS_S + rs) for j in range_constexpr(TR_VEC)] - vv = fx.Vector.from_elements(scalars, dtype=elem_ty) - valid = (ss < s) & (cc < c) - if valid: - go = fx.Int32(out_base + ss * c + cc) - buffer_ops.buffer_store(vv, out_rsrc, go) - - @flyc.jit - def launch_transpose(out: fx.Tensor, inp: fx.Tensor, stream: fx.Stream = fx.Stream(None)): - transpose_kernel(out, inp).launch( - grid=(grid_s, grid_c, n), - block=(TR_THREADS, 1, 1), - stream=stream, - ) - - return launch_transpose - - -def _ncdhw_to_ndhwc(x, stream): - """Fast NCDHW->NDHWC via the tiled transpose kernel; falls back to torch.""" - n, c, t, h, w = x.shape - s = t * h * w - if not (x.is_contiguous() and x.dtype == torch.bfloat16 and c % 8 == 0 and s % 8 == 0): - return x.permute(0, 2, 3, 4, 1).contiguous() - out = torch.empty((n, t, h, w, c), device=x.device, dtype=x.dtype) - exe = compile_transpose_ncdhw_ndhwc(n, c, s) - exe(out, x, torch.cuda.current_stream() if stream is None else stream) - return out - - -@functools.lru_cache(maxsize=64) -def compile_conv3d_implicit_8wave(n, c, d, h, w, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias=False, splitk=1): - do = (d + 2 * pt - kt) // st + 1 - ho = (h + 2 * ph - kh) // sh + 1 - wo = (w + 2 * pw - kw) // sw + 1 - dhw = do * ho * wo - hw_o = ho * wo - npq = n * dhw - crs = c * kt * kh * kw - k_tiles = (crs + TILE_K - 1) // TILE_K - - assert c % LDG_VEC == 0 - assert LDG_A_COUNT == 1 - - n_tail = k % TILE_N != 0 - grid_n = (k + TILE_N - 1) // TILE_N - - splitk = max(1, min(splitk, k_tiles)) - tiles_per_split = k_tiles // splitk - use_splitk = splitk > 1 - - grid_m = (npq + TILE_M - 1) // TILE_M - elem_ty = fx.BFloat16 - mfma_fn = rocdl.mfma_f32_16x16x32_bf16 - - @flyc.kernel(known_block_size=[BLOCK_THREADS, 1, 1]) - def conv3d_8wave_kernel(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias: fx.Tensor): - x_rsrc = buffer_ops.create_buffer_resource(x) - w_rsrc = buffer_ops.create_buffer_resource(weight) - y_rsrc = buffer_ops.create_buffer_resource(y) - if const_expr(has_bias): - bias_rsrc = buffer_ops.create_buffer_resource(bias) - - lds_alloc = fx.SharedAllocator(static=False) - a_lds = lds_alloc.allocate(fx.Array[elem_ty, LDS_A_SIZE, 16]).peek() - b_lds = lds_alloc.allocate(fx.Array[elem_ty, LDS_B_SIZE, 16]).peek() - - tid = fx.thread_idx.x - m_offset = fx.block_idx.x * TILE_M - n_offset = fx.block_idx.y * TILE_N - if const_expr(use_splitk): - k_off = fx.block_idx.z * (tiles_per_split * TILE_K) - else: - k_off = 0 - - wid = tid // WARP_SIZE - lane = tid % WARP_SIZE - wave_m = wid // WAVE_N - wave_n = wid % WAVE_N - - lane_m = lane % MFMA_M - lane_n = lane % MFMA_N - lane_k_a = lane // MFMA_M * MFMA_A_VALUES - lane_k_b = lane // MFMA_N * MFMA_B_VALUES - c_m_vec = lane // MFMA_N * MFMA_C_VALUES - c_n = lane % MFMA_N - - Vec = fx.Vector - - class Vec8Ty: - ir_type = Vec.make_type(8, elem_ty) - - acc0 = Vec.filled(MFMA_C_VALUES, 0.0, fx.Float32) - acc00 = [acc0 for _ in range_constexpr(N_SUB)] - acc01 = [acc0 for _ in range_constexpr(N_SUB)] - acc10 = [acc0 for _ in range_constexpr(N_SUB)] - acc11 = [acc0 for _ in range_constexpr(N_SUB)] - - zero8 = Vec.filled(8, 0.0, elem_ty) - - def barrier(vmcnt=0, lgkmcnt=None): - waits = [] - if vmcnt is not None: - waits.append(f"vmcnt({vmcnt})") - if lgkmcnt is not None: - waits.append(f"lgkmcnt({lgkmcnt})") - pre = ("s_waitcnt " + " ".join(waits) + "\n\t") if waits else "" - llvm.InlineAsmOp(None, [], f"{pre}s_barrier", "", has_side_effects=True) - - def waitcnt(vmcnt=None, lgkmcnt=None): - waits = [] - if vmcnt is not None: - waits.append(f"vmcnt({vmcnt})") - if lgkmcnt is not None: - waits.append(f"lgkmcnt({lgkmcnt})") - if waits: - llvm.InlineAsmOp(None, [], "s_waitcnt " + " ".join(waits), "", has_side_effects=True) - - def lds_ptr_at(lds_array, byte_offset): - lds_base = fx.Int64(fx.ptrtoint(lds_array.ptr)) + fx.Int64(byte_offset) - return buffer_ops.create_llvm_ptr(lds_base, address_space=3) - - def lds_store_vec8(lds_array, elem_offset, value): - llvm.StoreOp(value, lds_ptr_at(lds_array, elem_offset * 2), alignment=16) - - def lds_load_vec8(lds_array, elem_offset): - u8_ptr = fx.recast_iter(fx.Uint8, lds_array.ptr) - return fx.ptr_load(u8_ptr + fx.Int32(elem_offset * 2), result_type=Vec8Ty) - - def a_lds_off(stage, row, col): - return (fx.Index(stage) * TILE_M + row) * TILE_K + col - - def b_lds_off(stage, row, col): - return (fx.Index(stage) * TILE_N + row) * TILE_K + col - - def in_range(v, hi): - return (v >= 0) & (v < fx.Index(hi)) - - # ---- 3D im2col gather (global -> registers) ---- - def gather_a(k_base): - linear = tid * LDG_VEC - local_m = linear // TILE_K - local_k = linear % TILE_K - row = m_offset + local_m - row_valid = row < fx.Index(npq) - n_idx = row // dhw - rem = row % dhw - ot = rem // hw_o - rem2 = rem % hw_o - oh = rem2 // wo - ow = rem2 % wo - k_abs = fx.Index(k_base) + fx.Index(local_k) - cc = k_abs % c - ckk = k_abs // c - kw_i = ckk % kw - ckk2 = ckk // kw - kh_i = ckk2 % kh - kt_i = ckk2 // kh - in_t = ot * st + kt_i - pt - in_h = oh * sh + kh_i - ph - in_w = ow * sw + kw_i - pw - k_valid = k_abs < fx.Index(crs) - valid = row_valid & k_valid & in_range(in_t, d) & in_range(in_h, h) & in_range(in_w, w) - g_off = (((n_idx * d + in_t) * h + in_h) * w + in_w) * c + cc - g_off_i = fx.Int32(g_off) - safe = arith.select(valid, g_off_i, fx.Int32(0)) - raw = buffer_ops.buffer_load(x_rsrc, safe, vec_width=8, dtype=elem_ty) - return (raw, valid, local_m * TILE_K + local_k) - - def gather_b(k_base): - linear = tid * LDG_VEC - local_n = linear // TILE_K - local_k = linear % TILE_K - col = n_offset + fx.Index(local_n) - g_off = fx.Int32(col * crs + (fx.Index(k_base) + fx.Index(local_k))) - if const_expr(n_tail): - col_valid = col < fx.Index(k) - safe = arith.select(col_valid, g_off, fx.Int32(0)) - raw = buffer_ops.buffer_load(w_rsrc, safe, vec_width=8, dtype=elem_ty) - return (raw, col_valid, local_n * TILE_K + local_k) - raw = buffer_ops.buffer_load(w_rsrc, g_off, vec_width=8, dtype=elem_ty) - return (raw, None, local_n * TILE_K + local_k) - - def commit_a(stage, vo): - raw, valid, off = vo - val = arith.select(valid, raw, zero8) # mask consumed here (hidden behind MFMAs) - lds_store_vec8(a_lds, fx.Index(stage) * TILE_M * TILE_K + off, val) - - def commit_b(stage, vo): - raw, valid, off = vo - val = raw if const_expr(valid is None) else arith.select(valid, raw, zero8) - lds_store_vec8(b_lds, fx.Index(stage) * TILE_N * TILE_K + off, val) - - # ---- single-vec ds_read (LDS -> register) ---- - def read_a_vec(stage, m_half, wm): - a_row = m_half * HALF_M + wave_m * (HALF_M // WAVE_M) + wm * MFMA_M + lane_m - return lds_load_vec8(a_lds, a_lds_off(stage, fx.Index(a_row), fx.Index(lane_k_a))) - - def read_b_vec(stage, n_half, wn): - b_row = n_half * HALF_N + wave_n * (HALF_N // WAVE_N) + wn * MFMA_N + lane_n - return lds_load_vec8(b_lds, b_lds_off(stage, fx.Index(b_row), fx.Index(lane_k_b))) - - def mfma_one(a_frag, b_frag, c_frag): - out = mfma_fn( - T.vec(MFMA_C_VALUES, T.f32), - [a_frag, b_frag, c_frag, 0, 0, 0], - ) - rocdl.sched_mfma(1) - return out - - # phase_b_prefetch: compute C00 while prefetching B1 from LDS. - def phase_b_prefetch(read_stage, a0_0, a0_1, b0_0, acc): - out = [v for v in acc] - out[0] = mfma_one(a0_0, b0_0, out[0]) - b1_0 = read_b_vec(read_stage, 1, 0) - rocdl.sched_dsrd(1) - out[1] = mfma_one(a0_1, b0_0, out[1]) - return out, b1_0 - - # phase_a_prefetch: compute C01 while prefetching A1 from LDS. - def phase_a_prefetch(read_stage, a0_0, a0_1, b1_0, acc): - out = [v for v in acc] - out[0] = mfma_one(a0_0, b1_0, out[0]) - a1_0 = read_a_vec(read_stage, 1, 0) - rocdl.sched_dsrd(1) - out[1] = mfma_one(a0_1, b1_0, out[1]) - a1_1 = read_a_vec(read_stage, 1, 1) - rocdl.sched_dsrd(1) - return out, a1_0, a1_1 - - # phase_ab_prefetch: compute C11 while reading only the next tile's B0. - # A0 is read at the start of the next iteration to shorten VGPR lifetime. - def phase_ab_prefetch(read_stage, a1_0, a1_1, b1_0, acc): - out = [v for v in acc] - out[0] = mfma_one(a1_0, b1_0, out[0]) - next_b0_0 = read_b_vec(read_stage, 0, 0) - rocdl.sched_dsrd(1) - out[1] = mfma_one(a1_1, b1_0, out[1]) - return out, next_b0_0 - - def phase_compute(a1_0, a1_1, b_0, acc): - out = [v for v in acc] - out[0] = mfma_one(a1_0, b_0, out[0]) - out[1] = mfma_one(a1_1, b_0, out[1]) - return out - - def compute_prefetch_phases(read_stage, a0_0, a0_1, b0_0): - rocdl.s_setprio(1) - c00, b1_0 = phase_b_prefetch(read_stage, a0_0, a0_1, b0_0, acc00) - c01, a1_0, a1_1 = phase_a_prefetch(read_stage, a0_0, a0_1, b1_0, acc01) - rocdl.s_setprio(0) - return c00, c01, a1_0, a1_1, b1_0 - - # ---- prologue: tile 0 -> LDS, tile 1 -> VGPR prefetch ---- - stage = 0 - next_stage = 1 - commit_a(stage, gather_a(k_off)) - commit_b(stage, gather_b(k_off)) - if const_expr(tiles_per_split > 1): - pf_a = gather_a(k_off + TILE_K) - pf_b = gather_b(k_off + TILE_K) - rocdl.sched_vmem(2) - barrier(vmcnt=None, lgkmcnt=0) - - a0_0 = read_a_vec(stage, 0, 0) - a0_1 = read_a_vec(stage, 0, 1) - b0_0 = read_b_vec(stage, 0, 0) - rocdl.sched_dsrd(3) - - # ---- main loop: compute tile k, write prefetched k+1, load k+2 ---- - if const_expr(tiles_per_split > 2): - for kt_idx in range_constexpr(tiles_per_split - 2): - acc00, acc01, a1_0, a1_1, b1_0 = compute_prefetch_phases(stage, a0_0, a0_1, b0_0) - - # Extra lockstep barrier after the acc00/acc01 phase. - barrier(vmcnt=None, lgkmcnt=None) - - commit_a(next_stage, pf_a) - rocdl.sched_dswr(1) - pf_a = gather_a(k_off + (kt_idx + 2) * TILE_K) - rocdl.sched_vmem(1) - rocdl.s_setprio(1) - acc10[0] = mfma_one(a1_0, b0_0, acc10[0]) - - commit_b(next_stage, pf_b) - rocdl.sched_dswr(1) - pf_b = gather_b(k_off + (kt_idx + 2) * TILE_K) - rocdl.sched_vmem(1) - acc10[1] = mfma_one(a1_1, b0_0, acc10[1]) - rocdl.s_setprio(0) - - barrier(vmcnt=None, lgkmcnt=0) - - rocdl.s_setprio(1) - acc11, b0_0 = phase_ab_prefetch(next_stage, a1_0, a1_1, b1_0, acc11) - rocdl.s_setprio(0) - - # Extra lockstep barrier after the acc11 phase. - barrier(vmcnt=None, lgkmcnt=None) - - stage = next_stage - next_stage = (stage + 1) % STAGES - a0_0 = read_a_vec(stage, 0, 0) - a0_1 = read_a_vec(stage, 0, 1) - rocdl.sched_dsrd(2) - - # ---- peeled iteration: compute tile K-2, write final prefetched tile ---- - if const_expr(tiles_per_split >= 2): - acc00, acc01, a1_0, a1_1, b1_0 = compute_prefetch_phases(stage, a0_0, a0_1, b0_0) - - commit_a(next_stage, pf_a) - rocdl.sched_dswr(1) - rocdl.s_setprio(1) - acc10[0] = mfma_one(a1_0, b0_0, acc10[0]) - - commit_b(next_stage, pf_b) - rocdl.sched_dswr(1) - acc10[1] = mfma_one(a1_1, b0_0, acc10[1]) - rocdl.s_setprio(0) - - barrier(vmcnt=None, lgkmcnt=0) - - rocdl.s_setprio(1) - acc11, b0_0 = phase_ab_prefetch(next_stage, a1_0, a1_1, b1_0, acc11) - rocdl.s_setprio(0) - stage = next_stage - next_stage = (stage + 1) % STAGES - a0_0 = read_a_vec(stage, 0, 0) - a0_1 = read_a_vec(stage, 0, 1) - rocdl.sched_dsrd(2) - - # ---- epilogue: final tile, no more LDS overwrite or next-tile reads ---- - acc00, acc01, a1_0, a1_1, b1_0 = compute_prefetch_phases(stage, a0_0, a0_1, b0_0) - waitcnt(lgkmcnt=0) - rocdl.s_setprio(1) - acc10 = phase_compute(a1_0, a1_1, b0_0, acc10) - acc11 = phase_compute(a1_0, a1_1, b1_0, acc11) - rocdl.s_setprio(0) - - _row_chk = npq % TILE_M != 0 - _need_chk = _row_chk or n_tail - - def _valid_raw(row, col): - if const_expr(_row_chk and n_tail): - return arith.andi(row < fx.Index(npq), col < fx.Index(k)) - if const_expr(_row_chk): - v = row < fx.Index(npq) - return arith.andi(v, v) - v = col < fx.Index(k) - return arith.andi(v, v) - - def store_quad(acc, m_half, n_half): - for wm in range_constexpr(QM_STEPS): - row_base = m_offset + m_half * HALF_M + wave_m * (HALF_M // WAVE_M) + wm * MFMA_M + c_m_vec - for wn in range_constexpr(QN_STEPS): - col = n_offset + fx.Index(n_half * HALF_N + wave_n * (HALF_N // WAVE_N) + wn * MFMA_N + c_n) - a = Vec(acc[wm * QN_STEPS + wn]) - if const_expr(has_bias and not use_splitk): - col_i = fx.Int32(col) - if const_expr(n_tail): - col_i = arith.select(col < fx.Index(k), col_i, fx.Int32(0)) - bias_val = fx.Float32(buffer_ops.buffer_load(bias_rsrc, col_i, vec_width=1, dtype=fx.Float32)) - for i in range_constexpr(MFMA_C_VALUES): - row = fx.Index(row_base + i) - off_sk = row * k + col - - if const_expr(n == 1): - off_nk = col * dhw + row - else: - ni = row // dhw - sp = row % dhw - off_nk = ni * (k * dhw) + col * dhw + sp - - def _emit(): - if const_expr(use_splitk): - off_b = fx.Int32(off_sk * 4) - z0 = fx.Int32(0) - buffer_atomic_add(a[i], y_rsrc, off_b, z0, z0) - else: - cval = (a[i] + bias_val).to(elem_ty) if const_expr(has_bias) else a[i].to(elem_ty) - buffer_ops.buffer_store(cval, y_rsrc, off_nk) - - if const_expr(_need_chk): - if _valid_raw(row, col): - _emit() - else: - _emit() - - store_quad(acc00, 0, 0) - store_quad(acc01, 0, 1) - store_quad(acc10, 1, 0) - store_quad(acc11, 1, 1) - - @flyc.jit - def launch(y: fx.Tensor, x: fx.Tensor, weight: fx.Tensor, bias: fx.Tensor, stream: fx.Stream = fx.Stream(None)): - conv3d_8wave_kernel(y, x, weight, bias).launch( - grid=(grid_m, grid_n, splitk), block=(BLOCK_THREADS, 1, 1), stream=stream - ) - - return launch - - -def _choose_splitk(npq, crs, k, device): - grid_m = (npq + TILE_M - 1) // TILE_M - grid_n = (k + TILE_N - 1) // TILE_N - base = grid_m * grid_n - k_tiles = (crs + TILE_K - 1) // TILE_K - - if npq < 4096 or k_tiles < 16: - return 1 - if k % TILE_N != 0 or npq % TILE_M != 0 or crs % TILE_K != 0: # atomic path needs clean tiles - return 1 - try: - num_cu = torch.cuda.get_device_properties(device).multi_processor_count - except Exception: - num_cu = 256 - if base >= (3 * num_cu) // 4: # base grid already (nearly) fills the machine - return 1 - sk = min(4, max(1, num_cu // base), k_tiles) # aim to roughly fill the CUs - while sk > 1 and k_tiles % sk != 0: # prefer a divisor (no overhang) - sk -= 1 - return sk - - -def conv3d_implicit_8wave(x, weight, bias=None, stride=1, padding=0, splitk=None, stream=None): - # x: (N,C,D,H,W) bf16, weight: (K,C,T,R,S) bf16. splitk=None -> auto-dispatch. - n, c, d, h, w = x.shape - k, wc, kt, kh, kw = weight.shape - assert c == wc - assert x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16 - st, sh, sw = (stride, stride, stride) if isinstance(stride, int) else stride - pt, ph, pw = (padding, padding, padding) if isinstance(padding, int) else padding - do = (d + 2 * pt - kt) // st + 1 - ho = (h + 2 * ph - kh) // sh + 1 - wo = (w + 2 * pw - kw) // sw + 1 - npq = n * do * ho * wo - crs = c * kt * kh * kw - - sk = _choose_splitk(npq, crs, k, x.device) if splitk is None else max(1, splitk) - k_tiles = (crs + TILE_K - 1) // TILE_K - while sk > 1 and k_tiles % sk != 0: - sk -= 1 - use_splitk = sk > 1 - - # Fast fused NCDHW->NDHWC transpose + cached weight permute (reuse prod's - # helpers) instead of torch permute+contiguous per call. - x_ndhwc = _ncdhw_to_ndhwc(x, stream) - w_packed = _prep_weight(weight, k, kt, kh, kw, c) - if use_splitk: - y = torch.zeros((npq, k), device=x.device, dtype=torch.float32) - else: - y = torch.empty((n, k, do, ho, wo), device=x.device, dtype=torch.bfloat16) - has_bias = bias is not None - bias_arg = bias.to(torch.float32).contiguous() if has_bias else torch.empty(1, device=x.device, dtype=torch.float32) - exe = compile_conv3d_implicit_8wave(n, c, d, h, w, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias, sk) - exe(y, x_ndhwc, w_packed, bias_arg, torch.cuda.current_stream() if stream is None else stream) - if use_splitk: - if has_bias: - y = y + bias_arg.view(1, k) - y = y.to(torch.bfloat16) - return y.view(n, do, ho, wo, k).permute(0, 4, 1, 2, 3) - return y diff --git a/kernels/conv/conv3d_implicit_autotune.py b/kernels/conv/conv3d_implicit_autotune.py new file mode 100644 index 000000000..f130506de --- /dev/null +++ b/kernels/conv/conv3d_implicit_autotune.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Manual tile autotuner for the implicit-GEMM conv3d kernels.""" + +import json +import os +from pathlib import Path + +from flydsl.autotune import do_bench + + +def _device_fingerprint(): + """GPU arch string (e.g. 'gfx950'), or '' if unavailable.""" + try: + from flydsl.runtime.device import get_rocm_arch + + return str(get_rocm_arch()) + except Exception: + return "" + + +def _toolchain_fingerprint(): + """FlyDSL version, so a rebuild invalidates stale cached configs.""" + try: + import flydsl + + return str(getattr(flydsl, "__version__", "")) + except Exception: + return "" + + +# Curated legal candidates (TILE_M, TILE_N, WAVE_M, WAVE_N) from the enumerated +# constraint-satisfying space. Kept small to bound tuning time; illegal configs +# for a given shape are skipped at compile time (try/except in the sweep). +BF16_CANDIDATES = [ + (128, 128, 2, 4), + (128, 256, 2, 4), + (256, 128, 2, 4), + (256, 256, 2, 4), + (256, 256, 4, 4), + (128, 128, 4, 2), + (64, 128, 1, 4), + (64, 64, 2, 2), +] + +FP8_CANDIDATES = [ + (128, 128, 2, 4), + (128, 256, 2, 4), + (256, 128, 2, 4), + (256, 256, 2, 4), + (128, 128, 4, 2), + (64, 128, 1, 4), +] + +WGM_VALUES = [1, 4, 8] + +_MEM_CACHE = {} +_CACHE_SCHEMA_VERSION = 3 + + +def _cache_dir(): + return Path(os.environ.get("FLYDSL_AUTOTUNE_CACHE_DIR", os.path.expanduser("~/.flydsl/autotune"))) + + +def _cache_file(kind): + return _cache_dir() / f"conv3d_{kind}.json" + + +def _make_key(kind, shape, dtype_str, candidates): + def _canon(c): + if isinstance(c, tuple) and len(c) == 2 and isinstance(c[0], tuple): + return (tuple(c[0]), c[1]) + return tuple(c) + + return ( + kind, + tuple(shape), + dtype_str, + _device_fingerprint(), + _toolchain_fingerprint(), + _CACHE_SCHEMA_VERSION, + tuple(_canon(c) for c in candidates), + ) + + +def _load_disk(kind, key): + f = _cache_file(kind) + if not f.exists(): + return None + try: + data = json.loads(f.read_text()) + except Exception: + return None + ent = data.get(json.dumps(list(key))) + if ent is None: + return None + if isinstance(ent[0], list): + return (tuple(ent[0]), ent[1]) + return tuple(ent) + + +def _save_disk(kind, key, best): + f = _cache_file(kind) + f.parent.mkdir(parents=True, exist_ok=True) + data = {} + if f.exists(): + try: + data = json.loads(f.read_text()) + except Exception: + data = {} + if isinstance(best, tuple) and len(best) == 2 and isinstance(best[0], tuple): + data[json.dumps(list(key))] = [list(best[0]), best[1]] + else: + data[json.dumps(list(key))] = list(best) + f.write_text(json.dumps(data, indent=2)) + + +def autotune_conv3d(kind, shape, dtype_str, candidates, device, run_tile, warmup=5, rep=20): + """Return the best candidate for this problem shape. + + Each element of ``candidates`` is either a tile tuple (TILE_M, TILE_N, WAVE_M, WAVE_N) + or a (tile_tuple, wgm) pair when wgm sweep is enabled. ``run_tile(candidate)`` + must launch one full conv for the given candidate and return the output tensor. + The winning candidate is cached (in-memory + JSON on disk). + """ + key = _make_key(kind, shape, dtype_str, candidates) + if key in _MEM_CACHE: + return _MEM_CACHE[key] + disk = _load_disk(kind, key) + if disk is not None: + _MEM_CACHE[key] = disk + return disk + + results = [] + for tile in candidates: + try: + run_tile(tile) # dry run: triggers compile (lru_cache miss) outside do_bench + t = do_bench(lambda: run_tile(tile), warmup=warmup, rep=rep) + results.append((tile, t)) + except Exception: + pass # skip illegal / register-spilling / OOM configs + if not results: + raise RuntimeError(f"all conv3d {kind} autotune configs failed for shape {shape}") + + best = min(results, key=lambda x: x[1])[0] + _MEM_CACHE[key] = best + _save_disk(kind, key, best) + return best diff --git a/kernels/conv/conv3d_implicit_8wave_fp8.py b/kernels/conv/conv3d_implicit_fp8.py similarity index 96% rename from kernels/conv/conv3d_implicit_8wave_fp8.py rename to kernels/conv/conv3d_implicit_fp8.py index f8d57915c..f31d94808 100644 --- a/kernels/conv/conv3d_implicit_8wave_fp8.py +++ b/kernels/conv/conv3d_implicit_fp8.py @@ -194,9 +194,7 @@ def launch(out: fx.Tensor, weight: fx.Tensor, stream: fx.Stream = fx.Stream(None @functools.lru_cache(maxsize=64) -def compile_conv3d_implicit_8wave_fp8( - n, c, d, h, width, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias=False, splitk=1 -): +def compile_conv3d_implicit_fp8(n, c, d, h, width, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias=False, splitk=1): """Compile the FP8 conv: x is NDHWC FP8 bytes, weight is KTRSC FP8 bytes.""" do = (d + 2 * pt - kt) // st + 1 ho = (h + 2 * ph - kh) // sh + 1 @@ -421,9 +419,6 @@ def store_half_pair(acc0, acc1, m_half): for wn in range_constexpr(QN_STEPS): col = n_offset + fx.Index(n_half * HALF_N + wave_n * (HALF_N // WAVE_N) + wn * MFMA_N) + c_n col_valid = col < fx.Index(k) - # Under split-K the partial sums accumulate atomically into - # FP32; bias is a single per-output add left to the host - # post-pass (adding it per z-slice would scale it by splitk). if const_expr(has_bias and not use_splitk): bias_val = fx.Float32(buffer_ops.buffer_load(bias_rsrc, col, vec_width=1, dtype=fx.Float32)) acc_vec = Vec(acc[wm * QN_STEPS + wn]) @@ -440,14 +435,14 @@ def store_half_pair(acc0, acc1, m_half): else: if const_expr(has_bias): out = out + bias_val - # NCDHW output[ni, col, sp]: ni*(k*dhw) + col*dhw + sp. - # n==1 fast path: ni=0, sp=row, no integer division. + # NCDHW output[n_idx, col, sp]: n_idx*(k*dhw) + col*dhw + sp. + # n==1 fast path: n_idx=0, sp=row, no integer division. if const_expr(n == 1): off_ncdhw = col * dhw + row else: - ni = row // dhw + n_idx = row // dhw sp = row % dhw - off_ncdhw = ni * (k * dhw) + col * dhw + sp + off_ncdhw = n_idx * (k * dhw) + col * dhw + sp buffer_ops.buffer_store(out.to(fx.BFloat16), y_rsrc, off_ncdhw, mask=col_valid) store_half_pair(acc00, acc01, 0) @@ -554,7 +549,7 @@ def _prep_weight_fp8(weight: torch.Tensor, stream=None) -> torch.Tensor: return out -def conv3d_implicit_8wave_fp8(x, weight, bias=None, stride=1, padding=0, splitk=None, stream=None): +def conv3d_implicit_fp8(x, weight, bias=None, stride=1, padding=0, splitk=None, stream=None): """FP8 (E4M3FN) implicit conv3d. Same interface as the BF16 v6mb kernel. x: (N, C, D, H, W) bf16, weight: (K, C, T, R, S) bf16. Inputs are packed once @@ -594,7 +589,7 @@ def conv3d_implicit_8wave_fp8(x, weight, bias=None, stride=1, padding=0, splitk= y = torch.zeros((npq, k), device=x.device, dtype=torch.float32) else: y = torch.empty((n, k, do, ho, wo), device=x.device, dtype=torch.bfloat16) - exe = compile_conv3d_implicit_8wave_fp8(n, c, d, h, width, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias, sk) + exe = compile_conv3d_implicit_fp8(n, c, d, h, width, k, kt, kh, kw, st, sh, sw, pt, ph, pw, has_bias, sk) exe( flyc.from_torch_tensor(y.view(-1)), flyc.from_torch_tensor(x_arg), @@ -610,4 +605,4 @@ def conv3d_implicit_8wave_fp8(x, weight, bias=None, stride=1, padding=0, splitk= return y -__all__ = ["conv3d_implicit_8wave_fp8"] +__all__ = ["conv3d_implicit_fp8"] diff --git a/tests/kernels/test_conv3d_implicit.py b/tests/kernels/test_conv3d_implicit.py new file mode 100644 index 000000000..2eb117dbf --- /dev/null +++ b/tests/kernels/test_conv3d_implicit.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Correctness test for the bf16 implicit-GEMM conv3d kernel. + +Compares ``conv3d_implicit`` against ``torch.nn.functional.conv3d`` on +NCDHW/OIDHW bf16 inputs across stride/padding and M%TILE_M / K%TILE_N tail paths. +Channels must satisfy the kernel's ``c % 8 == 0`` and ``crs = c*kt*kh*kw`` a +multiple of TILE_K (32) constraints. +""" + +import pytest +import torch +import torch.nn.functional as F + +from flydsl.runtime.device import get_rocm_arch +from kernels.conv.conv3d_implicit import conv3d_implicit + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +_ARCH = get_rocm_arch() +# mfma_f32_16x16x32_bf16 is only available on CDNA4 (gfx95x) +_skip_non_cdna4 = pytest.mark.skipif( + not (isinstance(_ARCH, str) and _ARCH.startswith("gfx95")), + reason=f"conv3d BF16 needs mfma_f32_16x16x32_bf16 (CDNA4 gfx95x), got {_ARCH}", +) + + +# (N, C, T, H, W, K), kernel 3x3x3. Covers stride/padding and tile-tail paths. +@_skip_non_cdna4 +@pytest.mark.parametrize( + "n,c,t,h,w,k,stride,padding", + [ + (1, 32, 8, 16, 16, 64, 1, 0), + (1, 32, 9, 17, 17, 96, 1, 1), + (2, 64, 6, 18, 18, 192, 1, 1), + (1, 32, 10, 20, 20, 64, 2, 1), + # Partial K-tile: C=16 -> CRS=432, 432 % TILE_K(32) = 16 (masked). + (1, 16, 6, 16, 20, 16, 1, 1), + (1, 16, 4, 12, 16, 384, 1, 1), + ], +) +def test_conv3d_vs_torch(n, c, t, h, w, k, stride, padding): + torch.manual_seed(2000 + h + w + k) + x = torch.randn((n, c, t, h, w), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((k, c, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + bias = torch.randn((k,), device="cuda", dtype=torch.float32) + + y = conv3d_implicit(x, weight, bias=bias, stride=stride, padding=padding) + y_ref = F.conv3d(x, weight, bias=bias.to(torch.bfloat16), stride=stride, padding=padding) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape + assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) + + +@_skip_non_cdna4 +@pytest.mark.parametrize( + "kernel_shape,padding", + [ + ((1, 3, 3), (0, 1, 1)), + ((3, 1, 1), (1, 0, 0)), + ], +) +def test_conv3d_factorized_filters_vs_torch(kernel_shape, padding): + """Cover the spatial-only and temporal-only filter dispatch paths.""" + torch.manual_seed(3100 + sum(kernel_shape)) + n, c, t, h, w, k = 1, 64, 6, 18, 20, 128 + x = torch.randn((n, c, t, h, w), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((k, c, *kernel_shape), device="cuda", dtype=torch.bfloat16) + + y = conv3d_implicit(x, weight, stride=1, padding=padding) + y_ref = F.conv3d(x, weight, stride=1, padding=padding) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape + assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) + + +@_skip_non_cdna4 +@pytest.mark.parametrize("c", [16, 64]) +def test_conv3d_runtime_k_loop_short_problems(c): + """Exercise one- and two-K-tile runtime-pipeline epilogues.""" + torch.manual_seed(3200 + c) + n, t, h, w, k = 1, 3, 8, 8, 64 + x = torch.randn((n, c, t, h, w), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((k, c, 1, 1, 1), device="cuda", dtype=torch.bfloat16) + + y = conv3d_implicit(x, weight) + y_ref = F.conv3d(x, weight) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape + assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) + + +# Tile-size sweep: each forced (TILE_M, TILE_N, WAVE_M, WAVE_N) must stay correct. +@_skip_non_cdna4 +@pytest.mark.parametrize( + "tile", + [ + (128, 128, 2, 4), # default + (128, 256, 2, 4), + (256, 128, 2, 4), + (256, 256, 2, 4), + (256, 256, 4, 4), + (128, 128, 4, 2), + (64, 128, 1, 4), + (64, 64, 2, 2), + ], +) +def test_conv3d_tile_configs(tile): + torch.manual_seed(4000 + sum(tile)) + n, c, t, h, w, k, stride, padding = 2, 64, 6, 18, 18, 192, 1, 1 + x = torch.randn((n, c, t, h, w), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((k, c, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + bias = torch.randn((k,), device="cuda", dtype=torch.float32) + + y = conv3d_implicit(x, weight, bias=bias, stride=stride, padding=padding, tile=tile) + y_ref = F.conv3d(x, weight, bias=bias.to(torch.bfloat16), stride=stride, padding=padding) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape + assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) + + +@_skip_non_cdna4 +def test_conv3d_autotune(tmp_path, monkeypatch): + monkeypatch.setenv("FLYDSL_AUTOTUNE_CACHE_DIR", str(tmp_path / "at")) + from kernels.conv import conv3d_implicit_autotune + + conv3d_implicit_autotune._MEM_CACHE.clear() + + torch.manual_seed(4242) + n, c, t, h, w, k = 1, 128, 6, 40, 40, 128 + x = torch.randn((n, c, t, h, w), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((k, c, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + + y = conv3d_implicit(x, weight, stride=1, padding=1, autotune=True) + y_ref = F.conv3d(x, weight, stride=1, padding=1) + torch.cuda.synchronize() + assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) + + # A tile was chosen and persisted; the second call must hit the cache. + assert len(conv3d_implicit_autotune._MEM_CACHE) == 1 + calls = {"n": 0} + orig = conv3d_implicit_autotune.do_bench + + def _counting(*a, **kw): + calls["n"] += 1 + return orig(*a, **kw) + + monkeypatch.setattr(conv3d_implicit_autotune, "do_bench", _counting) + y2 = conv3d_implicit(x, weight, stride=1, padding=1, autotune=True) + torch.cuda.synchronize() + assert torch.allclose(y2, y_ref, rtol=2e-2, atol=2e-2) + assert calls["n"] == 0 # cached, no re-benchmark + + +# 2D conv via the depth-1 degenerate path through the 3D kernel. +@_skip_non_cdna4 +@pytest.mark.parametrize( + "kernel_shape,stride,padding", + [ + ((3, 3), 1, 1), + ((1, 1), 1, 0), # 1x1 -> temporal_only_fast-style vectorized epilogue + ((5, 5), 1, 2), + ((3, 3), 2, 1), + ], +) +def test_conv2d_vs_torch(kernel_shape, stride, padding): + torch.manual_seed(5000 + sum(kernel_shape) + stride + padding) + n, c, h, w, k = 2, 64, 24, 28, 128 + x = torch.randn((n, c, h, w), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((k, c, *kernel_shape), device="cuda", dtype=torch.bfloat16) + bias = torch.randn((k,), device="cuda", dtype=torch.float32) + + y = conv3d_implicit(x, weight, bias=bias, stride=stride, padding=padding) + y_ref = F.conv2d(x, weight, bias=bias.to(torch.bfloat16), stride=stride, padding=padding) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape + assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) + + +# 1D conv via the depth/height-1 degenerate path through the 3D kernel. +@_skip_non_cdna4 +@pytest.mark.parametrize( + "s,stride,padding", + [ + (3, 1, 1), + (1, 1, 0), + (5, 2, 2), + ], +) +def test_conv1d_vs_torch(s, stride, padding): + torch.manual_seed(6000 + s + stride + padding) + n, c, w, k = 2, 64, 96, 128 + x = torch.randn((n, c, w), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((k, c, s), device="cuda", dtype=torch.bfloat16) + bias = torch.randn((k,), device="cuda", dtype=torch.float32) + + y = conv3d_implicit(x, weight, bias=bias, stride=stride, padding=padding) + y_ref = F.conv1d(x, weight, bias=bias.to(torch.bfloat16), stride=stride, padding=padding) + torch.cuda.synchronize() + + assert y.shape == y_ref.shape + assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) diff --git a/tests/kernels/test_conv3d_implicit_8wave.py b/tests/kernels/test_conv3d_implicit_8wave.py deleted file mode 100644 index 9d05be183..000000000 --- a/tests/kernels/test_conv3d_implicit_8wave.py +++ /dev/null @@ -1,56 +0,0 @@ -#!/usr/bin/env python3 - -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""Correctness test for the bf16 8-wave implicit-GEMM conv3d kernel. - -Compares ``conv3d_implicit_8wave`` against ``torch.nn.functional.conv3d`` on -NCDHW/OIDHW bf16 inputs across stride/padding and M%TILE_M / K%TILE_N tail paths. -Channels must satisfy the kernel's ``c % 8 == 0`` and ``crs = c*kt*kh*kw`` a -multiple of TILE_K (32) constraints. -""" - -import pytest -import torch -import torch.nn.functional as F - -from flydsl.runtime.device import get_rocm_arch -from kernels.conv.conv3d_implicit_8wave import conv3d_implicit_8wave - -pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] - -_ARCH = get_rocm_arch() -# mfma_f32_16x16x32_bf16 is only available on CDNA4 (gfx95x) -_skip_non_cdna4 = pytest.mark.skipif( - not (isinstance(_ARCH, str) and _ARCH.startswith("gfx95")), - reason=f"conv3d 8-wave BF16 needs mfma_f32_16x16x32_bf16 (CDNA4 gfx95x), got {_ARCH}", -) - - -# (N, C, T, H, W, K), kernel 3x3x3. Covers stride/padding and tile-tail paths. -@_skip_non_cdna4 -@pytest.mark.parametrize( - "n,c,t,h,w,k,stride,padding", - [ - (1, 32, 8, 16, 16, 64, 1, 0), - (1, 32, 9, 17, 17, 96, 1, 1), - (2, 64, 6, 18, 18, 192, 1, 1), - (1, 32, 10, 20, 20, 64, 2, 1), - # Partial K-tile: C=16 -> CRS=432, 432 % TILE_K(32) = 16 (masked). - (1, 16, 6, 16, 20, 16, 1, 1), - (1, 16, 4, 12, 16, 384, 1, 1), - ], -) -def test_conv3d_vs_torch(n, c, t, h, w, k, stride, padding): - torch.manual_seed(2000 + h + w + k) - x = torch.randn((n, c, t, h, w), device="cuda", dtype=torch.bfloat16) - weight = torch.randn((k, c, 3, 3, 3), device="cuda", dtype=torch.bfloat16) - bias = torch.randn((k,), device="cuda", dtype=torch.float32) - - y = conv3d_implicit_8wave(x, weight, bias=bias, stride=stride, padding=padding) - y_ref = F.conv3d(x, weight, bias=bias.to(torch.bfloat16), stride=stride, padding=padding) - torch.cuda.synchronize() - - assert y.shape == y_ref.shape - assert torch.allclose(y, y_ref, rtol=2e-2, atol=2e-2) diff --git a/tests/kernels/test_conv3d_implicit_8wave_fp8.py b/tests/kernels/test_conv3d_implicit_fp8.py similarity index 94% rename from tests/kernels/test_conv3d_implicit_8wave_fp8.py rename to tests/kernels/test_conv3d_implicit_fp8.py index 9f3dde544..7064012da 100644 --- a/tests/kernels/test_conv3d_implicit_8wave_fp8.py +++ b/tests/kernels/test_conv3d_implicit_fp8.py @@ -17,7 +17,7 @@ import torch.nn.functional as F from flydsl.runtime.device import get_rocm_arch -from kernels.conv.conv3d_implicit_8wave_fp8 import conv3d_implicit_8wave_fp8 +from kernels.conv.conv3d_implicit_fp8 import conv3d_implicit_fp8 pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] @@ -48,7 +48,7 @@ def test_conv3d_fp8_vs_fp8cast_reference(n, c, t, h, w, k, stride, padding): x = torch.randn((n, c, t, h, w), device="cuda", dtype=torch.bfloat16) weight = torch.randn((k, c, 3, 3, 3), device="cuda", dtype=torch.bfloat16) - y = conv3d_implicit_8wave_fp8(x, weight, stride=stride, padding=padding) + y = conv3d_implicit_fp8(x, weight, stride=stride, padding=padding) ref = F.conv3d( x.to(torch.float8_e4m3fn).to(torch.bfloat16), weight.to(torch.float8_e4m3fn).to(torch.bfloat16),