refactor(conv): split conv_kernel into a package and fuse the backward fp16 casts

aether/custom_kernels/conv_kernel/__init__.py

@@ -0,0 +1,53 @@
+"""Matrix-core (WMMA) convolution kernels, one module per kernel.
+
+common   -- shape metadata, launch plans, tile constants, vendor dicts
+forward  -- conv_forward_wmma
+dweight  -- conv_backward_dweight_wmma and its split-K sizing
+dinput   -- conv_backward_dinput_wmma
+cast     -- conv_backward_cast_fp16, the fp16 operands for both backward kernels
+"""
+from aether.custom_kernels.conv_kernel.cast import (
+    _CONV_BACKWARD_CAST_CACHE,
+    cast_backward_operands_fp16,
+    get_compiled_backward_cast_kernel,
+)
+from aether.custom_kernels.conv_kernel.common import (
+    ConvLaunchPlan,
+    ConvShapeMeta,
+    get_is_conv_gpu_available,
+)
+from aether.custom_kernels.conv_kernel.dinput import (
+    _CONV_BACKWARD_DINPUT_CACHE,
+    get_compiled_backward_dinput_kernel,
+    plan_dinput_launch,
+)
+from aether.custom_kernels.conv_kernel.dweight import (
+    _CONV_BACKWARD_DWEIGHT_CACHE,
+    choose_dweight_split_k,
+    get_compiled_backward_dweight_kernel,
+    plan_dweight_launch,
+)
+from aether.custom_kernels.conv_kernel.forward import (
+    _CONV_FORWARD_KERNEL_CACHE,
+    get_compiled_forward_conv_kernel,
+    plan_forward_launch,
+)
+
+__all__ = [
+    "_CONV_BACKWARD_CAST_CACHE",
+    "_CONV_BACKWARD_DINPUT_CACHE",
+    "_CONV_BACKWARD_DWEIGHT_CACHE",
+    "_CONV_FORWARD_KERNEL_CACHE",
+    "ConvLaunchPlan",
+    "ConvShapeMeta",
+    "cast_backward_operands_fp16",
+    "choose_dweight_split_k",
+    "get_compiled_backward_cast_kernel",
+    "get_compiled_backward_dinput_kernel",
+    "get_compiled_backward_dweight_kernel",
+    "get_compiled_forward_conv_kernel",
+    "get_is_conv_gpu_available",
+    "plan_dinput_launch",
+    "plan_dweight_launch",
+    "plan_forward_launch",
+]

aether/custom_kernels/conv_kernel/cast.py

