Skip to content

Add support for x86 SIMD (SSE, AVX2, AVX512) #3019

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
dhiltgen wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

Add support for x86 SIMD (SSE, AVX2, AVX512) #3019

dhiltgen wants to merge 3 commits into from

Conversation

@dhiltgen
Copy link
Contributor

@dhiltgen dhiltgen commented Jan 19, 2026

Proposed changes

This implements CPU support for SSE, AVX2, and AVX512.

Performance benchmarks on Linux on Xeon Silver 4410T

Baseline

build/cpu/benchmarks/cpp/autograd 
Benchmarks for Device(cpu, 0)
Timing independent_value_and_grad ... 1045.1 msec
Timing combined_value_and_grad ... 820.53 msec
build/cpu/benchmarks/cpp/irregular_strides 
Benchmarks for Device(cpu, 0)
Timing (1D strided) mx::add ... 8.4961 msec
Timing (2D regular) mx::add ... 52.833 msec
Timing (2D mx::transpose) mx::add ... 164.15 msec
Timing (2D broadcast dim 0) mx::add ... 54.399 msec
Timing (2D broadcast dim 1) mx::add ... 65.15 msec
Timing (3D regular) mx::add ... 964.25 msec
Timing (3D mx::transpose) mx::add ... 125.86 msec
Timing (3D broadcast dim 0) mx::add ... 108.47 msec
Timing (3D broadcast dim 1) mx::add ... 121.23 msec
Timing (3D broadcast dim 2) mx::add ... 124.37 msec
Timing (3D broadcast dims 0, 1) mx::add ... 122.26 msec
Timing (3D broadcast dims 0, 2) mx::add ... 124.11 msec
Timing (3D broadcast dims 1, 2) mx::add ... 124.32 msec
Timing (4D regular) mx::add ... 243.4 msec
Timing (4D mx::transpose) mx::add ... 394.88 msec
Timing (4D broadcast dims 0) mx::add ... 245.68 msec
Timing (4D broadcast dims 0, 1) mx::add ... 246.26 msec
Timing (4D broadcast dims 0, 1, 2) mx::add ... 240.81 msec
Timing (4D broadcast dims 0, 1, 3) mx::add ... 247.51 msec
Timing (4D broadcast dims 0, 2) mx::add ... 241.58 msec
Timing (4D broadcast dims 0, 2, 3) mx::add ... 251.94 msec
Timing (4D broadcast dims 0, 3) mx::add ... 245.79 msec
Timing (4D broadcast dims 1) mx::add ... 242.37 msec
Timing (4D broadcast dims 1, 2) mx::add ... 243.04 msec
Timing (4D broadcast dims 1, 2, 3) mx::add ... 246.65 msec
Timing (4D broadcast dims 1, 3) mx::add ... 245.55 msec
Timing (4D broadcast dims 2) mx::add ... 239.72 msec
Timing (4D broadcast dims 2, 3) mx::add ... 250.05 msec
Timing (4D broadcast dims 3) mx::add ... 249.97 msec
Timing (3D contiguous) reshape_fn ... 0.029417 msec
Timing (3D mx::transpose) reshape_fn ... 15.293 msec
Timing (3D mx::transpose dims 1 2) reshape_fn ... 5.7285 msec
Timing (3D broadcast dim 0) reshape_fn ... 5.4893 msec
Timing (3D broadcast dim 1) reshape_fn ... 6.1788 msec
Timing (3D broadcast dim 2) reshape_fn ... 6.6059 msec
Timing (3D broadcast dims 0, 1) reshape_fn ... 6.4205 msec
Timing (3D broadcast dims 0, 2) reshape_fn ... 6.4971 msec
Timing (3D broadcast dims 1, 2) reshape_fn ... 5.8452 msec
Timing (3D broadcast dims 1, 2, 3) reshape_fn ... 0.069857 msec
Timing (1D strided) mx::astype ... 1.4186 msec
Timing (2D regular) mx::astype ... 7.7479 msec
Timing (2D mx::transpose) mx::astype ... 7.5769 msec
Timing (2D broadcast dim 0) mx::astype ... 0.037381 msec
Timing (2D broadcast dim 1) mx::astype ... 0.038433 msec
build/cpu/benchmarks/cpp/single_ops 
Benchmarks for Device(cpu, 0)
Timing full_fp32 ... 1.4704 msec
Timing zeros_fp32 ... 1.1875 msec
Timing ones_fp32 ... 1.1991 msec
Timing arange_fp32 ... 0.30664 msec
Timing (mx::float32 to mx::int32) mx::astype ... 1.8467 msec
Timing (mx::float32 to mx::uint32) mx::astype ... 1.7991 msec
Timing (mx::int32 to mx::float32) mx::astype ... 1.8071 msec
Timing (bool to mx::float32) mx::astype ... 1.9098 msec
Timing (bool to mx::int32) mx::astype ... 1.7379 msec
Timing (bool to mx::uint32) mx::astype ... 1.7207 msec
Timing mlx::core::abs ... 11.506 msec
Timing mx::negative ... 12.429 msec
Timing mx::sign ... 36.409 msec
Timing mx::square ... 9.8157 msec
Timing mlx::core::sqrt ... 24.597 msec
Timing mx::rsqrt ... 31.692 msec
Timing mlx::core::exp ... 172.39 msec
Timing mlx::core::log ... 15.905 msec
Timing mx::add ... 11.865 msec
Timing mx::subtract ... 12.049 msec
Timing mx::multiply ... 12.134 msec
Timing mx::divide ... 12.519 msec
Timing mx::maximum ... 22.939 msec
Timing mx::minimum ... 23.061 msec
Timing mx::where ... 29.546 msec
Timing (scalar) mx::add ... 12.806 msec
Timing (vector-scalar) mx::subtract ... 12.791 msec
Timing (scalar-vector) mx::subtract ... 12.687 msec
Timing (scalar) mx::multiply ... 12.574 msec
Timing (vector-scalar) mx::divide ... 13.364 msec
Timing (scalar-vector) mx::divide ... 13.894 msec
Timing (scalar-vector) mx::where ... 18.881 msec
Timing (scalar-scalar broadcast) mx::add ... 0.039647 msec
Timing (scalar-scalar broadcast) mx::subtract ... 0.040571 msec
Timing (scalar-scalar broadcast) mx::multiply ... 0.039894 msec
Timing (scalar-scalar broadcast) mx::divide ... 0.040526 msec
Timing (scalar-scalar broadcast) mx::where ... 0.044474 msec
Timing (non-strided) mx::add ... 73.777 msec
Timing (strided) mx::add ... 376.97 msec
Timing uniform ... 118.93 msec
Timing normal ... 306.51 msec
Timing mx::equal ... 11.684 msec
Timing mx::greater ... 11.375 msec
Timing mx::greater_equal ... 11.588 msec
Timing mx::less ... 11.567 msec
Timing mx::less_equal ... 11.426 msec
Timing matvec ... 0.12616 msec
Timing matvec_transpose ... 0.11048 msec
Timing mx::matmul ... 25.82 msec
Timing transpose_matmul ... 25.38 msec
Timing sum_all ... 120.85 msec
Timing sum_along_0 ... 122.71 msec
Timing sum_along_1 ... 121.51 msec
Timing prod_all ... 124.39 msec
Timing all_true ... 118.07 msec
Timing all_along_0 ... 147.34 msec
Timing all_along_1 ... 118.97 msec
Timing any_true ... 116.83 msec
Timing argmin_along_0 ... 65.201 msec
Timing argmin_along_1 ... 35.45 msec
Timing max_along_0 ... 174.95 msec
Timing max_along_1 ... 167.43 msec
Timing min_along_0 ... 173.38 msec
Timing min_along_1 ... 157.96 msec
Timing embedding_lookup ... 0.13422 msec
Timing single_element_lookup ... 31.588 msec
Timing embedding_update ... 10.008 msec
Timing embedding_add ... 10.602 msec
Timing single_element_update ... 40.991 msec
Timing single_element_add ... 55.524 msec
Timing divmod_fused ... 0.11231 msec
Timing divmod_separate ... 0.13845 msec

