Intermediate matrix form in `MOLFiniteDifference`

[valentinsulzer]

The code in MOLFiniteDifference is getting pretty messy as we add more functionality for nonlinear laplacian, upwinding, etc.

One way to simplify the code would be to have an intermediate matrix form - which then gets scalarized by the compiler - instead of directly going to the scalar form by looping over all indices. Ideally this would use DerivativeOperator to stop this package diverging into two almost entirely distinct packages (the two examples in the readme). Can DerivativeOperator be used like this? If so,

• can it already do nonlinear laplacian?
• how can this then be scalarized?
• Should boundary conditions remain as they are or use the *bc convolution?

I realize this is a little bit of a step backwards because that is closer to how it was done before the refactoring of MOLFiniteDifference, so worth discussing the pros and cons of each approach here a little bit before changing anything

[ChrisRackauckas]

With array symbolics now merged (https://symbolics.juliasymbolics.org/dev/manual/arrays/) we have a better idea of where this is all going. Pulling in @shashi, @chriselrod, and @YingboMa into this as well.

From https://tutorials.sciml.ai/html/introduction/03-optimizing_diffeq_code.html we know that fusing the operators is a pretty major effect on performance. So we do not want the final code to have sparse matrix multiplications, and instead want to make sure that this all lowers to something that uses LoopVectorization.jl over contiguous loop indices. In order to support huge equations, we probably do not want this to be scalarized. Also, we want to make sure that this use a cache-optimized or at least cache-oblivious loop iteration, which should be the case by using LoopVectorization.jl. It would be good for the code to be GPU friendly as well.

So on that note, it would be good to see a test of LoopVectorization against something like NNLib's conv for the nonlinear Laplacian. I have faith it could match or surpass it, but that's a thing we should benchmark.

Additionally, we need to think about how to lower symbolic maps and such down to something that could be GPU friendly, and what do about the boundary conditions in this case.

It would also be good to note the effect of the matrix-vector for m on the index reduction algorithm. IIUC one would need to write the boundary terms separate from the matvecs in order to allow for the tearing to work, since otherwise that would break the full incidence assumption. If the boundary conditions already have to be written out in index form, I don't know if that saves any simplicity.

@YingboMa can you share the suggested MTK code for the current form?

[chriselrod]

So on that note, it would be good to see a test of LoopVectorization against something like NNLib's conv for the nonlinear Laplacian. I have faith it could match or surpass it, but that's a thing we should benchmark.

Have an example?

using NNlib, Polyester, LoopVectorization, Static
function kernaxes(::DenseConvDims{2,K,C_in, C_out}) where {K,C_in, C_out} # LoopVectorization can take advantage of static size information
  K₁ =  StaticInt(1):StaticInt(K[1])
  K₂ =  StaticInt(1):StaticInt(K[2])
  Cᵢₙ =  StaticInt(1):StaticInt(C_in)
  Cₒᵤₜ = StaticInt(1):StaticInt(C_out)
  (K₁, K₂, Cᵢₙ, Cₒᵤₜ)
end
function convlayer!(
  out::AbstractArray{<:Any,4}, img, kern,
  dcd::DenseConvDims{2, <:Any, <:Any, <:Any, (1, 1), (0, 0, 0, 0), (1, 1), true}
  )
  (K₁, K₂, Cᵢₙ, Cₒᵤₜ) = kernaxes(dcd)
  @batch for d ∈ axes(out,4)
    @turbo for j₁ ∈ axes(out,1), j₂ ∈ axes(out,2), o ∈ Cₒᵤₜ
      s = zero(eltype(out))
      for k₁ ∈ K₁, k₂ ∈ K₂, i ∈ Cᵢₙ
        s += img[j₁ + k₁ - 1, j₂ + k₂ - 1, i, d] * kern[k₁, k₂, i, o]
      end
      out[j₁, j₂, o, d] = s
    end
  end
  out
end
img = rand(Float32, 260, 260, 3, 100);
kern = rand(Float32, 5, 5, 3, 6);
out1 = Array{Float32}(undef, size(img,1)+1-size(kern,1), size(img,2)+1-size(kern,2), size(kern,4), size(img,4));
out2 = similar(out1);
dcd = NNlib.DenseConvDims(img, kern, flipkernel = true);
@time NNlib.conv!(out1, img, kern, dcd);
@time convlayer!(out2, img, kern, dcd);
out1 ≈ out2
@benchmark NNlib.conv!($out1, $img, $kern, $dcd)
@benchmark convlayer!($out1, $img, $kern, $dcd)

LoopVectorization's algorithms scale very badly with the number of loops. 7 works, but takes a minute to compile. So I'm using 6 above, @batching across the seventh. For a CPU-optimized NNlib, consistent threading (e.g., always across batches) would also have the advantage of improving locality, i.e. data will hopefully be able to stay in the same L2 cache from one layer to the next (assuming the data is small enough to fit, as may be the case if training on the CPU).

This is what I get on an 18 core system with AVX512:

julia> @time NNlib.conv!(out1, img, kern, dcd);
  0.802161 seconds (2.06 M allocations: 790.165 MiB, 20.41% gc time)

julia> @time convlayer!(out2, img, kern, dcd);
  9.553409 seconds (21.51 M allocations: 1.149 GiB, 1.06% gc time)

julia> out1 ≈ out2
true

julia> @benchmark NNlib.conv!($out1, $img, $kern, $dcd)
samples: 27; evals/sample: 1; memory estimate: 675.02 MiB; allocs estimate: 195
ns

 (1.8e8 - 1.9e8]  ██████████████████████████████ 24
 (1.9e8 - 2.1e8]  █▎1
 (2.1e8 - 2.3e8]   0
 (2.3e8 - 2.5e8]   0
 (2.5e8 - 2.7e8]  █▎1
 (2.7e8 - 2.9e8]  █▎1

                  Counts

min: 175.652 ms (0.00% GC); mean: 191.870 ms (4.80% GC); median: 184.447 ms (0.44% GC); max: 288.455 ms (36.61% GC).

julia> @benchmark convlayer!($out1, $img, $kern, $dcd)
samples: 944; evals/sample: 1; memory estimate: 32 bytes; allocs estimate: 4
ns

 (5.233e6 - 5.256e6]  ██▍31
 (5.256e6 - 5.278e6]  ██████████████▏193
 (5.278e6 - 5.301e6]  ██████████████████████████████412
 (5.301e6 - 5.324e6]  ████████████████████▊284
 (5.324e6 - 5.347e6]  █▌20
 (5.347e6 - 5.37e6 ]   0
 (5.37e6  - 5.392e6]   0
 (5.392e6 - 5.415e6]  ▏1
 (5.415e6 - 5.438e6]   0
 (5.438e6 - 5.461e6]  ▏1
 (5.461e6 - 5.484e6]  ▎2

                  Counts

min: 5.233 ms (0.00% GC); mean: 5.292 ms (0.00% GC); median: 5.294 ms (0.00% GC); max: 5.484 ms (0.00% GC).

julia> versioninfo()
Julia Version 1.7.0-DEV.1208
Commit 0a133ff22b* (2021-05-30 11:28 UTC)
Platform Info:
  OS: Linux (x86_64-generic-linux)
  CPU: Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-12.0.0 (ORCJIT, cascadelake)
Environment:
  JULIA_NUM_THREADS = 36

LoopVectorization was over 30x faster.

This was also a generic convolution. The Laplacian has more structure you could exploit.
If it's only a 1-d laplacian, it'd probably be harder to do well. The more compute, the less it'll be bottle-necked by memory.

EDIT:
Just realized convlayer! is allocating...
Apparently it isn't making much use of threading -- only about 4x faster, so I guess it's already pretty constrained by memory:

julia> function convlayer_st!(
         out::AbstractArray{<:Any,4}, img, kern,
         dcd::DenseConvDims{2, <:Any, <:Any, <:Any, (1, 1), (0, 0, 0, 0), (1, 1), true}
         )
         (K₁, K₂, Cᵢₙ, Cₒᵤₜ) = kernaxes(dcd)
       for d ∈ axes(out,4)
           @turbo for j₁ ∈ axes(out,1), j₂ ∈ axes(out,2), o ∈ Cₒᵤₜ
             s = zero(eltype(out))
             for k₁ ∈ K₁, k₂ ∈ K₂, i ∈ Cᵢₙ
               s += img[j₁ + k₁ - 1, j₂ + k₂ - 1, i, d] * kern[k₁, k₂, i, o]
             end
             out[j₁, j₂, o, d] = s
           end
           out
         end
       end
convlayer_st! (generic function with 1 method)

julia> @benchmark convlayer_st!($out1, $img, $kern, $dcd)
samples: 180; evals/sample: 1; memory estimate: 0 bytes; allocs estimate: 0
ns

 (2.754e7 - 2.758e7]  ▋2
 (2.758e7 - 2.763e7]   0
 (2.763e7 - 2.768e7]   0
 (2.768e7 - 2.772e7]  █████▍19
 (2.772e7 - 2.777e7]  ██████████▍37
 (2.777e7 - 2.781e7]  ▋2
 (2.781e7 - 2.786e7]  ██▌9
 (2.786e7 - 2.79e7 ]  ██████████████████████████████ 108
 (2.79e7  - 2.795e7]  ▉3

                  Counts

min: 27.539 ms (0.00% GC); mean: 27.824 ms (0.00% GC); median: 27.865 ms (0.00% GC); max: 27.949 ms (0.00% GC).

Apparently Intel's next generation of server chips (Saphire Rapids) will come with up to 64 GiB of HBME2 memory. Maybe it'll only be a select few extremely expensive models, but if there're affordable options, that'd be exciting.
AMD also announced 3d stacking for 192 MiB of L3 cache, which isn't quite enough for problems of this size:

julia> (sizeof(out1) + sizeof(img) + sizeof(kern)) / (1<<20) # bytes of data / (bytes/MiB)
227.36377716064453

The three arrays here are 227 MiB. Making the problem just a little smaller would then let them all fit.
For comparison, the 10980XE I ran these benchmarks on has only 24.75 MiB of L3.
These were 100 260x260 input images (and 100 256x256 outputs), so the real use cases here are hopefully smaller.

[ChrisRackauckas]

Look at the PDE example in https://tutorials.sciml.ai/html/introduction/03-optimizing_diffeq_code.html for an example of it.

  mul!(Ayu,Ay,u)
  mul!(uAx,u,Ax)
  mul!(Ayv,Ay,v)
  mul!(vAx,v,Ax)

is just a cheap way to do 2D convolutions with dense matmuls, but it would be good to transcribe that into conv!

[chriselrod]

I played with it a bit, and haven't been able to get any major improvements over the explicit loop version:

using LinearAlgebra, OrdinaryDiffEq
# Generate the constants
p = (1.0,1.0,1.0,10.0,0.001,100.0) # a,α,ubar,β,D1,D2
N = 100
const Ax = Array(Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1]));
const Ay = copy(Ax);
Ax[2,1] = 2.0;
Ax[end-1,end] = 2.0;
Ay[1,2] = 2.0;
Ay[end,end-1] = 2.0;

