reducing while iterating over the index

[albertomercurio]

Hello,

I have a situation where I need to reduce and iterate over the indices of the array, rather than the array itself. A simple example:

x = rand(10);

AcceleratedKernels.mapreduce(+, eachindex(x)) do i
   i > 8 && return 0
   return x[i] + x[i+1]
end

ERROR: ArgumentError: Implement `KernelAbstractions.get_backend(::Base.OneTo{Int64})`
Stacktrace:
 [1] get_backend(A::Base.OneTo{Int64})
   @ KernelAbstractions ~/.julia/packages/KernelAbstractions/lGrz7/src/KernelAbstractions.jl:522
 [2] mapreduce(f::Function, op::Function, src::Base.OneTo{Int64})
   @ AcceleratedKernels ~/.julia/packages/AcceleratedKernels/AdYRJ/src/reduce/reduce.jl:156
 [3] top-level scope
   @ REPL[9]:1

This would be helpful on implementing the trace of a sparse CSC matrix. My current implementation uses the foreachindex function and atomic operations, which I don't think are optimal

function LinearAlgebra.tr(A::DeviceSparseMatrixCSC)
    m, n = size(A)
    m == n || throw(DimensionMismatch("Matrix must be square to compute the trace."))

    res = similar(nonzeros(A), 1)
    fill!(res, zero(eltype(A)))

    AcceleratedKernels.foreachindex(getcolptr(A)) do col
        col == n+1 && return
        for j in getcolptr(A)[col]:(getcolptr(A)[col+1]-1)
            if rowvals(A)[j] == col
                @atomic res[1] += nonzeros(A)[j]
            end
        end
    end

    return res[1]
end

[yolhan83]

When using iterators like each index you need to also specify the backend CPU or CudaBackend or others.

[albertomercurio]

Thanks, by specifying both backend and init it works.

However, it doesn't work if x is a JLArray

using JLArrays

dx = JLArray(x)

AcceleratedKernels.mapreduce(+, eachindex(dx), get_backend(dx); init=zero(eltype(dx))) do i
  i > 8 && return 0
  return dx[i] + dx[i+1]
end

ERROR: This kernel is unavailable for backend CPU
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] __mapreduce_block!(dev::JLBackend, sz::KernelAbstractions.NDIteration.StaticSize{…}, range::KernelAbstractions.NDIteration.DynamicSize)
   @ AcceleratedKernels ~/.julia/packages/KernelAbstractions/lGrz7/src/macros.jl:63
 [3] _mapreduce_block!(dev::JLBackend, size::Int64)
   @ AcceleratedKernels ~/.julia/packages/KernelAbstractions/lGrz7/src/macros.jl:68
 [4] mapreduce_1d_gpu(f::var"#5#6", op::typeof(+), src::Base.OneTo{…}, backend::JLBackend; init::Float64, neutral::Int64, max_tasks::Int64, min_elems::Int64, block_size::Int64, temp::Nothing, switch_below::Int64)
   @ AcceleratedKernels ~/.julia/packages/AcceleratedKernels/AdYRJ/src/reduce/mapreduce_1d_gpu.jl:94
 [5] mapreduce_1d_gpu
   @ ~/.julia/packages/AcceleratedKernels/AdYRJ/src/reduce/mapreduce_1d_gpu.jl:49 [inlined]
 [6] _mapreduce_impl(f::var"#5#6", op::typeof(+), src::Base.OneTo{…}, backend::JLBackend; init::Float64, neutral::Int64, dims::Nothing, max_tasks::Int64, min_elems::Int64, prefer_threads::Bool, block_size::Int64, temp::Nothing, switch_below::Int64)
   @ AcceleratedKernels ~/.julia/packages/AcceleratedKernels/AdYRJ/src/reduce/reduce.jl:187
 [7] #mapreduce#68
   @ ~/.julia/packages/AcceleratedKernels/AdYRJ/src/reduce/reduce.jl:161 [inlined]
 [8] top-level scope
   @ REPL[8]:1

I think it would be useful to support JLArrays.jl in order to test the kernels also in the CPU.

[yolhan83]

Not sure it works with JLArrays.jl indeed but the kernels on CPU are really well written in AK using complex OhMyThreads.jl so don't expect too mush

[yolhan83]

It may get available once KernelAbstractuon Pocl backend is released though and when it's performance will be on par with the ones here

[albertomercurio]

But aren't the methods implemented here based on KernelAbstractions?

Anyhow, I have also tried using Reactant.jl, which doesn't seem to work

using Reactant

dx = Reactant.to_rarray(x)

function test(dx)
  AcceleratedKernels.mapreduce(+, eachindex(dx), get_backend(dx); init=zero(eltype(dx))) do I
    i > 8 && return 0
    return dx[i] + dx[i+1]
  end
end

test(dx)

