Fixed the compile time flag to designate BRGEMM path

Author: bbhattar

BUILD.bazel

@@ -27,6 +27,13 @@ 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"],
@@ -318,7 +325,7 @@ cc_library(
         "ops/matmul.h",
     ],
     defines = select({
-        "@platforms//cpu:x86_64": ["GEMMA_ONEDNN=1", "DNNL_EXPERIMENTAL_UKERNEL"],
+        ":gemma_onednn_brgemm": ["GEMMA_ONEDNN_BRGEMM=1", "DNNL_EXPERIMENTAL_UKERNEL"],
         "//conditions:default": [],
     }),
     deps = [
@@ -332,7 +339,7 @@ cc_library(
         "@highway//:nanobenchmark",
         "@highway//:profiler",
     ] + select({
-        "@platforms//cpu:x86_64": ["@onednn//:onednn"],
+        ":gemma_onednn_brgemm": ["@onednn//:onednn"],
         "//conditions:default": [],
     }),
 )
@@ -359,7 +366,7 @@ cc_library(
         "@highway//:nanobenchmark",
         "@highway//:profiler",
     ] + select({
-        "@platforms//cpu:x86_64": ["@onednn//:onednn"],
+        ":gemma_onednn_brgemm": ["@onednn//:onednn"],
         "//conditions:default": [],
     }),
 )
@@ -396,7 +403,7 @@ cc_library(
         "@highway//:profiler",
         "@highway//:timer",
     ] + select({
-        "@platforms//cpu:x86_64": ["@onednn//:onednn"],
+        ":gemma_onednn_brgemm": ["@onednn//:onednn"],
         "//conditions:default": [],
     }),
 )

ops/bench_matmul.cc

