demo.jl

#!/bin/bash
# -*- mode: julia -*-
#=
JULIA="${JULIA:-julia}"
JULIA_CMD="${JULIA_CMD:-${JULIA} --color=yes --startup-file=no}"
export JULIA_PROJECT="$(dirname "$(realpath "${BASH_SOURCE[0]}")")"

set -ex
${JULIA_CMD} -e 'using Pkg; Pkg.instantiate()'
exec ${JULIA_CMD} "$@" "${BASH_SOURCE[0]}"
=#

using CUDA

if has_cuda_gpu()
    CUDA.allowscalar(false)
end

include("ka_mapreduce.jl")
@show ka_mapreduce(identity, +, CUDA.ones(3); device = CUDADevice())

ka_mapreduce.jl

using Core: Typeof
using Core.Compiler: return_type
using GPUArrays: @allowscalar
using InitialValues
using KernelAbstractions

@kernel function ka_mapreduce_kernel!(f, op, ys, xs, init, basesize)
    i, = @index(Global, NTuple)
    acc = op(init, f(xs[basesize*(i-1)+1]))
    for j in (basesize*(i-1)+2):min(basesize * i, size(xs)[1])
        acc = op(acc, f(xs[j]))
    end
    ys[i] = acc
end

function fake_mapreduce(f, op, xs, init)
    acc = op(init, f(first(xs)))
    for x in xs
        acc = op(acc, f(x))
    end
    return acc
end

function ka_mapreduce(
    f,
    op,
    xs::AbstractVector;
    init = InitialValue(op),
    basesize = 256,
    device = CPU(),
)
    isempty(xs) && return init
    acctype = return_type(
        fake_mapreduce,
        Tuple{Typeof(f),Typeof(op),Vector{eltype(xs)},Typeof(init)},
    )
    kernel! = ka_mapreduce_kernel!(device, basesize)  # basesize?

    ys = similar(xs, acctype, cld(length(xs), basesize))
    wait(kernel!(f, op, ys, xs, init, basesize; ndrange = size(ys)))
    while true
        length(ys) == 1 && return @allowscalar first(ys)
        zs = similar(ys, acctype, cld(length(ys), basesize))
        wait(kernel!(identity, op, zs, ys, init, basesize; ndrange = size(zs)))
        ys = zs
    end
end

Manifest.toml

# This file is machine-generated - editing it directly is not advised

