Use Accelerate for fast matmul blocks on macOS

FrancescAlted · FrancescAlted · commit 13063c25a815 · 2026-03-23T13:34:14.000+01:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -43,6 +43,8 @@ add_custom_command(
 
 # ...and add it to the target
 Python_add_library(blosc2_ext MODULE blosc2_ext.c WITH_SOABI)
+target_sources(blosc2_ext PRIVATE src/blosc2/matmul_kernels.c)
+target_include_directories(blosc2_ext PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src/blosc2)
 
 # We need to link against NumPy
 target_link_libraries(blosc2_ext PRIVATE Python::NumPy)
@@ -70,6 +72,9 @@ FetchContent_MakeAvailable(miniexpr)
 
 # Link against miniexpr static library
 target_link_libraries(blosc2_ext PRIVATE miniexpr_static)
+if(APPLE)
+  target_link_libraries(blosc2_ext PRIVATE "-framework Accelerate")
+endif()
 
 target_compile_features(blosc2_ext PRIVATE c_std_11)
 if(WIN32 AND CMAKE_C_COMPILER_ID STREQUAL "Clang")
diff --git a/bench/ndarray/matmul_path_compare.py b/bench/ndarray/matmul_path_compare.py
@@ -56,6 +56,7 @@ def set_path_mode(mode: str) -> bool:
 
 def run_case(
     mode: str,
+    block_backend: str,
     repeats: int,
     shape_a: tuple[int, ...],
     shape_b: tuple[int, ...],
@@ -73,8 +74,10 @@ def run_case(
         warnings.simplefilter("ignore", RuntimeWarning)
         expected = np.matmul(a_np, b_np)
     original_flag = set_path_mode(mode)
+    original_block_backend = blosc2.blosc2_ext.get_matmul_block_backend()
     original_set_pref_matmul = blosc2.NDArray._set_pref_matmul
     selected_paths = []
+    selected_block_backend = None
     times = []
     result = None
 
@@ -83,7 +86,9 @@ def wrapped_set_pref_matmul(self, inputs, fp_accuracy):
         return original_set_pref_matmul(self, inputs, fp_accuracy)
 
     blosc2.NDArray._set_pref_matmul = wrapped_set_pref_matmul
+    blosc2.blosc2_ext.set_matmul_block_backend(block_backend)
     try:
+        selected_block_backend = blosc2.blosc2_ext.get_selected_matmul_block_backend()
         for _ in range(repeats):
             before = len(selected_paths)
             t0 = time.perf_counter()
@@ -97,6 +102,7 @@ def wrapped_set_pref_matmul(self, inputs, fp_accuracy):
     finally:
         blosc2.NDArray._set_pref_matmul = original_set_pref_matmul
         linalg.try_miniexpr = original_flag
+        blosc2.blosc2_ext.set_matmul_block_backend(original_block_backend)
 
     if result is None:
         raise RuntimeError("matmul did not produce a result")
@@ -114,6 +120,8 @@ def wrapped_set_pref_matmul(self, inputs, fp_accuracy):
         "gflops_best": expected_gflops(shape_a, shape_b, best),
         "gflops_median": expected_gflops(shape_a, shape_b, median),
         "correct": True,
+        "configured_block_backend": block_backend,
+        "selected_block_backend": selected_block_backend,
         "selected_paths": selected_paths,
         "selected_path": selected_paths[0] if selected_paths and len(set(selected_paths)) == 1 else "mixed",
     }
@@ -132,6 +140,12 @@ def main() -> None:
     parser.add_argument("--blocks-out", default="100,100", help="Comma-separated block shape for output.")
     parser.add_argument("--repeats", type=int, default=250)
     parser.add_argument("--modes", nargs="+", default=["chunked", "fast", "auto"], choices=["chunked", "fast", "auto"])
+    parser.add_argument(
+        "--block-backend",
+        default="auto",
+        choices=["auto", "naive", "accelerate"],
+        help="Kernel backend for the fast matmul block path.",
+    )
     parser.add_argument("--json", action="store_true", help="Emit full JSON instead of a compact text summary.")
     args = parser.parse_args()
 
@@ -150,6 +164,7 @@ def main() -> None:
         results.append(
             run_case(
                 mode,
+                args.block_backend,
                 args.repeats,
                 shape_a,
                 shape_b,
@@ -173,6 +188,7 @@ def main() -> None:
         "blocks_b": blocks_b,
         "chunks_out": chunks_out,
         "blocks_out": blocks_out,
+        "block_backend": args.block_backend,
         "results": results,
     }
 
@@ -184,36 +200,27 @@ def main() -> None:
         print(json.dumps(summary, indent=2, sort_keys=True))
         return
 
-    print(
-        "case",
-        json.dumps(
-            {
-                "shape_a": shape_a,
-                "shape_b": shape_b,
-                "dtype": str(dtype),
-                "chunks_out": chunks_out,
-                "blocks_out": blocks_out,
-            },
-            sort_keys=True,
-        ),
-    )
+    print("Matmul path comparison")
+    print(f"  A shape: {shape_a}")
+    print(f"  B shape: {shape_b}")
+    print(f"  dtype: {dtype}")
+    print(f"  chunks A/B/out: {chunks_a} / {chunks_b} / {chunks_out}")
+    print(f"  blocks A/B/out: {blocks_a} / {blocks_b} / {blocks_out}")
+    print(f"  repeats: {args.repeats}")
+    print(f"  fast block backend: {args.block_backend}")
     for item in results:
+        gflops_best = "-" if item["gflops_best"] is None else f"{item['gflops_best']:.3f}"
         print(
-            "result",
-            json.dumps(
-                {
-                    "mode": item["mode"],
-                    "best_s": round(item["best_s"], 6),
-                    "median_s": round(item["median_s"], 6),
-                    "gflops_best": None if item["gflops_best"] is None else round(item["gflops_best"], 3),
-                    "correct": item["correct"],
-                    "selected_path": item["selected_path"],
-                },
-                sort_keys=True,
-            ),
+            f"{item['mode']:>7}: "
+            f"best={item['best_s']:.6f}s "
+            f"median={item['median_s']:.6f}s "
+            f"gflops={gflops_best} "
+            f"path={item['selected_path']} "
+            f"block_backend={item['selected_block_backend']} "
+            f"correct={item['correct']}"
         )
     if "speedup_fast_vs_chunked" in summary:
-        print("speedup", json.dumps({"fast_vs_chunked": round(summary["speedup_fast_vs_chunked"], 3)}, sort_keys=True))
+        print(f"Speedup fast vs chunked: {summary['speedup_fast_vs_chunked']:.3f}x")
 
 
 if __name__ == "__main__":
diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx
@@ -12,6 +12,8 @@ import dataclasses
 import ast
 import atexit
 import pathlib
+import time
+import warnings
 
 import _ctypes
 
@@ -62,6 +64,21 @@ ctypedef fused T:
 cdef extern from "<stdio.h>":
     int printf(const char *format, ...) nogil
 
+cdef extern from "matmul_kernels.h":
+    ctypedef enum b2_matmul_backend:
+        B2_MATMUL_BACKEND_AUTO
+        B2_MATMUL_BACKEND_NAIVE
+        B2_MATMUL_BACKEND_ACCELERATE
+
+    int b2_has_accelerate() nogil
+    void b2_set_matmul_backend(int backend) nogil
+    int b2_get_matmul_backend() nogil
+    int b2_get_selected_matmul_backend() nogil
+    const char *b2_get_matmul_backend_name() nogil
+    const char *b2_get_selected_matmul_backend_name() nogil
+    int b2_gemm_accelerate_f32(const float *a, const float *b, float *c, int m, int k, int n) nogil
+    int b2_gemm_accelerate_f64(const double *a, const double *b, double *c, int m, int k, int n) nogil
+
 cdef extern from "blosc2.h":
 
     ctypedef enum:
@@ -2384,6 +2401,7 @@ cdef int aux_matmul(mm_udata *udata, int64_t nchunk, int32_t nblock, void *param
     cdef int nchunk_ = nchunk
     cdef int coord, batch, batch_, batches = 1
     cdef int out_chunk_nrows, out_chunk_ncols, out_block_nrows, out_block_ncols
+    cdef int selected_backend = b2_get_selected_matmul_backend()
 
     # batches = sum(strides[i]*elcoords[i])
     for i in range(ndim - 2):
@@ -2487,9 +2505,43 @@ cdef int aux_matmul(mm_udata *udata, int64_t nchunk, int32_t nblock, void *param
                     offset += coord * udata.el_strides[0][i]
                 if typecode == 0:
                     if typesize == 4:
-                        rc = matmul_block_kernel[float](<float*>input_buffers[0] + offsetA, <float*>input_buffers[1] + offsetB, <float*>params_output + offset, p, q, r)
+                        if selected_backend == B2_MATMUL_BACKEND_ACCELERATE:
+                            rc = b2_gemm_accelerate_f32(
+                                <float*>input_buffers[0] + offsetA,
+                                <float*>input_buffers[1] + offsetB,
+                                <float*>params_output + offset,
+                                p,
+                                q,
+                                r,
+                            )
+                        else:
+                            rc = matmul_block_kernel[float](
+                                <float*>input_buffers[0] + offsetA,
+                                <float*>input_buffers[1] + offsetB,
+                                <float*>params_output + offset,
+                                p,
+                                q,
+                                r,
+                            )
                     else:
-                        rc = matmul_block_kernel[double](<double*>input_buffers[0] + offsetA, <double*>input_buffers[1] + offsetB, <double*>params_output + offset, p, q, r)
+                        if selected_backend == B2_MATMUL_BACKEND_ACCELERATE:
+                            rc = b2_gemm_accelerate_f64(
+                                <double*>input_buffers[0] + offsetA,
+                                <double*>input_buffers[1] + offsetB,
+                                <double*>params_output + offset,
+                                p,
+                                q,
+                                r,
+                            )
+                        else:
+                            rc = matmul_block_kernel[double](
+                                <double*>input_buffers[0] + offsetA,
+                                <double*>input_buffers[1] + offsetB,
+                                <double*>params_output + offset,
+                                p,
+                                q,
+                                r,
+                            )
                 elif typecode == 1:
                     if typesize == 4:
                         rc = matmul_block_kernel[int32_t](<int32_t*>input_buffers[0] + offsetA, <int32_t*>input_buffers[1] + offsetB, <int32_t*>params_output + offset, p, q, r)
@@ -3999,3 +4051,22 @@ def squeeze(arr1: NDArray, axis_mask: list[bool]) -> blosc2.NDArray:
     new_base = arr1 if arr1.base is None else arr1.base
     return blosc2.NDArray(_schunk=PyCapsule_New(view.sc, <char *> "blosc2_schunk*", NULL),
                         _array=PyCapsule_New(view, <char *> "b2nd_array_t*", NULL), _base=new_base)
+
+
+def set_matmul_block_backend(mode):
+    if mode == "auto":
+        b2_set_matmul_backend(B2_MATMUL_BACKEND_AUTO)
+    elif mode == "naive":
+        b2_set_matmul_backend(B2_MATMUL_BACKEND_NAIVE)
+    elif mode == "accelerate":
+        b2_set_matmul_backend(B2_MATMUL_BACKEND_ACCELERATE)
+    else:
+        raise ValueError("mode must be 'auto', 'naive', or 'accelerate'")
+
+
+def get_matmul_block_backend():
+    return b2_get_matmul_backend_name().decode("utf-8")
+
+
+def get_selected_matmul_block_backend():
+    return b2_get_selected_matmul_backend_name().decode("utf-8")
