pic pic
Personal
Website

8g. Type Stability with Higher-Order Functions

Martin Alfaro
PhD in Economics
Code Script
This section's scripts are available here, under the name allCode.jl. They've been tested under Julia 1.11.9.

Introduction

Functions in Julia are first-class objects, a concept also referred to as first-class citizens. This means that functions can be handled just like any other variable: we can define vectors of functions, have functions whose outputs are other functions, and compose them in flexible ways. In particular, the property makes it possible to define higher-order functions, which are functions that take another function as an argument. We've already worked with several of them, often in the form of anonymous functions passed as function arguments. A familiar example is map(<function>, <collection>), which applies <function> to every element of <collection>.

Our goal in this section is to understand when higher-order functions are type-stable. As we'll discover, these functions present some challenges in this regard.

Remark
Throughout the explanations, we'll often refer to the function passed as an argument as the callback function.

The Issue

In Julia, each function has its own unique concrete type. Moreover, each of these concrete types is a subtype of an abstract type called Function, which serves as an umbrella for all possible functions defined in Julia. This design creates challenges to specialize the computation method of higher-order functions. In particular, if Julia were to generate a distinct specialized method for every possible callback function, the number of methods could grow explosively.

To address this issue, Julia adopts a conservative strategy, often choosing not to specialize the methods of higher-order functions. When specialization is skipped, the performance can drop sharply, sometimes resembling the overhead of executing code in global scope.

It's therefore important to recognize when specialization is inhibited and to monitor its consequences. If you notice that performance is severely impaired, there are still ways to enforce specialization. In this section, we'll explore these strategies.

An Example of No Specialization

To understand when higher-order functions fail to specialize, consider a simple example: summing the transformed elements of a vector x. The only constraint we impose is that the transformation itself must remain generic, allowing us to supply different functions for the transformation.

Specifically, we define a function foo that broadcasts over x to transform its elements through some function f. To illustrate how foo works, we call it with the function abs, which provides absolute values.

x         = rand(100)

foo(f, x) = f.(x)
Output in REPL
julia>
@code_warntype foo(abs,x)

In the example, type instability arises because specialization is avoided when a callback function isn't explicitly called inside the function body. Specifically, the function f only enters foo as an argument of broadcasting, but there's no explicit line calling f. As a result, the compiler doesn't specialize foo for the specific function passed in.

Despite this, the macro @code_warntype doesn't report any type-stability issues. This happens because @code_warntype evaluates type stability under the assumption that specialization was attempted. In our case, that assumption simply doesn't hold, rendering @code_warntype of no use for diagnosing the problem.

To gather indirect evidence about the lack of specialization, we can compare the performance of the original foo function with a modified version that explicitly calls f.

x = rand(100)

function foo(f, x)
    
    f.(x)
end
Output in REPL
julia>
foo(abs, x)
100-element Vector{Float64}:
 0.9063
 0.443494
 â‹®
 0.121148
 0.20453

julia>
@btime foo(abs, $x)
  817.886 ns (12 allocations: 1.250 KiB)

x = rand(100)

function foo(f, x)
    f(1)                # irrelevant computation to force specialization
    f.(x)
end
Output in REPL
julia>
foo(abs, x)
100-element Vector{Float64}:
 0.9063
 0.443494
 â‹®
 0.121148
 0.20453

julia>
@btime foo(abs, $x)
  25.916 ns (2 allocations: 928 bytes)

The addition of f(1) forces the compiler to specialize the function, making the function call type stable. This is noticed through a significant decrease in execution time and memory allocations. Excessive allocations are often indicative of type instability.

Forcing Specialization

Warning!
Exercise caution when forcing specialization. Overly aggressive specialization can degrade performance severely, explaining why Julia's default approach is deliberately conservative. In particular, you should avoid specialization when your script repeatedly calls a higher-order function with many unique functions. [note] For discussions about the issue of excessive specialization, see here and here.

Explicitly calling the callback function to force specialization isn’t ideal, as it introduces an unnecessary computation. Fortunately, alternative solutions exist. One of them is to type-annotate f, thus providing Julia with a hint to specialize code for that function.

Another solution involves wrapping the function in a tuple, before passing it as an argument. This ensures the identification of the function's type, as tuples define a concrete type for each of their elements.

Below, we outline both approaches.

x     = rand(100)


function foo(f, x)    
    f.(x)
end
Output in REPL
julia>
foo(abs, x)
100-element Vector{Float64}:
 0.9063
 0.443494
 â‹®
 0.121148
 0.20453

julia>
@btime foo(abs, $x)
  807.353 ns (12 allocations: 1.250 KiB)

x     = rand(100)


function foo(f::F, x) where F
    f.(x)
end
Output in REPL
julia>
foo(abs, x)
100-element Vector{Float64}:
 0.9063
 0.443494
 â‹®
 0.121148
 0.20453

julia>
@btime foo(abs, $x)
  25.891 ns (2 allocations: 928 bytes)

x     = rand(100)
f_tup = (abs,)

function foo(f_tup, x)
    f_tup[1].(x)    
end
Output in REPL
julia>
foo(f_tup, x)
100-element Vector{Float64}:
 0.9063
 0.443494
 â‹®
 0.121148
 0.20453

julia>
@btime foo($f_tup, $x)
  26.825 ns (2 allocations: 928 bytes)