From df2a1eb4f094ed30868f82eabed885409c3dfe15 Mon Sep 17 00:00:00 2001 From: Jukka Aho Date: Mon, 11 May 2026 02:23:05 +0300 Subject: [PATCH] refactor(matrix-free): kernelless operators and README MatrixFreeOperator stores cache/asm/mesh only; forward Pass~1 keywords; deprecate redundant-kernel overloads via shared depwarn helper. --- src/assemblers/matrix_free/README.md | 17 +- src/assemblers/matrix_free/operator.jl | 730 ++++++++++++++++++++++--- 2 files changed, 669 insertions(+), 78 deletions(-) diff --git a/src/assemblers/matrix_free/README.md b/src/assemblers/matrix_free/README.md index 69a7f86..3e52527 100644 --- a/src/assemblers/matrix_free/README.md +++ b/src/assemblers/matrix_free/README.md @@ -1,3 +1,8 @@ + + # src/assemblers/matrix_free/ Declarative constraints, loads, preconditioners and the eigensolver @@ -9,15 +14,23 @@ assembled and the matrix-free solves. ## Files -- `dirichlet.jl` — `AbstractDirichletConstraint`, `PenaltyDirichlet`, `EliminatedDirichlet`. Drive both the assembled solve (`apply_constraint!(K, c)` / `apply_constraint!(K, b, c)`) and the matrix-free path (`apply_constraint_pre!` / `apply_constraint_post!` wrapped around `apply_K!`). Also re-exports `matrix_free_op` as a thin factory for `MatrixFreeOperator`. +- `dirichlet.jl` — `AbstractDirichletConstraint`, `PenaltyDirichlet`, `EliminatedDirichlet`. Drive both the assembled solve (`apply_constraint!(K, c)` / `apply_constraint!(K, b, c)`) and the matrix-free path (`apply_constraint_pre!` / `apply_constraint_post!` wrapped around `apply_K!`). `SparseMatrixCSC{Float64}` uses `O(nnz)` elimination and sparse-column RHS lifts; dense matrices keep the explicit row/column loops. Also re-exports `matrix_free_op` as a thin factory for `MatrixFreeOperator`. - `mpc.jl` — `AbstractMultipointConstraint`, `LinearMPC`. Penalty-enforced affine `u_s = sum(c_k * u_{m_k}) + g` constraints sharing the same hook protocol as Dirichlet, so they compose with both solve paths. The MPC contribution is folded into `MatrixFreeOperator` via the `mpc =` keyword. -- `operator.jl` — `AbstractMatrixFreeOperator`, `MatrixFreeOperator{C, A, K, M, D, P, L}`, `MatrixFreeMassOperator{C, A, K, M}`. Typed linear operators that implement `LinearAlgebra.mul!`, `eltype`, `size`, `*`, and a callable form (`op(y, x)`) so they plug into `LinearOperators.LinearOperator(...)`. Each operator owns its work buffer, so every mat-vec is allocation-free after warmup. The `MatrixFreeOperator` constructor accepts `dirichlet =`, `mpc =`, and `multiply_layout =` keywords; the layout drives `prepare_multiply_workspace!` before `apply_K!` (`LocalMultiplyLayout` copies `x`; future MPI layouts may fill ghost DOFs). Constraint hooks still compose as before. +- `operator.jl` — `AbstractMatrixFreeOperator`, `MatrixFreeOperator{C, A, K, M, D, P, L}`, `MatrixFreeMassOperator{C, A, K, M}`, `MatrixFreeOperatorKA`, `MatrixFreeMassOperatorKA`, `matrix_free_op_ka`, `matrix_free_mass_op_ka`. Typed linear operators that implement `LinearAlgebra.mul!`, `eltype`, `size`, `*`, and a callable form (`op(y, x)`) so they plug into `LinearOperators.LinearOperator(...)`. Each operator owns its work buffer, so every mat-vec is allocation-free after warmup. The `MatrixFreeOperator` constructor accepts `dirichlet =`, `mpc =`, `multiply_layout =`, and (for nonlinear continuum Pass~1) `configuration` / `global_material_cache` / `Δt`, forwarded into `apply_K!`; the layout drives `prepare_multiply_workspace!` before `apply_K!` (`LocalMultiplyLayout` copies `x`; future MPI layouts may fill ghost DOFs). Constraint hooks still compose as before. The KA variants wrap `apply_K!` / `apply_M!` on `DOFBasedCOOCacheKA` (volume-only stiffness for `MatrixFreeOperatorKA`; no Dirichlet/MPC on that path). The same file defines nonlinear callables `InternalForceOperator` / `NonlinearResidualOperator` (`internal_force_op` / `nonlinear_residual_op`) wrapping `assemble_internal_force!` / `nonlinear_equilibrium_residual!` (not `mul!`-based). - `loads.jl` — `AbstractNeumannLoad`, `NodalForce`, `UniformBodyForce`, `SurfaceLoad`, plus `apply_load!`. Reuses the cache's SoA `N_data` / `detJ_w` batches so body-force integration is allocation-free and shares Pass 1 with `apply_K!` / `apply_M!`. - `preconditioners.jl` — `JacobiPreconditioner`, `BlockJacobiPreconditioner`, `ICholPreconditioner` plus `compute_diagonal!` / `compute_block_diagonal!` and the `apply_constraint_*` diagonal hooks. Closes the conditioning gap of the penalty-Dirichlet matrix-free path so unpreconditioned CG still converges. - `eigensolve.jl` — `lowest_eigenpairs`, `solve_eigenproblem`. Subspace iteration with Rayleigh-Ritz built directly on `MatrixFreeOperator` and `MatrixFreeMassOperator`. Supports closure-style operators (for backward compatibility) and assembled `K`, `M` matrices. ## Design notes +Primary factories use `(cache, asm, mesh; …)` (no standalone volume `kernel`): +the cache already owns [`UniformKernelColumn`](@ref) / +[`PerElementKernelColumn`](@ref). Redundant-kernel overloads on operators, +`matrix_free_op`, `solve_eigenproblem`, `internal_force_op`, +`nonlinear_residual_op`, Jacobi / block-Jacobi / IChol constructors, and +`compute_block_diagonal!` forward to these forms and share one `Base.depwarn` per +session. + `MatrixFreeOperator` is the thread that ties this directory together. Every consumer (the eigensolver, the Krylov solves in `test/assemblers/`, downstream user code) goes through the same typed diff --git a/src/assemblers/matrix_free/operator.jl b/src/assemblers/matrix_free/operator.jl index 36932c5..60731e4 100644 --- a/src/assemblers/matrix_free/operator.jl +++ b/src/assemblers/matrix_free/operator.jl @@ -1,13 +1,13 @@ -# This file is a part of JuliaFEM. -# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md +# SPDX-FileCopyrightText: 2015-2026 Jukka Aho +# SPDX-License-Identifier: MIT """ AbstractMatrixFreeOperator Common supertype for typed matrix-free linear operators built on the -DOF-based assembler. Subtypes wrap a `DOFBasedCOOCache`, the kernel -they evaluate (`apply_K!` for the stiffness, `apply_M!` for the mass) -and any constraint hooks that should be folded into every mat-vec. +DOF-based assembler. Subtypes wrap a `DOFBasedCOOCache` (volume kernel +from `cache.kernel_column`), call [`apply_K!`](@ref) or [`apply_M!`](@ref), +and may attach constraint hooks folded into every mat-vec. The contract is: @@ -39,30 +39,36 @@ end # --------------------------------------------------------------------------- """ - MatrixFreeOperator{C, A, K, M, D, P, L} + MatrixFreeOperator{C, A, M, D, P, L} Typed matrix-free stiffness operator. -Encapsulates the `(cache, asm, kernel, mesh)` four-tuple required by -`apply_K!`, an optional `Dirichlet` constraint and an optional `MPC` -constraint, a length-`ndofs` work buffer that absorbs any non- -`Vector{Float64}` input column, and a second buffer for the 5-argument -`mul!` interface. Each `mul!` performs +Encapsulates `cache`, assembler tag, and `mesh` for [`apply_K!`](@ref). +The volume kernel is always read from `cache.kernel_column` (never from a +separate redundant field). Optional `Dirichlet` and `MPC` constraints, +length-`ndofs` work buffers for the multiply layout and the 5-argument +`mul!` interface, and nonlinear Pass~1 keywords round out the type. + +Each `mul!` performs 1. `prepare_multiply_workspace!(workbuf, x, multiply_layout)` 2. `apply_constraint_pre!(workbuf, x, dirichlet)` (if dirichlet) - 3. `apply_K!(y, cache, asm, kernel, mesh, workbuf)` - 4. `apply_constraint_post!(y, x, dirichlet)` (if dirichlet) - 5. `apply_constraint_post!(y, x, mpc)` (if mpc) + 3. `apply_K!(y, cache, asm, mesh, workbuf; configuration, …)` + 4. `apply_constraint_post!(y, x, dirichlet)` (if dirichlet) + 5. `apply_constraint_post!(y, x, mpc)` (if mpc) so the constrained operator `(K + λ·diag(eᵈ)) x`, `K_ff x[free] ⊕ x[fixed]`, etc. is materialised without ever forming `K`. +Optional fields `configuration`, `global_material_cache`, and `Δt` are +forwarded to [`apply_K!`](@ref) (Pass~1 material updates). Defaults match +the linear reference configuration. + # Examples ```julia -op = MatrixFreeOperator(cache, asm, kernel, mesh; dirichlet = bc) +op = MatrixFreeOperator(cache, asm, mesh; dirichlet = bc) mul!(y, op, x) # in-place y2 = op * x # allocating linop = LinearOperators.LinearOperator(Float64, size(op, 1), size(op, 2), @@ -71,52 +77,108 @@ linop = LinearOperators.LinearOperator(Float64, size(op, 1), size(op, 2), """ struct MatrixFreeOperator{C<:DOFBasedCOOCache, A<:DOFBasedCOOAssembler, - K<:AbstractKernel, M<:AbstractMesh, D, P, L<:AbstractMultiplyGhostLayout} <: AbstractMatrixFreeOperator cache::C asm::A - kernel::K mesh::M dirichlet::D mpc::P workbuf::Vector{Float64} mulbuf::Vector{Float64} multiply_layout::L + configuration::Union{Nothing,AbstractVector{Float64}} + global_material_cache::Union{Nothing,GlobalMaterialCache} + Δt::Float64 end -MatrixFreeOperator(cache::DOFBasedCOOCache, - asm::DOFBasedCOOAssembler, - kernel::AbstractKernel, - mesh::AbstractMesh, - dirichlet, - mpc, - workbuf::Vector{Float64}) = - MatrixFreeOperator(cache, asm, kernel, mesh, dirichlet, mpc, workbuf, similar(workbuf)) +MatrixFreeOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + mesh::AbstractMesh, + dirichlet, + mpc, + workbuf::Vector{Float64}, +) = + MatrixFreeOperator( + cache, asm, mesh, dirichlet, mpc, workbuf, similar(workbuf), + LocalMultiplyLayout(), nothing, nothing, 0.0, + ) -MatrixFreeOperator(cache::DOFBasedCOOCache, - asm::DOFBasedCOOAssembler, - kernel::AbstractKernel, - mesh::AbstractMesh, - dirichlet, - mpc, - workbuf::Vector{Float64}, - mulbuf::Vector{Float64}) = - MatrixFreeOperator(cache, asm, kernel, mesh, dirichlet, mpc, workbuf, mulbuf, LocalMultiplyLayout()) +MatrixFreeOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + mesh::AbstractMesh, + dirichlet, + mpc, + workbuf::Vector{Float64}, + mulbuf::Vector{Float64}, +) = + MatrixFreeOperator( + cache, asm, mesh, dirichlet, mpc, workbuf, mulbuf, + LocalMultiplyLayout(), nothing, nothing, 0.0, + ) -function MatrixFreeOperator(cache::DOFBasedCOOCache, - asm::DOFBasedCOOAssembler, - kernel::AbstractKernel, - mesh::AbstractMesh; - dirichlet = nothing, - mpc = nothing, - multiply_layout::L = LocalMultiplyLayout()) where L <: AbstractMultiplyGhostLayout +function MatrixFreeOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + mesh::AbstractMesh; + dirichlet = nothing, + mpc = nothing, + multiply_layout::L = LocalMultiplyLayout(), + configuration::Union{Nothing,AbstractVector{Float64}} = nothing, + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Real = 0.0, +) where {L <: AbstractMultiplyGhostLayout} workbuf = zeros(Float64, cache.ndofs) mulbuf = similar(workbuf) - return MatrixFreeOperator(cache, asm, kernel, mesh, dirichlet, mpc, workbuf, mulbuf, multiply_layout) + return MatrixFreeOperator( + cache, asm, mesh, dirichlet, mpc, workbuf, mulbuf, multiply_layout, + configuration, global_material_cache, Float64(Δt), + ) end +@inline function MatrixFreeOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh; + kwargs..., +) + _depwarn_redundant_kernel_arg!(:MatrixFreeOperator) + return MatrixFreeOperator(cache, asm, mesh; kwargs...) +end + +MatrixFreeOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh, + dirichlet, + mpc, + workbuf::Vector{Float64}, +) = + begin + _depwarn_redundant_kernel_arg!(:MatrixFreeOperator) + MatrixFreeOperator(cache, asm, mesh, dirichlet, mpc, workbuf) + end + +MatrixFreeOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh, + dirichlet, + mpc, + workbuf::Vector{Float64}, + mulbuf::Vector{Float64}, +) = + begin + _depwarn_redundant_kernel_arg!(:MatrixFreeOperator) + MatrixFreeOperator(cache, asm, mesh, dirichlet, mpc, workbuf, mulbuf) + end + @inline Base.size(op::MatrixFreeOperator) = (op.cache.ndofs, op.cache.ndofs) @inline Base.size(op::MatrixFreeOperator, d::Integer) = (d == 1 || d == 2) ? op.cache.ndofs : 1 @@ -127,7 +189,8 @@ LinearAlgebra.ishermitian(::MatrixFreeOperator) = true # Whether `K` is SPD is a property of the kernel; mixed / saddle-point # kernels override `operator_is_posdef` to `false` so Krylov stacks (CG, …) # do not pick the SPD branch. -@inline LinearAlgebra.isposdef(op::MatrixFreeOperator) = operator_is_posdef(op.kernel) +@inline LinearAlgebra.isposdef(op::MatrixFreeOperator) = + operator_is_posdef(prototype_kernel(op.cache.kernel_column)) function LinearAlgebra.mul!(y::AbstractVector{Float64}, op::MatrixFreeOperator, @@ -137,7 +200,18 @@ function LinearAlgebra.mul!(y::AbstractVector{Float64}, if op.dirichlet !== nothing apply_constraint_pre!(workbuf, x, op.dirichlet) end - apply_K!(y, op.cache, op.asm, op.kernel, op.mesh, workbuf) + # Avoid keyword `apply_K!` on the default linear path so `mul!` stays + # allocation-free after warmup (see `test_matrix_free_operator.jl`). + if op.configuration === nothing && op.global_material_cache === nothing && iszero(op.Δt) + apply_K!(y, op.cache, op.asm, op.mesh, workbuf) + else + apply_K!( + y, op.cache, op.asm, op.mesh, workbuf; + configuration = op.configuration, + global_material_cache = op.global_material_cache, + Δt = op.Δt, + ) + end if op.dirichlet !== nothing apply_constraint_post!(y, x, op.dirichlet) end @@ -173,9 +247,8 @@ function LinearAlgebra.mul!(y::AbstractVector{Float64}, end """ - matrix_free_op(cache, asm, kernel, mesh; dirichlet = nothing, - mpc = nothing) - -> MatrixFreeOperator + matrix_free_op(cache, asm, mesh; dirichlet = nothing, mpc = nothing, kwargs...) + matrix_free_op(cache, asm, kernel, mesh; …) Convenience wrapper that builds a `MatrixFreeOperator` for `K` (with the optional `dirichlet` and `mpc` constraints folded into every @@ -183,6 +256,15 @@ mat-vec). Returns the typed operator directly; the operator is callable (`op(y, x)` does an in-place mat-vec) so existing call sites that treat `matrix_free_op(...)` as a closure keep working. +The three-argument form `matrix_free_op(cache, asm, mesh; …)` is primary. +The four-argument form with a trailing `kernel` argument ignores that +kernel (backward compatibility; emits `Base.depwarn` once per session); the +kernel always comes from `cache.kernel_column`. + +Additional keywords `configuration`, `global_material_cache`, and `Δt` +are forwarded to [`MatrixFreeOperator`](@ref) and then into each +[`apply_K!`](@ref) during `mul!`. + The constraint type controls the constrained operator: * `PenaltyDirichlet` → `op(x) = K x + λ · diag(eᵈ) x` @@ -195,7 +277,7 @@ The constraint type controls the constrained operator: using LinearOperators, IterativeSolvers, JuliaFEM c = EliminatedDirichlet(fixed_dofs, û) -op = matrix_free_op(cache, asm, kernel, mesh; dirichlet = c) +op = matrix_free_op(cache, asm, mesh; dirichlet = c) linop = LinearOperator(Float64, size(op, 1), size(op, 2), true, true, op) T_mf = zeros(cache.ndofs) @@ -206,16 +288,50 @@ For new code prefer constructing `MatrixFreeOperator` directly; this helper is kept for backward compatibility and as a single-call factory that mirrors the kwargs the constraint and load APIs expect. """ -@inline function matrix_free_op(cache::DOFBasedCOOCache, - asm::DOFBasedCOOAssembler, - kernel::AbstractKernel, - mesh::AbstractMesh; - dirichlet::Union{AbstractDirichletConstraint,Nothing} = nothing, - mpc = nothing, - multiply_layout::AbstractMultiplyGhostLayout = LocalMultiplyLayout()) - return MatrixFreeOperator(cache, asm, kernel, mesh; - dirichlet = dirichlet, mpc = mpc, - multiply_layout = multiply_layout) +@inline function matrix_free_op( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + mesh::AbstractMesh; + dirichlet::Union{AbstractDirichletConstraint,Nothing} = nothing, + mpc = nothing, + multiply_layout::AbstractMultiplyGhostLayout = LocalMultiplyLayout(), + configuration::Union{Nothing,AbstractVector{Float64}} = nothing, + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Real = 0.0, +) + return MatrixFreeOperator( + cache, asm, mesh; + dirichlet = dirichlet, + mpc = mpc, + multiply_layout = multiply_layout, + configuration = configuration, + global_material_cache = global_material_cache, + Δt = Δt, + ) +end + +@inline function matrix_free_op( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh; + dirichlet::Union{AbstractDirichletConstraint,Nothing} = nothing, + mpc = nothing, + multiply_layout::AbstractMultiplyGhostLayout = LocalMultiplyLayout(), + configuration::Union{Nothing,AbstractVector{Float64}} = nothing, + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Real = 0.0, +) + _depwarn_redundant_kernel_arg!(:matrix_free_op) + return matrix_free_op( + cache, asm, mesh; + dirichlet = dirichlet, + mpc = mpc, + multiply_layout = multiply_layout, + configuration = configuration, + global_material_cache = global_material_cache, + Δt = Δt, + ) end # --------------------------------------------------------------------------- @@ -224,11 +340,11 @@ end # --------------------------------------------------------------------------- """ - MatrixFreeMassOperator{C, A, K, M} + MatrixFreeMassOperator{C, A, M} -Typed matrix-free mass operator wrapping the `(cache, asm, kernel, -mesh)` four-tuple plus a work buffer; each `mul!(y, op, x)` evaluates -`apply_M!(y, cache, asm, kernel, mesh, workbuf)`. +Typed matrix-free mass operator wrapping `cache`, assembler tag, and +`mesh`; each `mul!(y, op, x)` evaluates `apply_M!(y, cache, asm, mesh, workbuf)`. +The volume kernel is read from `cache.kernel_column`. Used by `solve_eigenproblem` so the lowest-eigenpair routine no longer needs ad-hoc `(y, x) -> apply_M!(...)` closures. A second buffer keeps @@ -236,32 +352,68 @@ the 5-argument `mul!` interface allocation-free. """ struct MatrixFreeMassOperator{C<:DOFBasedCOOCache, A<:DOFBasedCOOAssembler, - K<:AbstractKernel, M<:AbstractMesh} <: AbstractMatrixFreeOperator cache::C asm::A - kernel::K mesh::M workbuf::Vector{Float64} mulbuf::Vector{Float64} end -MatrixFreeMassOperator(cache::DOFBasedCOOCache, - asm::DOFBasedCOOAssembler, - kernel::AbstractKernel, - mesh::AbstractMesh, - workbuf::Vector{Float64}) = - MatrixFreeMassOperator(cache, asm, kernel, mesh, workbuf, similar(workbuf)) +MatrixFreeMassOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + mesh::AbstractMesh, + workbuf::Vector{Float64}, +) = + MatrixFreeMassOperator(cache, asm, mesh, workbuf, similar(workbuf)) -function MatrixFreeMassOperator(cache::DOFBasedCOOCache, - asm::DOFBasedCOOAssembler, - kernel::AbstractKernel, - mesh::AbstractMesh) +function MatrixFreeMassOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + mesh::AbstractMesh, +) workbuf = zeros(Float64, cache.ndofs) mulbuf = similar(workbuf) - return MatrixFreeMassOperator(cache, asm, kernel, mesh, workbuf, mulbuf) + return MatrixFreeMassOperator(cache, asm, mesh, workbuf, mulbuf) end +@inline MatrixFreeMassOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh, +) = + begin + _depwarn_redundant_kernel_arg!(:MatrixFreeMassOperator) + MatrixFreeMassOperator(cache, asm, mesh) + end + +MatrixFreeMassOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh, + workbuf::Vector{Float64}, +) = + begin + _depwarn_redundant_kernel_arg!(:MatrixFreeMassOperator) + MatrixFreeMassOperator(cache, asm, mesh, workbuf) + end + +MatrixFreeMassOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh, + workbuf::Vector{Float64}, + mulbuf::Vector{Float64}, +) = + begin + _depwarn_redundant_kernel_arg!(:MatrixFreeMassOperator) + MatrixFreeMassOperator(cache, asm, mesh, workbuf, mulbuf) + end + @inline Base.size(op::MatrixFreeMassOperator) = (op.cache.ndofs, op.cache.ndofs) @inline Base.size(op::MatrixFreeMassOperator, d::Integer) = (d == 1 || d == 2) ? op.cache.ndofs : 1 @@ -276,7 +428,7 @@ function LinearAlgebra.mul!(y::AbstractVector{Float64}, @inbounds @simd for i in eachindex(workbuf) workbuf[i] = x[i] end - apply_M!(y, op.cache, op.asm, op.kernel, op.mesh, workbuf) + apply_M!(y, op.cache, op.asm, op.mesh, workbuf) return y end @@ -301,3 +453,429 @@ function LinearAlgebra.mul!(y::AbstractVector{Float64}, end return y end + +# --------------------------------------------------------------------------- +# Mass operator on `DOFBasedCOOCacheKA` (KernelAbstractions path). +# --------------------------------------------------------------------------- + +""" + MatrixFreeMassOperatorKA{F, K, CKA} + +Matrix-free mass matvec on a [`DOFBasedCOOCacheKA`](@ref): each +`mul!(y, op, x)` calls [`apply_M!`](@ref)`(y, op.cache_ka, op.kernel, x)` +after copying `x` into an internal scratch (so `y` and `x` may alias the +same storage as long as `y !== x` is not required by callers; the copy matches +[`MatrixFreeMassOperator`](@ref) behaviour). + +Precision `F` matches `eltype(cache_ka.detJ_w_batch)` (typically `Float64` +or `Float32` after [`to_float32`](@ref)). Not a subtype of +[`AbstractMatrixFreeOperator`](@ref) because `eltype` is parametric and GPU +backends use `KernelAbstractions.get_backend(y)`. + +Use after Pass~1 on the CPU cache and [`sync_from_cpu!`](@ref); optionally +[`Adapt.adapt`](@ref)`(MetalBackend(), cache_ka)`, `Adapt.adapt(CUDABackend(), cache_ka)`, +`Adapt.adapt(ROCBackend(), cache_ka)`, or `Adapt.adapt(oneAPIBackend(), cache_ka)` when the +corresponding GPU package is loaded. + +# Example + +```julia +op_m = MatrixFreeMassOperatorKA(cache_ka, prototype_kernel(cpu_cache.kernel_column)) +mul!(y, op_m, x) # y, x same precision as cache_ka +``` +""" +struct MatrixFreeMassOperatorKA{F<:AbstractFloat, + K<:AbstractKernel, + CKA<:DOFBasedCOOCacheKA, + B<:AbstractVector{F}} + cache_ka::CKA + kernel::K + workbuf::B + mulbuf::B +end + +function MatrixFreeMassOperatorKA(cache_ka::DOFBasedCOOCacheKA, kernel::K) where {K<:AbstractKernel} + F = eltype(cache_ka.detJ_w_batch) + n = length(cache_ka.dof_counts) + be = KernelAbstractions.get_backend(cache_ka.detJ_w_batch) + z = Adapt.adapt(be, zeros(F, n)) + return MatrixFreeMassOperatorKA{F,K,typeof(cache_ka),typeof(z)}( + cache_ka, + kernel, + z, + similar(z), + ) +end + +@inline function matrix_free_mass_op_ka(cache_ka::DOFBasedCOOCacheKA, kernel::AbstractKernel) + return MatrixFreeMassOperatorKA(cache_ka, kernel) +end + +@inline Base.size(op::MatrixFreeMassOperatorKA) = + (length(op.cache_ka.dof_counts), length(op.cache_ka.dof_counts)) +@inline Base.size(op::MatrixFreeMassOperatorKA, d::Integer) = + (d == 1 || d == 2) ? length(op.cache_ka.dof_counts) : 1 +Base.eltype(::MatrixFreeMassOperatorKA{F}) where {F} = F + +LinearAlgebra.issymmetric(::MatrixFreeMassOperatorKA) = true +LinearAlgebra.ishermitian(::MatrixFreeMassOperatorKA) = true + +@inline function (op::MatrixFreeMassOperatorKA)(y::AbstractVector{F}, + x::AbstractVector{F}) where {F} + return LinearAlgebra.mul!(y, op, x) +end + +function Base.:*(op::MatrixFreeMassOperatorKA{F}, x::AbstractVector{F}) where {F} + y = similar(x, F, size(op, 1)) + return LinearAlgebra.mul!(y, op, x) +end + +function LinearAlgebra.mul!(y::AbstractVector{F}, + op::MatrixFreeMassOperatorKA{F}, + x::AbstractVector{F}) where {F<:AbstractFloat} + n = size(op, 1) + length(y) == n || throw(DimensionMismatch("y length $(length(y)); expected $n")) + length(x) == n || throw(DimensionMismatch("x length $(length(x)); expected $n")) + workbuf = op.workbuf + @inbounds @simd for i in eachindex(workbuf) + workbuf[i] = x[i] + end + apply_M!(y, op.cache_ka, op.kernel, workbuf) + return y +end + +function LinearAlgebra.mul!(y::AbstractVector{F}, + op::MatrixFreeMassOperatorKA{F}, + x::AbstractVector{F}, + α::Number, β::Number) where {F<:AbstractFloat} + n = size(op, 1) + if iszero(β) + fill!(y, zero(F)) + elseif !isone(β) + @inbounds @simd for i in 1:n + y[i] *= β + end + end + if !iszero(α) + scratch = op.mulbuf + LinearAlgebra.mul!(scratch, op, x) + @inbounds @simd for i in 1:n + y[i] += α * scratch[i] + end + end + return y +end + +# --------------------------------------------------------------------------- +# Stiffness operator on `DOFBasedCOOCacheKA` (KernelAbstractions path). +# --------------------------------------------------------------------------- + +""" + MatrixFreeOperatorKA{F, K, CKA, B, L} + +Matrix-free stiffness matvec on a [`DOFBasedCOOCacheKA`](@ref): each +`mul!(y, op, x)` fills the scratch with [`prepare_multiply_workspace!`](@ref) +(default [`LocalMultiplyLayout`](@ref): copy `x`), then calls +[`apply_K!`](@ref)`(y, op.cache_ka, op.kernel, workbuf)`. + +This is the KA analogue of unconstrained [`MatrixFreeOperator`](@ref) volume +matvecs: there are no Dirichlet / MPC hooks and no Pass~1 `configuration` / +`global_material_cache` / `Δt` forwarding, because the KA `apply_K!` entry +point is stiffness-only on the device batch. Use [`matrix_free_op`](@ref) +when you need constraints or nonlinear Pass~1 on the CPU cache. + +Precision `F` matches `eltype(cache_ka.detJ_w_batch)`. Not a subtype of +[`AbstractMatrixFreeOperator`](@ref); see [`MatrixFreeMassOperatorKA`](@ref). + +# Example + +```julia +op_k = MatrixFreeOperatorKA(cache_ka, prototype_kernel(cpu_cache.kernel_column)) +mul!(y, op_k, x) +``` +""" +struct MatrixFreeOperatorKA{F<:AbstractFloat, + K<:AbstractKernel, + CKA<:DOFBasedCOOCacheKA, + B<:AbstractVector{F}, + L<:AbstractMultiplyGhostLayout} + cache_ka::CKA + kernel::K + workbuf::B + mulbuf::B + multiply_layout::L +end + +function MatrixFreeOperatorKA( + cache_ka::DOFBasedCOOCacheKA, + kernel::K; + multiply_layout::L = LocalMultiplyLayout(), +) where {K<:AbstractKernel, L <: AbstractMultiplyGhostLayout} + F = eltype(cache_ka.detJ_w_batch) + n = length(cache_ka.dof_counts) + be = KernelAbstractions.get_backend(cache_ka.detJ_w_batch) + z = Adapt.adapt(be, zeros(F, n)) + return MatrixFreeOperatorKA{F,K,typeof(cache_ka),typeof(z),L}( + cache_ka, + kernel, + z, + similar(z), + multiply_layout, + ) +end + +@inline function matrix_free_op_ka( + cache_ka::DOFBasedCOOCacheKA, + kernel::AbstractKernel; + multiply_layout::L = LocalMultiplyLayout(), +) where {L <: AbstractMultiplyGhostLayout} + return MatrixFreeOperatorKA(cache_ka, kernel; multiply_layout = multiply_layout) +end + +@inline Base.size(op::MatrixFreeOperatorKA) = + (length(op.cache_ka.dof_counts), length(op.cache_ka.dof_counts)) +@inline Base.size(op::MatrixFreeOperatorKA, d::Integer) = + (d == 1 || d == 2) ? length(op.cache_ka.dof_counts) : 1 +Base.eltype(::MatrixFreeOperatorKA{F}) where {F} = F + +LinearAlgebra.issymmetric(::MatrixFreeOperatorKA) = true +LinearAlgebra.ishermitian(::MatrixFreeOperatorKA) = true + +@inline LinearAlgebra.isposdef(op::MatrixFreeOperatorKA) = + operator_is_posdef(op.kernel) + +@inline function (op::MatrixFreeOperatorKA)(y::AbstractVector{F}, + x::AbstractVector{F}) where {F} + return LinearAlgebra.mul!(y, op, x) +end + +function Base.:*(op::MatrixFreeOperatorKA{F}, x::AbstractVector{F}) where {F} + y = similar(x, F, size(op, 1)) + return LinearAlgebra.mul!(y, op, x) +end + +function LinearAlgebra.mul!(y::AbstractVector{F}, + op::MatrixFreeOperatorKA{F}, + x::AbstractVector{F}) where {F<:AbstractFloat} + n = size(op, 1) + length(y) == n || throw(DimensionMismatch("y length $(length(y)); expected $n")) + length(x) == n || throw(DimensionMismatch("x length $(length(x)); expected $n")) + workbuf = op.workbuf + prepare_multiply_workspace!(workbuf, x, op.multiply_layout) + apply_K!(y, op.cache_ka, op.kernel, workbuf) + return y +end + +function LinearAlgebra.mul!(y::AbstractVector{F}, + op::MatrixFreeOperatorKA{F}, + x::AbstractVector{F}, + α::Number, β::Number) where {F<:AbstractFloat} + n = size(op, 1) + if iszero(β) + fill!(y, zero(F)) + elseif !isone(β) + @inbounds @simd for i in 1:n + y[i] *= β + end + end + if !iszero(α) + scratch = op.mulbuf + LinearAlgebra.mul!(scratch, op, x) + @inbounds @simd for i in 1:n + y[i] += α * scratch[i] + end + end + return y +end + +# --------------------------------------------------------------------------- +# Nonlinear internal force and equilibrium residual (callable wrappers). +# --------------------------------------------------------------------------- + +""" + InternalForceOperator{C,A,M} + +Callable wrapper around [`assemble_internal_force!`](@ref): `(op)(f, u)` +zeros and fills `f` with ``f_{\\mathrm{int}}(u)`` using `configuration = u`. + +Holds optional Pass~1 `global_material_cache` and `Δt` (no separate +`configuration` field: the displacement state is always the second argument). + +Not a subtype of [`AbstractMatrixFreeOperator`](@ref); there is no `mul!` +contract for a nonlinear force map. + +`InternalForceOperator(cache, asm, kernel, mesh; …)` forwards to the +three-argument mesh form and ignores `kernel` (same deprecation hook as +[`matrix_free_op`](@ref)). +""" +struct InternalForceOperator{C<:DOFBasedCOOCache, + A<:DOFBasedCOOAssembler, + M<:AbstractMesh} + cache::C + asm::A + mesh::M + global_material_cache::Union{Nothing,GlobalMaterialCache} + Δt::Float64 +end + +function InternalForceOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + mesh::AbstractMesh; + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Real = 0.0, +) + return InternalForceOperator( + cache, asm, mesh, global_material_cache, Float64(Δt), + ) +end + +@inline function InternalForceOperator( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh; + kwargs..., +) + _depwarn_redundant_kernel_arg!(:InternalForceOperator) + return InternalForceOperator(cache, asm, mesh; kwargs...) +end + +@inline function (op::InternalForceOperator)( + f::AbstractVector{Float64}, + u::AbstractVector{Float64}, +) + return assemble_internal_force!( + f, op.cache, op.asm, op.mesh; + configuration = u, + global_material_cache = op.global_material_cache, + Δt = op.Δt, + ) +end + +""" + internal_force_op(cache, asm, mesh; global_material_cache = nothing, Δt = 0.0) + internal_force_op(cache, asm, kernel, mesh; …) + +Build an [`InternalForceOperator`](@ref). The four-argument form ignores `kernel` +(emits `Base.depwarn` once per session, same hook as [`matrix_free_op`](@ref)). +""" +@inline function internal_force_op( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + mesh::AbstractMesh; + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Real = 0.0, +) + return InternalForceOperator( + cache, asm, mesh; + global_material_cache = global_material_cache, + Δt = Δt, + ) +end + +@inline function internal_force_op( + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh; + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Real = 0.0, +) + _depwarn_redundant_kernel_arg!(:internal_force_op) + return internal_force_op( + cache, asm, mesh; + global_material_cache = global_material_cache, + Δt = Δt, + ) +end + +""" + NonlinearResidualOperator{C,A,M,F,W} + +Callable wrapper around [`nonlinear_equilibrium_residual!`](@ref): +`(op)(r, u)` sets `r = f_{\\mathrm{ext}} - f_{\\mathrm{int}}(u)`. + +References external load `f_ext` and a scratch vector `work` (same length as +`cache.ndofs`) owned by the operator or supplied at construction. +""" +struct NonlinearResidualOperator{C<:DOFBasedCOOCache, + A<:DOFBasedCOOAssembler, + M<:AbstractMesh, + F<:AbstractVector{Float64}, + W<:AbstractVector{Float64}} + cache::C + asm::A + mesh::M + f_ext::F + work::W + global_material_cache::Union{Nothing,GlobalMaterialCache} + Δt::Float64 +end + +function NonlinearResidualOperator( + f_ext::AbstractVector{Float64}, + work::AbstractVector{Float64}, + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + mesh::AbstractMesh; + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Real = 0.0, +) + nd = cache.ndofs + length(f_ext) == nd || + throw(DimensionMismatch("f_ext length $(length(f_ext)); expected $nd")) + length(work) == nd || + throw(DimensionMismatch("work length $(length(work)); expected $nd")) + return NonlinearResidualOperator( + cache, asm, mesh, f_ext, work, global_material_cache, Float64(Δt), + ) +end + +""" + nonlinear_residual_op(f_ext, cache, asm, mesh; kwargs…) + nonlinear_residual_op(f_ext, cache, asm, kernel, mesh; kwargs…) + +Build a [`NonlinearResidualOperator`](@ref). The five-argument form ignores +`kernel` (emits `Base.depwarn` once per session, same hook as [`matrix_free_op`](@ref)). +""" +function nonlinear_residual_op( + f_ext::AbstractVector{Float64}, + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + mesh::AbstractMesh; + work::Union{Nothing,AbstractVector{Float64}} = nothing, + global_material_cache::Union{Nothing,GlobalMaterialCache} = nothing, + Δt::Real = 0.0, +) + nd = cache.ndofs + w = work === nothing ? zeros(Float64, nd) : work + return NonlinearResidualOperator( + f_ext, w, cache, asm, mesh; + global_material_cache = global_material_cache, + Δt = Δt, + ) +end + +@inline function nonlinear_residual_op( + f_ext::AbstractVector{Float64}, + cache::DOFBasedCOOCache, + asm::DOFBasedCOOAssembler, + ::AbstractKernel, + mesh::AbstractMesh; + kwargs..., +) + _depwarn_redundant_kernel_arg!(:nonlinear_residual_op) + return nonlinear_residual_op(f_ext, cache, asm, mesh; kwargs...) +end + +@inline function (op::NonlinearResidualOperator)( + r::AbstractVector{Float64}, + u::AbstractVector{Float64}, +) + return nonlinear_equilibrium_residual!( + r, op.f_ext, op.work, op.cache, op.asm, op.mesh, u; + global_material_cache = op.global_material_cache, + Δt = op.Δt, + ) +end