Some small optimizations

Author: martaiborra

blosc2/blosc2_ext.pyx

@@ -509,13 +509,9 @@ ctypedef struct numba_udata:
     char* py_func
     uintptr_t inputs_id
     int output_cdtype
-    int32_t chunkshape
-    int ndim
-    int32_t *blockshape
-    int32_t *chunkshape_ndim
-    int64_t *shape
-    int64_t *ext_shape
     b2nd_array_t *array
+    int64_t chunks_in_array[B2ND_MAX_DIM]
+    int64_t blocks_in_chunk[B2ND_MAX_DIM]
 
 MAX_TYPESIZE = BLOSC_MAX_TYPESIZE
 MAX_BUFFERSIZE = BLOSC2_MAX_BUFFERSIZE
@@ -1566,21 +1562,12 @@ cdef int general_filler(blosc2_prefilter_params *params):
 
 cdef int general_numba(blosc2_prefilter_params *params):
     cdef numba_udata *udata = <numba_udata *> params.user_data
+
     cdef uint8_t nd = udata.array.ndim
-    # off caldrà sumar-li blockshape també
-    cdef int64_t offset = params.nchunk * udata.array.sc.chunksize + params.output_offset // params.output_typesize
     cdef int64_t chunk_ndim[B2ND_MAX_DIM]
-    cdef int64_t chunks_in_array[B2ND_MAX_DIM]
-    for i in range(nd):
-        chunks_in_array[i] = udata.array.extshape[i] // udata.array.chunkshape[i]
-    blosc2_unidim_to_multidim(nd, chunks_in_array, params.nchunk, chunk_ndim)
-
+    blosc2_unidim_to_multidim(nd, udata.chunks_in_array, params.nchunk, chunk_ndim)
     cdef int64_t block_ndim[B2ND_MAX_DIM]
-    cdef int64_t blocks_in_chunk[B2ND_MAX_DIM]
-    for i in range(nd):
-        blocks_in_chunk[i] = udata.array.extchunkshape[i] // udata.array.blockshape[i]
-    blosc2_unidim_to_multidim(nd, blocks_in_chunk, params.nblock, block_ndim)
-
+    blosc2_unidim_to_multidim(nd, udata.blocks_in_chunk, params.nblock, block_ndim)
     cdef int64_t start_ndim[B2ND_MAX_DIM]
     for i in range(nd):
         start_ndim[i] = chunk_ndim[i] * udata.array.chunkshape[i] + block_ndim[i] * udata.array.blockshape[i]
@@ -1589,7 +1576,7 @@ cdef int general_numba(blosc2_prefilter_params *params):
     padding = False
     blockshape = []
     for i in range(nd):
-        if start_ndim[i] + udata.blockshape[i] > udata.shape[i]:
+        if start_ndim[i] + udata.array.blockshape[i] > udata.array.shape[i]:
             padding = True
             blockshape.append(udata.array.shape[i] - start_ndim[i])
             if blockshape[i] <= 0:
@@ -1605,19 +1592,21 @@ cdef int general_numba(blosc2_prefilter_params *params):
         output = np.empty(blockshape, udata.array.dtype)
     else:
         output = np.PyArray_SimpleNewFromData(nd, dims, udata.output_cdtype, <void*>params.output)
+
     inputs_tuple = _ctypes.PyObj_FromPtr(udata.inputs_id)
     inputs = []
+    # Get slice of each operand
     if nd == 1:
-        # Enviar-ho a fer la mà quan nd = 1 ? o ho puc suportar sense problemes??
+        # When ndim == 1, schunks as operands are supported
         for obj, dtype in inputs_tuple:
             if isinstance(obj, blosc2.SChunk):
-                out = np.empty(udata.array.blockshape[0], dtype=dtype)
-                obj.get_slice(start=offset, stop=offset + dims[0], out=out)
+                out = np.empty(blockshape[0], dtype=dtype)
+                obj.get_slice(start=start_ndim[0], stop=start_ndim[0] + blockshape[0], out=out)
                 inputs.append(out)
             elif isinstance(obj, np.ndarray):
-                inputs.append(obj[offset : offset + dims[0]])
+                inputs.append(obj[start_ndim[0] : start_ndim[0] + blockshape[0]])
             elif isinstance(obj, (int, float, bool, complex)):
-                inputs.append(np.full(udata.array.blockshape[0], obj, dtype=dtype))
+                inputs.append(np.full(blockshape[0], obj, dtype=dtype))
             else:
                 raise ValueError("Unsupported operand")
     else:
@@ -1628,9 +1617,7 @@ cdef int general_numba(blosc2_prefilter_params *params):
             # print("slice dim ", i, " = inici = ", start_ndim[i], " final = ",  start_ndim[i] + blockshape[i])
         slices = tuple(l)
         for obj, dtype in inputs_tuple:
-            if isinstance(obj, blosc2.SChunk):
-                raise ValueError("Cannot mix unidim operands with multidim") # o sí?
-            elif isinstance(obj, np.ndarray):
+            if isinstance(obj, blosc2.NDArray):
                 inputs.append(obj[slices])
             elif isinstance(obj, np.ndarray):
                 inputs.append(obj[slices])
@@ -1639,13 +1626,13 @@ cdef int general_numba(blosc2_prefilter_params *params):
             else:
                 raise ValueError("Unsupported operand")
 
+    # Call numba function
     func_id = udata.py_func.decode("utf-8")
-    blosc2.prefilter_funcs[func_id](tuple(inputs), output, offset)
+    blosc2.prefilter_funcs[func_id](tuple(inputs), output, np.array(start_ndim))
 
     cdef int64_t start[B2ND_MAX_DIM]
     cdef int64_t slice_shape[B2ND_MAX_DIM]
     cdef int64_t blockshape_int64[B2ND_MAX_DIM]
-
     cdef Py_buffer *buf
     if padding:
         for i in range(nd):
@@ -2303,13 +2290,11 @@ cdef class NDArray:
         strcpy(pref_udata.py_func, func_id)
         pref_udata.inputs_id = inputs_id
         pref_udata.output_cdtype = np.dtype(self.dtype).num
-        pref_udata.chunkshape = self.array.sc.chunksize // self.array.sc.typesize
-        pref_udata.ndim = self.ndim
-        pref_udata.blockshape = self.array.blockshape
-        pref_udata.shape = self.array.shape
-        pref_udata.chunkshape_ndim = self.array.chunkshape
-        pref_udata.ext_shape = self.array.extshape
         pref_udata.array = self.array
+        # Save these in numba_udata to avoid computing them for each block
+        for i in range(self.array.ndim):
+            pref_udata.chunks_in_array[i] = pref_udata.array.extshape[i] // pref_udata.array.chunkshape[i]
+            pref_udata.blocks_in_chunk[i] = pref_udata.array.extchunkshape[i] // pref_udata.array.blockshape[i]
 
         preparams.user_data = pref_udata
         cparams.preparams = preparams