@@ -131,7 +131,7 @@ void BenchMatMul(size_t M, size_t K, size_t N, bool add, MatMulEnv& env) {
 
     // Only record times after autotuning finished.
     bool done = per_key->autotune.Best();
-#if GEMMA_ONEDNN
+#if GEMMA_ONEDNN_BRGEMM_BRGEMM
     done = done || per_key->brgemm_autotune.Best();
 #endif
     if (done) times.push_back(elapsed);

ops/brgemm-inl.h

@@ -15,7 +15,7 @@
 
 // BRGeMM dispatch. Included from matmul-inl.h inside gcpp::HWY_NAMESPACE.
 
-#if GEMMA_ONEDNN
+#if GEMMA_ONEDNN_BRGEMM_BRGEMM
 
 static bool MakeBrgemm(dnnl::ukernel::brgemm& brg, int64_t m, int64_t n,
                         int64_t k, int64_t batch, int64_t lda, int64_t ldb,
@@ -489,4 +489,4 @@ static HWY_NOINLINE void DoMatMul_BRGeMM(
   main_bufs.hw_ctx_kernel = nullptr;
 }
 
-#endif  // GEMMA_ONEDNN
+#endif  // GEMMA_ONEDNN_BRGEMM_BRGEMM

ops/brgemm.h

@@ -14,7 +14,7 @@
 // limitations under the License.
 
 // OneDNN BRGeMM micro-kernel integration for MatMul on Intel AMX/AVX-512.
-// Enabled at runtime via GEMMA_USE_ONEDNN_BRGEMM=1.
+// Enabled at compile time via GEMMA_ONEDNN_BRGEMM_BRGEMM=1 (Bazel: --define gemma_onednn_brgemm=1).
 
 #ifndef THIRD_PARTY_GEMMA_CPP_OPS_BRGEMM_H_
 #define THIRD_PARTY_GEMMA_CPP_OPS_BRGEMM_H_
@@ -23,30 +23,21 @@
 #include <stdint.h>
 
 #include <algorithm>
-#include <cstdlib>
 #include <unordered_map>
 #include <utility>
 #include <vector>
 
 #include "hwy/base.h"
 
-#if GEMMA_ONEDNN
+#if GEMMA_ONEDNN_BRGEMM_BRGEMM
 #include <sys/mman.h>
 
 #include "oneapi/dnnl/dnnl.hpp"
 #include "oneapi/dnnl/dnnl_ukernel.hpp"
-#endif  // GEMMA_ONEDNN
+#endif  // GEMMA_ONEDNN_BRGEMM_BRGEMM
 
 namespace gcpp {
 
-inline bool UseOneDnnBrgemm() {
-  static const bool enabled = [] {
-    const char* env = std::getenv("GEMMA_USE_ONEDNN_BRGEMM");
-    return env != nullptr && env[0] == '1' && env[1] == '\0';
-  }();
-  return enabled;
-}
-
 struct BRGeMMConfig {
   int64_t M_blk;
   int64_t N_blk;
@@ -55,7 +46,7 @@ struct BRGeMMConfig {
   int64_t par_m;
 };
 
-#if GEMMA_ONEDNN
+#if GEMMA_ONEDNN_BRGEMM_BRGEMM
 
 // Generates autotuning candidates. Fixed: N_blk=32, K_blk=32 (AMX BF16).
 // Tunable: M_blk in {32,64}, batch_size in {16,32,64,128,256}.
@@ -290,7 +281,7 @@ inline auto& GetBRGeMMPackedBCache() {
   return cache;
 }
 
-#endif  // GEMMA_ONEDNN
+#endif  // GEMMA_ONEDNN_BRGEMM_BRGEMM
 
 }  // namespace gcpp
 

ops/matmul-inl.h

@@ -47,9 +47,9 @@ namespace gcpp {
 namespace HWY_NAMESPACE {
 namespace hn = hwy::HWY_NAMESPACE;
 
-#if GEMMA_ONEDNN
+#if GEMMA_ONEDNN_BRGEMM_BRGEMM
 #include "ops/brgemm-inl.h"  // DoMatMul_BRGeMM
-#endif  // GEMMA_ONEDNN
+#endif  // GEMMA_ONEDNN_BRGEMM_BRGEMM
 
 // Like hn::PromoteOddTo, but uses assembly to avoid an extra vector register.
 template <class DF, class DBF = hn::Repartition<BF16, DF>>
@@ -1081,11 +1081,11 @@ HWY_NOINLINE MMPerKey* MatMul(const MatPtrT<TA>& A, const MatPtrT<TB>& B,
   MMPerKey& per_key = MMImpl::FindOrAddPerKey(
       M, K, N, num_B, cache.VectorBytes(), env.per_cluster[cluster_idx]);
 
-#if GEMMA_ONEDNN
+#if GEMMA_ONEDNN_BRGEMM_BRGEMM
   // BRGeMM path for BF16×BF16 on Intel AMX/AVX-512.
   // Requires M,N,K >= 32 and K % 32 == 0 (AMX tile constraint).
   if constexpr (IsBF16<TA>() && IsBF16<TB>()) {
-  if (UseOneDnnBrgemm() && M >= 32 && N >= 32 && K >= 32 && (K % 32) == 0) {
+  if (M >= 32 && N >= 32 && K >= 32 && (K % 32) == 0) {
     const float scale = A.Scale() * B.Scale();
     MMAutoTune<BRGeMMConfig>& brg_tuner = per_key.brgemm_autotune;
 
@@ -1119,7 +1119,7 @@ HWY_NOINLINE MMPerKey* MatMul(const MatPtrT<TA>& A, const MatPtrT<TB>& B,
     return &per_key;
   }
   }  // if constexpr BF16/float
-#endif  // GEMMA_ONEDNN
+#endif  // GEMMA_ONEDNN_BRGEMM_BRGEMM
 
   // (Also auto-tunes, hence outside the timed section to prevent interference.)
   const StridedViewBF A_view =

ops/matmul.h

@@ -32,7 +32,7 @@
 #include "hwy/base.h"
 #include "hwy/bit_set.h"
 #include "hwy/profiler.h"
-#include "ops/brgemm.h"  // BRGeMMConfig, GEMMA_ONEDNN
+#include "ops/brgemm.h"  // BRGeMMConfig, GEMMA_ONEDNN_BRGEMM_BRGEMM
 // IWYU pragma: end_exports
 
 namespace gcpp {
@@ -640,9 +640,9 @@ class MMKeys {
 struct MMPerKey {
   MMAutoTune<MMConfig> autotune;
   MMAutoTune<MMParA> autotune_par_a;
-#if GEMMA_ONEDNN
+#if GEMMA_ONEDNN_BRGEMM_BRGEMM
   MMAutoTune<BRGeMMConfig> brgemm_autotune;
-#endif  // GEMMA_ONEDNN
+#endif  // GEMMA_ONEDNN_BRGEMM_BRGEMM
 };
 
 // Stores state shared across MatMul calls. Non-copyable. `ctx` must outlive