SSE

build/cpu-sse/benchmarks/cpp/autograd 
Benchmarks for Device(cpu, 0)
Timing independent_value_and_grad ... 139.96 msec
Timing combined_value_and_grad ... 84.91 msec
build/cpu-sse/benchmarks/cpp/irregular_strides 
Benchmarks for Device(cpu, 0)
Timing (1D strided) mx::add ... 4.7285 msec
Timing (2D regular) mx::add ... 12.015 msec
Timing (2D mx::transpose) mx::add ... 67.34 msec
Timing (2D broadcast dim 0) mx::add ... 11.731 msec
Timing (2D broadcast dim 1) mx::add ... 10.165 msec
Timing (3D regular) mx::add ... 39.055 msec
Timing (3D mx::transpose) mx::add ... 100.49 msec
Timing (3D broadcast dim 0) mx::add ... 38.362 msec
Timing (3D broadcast dim 1) mx::add ... 38.756 msec
Timing (3D broadcast dim 2) mx::add ... 35.526 msec
Timing (3D broadcast dims 0, 1) mx::add ... 38.738 msec
Timing (3D broadcast dims 0, 2) mx::add ... 35.918 msec
Timing (3D broadcast dims 1, 2) mx::add ... 35.553 msec
Timing (4D regular) mx::add ... 78.329 msec
Timing (4D mx::transpose) mx::add ... 214.09 msec
Timing (4D broadcast dims 0) mx::add ... 77.425 msec
Timing (4D broadcast dims 0, 1) mx::add ... 77.135 msec
Timing (4D broadcast dims 0, 1, 2) mx::add ... 77.367 msec
Timing (4D broadcast dims 0, 1, 3) mx::add ... 71.45 msec
Timing (4D broadcast dims 0, 2) mx::add ... 77.459 msec
Timing (4D broadcast dims 0, 2, 3) mx::add ... 71.663 msec
Timing (4D broadcast dims 0, 3) mx::add ... 71.878 msec
Timing (4D broadcast dims 1) mx::add ... 77.451 msec
Timing (4D broadcast dims 1, 2) mx::add ... 77.567 msec
Timing (4D broadcast dims 1, 2, 3) mx::add ... 71.602 msec
Timing (4D broadcast dims 1, 3) mx::add ... 72.658 msec
Timing (4D broadcast dims 2) mx::add ... 77.39 msec
Timing (4D broadcast dims 2, 3) mx::add ... 71.081 msec
Timing (4D broadcast dims 3) mx::add ... 71.483 msec
Timing (3D contiguous) reshape_fn ... 0.07048 msec
Timing (3D mx::transpose) reshape_fn ... 7.3012 msec
Timing (3D mx::transpose dims 1 2) reshape_fn ... 4.3205 msec
Timing (3D broadcast dim 0) reshape_fn ... 4.0235 msec
Timing (3D broadcast dim 1) reshape_fn ... 4.2869 msec
Timing (3D broadcast dim 2) reshape_fn ... 4.3221 msec
Timing (3D broadcast dims 0, 1) reshape_fn ... 4.2899 msec
Timing (3D broadcast dims 0, 2) reshape_fn ... 4.3077 msec
Timing (3D broadcast dims 1, 2) reshape_fn ... 4.0351 msec
Timing (3D broadcast dims 1, 2, 3) reshape_fn ... 0.071582 msec
Timing (1D strided) mx::astype ... 1.0169 msec
Timing (2D regular) mx::astype ... 4.5678 msec
Timing (2D mx::transpose) mx::astype ... 4.4884 msec
Timing (2D broadcast dim 0) mx::astype ... 0.08198 msec
Timing (2D broadcast dim 1) mx::astype ... 0.082054 msec
build/cpu-sse/benchmarks/cpp/single_ops 
Benchmarks for Device(cpu, 0)
Timing full_fp32 ... 1.2566 msec
Timing zeros_fp32 ... 0.99179 msec
Timing ones_fp32 ... 0.96662 msec
Timing arange_fp32 ... 0.26623 msec
Timing (mx::float32 to mx::int32) mx::astype ... 1.0834 msec
Timing (mx::float32 to mx::uint32) mx::astype ... 1.193 msec
Timing (mx::int32 to mx::float32) mx::astype ... 1.2657 msec
Timing (bool to mx::float32) mx::astype ... 1.1877 msec
Timing (bool to mx::int32) mx::astype ... 1.0881 msec
Timing (bool to mx::uint32) mx::astype ... 1.0842 msec
Timing mlx::core::abs ... 2.4867 msec
Timing mx::negative ... 2.1311 msec
Timing mx::sign ... 5.7121 msec
Timing mx::square ... 2.0804 msec
Timing mlx::core::sqrt ... 1.922 msec
Timing mx::rsqrt ... 2.0272 msec
Timing mlx::core::exp ... 6.3711 msec
Timing mlx::core::log ... 7.2586 msec
Timing mx::add ... 2.8043 msec
Timing mx::subtract ... 2.8846 msec
Timing mx::multiply ... 2.7926 msec
Timing mx::divide ... 2.988 msec
Timing mx::maximum ... 6.601 msec
Timing mx::minimum ... 6.6399 msec
Timing mx::where ... 18.911 msec
Timing (scalar) mx::add ... 2.4642 msec
Timing (vector-scalar) mx::subtract ... 2.4544 msec
Timing (scalar-vector) mx::subtract ... 2.688 msec
Timing (scalar) mx::multiply ... 2.431 msec
Timing (vector-scalar) mx::divide ... 2.4385 msec
Timing (scalar-vector) mx::divide ... 2.6062 msec
Timing (scalar-vector) mx::where ... 10.712 msec
Timing (scalar-scalar broadcast) mx::add ... 0.087141 msec
Timing (scalar-scalar broadcast) mx::subtract ... 0.088643 msec
Timing (scalar-scalar broadcast) mx::multiply ... 0.086912 msec
Timing (scalar-scalar broadcast) mx::divide ... 0.088069 msec
Timing (scalar-scalar broadcast) mx::where ... 0.10348 msec
Timing (non-strided) mx::add ... 18.839 msec
Timing (strided) mx::add ... 109.83 msec
Timing uniform ... 52.687 msec
Timing normal ... 108.55 msec
Timing mx::equal ... 3.3605 msec
Timing mx::greater ... 3.416 msec
Timing mx::greater_equal ... 3.436 msec
Timing mx::less ... 3.3056 msec
Timing mx::less_equal ... 3.2785 msec
Timing matvec ... 0.13509 msec
Timing matvec_transpose ... 0.15689 msec
Timing mx::matmul ... 1.867 msec
Timing transpose_matmul ... 1.8566 msec
Timing sum_all ... 26.451 msec
Timing sum_along_0 ... 27.831 msec
Timing sum_along_1 ... 28.219 msec
Timing prod_all ... 25.506 msec
Timing all_true ... 68.328 msec
Timing all_along_0 ... 79.4 msec
Timing all_along_1 ... 68.867 msec
Timing any_true ... 63.689 msec
Timing argmin_along_0 ... 33.608 msec
Timing argmin_along_1 ... 19.081 msec
Timing max_along_0 ... 66.831 msec
Timing max_along_1 ... 64.081 msec
Timing min_along_0 ... 66.308 msec
Timing min_along_1 ... 64.064 msec
Timing embedding_lookup ... 0.1409 msec
Timing single_element_lookup ... 12.071 msec
Timing embedding_update ... 5.8033 msec
Timing embedding_add ... 5.7742 msec
Timing single_element_update ... 21.731 msec
Timing single_element_add ... 22.277 msec
Timing divmod_fused ... 0.15112 msec
Timing divmod_separate ... 0.17786 msec

