diff --git a/kernel-specification.yml b/kernel-specification.yml index a54e6cf2d1..b0ac312662 100644 --- a/kernel-specification.yml +++ b/kernel-specification.yml @@ -5184,6 +5184,83 @@ kernels: acc += fromptr[i] toptr[bin] = acc automatic-tests: true + - name: awkward_reduce_sumofsquares + specializations: + - name: awkward_reduce_sumofsquares_float64_int8_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[int8_t]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + - name: awkward_reduce_sumofsquares_float64_uint8_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[uint8_t]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + - name: awkward_reduce_sumofsquares_float64_int16_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[int16_t]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + - name: awkward_reduce_sumofsquares_float64_uint16_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[uint16_t]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + - name: awkward_reduce_sumofsquares_float64_int32_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[int32_t]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + - name: awkward_reduce_sumofsquares_float64_uint32_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[uint32_t]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + - name: awkward_reduce_sumofsquares_float64_int64_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[int64_t]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + - name: awkward_reduce_sumofsquares_float64_uint64_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[uint64_t]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + - name: awkward_reduce_sumofsquares_float64_bool_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[bool]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + - name: awkward_reduce_sumofsquares_float64_float32_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[float]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + - name: awkward_reduce_sumofsquares_float64_float64_64 + args: + - {name: toptr, type: "List[double]", dir: out} + - {name: fromptr, type: "Const[List[double]]", dir: in, role: reducer-fromptr} + - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets} + - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength} + description: null + definition: | + def awkward_reduce_sumofsquares(toptr, fromptr, offsets, outlength): + for bin in range(outlength): + acc = 0 + for i in range(offsets[bin], offsets[bin + 1]): + acc += fromptr[i] * fromptr[i] + toptr[bin] = acc + automatic-tests: true - name: awkward_reduce_sum_complex specializations: - name: awkward_reduce_sum_complex64_complex64_64 diff --git a/src/awkward/_backends/cupy.py b/src/awkward/_backends/cupy.py index d52849dd68..e01e133c18 100644 --- a/src/awkward/_backends/cupy.py +++ b/src/awkward/_backends/cupy.py @@ -64,6 +64,7 @@ def _supports_cuda_compute(self, kernel_name: str) -> bool: return kernel_name in ( # core reducers "awkward_reduce_sum", + "awkward_reduce_sumofsquares", "awkward_reduce_sum_bool", "awkward_reduce_sum_bool_complex", "awkward_reduce_sum_bool_complex64_64", # alias → _bool_complex @@ -183,6 +184,7 @@ def _get_cuda_compute_impl(self, kernel_name: str): return { "awkward_sort": cuda_compute.segmented_sort, "awkward_reduce_sum": cuda_compute.awkward_reduce_sum, + "awkward_reduce_sumofsquares": cuda_compute.awkward_reduce_sumofsquares, "awkward_reduce_sum_bool": cuda_compute.awkward_reduce_sum_bool, "awkward_reduce_sum_int32_bool_64": cuda_compute.awkward_reduce_sum_int32_bool_64, "awkward_reduce_sum_int64_bool_64": cuda_compute.awkward_reduce_sum_int64_bool_64, diff --git a/src/awkward/_connect/cuda/_compute.py b/src/awkward/_connect/cuda/_compute.py index 2bb070057b..5536e12f50 100644 --- a/src/awkward/_connect/cuda/_compute.py +++ b/src/awkward/_connect/cuda/_compute.py @@ -124,6 +124,24 @@ def _cast(x): return _cast +@functools.cache +def _make_square_to_float64(in_type): + """Interned ``x -> (double)x * (double)x`` map op. + + Widens each element to ``float64`` and squares it on the fly, so a + downstream segmented_reduce accumulates ``sum(x**2)`` in double precision + with no materialised ``x*x`` buffer and no integer/float32 overflow. Cached + per input dtype so cuda.compute builds one kernel per type. + """ + + def _sq(x): + v = float(x) + return v * v + + _sq.__annotations__ = {"x": in_type, "return": np.float64} + return _sq + + def _widen_for_reduce(input_data, out_dtype): """Fuse a widening cast into a downstream segmented_reduce. @@ -391,6 +409,33 @@ def awkward_reduce_sum( ) +def awkward_reduce_sumofsquares( + result, + input_data, + offsets_data, + outlength, +): + # sum(x**2) accumulated in float64. A TransformIterator squares each element + # (widened to double) on the fly, feeding a PLUS segmented_reduce -- no x*x + # buffer, no integer/float32 overflow. `result` is float64 (SumOfSquares). + d_input = TransformIterator( + input_data, _make_square_to_float64(input_data.dtype.type) + ) + start_o, end_o = make_segment_views(offsets_data) + + h_init = np.asarray(0, dtype=result.dtype) + + segmented_reduce( + d_in=d_input, + d_out=result, + num_segments=outlength, + start_offsets_in=start_o, + end_offsets_in=end_o, + op=OpKind.PLUS, + h_init=h_init, + ) + + def awkward_reduce_sum_bool( result, input_data, diff --git a/src/awkward/_reducers.py b/src/awkward/_reducers.py index 474fe832a0..308083ce10 100644 --- a/src/awkward/_reducers.py +++ b/src/awkward/_reducers.py @@ -451,6 +451,56 @@ def apply( return ak.contents.NumpyArray(result_array, backend=array.backend) +class SumOfSquares(KernelReducer): + """Per-segment ``sum(x**2)`` accumulated directly in ``float64``. + + Reads the input in its native dtype and squares each element in ``float64`` + inside the kernel, so integer and ``float32`` inputs neither overflow nor + lose precision, and no intermediate ``x*x`` buffer is allocated. Used by + ``ak.var``/``ak.std``/``ak.moment`` in place of ``ak.sum(x * x)``. + """ + + name: Final = "sumofsquares" + preferred_dtype: Final = np.float64 + needs_position: Final = False + + # No axis=None specialization: `_do.reduce` passes a single [0, length] + # segment for axis=None, so `apply` below handles it directly (and there is + # no optimized NumPy sum-of-squares to route to, unlike plain Sum). + + def apply( + self, + array: ak.contents.NumpyArray, + offsets: ak.index.Index, + starts: ak.index.Index, + shifts: ak.index.Index | None, + outlength: ShapeItem, + ) -> ak.contents.NumpyArray: + assert isinstance(array, ak.contents.NumpyArray) + if array.dtype.kind == "c": + raise TypeError( + f"cannot compute the sum-of-squares (ak.var/ak.std) of {array.dtype!r}" + ) + if array.dtype.kind == "M": + raise ValueError(f"cannot compute the sum-of-squares of {array.dtype!r}") + result = array.backend.nplike.empty(outlength, dtype=np.float64) + assert offsets.nplike is array.backend.nplike + array.backend.maybe_kernel_error( + array.backend[ + "awkward_reduce_sumofsquares", + np.float64, + array.dtype.type, + offsets.dtype.type, + ]( + result, + array.data, + offsets.data, + outlength, + ) + ) + return ak.contents.NumpyArray(result, backend=array.backend) + + class Prod(KernelReducer): name: Final = "prod" preferred_dtype: Final = np.float64 diff --git a/src/awkward/operations/__init__.py b/src/awkward/operations/__init__.py index c2e8723686..5c8ebc43e5 100644 --- a/src/awkward/operations/__init__.py +++ b/src/awkward/operations/__init__.py @@ -86,6 +86,7 @@ from awkward.operations.ak_std import * from awkward.operations.ak_strings_astype import * from awkward.operations.ak_sum import * +from awkward.operations.ak_sumofsquares import * from awkward.operations.ak_to_arrow import * from awkward.operations.ak_to_arrow_table import * from awkward.operations.ak_to_backend import * diff --git a/src/awkward/operations/ak_var.py b/src/awkward/operations/ak_var.py index 6bfe398fe8..226acb49c5 100644 --- a/src/awkward/operations/ak_var.py +++ b/src/awkward/operations/ak_var.py @@ -25,6 +25,56 @@ np = NumpyMetadata.instance() +def _has_complex_leaf(layout) -> bool: + """True if any NumpyArray leaf is complex (dtype kind 'c').""" + found = False + + def action(node, **kwargs): + nonlocal found + if node.is_numpy and node.dtype.kind == "c": + found = True + return node + return None + + ak._do.recursively_apply(layout, action, return_array=False) + return found + + +def _promote_products_to_float64(array): + """Promote bool/integer/float32 leaves to float64 so the *weighted* products + (``x*x*weight``, ``x*weight``) neither overflow nor lose precision. + + float64 and complex leaves are returned unchanged (no copy). Promoting ``x`` + alone is enough: ``float64 * weight`` is float64 for any real weight. Used + only on the weighted branch of ``ak.var``; the unweighted branch is copy-free + via the sum-of-squares reducer. (The fully-fused, copy-free weighted reduce + would need a two-input segmented reduce -- deferred.) + """ + if not hasattr(array, "layout"): + return array + + found = False + + def action(node, **kwargs): + nonlocal found + if node.is_numpy and (node.dtype.kind in "biu" or node.dtype == np.float32): + found = True + return node + return None + + ak._do.recursively_apply(array.layout, action, return_array=False) + if found: + return ak.operations.ak_values_astype._impl( + array, + np.float64, + including_unknown=False, + highlevel=True, + behavior=None, + attrs=None, + ) + return array + + @high_level_function() def var( x, @@ -230,16 +280,35 @@ def _impl(x, weight, ddof, axis, keepdims, mask_identity, highlevel, behavior, a behavior=ctx.behavior, attrs=ctx.attrs, ) - sumwxx = ak.operations.ak_sum._impl( - x * x, - axis, - keepdims=True, - mask_identity=True, - highlevel=True, - behavior=ctx.behavior, - attrs=ctx.attrs, - ) + # sum(x**2) accumulated in float64 directly from the input: no x*x + # buffer and no integer/float32 overflow (see ak_sumofsquares). The + # float64-only reducer doesn't cover complex, so complex keeps the + # original sum(x*x) path (which yields a complex E[x**2]). + if _has_complex_leaf(x.layout): + sumwxx = ak.operations.ak_sum._impl( + x * x, + axis, + keepdims=True, + mask_identity=True, + highlevel=True, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) + else: + sumwxx = ak.operations.ak_sumofsquares._impl( + x, + axis, + keepdims=True, + mask_identity=True, + highlevel=True, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) else: + # Promote x to float64 so the weighted products don't overflow / + # lose precision. Bounded copy (float64/complex left as-is); the + # copy-free fused path needs a two-input reduce (deferred). + x = _promote_products_to_float64(x) sumw = ak.operations.ak_sum._impl( x * 0 + weight, axis, diff --git a/tests-cuda/test_4232_cuda_var_sumofsquares.py b/tests-cuda/test_4232_cuda_var_sumofsquares.py new file mode 100644 index 0000000000..02d0af56ee --- /dev/null +++ b/tests-cuda/test_4232_cuda_var_sumofsquares.py @@ -0,0 +1,46 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +from __future__ import annotations + +import numpy as np +import pytest + +import awkward as ak + +cp = pytest.importorskip("cupy") + +# ak.var/ak.std on the cuda backend accumulate sum(x**2) in float64 via a +# TransformIterator (square) feeding segmented_reduce -- no x*x buffer, no +# integer/float32 overflow. Cross-checked against the CPU backend. + + +@pytest.mark.parametrize( + "dtype", ["int8", "uint8", "int16", "int32", "int64", "float32", "float64"] +) +def test_var_matches_cpu(dtype): + segs = [[3, 1, 2, 4], [], [40, 50], [7]] + cpu = ak.values_astype(ak.Array(segs), dtype) + gpu = ak.to_backend(cpu, "cuda") + + for op in (ak.var, ak.std): + out_cpu = ak.to_list(op(cpu, axis=-1)) + out_gpu = ak.to_list(ak.to_backend(op(gpu, axis=-1), "cpu")) + for a, b in zip(out_cpu, out_gpu, strict=True): + if a is None or (isinstance(a, float) and np.isnan(a)): + assert b is None or np.isnan(b) + else: + assert b == pytest.approx(a) + + +def test_var_int32_overflow_matches_numpy(): + # 100000**2 overflows int32; the GPU float64 accumulation must still be right. + cpu = ak.values_astype(ak.Array([[100000, 200000, 300000]]), np.int32) + gpu = ak.to_backend(cpu, "cuda") + result = ak.to_list(ak.to_backend(ak.var(gpu, axis=-1), "cpu"))[0] + assert result == pytest.approx(np.var([100000.0, 200000.0, 300000.0])) + + +def test_var_axis_none_matches_cpu(): + cpu = ak.values_astype(ak.Array([[1, 2, 3], [4, 5]]), np.int32) + gpu = ak.to_backend(cpu, "cuda") + assert ak.var(gpu) == pytest.approx(ak.var(cpu)) diff --git a/tests/test_4232_var_sumofsquares_overflow.py b/tests/test_4232_var_sumofsquares_overflow.py new file mode 100644 index 0000000000..b54ce4c114 --- /dev/null +++ b/tests/test_4232_var_sumofsquares_overflow.py @@ -0,0 +1,167 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +from __future__ import annotations + +import numpy as np +import pytest + +import awkward as ak + +# ak.var/ak.std accumulate sum(x**2) in float64 directly from the input (via the +# awkward_reduce_sumofsquares kernels), so integer and float32 squares neither +# overflow nor lose precision and no x*x buffer is allocated. + + +def test_var_int32_overflow(): + # 100000**2 = 1e10 overflows int32; the old sum(x*x) path wraps. + data = np.array([100000, 200000, 300000], dtype=np.int32) + assert ak.var(ak.Array(data)) == pytest.approx(np.var(data.astype(np.float64))) + assert ak.std(ak.Array(data)) == pytest.approx(np.std(data.astype(np.float64))) + + +def test_var_int64_overflow(): + # squares overflow int64 (the accumulator), so on main the variance goes + # negative and ak.std returns nan (issue #3525). + data = np.array([3_000_000_000, 4_000_000_000, 5_000_000_000], dtype=np.int64) + assert ak.var(ak.Array(data)) == pytest.approx(np.var(data.astype(np.float64))) + assert not np.isnan(ak.std(ak.Array(data))) + + +def test_var_float32_precision(): + # float32 squares lose precision / overflow; float64 accumulation matches. + data = np.array([1e20, 2e20, 3e20], dtype=np.float32) + result = ak.var(ak.Array(data)) + assert result == pytest.approx(np.var(data.astype(np.float64)), rel=1e-6) + + +def test_var_jagged_and_axis_none_agree(): + array = ak.values_astype( + ak.Array([[100000, 200000, 300000], [], [400000, 500000]]), np.int32 + ) + per_list = ak.var(array, axis=-1) + assert per_list[0] == pytest.approx(np.var([100000.0, 200000.0, 300000.0])) + assert np.isnan(per_list[1]) + assert per_list[2] == pytest.approx(np.var([400000.0, 500000.0])) + + flat = ak.values_astype(ak.Array([100000, 200000, 300000]), np.int32) + assert ak.var(flat) == pytest.approx(np.var([100000.0, 200000.0, 300000.0])) + + +def test_var_bool(): + array = ak.Array([[True, False, True], [True, True]]) + result = ak.var(array, axis=-1) + assert result[0] == pytest.approx(np.var([1.0, 0.0, 1.0])) + assert result[1] == pytest.approx(0.0) + + +def test_var_typetracer_is_float64(): + base = ak.values_astype(ak.Array([[1, 2, 3], [4, 5]]), np.int32) + tt = ak.to_backend(base, "typetracer") + assert str(ak.var(tt, axis=-1).type) == "2 * float64" + + +def test_var_float64_unchanged(): + data = np.array([1.5, 2.5, 3.5], dtype=np.float64) + assert ak.var(ak.Array(data)) == pytest.approx(np.var(data)) + + +def _np_weighted_var(x, w): + x = x.astype(np.float64) + w = w.astype(np.float64) + avg = np.average(x, weights=w) + return np.average((x - avg) ** 2, weights=w) + + +def test_weighted_var_int32_overflow(): + # x*x*weight overflows int32; promoting x to float64 keeps the products safe. + x = np.array([100000, 200000, 300000], dtype=np.int32) + w = np.array([1, 2, 3], dtype=np.int32) + result = ak.var(ak.Array(x), weight=ak.Array(w)) + assert result == pytest.approx(_np_weighted_var(x, w)) + + +def test_weighted_var_float32_precision(): + x = np.array([1e20, 2e20, 3e20], dtype=np.float32) + w = np.array([1.0, 2.0, 3.0], dtype=np.float32) + result = ak.var(ak.Array(x), weight=ak.Array(w)) + assert result == pytest.approx(_np_weighted_var(x, w), rel=1e-5) + + +def test_weighted_var_jagged(): + x = ak.values_astype( + ak.Array([[100000, 200000, 300000], [400000, 500000]]), np.int32 + ) + w = ak.values_astype(ak.Array([[1, 2, 3], [1, 1]]), np.int32) + result = ak.var(x, weight=w, axis=-1) + assert result[0] == pytest.approx( + _np_weighted_var(np.array([100000, 200000, 300000]), np.array([1, 2, 3])) + ) + assert result[1] == pytest.approx( + _np_weighted_var(np.array([400000, 500000]), np.array([1, 1])) + ) + + +def test_weighted_std_int32_overflow(): + x = np.array([100000, 200000, 300000], dtype=np.int32) + w = np.array([1, 2, 3], dtype=np.int32) + result = ak.std(ak.Array(x), weight=ak.Array(w)) + assert result == pytest.approx(np.sqrt(_np_weighted_var(x, w))) + + +def test_var_complex_preserved(): + # Complex is not covered by the float64 sum-of-squares reducer; var keeps the + # original sum(x*x) path, yielding a complex E[x**2] - E[x]**2 (unchanged). + data = np.array([1 + 2j, 3 + 4j, 5 + 1j]) + result = ak.var(ak.Array(data)) + expected = np.mean(data**2) - np.mean(data) ** 2 + assert complex(result) == pytest.approx(complex(expected)) + + +def test_var_uint_no_overflow(): + data = np.array([100000, 200000, 300000], dtype=np.uint32) + assert ak.var(ak.Array(data)) == pytest.approx(np.var(data.astype(np.float64))) + + +def test_var_negative_values(): + data = np.array([-100000, 200000, -300000], dtype=np.int32) + assert ak.var(ak.Array(data)) == pytest.approx(np.var(data.astype(np.float64))) + + +def test_std_int64_overflow(): + data = np.array([3_000_000_000, 4_000_000_000, 5_000_000_000], dtype=np.int64) + result = ak.std(ak.Array(data)) + assert result == pytest.approx(np.std(data.astype(np.float64))) + assert not np.isnan(result) + + +def test_var_keepdims(): + array = ak.values_astype(ak.Array([[100000, 200000, 300000]]), np.int32) + out = ak.var(array, axis=-1, keepdims=True) + assert out.to_list() == [[pytest.approx(np.var([1e5, 2e5, 3e5]))]] + + +def test_var_mask_identity_empty(): + array = ak.values_astype(ak.Array([[100000, 200000], [], [300000]]), np.int32) + out = ak.var(array, axis=-1, mask_identity=True) + assert out[1] is None + assert out[0] == pytest.approx(np.var([1e5, 2e5])) + + +def test_nanvar_integer(): + data = np.array([100000, 200000, 300000], dtype=np.int32) + assert ak.nanvar(ak.Array(data)) == pytest.approx(np.var(data.astype(np.float64))) + assert ak.nanstd(ak.Array(data)) == pytest.approx(np.std(data.astype(np.float64))) + + +def test_sumofsquares_reducer_rejects_complex(): + # The float64 reducer cannot represent complex; it raises before dispatch + # (var routes complex through the legacy sum(x*x) path instead). + layout = ak.Array(np.array([1 + 2j, 3 + 4j])).layout + with pytest.raises(TypeError): + ak._do.reduce( + layout, + ak._reducers.SumOfSquares(), + axis=None, + mask=False, + keepdims=False, + )