Skip to content

Integrate oneDNN matmul threadpool backend - #970

Open
romirdes wants to merge 4 commits into
google:devfrom
Intel-tensorflow:feature/onednn-matmul-threadpool
Open

Integrate oneDNN matmul threadpool backend#970
romirdes wants to merge 4 commits into
google:devfrom
Intel-tensorflow:feature/onednn-matmul-threadpool

Conversation

@romirdes

@romirdes romirdes commented Jul 29, 2026

Copy link
Copy Markdown

Summary

Adds an optional oneDNN matmul-primitive backend for MatMul, built against oneDNN's THREADPOOL CPU runtime so oneDNN parallelizes by calling back into gemma.cpp's existing thread pool instead of spawning its own threads.

The backend is opt-in and off by default:

  • Bazel: bazel build --define gemma_onednn_matmul=1 ...
  • CMake: cmake -DGEMMA_ONEDNN_MATMUL=ON ...

It is mutually exclusive with the existing BRGeMM backend (GEMMA_ONEDNN_BRGEMM), because oneDNN's CPU runtime (SEQ vs THREADPOOL) is a whole-library compile-time choice. Enabling both fails fast: a CMake FATAL_ERROR, an ambiguous select() in Bazel, and an #error guard in ops/onednn_matmul.h.

The path engages only for BF16×BF16 with M > 1 (M == 1 measured slower than stock gemma for some shapes) and falls through to the standard MatMul whenever oneDNN declines or fails, so behavior is unchanged when the flag is off.

Changes

File Description
ops/onednn_matmul.h New public header: build-flag guard, mutual-exclusion #error, and the threadpool adapter that lets oneDNN dispatch work onto gemma.cpp's thread pool.
ops/onednn_matmul-inl.h New DoMatMul_OneDnn() implementation: builds/caches the matmul primitive and reordered weights, applies scale and add, returns false to decline unsupported cases.
ops/matmul-inl.h Dispatch into the oneDNN path for BF16×BF16 with M > 1, guarded by GEMMA_ONEDNN_MATMUL, with fall-through to the existing MatMul.
ops/matmul.h Export ops/onednn_matmul.h; add MMPerKey::onednn_built to record that the oneDNN path handled a shape.
ops/bench_matmul.cc Exclude the first oneDNN call (JIT + weight reorder) from timings, since its cost is amortized by the primitive/weights caches; oneDNN has no autotune sweep.
util/zones.h Add Callers::kOneDnnMatMul profiler zone.
util/zones.cc Map the new zone to the "OneDnnMatMul" name.
BUILD.bazel Add the gemma_onednn_matmul config setting, GEMMA_ONEDNN_MATMUL=1 define, new headers, and the @onednn_tp dependency.
MODULE.bazel Add the @onednn_tp archive — same oneDNN v3.11 tarball/sha256 as @onednn (shared download cache), different build_file for the THREADPOOL runtime.
bazel/onednn.bzl New shared build logic: onednn_targets(cpu_runtime) derives the runtime-dependent config lines and conditional copt from its argument.
bazel/onednn.BUILD De-duplicated down to a single onednn_targets(cpu_runtime = "SEQ") call.
bazel/onednn_threadpool.BUILD New build file: a single onednn_targets(cpu_runtime = "THREADPOOL") call.
CMakeLists.txt Add the GEMMA_ONEDNN_MATMUL option, the mutual-exclusion FATAL_ERROR, and a FetchContent oneDNN v3.11 build configured with DNNL_CPU_RUNTIME=THREADPOOL.

Testing

Built and benchmarked with ops/bench_matmul under both Bazel and CMake, against the stock MatMul and the BRGeMM backend. Default (flag-off) builds are unaffected.

romirdes added 4 commits July 24, 2026 15:19
oneDNN backend that uses the high-level dnnl::matmul primitive
with the THREADPOOL CPU runtime, coexisting with but mutually exclusive
from the merged BRGeMM (SEQ ukernel) backend.
oneDNN picks the kernel and calls back into gemma.cpp's hwy::ThreadPool
through a threadpool_iface adapter for all parallelism.

- ops/onednn_matmul.h / -inl.h: engine singleton, HwyThreadPoolAdapter
  bridge, reordered-weights cache, thread_local
  user-managed scratchpad, and DoMatMul_OneDnn that falls back to HWY on failure.
- ops/matmul-inl.h: dispatch hook (BF16xBF16, M>1) with fallback to the
  stock path; MMPerKey bookkeeping field; zones caller label; bench
  fair-timing for the first JIT+reorder call.
- Mutual exclusion enforced at three layers: CMake FATAL_ERROR, Bazel
  ambiguous-select(), and a C++ #error guard.
- CMake + Bazel wiring to fetch/build oneDNN v3.11 (THREADPOOL runtime).
Move the shared body into bazel/onednn.bzl as onednn_targets(cpu_runtime),
which derives the runtime-dependent config lines and the one conditional
copt from its argument. Both BUILD files reduce to a single call:

    onednn_targets(cpu_runtime = "SEQ")          # bazel/onednn.BUILD
    onednn_targets(cpu_runtime = "THREADPOOL")   # bazel/onednn_threadpool.BUILD
…escriptor check.

Other minor, non-performance edits.
@romirdes
romirdes marked this pull request as ready for review July 29, 2026 21:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant