Skip to content

Integrate oneDNN matmul threadpool backend #970

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
romirdes wants to merge 4 commits into
base: dev
Choose a base branch
from
Open

Integrate oneDNN matmul threadpool backend #970

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b62704b
Add oneDNN matmul-primitive MatMul backend (threadpool runtime)
romirdes Jul 24, 2026
4f9ef09
De-dupcliate oneDNN bazel build files.
romirdes Jul 29, 2026
ee0d76e
Removed shape from cache key since caller already has a more robust d…
romirdes Jul 29, 2026
0acc653
Merge branch 'dev' of https://github.com/google/gemma.cpp into featur…
romirdes Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,17 @@ config_setting(
define_values = {"gemma_onednn_brgemm": "1"},
)

# To enable the OneDNN matmul-primitive backend (threadpool runtime), build with:
# bazel build --define gemma_onednn_matmul=1 ...
# This is mutually exclusive with gemma_onednn_brgemm: passing both --defines
# makes the select()s below match two conditions at once, which Bazel rejects
# with an "ambiguous select" error -- the intended failure mode, mirroring the
# CMake FATAL_ERROR and the #error guard in ops/onednn_matmul.h.
config_setting(
name = "gemma_onednn_matmul",
define_values = {"gemma_onednn_matmul": "1"},
)

