Cut the allocations on the calls that go inside a loop

CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Fixed
 
+- `with_local_array!` spent 23 allocations per call and did not infer its return type: it queried every vector's `PetscMemType` with a checkout of its own before checking the vectors out for real, and `memtype_backend` dispatched on a `Val` built from that run-time value. Host memory, the common case, now answers before the dispatch and the extra query is gone, so a one-vector checkout costs 2 allocations and infers. `with_field_views!` had the same fault plus a `read`/`write` variable that changed type, and goes from 24 allocations to 8. A SNES solve whose callbacks check arrays out drops from 210 allocations to 47. Device memory dispatches as before, so the CUDA extension is unaffected.
+- `v[i] = a` on a `PetscVec` and `v[i]` built their one-element index vector twice over. `v[i] = a` now allocates nothing and `v[i]` two.
+- Out-of-place broadcasting over a `PetscVec`, `w = 2 .* x`, read one entry at a time through `VecGetValues`, which cost 4003 allocations at 10 000 entries and failed on a distributed vector with a raw PETSc error. It runs on the local arrays now, at about 5 allocations whatever the length, and still gives a `Vector`. A distributed vector throws an `ArgumentError`, since a `Vector` of the whole vector only means something where one rank holds all of it.
 - Broadcasting into a `PetscVec` (`y .= 2 .* x`, a typical `MatShell` body) read and wrote the vectors one entry at a time, through `VecGetValues` and `VecSetValues`: at 10 000 entries a `MatShell` product took 139 000 allocations and was 500× slower than the same loop on the local arrays. The vectors now take part as their local arrays, at a fixed cost of about 10 allocations. On several ranks the broadcast runs over the entries each rank owns.
 - `LibPETSc.ISColoringGetIS` with `PETSC_OWN_POINTER` returned the index sets as borrowed, so they and the C array holding them leaked. The index sets are now owned by the caller and the array is freed. `ISColoringRestoreIS` accepts the vector `ISColoringGetIS` returns.
 - `LibPETSc.DMCreateFieldIS` leaked the field names and both C arrays.
@@ -21,6 +24,7 @@
 
 ### Added
 
+- A performance contract: `test/core/performance.jl` checks that 29 calls meant for a loop infer a concrete return type and allocate the same amount at two problem sizes, and `docs/src/man/performance.md` says how to keep that cost (array blocks, BLAS threads per rank, concretely typed objects, precompiling before `mpiexec`, `-log_view`).
 - `copyto!(dst, src)` between two `PetscVec`s.
 - Matrix methods: `fill!(A, 0)`, `zero_rows!` and `zero_rows_local!` for Dirichlet rows, `set_option!`, `diagonal!(d, A)` and `isassembled`.
 - `destroy!` on `IS`, `AO`, `PF` and `Tao` handles.

docs/make.jl

@@ -82,6 +82,7 @@ makedocs(;
         ],
         "C to Julia name index" => "man/api_index.md",
         "Utilities" => "man/utilities.md",
+        "Performance" => "man/performance.md",
         "Running on HPC Systems" => "man/hpc.md",
         "GPU Support (CUDA)" => "man/gpu.md",
         "FAQ"  => "man/FAQ.md",

docs/src/man/performance.md

