Merge #903: OneDNN BRGeMM Micro-Kernel Integration for BF16 MatMul

BUILD.bazel

@@ -27,13 +27,6 @@ exports_files([
     ".github/workflows/build.yml",
 ])
 
-# To enable OneDNN BRGeMM support, build with:
-#   bazel build --define gemma_onednn_brgemm=1 ...
-config_setting(
-    name = "gemma_onednn_brgemm",
-    define_values = {"gemma_onednn_brgemm": "1"},
-)
-
 cc_library(
     name = "basics",
     srcs = ["util/basics.cc"],
@@ -321,17 +314,7 @@ test_suite(
 cc_library(
     name = "matmul_env",
     srcs = ["ops/matmul.cc"],
-    hdrs = [
-        "ops/brgemm.h",
-        "ops/matmul.h",
-    ],
-    defines = select({
-        ":gemma_onednn_brgemm": [
-            "GEMMA_ONEDNN_BRGEMM=1",
-            "DNNL_EXPERIMENTAL_UKERNEL",
-        ],
-        "//conditions:default": [],
-    }),
+    hdrs = ["ops/matmul.h"],
     deps = [
         ":allocator",
         ":basics",
@@ -342,20 +325,14 @@ cc_library(
         "@highway//:hwy",
         "@highway//:nanobenchmark",
         "@highway//:profiler",
-    ] + select({
-        ":gemma_onednn_brgemm": ["@onednn"],
-        "//conditions:default": [],
-    }),
+    ],
 )
 
 cc_library(
     name = "matmul",
     # allow depending only on this target, without also matmul_env.
     hdrs = ["ops/matmul.h"],
-    textual_hdrs = [
-        "ops/brgemm-inl.h",
-        "ops/matmul-inl.h",
-    ],
+    textual_hdrs = ["ops/matmul-inl.h"],
     deps = [
         ":allocator",
         ":basics",
@@ -369,10 +346,7 @@ cc_library(
         "@highway//:hwy",
         "@highway//:nanobenchmark",
         "@highway//:profiler",
-    ] + select({
-        ":gemma_onednn_brgemm": ["@onednn"],
-        "//conditions:default": [],
-    }),
+    ],
 )
 
 cc_library(
@@ -389,7 +363,6 @@ cc_library(
         "ops/matmul_static.h",
     ],
     textual_hdrs = [
-        "ops/brgemm-inl.h",
         "ops/matmul_static-inl.h",
         "ops/matmul-inl.h",
     ],
@@ -406,10 +379,7 @@ cc_library(
         "@highway//:hwy",
         "@highway//:profiler",
         "@highway//:timer",
-    ] + select({
-        ":gemma_onednn_brgemm": ["@onednn"],
-        "//conditions:default": [],
-    }),
+    ],
 )
 
 cc_library(

CMakeLists.txt

@@ -22,10 +22,6 @@ set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
-# Optional: OneDNN BRGeMM micro-kernel support (x86-64 only).
-# Enable with: cmake -DGEMMA_ONEDNN_BRGEMM=ON ...
-option(GEMMA_ONEDNN_BRGEMM "Enable OneDNN BRGeMM micro-kernel for MatMul (x86-64)" OFF)
-
 if(EMSCRIPTEN)
   add_compile_options("-sMEMORY64")
   add_compile_options("-msimd128")
@@ -89,23 +85,6 @@ if(EMSCRIPTEN)
   target_compile_options(benchmark PRIVATE -Wno-c2y-extensions)
 endif()
 
-# OneDNN BRGeMM micro-kernel support (optional, x86-64 only).
-if(GEMMA_ONEDNN_BRGEMM)
-  set(DNNL_BUILD_TESTS OFF CACHE BOOL "" FORCE)
-  set(DNNL_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)
-  set(DNNL_CPU_RUNTIME "SEQ" CACHE STRING "" FORCE)
-  set(DNNL_GPU_RUNTIME "NONE" CACHE STRING "" FORCE)
-  set(DNNL_LIBRARY_TYPE "STATIC" CACHE STRING "" FORCE)
-  set(DNNL_EXPERIMENTAL_UKERNEL ON CACHE BOOL "" 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 BRGeMM micro-kernel support enabled")
-endif()
-
 # Base source files
 set(SOURCES
   compression/compress-inl.h
@@ -164,8 +143,6 @@ set(SOURCES
   ops/matmul-inl.h
   ops/matmul.cc
   ops/matmul.h
-  ops/brgemm.h
-  ops/brgemm-inl.h
   ops/ops-inl.h
   ops/ops.h
   ops/sum-inl.h
@@ -218,10 +195,6 @@ target_link_libraries(libgemma hwy hwy_contrib sentencepiece-static)
 target_include_directories(libgemma PUBLIC ${sentencepiece_SOURCE_DIR})
 target_compile_definitions(libgemma PRIVATE $<$<PLATFORM_ID:Windows>:_CRT_SECURE_NO_WARNINGS NOMINMAX>)
 target_compile_options(libgemma PRIVATE $<$<PLATFORM_ID:Windows>:-Wno-deprecated-declarations>)
-if(GEMMA_ONEDNN_BRGEMM)
-  target_compile_definitions(libgemma PUBLIC GEMMA_ONEDNN_BRGEMM=1 DNNL_EXPERIMENTAL_UKERNEL)
-  target_link_libraries(libgemma dnnl)
-endif()
 install(TARGETS libgemma DESTINATION lib)
 
 # Shared library target for C# interop
@@ -246,10 +219,6 @@ target_compile_definitions(gemma_shared
     $<$<PLATFORM_ID:Windows>:_CRT_SECURE_NO_WARNINGS NOMINMAX>
 )
 target_compile_options(gemma_shared PRIVATE $<$<PLATFORM_ID:Windows>:-Wno-deprecated-declarations>)
-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()
 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

@@ -25,17 +25,6 @@ git_override(
 
 http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
-# OneDNN v3.11 for BRGeMM micro-kernel support (optional, x86-64 only).
-http_archive(
-    name = "onednn",
-    build_file = "@//bazel:onednn.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",