Use new me_compile_nd/me_eval_nd for broadening scenarios for miniexpr

Author: FrancescAlted

CMakeLists.txt

@@ -58,7 +58,7 @@ set(MINIEXPR_BUILD_BENCH OFF CACHE BOOL "Build miniexpr benchmarks" FORCE)
 
 FetchContent_Declare(miniexpr
     GIT_REPOSITORY https://github.com/Blosc/miniexpr.git
-    GIT_TAG sleef  # latest SIMD additions
+    GIT_TAG ndim  # latest me_compile_nd()/me_eval_nd() APIs
     # In case you want to use a local copy of miniexpr for development, uncomment the line below
     # SOURCE_DIR "/Users/faltet/blosc/miniexpr"
 )

README_DEVELOPERS.md

@@ -59,15 +59,9 @@ brew install sccache ninja
 Then run:
 
 ```bash
-CMAKE_GENERATOR=Ninja \
-CMAKE_C_COMPILER=clang \
-CMAKE_CXX_COMPILER=clang++ \
 CMAKE_C_COMPILER_LAUNCHER=sccache \
-CMAKE_CXX_COMPILER_LAUNCHER=sccache \
-CMAKE_BUILD_PARALLEL_LEVEL=8 \
-SKBUILD_PARALLEL_LEVEL=8 \
 SKBUILD_BUILD_DIR=build \
-pip install -e .
+pip install -e . --no-build-isolation
 ```
 
 Using `SKBUILD_BUILD_DIR` keeps a stable build directory between runs, which

src/blosc2/blosc2_ext.pyx

@@ -573,6 +573,11 @@ cdef extern from "miniexpr.h":
     int me_compile(const char *expression, const me_variable *variables,
                    int var_count, me_dtype dtype, int *error, me_expr **out)
 
+    int me_compile_nd(const char *expression, const me_variable *variables,
+                      int var_count, me_dtype dtype, int ndims,
+                      const int64_t *shape, const int32_t *chunkshape,
+                      const int32_t *blockshape, int *error, me_expr **out)
+
     cdef enum me_compile_status:
         ME_COMPILE_SUCCESS
         ME_COMPILE_ERR_OOM
@@ -583,9 +588,26 @@ cdef extern from "miniexpr.h":
         ME_COMPILE_ERR_VAR_MIXED
         ME_COMPILE_ERR_VAR_UNSPECIFIED
         ME_COMPILE_ERR_INVALID_ARG_TYPE
+        ME_COMPILE_ERR_MIXED_TYPE_NESTED
+
+    cdef enum me_simd_ulp_mode:
+        ME_SIMD_ULP_DEFAULT
+        ME_SIMD_ULP_1
+        ME_SIMD_ULP_3_5
+
+    ctypedef struct me_eval_params:
+        c_bool disable_simd
+        me_simd_ulp_mode simd_ulp_mode
+
+    int me_eval(const me_expr *expr, const void **vars_block,
+                int n_vars, void *output_block, int chunk_nitems,
+                const me_eval_params *params) nogil
 
-    int me_eval(const me_expr *expr, const void ** vars_chunk,
-                int n_vars, void *output_chunk, int chunk_nitems) nogil
+    int me_eval_nd(const me_expr *expr, const void **vars_block,
+                   int n_vars, void *output_block, int block_nitems,
+                   int64_t nchunk, int64_t nblock, const me_eval_params *params) nogil
+
+    int me_nd_valid_nitems(const me_expr *expr, int64_t nchunk, int64_t nblock, int64_t *valid_nitems) nogil
 
     void me_print(const me_expr *n) nogil
     void me_free(me_expr *n) nogil
@@ -1860,10 +1882,8 @@ cdef int general_filler(blosc2_prefilter_params *params):
     return 0
 
 
-# Auxiliary function for just miniexpr as a prefilter
-# Only meant for (input and output) arrays that:
-# 1) Are blosc2.NDArray objects
-# 2) Do not have padding
+# Auxiliary function for miniexpr as a prefilter
+# Only meant for (input and output) arrays that are blosc2.NDArray objects.
 cdef int aux_miniexpr(me_udata *udata, int64_t nchunk, int32_t nblock,
                       c_bool is_postfilter, uint8_t *params_output, int32_t typesize) nogil:
     # Declare all C variables at the beginning
