fix: bypass kernels for 0-dim inputs

Project.toml

@@ -1,6 +1,6 @@
 name = "Strided"
 uuid = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
-version = "2.6.2"
+version = "2.6.3"
 authors = ["Lukas Devos <lukas.devos@ugent.be>", "Maarten Van Damme <maartenvd1994@gmail.com>", "Jutho Haegeman <jutho.haegeman@ugent.be>"]
 
 [deps]

src/mapreduce.jl

@@ -68,6 +68,7 @@ function _mapreduce(@nospecialize(f), @nospecialize(op), A::StridedView{T}, nt =
     dims = size(A)
     a = Base.mapreduce_first(f, op, first(A))
     a2 = nt === nothing ? a : op(a, nt.init)
+    ndims(A) == 0 && return a2 # 0-dim scalar: single-element reduction is the result
     out = similar(A, typeof(a2), (1,))
     if nt === nothing
         _init_reduction!(out, f, op, a)
@@ -162,12 +163,27 @@ function _mapreduce_order!(
         @nospecialize(f), @nospecialize(op), @nospecialize(initop),
         dims::Dims, arrays::Tuple{Vararg{StridedView}}
     )
+    isempty(dims) && return _mapreduce_scalar!(f, op, initop, arrays)
     dims, allstrides = order_and_fuse_dims(dims, map(strides, arrays))
     offsets = map(offset, arrays)
     costs = _computecosts(allstrides)
     return _mapreduce_block!(f, op, initop, dims, allstrides, offsets, costs, arrays)
 end
 
+# 0-dimensional fast path: bypass @generated kernel
+function _mapreduce_scalar!(@nospecialize(f), @nospecialize(op), @nospecialize(initop), arrays)
+    out = arrays[1]
+    iout = ParentIndex(offset(out) + 1)
+    v = f(map(a -> a[ParentIndex(offset(a) + 1)], Base.tail(arrays))...)
+    if op === nothing
+        out[iout] = v
+    else
+        o = initop === nothing ? out[iout] : initop(out[iout])
+        out[iout] = op(o, v)
+    end
+    return nothing
+end
+
 const MINTHREADLENGTH = 1 << 15 # minimal length before any kind of threading is applied
 function _mapreduce_block!(
         @nospecialize(f), @nospecialize(op), @nospecialize(initop),
@@ -269,6 +285,7 @@ function _mapreduce_threaded!(
 end
 
 function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int)
+    @assert N > 0 "Cannot generate 0-dim kernel"
 
     # many variables
     blockloopvars = Array{Symbol}(undef, N)

test/othertests.jl

@@ -155,6 +155,50 @@ end
     end
 end
 
+@testset "0-dimensional (scalar) StridedView" begin
+    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
+        R = fill(rand(T)) # 0-dimensional Array
+        A = StridedView(R)
+        @test ndims(A) == 0
+
+        # full reductions
+        @test sum(A) == sum(R)
+        @test prod(A) == prod(R)
+        @test mapreduce(abs2, +, A) == mapreduce(abs2, +, R)
+        @test maximum(abs, A) == maximum(abs, R)
+        @test minimum(abs, A) == minimum(abs, R)
+        @test sum(abs2, A) == sum(abs2, R)
+        @test mapreduce(identity, +, A; init = one(T)) ==
+            mapreduce(identity, +, R; init = one(T))
+
+        # map / map! / copy! / fill!
+        @test map(x -> 2x, A)[] == 2 * R[]
+        B = StridedView(fill(zero(T)))
+        map!(x -> x + one(T), B, A)
+        @test B[] == R[] + one(T)
+        copy!(B, A)
+        @test B[] == R[]
+        fill!(B, one(T))
+        @test B[] == one(T)
+
+        # offset handling: 0-dim views into a larger parent
+        Psrc = rand(T, 5)
+        Pdst = rand(T, 5)
+        s = sreshape(StridedView(Psrc)[4:4], ())
+        d = sreshape(StridedView(Pdst)[3:3], ())
+        @test sum(s) == Psrc[4]
+        copy!(d, s)
+        @test Pdst[3] == Psrc[4]
+
+        # low-level in-place reduction with a custom initop
+        Pd = rand(T, 5)
+        d2 = sreshape(StridedView(Pd)[2:2], ())
+        prev = Pd[2]
+        Strided._mapreducedim!(sin, +, identity, (), (d2, A))
+        @test Pd[2] == prev + sin(R[])
+    end
+end
+
 @testset "@strided macro" begin
     @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
         A1, A2, A3 = rand(T, (10,)), rand(T, (10, 10)), rand(T, (10, 10, 10))