CI run: Highway CPU kernels on Linux arm64 (not for merge)

.github/workflows/pull_request.yml

@@ -204,23 +204,33 @@ jobs:
           bash .github/scripts/build-linux-cuda-cmake.sh
 
   linux_swiftpm:
-    name: Linux SwiftPM (${{ matrix.arch }})
+    name: Linux SwiftPM (${{ matrix.arch }}, Highway ${{ matrix.highway }})
     # Advisory while it is new: builds the Linux profile of Package.swift (the LM libraries and
-    # their dependencies) and runs VMLXLinuxTests. GitHub-hosted runners only: the self-hosted pools
-    # queue for a day. python3 is for scripts/check-metal-only.sh.
+    # their dependencies) with Highway on x86-64 and arm64, and in a third leg without it, then
+    # runs VMLXLinuxTests and the CPU kernel tests. GitHub-hosted runners only: the self-hosted
+    # pools queue for a day. python3 is for scripts/check-metal-only.sh.
     continue-on-error: true
-    timeout-minutes: 60
+    timeout-minutes: 120
     strategy:
       fail-fast: false
       matrix:
         include:
           - host: ubuntu-24.04
             arch: x86_64
+            highway: "on"
+          - host: ubuntu-24.04-arm
+            arch: aarch64
+            highway: "on"
+          # Without Highway (VMLX_NO_HIGHWAY=1): MLX's scalar CPU code, as the aarch64 leg ran
+          # before it had Highway.
           - host: ubuntu-24.04-arm
             arch: aarch64
+            highway: "off"
     runs-on: ${{ matrix.host }}
+    env:
+      VMLX_NO_HIGHWAY: ${{ matrix.highway == 'off' && '1' || '' }}
     container:
-      image: swift:6.3-bookworm
+      image: swift:6.4-bookworm
     defaults:
       run:
         shell: bash
@@ -250,7 +260,88 @@ jobs:
           swift test --filter VMLXLinuxTests 2>&1 | tee vmlx-linux-tests.log
           grep -Eq 'Test run with [1-9][0-9]* tests? .*passed' vmlx-linux-tests.log
 
+      - name: CPU kernel tests
+        run: |
+          set -o pipefail
+          bash scripts/run-cpu-kernel-tests.sh 2>&1 | tee cpu-kernel-tests.log
+
       # Last, and with !cancelled(), so a broken checker never hides the build and test results.
       - name: METAL-ONLY checker test
         if: ${{ !cancelled() && hashFiles('scripts/check-metal-only-test.sh') != '' }}
         run: bash scripts/check-metal-only-test.sh
+
+  linux_cpu_kernels_sde:
+    name: Linux CPU kernels under SDE (Sapphire Rapids)
+    # Advisory: emulated AVX-512 is slow, so it runs one pool configuration and the tests that
+    # dispatch per target. The native x86 leg of linux_swiftpm runs everything.
+    continue-on-error: true
+    timeout-minutes: 180
+    runs-on: ubuntu-24.04
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v6
+        with:
+          submodules: recursive
+
+      - name: Install Intel SDE
+        uses: petarpetrovt/setup-sde@6d2773bc0ef3005fd52e0ae0bb91782321793aec # v6.0
+        with:
+          sdeVersion: 10.13.1
+
+      - name: Allow ptrace for SDE
+        run: sudo sysctl -w kernel.yama.ptrace_scope=0
+
+      - name: Kernel tests under SDE
+        run: |
+          docker run --rm --cap-add=SYS_PTRACE -v "$PWD:/work" -v "$SDE_PATH:/sde:ro" -w /work \
+            swift:6.4-bookworm bash -c '
+              set -euo pipefail
+              apt-get update
+              DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+                libopenblas-dev liblapack-dev liblapacke-dev gfortran g++ libsqlite3-dev python3
+              export PATH=/sde:$PATH
+              bash scripts/run-cpu-kernel-tests.sh --sde spr --once --stall 5400 \
+                --filter "VMLXCPUKernelTests\\.(TargetsTests|QuantizedMatmulTests|NormTests|RopeTests)|AttentionTests/eachPath"'
+
+  linux_cpu_kernels_tsan:
+    name: Linux CPU kernels under ThreadSanitizer (aarch64)
+    # Advisory: the CPU pool and the threaded kernels under ThreadSanitizer, on arm64. TSan checks
+    # the C++ memory model, not the CPU's, so one architecture suffices. Two races in upstream MLX's
+    # eval_impl are suppressed (scripts/tsan-suppressions.txt). One pool configuration, the threaded
+    # suites only.
+    continue-on-error: true
+    # 60 min: ~18 expected (4.4 min on a 12-CPU arm64 Linux VM x ~3.6 on this runner, plus setup),
+    # doubled, plus the runner's --stall window (1200 s, passed below), so that a hang ends in
+    # STALLED, not a timeout.
+    timeout-minutes: 60
+    runs-on: ubuntu-24.04-arm
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v6
+        with:
+          submodules: recursive
+
+      - name: Fewer random mmap bits, for ThreadSanitizer
+        run: |
+          sudo sysctl vm.mmap_rnd_bits
+          sudo sysctl -w vm.mmap_rnd_bits=28
+
+      - name: Kernel tests under ThreadSanitizer
+        run: |
+          docker run --rm -v "$PWD:/work" -w /work swift:6.4-bookworm bash -c '
+            set -euo pipefail
+            apt-get update
+            DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+              libopenblas-dev liblapack-dev liblapacke-dev gfortran g++ libsqlite3-dev python3
+            bash scripts/run-cpu-kernel-tests.sh --sanitize thread --once --stall 1200 \
+              --filter "VMLXCPUKernelTests\\.(QuantizedMatmulTests|GatherQuantizedMatmulTests|GemmTests|ReductionTests|IndexingTests)|threadedElementwise|threadedRows|threadedSoftmax"'
+
+      - name: Upload the ThreadSanitizer logs
+        if: failure()
+        uses: actions/upload-artifact@v4
+        with:
+          name: tsan-kernel-tests
+          path: .build/hwy-all-thread/kernel-tests-*.log
+          include-hidden-files: true
+          overwrite: true
+          retention-days: 7

