math funcs should go without np. prefix in dsl kernels

Author: FrancescAlted

bench/b2nd/jit-dsl.py

@@ -18,6 +18,12 @@
 import blosc2
 import numpy as np
 
+where = np.where
+sin = np.sin
+cos = np.cos
+exp = np.exp
+log = np.log
+
 
 @blosc2.dsl_kernel
 def k_dsl(x, y):
@@ -27,7 +33,7 @@ def k_dsl(x, y):
         if i == 0:
             acc = acc + y
         else:
-            acc = np.where(acc < y, acc + i, acc - i)
+            acc = where(acc < y, acc + i, acc - i)
         i = i + 1
     return acc
 
@@ -37,15 +43,15 @@ def k_heavy_dsl(x, y, niter):
     acc = x
     i = 0
     while i < niter:
-        t = np.sin(acc * 1.001 + y * 0.123)
-        u = np.cos(acc * 0.777 - y * 0.211)
-        v = np.exp(t * 0.25) - np.log(np.abs(u) + 1.0)
-        p = np.sin(v * 0.731 + acc * 0.071)
-        q = np.cos(v * 0.379 - y * 0.053)
-        r = np.exp((p - q) * 0.17) - np.log(np.abs(p + q) + 1.0)
-        w = np.sin((r + v) * 0.11) + np.cos((r - v) * 0.07)
+        t = sin(acc * 1.001 + y * 0.123)
+        u = cos(acc * 0.777 - y * 0.211)
+        v = exp(t * 0.25) - log(abs(u) + 1.0)
+        p = sin(v * 0.731 + acc * 0.071)
+        q = cos(v * 0.379 - y * 0.053)
+        r = exp((p - q) * 0.17) - log(abs(p + q) + 1.0)
+        w = sin((r + v) * 0.11) + cos((r - v) * 0.07)
         delta = v + r + w
-        acc = np.where((acc < y), (acc + delta), (acc - delta))
+        acc = where((acc < y), (acc + delta), (acc - delta))
         i = i + 1
     return acc
 

bench/ndarray/dsl-kernel-bench.py

@@ -14,16 +14,17 @@
 import importlib
 
 lazyexpr_mod = importlib.import_module("blosc2.lazyexpr")
+where = np.where
 
 
 @blosc2.dsl_kernel
 def kernel_loop1(x, y):
     acc = 0.0
     for i in range(1):
         if i % 2 == 0:
-            tmp = np.where(x < y, y + i, x - i)
+            tmp = where(x < y, y + i, x - i)
         else:
-            tmp = np.where(x > y, x + i, y - i)
+            tmp = where(x > y, x + i, y - i)
         acc = acc + tmp * (i + 1)
     return acc
 
@@ -33,9 +34,9 @@ def kernel_loop2(x, y):
     acc = 0.0
     for i in range(2):
         if i % 2 == 0:
-            tmp = np.where(x < y, y + i, x - i)
+            tmp = where(x < y, y + i, x - i)
         else:
-            tmp = np.where(x > y, x + i, y - i)
+            tmp = where(x > y, x + i, y - i)
         acc = acc + tmp * (i + 1)
     return acc
 
@@ -45,9 +46,9 @@ def kernel_loop4(x, y):
     acc = 0.0
     for i in range(4):
         if i % 2 == 0:
-            tmp = np.where(x < y, y + i, x - i)
+            tmp = where(x < y, y + i, x - i)
         else:
-            tmp = np.where(x > y, x + i, y - i)
+            tmp = where(x > y, x + i, y - i)
         acc = acc + tmp * (i + 1)
     return acc
 
@@ -57,9 +58,9 @@ def kernel_loop4_heavy(x, y):
     acc = 0.0
     for i in range(4):
         if i % 2 == 0:
-            tmp = np.where(x < y, y + i, x - i)
+            tmp = where(x < y, y + i, x - i)
         else:
-            tmp = np.where(x > y, x + i, y - i)
+            tmp = where(x > y, x + i, y - i)
         acc = acc + tmp * (i + 1) + (tmp * tmp) * 0.05
     return acc
 
@@ -70,9 +71,9 @@ def kernel_nested2(x, y):
     for i in range(2):
         for j in range(2):
             if (i + j) % 2 == 0:
-                tmp = np.where(x < y, y + i + j, x - i - j)
+                tmp = where(x < y, y + i + j, x - i - j)
             else:
-                tmp = np.where(x > y, x + i + j, y - i - j)
+                tmp = where(x > y, x + i + j, y - i - j)
             acc = acc + tmp * (i + j + 1)
     return acc
 
@@ -148,7 +149,7 @@ def bench_case(name, kernel, expr, a, b, dtype, gb):
         res_dsl = lazy_dsl.compute()
         dsl_time, _ = time_it(lambda: lazy_dsl.compute())
 
