From bbcb18293abf4d3aa811870d4b368bcfe5330fa9 Mon Sep 17 00:00:00 2001
From: Abdulmalik Alquwayfili <af.alquwayfili@gmail.com>
Date: Wed, 7 Oct 2026 14:05:54 +0300
Subject: [PATCH] DArray: Do not pass Base._InitialValue to partition
 reductions

When the caller gives no `init`, each partition was reduced with
`init=Base._InitialValue()`. `Base`'s own reductions read that as "no
`init`", but GPU arrays take it for the initial value, so the kernel
calls `add_sum(::Base._InitialValue, ::Float32)` and fails to compile:

    InvalidIRError: compiling MethodInstance for
    CUDACore.partial_mapreduce_grid(::typeof(identity), ::typeof(Base.add_sum), ::Base._InitialValue, ...)
    Reason: unsupported call to an unknown function (call to jl_f_throw_methoderror)

`sum`, `prod`, `mapreduce`, `minimum`, `maximum` and `extrema` all
failed this way on a `DArray` of GPU arrays. Leave `init` out when
there is none.
---
 src/array/map-reduce.jl | 13 +++++++++----
 test/gpu.jl             | 16 ++++++++++++++++
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl
index 3ea9d9fa8..739a7a448 100644
--- a/src/array/map-reduce.jl
+++ b/src/array/map-reduce.jl
@@ -72,16 +72,21 @@ end
 
 _mapped_type(f, op_inner, ::Type{T}) where T = Base._return_type(f, Tuple{T})
 
+# `Base._InitialValue` only means "no `init`" to `Base`'s own reductions. GPU
+# arrays take it for the initial value, so leave `init` out instead.
+_reduce_kwargs(dims, init) = init isa Base._InitialValue ? (; dims) : (; dims, init)
+
 function stage(ctx::Context, r::MapReduce{T,N}) where {T,N}
     inp = stage(ctx, r.input)
 
     # Reduce partitions individually
     dims = r.dims === nothing ? Colon() : r.dims
+    kwargs = _reduce_kwargs(dims, r.init)
     reduced_parts = map(chunks(inp)) do part
         if r.op_inner !== nothing
-            Dagger.@spawn r.op_inner(r.f, part; dims, init=r.init)
+            Dagger.@spawn r.op_inner(r.f, part; kwargs...)
         else
-            Dagger.@spawn mapreduce(r.f, r.op_outer, part; dims=dims, init=r.init)
+            Dagger.@spawn mapreduce(r.f, r.op_outer, part; kwargs...)
         end
     end
 
@@ -182,8 +187,8 @@ struct Extrema{T} <: Number
     max::T
 end
 Extrema(minmax::NTuple{2,T}) where T = Extrema{T}(minmax[1], minmax[2])
-function _extrema_inner(f, X; dims, init)
-    result = extrema(f, X; dims, init)
+function _extrema_inner(f, X; dims, init=Base._InitialValue())
+    result = extrema(f, X; _reduce_kwargs(dims, init)...)
     if dims === Colon()
         return Extrema(result)
     else
diff --git a/test/gpu.jl b/test/gpu.jl
index 7ace19cca..b6f236c7f 100644
--- a/test/gpu.jl
+++ b/test/gpu.jl
@@ -20,6 +20,22 @@ include(joinpath(@__DIR__, "array", "linalg", "matvec_defs.jl"))
     end
 end
 
+@testset "DArray reductions ($name)" for (name, scope) in GPU_SCOPES
+    Dagger.with_options(; scope) do
+        DA = rand(Blocks(4, 4), Float32, 16, 16)
+        A = collect(DA)
+        @test sum(DA) ≈ sum(A)
+        @test sum(DA; init=0f0) ≈ sum(A)
+        @test prod(DA) ≈ prod(A)
+        @test mapreduce(abs2, +, DA) ≈ mapreduce(abs2, +, A)
+        @test maximum(DA) == maximum(A)
+        @test extrema(DA) == extrema(A)
+        @test collect(sum(DA; dims=1)) ≈ sum(A; dims=1)
+        @test collect(maximum(DA; dims=2)) == maximum(A; dims=2)
+        @test collect(extrema(DA; dims=1)) == extrema(A; dims=1)
+    end
+end
+
 @everywhere begin
     function isongpu(X)
         return !(X isa Array)