Multithreaded matmul supported for ND

Author: lshaw8317

bench/ndarray/stringops_bench.py

@@ -12,7 +12,7 @@
 import time
 import numpy as np
 import blosc2
-from blosc2.lazyexpr import _toggle_miniexpr
+from blosc2.utils import _toggle_miniexpr
 
 # nparr = np.random.randint(low=0, high=128, size=(N, 10), dtype=np.uint32)
 # nparr = nparr.view('S40').astype('U10')

src/blosc2/blosc2_ext.pyx

@@ -2188,8 +2188,10 @@ cdef int aux_matmul(mm_udata *udata, int64_t nchunk, int32_t nblock, void *param
     cdef int ndim = out_arr.ndim
     cdef int nchunk_ = nchunk
     cdef int coord, batch, batch_, batches = 1
+
+    # batches = sum(strides[i]*elcoords[i])
     for i in range(ndim - 2):
-        batches *= out_arr.shape[i]
+        batches *= out_arr.blockshape[i]
 
     # nchunk = sum(strides[i]*chunkcoords[i])
     for i in range(ndim - 2):
@@ -2203,8 +2205,8 @@ cdef int aux_matmul(mm_udata *udata, int64_t nchunk, int32_t nblock, void *param
 
     i = nchunk_ // ncols # ncols * i + j
     j = nchunk_ % ncols
-    nchunkA = chunk_startA = nchunkA + i * ncols
-    nchunkB = chunk_startB = nchunkB + j
+    chunk_startA = nchunkA + i * ncols
+    chunk_startB = nchunkB + j
 
     # nblock = sum(strides[i]*blockcoords[i])
     cdef int nblock_ = nblock
@@ -2219,14 +2221,14 @@ cdef int aux_matmul(mm_udata *udata, int64_t nchunk, int32_t nblock, void *param
 
     block_i = nblock_ // block_ncols
     block_j = nblock_ % block_ncols
-    block_startA = nblockA = nblockA + i * block_ncols
-    block_startB = nblockB = nblockB + j
+    block_startA = nblockA + block_i * block_ncols
+    block_startB = nblockB + block_j
 
-    # batches = sum(strides[i]*elcoords[i])
     dctx = blosc2_create_dctx(BLOSC2_DPARAMS_DEFAULTS)
 
     first_run = True
-
+    nchunkA = chunk_startA
+    nchunkB = chunk_startB
     while True: # chunk loop
         for i in range(2):
             chunk_idx = nchunkA if i == 0 else nchunkB
@@ -2268,6 +2270,9 @@ cdef int aux_matmul(mm_udata *udata, int64_t nchunk, int32_t nblock, void *param
             if rc < 0:
                 raise ValueError("matmul: error decompressing the B chunk")
             batch = 0
+            offsetA = 0
+            offsetB = 0
+            offset = 0
             while batch < batches:
                 batch_ = batch
                 for i in range(ndim - 2):
@@ -3268,9 +3273,10 @@ cdef class NDArray:
             i = self.array.ndim - idx
             udata.chunks_strides[0][i] = udata.chunks_strides[0][i + 1] * udata.array.extshape[i + 1] // udata.array.chunkshape[i + 1]
             udata.blocks_strides[0][i] = udata.blocks_strides[0][i + 1] * udata.array.extchunkshape[i + 1] // udata.array.blockshape[i + 1]
+            udata.el_strides[0][i] = udata.el_strides[0][i + 1] * udata.array.blockshape[i + 1]
 
-        for j in range(2):
-            inp = inputs_[j]
+        for j in range(1, 3):
+            inp = inputs_[j - 1]
             cstrides = bstrides = estrides = 1
             for idx in range(2, self.array.ndim + 1):
                 i = inp.ndim - idx

src/blosc2/lazyexpr.py

@@ -68,6 +68,7 @@
     process_key,
     reducers,
     safe_numpy_globals,
+    try_miniexpr,
 )
 
 if not blosc2.IS_WASM:
@@ -76,14 +77,6 @@
 global safe_blosc2_globals
 safe_blosc2_globals = {}
 
-# Set this to False if miniexpr should not be tried out
-try_miniexpr = not blosc2.IS_WASM or getattr(blosc2, "_WASM_MINIEXPR_ENABLED", False)
-
-
-def _toggle_miniexpr(FLAG):
-    global try_miniexpr
-    try_miniexpr = FLAG
-
 
 def ne_evaluate(expression, local_dict=None, **kwargs):
     """Safely evaluate expressions using numexpr when possible, falling back to numpy."""

src/blosc2/linalg.py

@@ -17,7 +17,7 @@
 
 import blosc2
 
-from .utils import get_intersecting_chunks, nptranspose, npvecdot, slice_to_chunktuple
+from .utils import get_intersecting_chunks, nptranspose, npvecdot, slice_to_chunktuple, try_miniexpr
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -113,11 +113,14 @@ def matmul(x1: blosc2.Array, x2: blosc2.NDArray, **kwargs: Any) -> blosc2.NDArra
     result = blosc2.zeros(result_shape, dtype=blosc2.result_type(x1, x2), **kwargs)
 
     # multithreaded matmul
-    # TODO: handle a) type promotion, b) padding, c) (improved) >2D
+    # TODO: handle a) type promotion, b) padding (explicitly), c) (improved) >2D
     ops = (x1, x2, result)
     blocks = result.blocks
     all_ndarray = all(isinstance(value, blosc2.NDArray) and value.shape != () for value in ops)
-    use_miniexpr = True
+    global try_miniexpr
+
+    # Use a local copy so we don't modify the global
+    use_miniexpr = try_miniexpr
     if all_ndarray:
         if any(op.dtype != ops[0].dtype for op in ops):  # TODO: Remove this condition
             use_miniexpr = False
@@ -165,6 +168,7 @@ def matmul(x1: blosc2.Array, x2: blosc2.NDArray, **kwargs: Any) -> blosc2.NDArra
             if prefilter_set:
                 result.schunk.remove_prefilter("miniexpr")
     else:  # couldn't do multithreading
+        print("multithreading failed :( ")
         if 0 not in result.shape + x1.shape + x2.shape:  # if any array is empty, return array of 0s
             p, q = result.chunks[-2:]
             r = x2.chunks[-1]

src/blosc2/utils.py

@@ -19,6 +19,15 @@
 
 import blosc2
 
+# Set this to False if miniexpr should not be tried out
+try_miniexpr = not blosc2.IS_WASM or getattr(blosc2, "_WASM_MINIEXPR_ENABLED", False)
+
+
+def _toggle_miniexpr(FLAG):
+    global try_miniexpr
+    try_miniexpr = FLAG
+
+
 # NumPy version and a convenient boolean flag
 NUMPY_GE_2_0 = np.__version__ >= "2.0"
 # handle different numpy versions