ERROR: MethodError: no method matching ka_with_reactant(::Tuple{…}, ::Nothing, ::KernelAbstractions.Kernel{…}, ::UnitRange{…}, ::SubArray{…}, ::var"#5#6"{…}, ::typeof(+), ::Int64)
The function `ka_with_reactant` exists, but no method is defined for this combination of argument types.
Attempted to raise a KernelAbstractions kernel with Reactant but CUDA.jl is not loaded.
Load CUDA.jl using `using CUDA`. You might need to restart the Julia process (even if Revise.jl is loaded).
Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/Reactant/uU0IF/src/utils.jl:0 [inlined]
  [2] call_with_reactant(::typeof(Reactant.ka_with_reactant), ::Tuple{…}, ::Nothing, ::KernelAbstractions.Kernel{…}, ::UnitRange{…}, ::SubArray{…}, ::var"#5#6"{…}, ::typeof(+), ::Int64)
    @ Reactant ~/.julia/packages/Reactant/uU0IF/src/utils.jl:893
  [3] (::KernelAbstractions.Kernel{ReactantKernelAbstractionsExt.ReactantBackend})(::Any, ::Vararg{Any}; ndrange::Any, workgroupsize::Any)
    @ ReactantKernelAbstractionsExt ~/.julia/packages/Reactant/uU0IF/ext/ReactantKernelAbstractionsExt.jl:118
  [4] #kwcall
    @ ~/.julia/packages/Reactant/uU0IF/ext/ReactantKernelAbstractionsExt.jl:112 [inlined]
 ...

Do you plan to also support Reactant arrays?

[yolhan83]

1- yes but only the GPU ones for now
2- it's not like that we use Reactant use @jit test(dx) but not sure it works
ps : Reactant will try to raise the kernel to stablehlo ones so you won't win anything using AK just use base map reduce and Reactant will turn that to MLIR and run on CPU, GPU or TPU depending on the backend you use with Reactant and even use some stocastic to set the number of thread without even talking about the simd it will use
However AcceleratedKernels should support Reactant indeed

[albertomercurio]

I tried to both compile the function or using the simple mapreduce with Reactant, with it returns scalar indexing error.

[albertomercurio]

Hi @yolhan83, I have now a similar situation, where I need to iterate over different dimensions.

function test(x, M, y, i)
    N = size(x, 2)
    res = 0.0
    for α in 1:3
        for β in 1:3
            @inbounds for n in 1:N
                res += x[α, n] * M[n, i] * y[β, n]
            end
        end
    end
    return res
end

function test_gpu_v2(x, M, y, i)
    backend = get_backend(x)
    N = size(x, 2)

    indices = CartesianIndices((1:3, 1:3, 1:N))
    res = AcceleratedKernels.mapreduce(+, indices, backend; init=zero(eltype(x))) do idx
        α, β, n = Tuple(idx)
        @inbounds x[α, n] * M[n, i] * y[β, n]
    end
    return res
end

However I get an error for the GPU case

ERROR: InvalidIRError: compiling MethodInstance for AcceleratedKernels.gpu__mapreduce_block!(::KernelAbstractions.CompilerMetadata{…}, ::SubArray{…}, ::CuDeviceVector{…}, ::var"#27#28"{…}, ::typeof(+), ::CartesianIndex{…}) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to print_to_string(xs...) @ Base strings/io.jl:137)
Stacktrace:
  [1] string
    @ ./strings/io.jl:189
  [2] _throw_dmrs
    @ ./reshapedarray.jl:192
  [3] _reshape
    @ ./reshapedarray.jl:187
  [4] reshape
    @ ./reshapedarray.jl:122
  [5] adapt_structure
    @ ~/.julia/packages/Adapt/2UZ81/src/wrappers.jl:18
  [6] adapt
    @ ~/.julia/packages/Adapt/2UZ81/src/Adapt.jl:40
  [7] adapt_structure
    @ ~/.julia/packages/Adapt/2UZ81/src/wrappers.jl:10
  [8] adapt
    @ ~/.julia/packages/Adapt/2UZ81/src/Adapt.jl:40
  [9] constify
    @ ~/.julia/packages/KernelAbstractions/lGrz7/src/KernelAbstractions.jl:464
 [10] gpu__mapreduce_block!
    @ ./none:0
Reason: unsupported dynamic function invocation (call to print_to_string(xs...) @ Base strings/io.jl:137)
Stacktrace

I don't even know if this makes sense. But how to do it otherwise?

[yolhan83]

you just forgot the neutral

function test_gpu_v2(x, M, y, i)
           backend = get_backend(x)
           N = size(x, 2)
           indices =CartesianIndices((1:3, 1:3, 1:N))
           res = AcceleratedKernels.mapreduce(+, indices, backend; init=zero(eltype(x)),neutral = zero(eltype(x))) do idx
                   α, β, n = Tuple(idx)
                   x[α, n] * M[n, i] * y[β, n]
           end
           return res
       end

nevermind still not working on gpu , I was able to cheat sending it to CUDA but thats not perfect

