Type stable `mapreduce` on Tuples

General Usage Performance
,
1

Hello, I want to apply this very simple example of mapreduce involving eachindex on two tuples

function test_standard(x::Tuple, y::Tuple)
    length(x) == length(y) || throw(ArgumentError("Input tuples must have the same length."))
    mapreduce((a, b) -> (a[1] + b[1], (a[2]..., b[2]...)), eachindex(x)) do i
        (x[i]^2 + y[i]^2, (x[i]^3 + y[i]^3,))
    end
end

x = (0.2, 0.4, 0.5)
y = (0.1, 0.3, 0.4)

test_standard(x, y) # (0.71, (0.009000000000000003, 0.09100000000000001, 0.189))

But it is type unstable

MethodInstance for test_standard(::Tuple{Float64, Float64, Float64}, ::Tuple{Float64, Float64, Float64})
  from test_standard(x::Tuple, y::Tuple) @ Main ~/.julia/dev/QuantumToolbox/misc/script.jl:39
Arguments
  #self#::Core.Const(Main.test_standard)
  x::Tuple{Float64, Float64, Float64}
  y::Tuple{Float64, Float64, Float64}
Locals
  #16::var"#test_standard##2#test_standard##3"
  #15::var"#test_standard##0#test_standard##1"{Tuple{Float64, Float64, Float64}, Tuple{Float64, Float64, Float64}}