.gitmodules

@@ -4,3 +4,6 @@
 [submodule "submodules/mlx-c"]
 	path = Source/Cmlx/mlx-c
 	url = https://github.com/osaurus-ai/mlx-c
+[submodule "submodules/highway"]
+	path = Source/Cmlx/highway
+	url = https://github.com/google/highway

ACKNOWLEDGMENTS.md

@@ -296,3 +296,34 @@ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 
+## Google Highway
+
+Used under the BSD 3-Clause license, which Highway offers as an alternative to Apache 2.0.
+
+Copyright (c) The Highway Project Authors. All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+1.  Redistributions of source code must retain the above copyright notice, this
+    list of conditions and the following disclaimer.
+
+2.  Redistributions in binary form must reproduce the above copyright notice,
+    this list of conditions and the following disclaimer in the documentation
+    and/or other materials provided with the distribution.
+
+3.  Neither the name of the copyright holder nor the names of its
+    contributors may be used to endorse or promote products derived from
+    this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+

CMakeLists.txt

@@ -29,7 +29,8 @@ endif()
 # Use the same checked-out core/C ABI as SwiftPM. A stock mlx-c release
 # omits the fork's mmap and compiler-cache lifetime entry points.
 if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/mlx/CMakeLists.txt" OR
-   NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/mlx-c/CMakeLists.txt")
+   NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/mlx-c/CMakeLists.txt" OR
+   NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/highway/CMakeLists.txt")
   message(FATAL_ERROR "Initialize the pinned dependencies with git submodule update --init --recursive")
 endif()
 FetchContent_Declare(
@@ -38,6 +39,11 @@ FetchContent_Declare(
 FetchContent_Declare(
   mlx-c
   SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/mlx-c")
+# ml-explore/mlx#3019's Highway kernels take Google Highway from this submodule
+# instead of fetching it at build time.
+set(MLX_HIGHWAY_SOURCE_DIR
+    "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/highway"
+    CACHE PATH "Google Highway sources")
 FetchContent_MakeAvailable(mlx-c)
 
 # swift-numerics

Libraries/CmlxCPUShim/CmlxCPUShim.cpp

@@ -0,0 +1,157 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#include "CmlxCPUShim.h"
+
+#if defined(MLX_USE_HIGHWAY_KERNELS) || defined(__linux__)
+#include "mlx/backend/cpu/precision.h"
+#endif
+
+#if defined(MLX_USE_HIGHWAY_KERNELS)
+
+#include "mlx/backend/cpu/highway_info.h"
+#include "mlx/backend/cpu/threading/common.h"
+#include "mlx/backend/cpu/threading/config.h"
+
+namespace cpu = mlx::core::cpu;
+namespace hi = mlx::core::cpu::highway_info;
+
+// vmlx_cpu_family lists hi::Family's enumerators, in order.
+static constexpr bool same(vmlx_cpu_family family, hi::Family cpp) {
+  return static_cast<int>(family) == static_cast<int>(cpp);
+}
+static_assert(
+    same(VMLX_CPU_FAMILY_QMM_AFFINE_DEQUANT, hi::Family::QmmAffineDequant));
+static_assert(same(VMLX_CPU_FAMILY_QMM_AFFINE_INT8, hi::Family::QmmAffineInt8));
+static_assert(same(VMLX_CPU_FAMILY_QMM_FP, hi::Family::QmmFp));
+static_assert(same(VMLX_CPU_FAMILY_RMS_NORM, hi::Family::RmsNorm));
+static_assert(same(VMLX_CPU_FAMILY_LAYER_NORM, hi::Family::LayerNorm));
+static_assert(same(VMLX_CPU_FAMILY_ROPE, hi::Family::Rope));
+static_assert(same(VMLX_CPU_FAMILY_SDPA, hi::Family::Sdpa));
+static_assert(same(VMLX_CPU_FAMILY_COUNT, hi::Family::Count));
+
+// The counters of `family`, and none for a value outside vmlx_cpu_family.
+static hi::FamilyStats stats_of(vmlx_cpu_family family) {
+  const auto index = static_cast<unsigned>(family);
+  if (index >= static_cast<unsigned>(VMLX_CPU_FAMILY_COUNT)) {
+    return {0, 0, 0};
+  }
+  return hi::stats(static_cast<hi::Family>(index));
+}
+
+extern "C" bool vmlx_cpu_highway_enabled(void) {
+  return true;
+}
+extern "C" int64_t vmlx_cpu_highway_compiled_targets(void) {
+  return hi::compiled_targets();
+}
+extern "C" int64_t vmlx_cpu_highway_supported_targets(void) {
+  return hi::supported_targets();
+}
+extern "C" bool vmlx_cpu_highway_set_targets_for_test(int64_t targets) {
+  // Dispatching a target this CPU lacks would fault (SIGILL). Clearing the
+  // restriction first makes supported_targets() the CPU's own set.
+  hi::set_targets_for_test(0);
+  if ((targets & ~hi::supported_targets()) != 0) {
+    return false;
+  }
+  hi::set_targets_for_test(targets);
+  return true;
+}
+extern "C" const char* vmlx_cpu_highway_target_name(int64_t target) {
+  return hi::target_name(target);
+}
+extern "C" bool vmlx_cpu_quantized_int8(void) {
+  return cpu::quantized_int8();
+}
+extern "C" void vmlx_cpu_set_quantized_int8(bool enabled) {
+  cpu::set_quantized_int8(enabled);
+}
+extern "C" int32_t vmlx_cpu_thread_count(void) {
+  return cpu::ThreadPool::instance().max_threads();
+}
+extern "C" const char* vmlx_cpu_thread_count_reason(void) {
+  return cpu::thread_config().reason;
+}
+extern "C" const char* vmlx_cpu_thread_setting_error(void) {
+  return cpu::thread_config().error.c_str();
+}
+extern "C" bool vmlx_cpu_openblas_pinned(void) {
+  return cpu::openblas_pinned();
+}
+extern "C" int64_t vmlx_cpu_executed_targets(vmlx_cpu_family family) {
+  return stats_of(family).executed;
+}
+extern "C" uint64_t vmlx_cpu_highway_calls(vmlx_cpu_family family) {
+  return stats_of(family).highway_calls;
+}
+extern "C" uint64_t vmlx_cpu_fallback_calls(vmlx_cpu_family family) {
+  return stats_of(family).fallback_calls;
+}
+extern "C" uint64_t vmlx_cpu_sgemm_column_splits(void) {
+  return hi::sgemm_column_splits();
+}
+extern "C" void vmlx_cpu_reset_counters(void) {
+  hi::reset_stats();
+}
+
+#else
+
+// A build without Highway kernels: no targets, no pool, no counts.
+extern "C" bool vmlx_cpu_highway_enabled(void) {
+  return false;
+}
+extern "C" int64_t vmlx_cpu_highway_compiled_targets(void) {
+  return 0;
+}
+extern "C" int64_t vmlx_cpu_highway_supported_targets(void) {
+  return 0;
+}
+extern "C" bool vmlx_cpu_highway_set_targets_for_test(int64_t targets) {
+  return targets == 0;
+}
+extern "C" const char* vmlx_cpu_highway_target_name(int64_t) {
+  return "";
+}
+#if defined(__linux__)
+// Linux compiles precision.cpp with or without Highway kernels, so the shim
+// reports the core's own switch, which changes nothing without them.
+extern "C" bool vmlx_cpu_quantized_int8(void) {
+  return mlx::core::cpu::quantized_int8();
+}
+extern "C" void vmlx_cpu_set_quantized_int8(bool enabled) {
+  mlx::core::cpu::set_quantized_int8(enabled);
+}
+#else
+extern "C" bool vmlx_cpu_quantized_int8(void) {
+  return false;
+}
+extern "C" void vmlx_cpu_set_quantized_int8(bool) {}
+#endif
+extern "C" int32_t vmlx_cpu_thread_count(void) {
+  return 1;
+}
+extern "C" const char* vmlx_cpu_thread_count_reason(void) {
+  return "no pool";
+}
+extern "C" const char* vmlx_cpu_thread_setting_error(void) {
+  return "";
+}
+extern "C" bool vmlx_cpu_openblas_pinned(void) {
+  return false;
+}
+extern "C" int64_t vmlx_cpu_executed_targets(vmlx_cpu_family) {
+  return 0;
+}
+extern "C" uint64_t vmlx_cpu_highway_calls(vmlx_cpu_family) {
+  return 0;
+}
+extern "C" uint64_t vmlx_cpu_fallback_calls(vmlx_cpu_family) {
+  return 0;
+}
+extern "C" uint64_t vmlx_cpu_sgemm_column_splits(void) {
+  return 0;
+}
+extern "C" void vmlx_cpu_reset_counters(void) {}
+
+#endif