function basic_version!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = r[:,:,1]
  v = r[:,:,2]
  Du = D1*(Ay*u + u*Ax)
  Dv = D2*(Ay*v + v*Ax)
  dr[:,:,1] = Du .+ a.*u.*u./v .+ ubar .- α*u
  dr[:,:,2] = Dv .+ a.*u.*u .- β*v
end

a,α,ubar,β,D1,D2 = p;
uss = (ubar+β)/α;
vss = (a/β)*uss^2;
r0 = zeros(100,100,2);
r0[:,:,1] .= uss.+0.1.*rand.();
r0[:,:,2] .= vss;

prob = ODEProblem(basic_version!,r0,(0.0,0.1),p);
@benchmark solve($prob,Tsit5())
# samples: 123; evals/sample: 1; memory estimate: 186.80 MiB; allocs estimate: 4555
# ns

#  (3.0e7 - 5.0e7]  ██████████████████████████████ 121
#  (5.0e7 - 7.0e7]   0
#  (7.0e7 - 8.0e7]   0
#  (8.0e7 - 1.0e8]   0
#  (1.0e8 - 1.2e8]   0
#  (1.2e8 - 1.3e8]  ▎1
#  (1.3e8 - 1.5e8]   0
#  (1.5e8 - 1.6e8]  ▎1

