Merge branch 'main' into wsmoses-patch-4

Author: avik-pal

deps/build_local.jl

@@ -172,6 +172,7 @@ else
     push!(build_cmd_list, "--define=using_clang=true")
     push!(build_cmd_list, "--copt=-Wno-unused-command-line-argument")
 end
+push!(build_cmd_list, "--copt=-Wno-private-header")
 push!(build_cmd_list, "--color=$(parsed_args["color"])")
 push!(build_cmd_list, ":libReactantExtra.so")
 

ext/ReactantCUDAExt.jl

@@ -1489,6 +1489,7 @@ end
                 end
                 y = Reactant.ConcreteRArray([2.0]; client)
                 Reactant.Compiler.compile_mlir(square!, (y,); optimize=false)
+                finalize(y)
             end
         end
 

src/Ops.jl

@@ -1064,6 +1064,7 @@ function broadcast_in_dim(
     location=mlir_stacktrace("broadcast_in_dim", @__FILE__, @__LINE__),
 ) where {T,N}
     @assert length(dims) == N
+    @assert length(result_size) ≥ N
 
     res = MLIR.IR.result(
         stablehlo.broadcast_in_dim(

src/Precompile.jl

@@ -69,9 +69,11 @@ if Reactant_jll.is_available()
             @static if precompilation_supported()
                 x = ConcreteRNumber(2.0; client)
                 Reactant.compile(sin, (x,); client, optimize=:all)
+                finalize(x)
 
                 y = ConcreteRArray([2.0]; client)
                 Reactant.compile(Base.sum, (y,); client, optimize=:all)
+                finalize(y)
             end
         end
 

src/TracedRArray.jl

@@ -543,7 +543,14 @@ end
 function overloaded_mapreduce(
     @nospecialize(f), @nospecialize(op), @nospecialize(A); dims=:, init=nothing
 )
-    return overloaded_mapreduce(identity, op, unwrapped_broadcast(f, A); dims, init)
+    res = unwrapped_broadcast(f, A)
+    # This means we are unable to use the optimized dispatches. For now we will
+    # unroll the mapreduce.
+    if typeof(res) == typeof(A)
+        @assert dims == Colon() "dims not supported for mapreduce currently."
+        return foldl(op, res; init)
+    end
+    return overloaded_mapreduce(identity, op, res; dims=:, init)
 end
 
 function overloaded_mapreduce(
@@ -1436,16 +1443,27 @@ struct BroadcastIterator{F}
     f::F
 end
 
-(fn::BroadcastIterator)(args...) = fn.f((args...,))
+(fn::BroadcastIterator)(args...) = Reactant.call_with_reactant(fn.f, (args...,))
 
 function unwrapped_broadcast(f::F, x::Base.Iterators.Zip) where {F}
     min_length = Base.inferencebarrier(minimum)(length, x.is)
     itrs = [length(itr) > min_length ? itr[1:min_length] : itr for itr in x.is]
-    return (BroadcastIterator(f)).(itrs...)
+    if any(Base.Fix2(isa, AnyTracedRArray), itrs)
+        return (BroadcastIterator(f)).(itrs...)
+    else
+        fn = BroadcastIterator(f)
+        return [fn(Base.Fix2(getindex, i).(itrs)...) for i in 1:min_length]
+    end
 end
 
 function unwrapped_broadcast(f::F, x::Base.Iterators.Enumerate) where {F}
-    return (BroadcastIterator(f)).(1:length(x.itr), x.itr)
+    if x.itr isa AnyTracedRArray
+        return (BroadcastIterator(f)).(1:length(x.itr), x.itr)
+    else
+        return [f((i, x.itr[i])) for i in 1:length(x.itr)]
+    end
 end
 
+unwrapped_broadcast(f::F, xs::Vector) where {F} = [f(x) for x in xs]
+
 end

src/xla/Buffer.jl

@@ -1,5 +1,6 @@
 abstract type AbstractBuffer end
 
+function free_buffer end
 function synced_buffer end
 function buffer_on_cpu end
 function to_host end

src/xla/IFRT/Array.jl

@@ -3,7 +3,7 @@ mutable struct Array <: XLA.AbstractBuffer
 
     function Array(buffer::Ptr{Cvoid}, owned::Bool=true)
         !owned && return new(buffer)
-        return finalizer(free_ifrt_array, new(buffer))
+        return finalizer(XLA.free_buffer, new(buffer))
     end
 end
 
@@ -158,7 +158,7 @@ function Array(
     return Array(client, array, ifrt_sharding)
 end
 
-@inline function free_ifrt_array(buffer::Array)
+@inline function XLA.free_buffer(buffer::Array)
     if buffer.buffer != C_NULL
         @ccall MLIR.API.mlir_c.ifrt_free_array(buffer.buffer::Ptr{Cvoid})::Cvoid
     end

src/xla/PJRT/Buffer.jl

@@ -2,7 +2,7 @@ mutable struct Buffer <: XLA.AbstractBuffer
     buffer::Ptr{Cvoid}
 
     function Buffer(buffer::Ptr{Cvoid})
-        return finalizer(free_buffer, new(buffer))
+        return finalizer(XLA.free_buffer, new(buffer))
     end
 end
 
@@ -114,7 +114,7 @@ function Base.similar(a::Buffer, S::Type, dims::Dims)
     return Base.similar(Buffer, S, dims; client=XLA.client(a), device=XLA.device(a))
 end
 
-@inline function free_buffer(buffer::Buffer)
+@inline function XLA.free_buffer(buffer::Buffer)
     sbuffer = buffer.buffer
     if sbuffer != C_NULL
         @ccall MLIR.API.mlir_c.PjRtBufferFree(sbuffer::Ptr{Cvoid})::Cvoid

test/basic.jl

@@ -1444,6 +1444,20 @@ end
 zip_iterator(a, b) = mapreduce(splat(*), +, zip(a, b))
 enumerate_iterator(a) = mapreduce(splat(*), +, enumerate(a))
 
+function nested_mapreduce_zip(x, y)
+    return mapreduce(+, zip(eachcol(x), eachcol(y)); init=0.0f0) do (x, y)
+        return sum(abs2, x) + sum(abs2, y)
+    end
+end
+
+function nested_mapreduce_hcat(x, y)
+    return mapreduce(
+        hcat, zip(eachcol(x), eachcol(y)); init=similar(x, size(x, 1), 0)
+    ) do (x, y)
+        return x .+ y
+    end
+end
+
 @testset "Base.Iterators" begin
     @testset "zip" begin
         N = 10
@@ -1460,4 +1474,24 @@ enumerate_iterator(a) = mapreduce(splat(*), +, enumerate(a))
 
         @test @jit(enumerate_iterator(x_ra)) ≈ enumerate_iterator(x)
     end
+
+    @testset "nested mapreduce" begin
+        x = rand(Float32, 4, 3)
+        y = rand(Float32, 4, 3)
+
+        x_ra = Reactant.to_rarray(x)
+        y_ra = Reactant.to_rarray(y)
+
+        @test @jit(nested_mapreduce_zip(x_ra, y_ra)) ≈ nested_mapreduce_zip(x, y)
+    end
+
+    @testset "nested mapreduce hcat" begin
+        x = rand(Float32, 4, 3)
+        y = rand(Float32, 4, 3)
+
+        x_ra = Reactant.to_rarray(x)
+        y_ra = Reactant.to_rarray(y)
+
+        @test @jit(nested_mapreduce_hcat(x_ra, y_ra)) ≈ nested_mapreduce_hcat(x, y)
+    end
 end