9i. Lazy Broadcasting and Loop Fusion

Introduction

This section continues the analysis of lazy and eager operations as a means of reducing memory allocations. The focus now shifts to broadcasting operations, which strike a balance between code readability and performance.

Central to the upcoming discussion is the eager default behavior of broadcasting in Julia. This means that broadcasted operations compute their outputs immediately upon execution, inevitably leading to memory allocation when applied to objects such as vectors. This characteristic becomes particularly crucial in scenarios involving multiple intermediate broadcasted operations, resulting in multiple temporary memory allocations that could potentially be avoided.

To reduce allocations in such scenarios, next we present various approaches. We'll start highlighting the notion of loop fusion, which allows multiple broadcasting operations to be combined into a more efficient single operation. After this, we'll explore the LazyArrays package, which evaluates broadcasting operations in a lazy manner. The technique is particularly useful for reductions requiring intermediate transformations, entirely circumventing memory allocations for these intermediates.

Broadcasting: Loop Fusion

While eager broadcasting makes results readily available, their outputs may not be important in themselves. Instead, they could represent intermediate steps in a larger computation, with these outputs eventually being passed as inputs to subsequent operations.

In the following, we address scenarios like this, where broadcasting is employed for intermediate results. The first approach we explore leverages a technique called loop fusion, which combines multiple broadcasting operations into a single loop. By doing so, the compiler can perform all operations in a single pass over the data. This not only eliminates the creation of multiple intermediate vectors, but also provides the compiler with a holistic view of the operations, thus allowing for further optimizations.

Remarkably, when all broadcasting operations are nested within a single operation, the compiler automatically implements loop fusion. However, for complex expressions, writing a single lengthy expression can be impractical. To overcome this limitation, we'll show a method that enables us to break down operations into partial calculations, while still preserving loop fusion. This approach is based on the lazy design of functions definitions, making it possible to delay operations until their eventual combination.

x        = rand(100)

foo(x)   = x .* 2 .+ x .* 3     # or @. x * 2 + x * 3

  35.361 ns (1 allocation: 896 bytes)

x        = rand(100)

function foo(x)
    a      = x .* 2
    b      = x .* 3
    
    output = a .+ b
end

  124.420 ns (3 allocations: 2.62 KiB)

x        = rand(100)

term1(a) = a * 2
term2(a) = a * 3

foo(a)   = term1(a) + term2(a)

  34.330 ns (1 allocation: 896 bytes)

Vector Operations Allocate and Break Loop Fusion

A common scenario where loop fusion is prevented is when a single expression combines broadcasting and vector operations. This possibility arises because some vector operations yield similar results to their broadcasting equivalents, making it possible to combine these operations without dimensional mismatches. For instance, we show below that adding two vectors using + produces the same result as summing them element-wise by employing .+.

x        = [1, 2, 3]
y        = [4, 5, 6]

foo(x,y) = x .+ y

3-element Vector{Int64}:
 5
 7
 9

x        = [1, 2, 3]
y        = [4, 5, 6]

foo(x,y) = x + y

3-element Vector{Int64}:
 5
 7
 9

The same occurs with a vector product when one of the operands is a scalar.

x        = [1, 2, 3]
β        = 2

foo(x,β) = x .* β

3-element Vector{Int64}:
 2
 4
 6

x        = [1, 2, 3]
β        = 2

foo(x,β) = x * β

3-element Vector{Int64}:
 2
 4
 6

Omitting Dots Avoids Loop Fusion

Mixing vector operations and broadcasting is problematic for performance. The reason is that each vector operation will allocate memory, in a context that operations aren't fused. In particular, the issue arises when the following conditions are met:

• The final output requires combining multiple operations

• The operations yield the same result whether implemented through broadcasting or a vector operation

• We mix broadcasting and vector implementations, by omitting the inclusion of some dots

If all these conditions are satisfied, Julia will partition the operations and compute each separately. The consequence of this is the emergence of multiple temporary vectors, with each separately allocating memory.

The following example illustrates this possibility in the extreme case where all broadcasting dots . are omitted. It demonstrates that vector operations aren't fused, even when expressed in a single operation. Moreover, it establishes that vector operations are similar to obtaining the final result by separately calculating each operation.

x = rand(100)

foo(x)  = x * 2 + x * 3

  129.269 ns (3 allocations: 2.62 KiB)

x = rand(100)

function foo(x) 
    term1  = x * 2        
    term2  = x * 3
    
    output = term1 + term2
end

  130.798 ns (3 allocations: 2.62 KiB)

While the previous example exclusively consists of vector operations, the same principle applies when mixing broadcasting and non-broadcasting operations. In such cases, loop fusion is partially achieved, with only a subset of operations being internally computed through a single for-loop.

x      = rand(100)

foo(x) = x * 2 .+ x .* 3

  85.034 ns (2 allocations: 1.75 KiB)

x      = rand(100)

function foo(x)
    term1  = x * 2
    
    output = term1 .+ x .*3
end

  85.763 ns (2 allocations: 1.75 KiB)

Overall, the key takeaway from these examples is that guaranteeing loop fusion requires appending a dot to every operator and function to be broadcasted. Note that this can be error-prone, especially in large expressions where a single missing dot can be easily overlooked. Fortunately, there are two alternatives that mitigate this risk.

One option is to prefix the expression with @., as shown in the tab Equivalent 1 below. This ensures that all operators and functions are broadcasted. Alternatively, all operations could be combined into a scalar function, which you eventually broadcast. This is presented in the tab Equivalent 2 below.

x      = rand(100)

foo(x) = x .* 2 .+ x .* 3

  36.456 ns (1 allocation: 896 bytes)

x      = rand(100)

foo(x) = @. x * 2 + x * 3

  36.573 ns (1 allocation: 896 bytes)

x      = rand(100)

foo(a) = a * 2 + a * 3

  34.536 ns (1 allocation: 896 bytes)

When multiple long operations are combined, the need to split operations is inevitable. In this case, we can apply a similar trick as we did before, where we leverage that function definitions are inherently lazy. Specifically, this allows us to achieve loop fusion by defining each operation as a separate scalar function.

x        = rand(100)

term1(a) = a * 2
term2(a) = a * 3

foo(a)   = term1(a) + term2(a)

  35.346 ns (1 allocation: 896 bytes)