#                   Counts

# min: 34.770 ms (7.03% GC); mean: 40.722 ms (9.64% GC); median: 39.011 ms (5.79% GC); max: 164.571 ms (77.65% GC).
  
function gm2!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = @view r[:,:,1]
  v = @view r[:,:,2]
  du = @view dr[:,:,1]
  dv = @view dr[:,:,2]
  Du = D1*(Ay*u + u*Ax)
  Dv = D2*(Ay*v + v*Ax)
  @. du = Du + a.*u.*u./v + ubar - α*u
  @. dv = Dv + a.*u.*u - β*v
end
prob = ODEProblem(gm2!,r0,(0.0,0.1),p);
@benchmark solve($prob,Tsit5())
# samples: 171; evals/sample: 1; memory estimate: 119.44 MiB; allocs estimate: 2790
# ns

#  (2.68e7 - 2.75e7]  ▋1
#  (2.75e7 - 2.82e7]  ███▏6
#  (2.82e7 - 2.9e7 ]  ███████████████████████▍45
#  (2.9e7  - 2.97e7]  ██████████████████████████████▏58
#  (2.97e7 - 3.04e7]  ████████████████████████████▌55
#  (3.04e7 - 3.11e7]  ██▋5
#  (3.11e7 - 3.18e7]   0
#  (3.18e7 - 3.25e7]   0
#  (3.25e7 - 3.32e7]  ▋1

#                   Counts

# min: 26.824 ms (7.47% GC); mean: 29.393 ms (6.39% GC); median: 29.594 ms (7.36% GC); max: 33.241 ms (27.63% GC).


const Ayu = zeros(N,N);
const uAx = zeros(N,N);
const Du = zeros(N,N);
const Ayv = zeros(N,N);
const vAx = zeros(N,N);
const Dv = zeros(N,N);
function gm3!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = @view r[:,:,1]
  v = @view r[:,:,2]
  du = @view dr[:,:,1]
  dv = @view dr[:,:,2]
  mul!(Ayu,Ay,u)
  mul!(uAx,u,Ax)
  mul!(Ayv,Ay,v)
  mul!(vAx,v,Ax)
  @. Du = D1*(Ayu + uAx)
  @. Dv = D2*(Ayv + vAx)
  @. du = Du + a*u*u./v + ubar - α*u
  @. dv = Dv + a*u*u - β*v
end
prob_gm3 = ODEProblem(gm3!,r0,(0.0,0.1),p);
@benchmark solve($prob_gm3,Tsit5())
# samples: 181; evals/sample: 1; memory estimate: 29.63 MiB; allocs estimate: 438
# ns

#  (2.1e7 - 2.2e7]  ▌1
#  (2.2e7 - 2.3e7]  ▌1
#  (2.3e7 - 2.5e7]   0
#  (2.5e7 - 2.6e7]  █████▋14
#  (2.6e7 - 2.7e7]  ██████████████▎36
#  (2.7e7 - 2.8e7]  ██████████████████████████████ 76
#  (2.8e7 - 3.0e7]  ███████████████▌39
#  (3.0e7 - 3.1e7]  ████10
#  (3.1e7 - 3.2e7]  █▋4

#                   Counts

# min: 20.760 ms (0.00% GC); mean: 27.682 ms (2.07% GC); median: 27.566 ms (0.00% GC); max: 32.108 ms (0.00% GC).