@@ -0,0 +1,78 @@
+import warnings
+
+import aether.config as config
+from aether.custom_kernels.conv_kernel.common import _COMPILE_OPTIONS
+
+# One launch casts both fp16 operands of the conv backward kernels. dvalues is
+# clipped to the fp16 range first; inputs is not, so dweight sees exactly the
+# fp16 inputs the forward pass multiplied.
+CONV_BACKWARD_CAST_SOURCE = r'''
+$fp16_include
+
+#define HALF_MAX 65504.0f
+
+extern "C" __global__ void conv_backward_cast_fp16(
+    const float* __restrict__ dvalues,     // [n_dvalues] fp32, contiguous
+    half* __restrict__ dvalues_fp16,       // [n_dvalues]
+    const long long n_dvalues,
+    const float* __restrict__ inputs,      // [n_inputs] fp32, contiguous
+    half* __restrict__ inputs_fp16,        // [n_inputs]
+    const long long n_inputs
+) {
+    const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n_dvalues) {
+        const float v = dvalues[i];
+        // Comparisons rather than fminf/fmaxf, which would turn NaN into
+        // +-HALF_MAX: like xp.clip, NaN passes through and +-inf clamp.
+        dvalues_fp16[i] = __float2half(v > HALF_MAX ? HALF_MAX : (v < -HALF_MAX ? -HALF_MAX : v));
+    } else if (i < n_dvalues + n_inputs) {
+        const long long j = i - n_dvalues;
+        inputs_fp16[j] = __float2half(inputs[j]);
+    }
+}
+'''
+
+_CONV_BACKWARD_CAST_CACHE = {}
+
+
+def get_compiled_backward_cast_kernel(is_hip: bool):
+    if is_hip in _CONV_BACKWARD_CAST_CACHE:
+        return _CONV_BACKWARD_CAST_CACHE[is_hip]
+
+    fp16_include = "#include <hip/hip_fp16.h>" if is_hip else "#include <cuda_fp16.h>"
+    source = CONV_BACKWARD_CAST_SOURCE.replace("$fp16_include", fp16_include)
+
+    try:
+        module = config.cp.RawModule(code=source, options=_COMPILE_OPTIONS)
+        kernel = module.get_function('conv_backward_cast_fp16')
+    except (config.cp.cuda.compiler.CompileException, config.cp.cuda.driver.CUDADriverError) as e:
+        warnings.warn(f"[aether] conv_kernel: backward fp16 cast compile failed: {e}")
+        _CONV_BACKWARD_CAST_CACHE[is_hip] = None
+        return None
+
+    _CONV_BACKWARD_CAST_CACHE[is_hip] = kernel
+    return kernel
+
+
+def cast_backward_operands_fp16(dvalues, inputs):
+    """(dvalues clipped to +-65504, inputs) as fp16, in one launch.
+
+    Both must be contiguous fp32. Returns None when the kernel is unavailable,
+    so the caller can fall back to astype/clip.
+    """
+    kernel = get_compiled_backward_cast_kernel(config.IS_HIP)
+    if kernel is None:
+        return None
+
+    xp = config.cp
+    dvalues_fp16 = xp.empty(dvalues.shape, dtype=xp.float16)
+    inputs_fp16 = xp.empty(inputs.shape, dtype=xp.float16)
+    n_dvalues, n_inputs = dvalues.size, inputs.size
+
+    _, threads = config.resolve_gpu_launch_geometry()
+    blocks = -(-(n_dvalues + n_inputs) // threads)
+    kernel(
+        (blocks,), (threads,),
+        (dvalues, dvalues_fp16, xp.int64(n_dvalues), inputs, inputs_fp16, xp.int64(n_inputs)),
+    )
+    return dvalues_fp16, inputs_fp16

aether/custom_kernels/conv_kernel/common.py

@@ -0,0 +1,148 @@
+from dataclasses import dataclass
+
+import numpy as np
+
+import aether.config as config
+
+
+def get_is_conv_gpu_available():
+    """Check if GPU supports optimized tensor-core convolution kernels."""
+    return config.get_tensor_core_capable()
+
+
+# --- Shape metadata / cache key --------------------------------------------
+
+@dataclass(frozen=True)
+class ConvShapeMeta:
+    H_in: int
+    W_in: int
+    C_in: int
+    C_out: int
+    fH: int
+    fW: int
+    sH: int
+    sW: int
+    pad_h: int
+    pad_w: int
+    H_out: int
+    W_out: int
+    is_hip: bool
+
+    @property
+    def K_total(self):
+        return self.fH * self.fW * self.C_in
+
+    # pad_h/pad_w are the top/left pads the kernels offset by; the bottom/right
+    # pads are never materialized there, the kernels bounds-check them away.
+    @property
+    def pad_bottom(self):
+        return (self.H_out - 1) * self.sH + self.fH - self.H_in - self.pad_h
+
+    @property
+    def pad_right(self):
+        return (self.W_out - 1) * self.sW + self.fW - self.W_in - self.pad_w
+
+    @property
+    def cache_key(self):
+        return (
+            self.H_in, self.W_in, self.C_in, self.C_out,
+            self.fH, self.fW, self.sH, self.sW,
+            # H_out/W_out are not implied by the rest: 'same' and 'valid' can
+            # share a pad yet differ in output extent.
+            self.pad_h, self.pad_w, self.H_out, self.W_out, self.is_hip,
+        )
+
+    @classmethod
+    def create(cls, spatial_shape, C_out, filter_size, stride, padding, is_hip=None):
+        """Derive padding and output extents from a layer's configuration."""
+        H_in, W_in, C_in = spatial_shape
+        fH, fW = filter_size
+        sH, sW = stride
+
+        if padding == "same":
+            # TF convention: total pad covers ceil(in / stride) windows; an odd
+            # total puts the extra pixel on the bottom/right.
+            H_out, W_out = -(-H_in // sH), -(-W_in // sW)
+            pad_h = max((H_out - 1) * sH + fH - H_in, 0) // 2
+            pad_w = max((W_out - 1) * sW + fW - W_in, 0) // 2
+        elif padding == "valid":
+            H_out, W_out = (H_in - fH) // sH + 1, (W_in - fW) // sW + 1
+            pad_h, pad_w = 0, 0
+        else:
+            raise ValueError(
+                f"Unsupported padding mode '{padding}'. Use 'same' or 'valid'."
+            )
+
+        return cls(
+            H_in=H_in, W_in=W_in, C_in=C_in, C_out=C_out,
+            fH=fH, fW=fW, sH=sH, sW=sW, pad_h=pad_h, pad_w=pad_w,
+            H_out=H_out, W_out=W_out,
+            is_hip=config.IS_HIP if is_hip is None else is_hip,
+        )
+
+
+_BLOCK_TILE_M = 64
+_BLOCK_TILE_N = 64
+_WMMA_DIM = 16
+_WARP_SIZE = 32
+
+# RawModule options shared by every conv kernel. No -I for ROCm: hiprtc already
+# searches its own install's include dir (where rocWMMA lives), and a -I from
+# ROCM_PATH would be searched first, so a stale ROCM_PATH could shadow it.
+_COMPILE_OPTIONS = ('-std=c++17',)
+
+
+# --- Vendor substitution dictionaries --------------------------------------
+
+_CONV_CUDA = dict(
+    hip_include="#include <mma.h>\n#include <cuda_fp16.h>",
+    wmma_ns="nvcuda::wmma",
+    warp_size=_WARP_SIZE,
+    frag_layout_a="nvcuda::wmma::row_major",
+    frag_layout_a_col="nvcuda::wmma::col_major",
+    frag_layout_b="nvcuda::wmma::row_major",
+    mem_layout="nvcuda::wmma::mem_row_major",
+    matrix_a_tag="nvcuda::wmma::matrix_a",
+    matrix_b_tag="nvcuda::wmma::matrix_b",
+    accumulator_tag="nvcuda::wmma::accumulator",
+)
+
+_CONV_HIP = dict(
+    hip_include="#include <hip/hip_fp16.h>\n#include <rocwmma/rocwmma.hpp>",
+    wmma_ns="rocwmma",
+    warp_size=_WARP_SIZE,
+    frag_layout_a="rocwmma::row_major",
+    frag_layout_a_col="rocwmma::col_major",
+    frag_layout_b="rocwmma::row_major",
+    mem_layout="rocwmma::mem_row_major",
+    matrix_a_tag="rocwmma::matrix_a",
+    matrix_b_tag="rocwmma::matrix_b",
+    accumulator_tag="rocwmma::accumulator",
+)
+
+
+# --- Launch planning --------------------------------------------------------
+
+@dataclass(frozen=True)
+class ConvLaunchPlan:
+    """A compiled kernel bound to a concrete grid, block and M_total."""
+    kernel: object
+    grid: tuple
+    block: tuple
+    M_total: object
+    shape_meta: ConvShapeMeta
+
+    def __call__(self, *tensors):
+        """Launch. Every conv kernel takes M_total as its trailing argument."""
+        self.kernel(self.grid, self.block, (*tensors, self.M_total))
+
+    @property
+    def num_blocks(self):
+        n = 1
+        for d in self.grid:
+            n *= d
+        return n
+
+
+def _m_total(shape_meta: ConvShapeMeta, S: int):
+    return np.int64(S * shape_meta.H_out * shape_meta.W_out)