In-place mutation with Referenceables.jl

using CUDA
using Folds
using FoldsCUDA
using Referenceables: referenceable

Simple mutation

function increment_folds!(xs)
    Folds.foreach(referenceable(xs)) do x
        x[] += 1
    end
    return xs
end

if has_cuda_gpu()
    xs = CuArray(1:5)
else
    xs = Array(1:5)
end

collect(increment_folds!(xs))

5-element Vector{Int64}:
 2
 3
 4
 5
 6

This can also be written with FLoops.jl:

using FLoops

function increment_floops!(xs, ex = nothing)
    @floop ex for x in referenceable(xs)
        x[] += 1
    end
    return xs
end

collect(increment_floops!(xs))

5-element Vector{Int64}:
 3
 4
 5
 6
 7

Fusing reduction and mutationg

Computing sum(f, xs) and f.(xs) in one go:

function mutation_with_folds(f, xs)
    ys = similar(xs)
    s = Folds.sum(zip(referenceable(ys), xs)) do (r, x)
        r[] = y = f(x)
        return y
    end
    return s, ys
end

if has_cuda_gpu()
    xs = CuArray(1:5)
else
    xs = Array(1:5)
end

s, ys = mutation_with_folds(x -> x^2, xs)
s

55

collect(ys)

5-element Vector{Int64}:
  1
  4
  9
 16
 25

An equilvalent computaton with FLoops.jl:

using FLoops

function mutation_with_floop(f, xs)
    ys = similar(xs)
    z = zero(eltype(ys))
    @floop for (r, x) in zip(referenceable(ys), xs)
        r[] = y = f(x)
        @reduce(s = z + y)
    end
    return s, ys
end

@assert mutation_with_folds(x -> x^2, xs) == (s, ys)

This page was generated using Literate.jl.