Merge pull request #592 from Blosc/dsl-constructors

DSL constructors

Author: FrancescAlted

src/blosc2/dsl_kernel.py

@@ -105,7 +105,11 @@ def _remove_scalar_params_preserving_source(text: str, scalar_replacements: dict
     updated = f"{text[:pstart]}{', '.join(kept)}{text[pend:]}"
     body_start = 0
     if colon is not None:
-        body_start = _to_abs(_line_starts(updated), colon.end[0], colon.end[1])
+        # Signature shrink can move ':' to an earlier column, so recompute
+        # on the rewritten text to avoid skipping first-line body tokens.
+        _, _, updated_colon = _find_def_signature_span(updated)
+        if updated_colon is not None:
+            body_start = _to_abs(_line_starts(updated), updated_colon.end[0], updated_colon.end[1])
     return updated, body_start
 
 
@@ -162,7 +166,7 @@ def _replace_scalar_names_preserving_source(
     return out
 
 
-def _fold_numeric_cast_calls_preserving_source(text: str, body_start: int):
+def _fold_numeric_cast_calls_preserving_source(text: str, body_start: int):  # noqa: C901
     """Fold float(<number>) and int(<number>) calls into literals.
 
     miniexpr parses DSL function calls in a restricted way, and scalar specialization can
@@ -176,6 +180,20 @@ def _fold_numeric_cast_calls_preserving_source(text: str, body_start: int):
 
     line_starts = _line_starts(text)
     edits = []
+
+    def _numeric_literal_value(node):
+        if isinstance(node, ast.Constant) and isinstance(node.value, int | float | bool):
+            return node.value
+        if (
+            isinstance(node, ast.UnaryOp)
+            and isinstance(node.op, ast.UAdd | ast.USub)
+            and isinstance(node.operand, ast.Constant)
+            and isinstance(node.operand.value, int | float | bool)
+        ):
+            value = node.operand.value
+            return value if isinstance(node.op, ast.UAdd) else -value
+        return None
+
     for node in ast.walk(tree):
         if not isinstance(node, ast.Call):
             continue
@@ -185,7 +203,8 @@ def _fold_numeric_cast_calls_preserving_source(text: str, body_start: int):
             continue
 
         arg = node.args[0]
-        if not isinstance(arg, ast.Constant) or not isinstance(arg.value, int | float | bool):
+        value = _numeric_literal_value(arg)
+        if value is None:
             continue
 
         start_abs = _to_abs(line_starts, node.lineno, node.col_offset)
@@ -194,9 +213,9 @@ def _fold_numeric_cast_calls_preserving_source(text: str, body_start: int):
         end_abs = _to_abs(line_starts, node.end_lineno, node.end_col_offset)
 
         if node.func.id == "float":
-            repl = repr(float(arg.value))
+            repl = repr(float(value))
         else:
-            repl = repr(int(arg.value))
+            repl = repr(int(value))
         edits.append((start_abs, end_abs, repl))
 
     if not edits:

src/blosc2/ndarray.py

@@ -5454,6 +5454,10 @@ def arange_fill(inputs, output, offset):
         else:  # use linspace to have finer control over exclusion of endpoint for float types
             output[:] = np.linspace(start, stop, lout, endpoint=False, dtype=output.dtype)
 
+    @blosc2.dsl_kernel
+    def ramp_arange(start, step):
+        return start + _flat_idx * step  # noqa: F821  # DSL index/shape symbols resolved by miniexpr
+
     if step is None:  # not array-api compliant but for backwards compatibility
         step = 1
     if stop is None:
@@ -5478,14 +5482,19 @@ def arange_fill(inputs, output, offset):
         # We already have the dtype and shape, so return immediately
         return blosc2.zeros(shape, dtype=dtype, **kwargs)
 
-    lshape = (math.prod(shape),)
-    lazyarr = blosc2.lazyudf(arange_fill, (start, stop, step), dtype=dtype, shape=lshape)
+    # Windows and wasm32 does not support complex numbers in DSL
+    if blosc2.isdtype(dtype, "complex floating"):
+        lshape = (math.prod(shape),)
+        lazyarr = blosc2.lazyudf(arange_fill, (start, stop, step), dtype=dtype, shape=lshape)
 
-    if len(shape) == 1:
-        # C order is guaranteed, and no reshape is needed
-        return lazyarr.compute(**kwargs)
+        if len(shape) == 1:
+            # C order is guaranteed, and no reshape is needed
+            return lazyarr.compute(**kwargs)
 
-    return reshape(lazyarr, shape, c_order=c_order, **kwargs)
+        return reshape(lazyarr, shape, c_order=c_order, **kwargs)
+    else:
+        lazyarr = blosc2.lazyudf(ramp_arange, (start, step), dtype=dtype, shape=shape)
+        return lazyarr.compute(**kwargs)
 
 
 # Define a numpy linspace-like function
@@ -5550,6 +5559,10 @@ def linspace_fill(inputs, output, offset):
         else:
             output[:] = np.linspace(start_, stop_, lout, endpoint=False, dtype=output.dtype)
 
+    @blosc2.dsl_kernel
+    def ramp_linspace(start, step):
+        return float(start) + _flat_idx * float(step)  # noqa: F821  # DSL index/shape symbols resolved by miniexpr
+
     if shape is None:
         if num is None:
             raise ValueError("Either `shape` or `num` must be specified.")
@@ -5579,13 +5592,21 @@ def linspace_fill(inputs, output, offset):
         # We already have the dtype and shape, so return immediately
         return blosc2.zeros(shape, dtype=dtype, **kwargs)  # will return empty array for num == 0
 
-    inputs = (start, stop, num, endpoint)
-    lazyarr = blosc2.lazyudf(linspace_fill, inputs, dtype=dtype, shape=(num,))
-    if len(shape) == 1:
-        # C order is guaranteed, and no reshape is needed
-        return lazyarr.compute(**kwargs)
+    # Windows and wasm32 does not support complex numbers in DSL
+    if blosc2.isdtype(dtype, "complex floating"):
+        inputs = (start, stop, num, endpoint)
+        lazyarr = blosc2.lazyudf(linspace_fill, inputs, dtype=dtype, shape=(num,))
+        if len(shape) == 1:
+            # C order is guaranteed, and no reshape is needed
+            return lazyarr.compute(**kwargs)
 
-    return reshape(lazyarr, shape, c_order=c_order, **kwargs)
+        return reshape(lazyarr, shape, c_order=c_order, **kwargs)
+    else:
+        nitems = num - 1 if endpoint else num
+        step = (float(stop) - float(start)) / float(nitems) if nitems > 0 else 0.0
+        inputs = (start, step)
+        lazyarr = blosc2.lazyudf(ramp_linspace, inputs, dtype=dtype, shape=shape)
+        return lazyarr.compute(**kwargs)
 
 
 def eye(N, M=None, k=0, dtype=np.float64, **kwargs: Any) -> NDArray:

tests/ndarray/test_dsl_kernels.py

@@ -8,6 +8,9 @@
 
 import subprocess
 import sys
+import tempfile
+import textwrap
+from pathlib import Path
 
 import numpy as np
 import pytest
@@ -116,6 +119,18 @@ def kernel_scalar_start_step(start, step):
     return start + step * (_i0 * _n1 + _i1)  # noqa: F821  # DSL index/shape symbols resolved by miniexpr
 
 
+@blosc2.dsl_kernel
+def kernel_scalar_start_stop_nitems(start, stop, nitems):
+    step = (stop - start) / nitems
+    return start + _flat_idx * step  # noqa: F821  # DSL index/shape symbols resolved by miniexpr
+
+
+@blosc2.dsl_kernel
+def kernel_scalar_start_stop_nitems_float_cast(start, stop, nitems):
+    step = (float(stop) - float(start)) / float(nitems)
+    return float(start) + _flat_idx * step  # noqa: F821  # DSL index/shape symbols resolved by miniexpr
+
+
 @blosc2.dsl_kernel
 def kernel_fallback_kw_call(x, y):
     return clip(x + y, a_min=0.5, a_max=2.5)
@@ -545,6 +560,73 @@ def test_dsl_kernel_two_scalar_params_start_step_linear_ramp():
     np.testing.assert_allclose(res[...], expected, rtol=0.0, atol=0.0)
 
 
+def test_dsl_kernel_three_scalar_params_start_stop_nitems_ramp():
+    shape = (20, 25)
+    start = np.float64(1.0)
+    stop = np.float64(2.0)
+    nitems = np.int64(np.prod(shape))
+
+    expr = blosc2.lazyudf(
+        kernel_scalar_start_stop_nitems, (start, stop, nitems), dtype=np.float64, shape=shape
+    )
+    res = expr.compute()
+
+    step = (stop - start) / nitems
+    expected = (start + step * np.arange(np.prod(shape), dtype=np.float64)).reshape(shape)
+    np.testing.assert_allclose(res[...], expected, rtol=0.0, atol=0.0)
+
+
+def test_dsl_kernel_float_cast_with_negative_scalar_param():
+    shape = (10, 100)
+    start = -10
+    stop = 10
+    nitems = np.int64(np.prod(shape) - 1)
+
+    expr = blosc2.lazyudf(
+        kernel_scalar_start_stop_nitems_float_cast, (start, stop, nitems), dtype=np.float32, shape=shape
+    )
+    res = expr.compute()
+
+    expected = np.linspace(start, stop, np.prod(shape), dtype=np.float32).reshape(shape)
+    np.testing.assert_allclose(res[...], expected, rtol=1e-6, atol=1e-6)
+
+
+def test_dsl_kernel_float_cast_with_flat_idx_no_segfault_subprocess():
+    if blosc2.IS_WASM:
+        pytest.skip("subprocess is not supported on emscripten/wasm32")
+
+    code = textwrap.dedent(
+        """
+        import numpy as np
+        import blosc2
+
+        @blosc2.dsl_kernel
+        def kernel(start, stop, nitems):
+            step = (float(stop) - float(start)) / float(nitems)
+            return float(start) + _flat_idx * step  # noqa: F821
+
+        shape = (10, 100)
+        arr = blosc2.lazyudf(kernel, (-10, 10, 999), dtype=np.float32, shape=shape).compute()
+        exp = np.linspace(-10, 10, np.prod(shape), dtype=np.float32).reshape(shape)
+        np.testing.assert_allclose(arr, exp, rtol=1e-6, atol=1e-6)
+        print("ok")
+        """
+    )
+
+    # Run from a real .py file so inspect.getsource() can recover the DSL source.
+    with tempfile.TemporaryDirectory() as tmpdir:
+        script = Path(tmpdir) / "dsl_kernel_subprocess.py"
+        script.write_text(code, encoding="utf-8")
+        result = subprocess.run([sys.executable, str(script)], capture_output=True, text=True, check=False)
+
+    assert result.returncode == 0, (
+        "subprocess failed (possible segfault/regression in DSL float-cast path):\n"
+        f"stdout:\n{result.stdout}\n"
+        f"stderr:\n{result.stderr}"
+    )
+    assert "ok" in result.stdout
+
+
 def test_dsl_kernel_scalar_constant_subexpr_runtime_no_segfault(tmp_path):
     if blosc2.IS_WASM:
         pytest.skip("subprocess is not supported on emscripten/wasm32")