Avoiding runtime dispatch randomly selecting an algorithm in a loop

General Usage Performance
1

I have a scenario where I’m switching between different methods of one function in a loop that has a long lifetime.

I have a tuple of structs, where I’m selecting one randomly to be put into a function like so:

function somealgo(data, ::Type1)
  #do something to data
end

function somealgo(data, ::Type2)
  #do something else to data
end
function loop(data::SomeStruct, structs::Tuple, func)
  while data.somecondition
    selected = rand([1:length(structs);])
    somealgo(data, structs[selected})
  end
end

This is not so fast due to runtime dispatch. Inlining a switch would be much faster since the lifetime of this loop typically is long and is possible as the ordering, size and types of the structs are always known at compile time. The solution I could think of was using generated functions. What I do here is basically I define all the somealgo functions to dispatch on ::Type{Type1} etc. and to then inline what comes down to a switch statement. This is not super elegant and hard to maintain. I’m currently working with string manipulation (which I find a bit more transparent than expressions sometimes), but the code is still hard to read (also the situation in my real code is a bit more complicated, having some more arguments and types to take into account).

I was wondering if there’s some better way to do it. As far as I’m aware you can essentially do a “light” version of a generated function by writing if statements based on static parameters, where the unused branches get compiled away. Is there something similar that you can do for a for loop, which gets unrolled if the necessary details are known during compile time, or something like that?

2

As you noted, the issue with picking an argument from a mixed-type Tuple is that the argument to somealgo could be one of several (or many) different types. Beyond a few (I don’t remember the exact threshold, but I don’t think it’s bigger than 4), it gives up trying to keep track of the types and simply uses runtime dispatch.

What you want is to generate a function like

function GENERATED_callselected(f, data, structs::Tuple, pick)
    if pick == 1
        return f(data, structs[1])
    end
    if pick == 2
        return f(data, structs[2])
    end
    # repeat up to length(structs)...
    throw(BoundsError(structs, pick))) # invalid pick
end

The compiler doesn’t know what’s at structs[pick] for a variable pick, but it does know what’s at (e.g.) structs[1] so it knows exactly what it’s dealing with in each branch.

Note that if f has more than a few different return types with these different argument types, the caller of callselected will still infer the result as Any and use runtime dispatch later. But it looks like you aren’t using the return value of somealgo anyway, so that probably isn’t a problem.

Here’s a @generated to build this function:

@generated function callselected(f, data, structs::Tuple{Vararg{Any, N}}, pick) where N
	code = Expr(:block)
	for i in 1:N
		push!(code.args, :(if pick == $i; return f(data, structs[$i]); end))
	end
	push!(code.args, :(throw(BoundsError(structs, pick))))
	return code
end

# see that the code is type stable in every branch
@code_warntype callselected(+, 0e0, (1, 2//1, 3f0, 4e0), 2)

So instead of calling somealgo(data, structs[selected]), call callselected(somealgo, data, structs, selected).

1 Like