@@ -0,0 +1,74 @@
+# Performance
+
+PETSc.jl aims for a fixed cost per call: the work a high-level call does on the Julia side does not grow with the size of the vector, matrix or grid it acts on. `test/core/performance.jl` holds that as a contract, checking every call below at two problem sizes. If a call in this page starts allocating more at the larger size, that is a bug worth reporting.
+
+## Read and write arrays in a block
+
+[`with_local_array!`](@ref) and [`with_field_views!`](@ref) check the arrays out once, hand them to your function and give them back, so the cost is a few allocations per vector no matter how long it is. Entry-by-entry access through `v[i]` goes through `VecGetValues`, one C call per entry, and only reaches the entries this rank owns:
+
+```julia
+# one checkout, then plain Julia on the array
+PETSc.with_local_array!(x, y; write = (true, false)) do a, b
+    @inbounds for i in eachindex(a)
+        a[i] = 2b[i]
+    end
+end
+
+# one C call per entry, and wrong on several ranks
+for i in 1:length(x)
+    x[i] = 2y[i]
+end
+```
+
+Broadcasting does the checkout for you, in both directions. `x .= 2 .* y` writes the entries this rank owns. `w = 2 .* x` gives a plain `Vector`, and since that only means something where one rank holds the whole vector, it throws on a distributed one: broadcast into a `PetscVec` of the same layout instead.
+
+Every view a block hands you is indexed the way the DM numbers its points, so a view and [`stencil`](@ref) address the same entry, and the arrays stay concretely typed inside the block.
+
+## BLAS threads
+
+Every `PETSc_jll` build calls BLAS through libblastrampoline, which means PETSc's dense kernels run in Julia's own BLAS pool. PETSc cannot size that pool, so `-mat_mumps_*` and any other dense work inherits whatever Julia set.
+
+When several MPI ranks share a node, [`initialize`](@ref) sets the pool to one thread per rank, unless `-blas_num_threads`, `OPENBLAS_NUM_THREADS` or `OMP_NUM_THREADS` says otherwise. Without that, each rank runs a pool of busy-waiting threads competing for the same cores: a 4-rank DMStag Stokes solve measured 25 times slower. A serial run, or one rank per node, keeps Julia's default.
+
+To choose for yourself:
+
+```julia
+PETSc.initialize(petsclib; options = ["-blas_num_threads", "2"])
+```
+
+## Keep PETSc objects concretely typed
+
+A wrapper carries its library in its type, so a high-level call resolves the right method at compile time. Storing objects in a field typed `Any`, or in a `Ref{Any}`, throws that away and every call through it becomes a dynamic dispatch:
+
+```julia
+mutable struct Ctx
+    dm::Any            # every call on ctx.dm dispatches at run time
+end
+
+mutable struct Ctx{D}
+    dm::D              # resolved at compile time
+end
+```
+
+Where a context has to stay untyped, put a function barrier between it and the loop: read the fields once into local variables and pass them to a second function that does the work.
+
+## Precompile before an MPI launch
+
+Precompile serially, once, before running on several ranks:
+
+```bash
+julia --project=. -e 'using PETSc'
+mpiexec -n 4 julia --project=. run.jl
+```
+
+Ranks that start against a stale cache all try to precompile at once and wait on each other's lock. The run then sits at 0% CPU with no error.
+
+## Profiling a solve
+
+`-log_view` gives PETSc's own breakdown, which attributes time to `KSPSolve`, `MatMult`, `PCApply` and the rest, and counts the calls:
+
+```julia
+PETSc.initialize(petsclib; log_view = true)
+```
+
+It is the right first measurement for a slow solve, because it separates time spent in PETSc from time spent in your callbacks. Julia's `@profile` then covers the callbacks themselves.

scripts/renames.jl

