Added higher-order arithmetic functions: sum, prod, count, maximum, m…

…inium, cumsum, cumprod
JuliaGPU · Dec 30, 2024 · 8817cc7 · 8817cc7
1 parent 4fa4e44
commit 8817cc7
Show file tree

Hide file tree

Showing 10 changed files with 2,046 additions and 1,127 deletions.
diff --git a/README.md b/README.md
@@ -239,6 +239,8 @@ If you need other algorithms in your work that may be of general use, please ope
 | [Binary Search](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/binarysearch/) | `searchsortedfirst` `searchsortedfirst!`         | `std::lower_bound`                                        |
 |                                               | `searchsortedlast` `searchsortedlast!`           | `thrust::upper_bound`                                     |
 | [Predicates](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/predicates/) | `all` `any`                                      |                                                           |
+| [Arithmetics](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/arithmetics/) | `sum` `prod` `minimum` `maximum` `count` `cumsum` `cumprod`                                      |                                                           |
+
 
 
 ## 5. API and Examples

diff --git a/docs/make.jl b/docs/make.jl
@@ -29,6 +29,7 @@ makedocs(;
             "Accumulate" => "api/accumulate.md",
             "Binary Search" => "api/binarysearch.md",
             "Predicates" => "api/predicates.md",
+            "Arithmetics" => "api/arithmetics.md",
             "Custom Structs" => "api/custom_structs.md",
             "Task Partitioning" => "api/task_partition.md",
             "Utilities" => "api/utilities.md",

diff --git a/docs/src/api/arithmetics.md b/docs/src/api/arithmetics.md
@@ -0,0 +1,11 @@
+### Arithmetics
+
+```@docs
+AcceleratedKernels.sum
+AcceleratedKernels.prod
+AcceleratedKernels.minimum
+AcceleratedKernels.maximum
+AcceleratedKernels.count
+AcceleratedKernels.cumsum
+AcceleratedKernels.cumprod
+```
diff --git a/docs/src/api/mapreduce.md b/docs/src/api/mapreduce.md
@@ -3,8 +3,6 @@
 Equivalent to `reduce(op, map(f, iterable))`, without saving the intermediate mapped collection; can be used to e.g. split documents into words (map) and count the frequency thereof (reduce).
 - **Other names**: `transform_reduce`, some `fold` implementations include the mapping function too.
 
-**New in AcceleratedKernels 0.2.0: N-dimensional reductions via the `dims` keyword**
-
 ---
 
 ```@docs

diff --git a/docs/src/api/reduce.md b/docs/src/api/reduce.md
@@ -3,8 +3,6 @@
 Apply a custom binary operator reduction on all elements in an iterable; can be used to compute minima, sums, counts, etc.
 - **Other names**: `Kokkos:parallel_reduce`, `fold`, `aggregate`.
 
-**New in AcceleratedKernels 0.2.0: N-dimensional reductions via the `dims` keyword**
-
 ---
 
 ```@docs

diff --git a/ext/AcceleratedKernelsMetalExt.jl b/ext/AcceleratedKernelsMetalExt.jl
@@ -55,4 +55,60 @@ function AK.accumulate!(
 end
 
 
+function AK.cumsum(
+    src::AbstractArray, backend::MetalBackend;
+    init=zero(eltype(src)),
+    dims::Union{Nothing, Int}=nothing,
+
+    # Algorithm choice
+    alg::AK.AccumulateAlgorithm=AK.ScanPrefixes(),
+
+    # GPU settings
+    block_size::Int=256,
+    temp::Union{Nothing, AbstractArray}=nothing,
+    temp_flags::Union{Nothing, AbstractArray}=nothing,
+)
+    AK.accumulate(
+        +, src, backend;
+        init=init,
+        dims=dims,
+        inclusive=true,
+
+        alg=alg,
+
+        block_size=block_size,
+        temp=temp,
+        temp_flags=temp_flags,
+    )
+end
+
+
+function AK.cumprod(
+    src::AbstractArray, backend::MetalBackend;
+    init=one(eltype(src)),
+    dims::Union{Nothing, Int}=nothing,
+
+    # Algorithm choice
+    alg::AK.AccumulateAlgorithm=AK.ScanPrefixes(),
+
+    # GPU settings
+    block_size::Int=256,
+    temp::Union{Nothing, AbstractArray}=nothing,
+    temp_flags::Union{Nothing, AbstractArray}=nothing,
+)
+    AK.accumulate(
+        *, src, backend;
+        init=init,
+        dims=dims,
+        inclusive=true,
+
+        alg=alg,
+
+        block_size=block_size,
+        temp=temp,
+        temp_flags=temp_flags,
+    )
+end
+
+
 end   # module AcceleratedKernelsMetalExt
diff --git a/src/AcceleratedKernels.jl b/src/AcceleratedKernels.jl
@@ -35,6 +35,7 @@ include("reduce/reduce.jl")
 include("accumulate/accumulate.jl")
 include("searchsorted.jl")
 include("truth.jl")
+include("arithmetics.jl")
 
 
 end     # module AcceleratedKernels
diff --git a/src/accumulate/accumulate.jl b/src/accumulate/accumulate.jl
@@ -211,6 +211,10 @@ end
         dims::Union{Nothing, Int}=nothing,
         inclusive::Bool=true,
 
+        # Algorithm choice
+        alg::AccumulateAlgorithm=DecoupledLookback(),
+
+        # GPU settings
         block_size::Int=256,
         temp::Union{Nothing, AbstractArray}=nothing,
         temp_flags::Union{Nothing, AbstractArray}=nothing,
@@ -224,6 +228,10 @@ function accumulate(
     dims::Union{Nothing, Int}=nothing,
     inclusive::Bool=true,
 
+    # Algorithm choice
+    alg::AccumulateAlgorithm=DecoupledLookback(),
+
+    # GPU settings
     block_size::Int=256,
     temp::Union{Nothing, AbstractArray}=nothing,
     temp_flags::Union{Nothing, AbstractArray}=nothing,
@@ -236,6 +244,8 @@ function accumulate(
         init=init,
         dims=dims,
         inclusive=inclusive,
+
+        alg=alg,
 
         block_size=block_size,
         temp=temp,