Multiple Dispatch Misunderstanding?

General Usage
1

I’m trying to use multiple dispatch to let me use different data structures as a function input without having to change too much in a large codebase, but what I’ve come up with so far is significantly slower (~30%) than having just one definition of my data retrieval function.

function make_model(modeltype)
    retrieve_A2 = if modeltype == 1
        u::Array->u[:,1:10]
    elseif modeltype == 2
        u::Array->u[:,11:20]
    else
        error("This is not a valid modeltype")
    end
    
    retrieve_A(u::Array) = retrieve_A2(u)
    #ComponentArrays are easy and only need one function regardless of modeltype
    retrieve_A(u::ComponentArray) = u.A

    modelfn = function my_model(u)
        A = retrieve_A(u)
        operate_on_A(A)
    end
    return modelfn
end

model = make_model(n)
model_output = solve(model, u0)

Since modelfn is performance critical and called a lot, I wanted to try using multiple dispatch to avoid type checking in the function while still allowing me to swap u0 out for easy comparison of performance and features. Am I just pushing an if/else type check to a lower level where I can’t see it easily anymore, or am I misunderstanding something else about how functions should be defined that is causing the 30% slowdown?

2

It’s hard to tell exactly where your trouble here is coming from, but one issue I see is this:

retrieve_A2 ends up as one of two completely different functions (they may look similar to you, but Julia - like most languages - does not care). This may be causing runtime dispatch (which is somewhat slow). Instead, you want it to see it as only one function with different stored arguments.

Try something like

cols = if modeltype == 1
    1:10
elseif modeltype == 2
    11:20
else
    error("This is not a valid modeltype")
end
retrieve_A2 = let cols = cols # `let` helps ensure `cols` doesn't get needlessly boxed here
    u -> u[:, cols]
end

EDIT: my explanation may be slightly off-base here when it comes to the actual issue you’re reporting. It depends on whether retrieve_A2 gets boxed within the retrieve_A method or not. My suggestion might help or might not. A complete minimum working example would make it easier for someone to help you.

3

This MWE below has similar behavior to my full problem and fleshes out the code I initially posted. Your suggestion did not seem to make any changes to my full code unfortunately, and that is also shown in the MWE below.

using ComponentArrays, BenchmarkTools

function operate_on_A(A)
    sum(A.^2)
end

function mysolve(model, initial_conditions)
    for i in 1:100
        model(initial_conditions)
    end
end

function make_model(modeltype)
    retrieve_A2 = if modeltype == 1
        u::Array->@views u[:,1:10]
    elseif modeltype == 2
        u::Array->@views u[:,11:20]
    else
        error("This is not a valid modeltype")
    end
    
    retrieve_A(u::Array) = retrieve_A2(u)
    #ComponentArrays are easy and only need one function regardless of modeltype
    retrieve_A(u::ComponentArray) = @views u.A

    modelfn = function my_model(u)
        A = retrieve_A(u)
        operate_on_A(A)
    end
    return modelfn
end


function make_model2(modeltype)
    cols = if modeltype == 1
        1:10
    elseif modeltype == 2
        11:20
    else
        error("This is not a valid modeltype")
    end
    retrieve_A2 = let cols = cols # `let` helps ensure `cols` doesn't get needlessly boxed here
        u -> @views u[:, cols]
    end
    retrieve_A(u::ComponentArray) = @views u.A
    retrieve_A(u::Array) = retrieve_A2(u)
    
    modelfn = function my_model(u)
        A = retrieve_A(u)
        operate_on_A(A)
    end
    return modelfn
    
end



function modelCA(u::ComponentArray)
    @views A = u.A
    operate_on_A(A)
end

function modelArray(u::Array)
    @views A = u[:,1:10]
    operate_on_A(A)
end



data = rand(3,20);
u0_CA = ComponentArray(A = data[:,1:10], B = data[:,11:20]);
u0_array = deepcopy(data);


dispatch_model = make_model(1);
@btime mysolve(dispatch_model, u0_CA) #4.758 μs (500 allocations: 40.62 KiB)
@btime mysolve(dispatch_model, u0_array) #4.463 μs (400 allocations: 39.06 KiB)

#this seems the same as the original code, so it might be boxed as suggested?
dispatch_model2 = make_model2(1);
@btime mysolve(dispatch_model2, u0_CA) #4.380 μs (500 allocations: 40.62 KiB)
@btime mysolve(dispatch_model2, u0_array) #4.160 μs (400 allocations: 39.06 KiB)

@btime mysolve(modelCA, u0_CA) #2.828 μs (200 allocations: 31.25 KiB)
@btime mysolve(modelArray, u0_array) #2.503 μs (200 allocations: 31.25 KiB)
4 
julia> dump(dispatch_model2)
my_model (function of type var"#my_model#make_model2##7")
  retrieve_A: Core.Box
    contents: retrieve_A (function of type var"#retrieve_A#make_model2##6"{var"#make_model2##4#make_model2##5"{UnitRange{Int64}}})
      retrieve_A2: #make_model2##4 (function of type var"#make_model2##4#make_model2##5"{UnitRange{Int64}})
        #24#cols: UnitRange{Int64}
          start: Int64 1
          stop: Int64 10

so retrieve_A is getting boxed.

This version seems to fix it and match the performance you saw on your single-method functions:

function make_model3(modeltype)
    cols = if modeltype == 1
        1:10
    elseif modeltype == 2
        11:20
    else
        error("This is not a valid modeltype")
    end
    retrieve_A2 = let cols = cols # `let` helps ensure `cols` doesn't get needlessly boxed here
        u -> @views u[:, cols]
    end
    retrieve_A(u::ComponentArray) = @views u.A
    retrieve_A(u::Array) = retrieve_A2(u)
    
    modelfn = let retrieve_A = retrieve_A
        function my_model(u)
            A = retrieve_A(u)
            operate_on_A(A)
        end
    end
    return modelfn
end

I’m not exactly sure why (or whether there was a good reason) it thinks it needed to box it before, but a let block appears to resolve it. You could probably go back to the old non-let retrieve_A2 with this, but I haven’t tried.

5

Awesome, using a let block with a function more similar to the old non-let retrieve_A2 works great and recovered the performance of my non-MWE code. I should probably do more reading on scoping before doing more silly things like this though.

Thanks!

6

Opened Returning an anonymous function that uses a local multi-method function results in boxing · Issue #63706 · JuliaLang/julia · GitHub to see if this behavior should be expected.

7

I’m guessing that multiple method definitions, like assignments, are conservatively boxed in the current implementation. I think an exception can be made here because there isn’t actually a reassignment to a potentially different type, and method tables aren’t supposed to be dynamic during a function call, but the latter is a much bigger issue that hasn’t been fully hashed out.

8

Can anybody explain what is meant here by getting boxed, how dump() shows that and does how the let block help?

9

Boxing is a usually automatic (but sometimes manual) wrapping and referencing of an object. In some other languages with a huge line between reference types and value types (latter something like isbitstype), boxing is often used to wrap value types in reference types. In this case, boxing implements a closure capturing an outer local variable that can be reassigned; dump prints out the structure of the closure, including the internal Core.Box field that implements the captured variable. It’s possible in many cases to statically infer the type of that local variable, but the current implementation uninferrably boxes reassigned variables before that type inference can happen. There are some proposals to improve this, but it’s a much deeper problem and people tend to overestimate how well the captured variable can be inferred (often boils down to why most mainstream languages with generics demand explicit static parameters). The let block trick makes a new local variable that isn’t reassigned, which is often the intent anyway. In this case, a multimethod really doesn’t seem to warrant the same boxing as reassignments.