Apple GPU Tier-3: conv3d via im2col + GPU batched matmul - #22
Conversation
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 <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7e9b8d56a2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if (oD > 0 && oH > 0 && oW > 0) | ||
| std::memset(O, 0, static_cast<std::size_t>(N) * oD * oH * oW * Cout * 2); |
There was a problem hiding this comment.
On non-Apple hosts the Python dispatcher loads this stub and treats tessera_apple_gpu_conv3d_f16 as an available implementation, so any f16 conv3d returns an all-zero tensor instead of falling back to the reference math. This is not just theoretical: I checked tests/unit/test_apple_gpu_conv3d.py on Linux and test_conv3d_f16_native fails with 96% of elements mismatched because the stub zero-fills the output.
Useful? React with 👍 / 👎.
1 participant
MPSGraph has no 3-D convolution node, so conv3d lowers to the classic
im2col + GEMM decomposition with the dominant GEMM on-GPU.
What
tessera_apple_gpu_conv3d_{f32,f16}(NDHWC source[N,D,H,W,Cin], DHWIOweights
[kD,kH,kW,Cin/groups,Cout], optional bias[Cout]).conv3d_core_f32gathers patches into a per-group[groups, rows, K]columnmatrix (rows = N·oD·oH·oW, K = kD·kH·kW·Cin/groups), regroups weights to
[groups, K, Cout/groups], and runs the GEMM as a single cached MPSGraphbatched matmul (batch = groups, fp32 accumulation, RAII buffer-pool
acquires). Bias + scatter back to NDHWC on the host.
a host fp32 round-trip; f16 converts at the boundary.
runtime.pydispatcher + ctypes wrappers;tessera.conv3dwired into themetadata op-loop (conv2d/conv3d split by rank), the
_APPLE_GPU_CONV_OPSenvelope, and
driver.pygating (execution_mode="metal_runtime").Layout
Matches the eager
tessera.ops.conv3d(NDHWC / DHWIO) exactly — asserted bytest_conv3d_matches_reference_ops_conv3d.Tests
tests/unit/test_apple_gpu_conv3d.py— 12 tests, validated on AppleSilicon: f32 across 6 shape/stride/pad/dilation/groups/depthwise cases vs numpy
(rtol 1e-4), native f16, bf16 round-trip, symbol export, envelope membership,
and agreement with the eager reference.
Verification (local, Apple Silicon)
torch-import error remains)🤖 Generated with Claude Code