From 6e3438946122ca0771aafbe5d6643cd3a8a29261 Mon Sep 17 00:00:00 2001
From: serge-sans-paille <sguelton@mozilla.com>
Date: Tue, 6 Feb 2024 13:59:58 +0000
Subject: [PATCH] Bug 1878695 - Add avx512vnni support to firefox translation
 engine r=marco,gcp

Differential Revision: https://phabricator.services.mozilla.com/D200682
---
 build/moz.configure/toolchain.configure       | 11 +++++++++++
 js/src/intgemm/IntegerGemmIntrinsic.cpp       |  7 ++++++-
 js/src/intgemm/moz.build                      |  6 ++++++
 .../kernels/GemmologyEngineAVX512VNNI.cpp     | 19 +++++++++++++++++++
 4 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 third_party/gemmology/kernels/GemmologyEngineAVX512VNNI.cpp

diff --git a/build/moz.configure/toolchain.configure b/build/moz.configure/toolchain.configure
index d08b748db250..a8ab8bd698e4 100644
--- a/build/moz.configure/toolchain.configure
+++ b/build/moz.configure/toolchain.configure
@@ -3277,3 +3277,14 @@ set_config(
         flags=["-mavx512bw", "-mavx512f", "-mavx512dq", "-mavx512cd"],
     ),
 )
+
+# AVX512VNNI can be based on either avx512bw or avx512vbmi. We choose the
+# former.
+set_config(
+    "AVX512VNNI_FLAGS",
+    ["-mavx512vnni", "-mavx512bw", "-mavx512f", "-mavx512dq", "-mavx512cd"],
+    try_compile(
+        check_msg="for -mavx512vnni support",
+        flags=["-mavx512vnni", "-mavx512bw", "-mavx512f", "-mavx512dq", "-mavx512cd"],
+    ),
+)
diff --git a/js/src/intgemm/IntegerGemmIntrinsic.cpp b/js/src/intgemm/IntegerGemmIntrinsic.cpp
index d6547fdc32b7..aebb0f5f02f0 100644
--- a/js/src/intgemm/IntegerGemmIntrinsic.cpp
+++ b/js/src/intgemm/IntegerGemmIntrinsic.cpp
@@ -21,7 +21,12 @@
 #include "wasm/WasmLog.h"
 
 #if defined(USE_AVX512BW)
-#  if defined(USE_AVXVNNI)
+#  if defined(USE_AVX512VNNI)
+#    define SUPPORTED_ARCHS                                          \
+      xsimd::arch_list<xsimd::avx512vnni<xsimd::avx512bw>, \
+                       xsimd::avx512bw, xsimd::avx2, \
+                       xsimd::ssse3, xsimd::sse2>
+#  elif defined(USE_AVXVNNI)
 #    define SUPPORTED_ARCHS                                          \
       xsimd::arch_list<xsimd::avx512bw, xsimd::avxvnni, xsimd::avx2, \
                        xsimd::ssse3, xsimd::sse2>
diff --git a/js/src/intgemm/moz.build b/js/src/intgemm/moz.build
index daa51e1e2a73..9249938cbb26 100644
--- a/js/src/intgemm/moz.build
+++ b/js/src/intgemm/moz.build
@@ -49,6 +49,12 @@ if CONFIG["INTEL_ARCHITECTURE"]:
         SOURCES[
             "/third_party/gemmology/kernels/GemmologyEngineAVX512BW.cpp"
         ].flags += CONFIG["AVX512BW_FLAGS"]
+    if CONFIG["AVX512VNNI_FLAGS"]:
+        DEFINES["USE_AVX512VNNI"] = True
+        SOURCES += ["/third_party/gemmology/kernels/GemmologyEngineAVX512VNNI.cpp"]
+        SOURCES[
+            "/third_party/gemmology/kernels/GemmologyEngineAVX512VNNI.cpp"
+        ].flags += CONFIG["AVX512VNNI_FLAGS"]
 
 
 if CONFIG["TARGET_CPU"] == "aarch64":
diff --git a/third_party/gemmology/kernels/GemmologyEngineAVX512VNNI.cpp b/third_party/gemmology/kernels/GemmologyEngineAVX512VNNI.cpp
new file mode 100644
index 000000000000..80425fafd4f4
--- /dev/null
+++ b/third_party/gemmology/kernels/GemmologyEngineAVX512VNNI.cpp
@@ -0,0 +1,19 @@
+/* -*- mode: c++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* this source code form is subject to the terms of the mozilla public
+ * license, v. 2.0. if a copy of the mpl was not distributed with this file,
+ * You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include <gemmology.h>
+
+namespace gemmology {
+template struct Engine<xsimd::avx512vnni<xsimd::avx512bw>>;
+template void Engine<xsimd::avx512vnni<xsimd::avx512bw>>::SelectColumnsB(int8_t const*, int8_t*,
+                                                  size_t, uint32_t const*,
+                                                  uint32_t const*);
+template void Engine<xsimd::avx512vnni<xsimd::avx512bw>>::Shift::Multiply(
+    uint8_t const*, int8_t const*, size_t, size_t, size_t,
+    gemmology::callbacks::UnquantizeAndAddBiasAndWrite);
+template void Engine<xsimd::avx512vnni<xsimd::avx512bw>>::Shift::PrepareBias(
+    int8_t const*, size_t, size_t,
+    gemmology::callbacks::UnquantizeAndAddBiasAndWrite);
+}  // namespace gemmology