function test_gpu_v2(x, M, y, i)
           backend = get_backend(x)
           N = size(x, 2)
           indices =CartesianIndices((1:3, 1:3, 1:N)) |> CUDA.CuArray
           res = AcceleratedKernels.mapreduce(+, indices, backend; init=zero(eltype(x)),neutral = zero(eltype(x))) do idx
                   α, β, n = Tuple(idx)
                   x[α, n] * M[n, i] * y[β, n]
           end
           return res
       end

this is complex enough that you should write a kernel. the simplest one I can write is around 3x perf as the version where I send the indices to the gpu

using KernelAbstractions,CUDA,Atomix
@kernel function test_gpuker!(dest,x,M,y,i)
           α,β,n = @index(Global,NTuple)
           Atomix.@atomic dest[] += x[α, n] * M[n, i] * y[β, n]
 end

function test_gpu_v3(x,M,y,i)
   bck = get_backend(x)
   dest = allocate(bck,eltype(x),1)
   dest .= 0
   N = size(x,2)
   test_gpuker!(bck)(dest,x,M,y,i;ndrange=(3,3,N))
   CUDA.@allowscalar dest[]
end

and I think the 3 mul + all the non-alligned indexing is complex enough and a "standard" version of the reduction won't make this better.
note that this completly destroy the cpu perf (for now) , so I would add a dispatch on a version like

function test_parallel(x, M, y, i)
    N = size(x, 2)
    it = Iterators.partition(1:N, cld(N, 100))
    res = zeros(eltype(x),length(it))
    @sync for (id,I) in enumerate(it)
        Threads.@spawn for n in I
            for α in 1:3
                for β in 1:3
                    res[id] += x[α, n] * M[n, i] * y[β, n]
                end
            end
        end
    end
    return sum(res)
end

edit: I just tested with complex standard reduction

@kernel inbounds=true function test_gpuker!(dest, x, M, y, i)
    α, β, n = @index(Global, NTuple)
    lid     = @index(Local)    
    lsize   = @groupsize()[1] 
    smem = @localmem(Float32, 128)
    val = x[α, n] * M[n, i] * y[β, n]
    smem[lid] = val
    KernelAbstractions.@synchronize()
    stride = lsize >>> 1
    while stride > 0
        if lid <= stride
            smem[lid] += smem[lid + stride]
        end
        KernelAbstractions.@synchronize()
        stride >>>= 1
    end
    if lid == 1
        Atomix.@atomic dest[] += smem[1]
    end
end
function test_gpu_v3(x,M,y,i)
   bck = get_backend(x)
   dest = allocate(bck,eltype(x),1)
   dest .= 0 
   N = size(x,2)
   test_gpuker!(bck,128)(dest,x,M,y,i;ndrange=(3,3,N))
   CUDA.@allowscalar dest[]
end

[albertomercurio]

Thanks a lot. So you suggest to manually write the kernel reduction by hand. Is your last example the best way to perform reductions like this?

Your last implementation is the fastest for GPU, despite it has the same performances of the single core CPU case.

But this difference should decrease as I increase the size of the arrays.

[yolhan83]

Yes I think you should go for it this just go a little out of scope for AK, also I forgot to use KernelAbstractions.Extra.@unroll on the while loop

[yolhan83]

I'm sorry the kernel was really bad spawing a lot too much, need to use linear

using KernelAbstractions,CUDA,Atomix
import CUDA: i32
@kernel inbounds=true unsafe_indices=true function test_gpuker!(dest, @Const(x), @Const(M), @Const(y), @Const(i),@Const(sizes))
    I = @index(Global)
    lid    = @index(Local)         
    lsize  = @groupsize()[1]       
    smem = @localmem(eltype(x), lsize) 
    _, _, nd = sizes
    J = I - 1i32
    α = (J ÷ 1i32) % 3i32 + 1i32
    β = (J ÷ 3i32) % 3i32 + 1i32
    n = (J ÷ 9i32) + 1i32

    val = zero(eltype(x))
    if 1 <= α <= 3 && 1 <= β <= 3 && 1 <= n <= nd
        val = x[α, n] * M[n, i] * y[β, n]
    end
    smem[lid] = val
    KernelAbstractions.@synchronize()
    stride = lsize >>> 1i32
    while stride > 0i32
        if lid <= stride
            smem[lid] += smem[lid + stride]
        end
        KernelAbstractions.@synchronize()
        stride >>>= 1i32
    end
    if lid == 1i32
        Atomix.@atomic dest[] += smem[1]
    end
end


function test_gpu_v3(x,M,y,i)
   bck = get_backend(x)
   dest = KernelAbstractions.zeros(bck,eltype(x),1)
   N = size(x,2)
   test_gpuker!(bck,256)(dest,x,M,y,i,(3,3,N);ndrange=3*3*N)
   CUDA.@allowscalar dest[]
end

you can change the 3i32 to the alpha_d and beta_d from sizes if you want but since it allows to go from 32 regs to 20 I think its cool this way, really sorry for the poorly written one before, now its completly killing the cpu parallel version

[albertomercurio]

Thank you so much. So standard reduction over a linear indexing is the way.