Body::Tuple{Float64, Any}
1 ─       Core.NewvarNode(:(#16))
│         Core.NewvarNode(:(#15))
│   %3  = Main.:(==)::Core.Const(==)
│   %4  = Main.length::Core.Const(length)
│   %5  = (%4)(x)::Core.Const(3)
│   %6  = Main.length::Core.Const(length)
│   %7  = (%6)(y)::Core.Const(3)
│   %8  = (%3)(%5, %7)::Core.Const(true)
└──       goto #3 if not %8
2 ─       goto #4
3 ─       Core.Const(:(Main.throw))
│         Core.Const(:(Main.ArgumentError))
│         Core.Const(:((%12)("Input tuples must have the same length.")))
└──       Core.Const(:((%11)(%13)))
4 ┄ %15 = Main.mapreduce::Core.Const(mapreduce)
│   %16 = Main.:(var"#test_standard##0#test_standard##1")::Core.Const(var"#test_standard##0#test_standard##1")
│   %17 = Core._typeof_captured_variable(x)::Core.Const(Tuple{Float64, Float64, Float64})
│   %18 = Core._typeof_captured_variable(y)::Core.Const(Tuple{Float64, Float64, Float64})
│   %19 = Core.apply_type(%16, %17, %18)::Core.Const(var"#test_standard##0#test_standard##1"{Tuple{Float64, Float64, Float64}, Tuple{Float64, Float64, Float64}})
│         (#15 = %new(%19, x, y))
│   %21 = #15::var"#test_standard##0#test_standard##1"{Tuple{Float64, Float64, Float64}, Tuple{Float64, Float64, Float64}}
│   %22 = Main.:(var"#test_standard##2#test_standard##3")::Core.Const(var"#test_standard##2#test_standard##3")
│         (#16 = %new(%22))
│   %24 = #16::Core.Const(var"#test_standard##2#test_standard##3"())
│   %25 = Main.eachindex::Core.Const(eachindex)
│   %26 = (%25)(x)::Core.Const(Base.OneTo(3))
│   %27 = (%15)(%21, %24, %26)::Tuple{Float64, Any}
└──       return %27

I think the main problem is eachindex, which is not telling the compiler that the length is actually known at compile time.

I found a workaround with something like

tuple_mapreduce(f, op, iters::Tuple...) = reduce(op, map(f, iters...))
tuple_eachindex(x::Tuple) = ntuple(identity, Val(length(x)))

function test_custom(x::Tuple, y::Tuple)
    length(x) == length(y) || throw(ArgumentError("Input tuples must have the same length."))
    tuple_mapreduce((a, b) -> (a[1] + b[1], (a[2]..., b[2]...)), tuple_eachindex(x)) do i
        (x[i]^2 + y[i]^2, (x[i]^3 + y[i]^3,))
    end
end

test_custom(x, y) # (0.71, (0.009000000000000003, 0.09100000000000001, 0.189))

And it is type stable

MethodInstance for test_custom(::Tuple{Float64, Float64, Float64}, ::Tuple{Float64, Float64, Float64})
  from test_custom(x::Tuple, y::Tuple) @ Main ~/.julia/dev/QuantumToolbox/misc/script.jl:46
Arguments
  #self#::Core.Const(Main.test_custom)
  x::Tuple{Float64, Float64, Float64}
  y::Tuple{Float64, Float64, Float64}
Locals
  #18::var"#test_custom##2#test_custom##3"
  #17::var"#test_custom##0#test_custom##1"{Tuple{Float64, Float64, Float64}, Tuple{Float64, Float64, Float64}}
Body::Tuple{Float64, Tuple{Float64, Float64, Float64}}
1 ─       Core.NewvarNode(:(#18))
│         Core.NewvarNode(:(#17))
│   %3  = Main.:(==)::Core.Const(==)
│   %4  = Main.length::Core.Const(length)
│   %5  = (%4)(x)::Core.Const(3)
│   %6  = Main.length::Core.Const(length)
│   %7  = (%6)(y)::Core.Const(3)
│   %8  = (%3)(%5, %7)::Core.Const(true)
└──       goto #3 if not %8
2 ─       goto #4
3 ─       Core.Const(:(Main.throw))
│         Core.Const(:(Main.ArgumentError))
│         Core.Const(:((%12)("Input tuples must have the same length.")))
└──       Core.Const(:((%11)(%13)))
4 ┄ %15 = Main.tmapreduce::Core.Const(Main.tmapreduce)
│   %16 = Main.:(var"#test_custom##0#test_custom##1")::Core.Const(var"#test_custom##0#test_custom##1")
│   %17 = Core._typeof_captured_variable(x)::Core.Const(Tuple{Float64, Float64, Float64})
│   %18 = Core._typeof_captured_variable(y)::Core.Const(Tuple{Float64, Float64, Float64})
│   %19 = Core.apply_type(%16, %17, %18)::Core.Const(var"#test_custom##0#test_custom##1"{Tuple{Float64, Float64, Float64}, Tuple{Float64, Float64, Float64}})
│         (#17 = %new(%19, x, y))
│   %21 = #17::var"#test_custom##0#test_custom##1"{Tuple{Float64, Float64, Float64}, Tuple{Float64, Float64, Float64}}
│   %22 = Main.:(var"#test_custom##2#test_custom##3")::Core.Const(var"#test_custom##2#test_custom##3")
│         (#18 = %new(%22))
│   %24 = #18::Core.Const(var"#test_custom##2#test_custom##3"())
│   %25 = Main.tuple_eachindex::Core.Const(Main.tuple_eachindex)
│   %26 = (%25)(x)::Core.Const((1, 2, 3))
│   %27 = (%15)(%21, %24, %26)::Tuple{Float64, Tuple{Float64, Float64, Float64}}
└──       return %27

I was wondering if there are Base Julia functions that do this. I mean, this should be supported by the Julia Base library I would say.

2

Apologies for not answering the question but I would like to learn, which part of the error message hints at the instability?

3

Here is a type-stable version that first merges x and y to a single tuple.

function test_merge(x::Tuple, y::Tuple)
    length(x) == length(y) || throw(ArgumentError("Input tuples must have the same length."))
    z = map(tuple, x, y)
    mapreduce((a, b) -> (a[1] + b[1], (a[2]..., b[2]...)), z) do (xi, yi)
        (xi^2 + yi^2, (xi^3 + yi^3,))
    end
end

julia> @code_typed test_merge(x, y)
CodeInfo(
[...]
) => Tuple{Float64, Tuple{Float64, Float64, Float64}}

I first tried to use mapreduce with both x and y as arguments,

mapreduce((a, b) -> (a[1] + b[1], (a[2]..., b[2]...)), x, y) do xi, yi

but that was (to my surprise) not type-stable.

By the way, if the tuple b[2] always has a single element, then there is no need to form that tuple.

4

This seems to be a problem with the implementation of mapreduce for tuples. The “naive” version

function test_merge(x, y)
    @assert length(x) == length(y)
    mapreduce((a, b) -> (a[1] + b[1], (a[2]..., b[2]...)), x, y) do xi, yi
        (xi^2 + yi^2, (xi^3 + yi^3,))
    end
end

is type-stable if x and y are both SVectors from StaticArrays.jl (or FixedVectors from SmallCollections.jl).

5

It’s not very obvious with the way Discourse colours the output, but the Anys in

are the problem. When you use

julia> @code_warntype test_standard(x, y)

in the REPL, this is much clearer.