using LoopVectorization
function matmul!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix)
  @tturbo for n ∈ indices((C,B),2), m ∈ indices((C,A),1)
    Cₘₙ = zero(eltype(C))
    for k ∈ indices((A,B),(2,1))
      Cₘₙ += A[m,k]*B[k,n]
    end
    C[m,n] = Cₘₙ
  end
end
function matmuladd!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix)
  @tturbo for n ∈ indices((C,B),2), m ∈ indices((C,A),1)
    Cₘₙ = zero(eltype(C))
    for k ∈ indices((A,B),(2,1))
      Cₘₙ += A[m,k]*B[k,n]
    end
    C[m,n] += Cₘₙ
  end
end

function gm5!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = @view r[:,:,1];
  v = @view r[:,:,2];
  du = @view dr[:,:,1];
  dv = @view dr[:,:,2];
  matmul!(Ayu,Ay,u)
  matmuladd!(Ayu,u,Ax)
  # @turbo @. du = D1*Ayu + a*u*u./v + ubar - α*u
  @turbo for i ∈ eachindex(u)
    du[i] = D1 * Ayu[i] + a * u[i]*u[i] / v[i] + ubar - α*u[i]
  end
  matmul!(Ayu,Ay,v)
  matmuladd!(Ayu,v,Ax)
  # @turbo @. dv = D2*Ayu + a*u*u - β*v
  @turbo for i ∈ eachindex(u)
    dv[i] = D2 * Ayu[i] + a * u[i]*u[i] - β*v[i]
  end
end
prob_gm5 = ODEProblem(gm5!,r0,(0.0,0.1),p);
@benchmark solve($prob_gm5,Tsit5())
# samples: 306; evals/sample: 1; memory estimate: 29.62 MiB; allocs estimate: 433
# ns

#  (1.1e7 - 1.3e7]  ██████████████████████████████ 211
#  (1.3e7 - 1.5e7]  ██▎15
#  (1.5e7 - 1.7e7]  ▍2
#  (1.7e7 - 1.9e7]  ▍2
#  (1.9e7 - 2.1e7]   0
#  (2.1e7 - 2.3e7]   0
#  (2.3e7 - 2.5e7]  ▉6
#  (2.5e7 - 2.8e7]  ███████▊54
#  (2.8e7 - 3.0e7]  █7
#  (3.0e7 - 3.2e7]  █▍9

#                   Counts

# min: 10.999 ms (0.00% GC); mean: 16.355 ms (3.53% GC); median: 12.869 ms (0.00% GC); max: 31.714 ms (6.63% GC).

function matmul!(C::AbstractArray{<:Any,3}, A::AbstractMatrix, B::AbstractArray{<:Any,3})
  @tturbo for n ∈ indices((C,B),2), m ∈ indices((C,A),1), d ∈ 1:2
    Cₘₙ = zero(eltype(C))
    for k ∈ indices((A,B),(2,1))
      Cₘₙ += A[m,k]*B[k,n,d]
    end
    C[m,n,d] = Cₘₙ
  end
end
function matmuladd!(C::AbstractArray{<:Any,3}, A::AbstractArray{<:Any,3}, B::AbstractMatrix)
  @tturbo for n ∈ indices((C,B),2), m ∈ indices((C,A),1), d ∈ 1:2
    Cₘₙ = zero(eltype(C))
    for k ∈ indices((A,B),(2,1))
      Cₘₙ += A[m,k,d]*B[k,n]
    end
    C[m,n,d] += Cₘₙ
  end
end
const Ayuv = similar(Ayu, (size(Ayu,1),size(Ayu,2),2));
function gm6!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  matmul!(Ayuv, Ay, r)
  matmuladd!(Ayuv, r, Ax)
  N² = size(r,1)*size(r,2)
  u = reshape(r, (N²,2))
  du = reshape(dr, (N²,2))
  Ayuvr = reshape(Ayuv, (N²,2))
  @turbo for i ∈ 1:N²
    du[i,1] = D1 * Ayuvr[i,1] + a * u[i,1]*u[i,1] / u[i,2] + ubar - α*u[i,1]
    du[i,2] = D2 * Ayuvr[i,2] + a * u[i,1]*u[i,1] - β*u[i,2]
  end
end
prob_gm6 = ODEProblem(gm6!,r0,(0.0,0.1),p);
@benchmark solve($prob_gm6, Tsit5())
# samples: 310; evals/sample: 1; memory estimate: 29.67 MiB; allocs estimate: 1315
# ns

#  (1.1e7 - 1.3e7]  █████▌37
#  (1.3e7 - 1.4e7]  ██████████████████████████████205
#  (1.4e7 - 1.6e7]   0
#  (1.6e7 - 1.8e7]   0
#  (1.8e7 - 2.0e7]   0
#  (2.0e7 - 2.1e7]   0
#  (2.1e7 - 2.3e7]   0
#  (2.3e7 - 2.5e7]   0
#  (2.5e7 - 2.6e7]  ▎1
#  (2.6e7 - 2.8e7]  █████████▉67

#                   Counts

# min: 11.091 ms (0.00% GC); mean: 16.126 ms (3.49% GC); median: 13.182 ms (0.00% GC); max: 27.934 ms (9.08% GC).
  