@@ -459,6 +459,12 @@ const INTERNAL = Set{Symbol}([
     :petsclib_of,
     :make_local_array,
     :with_broadcast_arrays,
+    :aliases_dest,
+    :check_undistributed,
+    :PetscVecStyle,
+    :checkout_arrays,
+    :checked_array,
+    :flags,
     :to_petscint_tuple,
     :audit_walk,
     :audit_targets,

src/dmstag.jl

@@ -820,8 +820,17 @@ function with_field_views!(
     read::Union{Bool, NTuple{M, Bool}} = true,
     write::Union{Bool, NTuple{M, Bool}} = true,
 ) where {PetscLib, N, M}
-    read  isa Bool && (read  = ntuple(_ -> read,  Val(M)))
-    write isa Bool && (write = ntuple(_ -> write, Val(M)))
+    return with_field_views!(f, dm, vecs, fields, flags(read, Val(M)), flags(write, Val(M)))
+end
+
+function with_field_views!(
+    f,
+    dm::DMStag{PetscLib, N},
+    vecs::NTuple{M, AbstractPetscVec{PetscLib}},
+    fields,
+    read::NTuple{M, Bool},
+    write::NTuple{M, Bool},
+) where {PetscLib, N, M}
     lib = getlib(PetscLib)
     gc = ghost_corners(dm)
     q = Int(LibPETSc.DMStagGetEntriesPerElement(lib, dm))

src/vec.jl

@@ -156,11 +156,22 @@ end
 Base.ndims(::Type{<:AbstractPetscVec}) = 1
 Base.IndexStyle(::Type{<:AbstractPetscVec}) = IndexLinear()
 Base.axes(v::AbstractPetscVec) = (Base.OneTo(length(v)),)
-Base.BroadcastStyle(::Type{<:AbstractPetscVec}) = Broadcast.DefaultArrayStyle{1}()
 # A Vec takes part in a broadcast as itself; the default would `collect` it
 # entry by entry, since it iterates but is not an AbstractArray
 Base.Broadcast.broadcastable(v::AbstractPetscVec) = v
 
+"""
+    PetscVecStyle
+
+The broadcast style of a `PetscVec`: an expression with one in it runs on the
+local arrays rather than entry by entry. It wins over an ordinary `Vector`, so
+`x .+ v` uses it too.
+"""
+struct PetscVecStyle <: Broadcast.AbstractArrayStyle{1} end
+Base.BroadcastStyle(::Type{<:AbstractPetscVec}) = PetscVecStyle()
+PetscVecStyle(::Val{1}) = PetscVecStyle()
+PetscVecStyle(::Val{N}) where {N} = Broadcast.DefaultArrayStyle{N}()
+
 # The library a wrapper carries in its type parameter, as a value. Used where a
 # call has to recover `petsclib` from an object rather than take it as an
 # argument (§8), for instance `save_vtk(vecs, filename)`.
@@ -197,16 +208,16 @@ function set_type!(v::AbstractPetscVec{PetscLib}, type::Symbol) where {PetscLib}
     return v
 end
 
-function Base.getindex(v::AbstractPetscVec{PetscLib}, i::Integer) where {PetscLib} 
+function Base.getindex(v::AbstractPetscVec{PetscLib}, i::Integer) where {PetscLib}
     PetscInt = inttype(PetscLib)
-    val = LibPETSc.VecGetValues(PetscLib,v, PetscInt(1), PetscInt.([i-1]))
+    val = LibPETSc.VecGetValues(PetscLib, v, PetscInt(1), PetscInt[i - 1])
     return val[1]
 end
 
 # Range indexing
 function Base.getindex(v::AbstractPetscVec{PetscLib}, r::AbstractRange) where {PetscLib}
     PetscInt = inttype(PetscLib)
-    val = LibPETSc.VecGetValues(PetscLib,v, PetscInt(length(r)), PetscInt.(Vector(r) .- 1))
+    val = LibPETSc.VecGetValues(PetscLib, v, PetscInt(length(r)), PetscInt[i - 1 for i in r])
     return val
 end
 
@@ -218,16 +229,18 @@ end
 
 Base.isapprox(v::AbstractPetscVec{PetscLib}, w::AbstractPetscVec{PetscLib}; kwargs...) where {PetscLib} = all(v[:] .≈ w[:]; kwargs...)
 
-function Base.setindex!(v::AbstractPetscVec{PetscLib}, val, i::Integer) where {PetscLib} 
-     PetscInt = inttype(PetscLib)
-     PetscScalar = PETSc.scalartype(PetscLib)
-     LibPETSc.VecSetValues(PetscLib,v, PetscInt(1), PetscInt.([i-1]), [PetscScalar(val)], PETSc.INSERT_VALUES)
-     return v
+function Base.setindex!(v::AbstractPetscVec{PetscLib}, val, i::Integer) where {PetscLib}
+    PetscInt = inttype(PetscLib)
+    PetscScalar = scalartype(PetscLib)
+    LibPETSc.VecSetValues(PetscLib, v, PetscInt(1), PetscInt[i - 1],
+                          PetscScalar[val], PETSc.INSERT_VALUES)
+    return v
 end
 
-function Base.setindex!(v::AbstractPetscVec{PetscLib}, vals, r::AbstractRange) where {PetscLib} 
+function Base.setindex!(v::AbstractPetscVec{PetscLib}, vals, r::AbstractRange) where {PetscLib}
     PetscInt = inttype(PetscLib)
-    LibPETSc.VecSetValues(PetscLib,v, PetscInt(length(r)), PetscInt.(Vector(r) .- 1), vals, PETSc.INSERT_VALUES)
+    LibPETSc.VecSetValues(PetscLib, v, PetscInt(length(r)), PetscInt[i - 1 for i in r],
+                          vals, PETSc.INSERT_VALUES)
     return v
 end
 
@@ -275,14 +288,38 @@ function Base.copyto!(dest::AbstractPetscVec, bc::Base.Broadcast.Broadcasted)
     return dest
 end
 
-# Calls `f` with `args`, each `PetscVec` replaced by its local array:
-# `dest_arr` for `dest` itself, a read-only checkout for any other
+# Out-of-place broadcasting, `w = f.(args...)`, runs on the local arrays too and
+# gives a `Vector`, so `2 .* x` is what it was before, without the one C call per
+# entry. A `Vector` of the whole vector only means something where one rank holds
+# all of it, so a distributed vector throws instead.
+function Base.copy(bc::Broadcast.Broadcasted{PetscVecStyle})
+    flat = Broadcast.flatten(bc)
+    foreach(check_undistributed, flat.args)
+    return with_broadcast_arrays(nothing, nothing, flat.args) do args
+        copy(Broadcast.instantiate(Broadcast.Broadcasted(flat.f, args)))
+    end
+end
+
+check_undistributed(_) = nothing
+function check_undistributed(v::AbstractPetscVec{PetscLib}) where {PetscLib}
+    lib = getlib(PetscLib)
+    LibPETSc.VecGetLocalSize(lib, v) == LibPETSc.VecGetSize(lib, v) || throw(ArgumentError(
+        "a broadcast of a distributed Vec cannot give a Vector of the whole vector. " *
+        "Broadcast into a Vec of the same layout, or use with_local_array! for the " *
+        "entries this rank owns."
+    ))
+    return nothing
+end
+
+# Calls `f` with `args`, each `PetscVec` replaced by its local array: `dest_arr`
+# for `dest` itself, a read-only checkout for any other. `dest` is `nothing` when
+# the broadcast is out of place and nothing can alias.
 with_broadcast_arrays(f, dest, dest_arr, ::Tuple{}) = f(())
 function with_broadcast_arrays(f, dest, dest_arr, args::Tuple)
     a = first(args)
     rest = Base.tail(args)
     a isa AbstractPetscVec || return with_broadcast_arrays(r -> f((a, r...)), dest, dest_arr, rest)
-    a.ptr == dest.ptr && return with_broadcast_arrays(r -> f((dest_arr, r...)), dest, dest_arr, rest)
+    aliases_dest(a, dest) && return with_broadcast_arrays(r -> f((dest_arr, r...)), dest, dest_arr, rest)
     arr, raw, backend = acquire_local_array(a; read = true, write = false)
     try
         return with_broadcast_arrays(r -> f((arr, r...)), dest, dest_arr, rest)
@@ -291,6 +328,9 @@ function with_broadcast_arrays(f, dest, dest_arr, args::Tuple)
     end
 end
 
+aliases_dest(::AbstractPetscVec, ::Nothing) = false
+aliases_dest(a::AbstractPetscVec, dest::AbstractPetscVec) = a.ptr == dest.ptr
+
 Base.materialize!(dest::AbstractPetscVec, bc::Base.Broadcast.Broadcasted) = (Base.copyto!(dest, bc); dest)
 
 # Iterator interface
@@ -413,7 +453,11 @@ GPU extensions return their own singleton for device memory.
 memtype_backend(::Val{LibPETSc.PETSC_MEMTYPE_HOST}) = nothing
 memtype_backend(::Val{MT}) where {MT} =
     error("No GPU backend loaded for PetscMemType $MT — load CUDA.jl, AMDGPU.jl, …")
-memtype_backend(mt::LibPETSc.PetscMemType) = memtype_backend(Val(mt))
+# `Val(mt)` on a value PETSc returned at run time is a dynamic dispatch, so host
+# memory, the common case, answers before it: that keeps a checkout inferable and
+# allocation free. A device goes on dispatching, so extensions are unaffected.
+memtype_backend(mt::LibPETSc.PetscMemType) =
+    mt === LibPETSc.PETSC_MEMTYPE_HOST ? nothing : memtype_backend(Val(mt))
 
 # ── Device-aware local array access ───────────────────────────────────────────
 #
@@ -603,10 +647,10 @@ end
 function with_local_array!(
     f!,
     vecs::NTuple{N, AbstractPetscVec};
-    kwargs...,
+    read::Union{Bool, NTuple{N, Bool}} = true,
+    write::Union{Bool, NTuple{N, Bool}} = true,
 ) where {N}
-    A = memtype(vecs...)
-    return with_local_array!(f!, A, vecs; kwargs...)
+    return checkout_arrays(f!, nothing, vecs, flags(read, Val(N)), flags(write, Val(N)))
 end
 with_local_array!(f!, vecs...; kwargs...) = with_local_array!(f!, vecs; kwargs...)
 
@@ -617,32 +661,40 @@ function with_local_array!(
     read::Union{Bool, NTuple{N, Bool}} = true,
     write::Union{Bool, NTuple{N, Bool}} = true,
 ) where {A <: AbstractArray, N}
-    read isa NTuple{N, Bool} || (read = ntuple(_ -> read, N))
-    write isa NTuple{N, Bool} || (write = ntuple(_ -> write, N))
-    # Acquire all arrays first (no finalizers), then use try/finally for release.
-    # This avoids the Julia pitfall where Base.finalize + GC can both run the
-    # finalizer if the object becomes unreachable again (double-restore → crash).
+    return checkout_arrays(f!, A, vecs, flags(read, Val(N)), flags(write, Val(N)))
+end
+with_local_array!(f!, ::Type{A}, vecs...; kwargs...) where {A <: AbstractArray} =
+    with_local_array!(f!, A, vecs; kwargs...)
+
+flags(b::Bool, ::Val{N}) where {N} = ntuple(_ -> b, Val(N))
+flags(b::NTuple{N, Bool}, ::Val{N}) where {N} = b
+
+# Check out every vector, call `f!` with the arrays, and hand them all back, 
+# also when `f!` throws. `A` is the array type the caller asked for, or `nothing` to
+# take whatever the vectors are on. The arrays are acquired without a finalizer,
+# because after `Base.finalize` the GC may run the finalizer a second time once
+# the array becomes unreachable, and restoring twice crashes PETSc.
+function checkout_arrays(f!, A, vecs::NTuple{N, AbstractPetscVec}, read, write) where {N}
     acquired = map(vecs, read, write) do v, r, w
-        acquire_local_array(v; read=r, write=w)
+        acquire_local_array(v; read = r, write = w)
     end
     try
-        # Type check inside try so finally still releases on mismatch.
-        arrays = map(acquired) do (arr, cpu_arr, backend)
-            arr isa A || throw(ArgumentError(
-                "expected array of type $A but Vec returned $(typeof(arr)). " *
-                "Check that the Vec lives on the expected device."
-            ))
-            arr
-        end
+        # inside the try, so an array of the wrong type is still handed back
+        arrays = map(a -> checked_array(a[1], A), acquired)
         return f!(arrays...)
     finally
         foreach(vecs, acquired, read, write) do v, (_, cpu_arr, backend), r, w
-            release_local_array(cpu_arr, backend, v; read=r, write=w)
+            release_local_array(cpu_arr, backend, v; read = r, write = w)
         end
     end
 end
-with_local_array!(f!, ::Type{A}, vecs...; kwargs...) where {A <: AbstractArray} =
-    with_local_array!(f!, A, vecs; kwargs...)
+
+checked_array(arr, ::Nothing) = arr
+checked_array(arr, ::Type{A}) where {A} =
+    arr isa A ? arr : throw(ArgumentError(
+        "expected array of type $A but Vec returned $(typeof(arr)). " *
+        "Check that the Vec lives on the expected device."
+    ))
 
 
 """