-    np.testing.assert_allclose(res_dsl[...], res_base[...], rtol=1e-5, atol=1e-6)
+    np.testing.assert_allclose(res_dsl[...], res_base[...], rtol=1e-5, atol=2e-6)
 
     return {
         "case": name,

tests/ndarray/test_dsl_kernels.py

@@ -12,6 +12,9 @@
 from blosc2.dsl_kernel import DSLSyntaxError
 from blosc2.lazyexpr import _apply_jit_backend_pragma
 
+where = np.where
+clip = np.clip
+
 
 def _make_arrays(shape=(8, 8), chunks=(4, 4), blocks=(2, 2)):
     a = np.linspace(0, 1, num=np.prod(shape), dtype=np.float32).reshape(shape)
@@ -34,9 +37,9 @@ def kernel_loop(x, y):
     acc = 0.0
     for i in range(2):
         if i % 2 == 0:
-            tmp = np.where(x < y, y + i, x - i)
+            tmp = where(x < y, y + i, x - i)
         else:
-            tmp = np.where(x > y, x + i, y - i)
+            tmp = where(x > y, x + i, y - i)
         acc = acc + tmp * (i + 1)
     return acc
 
@@ -68,7 +71,7 @@ def kernel_control_flow_full(x, y):
         if i == 1:
             acc = acc - y
         else:
-            acc = np.where(acc < y, acc + i, acc - i)
+            acc = where(acc < y, acc + i, acc - i)
             if i == 3:
                 break
     return acc
@@ -79,7 +82,7 @@ def kernel_while_full(x, y):
     acc = x
     i = 0
     while i < 3:
-        acc = np.where(acc < y, acc + 1, acc - 1)
+        acc = where(acc < y, acc + 1, acc - 1)
         i = i + 1
     return acc
 
@@ -89,7 +92,7 @@ def kernel_loop_param(x, y, niter):
     acc = x
     # loop count comes from scalar niter
     for _i in range(niter):
-        acc = np.where(acc < y, acc + 1, acc - 1)
+        acc = where(acc < y, acc + 1, acc - 1)
     return acc
 
 
@@ -101,7 +104,7 @@ def kernel_scalar_float_cast(x, niter):
 
 @blosc2.dsl_kernel
 def kernel_fallback_kw_call(x, y):
-    return np.clip(x + y, a_min=0.5, a_max=2.5)
+    return clip(x + y, a_min=0.5, a_max=2.5)
 
 
 @blosc2.dsl_kernel
@@ -135,6 +138,16 @@ def kernel_index_ramp_float_cast(x):
     return float(_i0) * _n1 + _i1  # noqa: F821  # DSL index/shape symbols resolved by miniexpr
 
 
+@blosc2.dsl_kernel
+def kernel_index_ramp_int_cast(x):
+    return int(_i0 * _n1 + _i1)  # noqa: F821  # DSL index/shape symbols resolved by miniexpr
+
+
+@blosc2.dsl_kernel
+def kernel_bool_cast_numeric(x):
+    return bool(x)
+
+
 @blosc2.dsl_kernel
 def kernel_index_ramp_no_inputs():
     return _i0 * _n1 + _i1  # noqa: F821  # DSL index/shape symbols resolved by miniexpr
@@ -255,14 +268,60 @@ def test_dsl_kernel_index_symbols_float_cast_matches_expected_ramp():
     np.testing.assert_allclose(res, expected, rtol=0.0, atol=0.0)
 
 
+def test_dsl_kernel_index_symbols_float_cast_uses_miniexpr_fast_path(monkeypatch):
+    if blosc2.IS_WASM:
+        pytest.skip("miniexpr fast path is not available on WASM")
+
+    original_set_pref_expr = blosc2.NDArray._set_pref_expr
+    captured = {"calls": 0, "expr": None}
+
+    def wrapped_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, jit=None):
+        captured["calls"] += 1
+        captured["expr"] = expression.decode("utf-8") if isinstance(expression, bytes) else expression
+        return original_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc, jit=jit)
+
+    monkeypatch.setattr(blosc2.NDArray, "_set_pref_expr", wrapped_set_pref_expr)
+
+    shape = (16, 9)
+    x2 = blosc2.zeros(shape, dtype=np.float32)
+    expr = blosc2.lazyudf(kernel_index_ramp_float_cast, (x2,), dtype=np.float32)
+    res = expr[:]
+    expected = np.arange(np.prod(shape), dtype=np.float32).reshape(shape)
+
+    np.testing.assert_allclose(res, expected, rtol=0.0, atol=0.0)
+    assert captured["calls"] >= 1
+    assert "def kernel_index_ramp_float_cast(x):" in captured["expr"]
+    assert "float(_i0)" in captured["expr"]
+    assert "_n1" in captured["expr"]
+    assert "_i1" in captured["expr"]
+
+
+def test_dsl_kernel_index_symbols_int_cast_matches_expected_ramp():
+    shape = (32, 5)
+    x2 = blosc2.zeros(shape, dtype=np.float32)
+    expr = blosc2.lazyudf(kernel_index_ramp_int_cast, (x2,), dtype=np.int64)
+    res = expr[:]
+    expected = np.arange(np.prod(shape), dtype=np.int64).reshape(shape)
+    np.testing.assert_equal(res, expected)
+
+
+def test_dsl_kernel_bool_cast_numeric_matches_expected():
+    x = np.array([[0.0, 1.0, -2.0], [3.5, 0.0, -0.1]], dtype=np.float32)
+    x2 = blosc2.asarray(x, chunks=(2, 3), blocks=(1, 2))
+    expr = blosc2.lazyudf(kernel_bool_cast_numeric, (x2,), dtype=np.bool_)
+    res = expr[:]
+    expected = x != 0.0
+    np.testing.assert_equal(res, expected)
+
+
 def test_dsl_kernel_full_control_flow_kept_as_dsl_function():
     assert kernel_control_flow_full.dsl_source is not None
     assert "def kernel_control_flow_full(x, y):" in kernel_control_flow_full.dsl_source
     assert "for i in range(4):" in kernel_control_flow_full.dsl_source
     assert "if i == 1:" in kernel_control_flow_full.dsl_source
     assert "continue" in kernel_control_flow_full.dsl_source
     assert "break" in kernel_control_flow_full.dsl_source
-    assert "np.where(" in kernel_control_flow_full.dsl_source
+    assert "where(" in kernel_control_flow_full.dsl_source
 
     a, b, a2, b2 = _make_arrays()
     expr = blosc2.lazyudf(