Skip to content

Integrate oneDNN matmul threadpool backend - #970

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

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

Conversation

@romirdes

@romirdesromirdes 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

FileDescription
ops/onednn_matmul.hNew 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.hNew DoMatMul_OneDnn() implementation: builds/caches the matmul primitive and reordered weights, applies scale and add, returns false to decline unsupported cases.
ops/matmul-inl.hDispatch into the oneDNN path for BF16×BF16 with M > 1, guarded by GEMMA_ONEDNN_MATMUL, with fall-through to the existing MatMul.
ops/matmul.hExport ops/onednn_matmul.h; add MMPerKey::onednn_built to record that the oneDNN path handled a shape.
ops/bench_matmul.ccExclude 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.hAdd Callers::kOneDnnMatMul profiler zone.
util/zones.ccMap the new zone to the "OneDnnMatMul" name.
BUILD.bazelAdd the gemma_onednn_matmul config setting, GEMMA_ONEDNN_MATMUL=1 define, new headers, and the @onednn_tp dependency.
MODULE.bazelAdd the @onednn_tp archive — same oneDNN v3.11 tarball/sha256 as @onednn (shared download cache), different build_file for the THREADPOOL runtime.
bazel/onednn.bzlNew shared build logic: onednn_targets(cpu_runtime) derives the runtime-dependent config lines and conditional copt from its argument.
bazel/onednn.BUILDDe-duplicated down to a single onednn_targets(cpu_runtime = "SEQ") call.
bazel/onednn_threadpool.BUILDNew build file: a single onednn_targets(cpu_runtime = "THREADPOOL") call.
CMakeLists.txtAdd 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.

Loading
Sign up for freeto 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.

2 participants

@romirdes@jan-wassenberg