From 7e9b8d56a260d0815f1d6b61649feeb62d975305 Mon Sep 17 00:00:00 2001 From: Greg Stoner Date: Sat, 30 May 2026 06:57:17 -0600 Subject: [PATCH] Apple GPU Tier-3: conv3d via im2col + GPU batched matmul MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MPSGraph has no 3-D convolution node, so conv3d lowers to the classic im2col + GEMM decomposition with the dominant GEMM on-GPU. - apple_gpu_runtime.mm: conv3d_core_f32 (host im2col into a per-group [groups, rows, K] column matrix, weights regrouped to [groups, K, Cout/groups], bias + scatter on host) + mpsg_conv3d_batched_matmul_f32 (a single cached MPSGraph batched matmul, batch = groups, fp32 accumulation, RAII buffer-pool acquires) + tessera_apple_gpu_conv3d_{f32,f16} + conv3d_out_dim. NDHWC source, DHWIO weights, full stride/pad/dilation/groups, optional bias. The host GEMM is used as a fallback when Metal is unavailable. - apple_gpu_runtime_stub.cpp: non-Apple parity (f32 reference conv3d; f16 zero stub so python upcasts on fallback). - runtime.py: _apple_gpu_conv3d_{f32,f16} ctypes wrappers + a _apple_gpu_dispatch_conv3d dispatcher (f32/f16 native, bf16 host round-trip, None-fallback); tessera.conv3d added to _APPLE_GPU_CONV_OPS + the runtime envelope + the metadata op-dispatch loop (conv2d/conv3d split by rank). - driver.py: tessera.conv3d gated into _APPLE_GPU_CONV_OPS so a single conv3d plan reports execution_mode="metal_runtime". - tests/unit/test_apple_gpu_conv3d.py: 12 tests — f32 (6 shape/stride/pad/ dilation/groups/depthwise cases) vs numpy, native f16, bf16 round-trip, symbol export, envelope membership, agreement with eager ops.conv3d. - docs: plan conv3d row marked done; runtime_abi dashboard regenerated (113 -> 116 symbols). Co-Authored-By: Claude Opus 4.8 --- docs/apple_gpu_tier2_tier3_plan.md | 2 +- docs/audit/generated/runtime_abi.md | 7 +- python/tessera/compiler/driver.py | 5 +- python/tessera/runtime.py | 108 +++++++++- .../runtime/apple_gpu_runtime.mm | 198 ++++++++++++++++++ .../runtime/apple_gpu_runtime_stub.cpp | 66 ++++++ tests/unit/test_apple_gpu_conv3d.py | 146 +++++++++++++ 7 files changed, 523 insertions(+), 9 deletions(-) create mode 100644 tests/unit/test_apple_gpu_conv3d.py diff --git a/docs/apple_gpu_tier2_tier3_plan.md b/docs/apple_gpu_tier2_tier3_plan.md index af28bd6ae..e69a1944c 100644 --- a/docs/apple_gpu_tier2_tier3_plan.md +++ b/docs/apple_gpu_tier2_tier3_plan.md @@ -45,7 +45,7 @@ MPSGraph nodes cover most of what the original framing assumed needed bespoke MS | **Reductions** (`sum`/`mean`/`var`/`std`/`amax`/`amin`/`prod`/`argmax`/`argmin`/`cumsum`/`cumprod`) | **DONE** — `tessera_apple_gpu_mpsgraph_{reduce,argreduce,scan}_f32`; `runtime.py` normalizes arbitrary axis/keepdims/ddof by folding reduced axes to the last dim. `tests/unit/test_apple_gpu_reductions.py` (51). | Low–Med | **Done** | | **`dropout`, `rng_normal`, `rng_uniform`** | (a) MPSGraph random nodes — quick but **won't bit-match the CPU Philox stream** (breaks Decision #18); (b) hand-written Philox MSL for bit-exactness. | Med / Low | **Defer** (training-side); MSL if pursued | | **`conv2d`** | **DONE** — `tessera_apple_gpu_conv2d_{f32,f16}` via MPSGraph `convolution2DWithSourceTensor:weightsTensor:descriptor:` (NHWC source / HWIO weights, full stride/pad/dilation/groups, optional bias, fp32 internal accumulation; bf16 via host fp32 round-trip). Wired into the metadata op-loop + `_APPLE_GPU_CONV_OPS` envelope + driver gating; reference fallback in the stub. `tests/unit/test_apple_gpu_conv2d.py` (13). | Med | **Done** | -| **`conv3d`** | MPSGraph has no 3-D conv node → im2col + `bmm` fallback. | High | **Defer** unless 3-D vision enters scope | +| **`conv3d`** | **DONE** — MPSGraph has no 3-D conv node, so `tessera_apple_gpu_conv3d_{f32,f16}` lower to im2col + a single GPU MPSGraph **batched matmul** (batch = groups, fp32 accumulation): patches gathered to a per-group `[groups, rows, K]` column matrix, weights regrouped to `[groups, K, Cout/groups]`, the dominant GEMM on-GPU; bias + scatter on the host. NDHWC source / DHWIO weights, full stride/pad/dilation/groups, optional bias; bf16 via host fp32 round-trip. Wired into the metadata op-loop + `_APPLE_GPU_CONV_OPS` envelope + driver gating; reference fallback in the stub. `tests/unit/test_apple_gpu_conv3d.py` (12). | High | **Done** | ## Cross-cutting diff --git a/docs/audit/generated/runtime_abi.md b/docs/audit/generated/runtime_abi.md index b3b6ed0c3..0b1e7440c 100644 --- a/docs/audit/generated/runtime_abi.md +++ b/docs/audit/generated/runtime_abi.md @@ -4,9 +4,9 @@ Generated from `python/tessera/compiler/runtime_abi_audit.py`. Don't edit by ha ## Headline -- **113** unique `extern "C" tessera_*` C ABI symbols across all backends. +- **116** unique `extern "C" tessera_*` C ABI symbols across all backends. - **6 / 6** core runtime headers present. -- **55** Apple GPU kernel families with per-dtype variants. +- **56** Apple GPU kernel families with per-dtype variants. ## Core runtime headers @@ -23,7 +23,7 @@ Generated from `python/tessera/compiler/runtime_abi_audit.py`. Don't edit by ha | Backend | Unique tessera_* symbols | |---------|-------------------------:| -| `apple` | 102 | +| `apple` | 105 | | `nvidia` | 3 | | `x86` | 8 | @@ -51,6 +51,7 @@ Generated from `python/tessera/compiler/runtime_abi_audit.py`. Don't edit by ha | `complex_mul` | `f32` | | `complex_stereographic` | `f32` | | `conv2d` | `f16`, `f32` | +| `conv3d` | `f16`, `f32` | | `ebm_decode_init_noise_apply` | `f32` | | `ebm_ebt_tiny_refinement_argmin` | `f32` | | `ebm_energy_quadratic` | `f32` | diff --git a/python/tessera/compiler/driver.py b/python/tessera/compiler/driver.py index 633dcf445..b9866491a 100644 --- a/python/tessera/compiler/driver.py +++ b/python/tessera/compiler/driver.py @@ -427,8 +427,9 @@ def _is_apple_cpu_accelerate_executable(cpu_plan: CPUPlan | None) -> bool: "tessera.argmin", "tessera.cumsum", "tessera.cumprod", }) -# 2026-05-30 — Tier-3 conv2d via the MPSGraph convolution2D node (NHWC/HWIO). -_APPLE_GPU_CONV_OPS: frozenset[str] = frozenset({"tessera.conv2d"}) +# 2026-05-30 — Tier-3 convolutions: conv2d via the MPSGraph convolution2D node +# (NHWC/HWIO); conv3d via im2col + a GPU MPSGraph batched matmul (NDHWC/DHWIO). +_APPLE_GPU_CONV_OPS: frozenset[str] = frozenset({"tessera.conv2d", "tessera.conv3d"}) _APPLE_GPU_RUNTIME_OPS: frozenset[str] = ( _APPLE_GPU_MPS_OPS | _APPLE_GPU_MSL_OPS | _APPLE_GPU_MPSGRAPH_OPS diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index d6a10dc96..ef6c906ad 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -1897,8 +1897,9 @@ def _load_apple_cpu_runtime() -> ctypes.CDLL: "tessera.cumprod": ("scan", 1), } _APPLE_GPU_REDUCTION_OPS = frozenset(_APPLE_GPU_REDUCE_OPS) -# 2026-05-30 — Tier-3 conv2d via the MPSGraph convolution2D node (NHWC/HWIO). -_APPLE_GPU_CONV_OPS = frozenset({"tessera.conv2d"}) +# 2026-05-30 — Tier-3 convolutions: conv2d via the MPSGraph convolution2D node +# (NHWC/HWIO); conv3d via im2col + a GPU MPSGraph batched matmul (NDHWC/DHWIO). +_APPLE_GPU_CONV_OPS = frozenset({"tessera.conv2d", "tessera.conv3d"}) _APPLE_GPU_RUNTIME_OPS = ( _APPLE_GPU_MPS_OPS | _APPLE_GPU_MSL_OPS | _APPLE_GPU_MPSGRAPH_OPS | _APPLE_GPU_PROJECTION_OPS | _APPLE_GPU_REDUCTION_OPS | _APPLE_GPU_CONV_OPS @@ -2106,12 +2107,18 @@ def _execute_apple_gpu_mps_metadata(metadata: Mapping[str, Any], args: Any) -> A kwargs, np, ) - elif op_name in _APPLE_GPU_CONV_OPS: + elif op_name == "tessera.conv2d": values[str(result)] = _apple_gpu_dispatch_conv2d( [_as_numpy(values[name]) for name in operand_names], kwargs, np, ) + elif op_name == "tessera.conv3d": + values[str(result)] = _apple_gpu_dispatch_conv3d( + [_as_numpy(values[name]) for name in operand_names], + kwargs, + np, + ) else: # Phase 8.4.x will broaden further; today single-op gating in # driver.py is the authoritative envelope. A non-MPS, non-MSL op @@ -2691,6 +2698,101 @@ def _pair(v: Any) -> tuple[int, int]: return out.astype(out_dtype) +def _apple_gpu_conv3d_f32() -> Any: + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_conv3d_f32", None) + if sym is None: + return None + sym.argtypes = [ctypes.POINTER(ctypes.c_float)] * 4 + [ctypes.c_int32] * 19 + sym.restype = None + return sym + + +def _apple_gpu_conv3d_f16() -> Any: + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_conv3d_f16", None) + if sym is None: + return None + sym.argtypes = [ctypes.POINTER(ctypes.c_uint16)] * 4 + [ctypes.c_int32] * 19 + sym.restype = None + return sym + + +def _apple_gpu_dispatch_conv3d(operands: list[Any], kwargs: dict, np: Any) -> Any: + """Tier-3 3-D convolution via im2col + a GPU MPSGraph batched matmul + (NDHWC source, DHWIO weights). + + X is [N, D, H, W, Cin]; weight is [kD, kH, kW, Cin/groups, Cout]; optional + bias is [Cout]; output is [N, oD, oH, oW, Cout]. ``stride``/``padding``/ + ``dilation`` accept an int or a 3-tuple; ``groups`` defaults to 1. f32/f16 + run natively (fp32 GEMM accumulation); bf16 runs via a host fp32 round-trip; + any other dtype (or an unavailable runtime) returns None so the caller falls + back to the numpy reference.""" + X = np.asarray(operands[0]) + W = np.asarray(operands[1]) + bias = None + if len(operands) > 2 and operands[2] is not None: + bias = np.asarray(operands[2]) + if X.ndim != 5 or W.ndim != 5: + return None + + def _triple(v: Any) -> tuple[int, int, int]: + if isinstance(v, (tuple, list)): + return int(v[0]), int(v[1]), int(v[2]) + return int(v), int(v), int(v) + + sD, sH, sW = _triple(kwargs.get("stride", 1)) + pD, pH, pW = _triple(kwargs.get("padding", 0)) + dD, dH, dW = _triple(kwargs.get("dilation", 1)) + groups = int(kwargs.get("groups", 1)) + N, iD, iH, iW, Cin = (int(s) for s in X.shape) + kD, kH, kW, cinG, Cout = (int(s) for s in W.shape) + if groups <= 0 or Cin % groups or Cout % groups or cinG != Cin // groups: + return None + + def _out(i: int, k: int, s: int, p: int, d: int) -> int: + return (i + 2 * p - d * (k - 1) - 1) // s + 1 + + oD = _out(iD, kD, sD, pD, dD) + oH = _out(iH, kH, sH, pH, dH) + oW = _out(iW, kW, sW, pW, dW) + if oD <= 0 or oH <= 0 or oW <= 0: + return None + out_dtype = X.dtype + bf16 = _bfloat16_dtype() + iattrs = [ctypes.c_int32(v) for v in + (N, iD, iH, iW, Cin, Cout, kD, kH, kW, sD, sH, sW, pD, pH, pW, + dD, dH, dW, groups)] + + if out_dtype == np.float16: + sym = _apple_gpu_conv3d_f16() + if sym is None: + return None + up = lambda a: a.ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)) + xh = np.ascontiguousarray(X).view(np.uint16) + wh = np.ascontiguousarray(W).view(np.uint16) + bh = (np.ascontiguousarray(bias).view(np.uint16) + if bias is not None else None) + out = np.zeros((N, oD, oH, oW, Cout), dtype=np.uint16) + sym(up(xh), up(wh), up(bh) if bh is not None else None, up(out), *iattrs) + return out.view(np.float16) + + is_bf16 = bf16 is not None and out_dtype == bf16 + is_f32 = out_dtype == np.float32 + if not (is_f32 or is_bf16): + return None + sym = _apple_gpu_conv3d_f32() + if sym is None: + return None + fp = lambda a: a.ctypes.data_as(ctypes.POINTER(ctypes.c_float)) + xf = np.ascontiguousarray(X.astype(np.float32)) + wf = np.ascontiguousarray(W.astype(np.float32)) + bf = np.ascontiguousarray(bias.astype(np.float32)) if bias is not None else None + out = np.zeros((N, oD, oH, oW, Cout), dtype=np.float32) + sym(fp(xf), fp(wf), fp(bf) if bf is not None else None, fp(out), *iattrs) + return out.astype(out_dtype) + + def _apple_gpu_flash_attn_gqa_f32() -> Any: runtime = _load_apple_gpu_runtime() sym = getattr(runtime, "tessera_apple_gpu_flash_attn_gqa_f32", None) diff --git a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm index 1c0d814bc..f7cdaa168 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm +++ b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm @@ -8980,3 +8980,201 @@ static void reference_conv2d_f32(const float *X, const float *Wt, if (outH > 0 && outW > 0) std::memset(O, 0, (size_t)N * outH * outW * Cout * 2); } + +//===----------------------------------------------------------------------===// +// conv3d — im2col + MPSGraph batched matmul (NDHWC source, DHWIO weights) +// (2026-05-30) +// +// MPSGraph has no 3-D convolution node, so conv3d is lowered to the classic +// im2col + GEMM decomposition: the spatial patches are gathered on the host +// into a column matrix laid out per-group as [groups, rows, K] (rows = +// N*oD*oH*oW, K = kD*kH*kW*Cin/groups), the weights are regrouped to +// [groups, K, Cout/groups], and the dominant GEMM runs on-GPU as a single +// MPSGraph batched matmul (fp32 accumulation). Bias + scatter back to NDHWC +// happen on the host. f16 I/O converts to fp32 at the boundary. +//===----------------------------------------------------------------------===// + +namespace { + +// GPU batched matmul A[g,M,K] @ B[g,K,Ncols] -> O[g,M,Ncols], fp32, cached. +static bool mpsg_conv3d_batched_matmul_f32(MetalDeviceContext &ctx, + const float *A, const float *B, + float *O, int32_t G, int32_t M, + int32_t K, int32_t Ncols) { + if (G <= 0 || M <= 0 || K <= 0 || Ncols <= 0) return true; + @autoreleasepool { + size_t aBytes = (size_t)G * M * K * 4; + size_t bBytes = (size_t)G * K * Ncols * 4; + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufA, ctx, A, aBytes); + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufB, ctx, B, bBytes); + if (!bufA || !bufB) return false; + NSArray *aShape = @[ @(G), @(M), @(K) ]; + NSArray *bShape = @[ @(G), @(K), @(Ncols) ]; + NSString *key = [NSString + stringWithFormat:@"conv3dmm:%d:%d:%d:%d", G, M, K, Ncols]; + NSArray *entry = mpsg_cache_get(key); + MPSGraph *g; + MPSGraphTensor *pa, *pb, *y; + if (entry) { + g = entry[0]; + pa = ((NSArray *)entry[1])[0]; + pb = ((NSArray *)entry[1])[1]; + y = entry[2]; + } else { + g = [MPSGraph new]; + pa = [g placeholderWithShape:aShape dataType:MPSDataTypeFloat32 name:nil]; + pb = [g placeholderWithShape:bShape dataType:MPSDataTypeFloat32 name:nil]; + y = [g matrixMultiplicationWithPrimaryTensor:pa secondaryTensor:pb name:nil]; + mpsg_cache_put(key, @[ g, @[ pa, pb ], y ]); + } + MPSGraphTensorData *ad = [[MPSGraphTensorData alloc] initWithMTLBuffer:bufA shape:aShape dataType:MPSDataTypeFloat32]; + MPSGraphTensorData *bd = [[MPSGraphTensorData alloc] initWithMTLBuffer:bufB shape:bShape dataType:MPSDataTypeFloat32]; + NSDictionary *res = [g runWithMTLCommandQueue:ctx.queue + feeds:@{pa : ad, pb : bd} + targetTensors:@[ y ] + targetOperations:nil]; + MPSGraphTensorData *od = res[y]; + if (!od) return false; + [[od mpsndarray] readBytes:O strideBytes:nil]; + return true; + } +} + +static inline int32_t conv3d_out_dim(int32_t in, int32_t k, int32_t stride, + int32_t pad, int32_t dilation) { + return conv2d_out_dim(in, k, stride, pad, dilation); +} + +// fp32 core: host im2col + GPU GEMM + host bias/scatter. on_gpu=false runs a +// pure-host GEMM (reference path). Returns false only on a hard GPU failure. +static bool conv3d_core_f32(MetalDeviceContext *ctx, const float *X, + const float *Wt, const float *bias, float *O, + int32_t N, int32_t iD, int32_t iH, int32_t iW, + int32_t Cin, int32_t Cout, int32_t kD, int32_t kH, + int32_t kW, int32_t sD, int32_t sH, int32_t sW, + int32_t pD, int32_t pH, int32_t pW, int32_t dD, + int32_t dH, int32_t dW, int32_t groups) { + int32_t oD = conv3d_out_dim(iD, kD, sD, pD, dD); + int32_t oH = conv3d_out_dim(iH, kH, sH, pH, dH); + int32_t oW = conv3d_out_dim(iW, kW, sW, pW, dW); + if (oD <= 0 || oH <= 0 || oW <= 0 || groups <= 0 || Cin % groups || + Cout % groups) + return true; + int32_t cinG = Cin / groups, coutG = Cout / groups; + int32_t K = kD * kH * kW * cinG; + int32_t rows = N * oD * oH * oW; + if (K <= 0 || rows <= 0) return true; + + // im2col -> cols[g, r, kk]; weights -> wg[g, kk, oc'] + std::vector cols((size_t)groups * rows * K, 0.0f); + std::vector wg((size_t)groups * K * coutG); + for (int32_t g = 0; g < groups; ++g) + for (int32_t kk = 0; kk < K; ++kk) + for (int32_t oc = 0; oc < coutG; ++oc) + wg[((size_t)g * K + kk) * coutG + oc] = + Wt[(size_t)kk * Cout + g * coutG + oc]; + + for (int32_t n = 0; n < N; ++n) + for (int32_t od = 0; od < oD; ++od) + for (int32_t oh = 0; oh < oH; ++oh) + for (int32_t ow = 0; ow < oW; ++ow) { + int32_t r = ((n * oD + od) * oH + oh) * oW + ow; + for (int32_t kd = 0; kd < kD; ++kd) { + int32_t id = od * sD + kd * dD - pD; + if (id < 0 || id >= iD) continue; + for (int32_t kh = 0; kh < kH; ++kh) { + int32_t ih = oh * sH + kh * dH - pH; + if (ih < 0 || ih >= iH) continue; + for (int32_t kw = 0; kw < kW; ++kw) { + int32_t iw = ow * sW + kw * dW - pW; + if (iw < 0 || iw >= iW) continue; + int32_t kbase = ((kd * kH + kh) * kW + kw) * cinG; + for (int32_t g = 0; g < groups; ++g) { + const float *xp = + X + ((((size_t)n * iD + id) * iH + ih) * iW + iw) * Cin + + g * cinG; + float *cp = cols.data() + + ((size_t)g * rows + r) * K + kbase; + for (int32_t ic = 0; ic < cinG; ++ic) cp[ic] = xp[ic]; + } + } + } + } + } + + std::vector mm((size_t)groups * rows * coutG); + bool ran = false; + if (ctx && ctx->ok) + ran = mpsg_conv3d_batched_matmul_f32(*ctx, cols.data(), wg.data(), + mm.data(), groups, rows, K, coutG); + if (!ran) { + for (int32_t g = 0; g < groups; ++g) + for (int32_t r = 0; r < rows; ++r) + for (int32_t oc = 0; oc < coutG; ++oc) { + double acc = 0; + const float *cp = cols.data() + ((size_t)g * rows + r) * K; + const float *wp = wg.data() + (size_t)g * K * coutG; + for (int32_t kk = 0; kk < K; ++kk) acc += (double)cp[kk] * wp[kk * coutG + oc]; + mm[((size_t)g * rows + r) * coutG + oc] = (float)acc; + } + } + + // scatter mm[g,r,oc'] (+ bias) -> O[n,od,oh,ow, g*coutG+oc'] + for (int32_t g = 0; g < groups; ++g) + for (int32_t r = 0; r < rows; ++r) + for (int32_t oc = 0; oc < coutG; ++oc) { + int32_t ocAbs = g * coutG + oc; + float v = mm[((size_t)g * rows + r) * coutG + oc]; + if (bias) v += bias[ocAbs]; + O[(size_t)r * Cout + ocAbs] = v; + } + return true; +} + +} // namespace + +extern "C" int32_t tessera_apple_gpu_conv3d_out_dim(int32_t in, int32_t k, + int32_t stride, int32_t pad, + int32_t dilation) { + return conv2d_out_dim(in, k, stride, pad, dilation); +} + +extern "C" void tessera_apple_gpu_conv3d_f32( + const float *X, const float *Wt, const float *bias, float *O, int32_t N, + int32_t iD, int32_t iH, int32_t iW, int32_t Cin, int32_t Cout, int32_t kD, + int32_t kH, int32_t kW, int32_t sD, int32_t sH, int32_t sW, int32_t pD, + int32_t pH, int32_t pW, int32_t dD, int32_t dH, int32_t dW, int32_t groups) { + MetalDeviceContext &ctx = deviceContext(); + conv3d_core_f32(ctx.ok ? &ctx : nullptr, X, Wt, bias, O, N, iD, iH, iW, Cin, + Cout, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, groups); +} + +extern "C" void tessera_apple_gpu_conv3d_f16( + const uint16_t *X, const uint16_t *Wt, const uint16_t *bias, uint16_t *O, + int32_t N, int32_t iD, int32_t iH, int32_t iW, int32_t Cin, int32_t Cout, + int32_t kD, int32_t kH, int32_t kW, int32_t sD, int32_t sH, int32_t sW, + int32_t pD, int32_t pH, int32_t pW, int32_t dD, int32_t dH, int32_t dW, + int32_t groups) { + int32_t oD = conv2d_out_dim(iD, kD, sD, pD, dD); + int32_t oH = conv2d_out_dim(iH, kH, sH, pH, dH); + int32_t oW = conv2d_out_dim(iW, kW, sW, pW, dW); + if (oD <= 0 || oH <= 0 || oW <= 0 || groups <= 0 || Cin % groups || + Cout % groups) + return; + size_t xn = (size_t)N * iD * iH * iW * Cin; + size_t wn = (size_t)kD * kH * kW * (Cin / groups) * Cout; + size_t on = (size_t)N * oD * oH * oW * Cout; + std::vector xf(xn), wf(wn), of(on); + std::vector bf; + for (size_t i = 0; i < xn; ++i) xf[i] = half_to_float_gpu(X[i]); + for (size_t i = 0; i < wn; ++i) wf[i] = half_to_float_gpu(Wt[i]); + if (bias) { + bf.resize(Cout); + for (int32_t i = 0; i < Cout; ++i) bf[i] = half_to_float_gpu(bias[i]); + } + MetalDeviceContext &ctx = deviceContext(); + conv3d_core_f32(ctx.ok ? &ctx : nullptr, xf.data(), wf.data(), + bias ? bf.data() : nullptr, of.data(), N, iD, iH, iW, Cin, + Cout, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, groups); + for (size_t i = 0; i < on; ++i) O[i] = float_to_half_gpu(of[i]); +} diff --git a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp index db6290214..1e70a4cec 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp @@ -1298,4 +1298,70 @@ extern "C" void tessera_apple_gpu_conv2d_f16( std::memset(O, 0, static_cast(N) * outH * outW * Cout * 2); } +// ---- conv3d non-Apple reference (NDHWC source, DHWIO weights) (2026-05-30) -- +extern "C" int32_t tessera_apple_gpu_conv3d_out_dim(int32_t in, int32_t k, + int32_t stride, int32_t pad, + int32_t dilation) { + return conv2d_out_dim_stub(in, k, stride, pad, dilation); +} +static void reference_conv3d_f32_stub( + const float* X, const float* Wt, const float* bias, float* O, int32_t N, + int32_t iD, int32_t iH, int32_t iW, int32_t Cin, int32_t Cout, int32_t kD, + int32_t kH, int32_t kW, int32_t sD, int32_t sH, int32_t sW, int32_t pD, + int32_t pH, int32_t pW, int32_t dD, int32_t dH, int32_t dW, int32_t groups) { + int32_t oD = conv2d_out_dim_stub(iD, kD, sD, pD, dD); + int32_t oH = conv2d_out_dim_stub(iH, kH, sH, pH, dH); + int32_t oW = conv2d_out_dim_stub(iW, kW, sW, pW, dW); + if (oD <= 0 || oH <= 0 || oW <= 0 || groups <= 0 || Cin % groups || + Cout % groups) + return; + int32_t cinG = Cin / groups, coutG = Cout / groups; + for (int32_t n = 0; n < N; ++n) + for (int32_t od = 0; od < oD; ++od) + for (int32_t oh = 0; oh < oH; ++oh) + for (int32_t ow = 0; ow < oW; ++ow) + for (int32_t oc = 0; oc < Cout; ++oc) { + int32_t grp = oc / coutG; + double acc = bias ? static_cast(bias[oc]) : 0.0; + for (int32_t kd = 0; kd < kD; ++kd) { + int32_t id = od * sD + kd * dD - pD; + if (id < 0 || id >= iD) continue; + for (int32_t kh = 0; kh < kH; ++kh) { + int32_t ih = oh * sH + kh * dH - pH; + if (ih < 0 || ih >= iH) continue; + for (int32_t kw = 0; kw < kW; ++kw) { + int32_t iw = ow * sW + kw * dW - pW; + if (iw < 0 || iw >= iW) continue; + for (int32_t ic = 0; ic < cinG; ++ic) { + double xv = X[((((std::size_t)n * iD + id) * iH + ih) * iW + iw) * Cin + grp * cinG + ic]; + double wv = Wt[((((std::size_t)kd * kH + kh) * kW + kw) * cinG + ic) * Cout + oc]; + acc += xv * wv; + } + } + } + } + O[((((std::size_t)n * oD + od) * oH + oh) * oW + ow) * Cout + oc] = + static_cast(acc); + } +} +extern "C" void tessera_apple_gpu_conv3d_f32( + const float* X, const float* Wt, const float* bias, float* O, int32_t N, + int32_t iD, int32_t iH, int32_t iW, int32_t Cin, int32_t Cout, int32_t kD, + int32_t kH, int32_t kW, int32_t sD, int32_t sH, int32_t sW, int32_t pD, + int32_t pH, int32_t pW, int32_t dD, int32_t dH, int32_t dW, int32_t groups) { + reference_conv3d_f32_stub(X, Wt, bias, O, N, iD, iH, iW, Cin, Cout, kD, kH, kW, + sD, sH, sW, pD, pH, pW, dD, dH, dW, groups); +} +extern "C" void tessera_apple_gpu_conv3d_f16( + const uint16_t*, const uint16_t*, const uint16_t*, uint16_t* O, int32_t N, + int32_t iD, int32_t iH, int32_t iW, int32_t, int32_t Cout, int32_t kD, + int32_t kH, int32_t kW, int32_t sD, int32_t sH, int32_t sW, int32_t pD, + int32_t pH, int32_t pW, int32_t dD, int32_t dH, int32_t dW, int32_t) { + int32_t oD = conv2d_out_dim_stub(iD, kD, sD, pD, dD); + int32_t oH = conv2d_out_dim_stub(iH, kH, sH, pH, dH); + int32_t oW = conv2d_out_dim_stub(iW, kW, sW, pW, dW); + if (oD > 0 && oH > 0 && oW > 0) + std::memset(O, 0, static_cast(N) * oD * oH * oW * Cout * 2); +} + #endif // !__APPLE__ diff --git a/tests/unit/test_apple_gpu_conv3d.py b/tests/unit/test_apple_gpu_conv3d.py new file mode 100644 index 000000000..9ca0fec84 --- /dev/null +++ b/tests/unit/test_apple_gpu_conv3d.py @@ -0,0 +1,146 @@ +"""Apple GPU Tier-3 conv3d via im2col + a GPU MPSGraph batched matmul. + +MPSGraph has no 3-D convolution node, so `tessera_apple_gpu_conv3d_{f32,f16}` +lower conv3d to the classic im2col + GEMM decomposition: spatial patches are +gathered into a per-group column matrix and the dominant GEMM runs on-GPU as a +single MPSGraph batched matmul (batch = groups, fp32 accumulation). Bias + +scatter run on the host. Validated against a numpy reference. NDHWC source, +DHWIO weights. See docs/apple_gpu_tier2_tier3_plan.md. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from tessera import runtime as R + + +def _ref_conv3d(X, W, bias, stride, padding, dilation, groups): + """NDHWC X [N,D,H,W,Cin]; DHWIO weight [kD,kH,kW,Cin/groups,Cout].""" + X = X.astype(np.float64) + W = W.astype(np.float64) + sD, sH, sW = (stride,) * 3 if isinstance(stride, int) else stride + pD, pH, pW = (padding,) * 3 if isinstance(padding, int) else padding + dD, dH, dW = (dilation,) * 3 if isinstance(dilation, int) else dilation + N, iD, iH, iW, Cin = X.shape + kD, kH, kW, cinG, Cout = W.shape + oD = (iD + 2 * pD - dD * (kD - 1) - 1) // sD + 1 + oH = (iH + 2 * pH - dH * (kH - 1) - 1) // sH + 1 + oW = (iW + 2 * pW - dW * (kW - 1) - 1) // sW + 1 + coutG = Cout // groups + Xp = np.pad(X, ((0, 0), (pD, pD), (pH, pH), (pW, pW), (0, 0))) + O = np.zeros((N, oD, oH, oW, Cout), dtype=np.float64) + for n in range(N): + for od in range(oD): + for oh in range(oH): + for ow in range(oW): + for oc in range(Cout): + grp = oc // coutG + acc = float(bias[oc]) if bias is not None else 0.0 + for kd in range(kD): + iz = od * sD + kd * dD + for ky in range(kH): + iy = oh * sH + ky * dH + for kx in range(kW): + ix = ow * sW + kx * dW + for ic in range(cinG): + acc += (Xp[n, iz, iy, ix, grp * cinG + ic] + * W[kd, ky, kx, ic, oc]) + O[n, od, oh, ow, oc] = acc + return O + + +# N, D, H, W, Cin, Cout, kD, kH, kW, stride, pad, dil, groups, bias +_CASES = [ + pytest.param((1, 4, 5, 5, 3, 4, 2, 3, 3, 1, 0, 1, 1, True), id="basic"), + pytest.param((2, 5, 5, 4, 4, 6, 3, 3, 3, 1, 1, 1, 1, True), id="pad_bias"), + pytest.param((1, 6, 6, 6, 2, 4, 2, 2, 2, 2, 0, 1, 1, False), id="stride2"), + pytest.param((1, 7, 7, 7, 3, 3, 3, 3, 3, 1, 2, 2, 1, True), id="dilation2"), + pytest.param((1, 4, 4, 4, 4, 4, 2, 3, 3, 1, 1, 1, 2, True), id="groups2"), + pytest.param((1, 4, 4, 4, 4, 4, 1, 1, 1, 1, 0, 1, 4, False), id="depthwise"), +] + + +def _params(case): + (N, D, H, W, Cin, Cout, kD, kH, kW, + stride, pad, dil, groups, has_bias) = case + rng = np.random.RandomState(abs(hash(case)) % (2**31)) + X = rng.randn(N, D, H, W, Cin).astype(np.float32) + Wt = (rng.randn(kD, kH, kW, Cin // groups, Cout) * 0.3).astype(np.float32) + bias = rng.randn(Cout).astype(np.float32) if has_bias else None + kw = dict(stride=stride, padding=pad, dilation=dil, groups=groups) + return X, Wt, bias, kw + + +@pytest.mark.parametrize("case", _CASES) +def test_conv3d_f32_matches_numpy(case): + X, Wt, bias, kw = _params(case) + ops = [X, Wt] + ([bias] if bias is not None else []) + out = R._apple_gpu_dispatch_conv3d(ops, kw, np) + assert out is not None + ref = _ref_conv3d(X, Wt, bias, kw["stride"], kw["padding"], + kw["dilation"], kw["groups"]) + assert out.shape == ref.shape + np.testing.assert_allclose(out, ref, rtol=1e-4, atol=1e-4) + + +def test_conv3d_f16_native(): + N, D, H, W, Cin, Cout = 1, 4, 5, 5, 3, 4 + rng = np.random.RandomState(7) + X = (rng.randn(N, D, H, W, Cin) * 0.5).astype(np.float16) + Wt = (rng.randn(2, 3, 3, Cin, Cout) * 0.3).astype(np.float16) + bias = (rng.randn(Cout) * 0.1).astype(np.float16) + out = R._apple_gpu_dispatch_conv3d([X, Wt, bias], + dict(stride=1, padding=1), np) + assert out is not None and out.dtype == np.float16 + ref = _ref_conv3d(X.astype(np.float32), Wt.astype(np.float32), + bias.astype(np.float32), 1, 1, 1, 1) + np.testing.assert_allclose(out.astype(np.float32), ref, rtol=4e-2, atol=4e-2) + + +def test_conv3d_bf16_round_trip(): + ml_dtypes = pytest.importorskip("ml_dtypes") + bf16 = ml_dtypes.bfloat16 + rng = np.random.RandomState(11) + Xf = (rng.randn(1, 4, 5, 5, 3) * 0.5).astype(np.float32) + Wf = (rng.randn(2, 3, 3, 3, 4) * 0.3).astype(np.float32) + out = R._apple_gpu_dispatch_conv3d([Xf.astype(bf16), Wf.astype(bf16)], + dict(stride=1, padding=1), np) + assert out is not None and out.dtype == bf16 + ref = _ref_conv3d(Xf, Wf, None, 1, 1, 1, 1) + np.testing.assert_allclose(out.astype(np.float32), ref, rtol=8e-2, atol=8e-2) + + +def test_conv3d_symbols_exported(): + rt = R._load_apple_gpu_runtime() + assert hasattr(rt, "tessera_apple_gpu_conv3d_f32") + assert hasattr(rt, "tessera_apple_gpu_conv3d_f16") + assert R._apple_gpu_conv3d_f32() is not None + assert R._apple_gpu_conv3d_f16() is not None + + +def test_conv3d_rejects_non_5d(): + X = np.zeros((4, 5, 5, 3), np.float32) # 4-D (conv2d shape) + Wt = np.zeros((2, 3, 3, 3, 4), np.float32) + assert R._apple_gpu_dispatch_conv3d([X, Wt], {}, np) is None + + +def test_conv3d_in_runtime_envelope(): + assert "tessera.conv3d" in R._APPLE_GPU_CONV_OPS + assert "tessera.conv3d" in R._APPLE_GPU_RUNTIME_OPS + + +def test_conv3d_matches_reference_ops_conv3d(): + """The GPU dispatcher agrees with the eager numpy ops.conv3d reference.""" + import tessera + + rng = np.random.RandomState(3) + X = rng.randn(2, 5, 6, 6, 3).astype(np.float32) + Wt = (rng.randn(3, 3, 3, 3, 5) * 0.3).astype(np.float32) + bias = rng.randn(5).astype(np.float32) + ref = tessera.ops.conv3d(X, Wt, bias=bias, stride=1, padding=1) + out = R._apple_gpu_dispatch_conv3d([X, Wt, bias], + dict(stride=1, padding=1), np) + assert out is not None + np.testing.assert_allclose(out, ref, rtol=1e-4, atol=1e-4)