@@ -1880,9 +1900,29 @@ cdef int aux_miniexpr(me_udata *udata, int64_t nchunk, int32_t nblock,
     cdef void* src
     cdef int32_t chunk_nbytes, chunk_cbytes, block_nbytes
     cdef int start, blocknitems, expected_blocknitems
+    cdef int64_t valid_nitems
     cdef int32_t input_typesize
     cdef blosc2_context* dctx
     expected_blocknitems = -1
+    valid_nitems = 0
+
+    cdef me_expr* miniexpr_handle = udata.miniexpr_handle
+    cdef void* aux_reduc_ptr
+
+    if miniexpr_handle == NULL:
+        raise ValueError("miniexpr: handle not assigned")
+
+    # Query valid (unpadded) items for this block
+    rc = me_nd_valid_nitems(miniexpr_handle, nchunk, nblock, &valid_nitems)
+    if rc != 0:
+        raise RuntimeError(f"miniexpr: invalid block; error code: {rc}")
+    if valid_nitems <= 0:
+        # Nothing to compute for this block.
+        # For reductions, keep aux_reduc neutral values untouched.
+        if udata.aux_reduc_ptr == NULL:
+            memset(params_output, 0, udata.array.blocknitems * typesize)
+        free(input_buffers)
+        return 0
     for i in range(udata.ninputs):
         ndarr = udata.inputs[i]
         input_buffers[i] = malloc(ndarr.sc.blocksize)
@@ -1912,48 +1952,35 @@ cdef int aux_miniexpr(me_udata *udata, int64_t nchunk, int32_t nblock,
                 # In the future, perhaps one can create a specific (serial) context just for
                 # blosc2_getitem_ctx, but this is probably never going to be necessary.
                 dctx = blosc2_create_dctx(BLOSC2_DPARAMS_DEFAULTS)
-            if nchunk * ndarr.chunknitems + start + blocknitems > ndarr.nitems:
-                blocknitems = ndarr.nitems - (nchunk * ndarr.chunknitems + start)
-                if blocknitems <= 0:
-                    # Should never happen, but anyway
-                    continue
+            if valid_nitems > blocknitems:
+                raise ValueError("miniexpr: valid items exceed padded block size")
             rc = blosc2_getitem_ctx(dctx, src, chunk_cbytes, start, blocknitems,
                                     input_buffers[i], block_nbytes)
             blosc2_free_ctx(dctx)
             if rc < 0:
                 raise ValueError("miniexpr: error decompressing the chunk")
-
-    cdef me_expr* miniexpr_handle = udata.miniexpr_handle
-    cdef void* aux_reduc_ptr
     # For reduction operations, we need to track which block we're processing
     # The linear_block_index should be based on the INPUT array structure, not the output array
     # Get the first input array's chunk and block structure
     cdef b2nd_array_t* first_input = udata.inputs[0]
-    cdef int nblocks_per_chunk = (first_input.chunknitems + first_input.blocknitems - 1) // first_input.blocknitems
+    cdef int nblocks_per_chunk = 1
+    for i in range(first_input.ndim):
+        nblocks_per_chunk *= <int>udata.blocks_in_chunk[i]
     # Calculate the global linear block index: nchunk * blocks_per_chunk + nblock
     # This works because blocks never span chunks (chunks are padded to block boundaries)
     cdef int64_t linear_block_index = nchunk * nblocks_per_chunk + nblock
     cdef uintptr_t offset_bytes = typesize * linear_block_index
 
-    if miniexpr_handle == NULL:
-        raise ValueError("miniexpr: handle not assigned")
-
-    # Skip evaluation if blocknitems is invalid (can happen for padding blocks beyond data)
-    if blocknitems <= 0:
-        # Free resources
-        for i in range(udata.ninputs):
-            free(input_buffers[i])
-        free(input_buffers)
-        return 0
-
     # Call thread-safe miniexpr C API
     if udata.aux_reduc_ptr == NULL:
-        rc = me_eval(miniexpr_handle, <const void**>input_buffers, udata.ninputs,
-                     <void*>params_output, blocknitems)
+        rc = me_eval_nd(miniexpr_handle, <const void**>input_buffers, udata.ninputs,
+                        <void*>params_output, blocknitems, nchunk, nblock, NULL)
     else:
-        # Reduction operation
+        # Reduction operation: evaluate only valid items into a single output element.
+        # NOTE: miniexpr handles scalar outputs in me_eval_nd without touching tail bytes.
         aux_reduc_ptr = <void *> (<uintptr_t> udata.aux_reduc_ptr + offset_bytes)
-        rc = me_eval(miniexpr_handle, <const void**>input_buffers, udata.ninputs, aux_reduc_ptr, blocknitems)
+        rc = me_eval_nd(miniexpr_handle, <const void**>input_buffers, udata.ninputs,
+                        aux_reduc_ptr, blocknitems, nchunk, nblock, NULL)
     if rc != 0:
         raise RuntimeError(f"miniexpr: issues during evaluation; error code: {rc}")
 
@@ -2904,7 +2931,12 @@ cdef class NDArray:
         expression = expression.encode("utf-8") if isinstance(expression, str) else expression
         cdef me_dtype = me_dtype_from_numpy(self.dtype.num)
         cdef me_expr *out_expr
-        cdef int rc = me_compile(expression, variables, n, me_dtype, &error, &out_expr)
+        cdef int ndims = self.array.ndim
+        cdef int64_t* shape = &self.array.shape[0]
+        cdef int32_t* chunkshape = &self.array.chunkshape[0]
+        cdef int32_t* blockshape = &self.array.blockshape[0]
+        cdef int rc = me_compile_nd(expression, variables, n, me_dtype, ndims,
+                                    shape, chunkshape, blockshape, &error, &out_expr)
         if rc == ME_COMPILE_ERR_INVALID_ARG_TYPE:
             raise TypeError(f"miniexpr does not support operand or output dtype: {expression}")
         if rc != ME_COMPILE_SUCCESS:

src/blosc2/lazyexpr.py

@@ -1284,14 +1284,14 @@ def fast_eval(  # noqa: C901
 
     # Check whether we can use miniexpr
     if use_miniexpr:
-        # Avoid padding issues except for 1D arrays (contiguous along the only axis).
-        if len(shape) != 1 and builtins.any(s % c != 0 for s, c in zip(shape[1:], chunks[1:], strict=True)):
+        # Require aligned NDArray operands with identical chunk/block grid.
+        same_shape = all(hasattr(op, "shape") and op.shape == shape for op in operands.values())
+        same_chunks = all(hasattr(op, "chunks") and op.chunks == chunks for op in operands.values())
+        same_blocks = all(hasattr(op, "blocks") and op.blocks == blocks for op in operands.values())
+        if not (same_shape and same_chunks and same_blocks):
+            use_miniexpr = False
+        if not (all_ndarray and not any_persisted and out is None):
             use_miniexpr = False
-        for op in operands.values():
-            # Only NDArray in-memory operands
-            if not (isinstance(op, blosc2.NDArray) and op.urlpath is None and out is None):
-                use_miniexpr = False
-                break
 
     if use_miniexpr:
         cparams = kwargs.pop("cparams", blosc2.CParams())
@@ -1989,10 +1989,12 @@ def reduce_slices(  # noqa: C901
     if reduce_op in (ReduceOp.ARGMAX, ReduceOp.ARGMIN):
         use_miniexpr = False
 
-    # Only behaved partitions are supported in miniexpr reductions
+    # Check whether we can use miniexpr
     if use_miniexpr:
-        # Avoid padding issues except for 1D arrays (contiguous along the only axis).
-        if len(shape) != 1 and builtins.any(s % c != 0 for s, c in zip(shape[1:], chunks[1:], strict=True)):
+        same_shape = all(hasattr(op, "shape") and op.shape == shape for op in operands.values())
+        same_chunks = all(hasattr(op, "chunks") and op.chunks == chunks for op in operands.values())
+        same_blocks = all(hasattr(op, "blocks") and op.blocks == blocks for op in operands.values())
+        if not (same_shape and same_chunks and same_blocks):
             use_miniexpr = False
         if use_miniexpr and isinstance(expression, str):
             has_complex = any(
@@ -2001,12 +2003,6 @@ def reduce_slices(  # noqa: C901
             )
             if has_complex and any(tok in expression for tok in ("!=", "==", "<=", ">=", "<", ">")):
                 use_miniexpr = False
-        for op in operands.values():
-            # Ensure blocks fit exactly in chunks for the n-dim case, except for the first dimension
-            blocks_fit = builtins.all(c % b == 0 for c, b in zip(op.chunks[1:], op.blocks[1:], strict=True))
-            if len(op.shape) != 1 and not blocks_fit:
-                use_miniexpr = False
-                break
 
     if use_miniexpr:
         # Experiments say that not splitting is best (at least on Apple Silicon M4 Pro)

tests/ndarray/test_lazyexpr.py

@@ -270,7 +270,7 @@ def test_expression_with_constants(array_fixture):
     a1, a2, a3, a4, na1, na2, na3, na4 = array_fixture
     # Test with operands with same chunks and blocks
     expr = a1 + 2 - a3 * 3.14
-    nres = ne_evaluate("na1 + 2 - na3 * 3.14")
+    nres = na1 + 2 - na3 * 3.14
     res = expr.compute()
     if na1.dtype == np.float32:
         np.testing.assert_allclose(res[:], nres, rtol=1e-5)