AVX2

build/cpu-avx2/benchmarks/cpp/autograd 
Benchmarks for Device(cpu, 0)
Timing independent_value_and_grad ... 159.57 msec
Timing combined_value_and_grad ... 76.709 msec
build/cpu-avx2/benchmarks/cpp/irregular_strides 
Benchmarks for Device(cpu, 0)
Timing (1D strided) mx::add ... 5.0248 msec
Timing (2D regular) mx::add ... 8.3063 msec
Timing (2D mx::transpose) mx::add ... 95.716 msec
Timing (2D broadcast dim 0) mx::add ... 7.4883 msec
Timing (2D broadcast dim 1) mx::add ... 6.8755 msec
Timing (3D regular) mx::add ... 32.843 msec
Timing (3D mx::transpose) mx::add ... 109.13 msec
Timing (3D broadcast dim 0) mx::add ... 31.669 msec
Timing (3D broadcast dim 1) mx::add ... 32.777 msec
Timing (3D broadcast dim 2) mx::add ... 30.29 msec
Timing (3D broadcast dims 0, 1) mx::add ... 32.601 msec
Timing (3D broadcast dims 0, 2) mx::add ... 28.6 msec
Timing (3D broadcast dims 1, 2) mx::add ... 28.487 msec
Timing (4D regular) mx::add ... 66.448 msec
Timing (4D mx::transpose) mx::add ... 210.8 msec
Timing (4D broadcast dims 0) mx::add ... 62.956 msec
Timing (4D broadcast dims 0, 1) mx::add ... 61.08 msec
Timing (4D broadcast dims 0, 1, 2) mx::add ... 62.921 msec
Timing (4D broadcast dims 0, 1, 3) mx::add ... 58.33 msec
Timing (4D broadcast dims 0, 2) mx::add ... 63.231 msec
Timing (4D broadcast dims 0, 2, 3) mx::add ... 58.413 msec
Timing (4D broadcast dims 0, 3) mx::add ... 57.326 msec
Timing (4D broadcast dims 1) mx::add ... 61.786 msec
Timing (4D broadcast dims 1, 2) mx::add ... 61.027 msec
Timing (4D broadcast dims 1, 2, 3) mx::add ... 55.852 msec
Timing (4D broadcast dims 1, 3) mx::add ... 55.886 msec
Timing (4D broadcast dims 2) mx::add ... 60.811 msec
Timing (4D broadcast dims 2, 3) mx::add ... 55.67 msec
Timing (4D broadcast dims 3) mx::add ... 55.582 msec
Timing (3D contiguous) reshape_fn ... 0.067288 msec
Timing (3D mx::transpose) reshape_fn ... 7.3041 msec
Timing (3D mx::transpose dims 1 2) reshape_fn ... 4.3296 msec
Timing (3D broadcast dim 0) reshape_fn ... 4.0586 msec
Timing (3D broadcast dim 1) reshape_fn ... 4.1292 msec
Timing (3D broadcast dim 2) reshape_fn ... 4.0705 msec
Timing (3D broadcast dims 0, 1) reshape_fn ... 4.0959 msec
Timing (3D broadcast dims 0, 2) reshape_fn ... 4.076 msec
Timing (3D broadcast dims 1, 2) reshape_fn ... 4.0444 msec
Timing (3D broadcast dims 1, 2, 3) reshape_fn ... 0.069308 msec
Timing (1D strided) mx::astype ... 1.0742 msec
Timing (2D regular) mx::astype ... 4.5812 msec
Timing (2D mx::transpose) mx::astype ... 4.5891 msec
Timing (2D broadcast dim 0) mx::astype ... 0.086623 msec
Timing (2D broadcast dim 1) mx::astype ... 0.08862 msec
build/cpu-avx2/benchmarks/cpp/single_ops 
Benchmarks for Device(cpu, 0)
Timing full_fp32 ... 1.4496 msec
Timing zeros_fp32 ... 1.0436 msec
Timing ones_fp32 ... 1.0855 msec
Timing arange_fp32 ... 0.28111 msec
Timing (mx::float32 to mx::int32) mx::astype ... 1.1603 msec
Timing (mx::float32 to mx::uint32) mx::astype ... 1.2722 msec
Timing (mx::int32 to mx::float32) mx::astype ... 1.169 msec
Timing (bool to mx::float32) mx::astype ... 1.1483 msec
Timing (bool to mx::int32) mx::astype ... 1.1276 msec
Timing (bool to mx::uint32) mx::astype ... 1.0845 msec
Timing mlx::core::abs ... 1.4866 msec
Timing mx::negative ... 1.4133 msec
Timing mx::sign ... 3.5585 msec
Timing mx::square ... 1.3479 msec
Timing mlx::core::sqrt ... 1.354 msec
Timing mx::rsqrt ... 1.3551 msec
Timing mlx::core::exp ... 5.9196 msec
Timing mlx::core::log ... 6.1033 msec
Timing mx::add ... 1.9016 msec
Timing mx::subtract ... 1.9204 msec
Timing mx::multiply ... 1.8753 msec
Timing mx::divide ... 1.9132 msec
Timing mx::maximum ... 4.1294 msec
Timing mx::minimum ... 3.9875 msec
Timing mx::where ... 19.369 msec
Timing (scalar) mx::add ... 1.6116 msec
Timing (vector-scalar) mx::subtract ... 1.6221 msec
Timing (scalar-vector) mx::subtract ... 1.676 msec
Timing (scalar) mx::multiply ... 1.6324 msec
Timing (vector-scalar) mx::divide ... 1.6095 msec
Timing (scalar-vector) mx::divide ... 1.7195 msec
Timing (scalar-vector) mx::where ... 10.637 msec
Timing (scalar-scalar broadcast) mx::add ... 0.094488 msec
Timing (scalar-scalar broadcast) mx::subtract ... 0.09362 msec
Timing (scalar-scalar broadcast) mx::multiply ... 0.093646 msec
Timing (scalar-scalar broadcast) mx::divide ... 0.094245 msec
Timing (scalar-scalar broadcast) mx::where ... 0.10962 msec
Timing (non-strided) mx::add ... 11.898 msec
Timing (strided) mx::add ... 110.47 msec
Timing uniform ... 41.87 msec
Timing normal ... 72.403 msec
Timing mx::equal ... 4.1518 msec
Timing mx::greater ... 4.3036 msec
Timing mx::greater_equal ... 4.2984 msec
Timing mx::less ... 4.2256 msec
Timing mx::less_equal ... 4.1683 msec
Timing matvec ... 0.14507 msec
Timing matvec_transpose ... 0.16682 msec
Timing mx::matmul ... 1.9423 msec
Timing transpose_matmul ... 1.8716 msec
Timing sum_all ... 13.26 msec
Timing sum_along_0 ... 18.401 msec
Timing sum_along_1 ... 13.413 msec
Timing prod_all ... 13.464 msec
Timing all_true ... 59.969 msec
Timing all_along_0 ... 78.839 msec
Timing all_along_1 ... 60.077 msec
Timing any_true ... 61.031 msec
Timing argmin_along_0 ... 33.479 msec
Timing argmin_along_1 ... 19.232 msec
Timing max_along_0 ... 39.924 msec
Timing max_along_1 ... 39.488 msec
Timing min_along_0 ... 40.309 msec
Timing min_along_1 ... 39.403 msec
Timing embedding_lookup ... 0.14394 msec
Timing single_element_lookup ... 11.911 msec
Timing embedding_update ... 5.7638 msec
Timing embedding_add ... 5.6913 msec
Timing single_element_update ... 21.044 msec
Timing single_element_add ... 21.852 msec
Timing divmod_fused ... 0.21375 msec
Timing divmod_separate ... 0.28181 msec

