diff --git a/pyopencl/array.py b/pyopencl/array.py index 62e153c3e6f91423156341c4ec96d1b8443ee8e5..9466b41cbca5ca958eb6cc830ba6aa6f50b91f37 100644 --- a/pyopencl/array.py +++ b/pyopencl/array.py @@ -1400,6 +1400,7 @@ def minimum(a, b, out=None, queue=None): # }}} # {{{ reductions +_builtin_sum = sum _builtin_min = min _builtin_max = max @@ -1460,4 +1461,29 @@ subset_max.__doc__ = """.. versionadded:: 2011.1""" # }}} +# {{{ scans + +def cumsum(a, output_dtype=None, queue=None, wait_for=None, return_event=False): + # undocumented for now + + """ + .. versionadded:: 2013.1 + """ + + if output_dtype is None: + output_dtype = a.dtype + + result = a._new_like_me(output_dtype) + + from pyopencl.scan import get_cumsum_kernel + krnl = get_cumsum_kernel(a.context, a.dtype, output_dtype) + evt = krnl(a, result, queue=queue, wait_for=wait_for) + + if return_event: + return evt, result + else: + return result + +# }}} + # vim: foldmethod=marker diff --git a/pyopencl/scan.py b/pyopencl/scan.py index a2125f4bc55c93c083cc12621beeb821217f2d9c..8d99e802b3d1c9772a4fc454eb5d91b8ecadc4c0 100644 --- a/pyopencl/scan.py +++ b/pyopencl/scan.py @@ -35,7 +35,9 @@ import pyopencl as cl import pyopencl.array from pyopencl.tools import (dtype_to_ctype, bitlog2, KernelTemplateBase, _process_code_for_macro, - get_arg_list_scalar_arg_dtypes) + get_arg_list_scalar_arg_dtypes, + context_dependent_memoize, + ) import pyopencl._mymako as mako from pyopencl._cluda import CLUDA_PREAMBLE @@ -1570,4 +1572,23 @@ class ScanTemplate(KernelTemplateBase): # }}} +# {{{ 'canned' scan kernels + +@context_dependent_memoize +def get_cumsum_kernel(context, input_dtype, output_dtype): + from pyopencl.tools import VectorArg + return GenericScanKernel( + context, output_dtype, + arguments=[ + VectorArg(input_dtype, "input"), + VectorArg(output_dtype, "output"), + ], + input_expr="input[i]", + scan_expr="a+b", neutral="0", + output_statement=""" + output[i] = item; + """) + +# }}} + # vim: filetype=pyopencl:fdm=marker