Apple GPU Tier-3: conv3d via im2col + GPU batched matmul

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
 

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` |

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

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)

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<NSNumber *> *aShape = @[ @(G), @(M), @(K) ];
+    NSArray<NSNumber *> *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<float> cols((size_t)groups * rows * K, 0.0f);
+  std::vector<float> 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<float> 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<float> xf(xn), wf(wn), of(on);
+  std::vector<float> 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]);
+}