cc_library(
name = "basics",
srcs = ["util/basics.cc"],
Expand Down Expand Up @@ -363,12 +374,16 @@ cc_library(
hdrs = [
"ops/brgemm.h",
"ops/matmul.h",
"ops/onednn_matmul.h",
],
defines = select({
":gemma_onednn_brgemm": [
"GEMMA_ONEDNN_BRGEMM=1",
"DNNL_EXPERIMENTAL_UKERNEL",
],
":gemma_onednn_matmul": [
"GEMMA_ONEDNN_MATMUL=1",
],
"//conditions:default": [],
}),
deps = [
Expand All @@ -384,6 +399,7 @@ cc_library(
"@highway//:profiler",
] + select({
":gemma_onednn_brgemm": ["@onednn"],
":gemma_onednn_matmul": ["@onednn_tp//:onednn"],
"//conditions:default": [],
}),
)
Expand All @@ -395,6 +411,7 @@ cc_library(
textual_hdrs = [
"ops/brgemm-inl.h",
"ops/matmul-inl.h",
"ops/onednn_matmul-inl.h",
],
deps = [
":allocator",
Expand All @@ -412,6 +429,7 @@ cc_library(
"@highway//:profiler",
] + select({
":gemma_onednn_brgemm": ["@onednn"],
":gemma_onednn_matmul": ["@onednn_tp//:onednn"],
"//conditions:default": [],
}),
)
Expand All @@ -434,6 +452,7 @@ cc_library(
"ops/brgemm-inl.h",
"ops/matmul_static-inl.h",
"ops/matmul-inl.h",
"ops/onednn_matmul-inl.h",
],
deps = [
":allocator",
Expand All @@ -450,6 +469,7 @@ cc_library(
"@highway//:timer",
] + select({
":gemma_onednn_brgemm": ["@onednn"],
":gemma_onednn_matmul": ["@onednn_tp//:onednn"],
"//conditions:default": [],
}),
)
Expand Down
41 changes: 41 additions & 0 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,19 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
# Enable with: cmake -DGEMMA_ONEDNN_BRGEMM=ON ...
option(GEMMA_ONEDNN_BRGEMM "Enable OneDNN BRGeMM micro-kernel for MatMul (x86-64)" OFF)

# Optional: OneDNN matmul primitive via the threadpool runtime (x86-64 only).
# Enable with: cmake -DGEMMA_ONEDNN_MATMUL=ON ...
option(GEMMA_ONEDNN_MATMUL "Enable OneDNN matmul primitive (threadpool runtime) for MatMul (x86-64)" OFF)

# oneDNN's CPU runtime (SEQ vs THREADPOOL) is a whole-library compile-time
# choice, so the two oneDNN backends cannot share one oneDNN build.
if(GEMMA_ONEDNN_BRGEMM AND GEMMA_ONEDNN_MATMUL)
message(FATAL_ERROR
"GEMMA_ONEDNN_BRGEMM and GEMMA_ONEDNN_MATMUL are mutually exclusive: "
"oneDNN's CPU runtime (SEQ for BRGeMM vs THREADPOOL for the matmul "
"primitive) is a whole-library compile-time choice. Enable only one.")
endif()

if(EMSCRIPTEN)
add_compile_options("-sMEMORY64")
add_compile_options("-msimd128")
Expand Down Expand Up @@ -119,6 +132,24 @@ if(GEMMA_ONEDNN_BRGEMM)
message(STATUS "OneDNN BRGeMM micro-kernel support enabled")
endif()

# OneDNN matmul primitive support via the threadpool runtime (optional, x86-64).
# Same oneDNN source as the BRGeMM backend, but built with the THREADPOOL CPU
# runtime instead of SEQ so oneDNN can call back into gemma.cpp's thread pool.
if(GEMMA_ONEDNN_MATMUL)
set(DNNL_BUILD_TESTS OFF CACHE BOOL "" FORCE)
set(DNNL_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)
set(DNNL_CPU_RUNTIME "THREADPOOL" CACHE STRING "" FORCE)
set(DNNL_GPU_RUNTIME "NONE" CACHE STRING "" FORCE)
set(DNNL_LIBRARY_TYPE "STATIC" CACHE STRING "" FORCE)
FetchContent_Declare(onednn
GIT_REPOSITORY https://github.com/uxlfoundation/oneDNN.git
GIT_TAG v3.11
EXCLUDE_FROM_ALL
)
FetchContent_MakeAvailable(onednn)
message(STATUS "OneDNN matmul primitive (threadpool runtime) support enabled")
endif()

# Base source files
set(SOURCES
compression/compress-inl.h
Expand Down Expand Up @@ -193,6 +224,8 @@ set(SOURCES
ops/matmul.h
ops/brgemm.h
ops/brgemm-inl.h
ops/onednn_matmul.h
ops/onednn_matmul-inl.h
ops/ops-inl.h
ops/ops.h
ops/sum-inl.h
Expand Down Expand Up @@ -255,6 +288,10 @@ if(GEMMA_ONEDNN_BRGEMM)
target_compile_definitions(libgemma PUBLIC GEMMA_ONEDNN_BRGEMM=1 DNNL_EXPERIMENTAL_UKERNEL)
target_link_libraries(libgemma dnnl)
endif()
if(GEMMA_ONEDNN_MATMUL)
target_compile_definitions(libgemma PUBLIC GEMMA_ONEDNN_MATMUL=1)
target_link_libraries(libgemma dnnl)
endif()
install(TARGETS libgemma DESTINATION lib)

# Shared library target for C# interop
Expand Down Expand Up @@ -283,6 +320,10 @@ if(GEMMA_ONEDNN_BRGEMM)
target_compile_definitions(gemma_shared PUBLIC GEMMA_ONEDNN_BRGEMM=1 DNNL_EXPERIMENTAL_UKERNEL)
target_link_libraries(gemma_shared PRIVATE dnnl)
endif()
if(GEMMA_ONEDNN_MATMUL)
target_compile_definitions(gemma_shared PUBLIC GEMMA_ONEDNN_MATMUL=1)
target_link_libraries(gemma_shared PRIVATE dnnl)
endif()
install(TARGETS gemma_shared DESTINATION lib)
install(FILES gemma/c_api.h DESTINATION include/gemma)
install(FILES gemma/GemmaInterop.cs DESTINATION include/gemma)
Expand Down
18 changes: 18 additions & 0 deletions MODULE.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,24 @@ http_archive(
],
)

# Same oneDNN v3.11 source as @onednn, but built for the THREADPOOL CPU runtime
# (the matmul-primitive backend, GEMMA_ONEDNN_MATMUL) rather than SEQ (BRGeMM).
# oneDNN's CPU runtime is a whole-library compile-time choice, so the two cannot
# share one build; this is a separate repo with a different build_file. Both
# build_files are one call to onednn_targets() in bazel/onednn.bzl with a
# different cpu_runtime. The tarball and sha256 are identical, so Bazel's
# download cache is shared and only the arm selected by //:gemma_onednn_matmul
# actually compiles.
http_archive(
name = "onednn_tp",
build_file = "@//bazel:onednn_threadpool.BUILD",
sha256 = "04df98b18300daf6c3aa7cc2d5e7ce8a8f430fed1787151daed0254d8dd4e64e",
strip_prefix = "oneDNN-3.11",
urls = [
"https://github.com/uxlfoundation/oneDNN/archive/refs/tags/v3.11.tar.gz",
],
)

http_archive(
name = "com_google_absl_py",
sha256 = "8a3d0830e4eb4f66c4fa907c06edf6ce1c719ced811a12e26d9d3162f8471758",
Expand Down
Loading