9g. 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.

A key feature of broadcasting is its eager default behavior. This means that broadcast operations compute and materialize outputs immediately upon execution. Thus, it inevitably leads to memory allocation when applied to objects such as vectors. Such behavior becomes especially relevant in scenarios with intermediate broadcast operations, whose allocations are potentially avoidable.

To address the associated performance cost, we'll present two strategies. The first one highlights the notion of loop fusion. This enables the combination of multiple broadcasting operations into a single more efficient one. Its relevance lies in minimizing memory allocations, rather than their entire elimination. After this, we'll explore the LazyArrays package, which evaluates broadcasting operations lazily. When reductions require intermediate calculations, this technique can completely bypass memory allocations.

Loop Fusion

When working with long broadcast operations, splitting them into smaller intermediate steps can significantly improve readability and reduce the likelihood of errors. However, this approach comes at the cost of separately allocating each broadcasting operation.

To mitigate unnecessary memory allocations, it's essential to preserve loop fusion: a compiler optimization that merges multiple element-wise operations into a single loop over the data. Thus, with loop fusion enabled, 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 unlocking further optimizations.

Loop fusion is applied automatically when all operations are expressed as a single broadcast operation. Yet, as indicated before, writing a single lengthy monolithic expression is inconvenient for complex computations. To overcome this limitation, we can rely on the lazy design of function definitions. Below, we show how this approach can break down an operation into partial calculations, while still preserving loop fusion.

x        = rand(100)

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

  82.627 ns (6 allocations: 2.719 KiB)

x        = rand(100)

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

  26.787 ns (2 allocations: 928 bytes)

x        = rand(100)

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

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

  27.222 ns (2 allocations: 928 bytes)

Even with a single simple operation, certain coding patterns can inadvertently break loop fusion. When this occurs, Julia will internally fall back to evaluating sub-expressions in separate for-loops, which are eventually combined to present the final output. In the following, we present two of these cases. The key practical takeaway is that you should broadcast via the macro @. when possible, rather than manually adding dots to each operator and function within an expression.

Mixing Broadcasting with Vector Operations Breaks Loop Fusion

To understand the issue, there are two facts to consider. The first one is that some vector operations produce an equivalent result to their broadcast counterparts. For instance, adding two vectors with + yields the same result as summing them element-wise with .+.

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

Another example is with product operations.

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

The second fact to consider is that vector operations don't participate in loop fusion. Thus, even if all these operations were combined into a single expression, Julia will evaluate each sub-expression separately, allocating temporary vectors at every step.

x = rand(100)

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

  84.459 ns (6 allocations: 2.719 KiB)

x = rand(100)

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

  87.042 ns (6 allocations: 2.719 KiB)

Putting both facts together, mixing broadcasting with vector operations in the same expression may yield the correct result, but only achieve loop fusion partially. This means Julia will only fuse contiguous broadcast segments, internally partitioning the computation into separate for-loops when a vector operation is encountered. Each of these for-loops will produce a temporary intermediate vector.

x      = rand(100)

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

  67.221 ns (4 allocations: 1.812 KiB)

x      = rand(100)

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

  64.923 ns (4 allocations: 1.812 KiB)

This behavior leads to a clear and actionable guideline: for full loop fusion, every operator and function in the expression must be explicitly broadcast. Yet manually adding dots throughout a large expression is both tedious and error-prone: one missing dot is enough to reintroduce unnecessary allocations. There are two coding practices that mitigate this risk.

One option is to broadcast via the macro @., as shown below in the tab "Equivalent 1". By design, this adds a dot to all operators and functions within an expression, guaranteeing loop fusion. An alternative is to express the operation through a scalar function, which we then broadcast. This is presented below in the tab "Equivalent 2".

x      = rand(100)

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

  28.810 ns (2 allocations: 928 bytes)

x      = rand(100)

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

  28.215 ns (2 allocations: 928 bytes)

x      = rand(100)

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

  27.936 ns (2 allocations: 928 bytes)