pdevec = (1.0,1.0,1.0,10.0,0.001,100.0,N);
function fast_gm!(du,u,p,t)
  a,α,ubar,β,D1,D2,N = p

  @inbounds for j in 2:N-1, i in 2:N-1
    du[i,j,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end

  @inbounds for j in 2:N-1, i in 2:N-1
    du[i,j,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
            a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds for j in 2:N-1
    i = 1
    du[1,j,1] = D1*(2u[i+1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
            a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for j in 2:N-1
    i = 1
    du[1,j,2] = D2*(2u[i+1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
            a*u[i,j,1]^2 - β*u[i,j,2]
  end
  @inbounds for j in 2:N-1
    i = N
    du[end,j,1] = D1*(2u[i-1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
           a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for j in 2:N-1
    i = N
    du[end,j,2] = D2*(2u[i-1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
           a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds for i in 2:N-1
    j = 1
    du[i,1,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for i in 2:N-1
    j = 1
    du[i,1,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
              a*u[i,j,1]^2 - β*u[i,j,2]
  end
  @inbounds for i in 2:N-1
    j = N
    du[i,end,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for i in 2:N-1
    j = N
    du[i,end,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds begin
    i = 1; j = 1
    du[1,1,1] = D1*(2u[i+1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[1,1,2] = D2*(2u[i+1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
              a*u[i,j,1]^2 - β*u[i,j,2]

    i = 1; j = N
    du[1,N,1] = D1*(2u[i+1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[1,N,2] = D2*(2u[i+1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]

    i = N; j = 1
    du[N,1,1] = D1*(2u[i-1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[N,1,2] = D2*(2u[i-1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]

    i = N; j = N
    du[end,end,1] = D1*(2u[i-1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[end,end,2] = D2*(2u[i-1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]
   end
end
prob_devec = ODEProblem(fast_gm!,r0,(0.0,0.1),pdevec);
@benchmark solve($prob_devec,Tsit5())
# samples: 600; evals/sample: 1; memory estimate: 29.63 MiB; allocs estimate: 438
# ns

#  (7.2e6   - 7.74e6 ]  ██████████████████████████████▏337
#  (7.74e6  - 8.27e6 ]  ██▋29
#  (8.27e6  - 8.81e6 ]  ██████▏68
#  (8.81e6  - 9.34e6 ]  ██████▋74
#  (9.34e6  - 9.88e6 ]  ███33
#  (9.88e6  - 1.041e7]  ▍4
#  (1.041e7 - 1.095e7]   0
#  (1.095e7 - 1.149e7]  ██22
#  (1.149e7 - 1.202e7]  ▋7
#  (1.202e7 - 1.256e7]  ██▎24
#  (1.256e7 - 1.309e7]  ▎2

#                   Counts

# min: 7.200 ms (0.00% GC); mean: 8.335 ms (7.27% GC); median: 7.482 ms (0.00% GC); max: 13.092 ms (43.21% GC).

  
function turbo_gm!(du,u,p,t)
  a,α,ubar,β,D1,D2,N = p

  @turbo for j in 2:N-1, i in 2:N-1
    du[i,j,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end

  @turbo for j in 2:N-1, i in 2:N-1
    du[i,j,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
            a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds for j in 2:N-1
    i=1
    du[1,j,1] = D1*(2u[2,j,1] + u[1,j+1,1] + u[1,j-1,1] - 4u[1,j,1]) +
            a*u[1,j,1]^2/u[1,j,2] + ubar - α*u[1,j,1]
  end
  @inbounds for j in 2:N-1
    i = 1
    du[1,j,2] = D2*(2u[i+1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
            a*u[i,j,1]^2 - β*u[i,j,2]
  end
  @inbounds for j in 2:N-1
    i = N
    du[end,j,1] = D1*(2u[i-1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
           a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for j in 2:N-1
    i = N
    du[end,j,2] = D2*(2u[i-1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
           a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds for i in 2:N-1
    j = 1
    du[i,1,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for i in 2:N-1
    j = 1
    du[i,1,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
              a*u[i,j,1]^2 - β*u[i,j,2]
  end
  @inbounds for i in 2:N-1
    j = N
    du[i,end,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for i in 2:N-1
    j = N
    du[i,end,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds begin
    i = 1; j = 1
    du[1,1,1] = D1*(2u[i+1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[1,1,2] = D2*(2u[i+1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
              a*u[i,j,1]^2 - β*u[i,j,2]

    i = 1; j = N
    du[1,N,1] = D1*(2u[i+1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[1,N,2] = D2*(2u[i+1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]

    i = N; j = 1
    du[N,1,1] = D1*(2u[i-1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[N,1,2] = D2*(2u[i-1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]

    i = N; j = N
    du[end,end,1] = D1*(2u[i-1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[end,end,2] = D2*(2u[i-1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]
   end
end
prob_turbo_devec = ODEProblem(turbo_gm!,r0,(0.0,0.1),pdevec);
@benchmark solve($prob_turbo_devec,Tsit5())
# samples: 600; evals/sample: 1; memory estimate: 29.63 MiB; allocs estimate: 438
# ns

#  (7.11e6  - 7.78e6 ]  ██████████████████████████████ 324
#  (7.78e6  - 8.45e6 ]  ▍4
#  (8.45e6  - 9.12e6 ]  ████████▎88
#  (9.12e6  - 9.79e6 ]  ████████▉95
#  (9.79e6  - 1.046e7]  ██████▋71
#  (1.046e7 - 1.113e7]  ▌5
#  (1.113e7 - 1.18e7 ]  ▍3
#  (1.18e7  - 1.247e7]  ▍3
#  (1.247e7 - 1.314e7]  ▎2
#  (1.314e7 - 1.381e7]  ▎2
#  (1.381e7 - 1.448e7]  ▍3

#                   Counts

# min: 7.107 ms (0.00% GC); mean: 8.331 ms (7.72% GC); median: 7.374 ms (0.00% GC); max: 14.477 ms (0.00% GC).
function turbo_gm_fuse!(du,u,p,t)
  a,α,ubar,β,D1,D2,N = p

  @turbo for j in 2:N-1, i in 2:N-1
    du[i,j,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[i,j,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
            a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @turbo for j in 2:N-1
    du[1,j,1] = D1*(2u[2,j,1] + u[1,j+1,1] + u[1,j-1,1] - 4u[1,j,1]) +
            a*u[1,j,1]^2/u[1,j,2] + ubar - α*u[1,j,1]
  # end
  # @turbo for j in 2:N-1
    du[1,j,2] = D2*(2u[2,j,2] + u[1,j+1,2] + u[1,j-1,2] - 4u[1,j,2]) +
            a*u[1,j,1]^2 - β*u[1,j,2]
  # end
  # @turbo for j in 2:N-1
    du[N,j,1] = D1*(2u[N-1,j,1] + u[N,j+1,1] + u[N,j-1,1] - 4u[N,j,1]) +
           a*u[N,j,1]^2/u[N,j,2] + ubar - α*u[N,j,1]
  # end
  # @turbo for j in 2:N-1
    du[N,j,2] = D2*(2u[N-1,j,2] + u[N,j+1,2] + u[N,j-1,2] - 4u[N,j,2]) +
      a*u[N,j,1]^2 - β*u[N,j,2]
  end
  @turbo for j in 2:N-1
    du[j,1,1] = D1*(u[j-1,1,1] + u[j+1,1,1] + 2u[j,2,1] - 4u[j,1,1]) +
              a*u[j,1,1]^2/u[j,1,2] + ubar - α*u[j,1,1]
  # end
  # @turbo for j in 2:N-1
    du[j,1,2] = D2*(u[j-1,1,2] + u[j+1,1,2] + 2u[j,2,2] - 4u[j,1,2]) +
              a*u[j,1,1]^2 - β*u[j,1,2]
  # end
  # @turbo for j in 2:N-1
    du[j,N,1] = D1*(u[j-1,N,1] + u[j+1,N,1] + 2u[j,N-1,1] - 4u[j,N,1]) +
             a*u[j,N,1]^2/u[j,N,2] + ubar - α*u[j,N,1]
  # end
  # @turbo for j in 2:N-1
    du[j,N,2] = D2*(u[j-1,N,2] + u[j+1,N,2] + 2u[j,N-1,2] - 4u[j,N,2]) +
             a*u[j,N,1]^2 - β*u[j,N,2]
  end

  @inbounds begin
    i = 1; j = 1
    du[1,1,1] = D1*(2u[i+1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[1,1,2] = D2*(2u[i+1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
              a*u[i,j,1]^2 - β*u[i,j,2]

    i = 1; j = N
    du[1,N,1] = D1*(2u[i+1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[1,N,2] = D2*(2u[i+1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]

    i = N; j = 1
    du[N,1,1] = D1*(2u[i-1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[N,1,2] = D2*(2u[i-1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]

    i = N; j = N
    du[N,N,1] = D1*(2u[i-1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[N,N,2] = D2*(2u[i-1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]
   end
end
prob_turbo_fuse = ODEProblem(turbo_gm_fuse!,r0,(0.0,0.1),pdevec);
@benchmark solve($prob_turbo_fuse,Tsit5())
# samples: 629; evals/sample: 1; memory estimate: 29.62 MiB; allocs estimate: 433
# ns

#  (6.8e6  - 7.9e6 ]  ██████████████████████████████427
#  (7.9e6  - 8.9e6 ]  ██▋36
#  (8.9e6  - 1.0e7 ]  ████▊66
#  (1.0e7  - 1.11e7]  ██████▉97
#  (1.11e7 - 1.22e7]   0
#  (1.22e7 - 1.33e7]   0
#  (1.33e7 - 1.44e7]   0
#  (1.44e7 - 1.54e7]  ▏1
#  (1.54e7 - 1.65e7]  ▏1
#  (1.65e7 - 1.76e7]   0
#  (1.76e7 - 2.12e7]  ▏1

#                   Counts

# min: 6.786 ms (0.00% GC); mean: 7.938 ms (7.63% GC); median: 7.014 ms (0.00% GC); max: 21.180 ms (66.47% GC).

[chriselrod]

Those loops are easy enough for LLVM to SIMD, and there's not too much that can be done to make them faster.
I.e., there isn't much reuse you can get through blocking. You're just streaming through memory.

It is unrolling the j loop to try and save on j loads a bit (i.e., j+1 on one j iteration is j on the next, is j-1 on the next), but that doesn't seem to help much.

[ChrisRackauckas]

Yeah, I think you need to not have the nonlinear portion of the PDE (+ a*u[i,j,1]^2 - β*u[i,j,2]) for a stencil compiler to work well. This is what I was mentioning in https://discourse.julialang.org/t/how-to-tackle-2d-reaction-diffusion-equation/54110/15?u=chrisrackauckas . The authors of that seemed to get mad at me, but it seems to be the simple fact that once you have a term that is outside the original range, u[i,j,2], or a term which is sufficiently difficult to compute so that it's not too memory bound, u[i,j,1]^2, then the act of fusion is just better than having a good stencil compiler. That LoopVectorization demo above is precisely what I needed to see to fully prove that indeed the fused version outperforms even a good stencil compiler in that case, and that @turbo has an effect but it's not massive.

It would still be good to investigate multithreading between those when the space is really large though, but that's a separate thing and only matters for the huge PDE case. But essentially, those benchmarks are my argument for why we shouldn't be generating A*u operators but instead should be generating (and fusing) map operations.

[chriselrod]

Making them dense:

Ax = Array(Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1]))

adds a huge amount of computation, making it go from O(N^2) to O(N^3).
Seems like an obviously bad idea.
We could implement specialized Tridiagonal multiplication routines, e.g. like
JuliaSIMD/LoopVectorization.jl#266 (comment)
but at that point we may as well use the fusing map operations.

For the mapped version, 100x100 isn't large enough to benefit from multithreading yet, but 1000x1000 would be.

[ChrisRackauckas]

The first part is the same as an image filter.

[chriselrod]

or a term which is sufficiently difficult to compute so that it's not too memory bound, u[i,j,1]^2

Not sure if I follow here, but u[i,j,1] is easy to compute.
Fusion is great for memory bound applications though. Passing over the data once is better than passing over it twice.

We shouldn't be doing dense Ax::Matrix operations.
I assume ParallelStencil wouldn't be?

If we let Ax be tridiagonal, it'd do much better.
Still probably not as well as the loops, but better.

[ChrisRackauckas]

Not sure if I follow here, but u[i,j,1] is easy to compute.

I mean, when you're looping over u[i,j,2] and need to grab u[i,j,1], you very quickly expand the cache lines like that. And this is still a very small example with very few nonlinearities!

We shouldn't be doing dense Ax::Matrix operations. I assume ParallelStencil wouldn't be?

It wouldn't, I thought you were using the same setup as the image filter you had? A 2D Laplacian is just a 2D filter.

[chriselrod]

Ha, let me do that. Above it was just matmul!.

Note that the image filtering example doesn't do boundaries...

I mean, when you're looping over u[i,j,2] and need to grab u[i,j,1], you very quickly expand the cache lines like that.

Although reuse tends to be fairly immediate (e.g. u[i,j,1], u[i+1,j,1]), but then there aren't any more opportunities for re-use, so you don't mind these loads getting evicted.
Although there is concern over u[i,j,1], u[i,j+1,1]. It'd probably take some experimenting to see if some larger cache tiling helps more than the j unrolling does.

[chriselrod]

@ChrisRackauckas here is the example I think you were looking for:

using StaticArrays, LoopVectorization, Static
const LAPLACE_KERNEL = @MMatrix([0.0 1.0 0.0; 1.0 -4.0 1.0; 0.0 1.0 0.0]);
function filter2davx!(out::AbstractMatrix, A::AbstractMatrix, kern::MMatrix{3,3})
  N = size(out,1)
  @inbounds @fastmath let j_2 = j_1 = 1
    tmp  = A[1+j_1, 1+j_2] * kern[-1+2, -1+2]
    tmp += A[1+j_1, 0+j_2] * kern[-1+2, 0+2]
    tmp += A[1+j_1, 1+j_2] * kern[-1+2, 1+2]
    tmp += A[0+j_1, 1+j_2] * kern[0+2, -1+2]
    tmp += A[1+j_1, 1+j_2] * kern[1+2, -1+2]
    for i_1 ∈ 0:1, i_2 ∈ 0:1
      tmp += A[i_1+j_1, i_2+j_2] * kern[i_1+2, i_2+2]
    end
    out[j_1,j_2] = tmp
  end
  @inbounds @fastmath let j_1 = 1; j_2 = N
    tmp  = A[1+j_1, -1+j_2] * kern[-1+2, -1+2]
    tmp += A[1+j_1, 0+j_2] * kern[-1+2, 0+2]
    tmp += A[1+j_1, -1+j_2] * kern[-1+2, 1+2]
    tmp += A[0+j_1, -1+j_2] * kern[0+2, 1+2]
    tmp += A[1+j_1, -1+j_2] * kern[1+2, 1+2]
    for i_1 ∈ 0:1, i_2 ∈ -1:0
      tmp += A[i_1+j_1, i_2+j_2] * kern[i_1+2, i_2+2]
    end
    out[j_1,j_2] = tmp
  end
  @inbounds @fastmath let j_2 = 1; j_1 = N
    tmp  = A[-1+j_1, 1+j_2] * kern[1+2, -1+2]
    tmp += A[-1+j_1, 0+j_2] * kern[1+2, 0+2]
    tmp += A[-1+j_1, 1+j_2] * kern[-1+2, 1+2]
    tmp += A[0+j_1, 1+j_2] * kern[0+2, -1+2]
    tmp += A[-1+j_1, 1+j_2] * kern[1+2, -1+2]
    for i_1 ∈ -1:0, i_2 ∈ 0:1
      tmp += A[i_1+j_1, i_2+j_2] * kern[i_1+2, i_2+2]
    end
    out[j_1,j_2] = tmp
  end  
  @turbo for j ∈ 2:N-1
  # for j ∈ 2:N-1
    tmp_1 = zero(eltype(out))
    tmp_2 = zero(eltype(out))
    for i_1 ∈ -1:1
      tmp_1 += A[i_1+j, 2] * kern[i_1+2, 3]
      tmp_2 += A[2, i_1+j] * kern[3, i_1+2]
      for i_2 ∈ 0:1
        tmp_1 += A[i_1+j, 1+i_2] * kern[i_1+2, i_2+2]
        tmp_2 += A[1+i_2, i_1+j] * kern[i_2+2, i_1+2]
      end
    end
    out[j,1] = tmp_1
    out[1,j] = tmp_2
  end
  @turbo for j_2 ∈ 2:N-1, j_1 ∈ 2:N-1
  # for j_2 ∈ 2:N-1, j_1 ∈ 2:N-1
    tmp = zero(eltype(out))
    for i_1 ∈ -1:1, i_2 ∈ -1:1
      tmp += A[i_1+j_1, i_2+j_2] * kern[i_1+2, i_2+2]
    end
    out[j_1,j_2] = tmp
  end
  @turbo for j ∈ 2:N-1
  # for j ∈ 2:N-1
    tmp_1 = zero(eltype(out))
    tmp_2 = zero(eltype(out))
    for i_1 ∈ -1:1
      tmp_1 += A[i_1+j, N-1] * kern[i_1+2, 3]
      tmp_2 += A[N-1, i_1+j] * kern[3, i_1+2]
      for i_2 ∈ -1:0
        tmp_1 += A[i_1+j, N+i_2] * kern[i_1+2, i_2+2]
        tmp_2 += A[N+i_2, i_1+j] * kern[i_2+2, i_1+2]
      end
    end
    out[j,N] = tmp_1
    out[N,j] = tmp_2
  end
  @inbounds @fastmath let j_2 = j_1 = N
    tmp  = A[-1+j_1, -1+j_2] * kern[1+2, -1+2]
    tmp += A[-1+j_1, 0+j_2] * kern[1+2, 0+2]
    tmp += A[-1+j_1, -1+j_2] * kern[1+2, 1+2]
    tmp += A[-1+j_1, -1+j_2] * kern[-1+2, 1+2]
    tmp += A[0+j_1, -1+j_2] * kern[0+2, 1+2]
    for i_1 ∈ -1:0, i_2 ∈ -1:0
      tmp += A[i_1+j_1, i_2+j_2] * kern[i_1+2, i_2+2]
    end
    out[j_1,j_2] = tmp
  end
  nothing
end
function gm7!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = @view r[:,:,1];
  v = @view r[:,:,2];
  du = @view dr[:,:,1];
  dv = @view dr[:,:,2];
  filter2davx!(Ayu, u, LAPLACE_KERNEL)
  # @turbo @. du = D1*Ayu + a*u*u./v + ubar - α*u
  @turbo for i ∈ eachindex(u)
    du[i] = D1 * Ayu[i] + a * u[i]*u[i] / v[i] + ubar - α*u[i]
  end
  filter2davx!(Ayu, v, LAPLACE_KERNEL)
  # @turbo @. dv = D2*Ayu + a*u*u - β*v
  @turbo for i ∈ eachindex(u)
    dv[i] = D2 * Ayu[i] + a * u[i]*u[i] - β*v[i]
  end
end
prob_gm7 = ODEProblem(gm7!,r0,(0.0,0.1),p);
@benchmark solve($prob_gm7,Tsit5())
samples: 546; evals/sample: 1; memory estimate: 29.62 MiB; allocs estimate: 433
ns

 (7.96e6  - 8.63e6 ]  ██████████████████████████████311
 (8.63e6  - 9.29e6 ]  ██▋26
 (9.29e6  - 9.95e6 ]  █████████▊101
 (9.95e6  - 1.062e7]  █████51
 (1.062e7 - 1.128e7]  ▍3
 (1.128e7 - 1.194e7]  ██▌25
 (1.194e7 - 1.26e7 ]  ▏1
 (1.26e7  - 1.327e7]  ▏1
 (1.327e7 - 1.393e7]   0
 (1.393e7 - 1.459e7]  █▍13
 (1.459e7 - 1.525e7]  █▍14

                  Counts

min: 7.965 ms (0.00% GC); mean: 9.162 ms (6.71% GC); median: 8.165 ms (0.00% GC); max: 15.255 ms (10.53% GC).

If we wanted to optimize this a bit more, we could

1. take advantage of the sparsity in the laplacian kernel
2. make the values from the kernel constant
3. Fuse operations to pass over data less often.

Oops, we now have fast_gm!/turbo_gm_fuse!.

[ChrisRackauckas]

You can just hardcode that it's the stencil

0 1 0
1 -4 1
0 1 0

[chriselrod]

Sure, but at that point why not go the rest of the way to turbo_gm_fuse!?

[ChrisRackauckas]

Yeah, go all of the way.