New blosc2.validate_dsl_jit(); documented with examples

Author: FrancescAlted

src/blosc2/__init__.py

@@ -573,7 +573,7 @@ def _raise(exc):
 
 from .c2array import c2context, C2Array, URLPath
 
-from .dsl_kernel import DSLSyntaxError, DSLKernel, dsl_kernel, validate_dsl
+from .dsl_kernel import DSLSyntaxError, DSLKernel, dsl_kernel, validate_dsl, validate_dsl_jit
 from .lazyexpr import (
     LazyExpr,
     lazyudf,
@@ -976,6 +976,7 @@ def _raise(exc):
     "lazyexpr",
     "dsl_kernel",
     "validate_dsl",
+    "validate_dsl_jit",
     "lazyudf",
     "lazywhere",
     "less",

src/blosc2/blosc2_ext.pyx

@@ -735,6 +735,8 @@ cdef extern from "miniexpr.h":
     void me_print(const me_expr *n) nogil
     void me_free(me_expr *n) nogil
 
+    bint me_expr_has_jit_kernel(const me_expr *expr) nogil
+
     ctypedef int (*me_wasm_jit_instantiate_helper)(
         const unsigned char *wasm_bytes,
         int wasm_len,
@@ -4071,6 +4073,74 @@ cdef class NDArray:
         # alive for the duration of the prefilter-driven update_data loop.
         return cb
 
+    def _dsl_jit_status(self, expression, inputs):
+        """Compile *expression* with JIT against this array's grid and report whether
+        miniexpr produced a runtime JIT kernel (vs an interpreter fallback).
+
+        Compiles and immediately frees the program; it does not set a prefilter or
+        run the kernel.  Returns a dict with ``compiled`` (bool), ``jit`` (bool) and
+        ``status`` (miniexpr compile-status name).  ``inputs`` is a name -> NDArray
+        mapping; all operands must share this array's shape/chunks/blocks.
+        """
+        cdef Py_ssize_t n = len(inputs)
+        cdef me_variable* variables = NULL
+        if n > 0:
+            variables = <me_variable *> malloc(sizeof(me_variable) * n)
+            if variables == NULL:
+                raise MemoryError()
+        cdef me_variable *var
+        cdef np.dtype out_np_dtype = np.dtype(self.dtype)
+        cdef me_dtype me_output_dtype = _me_dtype_from_numpy_dtype(out_np_dtype)
+        if <int>me_output_dtype < 0:
+            free(variables)
+            raise TypeError(f"miniexpr does not support output dtype: {out_np_dtype}")
+
+        cdef np.dtype operand_dtype
+        cdef bytes var_name
+        for i, (k, v) in enumerate(inputs.items()):
+            var = &variables[i]
+            var_name = k.encode("utf-8") if isinstance(k, str) else k
+            var.name = <char *> malloc(strlen(var_name) + 1)
+            strcpy(var.name, var_name)
+            operand_dtype = np.dtype(v.dtype)
+            var.dtype = _me_dtype_from_numpy_dtype(operand_dtype)
+            if <int>var.dtype < 0:
+                for j in range(i + 1):
+                    free(variables[j].name)
+                free(variables)
+                raise TypeError(f"miniexpr does not support operand dtype '{operand_dtype}' for input '{k}'")
+            var.address = NULL
+            var.type = 0
+            var.context = NULL
+            var.itemsize = v.dtype.itemsize if v.dtype.num == 19 else 0
+
+        cdef bytes expression_bytes = (
+            (<str>expression).encode("utf-8") if isinstance(expression, str) else expression
+        )
+        cdef int error = 0
+        cdef me_expr *out_expr = NULL
+        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_jit(expression_bytes, variables, n, me_output_dtype, ndims,
+                                        shape, chunkshape, blockshape, ME_JIT_ON,
+                                        &error, &out_expr)
+        for i in range(n):
+            free(variables[i].name)
+        free(variables)
+
+        cdef bint jit = False
+        if rc == ME_COMPILE_SUCCESS and out_expr != NULL:
+            jit = me_expr_has_jit_kernel(out_expr)
+        if out_expr != NULL:
+            me_free(out_expr)
+        return {
+            "compiled": rc == ME_COMPILE_SUCCESS,
+            "jit": bool(jit),
+            "status": _me_compile_status_name(rc),
+        }
+
     def _set_pref_matmul(self, inputs, fp_accuracy):
         # Set prefilter for miniexpr
         cdef blosc2_cparams* cparams = self.array.sc.storage.cparams

src/blosc2/dsl_kernel.py

@@ -635,6 +635,31 @@ def validate_dsl(func):
         - ``dsl_source`` (str | None): extracted DSL source when valid
         - ``input_names`` (list[str] | None): input signature names when valid
         - ``error`` (str | None): user-facing error message when invalid
+
+    Examples
+    --------
+    >>> import blosc2
+    >>> @blosc2.dsl_kernel
+    ... def k(a, b):
+    ...     return a * a + b * b
+    >>> info = blosc2.validate_dsl(k)
+    >>> info["valid"]
+    True
+    >>> info["input_names"]
+    ['a', 'b']
+
+    An unsupported construct is reported instead of raising:
+
+    >>> @blosc2.dsl_kernel
+    ... def bad(a):
+    ...     return 1 if a > 0 else 0
+    >>> blosc2.validate_dsl(bad)["valid"]
+    False
+
+    See Also
+    --------
+    validate_dsl_jit : Additionally probe whether the kernel JIT-compiles (vs falls
+        back to the interpreter) for given operand/output dtypes.
     """
 
     kernel = func if isinstance(func, DSLKernel) else DSLKernel(func)
@@ -647,6 +672,115 @@ def validate_dsl(func):
     }
 
 
+def validate_dsl_jit(func, operands, out_dtype, *, shape=(64,), chunks=None, blocks=None):
+    """Report whether a DSL kernel JIT-compiles (vs. interpreter fallback).
+
+    Compiles a tiny probe of the kernel for the given operand/output dtypes and
+    queries miniexpr for whether it produced a runtime JIT kernel.  The kernel is
+    *not* run on real data, but compilation is dtype-specialized, so the operand and
+    output dtypes must be supplied.
+
+    Parameters
+    ----------
+    func
+        A Python callable or :class:`DSLKernel`.
+    operands
+        One spec per kernel input.  A NumPy dtype (e.g. ``np.float64``, ``"f8"``)
+        marks an *array* operand; a Python ``int``/``float``/``bool`` marks a *scalar*
+        parameter inlined as a constant (as happens at run time).  Either a dict
+        mapping input name -> spec, or a sequence matched positionally to the inputs.
+    out_dtype
+        Output dtype to specialize the codegen for.
+    shape
+        Probe shape; ``len(shape)`` sets the kernel dimensionality (kernels using
+        ``_i1`` etc. need a matching ndim).  Defaults to ``(64,)``.
+
+    Returns
+    -------
+    dict
+        - ``valid`` (bool): DSL syntax is valid
+        - ``jit`` (bool): a runtime JIT kernel was produced (vs interpreter fallback)
+        - ``compiled`` (bool): miniexpr accepted the kernel
+        - ``status`` (str | None): miniexpr compile-status name
+        - ``error`` (str | None): syntax error message when ``valid`` is False
+
+    Examples
+    --------
+    >>> import blosc2
+    >>> import numpy as np
+    >>> @blosc2.dsl_kernel
+    ... def k(a, b):
+    ...     return a * a + b * b
+    >>> status = blosc2.validate_dsl_jit(k, [np.float64, np.float64], np.float64)
+    >>> status["jit"]
+    True
+    >>> status["status"]
+    'ME_COMPILE_SUCCESS'
+
+    Pass array operands as dtypes and scalar parameters as values (the scalar is
+    inlined as a constant, as at run time):
+
+    >>> @blosc2.dsl_kernel
+    ... def scaled(a, n):
+    ...     return a * float(n)
+    >>> blosc2.validate_dsl_jit(scaled, {"a": np.float64, "n": 3}, np.float64)["jit"]
+    True
+
+    See Also
+    --------
+    validate_dsl : Static DSL-syntax check only (no compilation; no dtypes needed).
+    """
+    import numpy as np
+
+    import blosc2
+
+    kernel = func if isinstance(func, DSLKernel) else DSLKernel(func)
+    if kernel.dsl_error is not None:
+        return {
+            "valid": False,
+            "jit": False,
+            "compiled": False,
+            "status": None,
+            "error": str(kernel.dsl_error),
+        }
+
+    names = list(kernel.input_names or [])
+    if isinstance(operands, dict):
+        specs = dict(operands)
+    else:
+        operands = list(operands)
+        if len(operands) != len(names):
+            raise ValueError(f"expected {len(names)} operand specs for inputs {names}, got {len(operands)}")
+        specs = dict(zip(names, operands, strict=True))
+
+    scalars: dict = {}
+    arrays: dict = {}
+    for name in names:
+        if name not in specs:
+            raise ValueError(f"missing operand spec for input '{name}'")
+        spec = specs[name]
+        if isinstance(spec, bool | int | float):
+            scalars[name] = spec
+        else:
+            arrays[name] = np.dtype(spec)
+
+    source = kernel.dsl_source
+    if scalars:
+        source, _ = specialize_miniexpr_inputs(source, scalars)
+    if not arrays:
+        raise ValueError(
+            "validate_dsl_jit needs at least one array operand (pass a dtype spec); "
+            "all-scalar probes are not supported"
+        )
+
+    out = blosc2.empty(shape, dtype=np.dtype(out_dtype), chunks=chunks, blocks=blocks)
+    probe_inputs = {name: np.empty(1, dtype=dt) for name, dt in arrays.items()}
+    status = out._dsl_jit_status(source, probe_inputs)
+    status["valid"] = True
+    status["error"] = None
+    return status
+
+
 def kernel_from_source(source: str, name: str | None = None) -> DSLKernel:
     """Reconstruct a :class:`DSLKernel` from its stored source text.
 

tests/ndarray/test_dsl_kernels.py

@@ -1069,3 +1069,42 @@ def wrapped(self, expression, inputs, fp_accuracy, aux_reduc=None, jit=None):
     assert any(line.strip().startswith("out =") for line in captured["expr"].splitlines())
     expected = blosc2.lazyudf(k, (a2, a2, 32), dtype=np.float64, jit=False)[:]
     np.testing.assert_allclose(res, expected)
+
+
+def test_validate_dsl_jit_reports_compile_and_fallback(monkeypatch):
+    @blosc2.dsl_kernel
+    def simple(a, b):
+        return a * a + b * b
+
+    st = blosc2.validate_dsl_jit(simple, [np.float64, np.float64], np.float64)
+    assert st["valid"]
+    assert st["compiled"]
+    assert st["jit"]
+    assert st["status"] == "ME_COMPILE_SUCCESS"
+
+    # A scalar param is inlined (passed as a value); a variable named 'out' (the
+    # former silent-fallback collision) now JIT-compiles through to miniexpr.  Built
+    # from source so the linter does not rewrite the deliberate 'out' variable.
+    with_out = kernel_from_source(
+        "def withk(a, b, niter):\n    out = a + b * float(niter)\n    return out\n", "withk"
+    )
+    st = blosc2.validate_dsl_jit(with_out, {"a": np.float64, "b": np.float64, "niter": 3}, np.float64)
+    assert st["jit"]
+
+    # Invalid syntax -> not valid, nothing compiled.
+    invalid = kernel_from_source("def k(a, b):\n    a = a - b\n    return a\n", "k")
+    st = blosc2.validate_dsl_jit(invalid, [np.float64, np.float64], np.float64)
+    assert not st["valid"]
+    assert not st["jit"]
+    assert "input parameter 'a'" in st["error"]
+
+    # JIT disabled in the environment -> compiles but falls back to the interpreter.
+    monkeypatch.setenv("ME_DSL_JIT", "0")
+
+    @blosc2.dsl_kernel
+    def other(a, b):
+        return a - b * 0.5
+
+    st = blosc2.validate_dsl_jit(other, [np.float64, np.float64], np.float64)
+    assert st["compiled"]
+    assert not st["jit"]