Add oneDNN matmul-primitive MatMul backend (threadpool runtime)

Second, build-time-selectable 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, self-validating reordered-weights cache, thread_local
  user-managed scratchpad, and DoMatMul_OneDnn with native scale/add/
  demote fusion (writes final values straight into C, clean fallback on
  any 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).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: romirdes

BUILD.bazel

@@ -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"],
@@ -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 = [
@@ -384,6 +399,7 @@ cc_library(
         "@highway//:profiler",
     ] + select({
         ":gemma_onednn_brgemm": ["@onednn"],
+        ":gemma_onednn_matmul": ["@onednn_tp//:onednn"],
         "//conditions:default": [],
     }),
 )
@@ -395,6 +411,7 @@ cc_library(
     textual_hdrs = [
         "ops/brgemm-inl.h",
         "ops/matmul-inl.h",
+        "ops/onednn_matmul-inl.h",
     ],
     deps = [
         ":allocator",
@@ -412,6 +429,7 @@ cc_library(
         "@highway//:profiler",
     ] + select({
         ":gemma_onednn_brgemm": ["@onednn"],
+        ":gemma_onednn_matmul": ["@onednn_tp//:onednn"],
         "//conditions:default": [],
     }),
 )
@@ -433,6 +451,7 @@ cc_library(
         "ops/brgemm-inl.h",
         "ops/matmul_static-inl.h",
         "ops/matmul-inl.h",
+        "ops/onednn_matmul-inl.h",
     ],
     deps = [
         ":allocator",
@@ -449,6 +468,7 @@ cc_library(
         "@highway//:timer",
     ] + select({
         ":gemma_onednn_brgemm": ["@onednn"],
+        ":gemma_onednn_matmul": ["@onednn_tp//:onednn"],
         "//conditions:default": [],
     }),
 )

CMakeLists.txt

@@ -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")
@@ -113,6 +126,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
@@ -185,6 +216,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
@@ -247,6 +280,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
@@ -275,6 +312,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)

MODULE.bazel

@@ -36,6 +36,22 @@ 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. 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",