# Reducing allocations in a stencil function

**URL:** <https://discourse.julialang.org/t/reducing-allocations-in-a-stencil-function/87733>\
**Category:** Performance\
**Created:** [September 24, 2022, 2:29am UTC](https://discourse.julialang.org/t/reducing-allocations-in-a-stencil-function/87733 "2022-09-24T02:29:55Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![maxkapur](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maxkapur/32/21208_2.png) [@maxkapur](https://discourse.julialang.org/u/maxkapur)\
**Post date:** [September 24, 2022, 2:29am UTC](https://discourse.julialang.org/t/reducing-allocations-in-a-stencil-function/87733/1 "2022-09-24T02:29:55Z")

</div>

I’m writing a function that takes a matrix `X` and returns a matrix `Y` of the same dimensions such that `Y[i, j]` is the sum of the values in the eight cells surrounding `X[i, j]`. This function is the “number of neighbors” function from Conway’s Game of Life. I seem to recall that this kind of function is also used in image processing, where they call it a “stencil” function.

For the edge cases, I will use “torus topology,” i.e. you wrap indices around to the opposite end of the grid. For example, the cell immediately “above” `X[1, 5]` is `X[end, 5]`.

I have written two implementations of this function:

- `neighborhood_density_loop(X)` loops manually through all the indices in a way that tries to improve memory localization
- `neighborhood_density_offset(X)` uses array slicing in a way that tries to lean on Julia’s optimized built-in array summation

In both cases, the code is pretty verbose because you have to manually code the edge and corner cases (see below). I prefer the `neighborhood_density_offset` implementation because the code is much more legible and self-explanatory. However, `neighborhood_density_loop` demonstrated much better performance in my benchmark:

```julia
VERSION = v"1.8.1"
[ Info: Test OK
neighborhood_density_loop 234.064 ns (1 allocation: 448 bytes)
neighborhood_density_offset 3.876 μs (41 allocations: 8.75 KiB)

```

**Is it possible to improve the code of `neighborhood_density_offset` to reduce the number of allocations and achieve similar performance to the for loop?** Following a common performance tip, I’ve used `@view` on all the right-hand array slices, but there is still a lot of allocation going on.

Here is my code:

```julia
using BenchmarkTools

"""
Compute the neighborhood density for each cell using a for loop.
"""
function neighborhood_density_loop(X::Matrix{T})::Matrix{Float64} where T<:Number
    m, n = size(X)
    Y = zeros(Float64, m, n)

    for j in 2:n-1
        # Inner part
        for i in 2:m-1
            Y[i, j] = begin
                X[i-1, j-1] + 
                X[i , j-1] + 
                X[i+1, j-1] + 
                X[i-1, j] + 
                X[i+1, j] + 
                X[i-1, j+1] + 
                X[i , j+1] + 
                X[i+1, j+1]
            end
        end

        # Top edge
        Y[1, j] = begin
            X[m, j-1] + 
            X[1, j-1] + 
            X[2, j-1] + 
            X[m, j] + 
            X[2, j] + 
            X[m, j+1] + 
            X[1, j+1] + 
            X[2, j+1]
        end

        # Bottom edge
        Y[m, j] = begin
            X[m-1, j-1] + 
            X[m , j-1] + 
            X[1 , j-1] + 
            X[m-1, j] + 
            X[1 , j] + 
            X[m-1, j+1] + 
            X[m , j+1] + 
            X[1 , j+1]
        end
    end

    for i in 2:m-1
        # Left edge
        Y[i, 1] = begin
            X[i-1, n] + 
            X[i , n] + 
            X[i+1, n] + 
            X[i-1, 1] + 
            X[i+1, 1] + 
            X[i-1, 2] + 
            X[i , 2] + 
            X[i+1, 2]
        end
        
        # Right edge
        Y[i, n] = begin
            X[i-1, n-1] + 
            X[i , n-1] + 
            X[i+1, n-1] + 
            X[i-1, n] + 
            X[i+1, n] + 
            X[i-1, 1] + 
            X[i , 1] + 
            X[i+1, 1]
        end
    end 
    
    # Corners
    Y[1, 1] = begin
        X[m, n] + 
        X[1, n] + 
        X[2, n] + 
        X[m, 1] + 
        X[2, 1] + 
        X[m, 2] + 
        X[1, 2] + 
        X[2, 2]
    end

    Y[1, n] = begin
        X[m, n-1] + 
        X[1, n-1] + 
        X[2, n-1] + 
        X[m, n] + 
        X[2, n] + 
        X[m, 1] + 
        X[1, 1] + 
        X[2, 1]
    end

    Y[m, 1] = begin
        X[m-1, n] + 
        X[m , n] + 
        X[1 , n] + 
        X[m-1, 1] + 
        X[1 , 1] + 
        X[m-1, 2] + 
        X[m , 2] + 
        X[1 , 2]
    end

    Y[m, n] = begin
        X[m-1, n-1] + 
        X[m , n-1] + 
        X[1 , n-1] + 
        X[m-1, n] + 
        X[1 , n] + 
        X[m-1, 1] + 
        X[m , 1] + 
        X[1 , 1]
    end

    return Y
end

"""
Compute the neighborhood density for each cell using offset array slices.
"""
function neighborhood_density_offset(X::Matrix{T})::Matrix{Float64} where T<:Number
    m, n = size(X)
    Y = zeros(Float64, m, n)

    # Right neighbors
    Y[:, 1:n-1] += @view X[:, 2:n]
    # Left neighbors
    Y[:, 2:n] += @view X[:, 1:n-1]
    # Lower neighbors
    Y[1:m-1, :] += @view X[2:m, :]
    # Upper neighbors
    Y[2:m, :] += @view X[1:m-1, :]
    
    # Diagonal neighbors
    Y[1:m-1, 1:n-1] += @view X[2:m, 2:n]
    Y[1:m-1, 2:n] += @view X[2:m, 1:n-1]
    Y[2:m, 1:n-1] += @view X[1:m-1, 2:n]
    Y[2:m, 2:n] += @view X[1:m-1, 1:n-1]

    # Right, left neighbors of right and left edges
    Y[:, n] += @view X[:, 1]
    Y[:, 1] += @view X[:, n]
    # Lower, upper neighbors of bottom and top edges
    Y[m, :] += @view X[1, :]
    Y[1, :] += @view X[m, :]

    # Diagonal neighbors along right edge
    Y[1:m-1, n] += @view X[2:m, 1]
    Y[2:m, n] += @view X[1:m-1, 1]
    # Diagonal neighbors along left edge
    Y[1:m-1, 1] += @view X[2:m, n]
    Y[2:m, 1] += @view X[1:m-1, n]
    # Diagonal neighbors along bottom edge
    Y[m, 1:n-1] += @view X[1, 2:n]
    Y[m, 2:n] += @view X[1, 1:n-1]
    # Diagonal neighbors along top edge
    Y[1, 1:n-1] += @view X[m, 2:n]
    Y[1, 2:n] += @view X[m, 1:n-1]

    # Diagonals pointing outward from each corner
    Y[1, 1] += X[m, n]
    Y[1, n] += X[m, 1]
    Y[m, 1] += X[1, n]
    Y[m, n] += X[1, 1]

    return Y
end

function test(samples=500, m=6, n=8)
    for _ in 1:samples
        X = rand(m, n)
        @assert all(
            neighborhood_density_loop(X) .≈ neighborhood_density_offset(X)
        )
    end

    @info "Test OK"
end

function benchmark(m=6, n=8)
    print("neighborhood_density_loop")
    @btime neighborhood_density_loop(X) setup=(X=rand($m, $n))
    print("neighborhood_density_offset")
    @btime neighborhood_density_offset(X) setup=(X=rand($m, $n))
    nothing
end

function main()
    @show VERSION

    test()
    benchmark()
end

main()

```

Nitpicky comments and general review are also welcome.

---

<div class="post-metadata">

**Author:** ![Oscar\_Smith](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oscar_smith/32/25343_2.png) [@Oscar\_Smith](https://discourse.julialang.org/u/Oscar_Smith)\
**Post date:** [September 24, 2022, 2:42am UTC](https://discourse.julialang.org/t/reducing-allocations-in-a-stencil-function/87733/2 "2022-09-24T02:42:50Z")

</div>

You want to use `.+=` instead of `+=` when adding vectors.

---

<div class="post-metadata">

**Author:** ![Elrod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elrod/32/22461_2.png) [@Elrod](https://discourse.julialang.org/u/Elrod)\
**Post date:** [September 24, 2022, 2:42am UTC](https://discourse.julialang.org/t/reducing-allocations-in-a-stencil-function/87733/3 "2022-09-24T02:42:55Z")

</div>

Use `@views` in front of the offset function. Or, at least in front of every line like

```julia
Y[m, 1:n-1] += @view X[1, 2:n]

```

because this line is equivalent to

```julia
Y[m, 1:n-1] = Y[m, 1:n-1] + @view X[1, 2:n]

```

meaning you’re slicing without a view.

Also, as Oscar said, use `@.` or just add `.`s manually.  
That is, you want something like

```julia
@views Y[m, 1:n-1] .= Y[m, 1:n-1] .+ X[1, 2:n]

```

---

<div class="post-metadata">

**Author:** ![maxkapur](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maxkapur/32/21208_2.png) [@maxkapur](https://discourse.julialang.org/u/maxkapur)\
**Post date:** [September 24, 2022, 6:40am UTC](https://discourse.julialang.org/t/reducing-allocations-in-a-stencil-function/87733/4 "2022-09-24T06:40:00Z")

</div>

Thank you! I changed it to this pattern:

```julia
@views Y[m, 1:n-1] .+= X[1, 2:n]

```

The new results are much closer:

```julia
VERSION = v"1.8.1"
[ Info: Test OK
neighborhood_density_loop 235.694 ns (1 allocation: 448 bytes)
neighborhood_density_offset 963.833 ns (1 allocation: 448 bytes)

```

Is it safe, now, to attribute the 3x speedup with the loop version to memory layout, or are there further optimizations to be made to the slice version?

---

<div class="post-metadata">

**Author:** ![baggepinnen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/baggepinnen/32/693_2.png) [@baggepinnen](https://discourse.julialang.org/u/baggepinnen)\
**Post date:** [September 24, 2022, 7:04am UTC](https://discourse.julialang.org/t/reducing-allocations-in-a-stencil-function/87733/5 "2022-09-24T07:04:45Z")

</div>

Adding these annotations to the loop version improves the timing for me from

```julia
    @inbounds @fastmath for j in 2:n-1
        # Inner part
        @simd for i in 2:m-1

```

neighborhood\_density\_loop 99.433 ns (1 allocation: 448 bytes)  
to  
neighborhood\_density\_loop 77.053 ns (1 allocation: 448 bytes)

---

<div class="post-metadata">

**Author:** ![baggepinnen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/baggepinnen/32/693_2.png) [@baggepinnen](https://discourse.julialang.org/u/baggepinnen)\
**Post date:** [September 24, 2022, 7:17am UTC](https://discourse.julialang.org/t/reducing-allocations-in-a-stencil-function/87733/6 "2022-09-24T07:17:29Z")

</div>

For fun, here’s a version that does the allocations of `Y` outside of the benchmarked functions

```julia
using BenchmarkTools

"""
Compute the neighborhood density for each cell using a for loop.
"""
function neighborhood_density_loop(Y, X::Matrix{T}) where T<:Number
    m, n = size(X)
    

    @inbounds @fastmath for j in 2:n-1
        # Inner part
        @simd for i in 2:m-1
            Y[i, j] = begin
                X[i-1, j-1] + 
                X[i , j-1] + 
                X[i+1, j-1] + 
                X[i-1, j] + 
                X[i+1, j] + 
                X[i-1, j+1] + 
                X[i , j+1] + 
                X[i+1, j+1]
            end
        end

        # Top edge
        Y[1, j] = begin
            X[m, j-1] + 
            X[1, j-1] + 
            X[2, j-1] + 
            X[m, j] + 
            X[2, j] + 
            X[m, j+1] + 
            X[1, j+1] + 
            X[2, j+1]
        end

        # Bottom edge
        Y[m, j] = begin
            X[m-1, j-1] + 
            X[m , j-1] + 
            X[1 , j-1] + 
            X[m-1, j] + 
            X[1 , j] + 
            X[m-1, j+1] + 
            X[m , j+1] + 
            X[1 , j+1]
        end
    end

    for i in 2:m-1
        # Left edge
        Y[i, 1] = begin
            X[i-1, n] + 
            X[i , n] + 
            X[i+1, n] + 
            X[i-1, 1] + 
            X[i+1, 1] + 
            X[i-1, 2] + 
            X[i , 2] + 
            X[i+1, 2]
        end
        
        # Right edge
        Y[i, n] = begin
            X[i-1, n-1] + 
            X[i , n-1] + 
            X[i+1, n-1] + 
            X[i-1, n] + 
            X[i+1, n] + 
            X[i-1, 1] + 
            X[i , 1] + 
            X[i+1, 1]
        end
    end 
    
    # Corners
    Y[1, 1] = begin
        X[m, n] + 
        X[1, n] + 
        X[2, n] + 
        X[m, 1] + 
        X[2, 1] + 
        X[m, 2] + 
        X[1, 2] + 
        X[2, 2]
    end

    Y[1, n] = begin
        X[m, n-1] + 
        X[1, n-1] + 
        X[2, n-1] + 
        X[m, n] + 
        X[2, n] + 
        X[m, 1] + 
        X[1, 1] + 
        X[2, 1]
    end

    Y[m, 1] = begin
        X[m-1, n] + 
        X[m , n] + 
        X[1 , n] + 
        X[m-1, 1] + 
        X[1 , 1] + 
        X[m-1, 2] + 
        X[m , 2] + 
        X[1 , 2]
    end

    Y[m, n] = begin
        X[m-1, n-1] + 
        X[m , n-1] + 
        X[1 , n-1] + 
        X[m-1, n] + 
        X[1 , n] + 
        X[m-1, 1] + 
        X[m , 1] + 
        X[1 , 1]
    end

    nothing
end

"""
Compute the neighborhood density for each cell using offset array slices.
"""
@views @fastmath function neighborhood_density_offset(Y, X::Matrix{T}) where T<:Number
    m, n = size(X)

    # Right neighbors
    @. Y[:, 1:n-1] += X[:, 2:n]
    # Left neighbors
    @. Y[:, 2:n] += X[:, 1:n-1]
    # Lower neighbors
    @. Y[1:m-1, :] += X[2:m, :]
    # Upper neighbors
    @. Y[2:m, :] += X[1:m-1, :]
    
    # Diagonal neighbors
    @. Y[1:m-1, 1:n-1] += X[2:m, 2:n]
    @. Y[1:m-1, 2:n] += X[2:m, 1:n-1]
    @. Y[2:m, 1:n-1] += X[1:m-1, 2:n]
    @. Y[2:m, 2:n] += X[1:m-1, 1:n-1]

    # Right, left neighbors of right and left edges
    @. Y[:, n] += X[:, 1]
    @. Y[:, 1] += X[:, n]
    # Lower, upper neighbors of bottom and top edges
    @. Y[m, :] += X[1, :]
    @. Y[1, :] += X[m, :]

    # Diagonal neighbors along right edge
    @. Y[1:m-1, n] += X[2:m, 1]
    @. Y[2:m, n] += X[1:m-1, 1]
    # Diagonal neighbors along left edge
    @. Y[1:m-1, 1] += X[2:m, n]
    @. Y[2:m, 1] += X[1:m-1, n]
    # Diagonal neighbors along bottom edge
    @. Y[m, 1:n-1] += X[1, 2:n]
    @. Y[m, 2:n] += X[1, 1:n-1]
    # Diagonal neighbors along top edge
    @. Y[1, 1:n-1] += X[m, 2:n]
    @. Y[1, 2:n] += X[m, 1:n-1]

    # Diagonals pointing outward from each corner
    Y[1, 1] += X[m, n]
    Y[1, n] += X[m, 1]
    Y[m, 1] += X[1, n]
    Y[m, n] += X[1, 1]

    nothing
end

function test(samples=500, m=6, n=8)
    for _ in 1:samples
        X = rand(Float64, m, n)
        Y1 = zeros(Float64, m, n)
        Y2 = zeros(Float64, m, n)
        neighborhood_density_loop(Y1, X)
        neighborhood_density_offset(Y2, X)
        @assert Y1 ≈ Y2
    end

    @info "Test OK"
end

function benchmark(m=6, n=8)
    print("neighborhood_density_loop")
    @btime neighborhood_density_loop(Y, X) setup=(X=rand(Float64, $m, $n); Y=rand(Float64, $m, $n))
    print("neighborhood_density_offset")
    @btime neighborhood_density_offset(Y, X) setup=(X=rand(Float64, $m, $n); Y=rand(Float64, $m, $n))
    nothing
end

function main()
    @show VERSION

    test()
    benchmark()
end

main()

```

```julia
VERSION = v"1.8.1"
[ Info: Test OK
neighborhood_density_loop 55.675 ns (0 allocations: 0 bytes)
neighborhood_density_offset 350.665 ns (0 allocations: 0 bytes)

```

---

<div class="post-metadata">

**Author:** ![Raf](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/raf/32/3383_2.png) [@Raf](https://discourse.julialang.org/u/Raf)\
**Post date:** [September 24, 2022, 12:20pm UTC](https://discourse.julialang.org/t/reducing-allocations-in-a-stencil-function/87733/7 "2022-09-24T12:20:38Z")

</div>

The neighborhood handling in DynamicGrids.jl is pretty fast for this, and lets you construct layered custom neighborhood shapes you can just `map` over instead of hand coding. It uses a @generated neighborhood of arbitrary shape and size known at compile time. You would use `Moore{1}` for your use case. Using a view is much slower than an `@generated` custom neighborhood, probably because of bad SIMD (@elrod would know better than me here). But you can do much better by explicitly extracting the neighborhood as a tuple than you can by broadcasting over a view. This is also true on GPUs, where I saw a 10x improvment using generated `Tuple` neighborhoods rather than views.

But these tools aren’t that easily useable, so I’m in the (admittedly slow) process of abstracting them into a Neighborhoods.jl package (currently a submodule of DynamicGrids.jl).

If you actually need game of life style simulations, DynamicGrids.jl is pretty good for this as-is, and will do your torus with `boundary=Wrap()`.