AVX512

build/cpu-avx512/benchmarks/cpp/autograd 
Benchmarks for Device(cpu, 0)
Timing independent_value_and_grad ... 109.05 msec
Timing combined_value_and_grad ... 63.535 msec
build/cpu-avx512/benchmarks/cpp/irregular_strides 
Benchmarks for Device(cpu, 0)
Timing (1D strided) mx::add ... 4.8721 msec
Timing (2D regular) mx::add ... 5.243 msec
Timing (2D mx::transpose) mx::add ... 67.211 msec
Timing (2D broadcast dim 0) mx::add ... 4.9516 msec
Timing (2D broadcast dim 1) mx::add ... 4.0123 msec
Timing (3D regular) mx::add ... 25.265 msec
Timing (3D mx::transpose) mx::add ... 106.37 msec
Timing (3D broadcast dim 0) mx::add ... 24.588 msec
Timing (3D broadcast dim 1) mx::add ... 25.514 msec
Timing (3D broadcast dim 2) mx::add ... 28.413 msec
Timing (3D broadcast dims 0, 1) mx::add ... 26.843 msec
Timing (3D broadcast dims 0, 2) mx::add ... 24.925 msec
Timing (3D broadcast dims 1, 2) mx::add ... 24.157 msec
Timing (4D regular) mx::add ... 49.665 msec
Timing (4D mx::transpose) mx::add ... 218.96 msec
Timing (4D broadcast dims 0) mx::add ... 50.408 msec
Timing (4D broadcast dims 0, 1) mx::add ... 49.538 msec
Timing (4D broadcast dims 0, 1, 2) mx::add ... 49.76 msec
Timing (4D broadcast dims 0, 1, 3) mx::add ... 45.57 msec
Timing (4D broadcast dims 0, 2) mx::add ... 51.235 msec
Timing (4D broadcast dims 0, 2, 3) mx::add ... 47.242 msec
Timing (4D broadcast dims 0, 3) mx::add ... 47.109 msec
Timing (4D broadcast dims 1) mx::add ... 50.02 msec
Timing (4D broadcast dims 1, 2) mx::add ... 49.516 msec
Timing (4D broadcast dims 1, 2, 3) mx::add ... 46.862 msec
Timing (4D broadcast dims 1, 3) mx::add ... 45.302 msec
Timing (4D broadcast dims 2) mx::add ... 49.689 msec
Timing (4D broadcast dims 2, 3) mx::add ... 44.973 msec
Timing (4D broadcast dims 3) mx::add ... 47.217 msec
Timing (3D contiguous) reshape_fn ... 0.069337 msec
Timing (3D mx::transpose) reshape_fn ... 7.3935 msec
Timing (3D mx::transpose dims 1 2) reshape_fn ... 4.169 msec
Timing (3D broadcast dim 0) reshape_fn ... 4.0922 msec
Timing (3D broadcast dim 1) reshape_fn ... 4.1342 msec
Timing (3D broadcast dim 2) reshape_fn ... 4.1673 msec
Timing (3D broadcast dims 0, 1) reshape_fn ... 4.1226 msec
Timing (3D broadcast dims 0, 2) reshape_fn ... 4.1619 msec
Timing (3D broadcast dims 1, 2) reshape_fn ... 4.0794 msec
Timing (3D broadcast dims 1, 2, 3) reshape_fn ... 0.071931 msec
Timing (1D strided) mx::astype ... 1.0185 msec
Timing (2D regular) mx::astype ... 4.5763 msec
Timing (2D mx::transpose) mx::astype ... 4.5577 msec
Timing (2D broadcast dim 0) mx::astype ... 0.089004 msec
Timing (2D broadcast dim 1) mx::astype ... 0.086812 msec
build/cpu-avx512/benchmarks/cpp/single_ops 
Benchmarks for Device(cpu, 0)
Timing full_fp32 ... 1.2871 msec
Timing zeros_fp32 ... 1.4398 msec
Timing ones_fp32 ... 1.4539 msec
Timing arange_fp32 ... 0.27884 msec
Timing (mx::float32 to mx::int32) mx::astype ... 1.1837 msec
Timing (mx::float32 to mx::uint32) mx::astype ... 1.1805 msec
Timing (mx::int32 to mx::float32) mx::astype ... 1.0768 msec
Timing (bool to mx::float32) mx::astype ... 1.1042 msec
Timing (bool to mx::int32) mx::astype ... 1.0646 msec
Timing (bool to mx::uint32) mx::astype ... 1.0665 msec
Timing mlx::core::abs ... 0.94824 msec
Timing mx::negative ... 1.0396 msec
Timing mx::sign ... 1.8795 msec
Timing mx::square ... 0.94696 msec
Timing mlx::core::sqrt ... 0.97706 msec
Timing mx::rsqrt ... 1.0015 msec
Timing mlx::core::exp ... 6.0047 msec
Timing mlx::core::log ... 6.3887 msec
Timing mx::add ... 1.2969 msec
Timing mx::subtract ... 1.2823 msec
Timing mx::multiply ... 1.2448 msec
Timing mx::divide ... 1.2415 msec
Timing mx::maximum ... 2.1673 msec
Timing mx::minimum ... 2.1821 msec
Timing mx::where ... 19.398 msec
Timing (scalar) mx::add ... 1.0311 msec
Timing (vector-scalar) mx::subtract ... 1.0203 msec
Timing (scalar-vector) mx::subtract ... 1.0587 msec
Timing (scalar) mx::multiply ... 1.1047 msec
Timing (vector-scalar) mx::divide ... 1.1687 msec
Timing (scalar-vector) mx::divide ... 1.1195 msec
Timing (scalar-vector) mx::where ... 10.633 msec
Timing (scalar-scalar broadcast) mx::add ... 0.097409 msec
Timing (scalar-scalar broadcast) mx::subtract ... 0.095516 msec
Timing (scalar-scalar broadcast) mx::multiply ... 0.095917 msec
Timing (scalar-scalar broadcast) mx::divide ... 0.095564 msec
Timing (scalar-scalar broadcast) mx::where ... 0.1068 msec
Timing (non-strided) mx::add ... 9.2285 msec
Timing (strided) mx::add ... 106.65 msec
Timing uniform ... 43.132 msec
Timing normal ... 64.343 msec
Timing mx::equal ... 2.8174 msec
Timing mx::greater ... 3.0892 msec
Timing mx::greater_equal ... 2.7417 msec
Timing mx::less ... 2.9577 msec
Timing mx::less_equal ... 2.8655 msec
Timing matvec ... 0.1461 msec
Timing matvec_transpose ... 0.16985 msec
Timing mx::matmul ... 1.8523 msec
Timing transpose_matmul ... 1.8644 msec
Timing sum_all ... 7.5322 msec
Timing sum_along_0 ... 11.248 msec
Timing sum_along_1 ... 8.4241 msec
Timing prod_all ... 7.9362 msec
Timing all_true ... 66.383 msec
Timing all_along_0 ... 75.197 msec
Timing all_along_1 ... 62.047 msec
Timing any_true ... 62.982 msec
Timing argmin_along_0 ... 33.948 msec
Timing argmin_along_1 ... 19.316 msec
Timing max_along_0 ... 40.734 msec
Timing max_along_1 ... 41.681 msec
Timing min_along_0 ... 40.69 msec
Timing min_along_1 ... 41.707 msec
Timing embedding_lookup ... 0.14615 msec
Timing single_element_lookup ... 11.918 msec
Timing embedding_update ... 5.6884 msec
Timing embedding_add ... 5.843 msec
Timing single_element_update ... 21.232 msec
Timing single_element_add ... 21.481 msec
Timing divmod_fused ... 0.14024 msec
Timing divmod_separate ... 0.19276 msec

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@dhiltgen dhiltgen changed the title Add support for x85 SIMD (SSE, AVX2, AVX512) Add support for x86 SIMD (SSE, AVX2, AVX512) Jan 19, 2026
awni
awni reviewed Jan 21, 2026
View reviewed changes
mlx/backend/cpu/simd/type.h
Comment on lines +15 to +23
#if !defined(MLX_USE_ACCELERATE)
#if defined(__AVX512F__)
#include "mlx/backend/cpu/simd/avx512_simd.h"
#elif defined(__AVX2__)
#include "mlx/backend/cpu/simd/avx_simd.h"
#elif defined(__SSE4_2__)
#include "mlx/backend/cpu/simd/sse_simd.h"
#endif
#endif
Copy link
Member

@awni awni Jan 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm wondering if this will break our linux x86 distribution in some cases. If we build with avx512 then someone tries to run it on a machine which doesn't support avx512 it will crash right?

Copy link
Member

@awni awni Jan 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually it looks like just the lowest level is enabled by default. So we should be ok.

@awni
Copy link
Member

awni commented Jan 23, 2026

@dhiltgen what are you thinking for next steps here?

I might suggest we split this out into multiple PRs to make it easier to review and incorporate. The first PR could be the basic SSE backend for X86 which we should definitely integrate. Following that we could add the extra back-ends (there is a question of how to tests those as well).

We will probably also want a neon-only back-end for linux ARM (i.e. no through accelerate).

@dhiltgen
Copy link
Contributor Author

dhiltgen commented Jan 27, 2026

Splitting up to smaller chunks sounds like a reasonable approach.

I'll probably keep this in draft for a bit, while we focus on full GPU load for best performance.

@awni
Copy link
Member

awni commented Jan 28, 2026

Sounds good!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@awni awni awni left review comments

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@dhiltgen @awni