[[AbstractFFTs]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "051c95d6836228d120f5f4b984dd5aba1624f716"
uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
version = "0.5.0"

[[Adapt]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "42c42f2221906892ceb765dbcb1a51deeffd86d7"
uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
version = "2.3.0"

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[BinaryProvider]]
deps = ["Libdl", "Logging", "SHA"]
git-tree-sha1 = "ecdec412a9abc8db54c0efc5548c64dfce072058"
uuid = "b99e7846-7c00-51b0-8f62-c81ae34c0232"
version = "0.5.10"

[[CEnum]]
git-tree-sha1 = "215a9aa4a1f23fbd05b92769fdd62559488d70e9"
uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82"
version = "0.4.1"

[[CUDA]]
deps = ["AbstractFFTs", "Adapt", "BinaryProvider", "CEnum", "DataStructures", "ExprTools", "GPUArrays", "GPUCompiler", "LLVM", "Libdl", "LinearAlgebra", "Logging", "MacroTools", "NNlib", "Pkg", "Printf", "Random", "Reexport", "Requires", "SparseArrays", "Statistics", "TimerOutputs"]
git-tree-sha1 = "83bfd180e2f842f6d4ee315a6db8665e9aa0c19b"
uuid = "052768ef-5323-5732-b1bb-66c8b64840ba"
version = "1.3.3"

[[Cassette]]
git-tree-sha1 = "ff6f5109371926beb67ec3101be17d2c211e497d"
uuid = "7057c7e9-c182-5462-911a-8362d720325c"
version = "0.3.3"

[[Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "cf03b37436c6bc162e7c8943001568b4cad4bee3"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.19.0"

[[CompilerSupportLibraries_jll]]
deps = ["Libdl", "Pkg"]
git-tree-sha1 = "7c4f882c41faa72118841185afc58a2eb00ef612"
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
version = "0.3.3+0"

[[DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "db07bb22795762895b60e44d62b34b16c982a687"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.7"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[ExprTools]]
git-tree-sha1 = "7fce513fcda766962ff67c5596cb16c463dfd371"
uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
version = "0.1.2"

[[GPUArrays]]
deps = ["AbstractFFTs", "Adapt", "LinearAlgebra", "Printf", "Random", "Serialization"]
git-tree-sha1 = "da6398282abd2a8c0dc3e55b49d984fcc2c582e5"
uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
version = "5.2.1"

[[GPUCompiler]]
deps = ["DataStructures", "InteractiveUtils", "LLVM", "Libdl", "TimerOutputs", "UUIDs"]
git-tree-sha1 = "05097d81898c527e3bf218bb083ad0ead4378e5f"
uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
version = "0.6.1"

[[InitialValues]]
git-tree-sha1 = "26c8832afd63ac558b98a823265856670d898b6c"
uuid = "22cec73e-a1b8-11e9-2c92-598750a2cf9c"
version = "0.2.10"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[KernelAbstractions]]
deps = ["Adapt", "CUDA", "Cassette", "InteractiveUtils", "MacroTools", "SpecialFunctions", "StaticArrays"]
git-tree-sha1 = "9262c704a2897c7e4837e75638ec4fc522a0cc19"
uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
version = "0.4.3"

[[LLVM]]
deps = ["CEnum", "Libdl", "Printf", "Unicode"]
git-tree-sha1 = "a662366a5d485dee882077e8da3e1a95a86d097f"
uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
version = "2.0.0"

[[LibGit2]]
deps = ["Printf"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "f7d2e3f654af75f01ec49be82c231c382214223a"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.5"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[NNlib]]
deps = ["Libdl", "LinearAlgebra", "Pkg", "Requires", "Statistics"]
git-tree-sha1 = "1ef04283efe283be08e2d0de842f5e5286dd0b7a"
uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
version = "0.7.5"

[[OpenSpecFun_jll]]
deps = ["CompilerSupportLibraries_jll", "Libdl", "Pkg"]
git-tree-sha1 = "d51c416559217d974a1113522d5919235ae67a87"
uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
version = "0.5.3+3"

[[OrderedCollections]]
git-tree-sha1 = "16c08bf5dba06609fe45e30860092d6fa41fde7b"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.3.1"

[[Pkg]]
deps = ["Dates", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "UUIDs"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[Reexport]]
deps = ["Pkg"]
git-tree-sha1 = "7b1d07f411bc8ddb7977ec7f377b97b158514fe0"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "0.2.0"

[[Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "28faf1c963ca1dc3ec87f166d92982e3c4a1f66d"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.1.0"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[SpecialFunctions]]
deps = ["OpenSpecFun_jll"]
git-tree-sha1 = "d8d8b8a9f4119829410ecd706da4cc8594a1e020"
uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
version = "0.10.3"

[[StaticArrays]]
deps = ["LinearAlgebra", "Random", "Statistics"]
git-tree-sha1 = "016d1e1a00fabc556473b07161da3d39726ded35"
uuid = "90137ffa-7385-5640-81b9-e52037218182"
version = "0.12.4"

[[Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[Test]]
deps = ["Distributed", "InteractiveUtils", "Logging", "Random"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[TimerOutputs]]
deps = ["Printf"]
git-tree-sha1 = "f458ca23ff80e46a630922c555d838303e4b9603"
uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
version = "0.5.6"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

Project.toml

[deps]
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
InitialValues = "22cec73e-a1b8-11e9-2c92-598750a2cf9c"
KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"