From c31ab20ef43fb511c1a9db10a863e66664efdfc9 Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Wed, 25 Feb 2026 13:00:42 -0600 Subject: [PATCH 01/50] Add BatchExaModel with NLPModels batch API --- Project.toml | 2 +- docs/make.jl | 2 + docs/src/batch.jl | 133 +++++++++ ext/ExaModelsKernelAbstractions.jl | 8 + src/ExaModels.jl | 5 +- src/batch.jl | 433 +++++++++++++++++++++++++++ test/BatchTest/BatchTest.jl | 454 +++++++++++++++++++++++++++++ test/Project.toml | 1 + test/runtests.jl | 3 + 9 files changed, 1039 insertions(+), 2 deletions(-) create mode 100644 docs/src/batch.jl create mode 100644 src/batch.jl create mode 100644 test/BatchTest/BatchTest.jl diff --git a/Project.toml b/Project.toml index fa522f0c2..ff506c38c 100644 --- a/Project.toml +++ b/Project.toml @@ -43,7 +43,7 @@ JuMP = "1" KernelAbstractions = "0.9" MadNLP = "0.9" MathOptInterface = "1.19" -NLPModels = "0.21" +NLPModels = "0.21.10" NLPModelsIpopt = "0.11" OpenCL = "0.10" SolverCore = "0.3" diff --git a/docs/make.jl b/docs/make.jl index fd32c61da..0c1cae3c3 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -14,6 +14,7 @@ if !(@isdefined _PAGES) "gpu.md", "parameters.md", "two_stage.md", + "batch.md", "develop.md", "quad.md", "distillation.md", @@ -36,6 +37,7 @@ if !(@isdefined _JL_FILENAMES) "performance.jl", "parameters.jl", "two_stage.jl", + "batch.jl", ] end diff --git a/docs/src/batch.jl b/docs/src/batch.jl new file mode 100644 index 000000000..1f7f9d9d4 --- /dev/null +++ b/docs/src/batch.jl @@ -0,0 +1,133 @@ +# # [Batch Optimization](@id batch) +# ExaModels supports batch optimization through the `BatchExaModel`. This feature +# enables efficient evaluation of multiple fully independent optimization scenarios +# that share identical structure but differ in parameter values. +# +# Unlike `TwoStageExaModel`, which couples scenarios through shared design variables, +# `BatchExaModel` treats each scenario as completely independent. The key advantage +# is that all scenarios share one compiled expression pattern and are fused into a +# single model for efficient SIMD evaluation. + +# ## Problem Formulation +# A batch optimization problem solves `ns` independent scenarios simultaneously: +# ```math +# \begin{aligned} +# \min_{v_i} \quad & f(v_i; \theta_i), \quad i = 1, \ldots, S \\ +# \text{s.t.} \quad & h(v_i; \theta_i) = 0 \\ +# & v_i \in \mathcal{V} +# \end{aligned} +# ``` +# where each scenario has the same structure but different parameters $\theta_i$. + +# ## Building a Batch Model +# The builder function defines expressions for a **single scenario**. `BatchExaModel` +# calls it `ns` times internally with per-scenario parameter handles. +# This means you never have to compute global index offsets manually. + +using ExaModels, MadNLP + +# Define the problem dimensions and scenario parameters as a matrix of size `(nθ, ns)`: +ns = 3 ## number of scenarios +nv = 1 ## variables per scenario +θ_data = [2.0 4.0 6.0] ## (1, 3) matrix: θ₁=2, θ₂=4, θ₃=6 + +# Build the model. First create an `ExaCore`, then pass it along with the parameter +# matrix to `BatchExaModel`: +c = ExaCore() +model = BatchExaModel(c, ns, θ_data) do c, θ + ## Create variables — this is called once per scenario, offsets are automatic + v = variable(c, nv) + ## Objective: minimize (v - θ)² + objective(c, (v[1] - θ[1])^2) + ## Constraint: v ≥ 0 + constraint(c, v[1]; lcon = 0.0, ucon = Inf) +end + +# The builder function receives: +# - `c`: the `ExaCore` — use `variable(c, ...)`, `objective(c, ...)`, `constraint(c, ...)` as usual +# - `θ`: a per-scenario parameter handle (indices 1:nθ) +# +# Variable creation via `variable(c, ...)` works exactly like in a regular `ExaModel`. +# You can set start values, lower/upper bounds, etc. + +# ## Batch API (NLPModels) +# `BatchExaModel` implements the `AbstractBatchNLPModel` interface from NLPModels.jl. +# All evaluation functions use matrices of size `(dim, ns)`: +import NLPModels + +println("Variables per scenario: ", NLPModels.get_nvar(model)) +println("Constraints per scenario: ", NLPModels.get_ncon(model)) +println("Number of scenarios: ", NLPModels.get_nbatch(model)) + +# Evaluate objectives for all scenarios at once: +bx = reshape([1.0, 3.0, 5.0], nv, ns) +bf = zeros(ns) +NLPModels.obj!(model, bx, bf) +println("\nObjective values: ", bf) +## scenario 1: (1-2)² = 1, scenario 2: (3-4)² = 1, scenario 3: (5-6)² = 1 + +# Evaluate gradients: +bg = zeros(nv, ns) +NLPModels.grad!(model, bx, bg) +println("Gradients: ", bg) + +# Evaluate constraints: +bc = zeros(NLPModels.get_ncon(model), ns) +NLPModels.cons!(model, bx, bc) +println("Constraints: ", bc) + +# ## Solving via the Fused Model +# For solving, access the underlying fused `ExaModel` and use any NLPModels-compatible +# solver (e.g., MadNLP): +result = madnlp(ExaModels.get_model(model); print_level = MadNLP.ERROR) +println("\nSolution status: ", result.status) +println("Optimal objective: ", round(result.objective, digits = 4)) + +# Extract per-scenario solutions: +x_sol = result.solution +for i in 1:ns + v_sol = x_sol[ExaModels.var_indices(model, i)] + println("Scenario $i: v* = ", round(v_sol[1], digits = 4)) +end + +# ## A More Complex Example +# Here's a batch model with multiple variables, objectives, and constraints per scenario: +ns2, nv2 = 2, 3 +θ_data2 = [1.0 4.0; 2.0 5.0; 3.0 6.0] ## (3, 2) matrix + +c2 = ExaCore() +model2 = BatchExaModel(c2, ns2, θ_data2) do c, θ + v = variable(c, nv2; start = 1.0, lvar = 0.0, uvar = 10.0) + ## Objective: Σⱼ (vⱼ - θⱼ)² + objective(c, (v[j] - θ[j])^2 for j in 1:nv2) + ## Constraints: sum of all variables ≤ 20 + constraint(c, sum(v[j] for j in 1:nv2); ucon = 20.0) +end + +result2 = madnlp(ExaModels.get_model(model2); print_level = MadNLP.ERROR) +println("\nMulti-variable example:") +println("Status: ", result2.status) +x_sol2 = result2.solution +for i in 1:ns2 + v_sol = x_sol2[ExaModels.var_indices(model2, i)] + println("Scenario $i: v* = ", round.(v_sol, digits = 4)) +end + +# ## Updating Parameters +# You can update scenario parameters and re-solve without rebuilding the model: + +# Update a single scenario: +ExaModels.set_scenario_parameters!(model, 1, [10.0]) + +# Or update all scenarios at once: +ExaModels.set_all_scenario_parameters!(model, [[10.0], [12.0], [14.0]]) + +# Re-solve with new parameters: +result3 = madnlp(ExaModels.get_model(model); print_level = MadNLP.ERROR) +println("\nAfter parameter update:") +println("Status: ", result3.status) +x_sol3 = result3.solution +for i in 1:ns + v_sol = x_sol3[ExaModels.var_indices(model, i)] + println("Scenario $i: v* = ", round(v_sol[1], digits = 4)) +end diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index 111f4b759..d2d1adacb 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -204,6 +204,14 @@ function _obj(backend, objbuffer, obj, x, θ) end function _obj(backend, objbuffer, f::ExaModels.ObjectiveNull, x, θ) end +function ExaModels._eval_objbuffer!( + objbuffer, m::ExaModels.ExaModel{T, VT, E}, x + ) where {T, VT, E <: KAExtension} + return if !isempty(objbuffer) + _obj(m.ext.backend, objbuffer, m.objs, x, m.θ) + end +end + function ExaModels.cons_nln!( m::ExaModels.AbstractExaModel{T,VT,E}, x::AbstractVector, diff --git a/src/ExaModels.jl b/src/ExaModels.jl index f5d0b5390..ca29c8f7b 100644 --- a/src/ExaModels.jl +++ b/src/ExaModels.jl @@ -10,6 +10,7 @@ module ExaModels import NLPModels: NLPModels, obj, + obj!, cons!, grad!, jac_coord!, @@ -37,6 +38,7 @@ include("hessian.jl") include("nlp.jl") include("tags.jl") include("utils.jl") +include("batch.jl") export ExaModel, ExaCore, @@ -57,6 +59,7 @@ export ExaModel, multipliers_L, multipliers_U, @register_univariate, - @register_bivariate + @register_bivariate, + BatchExaModel end # module ExaModels diff --git a/src/batch.jl b/src/batch.jl new file mode 100644 index 000000000..23f1094bd --- /dev/null +++ b/src/batch.jl @@ -0,0 +1,433 @@ +# ============================================================================ +# BatchExaModel — single-scenario builder + NLPModels matrix batch API +# ============================================================================ + +""" + BatchExaModel{T, VT, MT, M} <: NLPModels.AbstractBatchNLPModel{T, MT} + +Parametric optimization model where multiple fully independent scenarios are fused +into a single ExaModel and evaluated simultaneously using shared compiled expression +patterns. + +All scenarios share identical sparsity structures. The builder defines expressions +for **one scenario**; `BatchExaModel` calls it `ns` times internally with offset +variable/parameter handles. + +# Dimensions + +- `ns`: number of scenarios (batch size) +- `nv`: number of variables per scenario +- `nc`: number of constraints per scenario +- `nθ`: number of parameters per scenario +""" +struct BatchExaModel{T, VT <: AbstractVector{T}, MT <: AbstractMatrix{T}, M} <: + NLPModels.AbstractBatchNLPModel{T, MT} + meta::NLPModels.BatchNLPModelMeta{T, MT} + model::M + objbuffer::VT + hess_perm::Vector{Int} + hess_buffer::VT + ns::Int + nv::Int + nc::Int + nθ::Int + nobj_per::Int + nnzj_per::Int + nnzh_per::Int + nnzh_obj_per::Int + nnzh_con_per::Int +end + +function Base.show(io::IO, m::BatchExaModel{T, VT}) where {T, VT} + println(io, "BatchExaModel{$T, $VT}") + println(io, " Scenarios: $(m.ns)") + println(io, " Variables per scenario: $(m.nv)") + println(io, " Constraints per scenario: $(m.nc)") + println(io, " Parameters per scenario: $(m.nθ)") + println(io, " Total variables: $(m.ns * m.nv)") + println(io, " Total constraints: $(m.ns * m.nc)") + println(io, " Jacobian nnz per scenario: $(m.nnzj_per)") + return println(io, " Hessian nnz per scenario: $(m.nnzh_per)") +end + +# ============================================================================ +# Helpers +# ============================================================================ + +_count_hess_nnz(::ObjectiveNull) = 0 +_count_hess_nnz(::ConstraintNull) = 0 +_count_hess_nnz(node) = _count_hess_nnz(node.inner) + node.f.o2step * length(node.itr) + +function _build_hess_perm(ns, nnzh_obj_per, nnzh_con_per) + nnzh_per = nnzh_obj_per + nnzh_con_per + perm = Vector{Int}(undef, ns * nnzh_per) + for s in 1:ns + base = (s - 1) * nnzh_per + for k in 1:nnzh_obj_per + perm[base + k] = (s - 1) * nnzh_obj_per + k + end + for k in 1:nnzh_con_per + perm[base + nnzh_obj_per + k] = ns * nnzh_obj_per + (s - 1) * nnzh_con_per + k + end + end + return perm +end + +function _to_matrix(v::AbstractVector, nrows::Int, ncols::Int) + mat = similar(v, nrows, ncols) + copyto!(vec(mat), v) + return mat +end + +# ============================================================================ +# Constructor +# ============================================================================ + +""" + BatchExaModel(build, c::ExaCore, ns::Int, θ_data::AbstractMatrix) + +Build a batch model from a single-scenario builder function. + +The user creates an `ExaCore` and passes parameter data as a matrix of size `(nθ, ns)`. +The builder defines a **single scenario** — creating variables, objectives, and +constraints using standard ExaModels calls. `BatchExaModel` invokes it `ns` times, +each time with a per-scenario parameter handle `θ`. Variable creation via +`variable(c, ...)` works normally and automatically gets the correct offsets. + +# Arguments +- `build::Function`: Function `(c, θ) -> nothing` + - `c`: ExaCore — use `variable(c, ...)`, `objective(c, ...)`, `constraint(c, ...)` as usual + - `θ`: Parameter handle for this scenario's parameters (indices 1:nθ) +- `c::ExaCore`: ExaCore instance (parameters will be registered internally) +- `ns::Int`: Number of scenarios +- `θ_data::AbstractMatrix`: Parameter matrix of size `(nθ, ns)` + +# Example +```julia +ns, nθ = 100, 3 +θ_data = rand(nθ, ns) + +c = ExaCore() +model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, 5; start = 1.0, lvar = 0.0, uvar = 10.0) + objective(c, θ[j] * v[j]^2 for j in 1:5) + constraint(c, v[j] - θ[j] for j in 1:3) +end +``` +""" +function BatchExaModel( + build::Function, + c::ExaCore, + ns::Int, + θ_data::AbstractMatrix, + ) + size(θ_data, 2) == ns || throw( + ArgumentError("θ_data must have ns=$ns columns, got $(size(θ_data, 2))"), + ) + nθ = size(θ_data, 1) + + # Register parameters as a flat vector (column-major: scenario 1, scenario 2, ...) + parameter(c, vec(θ_data)) + + # Call builder once per scenario with per-scenario θ handle. + # The builder calls variable(c, ...) itself — offsets are automatic. + for s in 1:ns + θ_s = Parameter(nθ, nθ, (s - 1) * nθ) + build(c, θ_s) + end + + # Infer per-scenario dimensions + nv = c.nvar ÷ ns + nv * ns != c.nvar && throw( + DimensionMismatch( + "Total variables ($(c.nvar)) not evenly divisible by ns ($ns)", + ), + ) + + nc_total = c.ncon + nc = nc_total ÷ ns + nc * ns != nc_total && throw( + DimensionMismatch( + "Total constraints ($nc_total) not evenly divisible by ns ($ns)", + ), + ) + + nobj_total = c.nobj + nobj_per = nobj_total ÷ ns + nobj_per * ns != nobj_total && throw( + DimensionMismatch( + "Total objective entries ($nobj_total) not evenly divisible by ns ($ns)", + ), + ) + + objbuffer = similar(c.x0, nobj_total) + + model = ExaModel(c) + + # Per-scenario sparsity counts + total_nnzj = NLPModels.get_nnzj(model) + total_nnzh = NLPModels.get_nnzh(model) + nnzj_per = total_nnzj ÷ ns + nnzh_per = total_nnzh ÷ ns + + # Obj/con hessian split + nnzh_obj_total = _count_hess_nnz(model.objs) + nnzh_con_total = _count_hess_nnz(model.cons) + nnzh_obj_per = nnzh_obj_total ÷ ns + nnzh_con_per = nnzh_con_total ÷ ns + + # Hessian permutation + hess_perm = _build_hess_perm(ns, nnzh_obj_per, nnzh_con_per) + + # Hessian buffer + hess_buffer = similar(c.x0, total_nnzh) + + # Build BatchNLPModelMeta with matrices + T = eltype(c.x0) + x0_mat = _to_matrix(model.meta.x0, nv, ns) + lvar_mat = _to_matrix(model.meta.lvar, nv, ns) + uvar_mat = _to_matrix(model.meta.uvar, nv, ns) + y0_mat = _to_matrix(model.meta.y0, nc, ns) + lcon_mat = _to_matrix(model.meta.lcon, nc, ns) + ucon_mat = _to_matrix(model.meta.ucon, nc, ns) + + MT = typeof(x0_mat) + meta = NLPModels.BatchNLPModelMeta{T, MT}( + ns, + nv; + x0 = x0_mat, + lvar = lvar_mat, + uvar = uvar_mat, + ncon = nc, + y0 = y0_mat, + lcon = lcon_mat, + ucon = ucon_mat, + nnzj = nnzj_per, + nnzh = nnzh_per, + minimize = model.meta.minimize, + ) + + VT = typeof(c.x0) + return BatchExaModel{T, VT, MT, typeof(model)}( + meta, + model, + objbuffer, + hess_perm, + hess_buffer, + ns, + nv, + nc, + nθ, + nobj_per, + nnzj_per, + nnzh_per, + nnzh_obj_per, + nnzh_con_per, + ) +end + +# ============================================================================ +# Accessors +# ============================================================================ + +num_scenarios(m::BatchExaModel) = m.ns + +""" + get_model(model::BatchExaModel) + +Get the underlying fused ExaModel for direct NLPModels interface usage (e.g. Ipopt). +""" +get_model(model::BatchExaModel) = model.model + +""" + var_indices(model::BatchExaModel, i) -> UnitRange + +Index range for variables of scenario `i` in the global (fused) variable vector. +""" +function var_indices(model::BatchExaModel, i::Int) + nv = model.nv + return ((i - 1) * nv + 1):(i * nv) +end + +""" + cons_block_indices(model::BatchExaModel, i) -> UnitRange + +Index range for constraints of scenario `i` in the global (fused) constraint vector. +""" +function cons_block_indices(model::BatchExaModel, i::Int) + nc = model.nc + return ((i - 1) * nc + 1):(i * nc) +end + +# ============================================================================ +# Parameter Updates +# ============================================================================ + +function set_scenario_parameters!(model::BatchExaModel, i::Int, θ_new::AbstractVector) + nθ = model.nθ + length(θ_new) != nθ && throw( + DimensionMismatch("Parameter size mismatch: expected $nθ, got $(length(θ_new))"), + ) + θ_start = (i - 1) * nθ + 1 + θ_end = i * nθ + copyto!(view(model.model.θ, θ_start:θ_end), θ_new) + return nothing +end + +function set_all_scenario_parameters!( + model::BatchExaModel, + θ_sets::Vector{<:AbstractVector}, + ) + length(θ_sets) == model.ns || + throw(ArgumentError("θ_sets must have length $(model.ns)")) + for i in 1:(model.ns) + set_scenario_parameters!(model, i, θ_sets[i]) + end + return nothing +end + +# ============================================================================ +# Objective buffer evaluation (CPU) +# ============================================================================ + +function _eval_objbuffer!(objbuffer, objs, x, θ) + _eval_objbuffer!(objbuffer, objs.inner, x, θ) + for i in eachindex(objs.itr) + objbuffer[offset0(objs, i)] = objs.f(objs.itr[i], x, θ) + end + return +end +_eval_objbuffer!(objbuffer, ::ObjectiveNull, x, θ) = nothing + +function _eval_objbuffer!(objbuffer, m::ExaModel, x) + return _eval_objbuffer!(objbuffer, m.objs, x, m.θ) +end + +# ============================================================================ +# Batch API: obj! +# ============================================================================ + +function obj!(m::BatchExaModel, bx::AbstractMatrix, bf::AbstractVector) + _eval_objbuffer!(m.objbuffer, m.model, vec(bx)) + obj_mat = reshape(m.objbuffer, m.nobj_per, m.ns) + bf .= vec(sum(obj_mat; dims = 1)) + return bf +end + +# ============================================================================ +# Batch API: grad! +# ============================================================================ + +function grad!(m::BatchExaModel, bx::AbstractMatrix, bg::AbstractMatrix) + grad!(m.model, vec(bx), vec(bg)) + return bg +end + +# ============================================================================ +# Batch API: cons! +# ============================================================================ + +function cons!(m::BatchExaModel, bx::AbstractMatrix, bc::AbstractMatrix) + cons_nln!(m.model, vec(bx), vec(bc)) + return bc +end + +# ============================================================================ +# Batch API: jac_structure! +# ============================================================================ + +function jac_structure!( + m::BatchExaModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, + ) + total_nnzj = NLPModels.get_nnzj(m.model) + full_rows = zeros(Int, total_nnzj) + full_cols = zeros(Int, total_nnzj) + jac_structure!(m.model, full_rows, full_cols) + + # Scenario 1 uses zero offset → its entries are already local (1:nc, 1:nv) + for k in 1:(m.nnzj_per) + rows[k] = full_rows[k] + cols[k] = full_cols[k] + end + return rows, cols +end + +# ============================================================================ +# Batch API: jac_coord! +# ============================================================================ + +function jac_coord!(m::BatchExaModel, bx::AbstractMatrix, bjvals::AbstractMatrix) + jac_coord!(m.model, vec(bx), vec(bjvals)) + return bjvals +end + +# ============================================================================ +# Batch API: hess_structure! +# ============================================================================ + +function hess_structure!( + m::BatchExaModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, + ) + total_nnzh = NLPModels.get_nnzh(m.model) + full_rows = zeros(Int, total_nnzh) + full_cols = zeros(Int, total_nnzh) + hess_structure!(m.model, full_rows, full_cols) + + # Extract scenario 1's entries using hess_perm (interleaves obj+con) + for k in 1:(m.nnzh_per) + idx = m.hess_perm[k] + rows[k] = full_rows[idx] + cols[k] = full_cols[idx] + end + return rows, cols +end + +# ============================================================================ +# Batch API: hess_coord! +# ============================================================================ + +function hess_coord!( + m::BatchExaModel, + bx::AbstractMatrix, + by::AbstractMatrix, + bobj_weight::AbstractVector, + bhvals::AbstractMatrix, + ) + x_flat = vec(bx) + y_flat = vec(by) + nh = m.nnzh_per + perm = m.hess_perm + + if allequal(bobj_weight) + # Common case: uniform obj_weight → single fused call + permute + w = bobj_weight[1] + hess_coord!(m.model, x_flat, y_flat, m.hess_buffer; obj_weight = w) + bhvals_flat = vec(bhvals) + for i in eachindex(perm) + bhvals_flat[i] = m.hess_buffer[perm[i]] + end + else + # Varying weights: 2-pass approach + # Pass 1: objective hessian only (y=0, obj_weight=1) + y_zero = similar(y_flat) + fill!(y_zero, zero(eltype(y_flat))) + hess_obj = similar(m.hess_buffer) + hess_coord!(m.model, x_flat, y_zero, hess_obj; obj_weight = one(eltype(x_flat))) + + # Pass 2: constraint hessian only (obj_weight=0) + hess_coord!(m.model, x_flat, y_flat, m.hess_buffer; obj_weight = zero(eltype(x_flat))) + + # Combine per scenario + bhvals_flat = vec(bhvals) + for i in eachindex(perm) + s = (i - 1) ÷ nh + 1 + bhvals_flat[i] = + bobj_weight[s] * hess_obj[perm[i]] + m.hess_buffer[perm[i]] + end + end + return bhvals +end diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl new file mode 100644 index 000000000..a2420ac20 --- /dev/null +++ b/test/BatchTest/BatchTest.jl @@ -0,0 +1,454 @@ +module BatchTest + +using Test +using ExaModels +import NLPModels +import NLPModels: + obj, obj!, cons!, cons_nln!, grad!, jac_coord!, hess_coord!, jac_structure!, + hess_structure! +import ExaModels: + num_scenarios, set_scenario_parameters!, set_all_scenario_parameters!, + var_indices, cons_block_indices, get_model + +import NLPModelsIpopt: ipopt + +import ..BACKENDS +using Adapt + +function runtests() + return @testset "BatchExaModel" begin + + @testset "Construction and dimensions" begin + ns, nv = 3, 2 + nθ = 2 + θ_data = [1.0 3.0 5.0; 2.0 4.0 6.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, θ[j] * v[j]^2 for j in 1:nv) + constraint(c, v[j] for j in 1:nv) + end + + @test num_scenarios(model) == 3 + @test NLPModels.get_nvar(model) == nv + @test NLPModels.get_ncon(model) == nv + @test NLPModels.get_nbatch(model) == ns + @test model.nv == 2 + @test model.nc == 2 + end + + @testset "Batch obj! evaluation" begin + ns, nv = 2, 2 + θ_data = [2.0 3.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, θ[1] * v[j]^2 for j in 1:nv) + constraint(c, v[j] for j in 1:nv) + end + + # bx: (nv, ns) matrix + bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) + + bf = zeros(ns) + obj!(model, bx, bf) + + # scenario1: θ=2, v=[1,2], obj = 2*(1 + 4) = 10 + # scenario2: θ=3, v=[3,4], obj = 3*(9 + 16) = 75 + @test bf[1] ≈ 10.0 + @test bf[2] ≈ 75.0 + + # Consistency: sum(bf) ≈ obj(get_model(m), vec(bx)) + @test sum(bf) ≈ obj(get_model(model), vec(bx)) + + # Convenience obj() also works + bf2 = obj(model, bx) + @test bf2 ≈ bf + end + + @testset "Batch grad! evaluation" begin + ns, nv = 2, 2 + θ_data = [2.0 3.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, θ[1] * v[j]^2 for j in 1:nv) + constraint(c, v[j] for j in 1:nv) + end + + bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) + bg = zeros(nv, ns) + grad!(model, bx, bg) + + # ∂(θ*v²)/∂v = 2*θ*v + # s1: [2*2*1, 2*2*2] = [4, 8] + # s2: [2*3*3, 2*3*4] = [18, 24] + @test bg[:, 1] ≈ [4.0, 8.0] + @test bg[:, 2] ≈ [18.0, 24.0] + + # Consistency with fused model + g_flat = zeros(ns * nv) + grad!(get_model(model), vec(bx), g_flat) + @test vec(bg) ≈ g_flat + end + + @testset "Batch cons! evaluation" begin + ns, nv = 2, 2 + θ_data = [1.0 2.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, v[j]^2 for j in 1:nv) + constraint(c, v[j] - θ[1] for j in 1:nv) + end + + bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) + bc = zeros(nv, ns) + cons!(model, bx, bc) + + # s1: v=[1,2], θ=1 → [0, 1] + # s2: v=[3,4], θ=2 → [1, 2] + @test bc[:, 1] ≈ [0.0, 1.0] + @test bc[:, 2] ≈ [1.0, 2.0] + + # Consistency with fused model + c_flat = zeros(ns * nv) + cons_nln!(get_model(model), vec(bx), c_flat) + @test vec(bc) ≈ c_flat + end + + @testset "Batch jac_structure! and jac_coord!" begin + ns, nv = 2, 2 + θ_data = [1.0 2.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, v[j]^2 for j in 1:nv) + constraint(c, v[j] for j in 1:nv) + end + + nnzj = NLPModels.get_nnzj(model) + @test nnzj > 0 + + rows = zeros(Int, nnzj) + cols = zeros(Int, nnzj) + jac_structure!(model, rows, cols) + + # Per-scenario local indices: rows ∈ 1:nc, cols ∈ 1:nv + @test all(r -> 1 <= r <= model.nc, rows) + @test all(c -> 1 <= c <= model.nv, cols) + + # Evaluate Jacobian: bjvals is (nnzj, ns) + bx = reshape(ones(nv * ns), nv, ns) + bjvals = zeros(nnzj, ns) + jac_coord!(model, bx, bjvals) + + # Linear constraints → all values should be 1 + @test all(v -> v ≈ 1.0, bjvals) + end + + @testset "Batch hess_structure! and hess_coord!" begin + ns, nv = 2, 2 + θ_data = [2.0 3.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, θ[1] * v[j]^2 for j in 1:nv) + constraint(c, v[j] for j in 1:nv) + end + + nnzh = NLPModels.get_nnzh(model) + @test nnzh > 0 + + rows = zeros(Int, nnzh) + cols = zeros(Int, nnzh) + hess_structure!(model, rows, cols) + + # Per-scenario local indices + @test all(r -> 1 <= r <= model.nv, rows) + @test all(c -> 1 <= c <= model.nv, cols) + + # Evaluate with uniform obj_weight + bx = reshape(ones(nv * ns), nv, ns) + by = zeros(model.nc, ns) + bobj_weight = ones(ns) + bhvals = zeros(nnzh, ns) + hess_coord!(model, bx, by, bobj_weight, bhvals) + + # Hessian of θ*v[j]^2 is 2*θ on diagonal + # s1: 2*2 = 4, s2: 2*3 = 6 + @test any(v -> v ≈ 4.0, bhvals[:, 1]) + @test any(v -> v ≈ 6.0, bhvals[:, 2]) + end + + @testset "hess_coord! with varying obj_weight" begin + ns, nv = 2, 2 + θ_data = [2.0 3.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, θ[1] * v[j]^2 for j in 1:nv) + constraint(c, v[j]^3 for j in 1:nv) + end + + nnzh = NLPModels.get_nnzh(model) + bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) + by = ones(model.nc, ns) + + # Uniform weight for reference + bhvals_uniform = zeros(nnzh, ns) + hess_coord!(model, bx, by, [1.0, 1.0], bhvals_uniform) + + # Varying weights + bhvals_varying = zeros(nnzh, ns) + hess_coord!(model, bx, by, [2.0, 0.5], bhvals_varying) + + # Compute reference via fused model for each scenario + # With varying weights, obj part is scaled differently per scenario + inner = get_model(model) + total_nnzh = NLPModels.get_nnzh(inner) + + # obj-only hessian + hess_obj = zeros(total_nnzh) + hess_coord!(inner, vec(bx), zeros(ns * model.nc), hess_obj; obj_weight = 1.0) + + # con-only hessian + hess_con = zeros(total_nnzh) + hess_coord!(inner, vec(bx), vec(by), hess_con; obj_weight = 0.0) + + # Verify per-scenario reconstruction + perm = model.hess_perm + for s in 1:ns + for k in 1:nnzh + idx = perm[(s - 1) * nnzh + k] + expected = [2.0, 0.5][s] * hess_obj[idx] + hess_con[idx] + @test bhvals_varying[k, s] ≈ expected + end + end + end + + @testset "Multiple constraint() calls" begin + ns, nv = 2, 3 + θ_data = reshape([1.0, 2.0], 1, ns) + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, θ[1] * v[j]^2 for j in 1:nv) + # Two separate constraint() calls per scenario + constraint(c, v[j] - θ[1] for j in 1:nv) + constraint(c, v[1] + v[2] + v[3]; ucon = 10.0) + end + + # nc = nv + 1 = 4 per scenario + @test model.nc == nv + 1 + @test NLPModels.get_ncon(model) == nv + 1 + @test NLPModels.get_nbatch(model) == ns + + bx = reshape([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], nv, ns) + bc = zeros(model.nc, ns) + cons!(model, bx, bc) + + # s1: v=[1,2,3], θ=1 → [1-1, 2-1, 3-1, 1+2+3] = [0, 1, 2, 6] + # s2: v=[4,5,6], θ=2 → [4-2, 5-2, 6-2, 4+5+6] = [2, 3, 4, 15] + @test bc[:, 1] ≈ [0.0, 1.0, 2.0, 6.0] + @test bc[:, 2] ≈ [2.0, 3.0, 4.0, 15.0] + + # Consistency with fused model + c_flat = zeros(model.nc * ns) + cons_nln!(get_model(model), vec(bx), c_flat) + @test vec(bc) ≈ c_flat + + # Jacobian structure should have local indices + nnzj = NLPModels.get_nnzj(model) + rows = zeros(Int, nnzj) + cols = zeros(Int, nnzj) + jac_structure!(model, rows, cols) + @test all(r -> 1 <= r <= model.nc, rows) + @test all(c -> 1 <= c <= model.nv, cols) + + # Hessian with both obj and con contributions + nnzh = NLPModels.get_nnzh(model) + hrows = zeros(Int, nnzh) + hcols = zeros(Int, nnzh) + hess_structure!(model, hrows, hcols) + @test all(r -> 1 <= r <= model.nv, hrows) + @test all(c -> 1 <= c <= model.nv, hcols) + + # Evaluate hessian — both obj and con have nonzero second derivatives + by = ones(model.nc, ns) + bobj_weight = ones(ns) + bhvals = zeros(nnzh, ns) + hess_coord!(model, bx, by, bobj_weight, bhvals) + @test any(v -> v != 0.0, bhvals[:, 1]) + @test any(v -> v != 0.0, bhvals[:, 2]) + end + + @testset "Parameter updates" begin + ns, nv = 2, 2 + θ_data = [1.0 2.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, θ[1] * v[j]^2 for j in 1:nv) + constraint(c, v[j] for j in 1:nv) + end + + x_global = ones(ns * nv) + + # Initial: θ1=1, θ2=2, so obj = 1*2 + 2*2 = 6 + @test obj(get_model(model), x_global) ≈ 6.0 + + # Update scenario 1: θ1 = 5 + set_scenario_parameters!(model, 1, [5.0]) + @test obj(get_model(model), x_global) ≈ 14.0 + + # set_all_scenario_parameters! + set_all_scenario_parameters!(model, [[10.0], [20.0]]) + @test obj(get_model(model), x_global) ≈ 60.0 + end + + @testset "Get underlying model" begin + ns, nv = 2, 2 + θ_data = [1.0 2.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, v[j]^2 for j in 1:nv) + constraint(c, v[j] for j in 1:nv) + end + + inner = get_model(model) + @test inner isa ExaModels.ExaModel + @test NLPModels.get_nvar(inner) == ns * nv + @test NLPModels.get_ncon(inner) == ns * nv + end + + @testset "Variable bounds and start values" begin + ns, nv = 2, 2 + θ_data = [1.0 2.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv; start = 0.5, lvar = 0.0, uvar = 10.0) + objective(c, v[j]^2 for j in 1:nv) + constraint(c, v[j] for j in 1:nv; lcon = 0.0, ucon = 100.0) + end + + # Check that meta matrices have correct shape and values + @test size(model.meta.x0) == (nv, ns) + @test all(model.meta.x0 .== 0.5) + + @test all(model.meta.lvar .== 0.0) + @test all(model.meta.uvar .== 10.0) + end + + @testset "Ipopt solver with known solution" begin + ns, nv = 3, 1 + θ_vals = [2.0, 4.0, 6.0] + θ_data = reshape(θ_vals, 1, ns) + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, (v[1] - θ[1])^2) + constraint(c, v[1]; lcon = 0.0, ucon = Inf) + end + + # Solve via fused model + result = ipopt(get_model(model); print_level = 0) + @test result.status == :first_order + + x_sol = result.solution + for (i, θ) in enumerate(θ_vals) + @test x_sol[var_indices(model, i)] ≈ [θ] atol = 1.0e-5 + end + @test result.objective ≈ 0.0 atol = 1.0e-8 + end + + @testset "Ipopt solver - multiple variables per scenario" begin + ns, nv = 2, 2 + θ_data = [1.0 2.0; 3.0 2.0] + + c = ExaCore() + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, (v[j] - θ[j])^2 for j in 1:nv) + constraint(c, v[1] + v[2]; ucon = 10.0) + end + + result = ipopt(get_model(model); print_level = 0) + @test result.status == :first_order + + x_sol = result.solution + @test x_sol[var_indices(model, 1)] ≈ [1.0, 3.0] atol = 1.0e-5 + @test x_sol[var_indices(model, 2)] ≈ [2.0, 2.0] atol = 1.0e-5 + @test result.objective ≈ 0.0 atol = 1.0e-8 + end + + @testset "Multi-backend: $backend" for backend in BACKENDS + ns, nv = 2, 2 + θ_data = [2.0 3.0] + + c = ExaCore(; backend = backend) + model = BatchExaModel(c, ns, θ_data) do c, θ + v = variable(c, nv) + objective(c, θ[1] * v[j]^2 for j in 1:nv) + constraint(c, v[j] - θ[1] for j in 1:nv) + end + + # Create test matrix on the right device + bx_cpu = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) + bx = backend === nothing ? bx_cpu : adapt(backend, bx_cpu) + + # Batch obj! + bf = similar(bx, ns) + obj!(model, bx, bf) + @test Array(bf) ≈ [10.0, 75.0] + @test sum(Array(bf)) ≈ 85.0 + + # Consistency with fused model + @test sum(Array(bf)) ≈ obj(get_model(model), vec(bx)) + + # Batch grad! + bg = similar(bx, nv, ns) + grad!(model, bx, bg) + @test Array(bg) ≈ [4.0 18.0; 8.0 24.0] + + # Batch cons! + bc = similar(bx, model.nc, ns) + cons!(model, bx, bc) + @test Array(bc) ≈ [-1.0 0.0; 0.0 1.0] + + # Batch jac_coord! + nnzj = NLPModels.get_nnzj(model) + bjvals = similar(bx, nnzj, ns) + jac_coord!(model, bx, bjvals) + @test all(v -> v ≈ 1.0, Array(bjvals)) + + # Batch hess_coord! + nnzh = NLPModels.get_nnzh(model) + by = similar(bx, model.nc, ns) + fill!(by, zero(eltype(by))) + bobj_weight = similar(bx, ns) + fill!(bobj_weight, one(eltype(bx))) + bhvals = similar(bx, nnzh, ns) + hess_coord!(model, bx, by, bobj_weight, bhvals) + hv = Array(bhvals) + @test any(v -> v ≈ 4.0, hv[:, 1]) + @test any(v -> v ≈ 6.0, hv[:, 2]) + end + + end +end + +end # module diff --git a/test/Project.toml b/test/Project.toml index 64d271bfc..31437441b 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" diff --git a/test/runtests.jl b/test/runtests.jl index fa1a70570..688617794 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,6 +11,7 @@ include("JuMPTest/JuMPTest.jl") include("UtilsTest/UtilsTest.jl") include("TwoStageTest/TwoStageTest.jl") include("LinAlgTest/LinAlgTest.jl") +include("BatchTest/BatchTest.jl") @testset "ExaModels test" begin @info "Running AD Test" @@ -30,4 +31,6 @@ include("LinAlgTest/LinAlgTest.jl") @info "Running LinAlg Test" LinAlgTest.runtests() + @info "Running Batch Test" + BatchTest.runtests() end From 4d1a531a5845f9848c8cbdfdae9c0a78d5441ab4 Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Wed, 25 Feb 2026 13:12:54 -0600 Subject: [PATCH 02/50] Make docs work --- src/batch.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/batch.jl b/src/batch.jl index 23f1094bd..3991c299a 100644 --- a/src/batch.jl +++ b/src/batch.jl @@ -121,10 +121,10 @@ function BatchExaModel( ns::Int, θ_data::AbstractMatrix, ) - size(θ_data, 2) == ns || throw( - ArgumentError("θ_data must have ns=$ns columns, got $(size(θ_data, 2))"), + Base.size(θ_data, 2) == ns || throw( + ArgumentError("θ_data must have ns=$ns columns, got $(Base.size(θ_data, 2))"), ) - nθ = size(θ_data, 1) + nθ = Base.size(θ_data, 1) # Register parameters as a flat vector (column-major: scenario 1, scenario 2, ...) parameter(c, vec(θ_data)) From ea364802a6af3804e92152a2410a728f2dbe2d3d Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Wed, 25 Feb 2026 13:53:16 -0600 Subject: [PATCH 03/50] Fix --- src/batch.jl | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/src/batch.jl b/src/batch.jl index 3991c299a..6df89f392 100644 --- a/src/batch.jl +++ b/src/batch.jl @@ -402,14 +402,17 @@ function hess_coord!( nh = m.nnzh_per perm = m.hess_perm - if allequal(bobj_weight) + # Move perm to device for GPU-compatible gather indexing + perm_dev = similar(m.hess_buffer, Int, length(perm)) + copyto!(perm_dev, perm) + + bobj_weight_cpu = Array(bobj_weight) + if allequal(bobj_weight_cpu) # Common case: uniform obj_weight → single fused call + permute - w = bobj_weight[1] + w = bobj_weight_cpu[1] hess_coord!(m.model, x_flat, y_flat, m.hess_buffer; obj_weight = w) bhvals_flat = vec(bhvals) - for i in eachindex(perm) - bhvals_flat[i] = m.hess_buffer[perm[i]] - end + bhvals_flat .= m.hess_buffer[perm_dev] else # Varying weights: 2-pass approach # Pass 1: objective hessian only (y=0, obj_weight=1) @@ -421,13 +424,14 @@ function hess_coord!( # Pass 2: constraint hessian only (obj_weight=0) hess_coord!(m.model, x_flat, y_flat, m.hess_buffer; obj_weight = zero(eltype(x_flat))) - # Combine per scenario + # Build per-element weight vector: element i belongs to scenario (i-1)÷nh+1 + w_cpu = [bobj_weight_cpu[(i - 1) ÷ nh + 1] for i in 1:length(perm)] + w_dev = similar(bobj_weight, length(perm)) + copyto!(w_dev, w_cpu) + + # Combine per scenario (vectorized for GPU) bhvals_flat = vec(bhvals) - for i in eachindex(perm) - s = (i - 1) ÷ nh + 1 - bhvals_flat[i] = - bobj_weight[s] * hess_obj[perm[i]] + m.hess_buffer[perm[i]] - end + bhvals_flat .= w_dev .* hess_obj[perm_dev] .+ m.hess_buffer[perm_dev] end return bhvals end From cb3a523de9f2fb5cab9a3f218fadd8b13df2b3ce Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Fri, 6 Mar 2026 11:18:40 -0600 Subject: [PATCH 04/50] Use get_nbatch() --- src/batch.jl | 2 -- test/BatchTest/BatchTest.jl | 4 ++-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/batch.jl b/src/batch.jl index 6df89f392..17f9d6db8 100644 --- a/src/batch.jl +++ b/src/batch.jl @@ -230,8 +230,6 @@ end # Accessors # ============================================================================ -num_scenarios(m::BatchExaModel) = m.ns - """ get_model(model::BatchExaModel) diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index a2420ac20..9967e842c 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -7,7 +7,7 @@ import NLPModels: obj, obj!, cons!, cons_nln!, grad!, jac_coord!, hess_coord!, jac_structure!, hess_structure! import ExaModels: - num_scenarios, set_scenario_parameters!, set_all_scenario_parameters!, + set_scenario_parameters!, set_all_scenario_parameters!, var_indices, cons_block_indices, get_model import NLPModelsIpopt: ipopt @@ -30,7 +30,7 @@ function runtests() constraint(c, v[j] for j in 1:nv) end - @test num_scenarios(model) == 3 + @test NLPModels.get_nbatch(model) == 3 @test NLPModels.get_nvar(model) == nv @test NLPModels.get_ncon(model) == nv @test NLPModels.get_nbatch(model) == ns From bce56ca55a6394a0718c5450317d6731ebe57635 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 12:35:02 -0400 Subject: [PATCH 05/50] Rewrite BatchExaModel using BatchExaCore + EachInstance() API - BatchExaCore(nbatch) creates ExaCore with Matrix storage - EachInstance() marker for per-instance add_var/add_par/add_con - Unified append! for Vector and Matrix (no separate methods) - ExaModel(c) handles both batch and non-batch via _meta_dims dispatch - Removed BatchExt: no fused_model, objbuffer, or hess_buffer storage - get_model() constructs flat model on-the-fly for solver compatibility - hess_perm computed on-the-fly in hess_coord! - BatchNLPModelMeta stores per-instance nvar/ncon/nnzj/nnzh - Relaxed ExaCore VT constraint from AbstractVector to AbstractArray Co-Authored-By: Claude Opus 4.6 --- docs/src/batch.jl | 66 ++-- src/BatchNLPModels.jl | 353 ++++++++++++++++++ src/ExaModels.jl | 14 +- src/batch.jl | 710 ++++++++++++++++++++---------------- src/deprecated.jl | 2 +- src/nlp.jl | 97 ++--- test/BatchTest/BatchTest.jl | 430 +++++++++------------- 7 files changed, 1021 insertions(+), 651 deletions(-) create mode 100644 src/BatchNLPModels.jl diff --git a/docs/src/batch.jl b/docs/src/batch.jl index 1f7f9d9d4..f59ce3229 100644 --- a/docs/src/batch.jl +++ b/docs/src/batch.jl @@ -1,15 +1,15 @@ # # [Batch Optimization](@id batch) -# ExaModels supports batch optimization through the `BatchExaModel`. This feature -# enables efficient evaluation of multiple fully independent optimization scenarios +# ExaModels supports batch optimization through the `ExaModel`. This feature +# enables efficient evaluation of multiple fully independent optimization instances # that share identical structure but differ in parameter values. # -# Unlike `TwoStageExaModel`, which couples scenarios through shared design variables, -# `BatchExaModel` treats each scenario as completely independent. The key advantage -# is that all scenarios share one compiled expression pattern and are fused into a +# Unlike `TwoStageExaModel`, which couples instances through shared design variables, +# `ExaModel` treats each instance as completely independent. The key advantage +# is that all instances share one compiled expression pattern and are fused into a # single model for efficient SIMD evaluation. # ## Problem Formulation -# A batch optimization problem solves `ns` independent scenarios simultaneously: +# A batch optimization problem solves `ns` independent instances simultaneously: # ```math # \begin{aligned} # \min_{v_i} \quad & f(v_i; \theta_i), \quad i = 1, \ldots, S \\ @@ -17,25 +17,25 @@ # & v_i \in \mathcal{V} # \end{aligned} # ``` -# where each scenario has the same structure but different parameters $\theta_i$. +# where each instance has the same structure but different parameters $\theta_i$. # ## Building a Batch Model -# The builder function defines expressions for a **single scenario**. `BatchExaModel` -# calls it `ns` times internally with per-scenario parameter handles. +# The builder function defines expressions for a **single instance**. `ExaModel` +# calls it `ns` times internally with per-instance parameter handles. # This means you never have to compute global index offsets manually. using ExaModels, MadNLP -# Define the problem dimensions and scenario parameters as a matrix of size `(nθ, ns)`: -ns = 3 ## number of scenarios -nv = 1 ## variables per scenario +# Define the problem dimensions and instance parameters as a matrix of size `(nθ, ns)`: +ns = 3 ## number of instances +nv = 1 ## variables per instance θ_data = [2.0 4.0 6.0] ## (1, 3) matrix: θ₁=2, θ₂=4, θ₃=6 # Build the model. First create an `ExaCore`, then pass it along with the parameter -# matrix to `BatchExaModel`: +# matrix to `ExaModel`: c = ExaCore() -model = BatchExaModel(c, ns, θ_data) do c, θ - ## Create variables — this is called once per scenario, offsets are automatic +model = ExaModel(c, ns, θ_data) do c, θ + ## Create variables — this is called once per instance, offsets are automatic v = variable(c, nv) ## Objective: minimize (v - θ)² objective(c, (v[1] - θ[1])^2) @@ -45,26 +45,26 @@ end # The builder function receives: # - `c`: the `ExaCore` — use `variable(c, ...)`, `objective(c, ...)`, `constraint(c, ...)` as usual -# - `θ`: a per-scenario parameter handle (indices 1:nθ) +# - `θ`: a per-instance parameter handle (indices 1:nθ) # # Variable creation via `variable(c, ...)` works exactly like in a regular `ExaModel`. # You can set start values, lower/upper bounds, etc. # ## Batch API (NLPModels) -# `BatchExaModel` implements the `AbstractBatchNLPModel` interface from NLPModels.jl. +# `ExaModel` implements the `AbstractBatchNLPModel` interface from NLPModels.jl. # All evaluation functions use matrices of size `(dim, ns)`: import NLPModels -println("Variables per scenario: ", NLPModels.get_nvar(model)) -println("Constraints per scenario: ", NLPModels.get_ncon(model)) -println("Number of scenarios: ", NLPModels.get_nbatch(model)) +println("Variables per instance: ", NLPModels.get_nvar(model)) +println("Constraints per instance: ", NLPModels.get_ncon(model)) +println("Number of instances: ", NLPModels.get_nbatch(model)) -# Evaluate objectives for all scenarios at once: +# Evaluate objectives for all instances at once: bx = reshape([1.0, 3.0, 5.0], nv, ns) bf = zeros(ns) NLPModels.obj!(model, bx, bf) println("\nObjective values: ", bf) -## scenario 1: (1-2)² = 1, scenario 2: (3-4)² = 1, scenario 3: (5-6)² = 1 +## instance 1: (1-2)² = 1, instance 2: (3-4)² = 1, instance 3: (5-6)² = 1 # Evaluate gradients: bg = zeros(nv, ns) @@ -83,20 +83,20 @@ result = madnlp(ExaModels.get_model(model); print_level = MadNLP.ERROR) println("\nSolution status: ", result.status) println("Optimal objective: ", round(result.objective, digits = 4)) -# Extract per-scenario solutions: +# Extract per-instance solutions: x_sol = result.solution for i in 1:ns v_sol = x_sol[ExaModels.var_indices(model, i)] - println("Scenario $i: v* = ", round(v_sol[1], digits = 4)) + println("Instance $i: v* = ", round(v_sol[1], digits = 4)) end # ## A More Complex Example -# Here's a batch model with multiple variables, objectives, and constraints per scenario: +# Here's a batch model with multiple variables, objectives, and constraints per instance: ns2, nv2 = 2, 3 θ_data2 = [1.0 4.0; 2.0 5.0; 3.0 6.0] ## (3, 2) matrix c2 = ExaCore() -model2 = BatchExaModel(c2, ns2, θ_data2) do c, θ +model2 = ExaModel(c2, ns2, θ_data2) do c, θ v = variable(c, nv2; start = 1.0, lvar = 0.0, uvar = 10.0) ## Objective: Σⱼ (vⱼ - θⱼ)² objective(c, (v[j] - θ[j])^2 for j in 1:nv2) @@ -110,17 +110,17 @@ println("Status: ", result2.status) x_sol2 = result2.solution for i in 1:ns2 v_sol = x_sol2[ExaModels.var_indices(model2, i)] - println("Scenario $i: v* = ", round.(v_sol, digits = 4)) + println("Instance $i: v* = ", round.(v_sol, digits = 4)) end # ## Updating Parameters -# You can update scenario parameters and re-solve without rebuilding the model: +# You can update instance parameters and re-solve without rebuilding the model: -# Update a single scenario: -ExaModels.set_scenario_parameters!(model, 1, [10.0]) +# Update a single instance: +ExaModels.set_instance_parameters!(model, 1, [10.0]) -# Or update all scenarios at once: -ExaModels.set_all_scenario_parameters!(model, [[10.0], [12.0], [14.0]]) +# Or update all instances at once: +ExaModels.set_all_instance_parameters!(model, [[10.0], [12.0], [14.0]]) # Re-solve with new parameters: result3 = madnlp(ExaModels.get_model(model); print_level = MadNLP.ERROR) @@ -129,5 +129,5 @@ println("Status: ", result3.status) x_sol3 = result3.solution for i in 1:ns v_sol = x_sol3[ExaModels.var_indices(model, i)] - println("Scenario $i: v* = ", round(v_sol[1], digits = 4)) + println("Instance $i: v* = ", round(v_sol[1], digits = 4)) end diff --git a/src/BatchNLPModels.jl b/src/BatchNLPModels.jl new file mode 100644 index 000000000..6db579c83 --- /dev/null +++ b/src/BatchNLPModels.jl @@ -0,0 +1,353 @@ +""" + BatchNLPModels + +Template module for batched NLP models. Defines abstract types, metadata, +and generic API functions following the NLPModels.jl pattern. + +Key design: `AbstractBatchNLPModel <: NLPModels.AbstractNLPModel`, so batch +models participate in the standard NLPModels dispatch hierarchy. +""" +module BatchNLPModels + +import NLPModels: + NLPModels, + AbstractNLPModel, + AbstractNLPModelMeta + +# ============================================================================ +# Abstract types +# ============================================================================ + +""" + AbstractBatchNLPModel{T, S} <: AbstractNLPModel{T, S} + +Abstract type for batched NLP models. Subtypes `AbstractNLPModel` so that +batch models participate in the standard NLPModels dispatch hierarchy. + +Implementations must provide: +- `meta` field of type `<: AbstractBatchNLPModelMeta` +- `counters` field of type `NLPModels.Counters` +- Batch API methods: `obj!`, `grad!`, `cons!`, `jac_structure!`, `jac_coord!`, + `hess_structure!`, `hess_coord!` +""" +abstract type AbstractBatchNLPModel{T, S} <: AbstractNLPModel{T, S} end + +# ============================================================================ +# BatchNLPModelMeta +# ============================================================================ + +""" + BatchNLPModelMeta{T, VT, VI} <: AbstractNLPModelMeta{T, VT} + +Metadata for a batched NLP where `nbatch` independent problems share +identical structure (dimensions, sparsity patterns). Extends the standard +`NLPModelMeta` interface with a batch dimension. + +All `VT`-typed arrays are either vectors (nbatch=1) or matrices with +columns indexing instances. + +# Type parameters +- `T`: element type (e.g. `Float64`) +- `VT`: storage type — `Matrix{T}` when `nbatch > 1` (columns = instances), + `Vector{T}` when `nbatch == 1` +- `VI`: integer index vector type (typically `Vector{Int}`) +""" +struct BatchNLPModelMeta{T, VT, VI} <: AbstractNLPModelMeta{T, VT} + nvar::Int + x0::VT + lvar::VT + uvar::VT + ifix::VI + ilow::VI + iupp::VI + irng::VI + ifree::VI + iinf::VI + nlvb::Int + nlvo::Int + nlvc::Int + ncon::Int + y0::VT + lcon::VT + ucon::VT + jfix::VI + jlow::VI + jupp::VI + jrng::VI + jfree::VI + jinf::VI + nnzo::Int + nnzj::Int + lin_nnzj::Int + nln_nnzj::Int + nnzh::Int + nlin::Int + nnln::Int + lin::VI + nln::VI + minimize::Bool + islp::Bool + name::String + variable_bounds_analysis::Bool + constraint_bounds_analysis::Bool + sparse_jacobian::Bool + sparse_hessian::Bool + grad_available::Bool + jac_available::Bool + hess_available::Bool + jprod_available::Bool + jtprod_available::Bool + hprod_available::Bool +end + +# ============================================================================ +# Accessors — auto-generate get_* for all fields +# ============================================================================ + +for field in fieldnames(BatchNLPModelMeta) + meth = Symbol("get_", field) + # Extend NLPModels accessor if it exists, otherwise define locally + if isdefined(NLPModels, meth) + @eval begin + NLPModels.$meth(meta::BatchNLPModelMeta) = getproperty(meta, $(QuoteNode(field))) + NLPModels.$meth(m::AbstractBatchNLPModel) = NLPModels.$meth(m.meta) + end + else + @eval begin + $meth(meta::BatchNLPModelMeta) = getproperty(meta, $(QuoteNode(field))) + $meth(m::AbstractBatchNLPModel) = $meth(m.meta) + export $meth + end + end +end + +# get_nbatch is derived from the VT type, not stored as a field +get_nbatch(meta::BatchNLPModelMeta{T, <:AbstractMatrix}) where {T} = Base.size(meta.x0, 2) +get_nbatch(meta::BatchNLPModelMeta) = 1 +get_nbatch(m::AbstractBatchNLPModel) = get_nbatch(m.meta) +export get_nbatch + +# ============================================================================ +# Constructor helpers +# ============================================================================ + +""" + _first_instance(v) + +Extract first-instance data for bounds classification. +""" +_first_instance(v::AbstractVector) = v +_first_instance(m::AbstractMatrix) = @view m[:, 1] + +function _classify_bounds(lb, ub, ::Type{T}) where {T} + ifix = findall(lb .== ub) + ilow = findall((lb .> T(-Inf)) .& (ub .== T(Inf))) + iupp = findall((lb .== T(-Inf)) .& (ub .< T(Inf))) + irng = findall((lb .> T(-Inf)) .& (ub .< T(Inf)) .& (lb .< ub)) + ifree = findall((lb .== T(-Inf)) .& (ub .== T(Inf))) + iinf = findall(lb .> ub) + return ifix, ilow, iupp, irng, ifree, iinf +end + +""" + BatchNLPModelMeta(nbatch, nvar, x0, lvar, uvar, ncon, y0, lcon, ucon; kwargs...) + +Construct batch NLP metadata. Bounds analysis (variable/constraint +classification) is performed on the first instance. +""" +function BatchNLPModelMeta( + nvar::Int, + x0::VT, + lvar::VT, + uvar::VT, + ncon::Int, + y0::VT, + lcon::VT, + ucon::VT; + nnzj::Int = nvar * ncon, + nnzh::Int = nvar * (nvar + 1) ÷ 2, + minimize::Bool = true, + islp::Bool = false, + name::String = "Generic", +) where {VT} + T = eltype(VT) + + # Variable bounds analysis (first instance) + ifix, ilow, iupp, irng, ifree, iinf = _classify_bounds( + _first_instance(lvar), _first_instance(uvar), T, + ) + + # Constraint bounds analysis (first instance) + if ncon > 0 + jfix, jlow, jupp, jrng, jfree, jinf = _classify_bounds( + _first_instance(lcon), _first_instance(ucon), T, + ) + else + jfix = jlow = jupp = jrng = jfree = jinf = Int[] + end + + nln = collect(1:ncon) + VI = Vector{Int} + + return BatchNLPModelMeta{T, VT, VI}( + nvar, + x0, lvar, uvar, + ifix, ilow, iupp, irng, ifree, iinf, + nvar, nvar, nvar, # nlvb, nlvo, nlvc + ncon, + y0, lcon, ucon, + jfix, jlow, jupp, jrng, jfree, jinf, + nvar, # nnzo + nnzj, 0, nnzj, # nnzj, lin_nnzj, nln_nnzj + nnzh, + 0, ncon, # nlin, nnln + Int[], nln, # lin, nln + minimize, islp, name, + true, true, # variable/constraint_bounds_analysis + true, true, # sparse_jacobian/hessian + true, # grad_available + ncon > 0, # jac_available + true, # hess_available + ncon > 0, # jprod_available + ncon > 0, # jtprod_available + true, # hprod_available + ) +end + +# ============================================================================ +# Generic batch API — function stubs +# ============================================================================ +# +# Concrete implementations must define the `!` (in-place) methods. +# Allocating wrappers are provided as defaults. + +""" + obj!(m::AbstractBatchNLPModel, bx::AbstractMatrix, bf::AbstractVector) + +Evaluate per-instance objectives. `bx` is `(nvar, nbatch)`, `bf` is `(nbatch,)`. +""" +function obj! end + +""" + obj(m::AbstractBatchNLPModel, bx::AbstractMatrix) -> Vector + +Allocating version of `obj!`. +""" +function NLPModels.obj(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} + bf = Vector{T}(undef, get_nbatch(m)) + obj!(m, bx, bf) + return bf +end + +""" + grad!(m::AbstractBatchNLPModel, bx::AbstractMatrix, bg::AbstractMatrix) + +Evaluate per-instance gradients. `bx` and `bg` are `(nvar, nbatch)`. +""" +function NLPModels.grad!(m::AbstractBatchNLPModel, bx::AbstractMatrix, bg::AbstractMatrix) + error("grad! not implemented for $(typeof(m))") +end + +""" + grad(m::AbstractBatchNLPModel, bx::AbstractMatrix) -> Matrix + +Allocating version of `grad!`. +""" +function NLPModels.grad(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} + bg = Matrix{T}(undef, get_nvar(m), get_nbatch(m)) + NLPModels.grad!(m, bx, bg) + return bg +end + +""" + cons!(m::AbstractBatchNLPModel, bx::AbstractMatrix, bc::AbstractMatrix) + +Evaluate per-instance constraints. `bx` is `(nvar, nbatch)`, `bc` is `(ncon, nbatch)`. +""" +function NLPModels.cons!(m::AbstractBatchNLPModel, bx::AbstractMatrix, bc::AbstractMatrix) + error("cons! not implemented for $(typeof(m))") +end + +""" + cons(m::AbstractBatchNLPModel, bx::AbstractMatrix) -> Matrix + +Allocating version of `cons!`. +""" +function NLPModels.cons(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} + bc = Matrix{T}(undef, get_ncon(m), get_nbatch(m)) + NLPModels.cons!(m, bx, bc) + return bc +end + +""" + jac_structure!(m::AbstractBatchNLPModel, rows, cols) + +Per-instance Jacobian sparsity pattern (local indices). `rows` and `cols` +have length `nnzj` (per instance). +""" +function NLPModels.jac_structure!( + m::AbstractBatchNLPModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + error("jac_structure! not implemented for $(typeof(m))") +end + +""" + jac_coord!(m::AbstractBatchNLPModel, bx::AbstractMatrix, jvals::AbstractVector) + +Evaluate batch Jacobian values. `jvals` is a flat vector of length +`nnzj * nbatch`, laid out instance by instance. +""" +function NLPModels.jac_coord!( + m::AbstractBatchNLPModel, + bx::AbstractMatrix, + jvals::AbstractVector, +) + error("jac_coord! not implemented for $(typeof(m))") +end + +""" + hess_structure!(m::AbstractBatchNLPModel, rows, cols) + +Per-instance Hessian sparsity pattern (local indices). `rows` and `cols` +have length `nnzh` (per instance). +""" +function NLPModels.hess_structure!( + m::AbstractBatchNLPModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + error("hess_structure! not implemented for $(typeof(m))") +end + +""" + hess_coord!(m::AbstractBatchNLPModel, bx, by, bobj_weight, hvals) + +Evaluate batch Hessian values. `hvals` is a flat vector of length +`nnzh * nbatch`, laid out instance by instance. + +- `bx`: `(nvar, nbatch)` primal values +- `by`: `(ncon, nbatch)` dual values +- `bobj_weight`: `(nbatch,)` per-instance objective weights +- `hvals`: `(nnzh * nbatch,)` flat Hessian values +""" +function NLPModels.hess_coord!( + m::AbstractBatchNLPModel, + bx::AbstractMatrix, + by::AbstractMatrix, + bobj_weight::AbstractVector, + hvals::AbstractVector, +) + error("hess_coord! not implemented for $(typeof(m))") +end + +# ============================================================================ +# Exports +# ============================================================================ + +export AbstractBatchNLPModel, + BatchNLPModelMeta, + obj! + +end # module BatchNLPModels diff --git a/src/ExaModels.jl b/src/ExaModels.jl index 955688526..f75b339d6 100644 --- a/src/ExaModels.jl +++ b/src/ExaModels.jl @@ -59,8 +59,10 @@ include("hessian.jl") include("nlp.jl") include("deprecated.jl") include("utils.jl") -include("batch.jl") include("tags.jl") +include("BatchNLPModels.jl") +using .BatchNLPModels +include("batch.jl") export ExaModel, ExaCore, @@ -90,6 +92,14 @@ export ExaModel, exa_prod, @register_univariate, @register_bivariate, - BatchExaModel + AbstractBatchNLPModel, + BatchNLPModelMeta, + BatchExaCore, + BatchExaModel, + EachInstance, + get_nbatch, + get_model, + var_indices, + cons_block_indices end # module ExaModels diff --git a/src/batch.jl b/src/batch.jl index 17f9d6db8..9df964c60 100644 --- a/src/batch.jl +++ b/src/batch.jl @@ -1,323 +1,339 @@ # ============================================================================ -# BatchExaModel — single-scenario builder + NLPModels matrix batch API +# Batch ExaModel — dispatch on VT <: AbstractMatrix # ============================================================================ +# +# Following the same pattern as TwoStageExaModel: +# - BatchExaModelTag stored in ExaCore's tag field +# - BatchExaCore = ExaCore{T, <:AbstractMatrix, B, <:BatchExaModelTag} +# - EachInstance() marker for per-instance declarations +# - ExaModel(c) handles both batch and non-batch via dispatch -""" - BatchExaModel{T, VT, MT, M} <: NLPModels.AbstractBatchNLPModel{T, MT} +import .BatchNLPModels: obj! -Parametric optimization model where multiple fully independent scenarios are fused -into a single ExaModel and evaluated simultaneously using shared compiled expression -patterns. +# ============================================================================ +# Tag and marker types +# ============================================================================ -All scenarios share identical sparsity structures. The builder defines expressions -for **one scenario**; `BatchExaModel` calls it `ns` times internally with offset -variable/parameter handles. +""" + EachInstance -# Dimensions +Marker type used with [`add_var`](@ref), [`add_par`](@ref), and [`add_con`](@ref) +to indicate that the declaration is replicated for each instance in a batch model. -- `ns`: number of scenarios (batch size) -- `nv`: number of variables per scenario -- `nc`: number of constraints per scenario -- `nθ`: number of parameters per scenario +## Example +```julia +core = BatchExaCore(3) +c, v = add_var(core, EachInstance(), 10) +c, _ = add_obj(c, v[i, s]^2 for i in 1:10, s in 1:3) +``` """ -struct BatchExaModel{T, VT <: AbstractVector{T}, MT <: AbstractMatrix{T}, M} <: - NLPModels.AbstractBatchNLPModel{T, MT} - meta::NLPModels.BatchNLPModelMeta{T, MT} - model::M - objbuffer::VT - hess_perm::Vector{Int} - hess_buffer::VT - ns::Int - nv::Int - nc::Int - nθ::Int - nobj_per::Int - nnzj_per::Int - nnzh_per::Int - nnzh_obj_per::Int - nnzh_con_per::Int -end +struct EachInstance end -function Base.show(io::IO, m::BatchExaModel{T, VT}) where {T, VT} - println(io, "BatchExaModel{$T, $VT}") - println(io, " Scenarios: $(m.ns)") - println(io, " Variables per scenario: $(m.nv)") - println(io, " Constraints per scenario: $(m.nc)") - println(io, " Parameters per scenario: $(m.nθ)") - println(io, " Total variables: $(m.ns * m.nv)") - println(io, " Total constraints: $(m.ns * m.nc)") - println(io, " Jacobian nnz per scenario: $(m.nnzj_per)") - return println(io, " Hessian nnz per scenario: $(m.nnzh_per)") +struct BatchExaModelTag <: AbstractExaModelTag + nbatch::Int end # ============================================================================ -# Helpers +# Type aliases # ============================================================================ -_count_hess_nnz(::ObjectiveNull) = 0 -_count_hess_nnz(::ConstraintNull) = 0 -_count_hess_nnz(node) = _count_hess_nnz(node.inner) + node.f.o2step * length(node.itr) - -function _build_hess_perm(ns, nnzh_obj_per, nnzh_con_per) - nnzh_per = nnzh_obj_per + nnzh_con_per - perm = Vector{Int}(undef, ns * nnzh_per) - for s in 1:ns - base = (s - 1) * nnzh_per - for k in 1:nnzh_obj_per - perm[base + k] = (s - 1) * nnzh_obj_per + k - end - for k in 1:nnzh_con_per - perm[base + nnzh_obj_per + k] = ns * nnzh_obj_per + (s - 1) * nnzh_con_per + k - end - end - return perm -end +""" + BatchExaCore{T,VT,B} -function _to_matrix(v::AbstractVector, nrows::Int, ncols::Int) - mat = similar(v, nrows, ncols) - copyto!(vec(mat), v) - return mat -end +Type alias for an [`ExaCore`](@ref) whose `tag` is a [`BatchExaModelTag`] +and whose storage arrays are matrices (columns = instances). +""" +const BatchExaCore{T,VT,B} = ExaCore{T,VT,B,<:BatchExaModelTag} -# ============================================================================ -# Constructor -# ============================================================================ +""" + BatchExaModel{T,VT,E,V,P,O,C,R,M} +Type alias for an [`ExaModel`](@ref) built from a [`BatchExaCore`](@ref). """ - BatchExaModel(build, c::ExaCore, ns::Int, θ_data::AbstractMatrix) +const BatchExaModel{T,VT,E,V,P,O,C,R,M} = ExaModel{T,VT,E,V,P,O,C,<:BatchExaModelTag,R,M} -Build a batch model from a single-scenario builder function. +# ============================================================================ +# get_nbatch +# ============================================================================ -The user creates an `ExaCore` and passes parameter data as a matrix of size `(nθ, ns)`. -The builder defines a **single scenario** — creating variables, objectives, and -constraints using standard ExaModels calls. `BatchExaModel` invokes it `ns` times, -each time with a per-scenario parameter handle `θ`. Variable creation via -`variable(c, ...)` works normally and automatically gets the correct offsets. +get_nbatch(m::BatchExaModel) = m.tag.nbatch +get_nbatch(c::BatchExaCore) = c.tag.nbatch +get_nbatch(m::AbstractExaModel) = 1 -# Arguments -- `build::Function`: Function `(c, θ) -> nothing` - - `c`: ExaCore — use `variable(c, ...)`, `objective(c, ...)`, `constraint(c, ...)` as usual - - `θ`: Parameter handle for this scenario's parameters (indices 1:nθ) -- `c::ExaCore`: ExaCore instance (parameters will be registered internally) -- `ns::Int`: Number of scenarios -- `θ_data::AbstractMatrix`: Parameter matrix of size `(nθ, ns)` +# ============================================================================ +# _meta_dims — per-instance dimensions for batch +# ============================================================================ -# Example -```julia -ns, nθ = 100, 3 -θ_data = rand(nθ, ns) - -c = ExaCore() -model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, 5; start = 1.0, lvar = 0.0, uvar = 10.0) - objective(c, θ[j] * v[j]^2 for j in 1:5) - constraint(c, v[j] - θ[j] for j in 1:3) +function _meta_dims(c::C) where {T, VT <: AbstractArray{T}, C <: BatchExaCore{T, VT}} + nb = c.tag.nbatch + return (c.nvar ÷ nb, c.ncon ÷ nb, c.nnzj ÷ nb, c.nnzh ÷ nb) end -``` -""" -function BatchExaModel( - build::Function, - c::ExaCore, - ns::Int, - θ_data::AbstractMatrix, - ) - Base.size(θ_data, 2) == ns || throw( - ArgumentError("θ_data must have ns=$ns columns, got $(Base.size(θ_data, 2))"), - ) - nθ = Base.size(θ_data, 1) - # Register parameters as a flat vector (column-major: scenario 1, scenario 2, ...) - parameter(c, vec(θ_data)) - - # Call builder once per scenario with per-scenario θ handle. - # The builder calls variable(c, ...) itself — offsets are automatic. - for s in 1:ns - θ_s = Parameter(nθ, nθ, (s - 1) * nθ) - build(c, θ_s) - end +# ============================================================================ +# BatchExaCore constructor +# ============================================================================ - # Infer per-scenario dimensions - nv = c.nvar ÷ ns - nv * ns != c.nvar && throw( - DimensionMismatch( - "Total variables ($(c.nvar)) not evenly divisible by ns ($ns)", - ), - ) +""" + BatchExaCore(nbatch; T = Float64, backend = nothing, kwargs...) - nc_total = c.ncon - nc = nc_total ÷ ns - nc * ns != nc_total && throw( - DimensionMismatch( - "Total constraints ($nc_total) not evenly divisible by ns ($ns)", - ), - ) +Create an [`ExaCore`](@ref) for building batch optimization models with +`nbatch` independent instances. - nobj_total = c.nobj - nobj_per = nobj_total ÷ ns - nobj_per * ns != nobj_total && throw( - DimensionMismatch( - "Total objective entries ($nobj_total) not evenly divisible by ns ($ns)", - ), - ) +Storage arrays (`x0`, `lvar`, `uvar`, `θ`, `y0`, `lcon`, `ucon`) are +`Matrix{T}` with `nbatch` columns. Use [`add_var`](@ref), [`add_par`](@ref), +[`add_obj`](@ref), and [`add_con`](@ref) with [`EachInstance()`](@ref) to +declare per-instance components. - objbuffer = similar(c.x0, nobj_total) - - model = ExaModel(c) - - # Per-scenario sparsity counts - total_nnzj = NLPModels.get_nnzj(model) - total_nnzh = NLPModels.get_nnzh(model) - nnzj_per = total_nnzj ÷ ns - nnzh_per = total_nnzh ÷ ns - - # Obj/con hessian split - nnzh_obj_total = _count_hess_nnz(model.objs) - nnzh_con_total = _count_hess_nnz(model.cons) - nnzh_obj_per = nnzh_obj_total ÷ ns - nnzh_con_per = nnzh_con_total ÷ ns - - # Hessian permutation - hess_perm = _build_hess_perm(ns, nnzh_obj_per, nnzh_con_per) - - # Hessian buffer - hess_buffer = similar(c.x0, total_nnzh) - - # Build BatchNLPModelMeta with matrices - T = eltype(c.x0) - x0_mat = _to_matrix(model.meta.x0, nv, ns) - lvar_mat = _to_matrix(model.meta.lvar, nv, ns) - uvar_mat = _to_matrix(model.meta.uvar, nv, ns) - y0_mat = _to_matrix(model.meta.y0, nc, ns) - lcon_mat = _to_matrix(model.meta.lcon, nc, ns) - ucon_mat = _to_matrix(model.meta.ucon, nc, ns) - - MT = typeof(x0_mat) - meta = NLPModels.BatchNLPModelMeta{T, MT}( - ns, - nv; - x0 = x0_mat, - lvar = lvar_mat, - uvar = uvar_mat, - ncon = nc, - y0 = y0_mat, - lcon = lcon_mat, - ucon = ucon_mat, - nnzj = nnzj_per, - nnzh = nnzh_per, - minimize = model.meta.minimize, - ) - - VT = typeof(c.x0) - return BatchExaModel{T, VT, MT, typeof(model)}( - meta, - model, - objbuffer, - hess_perm, - hess_buffer, - ns, - nv, - nc, - nθ, - nobj_per, - nnzj_per, - nnzh_per, - nnzh_obj_per, - nnzh_con_per, +## Example +```julia +core = BatchExaCore(3) +c, v = add_var(core, EachInstance(), 10; start = 1.0, lvar = 0.0, uvar = 10.0) +nb = get_nbatch(c) +c, _ = add_obj(c, v[i, s]^2 for i in 1:10, s in 1:nb) +model = ExaModel(c) +``` +""" +function BatchExaCore(nbatch::Integer; T::Type{<:AbstractFloat} = Float64, backend = nothing, kwargs...) + x0 = convert_array(zeros(T, 0, nbatch), backend) + return ExaCore( + :Generic, + backend, + (), # var + (), # par + (), # obj + (), # cons + 0, 0, 0, 0, 0, # nvar, npar, ncon, nconaug, nobj + 0, 0, 0, 0, # nnzc, nnzg, nnzj, nnzh + x0, + similar(x0), # θ + similar(x0), # lvar + similar(x0), # uvar + similar(x0), # y0 + similar(x0), # lcon + similar(x0), # ucon + true, # minimize + BatchExaModelTag(nbatch), + (;), # refs ) end # ============================================================================ -# Accessors +# add_var / add_par / add_con for BatchExaCore # ============================================================================ """ - get_model(model::BatchExaModel) + add_var(core::BatchExaCore, ::EachInstance, dims...; start = 0, lvar = -Inf, uvar = Inf, name = nothing) -Get the underlying fused ExaModel for direct NLPModels interface usage (e.g. Ipopt). +Add per-instance variables to a batch core. Creates `prod(dims) * nbatch` +variables total — one copy of the block per instance. The returned `Variable` +has dimensions `(dims..., nbatch)`. """ -get_model(model::BatchExaModel) = model.model +function add_var( + c::C, + ::EachInstance, + ns...; + name = nothing, + start = zero(T), + lvar = T(-Inf), + uvar = T(Inf), +) where {T, VT <: AbstractArray{T}, C <: BatchExaCore{T, VT}} + nbatch = c.tag.nbatch + len_per = total(ns) # per-instance count + len_total = len_per * nbatch # fused total + nvar = c.nvar + len_total + + # Append per-instance rows to matrix storage + x0 = append!(c.backend, c.x0, start, len_per) + lv = append!(c.backend, c.lvar, lvar, len_per) + uv = append!(c.backend, c.uvar, uvar, len_per) + + v = Variable((ns..., nbatch), len_total, c.nvar, _val_name(name), nothing) + (ExaCore(c; var = (v, c.var...), nvar = nvar, x0 = x0, lvar = lv, uvar = uv, + refs = add_refs(c.refs, name, v)), v) +end """ - var_indices(model::BatchExaModel, i) -> UnitRange + add_par(core::BatchExaCore, ::EachInstance, value::AbstractArray; name = nothing) + add_par(core::BatchExaCore, ::EachInstance, dims...; name = nothing, start = 0) -Index range for variables of scenario `i` in the global (fused) variable vector. +Add per-instance parameters to a batch core. The parameter values are +replicated for each instance. """ -function var_indices(model::BatchExaModel, i::Int) - nv = model.nv - return ((i - 1) * nv + 1):(i * nv) +function add_par( + c::C, + ::EachInstance, + value::AbstractArray; + name = nothing, +) where {T, VT <: AbstractArray{T}, C <: BatchExaCore{T, VT}} + nbatch = c.tag.nbatch + len_per = length(value) + len_total = len_per * nbatch + npar = c.npar + len_total + θ = append!(c.backend, c.θ, value, len_per) + p = Parameter((Base.size(value)..., nbatch), len_total, c.npar, nothing) + (ExaCore(c; par = (p, c.par...), θ = θ, npar = npar, refs = add_refs(c.refs, name, p)), p) +end + +function add_par( + c::C, + ::EachInstance, + ns...; + name = nothing, + start = zero(T), +) where {T, VT <: AbstractArray{T}, C <: BatchExaCore{T, VT}} + nbatch = c.tag.nbatch + len_per = total(ns) + len_total = len_per * nbatch + npar = c.npar + len_total + θ = append!(c.backend, c.θ, start, len_per) + p = Parameter((ns..., nbatch), len_total, c.npar, nothing) + (ExaCore(c; par = (p, c.par...), θ = θ, npar = npar, refs = add_refs(c.refs, name, p)), p) end """ - cons_block_indices(model::BatchExaModel, i) -> UnitRange + add_con(core::BatchExaCore, ::EachInstance, gen; start = 0, lcon = 0, ucon = 0, name = nothing) -Index range for constraints of scenario `i` in the global (fused) constraint vector. +Add per-instance constraints to a batch core. """ -function cons_block_indices(model::BatchExaModel, i::Int) - nc = model.nc - return ((i - 1) * nc + 1):(i * nc) +function add_con( + c::C, + ::EachInstance, + ns...; + name = nothing, + tag = nothing, + start = zero(T), + lcon = zero(T), + ucon = zero(T), +) where {T, VT <: AbstractArray{T}, C <: BatchExaCore{T, VT}} + gen = _get_generator(ns) + dims = _get_con_dims(ns) + gen = _adapt_gen(gen) + f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) + pars = gen.iter + + nitr = length(pars) + o = c.ncon + ncon = c.ncon + nitr + nnzj = c.nnzj + nitr * f.o1step + nnzh = c.nnzh + nitr * f.o2step + + # Append per-instance rows to matrix storage + nbatch = c.tag.nbatch + nitr_per = nitr ÷ nbatch + y0 = append!(c.backend, c.y0, start, nitr_per) + lc = append!(c.backend, c.lcon, lcon, nitr_per) + uc = append!(c.backend, c.ucon, ucon, nitr_per) + + con = Constraint(f, convert_array(pars, c.backend), o, dims, nothing) + (ExaCore(c; ncon = ncon, nnzj = nnzj, nnzh = nnzh, y0 = y0, lcon = lc, ucon = uc, + cons = (con, c.cons...), refs = add_refs(c.refs, name, con)), con) end # ============================================================================ -# Parameter Updates +# Batch show # ============================================================================ -function set_scenario_parameters!(model::BatchExaModel, i::Int, θ_new::AbstractVector) - nθ = model.nθ - length(θ_new) != nθ && throw( - DimensionMismatch("Parameter size mismatch: expected $nθ, got $(length(θ_new))"), - ) - θ_start = (i - 1) * nθ + 1 - θ_end = i * nθ - copyto!(view(model.model.θ, θ_start:θ_end), θ_new) - return nothing +function Base.show(io::IO, m::BatchExaModel{T, VT}) where {T, VT} + nb = get_nbatch(m) + nv = NLPModels.get_nvar(m) + nc = NLPModels.get_ncon(m) + println(io, "An ExaModel{$T, $VT, ...} (batch)") + println(io, " Instances: $nb") + println(io, " Variables per instance: $nv") + println(io, " Constraints per instance: $nc") + println(io, " Total variables: $(nb * nv)") + return println(io, " Total constraints: $(nb * nc)") end -function set_all_scenario_parameters!( - model::BatchExaModel, - θ_sets::Vector{<:AbstractVector}, +# ============================================================================ +# Accessors +# ============================================================================ + +""" + get_model(model) + +For batch models, returns a flat (Vector-based) ExaModel for direct NLPModels +interface usage (e.g. MadNLP/Ipopt). For regular models, returns self. +""" +get_model(model::ExaModel) = model +function get_model(model::BatchExaModel) + nb = get_nbatch(model) + nvar = NLPModels.get_nvar(model) * nb + ncon = NLPModels.get_ncon(model) * nb + nnzj = NLPModels.get_nnzj(model) * nb + nnzh = NLPModels.get_nnzh(model) * nb + meta = BatchNLPModelMeta( + nvar, vec(model.meta.x0), vec(model.meta.lvar), vec(model.meta.uvar), + ncon, vec(model.meta.y0), vec(model.meta.lcon), vec(model.meta.ucon); + nnzj = nnzj, nnzh = nnzh, minimize = model.meta.minimize, + ) + return ExaModel( + model.name, model.vars, model.pars, model.objs, model.cons, + vec(model.θ), meta, NLPModels.Counters(), nothing, nothing, model.refs, ) - length(θ_sets) == model.ns || - throw(ArgumentError("θ_sets must have length $(model.ns)")) - for i in 1:(model.ns) - set_scenario_parameters!(model, i, θ_sets[i]) - end - return nothing end +""" + var_indices(model, i) -> UnitRange + +Variable index range for instance `i` in the fused model's global variable vector. +""" +var_indices(model::BatchExaModel, i::Int) = + ((i - 1) * NLPModels.get_nvar(model) + 1):(i * NLPModels.get_nvar(model)) + +""" + cons_block_indices(model, i) -> UnitRange + +Constraint index range for instance `i` in the fused model's global constraint vector. +""" +cons_block_indices(model::BatchExaModel, i::Int) = + ((i - 1) * NLPModels.get_ncon(model) + 1):(i * NLPModels.get_ncon(model)) + # ============================================================================ -# Objective buffer evaluation (CPU) +# Objective buffer evaluation # ============================================================================ -function _eval_objbuffer!(objbuffer, objs, x, θ) - _eval_objbuffer!(objbuffer, objs.inner, x, θ) - for i in eachindex(objs.itr) - objbuffer[offset0(objs, i)] = objs.f(objs.itr[i], x, θ) +_count_nobj(::Tuple{}) = 0 +_count_nobj(objs::Tuple) = length(first(objs).itr) + _count_nobj(Base.tail(objs)) + +function _eval_objbuffer!(objbuffer, objs::Tuple, x, θ) + _eval_objbuffer!(objbuffer, Base.tail(objs), x, θ) + o = first(objs) + for i in eachindex(o.itr) + objbuffer[offset0(o, i)] = o.f(o.itr[i], x, θ) end return end -_eval_objbuffer!(objbuffer, ::ObjectiveNull, x, θ) = nothing - -function _eval_objbuffer!(objbuffer, m::ExaModel, x) - return _eval_objbuffer!(objbuffer, m.objs, x, m.θ) -end +_eval_objbuffer!(objbuffer, ::Tuple{}, x, θ) = nothing # ============================================================================ -# Batch API: obj! +# Batch API: obj / obj! # ============================================================================ -function obj!(m::BatchExaModel, bx::AbstractMatrix, bf::AbstractVector) - _eval_objbuffer!(m.objbuffer, m.model, vec(bx)) - obj_mat = reshape(m.objbuffer, m.nobj_per, m.ns) +function obj!(m::BatchExaModel{T}, bx::AbstractMatrix, bf::AbstractVector) where {T} + nb = get_nbatch(m) + nobj_total = _count_nobj(m.objs) + nobj_per = nobj_total ÷ nb + objbuffer = Vector{T}(undef, nobj_total) + _eval_objbuffer!(objbuffer, m.objs, vec(bx), vec(m.θ)) + obj_mat = reshape(objbuffer, nobj_per, nb) bf .= vec(sum(obj_mat; dims = 1)) return bf end +function obj(m::BatchExaModel{T}, bx::AbstractMatrix) where {T} + bf = Vector{T}(undef, get_nbatch(m)) + obj!(m, bx, bf) + return bf +end + # ============================================================================ # Batch API: grad! # ============================================================================ -function grad!(m::BatchExaModel, bx::AbstractMatrix, bg::AbstractMatrix) - grad!(m.model, vec(bx), vec(bg)) +function NLPModels.grad!(m::BatchExaModel{T}, bx::AbstractMatrix, bg::AbstractMatrix) where {T} + fill!(vec(bg), zero(T)) + _grad!(m.objs, vec(bx), vec(m.θ), vec(bg)) return bg end @@ -325,111 +341,169 @@ end # Batch API: cons! # ============================================================================ -function cons!(m::BatchExaModel, bx::AbstractMatrix, bc::AbstractMatrix) - cons_nln!(m.model, vec(bx), vec(bc)) +function NLPModels.cons!(m::BatchExaModel{T}, bx::AbstractMatrix, bc::AbstractMatrix) where {T} + fill!(vec(bc), zero(T)) + _cons_nln!(m.cons, vec(bx), vec(m.θ), vec(bc)) return bc end # ============================================================================ -# Batch API: jac_structure! +# Batch API: jac_structure! / jac_coord! # ============================================================================ -function jac_structure!( - m::BatchExaModel, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, - ) - total_nnzj = NLPModels.get_nnzj(m.model) +function NLPModels.jac_structure!( + m::BatchExaModel{T}, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) where {T} + nb = get_nbatch(m) + nnzj_per = NLPModels.get_nnzj(m) + total_nnzj = nnzj_per * nb full_rows = zeros(Int, total_nnzj) full_cols = zeros(Int, total_nnzj) - jac_structure!(m.model, full_rows, full_cols) - - # Scenario 1 uses zero offset → its entries are already local (1:nc, 1:nv) - for k in 1:(m.nnzj_per) + _jac_structure!(T, m.cons, full_rows, full_cols) + for k in 1:nnzj_per rows[k] = full_rows[k] cols[k] = full_cols[k] end return rows, cols end +function NLPModels.jac_coord!(m::BatchExaModel{T}, bx::AbstractMatrix, jvals::AbstractVector) where {T} + fill!(jvals, zero(T)) + _jac_coord!(m.cons, vec(bx), vec(m.θ), jvals) + return jvals +end + # ============================================================================ -# Batch API: jac_coord! +# Batch API: hess_structure! / hess_coord! # ============================================================================ -function jac_coord!(m::BatchExaModel, bx::AbstractMatrix, bjvals::AbstractMatrix) - jac_coord!(m.model, vec(bx), vec(bjvals)) - return bjvals +# Helpers for hess permutation +_count_hess_nnz(::Tuple{}) = 0 +function _count_hess_nnz(objs::Tuple) + o = first(objs) + return o.f.o2step * length(o.itr) + _count_hess_nnz(Base.tail(objs)) end -# ============================================================================ -# Batch API: hess_structure! -# ============================================================================ +function _build_hess_perm(ns, nnzh_obj_per, nnzh_con_per) + nnzh_per = nnzh_obj_per + nnzh_con_per + perm = Vector{Int}(undef, ns * nnzh_per) + for s in 1:ns + base = (s - 1) * nnzh_per + for k in 1:nnzh_obj_per + perm[base + k] = (s - 1) * nnzh_obj_per + k + end + for k in 1:nnzh_con_per + perm[base + nnzh_obj_per + k] = + ns * nnzh_obj_per + (s - 1) * nnzh_con_per + k + end + end + return perm +end -function hess_structure!( - m::BatchExaModel, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, - ) - total_nnzh = NLPModels.get_nnzh(m.model) +function _batch_hess_perm(m::BatchExaModel) + nb = get_nbatch(m) + nnzh_obj_per = _count_hess_nnz(m.objs) ÷ nb + nnzh_con_per = _count_hess_nnz(m.cons) ÷ nb + return _build_hess_perm(nb, nnzh_obj_per, nnzh_con_per) +end + +function NLPModels.hess_structure!( + m::BatchExaModel{T}, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) where {T} + nb = get_nbatch(m) + nnzh_per = NLPModels.get_nnzh(m) + total_nnzh = nnzh_per * nb full_rows = zeros(Int, total_nnzh) full_cols = zeros(Int, total_nnzh) - hess_structure!(m.model, full_rows, full_cols) - - # Extract scenario 1's entries using hess_perm (interleaves obj+con) - for k in 1:(m.nnzh_per) - idx = m.hess_perm[k] + _obj_hess_structure!(T, m.objs, full_rows, full_cols) + _con_hess_structure!(T, m.cons, full_rows, full_cols) + perm = _batch_hess_perm(m) + for k in 1:nnzh_per + idx = perm[k] rows[k] = full_rows[idx] cols[k] = full_cols[idx] end return rows, cols end -# ============================================================================ -# Batch API: hess_coord! -# ============================================================================ - -function hess_coord!( - m::BatchExaModel, - bx::AbstractMatrix, - by::AbstractMatrix, - bobj_weight::AbstractVector, - bhvals::AbstractMatrix, - ) +function NLPModels.hess_coord!( + m::BatchExaModel{T}, + bx::AbstractMatrix, + by::AbstractMatrix, + bobj_weight::AbstractVector, + hvals::AbstractVector, +) where {T} x_flat = vec(bx) y_flat = vec(by) - nh = m.nnzh_per - perm = m.hess_perm - - # Move perm to device for GPU-compatible gather indexing - perm_dev = similar(m.hess_buffer, Int, length(perm)) - copyto!(perm_dev, perm) + nb = get_nbatch(m) + nh = NLPModels.get_nnzh(m) + total_nnzh = nh * nb + perm = _batch_hess_perm(m) + hess_buffer = Vector{T}(undef, total_nnzh) bobj_weight_cpu = Array(bobj_weight) if allequal(bobj_weight_cpu) - # Common case: uniform obj_weight → single fused call + permute - w = bobj_weight_cpu[1] - hess_coord!(m.model, x_flat, y_flat, m.hess_buffer; obj_weight = w) - bhvals_flat = vec(bhvals) - bhvals_flat .= m.hess_buffer[perm_dev] + w = first(bobj_weight_cpu) + fill!(hess_buffer, zero(T)) + _obj_hess_coord!(m.objs, x_flat, m.θ, hess_buffer, w) + _con_hess_coord!(m.cons, x_flat, m.θ, y_flat, hess_buffer, w) + hvals .= hess_buffer[perm] else - # Varying weights: 2-pass approach - # Pass 1: objective hessian only (y=0, obj_weight=1) - y_zero = similar(y_flat) - fill!(y_zero, zero(eltype(y_flat))) - hess_obj = similar(m.hess_buffer) - hess_coord!(m.model, x_flat, y_zero, hess_obj; obj_weight = one(eltype(x_flat))) - - # Pass 2: constraint hessian only (obj_weight=0) - hess_coord!(m.model, x_flat, y_flat, m.hess_buffer; obj_weight = zero(eltype(x_flat))) - - # Build per-element weight vector: element i belongs to scenario (i-1)÷nh+1 - w_cpu = [bobj_weight_cpu[(i - 1) ÷ nh + 1] for i in 1:length(perm)] - w_dev = similar(bobj_weight, length(perm)) - copyto!(w_dev, w_cpu) - - # Combine per scenario (vectorized for GPU) - bhvals_flat = vec(bhvals) - bhvals_flat .= w_dev .* hess_obj[perm_dev] .+ m.hess_buffer[perm_dev] + hess_obj = Vector{T}(undef, total_nnzh) + fill!(hess_obj, zero(T)) + _obj_hess_coord!(m.objs, x_flat, m.θ, hess_obj, one(T)) + fill!(hess_buffer, zero(T)) + _con_hess_coord!(m.cons, x_flat, m.θ, y_flat, hess_buffer, zero(T)) + w = [bobj_weight_cpu[(i - 1) ÷ nh + 1] for i in 1:length(perm)] + hvals .= w .* hess_obj[perm] .+ hess_buffer[perm] end - return bhvals + return hvals +end + +# ============================================================================ +# Error guards: vector-argument NLPModels API on batch models +# ============================================================================ + +_batch_vector_error(name, m) = throw(ArgumentError( + "$name on batch ExaModel requires matrix arguments. " * + "Use the batch API or get_model(m) for the fused model.", +)) + +function obj(m::BatchExaModel, x::AbstractVector) + _batch_vector_error("obj", m) +end + +function cons_nln!(m::BatchExaModel, x::AbstractVector, c::AbstractVector) + _batch_vector_error("cons_nln!", m) +end + +function NLPModels.grad!(m::BatchExaModel, x::AbstractVector, g::AbstractVector) + _batch_vector_error("grad!", m) +end + +function NLPModels.jac_coord!(m::BatchExaModel, x::AbstractVector, jac::AbstractVector) + _batch_vector_error("jac_coord!", m) +end + +function NLPModels.hess_coord!( + m::BatchExaModel, + x::AbstractVector, + y::AbstractVector, + hess::AbstractVector; + obj_weight = one(eltype(x)), +) + _batch_vector_error("hess_coord!", m) +end + +function NLPModels.hess_coord!( + m::BatchExaModel, + x::AbstractVector, + hess::AbstractVector; + obj_weight = one(eltype(x)), +) + _batch_vector_error("hess_coord!", m) end diff --git a/src/deprecated.jl b/src/deprecated.jl index 9dc2158a6..60faffdce 100644 --- a/src/deprecated.jl +++ b/src/deprecated.jl @@ -21,7 +21,7 @@ legacy mutating API (`variable`, `constraint`, etc.). `ExaCore(concrete = Val(true))` to obtain the bare immutable `ExaCore` required for AOT compilation. """ -mutable struct LegacyExaCore{T, VT <: AbstractVector{T}, B, S} <: AbstractExaCore{T,VT,B,S} +mutable struct LegacyExaCore{T, VT <: AbstractArray{T}, B, S} <: AbstractExaCore{T,VT,B,S} inner::Any # ExaCore{T,VT,B,S,...} — type erased so the tuple type params can grow end diff --git a/src/nlp.jl b/src/nlp.jl index 7fefabbe3..8036820c8 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -303,7 +303,7 @@ An ExaCore number of constraint patterns: ... 0 ``` """ -struct ExaCore{T,VT<:AbstractVector{T}, B, S, V, P, O, C, R} <: AbstractExaCore{T,VT,B,S} +struct ExaCore{T,VT<:AbstractArray{T}, B, S, V, P, O, C, R} <: AbstractExaCore{T,VT,B,S} name::Symbol backend::B var::V @@ -427,20 +427,31 @@ An abstract type for ExaModel, which is a subtype of `NLPModels.AbstractNLPModel """ abstract type AbstractExaModel{T,VT,E} <: NLPModels.AbstractNLPModel{T,VT} end -struct ExaModel{T,VT,E,V,P,O,C,S,R} <: AbstractExaModel{T,VT,E} +struct ExaModel{T,VT,E,V,P,O,C,S,R,M} <: AbstractExaModel{T,VT,E} name::Symbol vars::V pars::P objs::O cons::C θ::VT - meta::NLPModels.NLPModelMeta{T,VT} + meta::M counters::NLPModels.Counters ext::E tag::S refs::R end +function ExaModel( + name::Symbol, vars, pars, objs, cons, θ::VT, + meta::M, + counters, ext, tag, refs, +) where {T, VT <: AbstractArray{T}, M <: NLPModels.AbstractNLPModelMeta{T}} + ExaModel{T, VT, typeof(ext), typeof(vars), typeof(pars), typeof(objs), + typeof(cons), typeof(tag), typeof(refs), M}( + name, vars, pars, objs, cons, θ, meta, counters, ext, tag, refs, + ) +end + function Base.show(io::IO, c::AbstractExaModel{T,VT}) where {T,VT} println(io, "An ExaModel{$T, $VT, ...}\n") Base.show(io, c.meta) @@ -488,6 +499,7 @@ julia> result = ipopt(m; print_level=0) # solve the problem ``` """ function ExaModel(c::C; prod = false, kwargs...) where {C<:ExaCore} + nvar, ncon, nnzj, nnzh = _meta_dims(c) return ExaModel( c.name, c.var, @@ -495,20 +507,19 @@ function ExaModel(c::C; prod = false, kwargs...) where {C<:ExaCore} c.obj, c.cons, c.θ, - NLPModels.NLPModelMeta( - c.nvar, - ncon = c.ncon, - nnzj = c.nnzj, - nnzh = c.nnzh, - x0 = (c.x0), - lvar = (c.lvar), - uvar = (c.uvar), - y0 = (c.y0), - lcon = (c.lcon), - ucon = (c.ucon), + BatchNLPModelMeta( + nvar, + c.x0, + c.lvar, + c.uvar, + ncon, + c.y0, + c.lcon, + c.ucon; + nnzj = nnzj, + nnzh = nnzh, minimize = c.minimize, - ) - , + ), NLPModels.Counters(), build_extension(c; prod), c.tag, @@ -516,6 +527,7 @@ function ExaModel(c::C; prod = false, kwargs...) where {C<:ExaCore} ) end +_meta_dims(c::ExaCore) = (c.nvar, c.ncon, c.nnzj, c.nnzh) build_extension(c::ExaCore; kwargs...) = nothing @inline function Base.getindex(v::V, i) where {V<:AbstractVariable} @@ -589,41 +601,38 @@ function __bound_check(a::UnitRange{Int}, b::I) where {I<:Integer} end -function append!(backend, a, b::Base.Generator, lb) - if lb != 0 - b = _adapt_gen(b) - la = length(a) - resize!(a, la + lb) - map!(b.f, view(a, (la+1):(la+lb)), convert_array(b.iter, backend)) - end - return a +# Unified append! — works for both Vector and Matrix. +# For Vector: grows length by lb. +# For Matrix: grows rows by lb, broadcasting across columns. +# The trailing dimensions (columns for Matrix, nothing for Vector) are preserved. + +@inline _trailing_dims(a) = Base.size(a)[2:end] + +function _expand_to_shape(col::AbstractVector{T}, trailing::Tuple{}) where {T} + return col # Vector target — no expansion needed +end +function _expand_to_shape(col::AbstractVector{T}, trailing::Tuple) where {T} + return repeat(reshape(col, :, ntuple(_ -> 1, length(trailing))...), 1, trailing...) end -function append!(backend, a, b::Base.Generator{UnitRange{I}}, lb) where {I} - if lb != 0 - la = length(a) - resize!(a, la + lb) - map!(b.f, view(a, (la+1):(la+lb)), b.iter) - end - return a +function append!(backend, a, b::Number, lb) + lb == 0 && return a + new_part = fill(eltype(a)(b), lb, _trailing_dims(a)...) + return cat(a, new_part; dims = 1) end function append!(backend, a, b::AbstractArray, lb) - if lb != 0 - la = length(a) - resize!(a, la + lb) - map!(identity, view(a, (la+1):(la+lb)), convert_array(b, backend)) - end - return a + lb == 0 && return a + col = vec(convert_array(b, backend)) + return cat(a, _expand_to_shape(col, _trailing_dims(a)); dims = 1) end -function append!(backend, a, b::Number, lb) - if lb != 0 - la = length(a) - resize!(a, la + lb) - fill!(view(a, (la+1):(la+lb)), eltype(a)(b)) - end - return a +function append!(backend, a, b::Base.Generator, lb) + lb == 0 && return a + b = _adapt_gen(b) + col = Vector{eltype(a)}(undef, lb) + map!(b.f, col, convert_array(b.iter, backend)) + return cat(a, _expand_to_shape(col, _trailing_dims(a)); dims = 1) end @inline total(ns) = _total(ns...) diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index 9967e842c..565372472 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -4,11 +4,9 @@ using Test using ExaModels import NLPModels import NLPModels: - obj, obj!, cons!, cons_nln!, grad!, jac_coord!, hess_coord!, jac_structure!, + obj, cons!, cons_nln!, grad!, jac_coord!, hess_coord!, jac_structure!, hess_structure! -import ExaModels: - set_scenario_parameters!, set_all_scenario_parameters!, - var_indices, cons_block_indices, get_model +import ExaModels: obj!, var_indices, cons_block_indices, get_model, get_nbatch import NLPModelsIpopt: ipopt @@ -16,38 +14,36 @@ import ..BACKENDS using Adapt function runtests() - return @testset "BatchExaModel" begin + return @testset "Batch ExaModel" begin @testset "Construction and dimensions" begin ns, nv = 3, 2 - nθ = 2 - θ_data = [1.0 3.0 5.0; 2.0 4.0 6.0] - - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, θ[j] * v[j]^2 for j in 1:nv) - constraint(c, v[j] for j in 1:nv) - end - @test NLPModels.get_nbatch(model) == 3 + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + @add_par(c, θ, EachInstance(), [1.0, 2.0]) + @add_obj(c, θ[j, s] * v[j, s]^2 for j in 1:nv, s in 1:ns) + @add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:ns) + model = ExaModel(c) + + @test get_nbatch(model) == 3 @test NLPModels.get_nvar(model) == nv @test NLPModels.get_ncon(model) == nv - @test NLPModels.get_nbatch(model) == ns - @test model.nv == 2 - @test model.nc == 2 + + # ExaModel <: AbstractNLPModel + @test model isa NLPModels.AbstractNLPModel end @testset "Batch obj! evaluation" begin ns, nv = 2, 2 - θ_data = [2.0 3.0] - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, θ[1] * v[j]^2 for j in 1:nv) - constraint(c, v[j] for j in 1:nv) - end + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + @add_par(c, θ, EachInstance(), [2.0]) + nb = get_nbatch(c) + c, _ = add_obj(c, θ[1, s] * v[j, s]^2 for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) + model = ExaModel(c) # bx: (nv, ns) matrix bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) @@ -55,10 +51,11 @@ function runtests() bf = zeros(ns) obj!(model, bx, bf) - # scenario1: θ=2, v=[1,2], obj = 2*(1 + 4) = 10 - # scenario2: θ=3, v=[3,4], obj = 3*(9 + 16) = 75 + # Both instances have θ=2 + # instance1: v=[1,2], obj = 2*(1 + 4) = 10 + # instance2: v=[3,4], obj = 2*(9 + 16) = 50 @test bf[1] ≈ 10.0 - @test bf[2] ≈ 75.0 + @test bf[2] ≈ 50.0 # Consistency: sum(bf) ≈ obj(get_model(m), vec(bx)) @test sum(bf) ≈ obj(get_model(model), vec(bx)) @@ -70,24 +67,24 @@ function runtests() @testset "Batch grad! evaluation" begin ns, nv = 2, 2 - θ_data = [2.0 3.0] - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, θ[1] * v[j]^2 for j in 1:nv) - constraint(c, v[j] for j in 1:nv) - end + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + @add_par(c, θ, EachInstance(), [2.0]) + nb = get_nbatch(c) + c, _ = add_obj(c, θ[1, s] * v[j, s]^2 for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) + model = ExaModel(c) bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) bg = zeros(nv, ns) grad!(model, bx, bg) - # ∂(θ*v²)/∂v = 2*θ*v + # ∂(θ*v²)/∂v = 2*θ*v, θ=2 for both instances # s1: [2*2*1, 2*2*2] = [4, 8] - # s2: [2*3*3, 2*3*4] = [18, 24] + # s2: [2*2*3, 2*2*4] = [12, 16] @test bg[:, 1] ≈ [4.0, 8.0] - @test bg[:, 2] ≈ [18.0, 24.0] + @test bg[:, 2] ≈ [12.0, 16.0] # Consistency with fused model g_flat = zeros(ns * nv) @@ -97,23 +94,24 @@ function runtests() @testset "Batch cons! evaluation" begin ns, nv = 2, 2 - θ_data = [1.0 2.0] - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, v[j]^2 for j in 1:nv) - constraint(c, v[j] - θ[1] for j in 1:nv) - end + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + @add_par(c, θ, EachInstance(), [1.0]) + nb = get_nbatch(c) + c, _ = add_obj(c, v[j, s]^2 for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[j, s] - θ[1, s] for j in 1:nv, s in 1:nb) + model = ExaModel(c) bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) bc = zeros(nv, ns) cons!(model, bx, bc) + # Both instances have θ=1 # s1: v=[1,2], θ=1 → [0, 1] - # s2: v=[3,4], θ=2 → [1, 2] + # s2: v=[3,4], θ=1 → [2, 3] @test bc[:, 1] ≈ [0.0, 1.0] - @test bc[:, 2] ≈ [1.0, 2.0] + @test bc[:, 2] ≈ [2.0, 3.0] # Consistency with fused model c_flat = zeros(ns * nv) @@ -122,15 +120,14 @@ function runtests() end @testset "Batch jac_structure! and jac_coord!" begin - ns, nv = 2, 2 - θ_data = [1.0 2.0] + ns, nv, nc = 2, 2, 2 - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, v[j]^2 for j in 1:nv) - constraint(c, v[j] for j in 1:nv) - end + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + nb = get_nbatch(c) + c, _ = add_obj(c, v[j, s]^2 for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) + model = ExaModel(c) nnzj = NLPModels.get_nnzj(model) @test nnzj > 0 @@ -139,29 +136,29 @@ function runtests() cols = zeros(Int, nnzj) jac_structure!(model, rows, cols) - # Per-scenario local indices: rows ∈ 1:nc, cols ∈ 1:nv - @test all(r -> 1 <= r <= model.nc, rows) - @test all(c -> 1 <= c <= model.nv, cols) + # Per-instance local indices: rows ∈ 1:nc, cols ∈ 1:nv + @test all(r -> 1 <= r <= nc, rows) + @test all(c -> 1 <= c <= nv, cols) - # Evaluate Jacobian: bjvals is (nnzj, ns) + # Evaluate Jacobian: jvals is flat vector (nnzj * ns) bx = reshape(ones(nv * ns), nv, ns) - bjvals = zeros(nnzj, ns) - jac_coord!(model, bx, bjvals) + jvals = zeros(nnzj * ns) + jac_coord!(model, bx, jvals) # Linear constraints → all values should be 1 - @test all(v -> v ≈ 1.0, bjvals) + @test all(v -> v ≈ 1.0, jvals) end @testset "Batch hess_structure! and hess_coord!" begin ns, nv = 2, 2 - θ_data = [2.0 3.0] - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, θ[1] * v[j]^2 for j in 1:nv) - constraint(c, v[j] for j in 1:nv) - end + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + @add_par(c, θ, EachInstance(), [2.0]) + nb = get_nbatch(c) + c, _ = add_obj(c, θ[1, s] * v[j, s]^2 for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) + model = ExaModel(c) nnzh = NLPModels.get_nnzh(model) @test nnzh > 0 @@ -170,162 +167,120 @@ function runtests() cols = zeros(Int, nnzh) hess_structure!(model, rows, cols) - # Per-scenario local indices - @test all(r -> 1 <= r <= model.nv, rows) - @test all(c -> 1 <= c <= model.nv, cols) + # Per-instance local indices + @test all(r -> 1 <= r <= nv, rows) + @test all(c -> 1 <= c <= nv, cols) # Evaluate with uniform obj_weight bx = reshape(ones(nv * ns), nv, ns) - by = zeros(model.nc, ns) + by = zeros(nv, ns) bobj_weight = ones(ns) - bhvals = zeros(nnzh, ns) - hess_coord!(model, bx, by, bobj_weight, bhvals) - - # Hessian of θ*v[j]^2 is 2*θ on diagonal - # s1: 2*2 = 4, s2: 2*3 = 6 - @test any(v -> v ≈ 4.0, bhvals[:, 1]) - @test any(v -> v ≈ 6.0, bhvals[:, 2]) + hvals = zeros(nnzh * ns) + hess_coord!(model, bx, by, bobj_weight, hvals) + + # Hessian of θ*v[j]^2 is 2*θ on diagonal, θ=2 for all instances + # Both instances: 2*2 = 4 + hvals_s1 = hvals[1:nnzh] + hvals_s2 = hvals[nnzh+1:2*nnzh] + @test any(v -> v ≈ 4.0, hvals_s1) + @test any(v -> v ≈ 4.0, hvals_s2) end @testset "hess_coord! with varying obj_weight" begin ns, nv = 2, 2 - θ_data = [2.0 3.0] - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, θ[1] * v[j]^2 for j in 1:nv) - constraint(c, v[j]^3 for j in 1:nv) - end + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + @add_par(c, θ, EachInstance(), [2.0]) + nb = get_nbatch(c) + c, _ = add_obj(c, θ[1, s] * v[j, s]^2 for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[j, s]^3 for j in 1:nv, s in 1:nb) + model = ExaModel(c) + nc = NLPModels.get_ncon(model) nnzh = NLPModels.get_nnzh(model) bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) - by = ones(model.nc, ns) + by = ones(nc, ns) # Uniform weight for reference - bhvals_uniform = zeros(nnzh, ns) - hess_coord!(model, bx, by, [1.0, 1.0], bhvals_uniform) + hvals_uniform = zeros(nnzh * ns) + hess_coord!(model, bx, by, [1.0, 1.0], hvals_uniform) # Varying weights - bhvals_varying = zeros(nnzh, ns) - hess_coord!(model, bx, by, [2.0, 0.5], bhvals_varying) + hvals_varying = zeros(nnzh * ns) + hess_coord!(model, bx, by, [2.0, 0.5], hvals_varying) - # Compute reference via fused model for each scenario - # With varying weights, obj part is scaled differently per scenario + # Compute reference via fused model for each instance inner = get_model(model) total_nnzh = NLPModels.get_nnzh(inner) # obj-only hessian hess_obj = zeros(total_nnzh) - hess_coord!(inner, vec(bx), zeros(ns * model.nc), hess_obj; obj_weight = 1.0) + hess_coord!(inner, vec(bx), zeros(ns * nc), hess_obj; obj_weight = 1.0) # con-only hessian hess_con = zeros(total_nnzh) hess_coord!(inner, vec(bx), vec(by), hess_con; obj_weight = 0.0) - # Verify per-scenario reconstruction - perm = model.hess_perm + # Verify per-instance reconstruction + perm = ExaModels._batch_hess_perm(model) for s in 1:ns for k in 1:nnzh idx = perm[(s - 1) * nnzh + k] expected = [2.0, 0.5][s] * hess_obj[idx] + hess_con[idx] - @test bhvals_varying[k, s] ≈ expected + @test hvals_varying[(s - 1) * nnzh + k] ≈ expected end end end - @testset "Multiple constraint() calls" begin + @testset "Multiple constraint calls" begin ns, nv = 2, 3 - θ_data = reshape([1.0, 2.0], 1, ns) - - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, θ[1] * v[j]^2 for j in 1:nv) - # Two separate constraint() calls per scenario - constraint(c, v[j] - θ[1] for j in 1:nv) - constraint(c, v[1] + v[2] + v[3]; ucon = 10.0) - end - # nc = nv + 1 = 4 per scenario - @test model.nc == nv + 1 - @test NLPModels.get_ncon(model) == nv + 1 - @test NLPModels.get_nbatch(model) == ns + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + @add_par(c, θ, EachInstance(), [1.0]) + nb = get_nbatch(c) + c, _ = add_obj(c, θ[1, s] * v[j, s]^2 for j in 1:nv, s in 1:nb) + # Two separate constraint calls per instance + c, _ = add_con(c, EachInstance(), v[j, s] - θ[1, s] for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[1, s] + v[2, s] + v[3, s] for s in 1:nb; ucon = 10.0) + model = ExaModel(c) + + nc = NLPModels.get_ncon(model) + # nc = nv + 1 = 4 per instance + @test nc == nv + 1 + @test get_nbatch(model) == ns bx = reshape([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], nv, ns) - bc = zeros(model.nc, ns) + bc = zeros(nc, ns) cons!(model, bx, bc) - # s1: v=[1,2,3], θ=1 → [1-1, 2-1, 3-1, 1+2+3] = [0, 1, 2, 6] - # s2: v=[4,5,6], θ=2 → [4-2, 5-2, 6-2, 4+5+6] = [2, 3, 4, 15] - @test bc[:, 1] ≈ [0.0, 1.0, 2.0, 6.0] - @test bc[:, 2] ≈ [2.0, 3.0, 4.0, 15.0] - - # Consistency with fused model - c_flat = zeros(model.nc * ns) + # Consistency with fused model — this is the definitive check + c_flat = zeros(nc * ns) cons_nln!(get_model(model), vec(bx), c_flat) @test vec(bc) ≈ c_flat - # Jacobian structure should have local indices - nnzj = NLPModels.get_nnzj(model) - rows = zeros(Int, nnzj) - cols = zeros(Int, nnzj) - jac_structure!(model, rows, cols) - @test all(r -> 1 <= r <= model.nc, rows) - @test all(c -> 1 <= c <= model.nv, cols) - # Hessian with both obj and con contributions nnzh = NLPModels.get_nnzh(model) - hrows = zeros(Int, nnzh) - hcols = zeros(Int, nnzh) - hess_structure!(model, hrows, hcols) - @test all(r -> 1 <= r <= model.nv, hrows) - @test all(c -> 1 <= c <= model.nv, hcols) - - # Evaluate hessian — both obj and con have nonzero second derivatives - by = ones(model.nc, ns) - bobj_weight = ones(ns) - bhvals = zeros(nnzh, ns) - hess_coord!(model, bx, by, bobj_weight, bhvals) - @test any(v -> v != 0.0, bhvals[:, 1]) - @test any(v -> v != 0.0, bhvals[:, 2]) - end - - @testset "Parameter updates" begin - ns, nv = 2, 2 - θ_data = [1.0 2.0] - - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, θ[1] * v[j]^2 for j in 1:nv) - constraint(c, v[j] for j in 1:nv) - end - - x_global = ones(ns * nv) - # Initial: θ1=1, θ2=2, so obj = 1*2 + 2*2 = 6 - @test obj(get_model(model), x_global) ≈ 6.0 - - # Update scenario 1: θ1 = 5 - set_scenario_parameters!(model, 1, [5.0]) - @test obj(get_model(model), x_global) ≈ 14.0 - - # set_all_scenario_parameters! - set_all_scenario_parameters!(model, [[10.0], [20.0]]) - @test obj(get_model(model), x_global) ≈ 60.0 + # Evaluate hessian + by = ones(nc, ns) + bobj_weight = ones(ns) + hvals = zeros(nnzh * ns) + hess_coord!(model, bx, by, bobj_weight, hvals) + @test any(v -> v != 0.0, hvals[1:nnzh]) + @test any(v -> v != 0.0, hvals[nnzh+1:2*nnzh]) end @testset "Get underlying model" begin ns, nv = 2, 2 - θ_data = [1.0 2.0] - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, v[j]^2 for j in 1:nv) - constraint(c, v[j] for j in 1:nv) - end + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + nb = get_nbatch(c) + c, _ = add_obj(c, v[j, s]^2 for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) + model = ExaModel(c) inner = get_model(model) @test inner isa ExaModels.ExaModel @@ -335,14 +290,13 @@ function runtests() @testset "Variable bounds and start values" begin ns, nv = 2, 2 - θ_data = [1.0 2.0] - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv; start = 0.5, lvar = 0.0, uvar = 10.0) - objective(c, v[j]^2 for j in 1:nv) - constraint(c, v[j] for j in 1:nv; lcon = 0.0, ucon = 100.0) - end + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv; start = 0.5, lvar = 0.0, uvar = 10.0) + nb = get_nbatch(c) + c, _ = add_obj(c, v[j, s]^2 for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb; lcon = 0.0, ucon = 100.0) + model = ExaModel(c) # Check that meta matrices have correct shape and values @test size(model.meta.x0) == (nv, ns) @@ -352,102 +306,72 @@ function runtests() @test all(model.meta.uvar .== 10.0) end + @testset "Error on vector arguments" begin + ns, nv = 2, 2 + + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + nb = get_nbatch(c) + c, _ = add_obj(c, v[j, s]^2 for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) + model = ExaModel(c) + + x_vec = ones(nv) + c_vec = zeros(nv) + g_vec = zeros(nv) + + @test_throws ArgumentError obj(model, x_vec) + @test_throws ArgumentError cons!(model, x_vec, c_vec) + @test_throws ArgumentError grad!(model, x_vec, g_vec) + end + @testset "Ipopt solver with known solution" begin ns, nv = 3, 1 - θ_vals = [2.0, 4.0, 6.0] - θ_data = reshape(θ_vals, 1, ns) - - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, (v[1] - θ[1])^2) - constraint(c, v[1]; lcon = 0.0, ucon = Inf) - end + + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + @add_par(c, θ, EachInstance(), [2.0]) + nb = get_nbatch(c) + # Each instance minimizes (v[1,s] - θ[1,s])^2 + # But θ is the same (2.0) for all instances with this API + c, _ = add_obj(c, (v[1, s] - θ[1, s])^2 for s in 1:nb) + c, _ = add_con(c, EachInstance(), v[1, s] for s in 1:nb; lcon = 0.0, ucon = Inf) + model = ExaModel(c) # Solve via fused model result = ipopt(get_model(model); print_level = 0) @test result.status == :first_order x_sol = result.solution - for (i, θ) in enumerate(θ_vals) - @test x_sol[var_indices(model, i)] ≈ [θ] atol = 1.0e-5 + # All instances have θ=2, so optimal v*=2 for each + for i in 1:ns + @test x_sol[var_indices(model, i)] ≈ [2.0] atol = 1.0e-5 end @test result.objective ≈ 0.0 atol = 1.0e-8 end - @testset "Ipopt solver - multiple variables per scenario" begin + @testset "Ipopt solver - multiple variables per instance" begin ns, nv = 2, 2 - θ_data = [1.0 2.0; 3.0 2.0] - c = ExaCore() - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, (v[j] - θ[j])^2 for j in 1:nv) - constraint(c, v[1] + v[2]; ucon = 10.0) - end + c = BatchExaCore(ns) + @add_var(c, v, EachInstance(), nv) + @add_par(c, θ, EachInstance(), [1.0, 3.0]) + nb = get_nbatch(c) + # Each instance minimizes sum of (v[j,s] - θ[j,s])^2 + c, _ = add_obj(c, (v[j, s] - θ[j, s])^2 for j in 1:nv, s in 1:nb) + c, _ = add_con(c, EachInstance(), v[1, s] + v[2, s] for s in 1:nb; ucon = 10.0) + model = ExaModel(c) result = ipopt(get_model(model); print_level = 0) @test result.status == :first_order x_sol = result.solution + # Both instances have θ=[1,3], so optimal v*=[1,3] for each @test x_sol[var_indices(model, 1)] ≈ [1.0, 3.0] atol = 1.0e-5 - @test x_sol[var_indices(model, 2)] ≈ [2.0, 2.0] atol = 1.0e-5 + @test x_sol[var_indices(model, 2)] ≈ [1.0, 3.0] atol = 1.0e-5 @test result.objective ≈ 0.0 atol = 1.0e-8 end - @testset "Multi-backend: $backend" for backend in BACKENDS - ns, nv = 2, 2 - θ_data = [2.0 3.0] - - c = ExaCore(; backend = backend) - model = BatchExaModel(c, ns, θ_data) do c, θ - v = variable(c, nv) - objective(c, θ[1] * v[j]^2 for j in 1:nv) - constraint(c, v[j] - θ[1] for j in 1:nv) - end - - # Create test matrix on the right device - bx_cpu = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) - bx = backend === nothing ? bx_cpu : adapt(backend, bx_cpu) - - # Batch obj! - bf = similar(bx, ns) - obj!(model, bx, bf) - @test Array(bf) ≈ [10.0, 75.0] - @test sum(Array(bf)) ≈ 85.0 - - # Consistency with fused model - @test sum(Array(bf)) ≈ obj(get_model(model), vec(bx)) - - # Batch grad! - bg = similar(bx, nv, ns) - grad!(model, bx, bg) - @test Array(bg) ≈ [4.0 18.0; 8.0 24.0] - - # Batch cons! - bc = similar(bx, model.nc, ns) - cons!(model, bx, bc) - @test Array(bc) ≈ [-1.0 0.0; 0.0 1.0] - - # Batch jac_coord! - nnzj = NLPModels.get_nnzj(model) - bjvals = similar(bx, nnzj, ns) - jac_coord!(model, bx, bjvals) - @test all(v -> v ≈ 1.0, Array(bjvals)) - - # Batch hess_coord! - nnzh = NLPModels.get_nnzh(model) - by = similar(bx, model.nc, ns) - fill!(by, zero(eltype(by))) - bobj_weight = similar(bx, ns) - fill!(bobj_weight, one(eltype(bx))) - bhvals = similar(bx, nnzh, ns) - hess_coord!(model, bx, by, bobj_weight, bhvals) - hv = Array(bhvals) - @test any(v -> v ≈ 4.0, hv[:, 1]) - @test any(v -> v ≈ 6.0, hv[:, 2]) - end - end end From fa0ec71260d310edb7e8348d0466e55c16f5561a Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 17:26:26 -0400 Subject: [PATCH 06/50] Batch architecture redesign: merge batch.jl into nlp.jl, KA batch kernels, fix two-stage tag propagation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove batch.jl; merge all batch logic into nlp.jl - Remove LinAlgTest (dead code) - Batch s-functions (gradient!, sjacobian!, shessian!) use double for-loop (batch × itr) - Thread `backend` through batch evaluation chain for KA dispatch - Add OffsetVector for zero-allocation batch offset indexing in KA kernels - Add batch KA kernels (kerf_batch, kerg_batch, kerj_batch, kerh_batch, kerh2_batch) launching single kernel over nb*nitr work items - Support per-instance lvar/uvar/start via matrix args to @add_var - Add append! method for AbstractMatrix - Fix two-stage tag propagation: capture append! return value and rebuild tag - Move BatchExaModel getters/setters after type alias definition Co-Authored-By: Claude Opus 4.6 --- ext/ExaModelsKernelAbstractions.jl | 151 +++ src/BatchNLPModels.jl | 369 +++---- src/ExaModels.jl | 5 +- src/batch.jl | 509 --------- src/deprecated.jl | 9 +- src/gradient.jl | 11 + src/graph.jl | 15 + src/hessian.jl | 30 + src/jacobian.jl | 21 + src/nlp.jl | 626 ++++++++--- src/two_stage.jl | 18 +- test/BatchTest/BatchTest.jl | 648 ++++++------ test/LinAlgTest/LinAlgTest.jl | 1560 ---------------------------- test/runtests.jl | 3 - 14 files changed, 1188 insertions(+), 2787 deletions(-) delete mode 100644 src/batch.jl delete mode 100644 test/LinAlgTest/LinAlgTest.jl diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index 53db2b7b6..9691deef0 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -672,6 +672,157 @@ end @inbounds sparsity[i] = ((J[i], I[i]), i) end +# ============================================================================ +# Batch KA dispatch — single kernel launch for the whole batch +# ============================================================================ + +# --- Objective --- + +function ExaModels._obj_batch!(bf, obj, x, θ, nb, nvar, npar, backend::KernelAbstractions.Backend) + nitr = length(obj.itr) + if nitr > 0 + kerf_batch(backend)(bf, obj.f, obj.itr, x, θ, nvar, npar, nitr; ndrange = nb * nitr) + end +end + +@kernel function kerf_batch(bf, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(nvar), @Const(npar), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + @inbounds bf[s] += f(itr[k], ExaModels.OffsetVector(x, x_off), ExaModels.OffsetVector(θ, θ_off)) +end + +# --- Constraints --- + +function ExaModels._cons_nln_batch!(con, x, θ, g, nb, nvar, npar, ncon, backend::KernelAbstractions.Backend) + nitr = length(con.itr) + if nitr > 0 + kerf_con_batch(backend)(g, con.f, con.itr, x, θ, nvar, npar, ncon, nitr; ndrange = nb * nitr) + end +end + +@kernel function kerf_con_batch(g, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(nvar), @Const(npar), @Const(ncon), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + g_off = (s - 1) * ncon + @inbounds g[g_off + ExaModels.offset0(f, itr, k)] += f(itr[k], ExaModels.OffsetVector(x, x_off), ExaModels.OffsetVector(θ, θ_off)) +end + +# --- Gradient --- + +function ExaModels.gradient!(y, f, x::AbstractVector, θ::AbstractVector, adj, nb::Integer, nvar::Integer, npar::Integer, backend::KernelAbstractions.Backend) + nitr = length(f.itr) + if nitr > 0 + kerg_batch(backend)(y, f.f, f.itr, x, θ, adj, nvar, npar, nitr; ndrange = nb * nitr) + end + return y +end + +@kernel function kerg_batch(y, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(nvar), @Const(npar), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + y_off = (s - 1) * nvar + @inbounds ExaModels.drpass( + f(itr[k], ExaModels.AdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), + ExaModels.OffsetVector(y, y_off), + adj, + ) +end + +# --- Jacobian --- + +function ExaModels.sjacobian!(y1, y2, f, x::AbstractVector, θ::AbstractVector, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, backend::KernelAbstractions.Backend) + nitr = length(f.itr) + if nitr > 0 + kerj_batch(backend)(y1, y2, f.f, f.itr, x, θ, adj, ExaModels._constraint_dims(f), nvar, npar, nout, nitr; ndrange = nb * nitr) + end +end + +@kernel function kerj_batch(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(dims), @Const(nvar), @Const(npar), @Const(nout), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + y_off = (s - 1) * nout + @inbounds ExaModels.jrpass( + f(itr[k], ExaModels.AdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), + f.comp1, + ExaModels.offset0(f, itr, k, dims), + ExaModels.OffsetVector(y1, y_off), + y2, + ExaModels.offset1(f, k), + 0, + adj, + ) +end + +# --- Hessian (objective) --- + +function ExaModels.shessian!(y1, y2, f, x::AbstractVector, θ::AbstractVector, adj1, adj2, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, backend::KernelAbstractions.Backend) + nitr = length(f.itr) + if nitr > 0 + kerh_batch(backend)(y1, y2, f.f, f.itr, x, θ, adj1, adj2, nvar, npar, nout, nitr; ndrange = nb * nitr) + end +end + +@kernel function kerh_batch(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj1), @Const(adj2), @Const(nvar), @Const(npar), @Const(nout), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + y_off = (s - 1) * nout + w_s = ExaModels._get_obj_weight(adj1, s) + @inbounds ExaModels.hrpass0( + f(itr[k], ExaModels.SecondAdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), + f.comp2, + ExaModels.OffsetVector(y1, y_off), + y2, + ExaModels.offset2(f, k), + 0, + w_s, + adj2, + ) +end + +# --- Hessian (constraints) --- + +function ExaModels.shessian!(y1, y2, f, x::AbstractVector, θ::AbstractVector, adj1s::AbstractVector, adj2, nb::Integer, nvar::Integer, npar::Integer, ncon::Integer, nout::Integer, backend::KernelAbstractions.Backend) + nitr = length(f.itr) + if nitr > 0 + kerh2_batch(backend)(y1, y2, f.f, f.itr, x, θ, adj1s, adj2, ExaModels._constraint_dims(f), nvar, npar, ncon, nout, nitr; ndrange = nb * nitr) + end +end + +@kernel function kerh2_batch(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj1s), @Const(adj2), @Const(dims), @Const(nvar), @Const(npar), @Const(ncon), @Const(nout), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + y_off = (s - 1) * nout + a_off = (s - 1) * ncon + @inbounds ExaModels.hrpass0( + f(itr[k], ExaModels.SecondAdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), + f.comp2, + ExaModels.OffsetVector(y1, y_off), + y2, + ExaModels.offset2(f, k), + 0, + adj1s[a_off + ExaModels.offset0(f, itr, k, dims)], + adj2, + ) +end + end # module ExaModelsKernelAbstractions diff --git a/src/BatchNLPModels.jl b/src/BatchNLPModels.jl index 6db579c83..17e3bbea7 100644 --- a/src/BatchNLPModels.jl +++ b/src/BatchNLPModels.jl @@ -1,8 +1,8 @@ """ BatchNLPModels -Template module for batched NLP models. Defines abstract types, metadata, -and generic API functions following the NLPModels.jl pattern. +Template module for batched NLP models. Defines abstract types and +generic API functions following the NLPModels.jl pattern. Key design: `AbstractBatchNLPModel <: NLPModels.AbstractNLPModel`, so batch models participate in the standard NLPModels dispatch hierarchy. @@ -12,7 +12,9 @@ module BatchNLPModels import NLPModels: NLPModels, AbstractNLPModel, - AbstractNLPModelMeta + AbstractNLPModelMeta, + NLPModelMeta, + obj! # ============================================================================ # Abstract types @@ -25,7 +27,7 @@ Abstract type for batched NLP models. Subtypes `AbstractNLPModel` so that batch models participate in the standard NLPModels dispatch hierarchy. Implementations must provide: -- `meta` field of type `<: AbstractBatchNLPModelMeta` +- `meta` field of type `NLPModelMeta` - `counters` field of type `NLPModels.Counters` - Batch API methods: `obj!`, `grad!`, `cons!`, `jac_structure!`, `jac_coord!`, `hess_structure!`, `hess_coord!` @@ -33,205 +35,22 @@ Implementations must provide: abstract type AbstractBatchNLPModel{T, S} <: AbstractNLPModel{T, S} end # ============================================================================ -# BatchNLPModelMeta -# ============================================================================ - -""" - BatchNLPModelMeta{T, VT, VI} <: AbstractNLPModelMeta{T, VT} - -Metadata for a batched NLP where `nbatch` independent problems share -identical structure (dimensions, sparsity patterns). Extends the standard -`NLPModelMeta` interface with a batch dimension. - -All `VT`-typed arrays are either vectors (nbatch=1) or matrices with -columns indexing instances. - -# Type parameters -- `T`: element type (e.g. `Float64`) -- `VT`: storage type — `Matrix{T}` when `nbatch > 1` (columns = instances), - `Vector{T}` when `nbatch == 1` -- `VI`: integer index vector type (typically `Vector{Int}`) -""" -struct BatchNLPModelMeta{T, VT, VI} <: AbstractNLPModelMeta{T, VT} - nvar::Int - x0::VT - lvar::VT - uvar::VT - ifix::VI - ilow::VI - iupp::VI - irng::VI - ifree::VI - iinf::VI - nlvb::Int - nlvo::Int - nlvc::Int - ncon::Int - y0::VT - lcon::VT - ucon::VT - jfix::VI - jlow::VI - jupp::VI - jrng::VI - jfree::VI - jinf::VI - nnzo::Int - nnzj::Int - lin_nnzj::Int - nln_nnzj::Int - nnzh::Int - nlin::Int - nnln::Int - lin::VI - nln::VI - minimize::Bool - islp::Bool - name::String - variable_bounds_analysis::Bool - constraint_bounds_analysis::Bool - sparse_jacobian::Bool - sparse_hessian::Bool - grad_available::Bool - jac_available::Bool - hess_available::Bool - jprod_available::Bool - jtprod_available::Bool - hprod_available::Bool -end - -# ============================================================================ -# Accessors — auto-generate get_* for all fields +# get_nbatch — derived from the VT type of the meta # ============================================================================ -for field in fieldnames(BatchNLPModelMeta) - meth = Symbol("get_", field) - # Extend NLPModels accessor if it exists, otherwise define locally - if isdefined(NLPModels, meth) - @eval begin - NLPModels.$meth(meta::BatchNLPModelMeta) = getproperty(meta, $(QuoteNode(field))) - NLPModels.$meth(m::AbstractBatchNLPModel) = NLPModels.$meth(m.meta) - end - else - @eval begin - $meth(meta::BatchNLPModelMeta) = getproperty(meta, $(QuoteNode(field))) - $meth(m::AbstractBatchNLPModel) = $meth(m.meta) - export $meth - end - end -end +get_nbatch(meta::NLPModelMeta{T, <:AbstractMatrix}) where {T} = Base.size(meta.x0, 2) +get_nbatch(meta::NLPModelMeta) = 1 +get_nbatch(m::AbstractNLPModel) = get_nbatch(m.meta) -# get_nbatch is derived from the VT type, not stored as a field -get_nbatch(meta::BatchNLPModelMeta{T, <:AbstractMatrix}) where {T} = Base.size(meta.x0, 2) -get_nbatch(meta::BatchNLPModelMeta) = 1 -get_nbatch(m::AbstractBatchNLPModel) = get_nbatch(m.meta) -export get_nbatch - -# ============================================================================ -# Constructor helpers -# ============================================================================ - -""" - _first_instance(v) - -Extract first-instance data for bounds classification. -""" -_first_instance(v::AbstractVector) = v -_first_instance(m::AbstractMatrix) = @view m[:, 1] - -function _classify_bounds(lb, ub, ::Type{T}) where {T} - ifix = findall(lb .== ub) - ilow = findall((lb .> T(-Inf)) .& (ub .== T(Inf))) - iupp = findall((lb .== T(-Inf)) .& (ub .< T(Inf))) - irng = findall((lb .> T(-Inf)) .& (ub .< T(Inf)) .& (lb .< ub)) - ifree = findall((lb .== T(-Inf)) .& (ub .== T(Inf))) - iinf = findall(lb .> ub) - return ifix, ilow, iupp, irng, ifree, iinf -end - -""" - BatchNLPModelMeta(nbatch, nvar, x0, lvar, uvar, ncon, y0, lcon, ucon; kwargs...) - -Construct batch NLP metadata. Bounds analysis (variable/constraint -classification) is performed on the first instance. -""" -function BatchNLPModelMeta( - nvar::Int, - x0::VT, - lvar::VT, - uvar::VT, - ncon::Int, - y0::VT, - lcon::VT, - ucon::VT; - nnzj::Int = nvar * ncon, - nnzh::Int = nvar * (nvar + 1) ÷ 2, - minimize::Bool = true, - islp::Bool = false, - name::String = "Generic", -) where {VT} - T = eltype(VT) - - # Variable bounds analysis (first instance) - ifix, ilow, iupp, irng, ifree, iinf = _classify_bounds( - _first_instance(lvar), _first_instance(uvar), T, - ) - - # Constraint bounds analysis (first instance) - if ncon > 0 - jfix, jlow, jupp, jrng, jfree, jinf = _classify_bounds( - _first_instance(lcon), _first_instance(ucon), T, - ) - else - jfix = jlow = jupp = jrng = jfree = jinf = Int[] - end - - nln = collect(1:ncon) - VI = Vector{Int} - - return BatchNLPModelMeta{T, VT, VI}( - nvar, - x0, lvar, uvar, - ifix, ilow, iupp, irng, ifree, iinf, - nvar, nvar, nvar, # nlvb, nlvo, nlvc - ncon, - y0, lcon, ucon, - jfix, jlow, jupp, jrng, jfree, jinf, - nvar, # nnzo - nnzj, 0, nnzj, # nnzj, lin_nnzj, nln_nnzj - nnzh, - 0, ncon, # nlin, nnln - Int[], nln, # lin, nln - minimize, islp, name, - true, true, # variable/constraint_bounds_analysis - true, true, # sparse_jacobian/hessian - true, # grad_available - ncon > 0, # jac_available - true, # hess_available - ncon > 0, # jprod_available - ncon > 0, # jtprod_available - true, # hprod_available - ) -end # ============================================================================ # Generic batch API — function stubs # ============================================================================ -# -# Concrete implementations must define the `!` (in-place) methods. -# Allocating wrappers are provided as defaults. - -""" - obj!(m::AbstractBatchNLPModel, bx::AbstractMatrix, bf::AbstractVector) - -Evaluate per-instance objectives. `bx` is `(nvar, nbatch)`, `bf` is `(nbatch,)`. -""" -function obj! end """ obj(m::AbstractBatchNLPModel, bx::AbstractMatrix) -> Vector -Allocating version of `obj!`. +Allocating version of `NLPModels.obj!`. """ function NLPModels.obj(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} bf = Vector{T}(undef, get_nbatch(m)) @@ -254,7 +73,7 @@ end Allocating version of `grad!`. """ function NLPModels.grad(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} - bg = Matrix{T}(undef, get_nvar(m), get_nbatch(m)) + bg = Matrix{T}(undef, NLPModels.get_nvar(m), get_nbatch(m)) NLPModels.grad!(m, bx, bg) return bg end @@ -274,7 +93,7 @@ end Allocating version of `cons!`. """ function NLPModels.cons(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} - bc = Matrix{T}(undef, get_ncon(m), get_nbatch(m)) + bc = Matrix{T}(undef, NLPModels.get_ncon(m), get_nbatch(m)) NLPModels.cons!(m, bx, bc) return bc end @@ -282,8 +101,7 @@ end """ jac_structure!(m::AbstractBatchNLPModel, rows, cols) -Per-instance Jacobian sparsity pattern (local indices). `rows` and `cols` -have length `nnzj` (per instance). +Per-instance Jacobian sparsity pattern (local indices). """ function NLPModels.jac_structure!( m::AbstractBatchNLPModel, @@ -296,8 +114,7 @@ end """ jac_coord!(m::AbstractBatchNLPModel, bx::AbstractMatrix, jvals::AbstractVector) -Evaluate batch Jacobian values. `jvals` is a flat vector of length -`nnzj * nbatch`, laid out instance by instance. +Evaluate batch Jacobian values. """ function NLPModels.jac_coord!( m::AbstractBatchNLPModel, @@ -310,8 +127,7 @@ end """ hess_structure!(m::AbstractBatchNLPModel, rows, cols) -Per-instance Hessian sparsity pattern (local indices). `rows` and `cols` -have length `nnzh` (per instance). +Per-instance Hessian sparsity pattern (local indices). """ function NLPModels.hess_structure!( m::AbstractBatchNLPModel, @@ -322,32 +138,157 @@ function NLPModels.hess_structure!( end """ - hess_coord!(m::AbstractBatchNLPModel, bx, by, bobj_weight, hvals) - -Evaluate batch Hessian values. `hvals` is a flat vector of length -`nnzh * nbatch`, laid out instance by instance. + hess_coord!(m::AbstractBatchNLPModel, bx, by, hvals; obj_weight = 1) -- `bx`: `(nvar, nbatch)` primal values -- `by`: `(ncon, nbatch)` dual values -- `bobj_weight`: `(nbatch,)` per-instance objective weights -- `hvals`: `(nnzh * nbatch,)` flat Hessian values +Evaluate batch Hessian values. """ function NLPModels.hess_coord!( - m::AbstractBatchNLPModel, + m::AbstractBatchNLPModel{T}, bx::AbstractMatrix, by::AbstractMatrix, - bobj_weight::AbstractVector, - hvals::AbstractVector, -) + hvals::AbstractVector; + obj_weight = one(T), +) where {T} error("hess_coord! not implemented for $(typeof(m))") end +# ============================================================================ +# FlattenNLPModel +# ============================================================================ + +""" + FlattenNLPModel{T, M} <: AbstractNLPModel{T, Vector{T}} + +Wrapper that presents a batch NLP model as a flat (Vector-based) NLP model. +All NLPModels callbacks delegate to the underlying batch model's matrix API. + + FlattenNLPModel(model::AbstractNLPModel) + +Construct a flat model from a batch model whose `meta.x0` is a matrix. +""" +struct FlattenNLPModel{T, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, Vector{T}} + batch::M + meta::NLPModelMeta{T, Vector{T}} + counters::NLPModels.Counters +end + +function FlattenNLPModel(model::AbstractNLPModel{T}) where {T} + nb = get_nbatch(model) + nvar = NLPModels.get_nvar(model) * nb + ncon = NLPModels.get_ncon(model) * nb + nnzj = NLPModels.get_nnzj(model) * nb + nnzh = NLPModels.get_nnzh(model) * nb + meta = NLPModelMeta{T, Vector{T}}( + nvar, + vec(model.meta.x0), vec(model.meta.lvar), vec(model.meta.uvar), + Int[], Int[], Int[], Int[], collect(1:nvar), Int[], + nvar, nvar, nvar, + ncon, + vec(model.meta.y0), vec(model.meta.lcon), vec(model.meta.ucon), + Int[], Int[], Int[], Int[], Int[], Int[], + nvar, nnzj, 0, nnzj, nnzh, + 0, ncon, Int[], collect(1:ncon), + model.meta.minimize, false, String(model.meta.name), + false, false, true, true, true, ncon > 0, true, ncon > 0, ncon > 0, true, + ) + return FlattenNLPModel(model, meta, NLPModels.Counters()) +end + +function NLPModels.obj(m::FlattenNLPModel{T}, x::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + bx = reshape(x, nvar, nb) + bf = Vector{T}(undef, nb) + obj!(m.batch, bx, bf) + return sum(bf) +end + +function NLPModels.grad!(m::FlattenNLPModel{T}, x::AbstractVector, g::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + NLPModels.grad!(m.batch, reshape(x, nvar, nb), reshape(g, nvar, nb)) + return g +end + +function NLPModels.cons_nln!(m::FlattenNLPModel{T}, x::AbstractVector, c::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + ncon = NLPModels.get_ncon(m.batch) + NLPModels.cons!(m.batch, reshape(x, nvar, nb), reshape(c, ncon, nb)) + return c +end + +function NLPModels.jac_structure!(m::FlattenNLPModel, rows::AbstractVector, cols::AbstractVector) + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + ncon = NLPModels.get_ncon(m.batch) + nnzj = NLPModels.get_nnzj(m.batch) + + # Get per-instance structure + r1 = @view rows[1:nnzj] + c1 = @view cols[1:nnzj] + NLPModels.jac_structure!(m.batch, r1, c1) + + # Replicate for each instance with shifted indices + for s in 2:nb + offset = (s - 1) * nnzj + row_shift = (s - 1) * ncon + col_shift = (s - 1) * nvar + for k in 1:nnzj + @inbounds rows[offset + k] = r1[k] + row_shift + @inbounds cols[offset + k] = c1[k] + col_shift + end + end + return rows, cols +end + +function NLPModels.jac_coord!(m::FlattenNLPModel{T}, x::AbstractVector, jvals::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + nnzj = NLPModels.get_nnzj(m.batch) + NLPModels.jac_coord!(m.batch, reshape(x, nvar, nb), reshape(jvals, nnzj, nb)) + return jvals +end + +function NLPModels.hess_structure!(m::FlattenNLPModel, rows::AbstractVector, cols::AbstractVector) + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + nnzh = NLPModels.get_nnzh(m.batch) + + # Get per-instance structure + r1 = @view rows[1:nnzh] + c1 = @view cols[1:nnzh] + NLPModels.hess_structure!(m.batch, r1, c1) + + # Replicate for each instance with shifted indices + for s in 2:nb + offset = (s - 1) * nnzh + shift = (s - 1) * nvar + for k in 1:nnzh + @inbounds rows[offset + k] = r1[k] + shift + @inbounds cols[offset + k] = c1[k] + shift + end + end + return rows, cols +end + +function NLPModels.hess_coord!( + m::FlattenNLPModel{T}, x::AbstractVector, y::AbstractVector, + hvals::AbstractVector; obj_weight = one(T), +) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + ncon = NLPModels.get_ncon(m.batch) + nnzh = NLPModels.get_nnzh(m.batch) + NLPModels.hess_coord!(m.batch, reshape(x, nvar, nb), reshape(y, ncon, nb), reshape(hvals, nnzh, nb); obj_weight) + return hvals +end + # ============================================================================ # Exports # ============================================================================ export AbstractBatchNLPModel, - BatchNLPModelMeta, - obj! + FlattenNLPModel end # module BatchNLPModels diff --git a/src/ExaModels.jl b/src/ExaModels.jl index 30ea45944..8dda14eee 100644 --- a/src/ExaModels.jl +++ b/src/ExaModels.jl @@ -63,7 +63,7 @@ include("tags.jl") include("two_stage.jl") include("BatchNLPModels.jl") using .BatchNLPModels -include("batch.jl") +get_model(model::BatchExaModel) = BatchNLPModels.FlattenNLPModel(model) export ExaModel, ExaCore, @@ -113,12 +113,11 @@ export ExaModel, get_ucon, set_ucon!, AbstractBatchNLPModel, - BatchNLPModelMeta, BatchExaCore, BatchExaModel, - EachInstance, get_nbatch, get_model, + flatten_model, var_indices, cons_block_indices diff --git a/src/batch.jl b/src/batch.jl deleted file mode 100644 index 9df964c60..000000000 --- a/src/batch.jl +++ /dev/null @@ -1,509 +0,0 @@ -# ============================================================================ -# Batch ExaModel — dispatch on VT <: AbstractMatrix -# ============================================================================ -# -# Following the same pattern as TwoStageExaModel: -# - BatchExaModelTag stored in ExaCore's tag field -# - BatchExaCore = ExaCore{T, <:AbstractMatrix, B, <:BatchExaModelTag} -# - EachInstance() marker for per-instance declarations -# - ExaModel(c) handles both batch and non-batch via dispatch - -import .BatchNLPModels: obj! - -# ============================================================================ -# Tag and marker types -# ============================================================================ - -""" - EachInstance - -Marker type used with [`add_var`](@ref), [`add_par`](@ref), and [`add_con`](@ref) -to indicate that the declaration is replicated for each instance in a batch model. - -## Example -```julia -core = BatchExaCore(3) -c, v = add_var(core, EachInstance(), 10) -c, _ = add_obj(c, v[i, s]^2 for i in 1:10, s in 1:3) -``` -""" -struct EachInstance end - -struct BatchExaModelTag <: AbstractExaModelTag - nbatch::Int -end - -# ============================================================================ -# Type aliases -# ============================================================================ - -""" - BatchExaCore{T,VT,B} - -Type alias for an [`ExaCore`](@ref) whose `tag` is a [`BatchExaModelTag`] -and whose storage arrays are matrices (columns = instances). -""" -const BatchExaCore{T,VT,B} = ExaCore{T,VT,B,<:BatchExaModelTag} - -""" - BatchExaModel{T,VT,E,V,P,O,C,R,M} - -Type alias for an [`ExaModel`](@ref) built from a [`BatchExaCore`](@ref). -""" -const BatchExaModel{T,VT,E,V,P,O,C,R,M} = ExaModel{T,VT,E,V,P,O,C,<:BatchExaModelTag,R,M} - -# ============================================================================ -# get_nbatch -# ============================================================================ - -get_nbatch(m::BatchExaModel) = m.tag.nbatch -get_nbatch(c::BatchExaCore) = c.tag.nbatch -get_nbatch(m::AbstractExaModel) = 1 - -# ============================================================================ -# _meta_dims — per-instance dimensions for batch -# ============================================================================ - -function _meta_dims(c::C) where {T, VT <: AbstractArray{T}, C <: BatchExaCore{T, VT}} - nb = c.tag.nbatch - return (c.nvar ÷ nb, c.ncon ÷ nb, c.nnzj ÷ nb, c.nnzh ÷ nb) -end - -# ============================================================================ -# BatchExaCore constructor -# ============================================================================ - -""" - BatchExaCore(nbatch; T = Float64, backend = nothing, kwargs...) - -Create an [`ExaCore`](@ref) for building batch optimization models with -`nbatch` independent instances. - -Storage arrays (`x0`, `lvar`, `uvar`, `θ`, `y0`, `lcon`, `ucon`) are -`Matrix{T}` with `nbatch` columns. Use [`add_var`](@ref), [`add_par`](@ref), -[`add_obj`](@ref), and [`add_con`](@ref) with [`EachInstance()`](@ref) to -declare per-instance components. - -## Example -```julia -core = BatchExaCore(3) -c, v = add_var(core, EachInstance(), 10; start = 1.0, lvar = 0.0, uvar = 10.0) -nb = get_nbatch(c) -c, _ = add_obj(c, v[i, s]^2 for i in 1:10, s in 1:nb) -model = ExaModel(c) -``` -""" -function BatchExaCore(nbatch::Integer; T::Type{<:AbstractFloat} = Float64, backend = nothing, kwargs...) - x0 = convert_array(zeros(T, 0, nbatch), backend) - return ExaCore( - :Generic, - backend, - (), # var - (), # par - (), # obj - (), # cons - 0, 0, 0, 0, 0, # nvar, npar, ncon, nconaug, nobj - 0, 0, 0, 0, # nnzc, nnzg, nnzj, nnzh - x0, - similar(x0), # θ - similar(x0), # lvar - similar(x0), # uvar - similar(x0), # y0 - similar(x0), # lcon - similar(x0), # ucon - true, # minimize - BatchExaModelTag(nbatch), - (;), # refs - ) -end - -# ============================================================================ -# add_var / add_par / add_con for BatchExaCore -# ============================================================================ - -""" - add_var(core::BatchExaCore, ::EachInstance, dims...; start = 0, lvar = -Inf, uvar = Inf, name = nothing) - -Add per-instance variables to a batch core. Creates `prod(dims) * nbatch` -variables total — one copy of the block per instance. The returned `Variable` -has dimensions `(dims..., nbatch)`. -""" -function add_var( - c::C, - ::EachInstance, - ns...; - name = nothing, - start = zero(T), - lvar = T(-Inf), - uvar = T(Inf), -) where {T, VT <: AbstractArray{T}, C <: BatchExaCore{T, VT}} - nbatch = c.tag.nbatch - len_per = total(ns) # per-instance count - len_total = len_per * nbatch # fused total - nvar = c.nvar + len_total - - # Append per-instance rows to matrix storage - x0 = append!(c.backend, c.x0, start, len_per) - lv = append!(c.backend, c.lvar, lvar, len_per) - uv = append!(c.backend, c.uvar, uvar, len_per) - - v = Variable((ns..., nbatch), len_total, c.nvar, _val_name(name), nothing) - (ExaCore(c; var = (v, c.var...), nvar = nvar, x0 = x0, lvar = lv, uvar = uv, - refs = add_refs(c.refs, name, v)), v) -end - -""" - add_par(core::BatchExaCore, ::EachInstance, value::AbstractArray; name = nothing) - add_par(core::BatchExaCore, ::EachInstance, dims...; name = nothing, start = 0) - -Add per-instance parameters to a batch core. The parameter values are -replicated for each instance. -""" -function add_par( - c::C, - ::EachInstance, - value::AbstractArray; - name = nothing, -) where {T, VT <: AbstractArray{T}, C <: BatchExaCore{T, VT}} - nbatch = c.tag.nbatch - len_per = length(value) - len_total = len_per * nbatch - npar = c.npar + len_total - θ = append!(c.backend, c.θ, value, len_per) - p = Parameter((Base.size(value)..., nbatch), len_total, c.npar, nothing) - (ExaCore(c; par = (p, c.par...), θ = θ, npar = npar, refs = add_refs(c.refs, name, p)), p) -end - -function add_par( - c::C, - ::EachInstance, - ns...; - name = nothing, - start = zero(T), -) where {T, VT <: AbstractArray{T}, C <: BatchExaCore{T, VT}} - nbatch = c.tag.nbatch - len_per = total(ns) - len_total = len_per * nbatch - npar = c.npar + len_total - θ = append!(c.backend, c.θ, start, len_per) - p = Parameter((ns..., nbatch), len_total, c.npar, nothing) - (ExaCore(c; par = (p, c.par...), θ = θ, npar = npar, refs = add_refs(c.refs, name, p)), p) -end - -""" - add_con(core::BatchExaCore, ::EachInstance, gen; start = 0, lcon = 0, ucon = 0, name = nothing) - -Add per-instance constraints to a batch core. -""" -function add_con( - c::C, - ::EachInstance, - ns...; - name = nothing, - tag = nothing, - start = zero(T), - lcon = zero(T), - ucon = zero(T), -) where {T, VT <: AbstractArray{T}, C <: BatchExaCore{T, VT}} - gen = _get_generator(ns) - dims = _get_con_dims(ns) - gen = _adapt_gen(gen) - f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) - pars = gen.iter - - nitr = length(pars) - o = c.ncon - ncon = c.ncon + nitr - nnzj = c.nnzj + nitr * f.o1step - nnzh = c.nnzh + nitr * f.o2step - - # Append per-instance rows to matrix storage - nbatch = c.tag.nbatch - nitr_per = nitr ÷ nbatch - y0 = append!(c.backend, c.y0, start, nitr_per) - lc = append!(c.backend, c.lcon, lcon, nitr_per) - uc = append!(c.backend, c.ucon, ucon, nitr_per) - - con = Constraint(f, convert_array(pars, c.backend), o, dims, nothing) - (ExaCore(c; ncon = ncon, nnzj = nnzj, nnzh = nnzh, y0 = y0, lcon = lc, ucon = uc, - cons = (con, c.cons...), refs = add_refs(c.refs, name, con)), con) -end - -# ============================================================================ -# Batch show -# ============================================================================ - -function Base.show(io::IO, m::BatchExaModel{T, VT}) where {T, VT} - nb = get_nbatch(m) - nv = NLPModels.get_nvar(m) - nc = NLPModels.get_ncon(m) - println(io, "An ExaModel{$T, $VT, ...} (batch)") - println(io, " Instances: $nb") - println(io, " Variables per instance: $nv") - println(io, " Constraints per instance: $nc") - println(io, " Total variables: $(nb * nv)") - return println(io, " Total constraints: $(nb * nc)") -end - -# ============================================================================ -# Accessors -# ============================================================================ - -""" - get_model(model) - -For batch models, returns a flat (Vector-based) ExaModel for direct NLPModels -interface usage (e.g. MadNLP/Ipopt). For regular models, returns self. -""" -get_model(model::ExaModel) = model -function get_model(model::BatchExaModel) - nb = get_nbatch(model) - nvar = NLPModels.get_nvar(model) * nb - ncon = NLPModels.get_ncon(model) * nb - nnzj = NLPModels.get_nnzj(model) * nb - nnzh = NLPModels.get_nnzh(model) * nb - meta = BatchNLPModelMeta( - nvar, vec(model.meta.x0), vec(model.meta.lvar), vec(model.meta.uvar), - ncon, vec(model.meta.y0), vec(model.meta.lcon), vec(model.meta.ucon); - nnzj = nnzj, nnzh = nnzh, minimize = model.meta.minimize, - ) - return ExaModel( - model.name, model.vars, model.pars, model.objs, model.cons, - vec(model.θ), meta, NLPModels.Counters(), nothing, nothing, model.refs, - ) -end - -""" - var_indices(model, i) -> UnitRange - -Variable index range for instance `i` in the fused model's global variable vector. -""" -var_indices(model::BatchExaModel, i::Int) = - ((i - 1) * NLPModels.get_nvar(model) + 1):(i * NLPModels.get_nvar(model)) - -""" - cons_block_indices(model, i) -> UnitRange - -Constraint index range for instance `i` in the fused model's global constraint vector. -""" -cons_block_indices(model::BatchExaModel, i::Int) = - ((i - 1) * NLPModels.get_ncon(model) + 1):(i * NLPModels.get_ncon(model)) - -# ============================================================================ -# Objective buffer evaluation -# ============================================================================ - -_count_nobj(::Tuple{}) = 0 -_count_nobj(objs::Tuple) = length(first(objs).itr) + _count_nobj(Base.tail(objs)) - -function _eval_objbuffer!(objbuffer, objs::Tuple, x, θ) - _eval_objbuffer!(objbuffer, Base.tail(objs), x, θ) - o = first(objs) - for i in eachindex(o.itr) - objbuffer[offset0(o, i)] = o.f(o.itr[i], x, θ) - end - return -end -_eval_objbuffer!(objbuffer, ::Tuple{}, x, θ) = nothing - -# ============================================================================ -# Batch API: obj / obj! -# ============================================================================ - -function obj!(m::BatchExaModel{T}, bx::AbstractMatrix, bf::AbstractVector) where {T} - nb = get_nbatch(m) - nobj_total = _count_nobj(m.objs) - nobj_per = nobj_total ÷ nb - objbuffer = Vector{T}(undef, nobj_total) - _eval_objbuffer!(objbuffer, m.objs, vec(bx), vec(m.θ)) - obj_mat = reshape(objbuffer, nobj_per, nb) - bf .= vec(sum(obj_mat; dims = 1)) - return bf -end - -function obj(m::BatchExaModel{T}, bx::AbstractMatrix) where {T} - bf = Vector{T}(undef, get_nbatch(m)) - obj!(m, bx, bf) - return bf -end - -# ============================================================================ -# Batch API: grad! -# ============================================================================ - -function NLPModels.grad!(m::BatchExaModel{T}, bx::AbstractMatrix, bg::AbstractMatrix) where {T} - fill!(vec(bg), zero(T)) - _grad!(m.objs, vec(bx), vec(m.θ), vec(bg)) - return bg -end - -# ============================================================================ -# Batch API: cons! -# ============================================================================ - -function NLPModels.cons!(m::BatchExaModel{T}, bx::AbstractMatrix, bc::AbstractMatrix) where {T} - fill!(vec(bc), zero(T)) - _cons_nln!(m.cons, vec(bx), vec(m.θ), vec(bc)) - return bc -end - -# ============================================================================ -# Batch API: jac_structure! / jac_coord! -# ============================================================================ - -function NLPModels.jac_structure!( - m::BatchExaModel{T}, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, -) where {T} - nb = get_nbatch(m) - nnzj_per = NLPModels.get_nnzj(m) - total_nnzj = nnzj_per * nb - full_rows = zeros(Int, total_nnzj) - full_cols = zeros(Int, total_nnzj) - _jac_structure!(T, m.cons, full_rows, full_cols) - for k in 1:nnzj_per - rows[k] = full_rows[k] - cols[k] = full_cols[k] - end - return rows, cols -end - -function NLPModels.jac_coord!(m::BatchExaModel{T}, bx::AbstractMatrix, jvals::AbstractVector) where {T} - fill!(jvals, zero(T)) - _jac_coord!(m.cons, vec(bx), vec(m.θ), jvals) - return jvals -end - -# ============================================================================ -# Batch API: hess_structure! / hess_coord! -# ============================================================================ - -# Helpers for hess permutation -_count_hess_nnz(::Tuple{}) = 0 -function _count_hess_nnz(objs::Tuple) - o = first(objs) - return o.f.o2step * length(o.itr) + _count_hess_nnz(Base.tail(objs)) -end - -function _build_hess_perm(ns, nnzh_obj_per, nnzh_con_per) - nnzh_per = nnzh_obj_per + nnzh_con_per - perm = Vector{Int}(undef, ns * nnzh_per) - for s in 1:ns - base = (s - 1) * nnzh_per - for k in 1:nnzh_obj_per - perm[base + k] = (s - 1) * nnzh_obj_per + k - end - for k in 1:nnzh_con_per - perm[base + nnzh_obj_per + k] = - ns * nnzh_obj_per + (s - 1) * nnzh_con_per + k - end - end - return perm -end - -function _batch_hess_perm(m::BatchExaModel) - nb = get_nbatch(m) - nnzh_obj_per = _count_hess_nnz(m.objs) ÷ nb - nnzh_con_per = _count_hess_nnz(m.cons) ÷ nb - return _build_hess_perm(nb, nnzh_obj_per, nnzh_con_per) -end - -function NLPModels.hess_structure!( - m::BatchExaModel{T}, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, -) where {T} - nb = get_nbatch(m) - nnzh_per = NLPModels.get_nnzh(m) - total_nnzh = nnzh_per * nb - full_rows = zeros(Int, total_nnzh) - full_cols = zeros(Int, total_nnzh) - _obj_hess_structure!(T, m.objs, full_rows, full_cols) - _con_hess_structure!(T, m.cons, full_rows, full_cols) - perm = _batch_hess_perm(m) - for k in 1:nnzh_per - idx = perm[k] - rows[k] = full_rows[idx] - cols[k] = full_cols[idx] - end - return rows, cols -end - -function NLPModels.hess_coord!( - m::BatchExaModel{T}, - bx::AbstractMatrix, - by::AbstractMatrix, - bobj_weight::AbstractVector, - hvals::AbstractVector, -) where {T} - x_flat = vec(bx) - y_flat = vec(by) - nb = get_nbatch(m) - nh = NLPModels.get_nnzh(m) - total_nnzh = nh * nb - perm = _batch_hess_perm(m) - hess_buffer = Vector{T}(undef, total_nnzh) - - bobj_weight_cpu = Array(bobj_weight) - if allequal(bobj_weight_cpu) - w = first(bobj_weight_cpu) - fill!(hess_buffer, zero(T)) - _obj_hess_coord!(m.objs, x_flat, m.θ, hess_buffer, w) - _con_hess_coord!(m.cons, x_flat, m.θ, y_flat, hess_buffer, w) - hvals .= hess_buffer[perm] - else - hess_obj = Vector{T}(undef, total_nnzh) - fill!(hess_obj, zero(T)) - _obj_hess_coord!(m.objs, x_flat, m.θ, hess_obj, one(T)) - fill!(hess_buffer, zero(T)) - _con_hess_coord!(m.cons, x_flat, m.θ, y_flat, hess_buffer, zero(T)) - w = [bobj_weight_cpu[(i - 1) ÷ nh + 1] for i in 1:length(perm)] - hvals .= w .* hess_obj[perm] .+ hess_buffer[perm] - end - return hvals -end - -# ============================================================================ -# Error guards: vector-argument NLPModels API on batch models -# ============================================================================ - -_batch_vector_error(name, m) = throw(ArgumentError( - "$name on batch ExaModel requires matrix arguments. " * - "Use the batch API or get_model(m) for the fused model.", -)) - -function obj(m::BatchExaModel, x::AbstractVector) - _batch_vector_error("obj", m) -end - -function cons_nln!(m::BatchExaModel, x::AbstractVector, c::AbstractVector) - _batch_vector_error("cons_nln!", m) -end - -function NLPModels.grad!(m::BatchExaModel, x::AbstractVector, g::AbstractVector) - _batch_vector_error("grad!", m) -end - -function NLPModels.jac_coord!(m::BatchExaModel, x::AbstractVector, jac::AbstractVector) - _batch_vector_error("jac_coord!", m) -end - -function NLPModels.hess_coord!( - m::BatchExaModel, - x::AbstractVector, - y::AbstractVector, - hess::AbstractVector; - obj_weight = one(eltype(x)), -) - _batch_vector_error("hess_coord!", m) -end - -function NLPModels.hess_coord!( - m::BatchExaModel, - x::AbstractVector, - hess::AbstractVector; - obj_weight = one(eltype(x)), -) - _batch_vector_error("hess_coord!", m) -end diff --git a/src/deprecated.jl b/src/deprecated.jl index 2f34ddd17..851dd1952 100644 --- a/src/deprecated.jl +++ b/src/deprecated.jl @@ -26,11 +26,18 @@ mutable struct LegacyExaCore{T, VT <: AbstractArray{T}, B, S} <: AbstractExaCore end # Override the Val{false} dispatch defined in nlp.jl so ExaCore() returns a LegacyExaCore. -@inline function _make_exacore(::Val{false}, ::Type{T}, backend; kwargs...) where {T} +@inline function _make_exacore(::Val{false}, ::Type{T}, backend, ::Val{1}; kwargs...) where {T} @warn "`ExaCore()` is deprecated, and will be removed in v0.11. Use `ExaCore(concrete = Val(true))` for the immutable ExaCore. The default behavior for `ExaCore()` will change to return the immutable ExaCore in v0.11." inner = _exa_core(; x0 = convert_array(zeros(T, 0), backend), backend, kwargs...) return LegacyExaCore{T, typeof(inner.x0), typeof(backend), typeof(inner.tag)}(inner) end +@inline function _make_exacore(::Val{false}, ::Type{T}, backend, ::Val{NB}; kwargs...) where {T, NB} + @warn "`ExaCore()` is deprecated, and will be removed in v0.11. Use `ExaCore(concrete = Val(true))` for the immutable ExaCore. The default behavior for `ExaCore()` will change to return the immutable ExaCore in v0.11." + x0 = convert_array(zeros(T, 0, NB), backend) + inner = _exa_core(; x0, θ = similar(x0), lvar = similar(x0), uvar = similar(x0), + y0 = similar(x0), lcon = similar(x0), ucon = similar(x0), backend, kwargs...) + return LegacyExaCore{T, typeof(inner.x0), typeof(backend), typeof(inner.tag)}(inner) +end # --------------------------------------------------------------------------- # Property forwarding for LegacyExaCore diff --git a/src/gradient.jl b/src/gradient.jl index 3bf806cf4..18c75bc3a 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -42,6 +42,17 @@ function gradient!(y, f, x, θ, adj) end return y end +function gradient!(y, f, x::AbstractArray, θ::AbstractArray, adj, nb::Integer, nvar::Integer, npar::Integer, ::Nothing = nothing) + for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y_s = @view y[(s-1)*nvar+1 : s*nvar] + @simd for k in eachindex(f.itr) + @inbounds gradient!(y_s, f.f, x_s, θ_s, f.itr[k], adj) + end + end + return y +end function gradient!(y, f, x, θ, p, adj) graph = f(p, AdjointNodeSource(x), θ) drpass(graph, y, adj) diff --git a/src/graph.jl b/src/graph.jl index d38fa433a..20e016a7f 100644 --- a/src/graph.jl +++ b/src/graph.jl @@ -1,3 +1,18 @@ +# ── OffsetVector — zero-allocation offset indexing for batch kernels ────────── + +""" + OffsetVector(data, offset) + +Lightweight wrapper: `ov[i]` returns `data[offset + i]`. +Used in batch KA kernels to avoid GPU view allocations. +""" +struct OffsetVector{V} + data::V + offset::Int +end +@inline Base.getindex(ov::OffsetVector, i) = @inbounds ov.data[ov.offset + i] +@inline Base.setindex!(ov::OffsetVector, v, i) = @inbounds ov.data[ov.offset + i] = v + # ── Abstract node types ─────────────────────────────────────────────────────── """ diff --git a/src/hessian.jl b/src/hessian.jl index 64f92a07f..437a98358 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -1,3 +1,6 @@ +@inline _get_obj_weight(w::Number, s) = w +@inline _get_obj_weight(w::AbstractVector, s) = @inbounds w[s] + """ hdrpass(t1::T1, t2::T2, comp, y1, y2, o2, cnt, adj) @@ -694,6 +697,19 @@ function shessian!(y1, y2, f, x, θ, adj1, adj2) ) end end +function shessian!(y1, y2, f, x::AbstractArray, θ::AbstractArray, adj1, adj2, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, ::Nothing = nothing) + for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y1_s = @view y1[(s-1)*nout+1 : s*nout] + w_s = _get_obj_weight(adj1, s) + @simd for k in eachindex(f.itr) + @inbounds shessian!( + y1_s, y2, f.f, f.itr[k], x_s, θ_s, f.f.comp2, offset2(f, k), w_s, adj2, + ) + end + end +end function shessian!(y1, y2, f, x, θ, adj1s::V, adj2) where {V<:AbstractVector} @simd for k in eachindex(f.itr) @inbounds shessian!( @@ -710,6 +726,20 @@ function shessian!(y1, y2, f, x, θ, adj1s::V, adj2) where {V<:AbstractVector} ) end end +function shessian!(y1, y2, f, x::AbstractArray, θ::AbstractArray, adj1s::AbstractVector, adj2, nb::Integer, nvar::Integer, npar::Integer, ncon::Integer, nout::Integer, ::Nothing = nothing) + for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y1_s = @view y1[(s-1)*nout+1 : s*nout] + a_s = @view adj1s[(s-1)*ncon+1 : s*ncon] + @simd for k in eachindex(f.itr) + @inbounds shessian!( + y1_s, y2, f.f, f.itr[k], x_s, θ_s, f.f.comp2, offset2(f, k), + a_s[offset0(f, k)], adj2, + ) + end + end +end function shessian!(y1, y2, f, p, x, θ, comp, o2, adj1, adj2) graph = f(p, SecondAdjointNodeSource(x), θ) diff --git a/src/jacobian.jl b/src/jacobian.jl index bc60abcb2..aa05b5697 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -125,6 +125,27 @@ function sjacobian!(y1, y2, f, x, θ, adj) ) end end +function sjacobian!(y1, y2, f, x::AbstractArray, θ::AbstractArray, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, ::Nothing = nothing) + for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y1_s = @view y1[(s-1)*nout+1 : s*nout] + @simd for i in eachindex(f.itr) + @inbounds sjacobian!( + y1_s, + y2, + f.f, + f.itr[i], + x_s, + θ_s, + f.f.comp1, + offset0(f, i), + offset1(f, i), + adj, + ) + end + end +end function sjacobian!(y1, y2, f, p, x, θ, comp, o0, o1, adj) graph = f(p, AdjointNodeSource(x), θ) diff --git a/src/nlp.jl b/src/nlp.jl index 8f89e8ef2..ae67cdf12 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -389,15 +389,25 @@ end ) end -@inline ExaCore(::Type{T}; backend = nothing, concrete = Val(false), kwargs...) where {T<:AbstractFloat} = - _make_exacore(concrete, T, backend; kwargs...) -@inline ExaCore(; backend = nothing, concrete = Val(false), kwargs...) = ExaCore(default_T(backend); backend, concrete, kwargs...) -@inline _make_exacore(::Val{true}, ::Type{T}, backend; kwargs...) where {T} = +@inline ExaCore(::Type{T}; backend = nothing, concrete = Val(false), nbatch = Val(1), kwargs...) where {T<:AbstractFloat} = + _make_exacore(concrete, T, backend, nbatch; kwargs...) +@inline ExaCore(; backend = nothing, concrete = Val(false), nbatch = Val(1), kwargs...) = ExaCore(default_T(backend); backend, concrete, nbatch, kwargs...) +@inline _make_exacore(::Val{true}, ::Type{T}, backend, ::Val{1}; kwargs...) where {T} = _exa_core(; x0 = convert_array(zeros(T, 0), backend), backend, kwargs...) +@inline function _make_exacore(::Val{true}, ::Type{T}, backend, ::Val{NB}; kwargs...) where {T, NB} + x0 = convert_array(zeros(T, 0, NB), backend) + _exa_core(; x0, θ = similar(x0), lvar = similar(x0), uvar = similar(x0), + y0 = similar(x0), lcon = similar(x0), ucon = similar(x0), backend, kwargs...) +end # Val{false} is overridden in deprecated.jl once LegacyExaCore is defined; # this fallback handles any other Val value by returning a concrete ExaCore. -@inline _make_exacore(::Val, ::Type{T}, backend; kwargs...) where {T} = +@inline _make_exacore(::Val, ::Type{T}, backend, ::Val{1}; kwargs...) where {T} = _exa_core(; x0 = convert_array(zeros(T, 0), backend), backend, kwargs...) +@inline function _make_exacore(::Val, ::Type{T}, backend, ::Val{NB}; kwargs...) where {T, NB} + x0 = convert_array(zeros(T, 0, NB), backend) + _exa_core(; x0, θ = similar(x0), lvar = similar(x0), uvar = similar(x0), + y0 = similar(x0), lcon = similar(x0), ucon = similar(x0), backend, kwargs...) +end @inline ExaCore(c::C; kwargs...) where C <: ExaCore = _exa_core( ; zip(fieldnames(C), ntuple(i -> getfield(c, i), Val(fieldcount(C))))..., @@ -452,9 +462,11 @@ function ExaModel( ) end -function Base.show(io::IO, c::AbstractExaModel{T,VT}) where {T,VT} - println(io, "An ExaModel{$T, $VT, ...}\n") - Base.show(io, c.meta) +function Base.show(io::IO, m::AbstractExaModel{T,VT}) where {T,VT} + nb = get_nbatch(m) + batch_str = nb > 1 ? " (batch, $nb instances)" : "" + println(io, "An ExaModel{$T, $VT, ...}$batch_str\n") + Base.show(io, m.meta) end """ @@ -507,7 +519,7 @@ function ExaModel(c::C; prod = false, kwargs...) where {C<:ExaCore} c.obj, c.cons, c.θ, - BatchNLPModelMeta( + _build_meta( nvar, c.x0, c.lvar, @@ -530,6 +542,55 @@ end _meta_dims(c::ExaCore) = (c.nvar, c.ncon, c.nnzj, c.nnzh) build_extension(c::ExaCore; kwargs...) = nothing +# ============================================================================ +# _build_meta: construct NLPModelMeta supporting both Vector and Matrix VT +# ============================================================================ + +_first_instance(v::AbstractVector) = v +_first_instance(m::AbstractMatrix) = @view m[:, 1] + +function _classify_bounds(lb, ub, ::Type{T}) where {T} + ifix = findall(lb .== ub) + ilow = findall((lb .> T(-Inf)) .& (ub .== T(Inf))) + iupp = findall((lb .== T(-Inf)) .& (ub .< T(Inf))) + irng = findall((lb .> T(-Inf)) .& (ub .< T(Inf)) .& (lb .< ub)) + ifree = findall((lb .== T(-Inf)) .& (ub .== T(Inf))) + iinf = findall(lb .> ub) + return ifix, ilow, iupp, irng, ifree, iinf +end + +function _build_meta( + nvar::Int, x0::VT, lvar::VT, uvar::VT, + ncon::Int, y0::VT, lcon::VT, ucon::VT; + nnzj::Int = nvar * ncon, + nnzh::Int = nvar * (nvar + 1) ÷ 2, + minimize::Bool = true, + islp::Bool = false, + name::String = "Generic", +) where {VT} + T = eltype(VT) + ifix, ilow, iupp, irng, ifree, iinf = _classify_bounds( + _first_instance(lvar), _first_instance(uvar), T) + if ncon > 0 + jfix, jlow, jupp, jrng, jfree, jinf = _classify_bounds( + _first_instance(lcon), _first_instance(ucon), T) + else + jfix = jlow = jupp = jrng = jfree = jinf = Int[] + end + nln = collect(1:ncon) + return NLPModels.NLPModelMeta{T, VT}( + nvar, x0, lvar, uvar, + ifix, ilow, iupp, irng, ifree, iinf, + nvar, nvar, nvar, + ncon, y0, lcon, ucon, + jfix, jlow, jupp, jrng, jfree, jinf, + nvar, nnzj, 0, nnzj, nnzh, + 0, ncon, Int[], nln, + minimize, islp, name, + true, true, true, true, true, ncon > 0, true, ncon > 0, ncon > 0, true, + ) +end + @inline function Base.getindex(v::V, i) where {V<:AbstractVariable} _bound_check(v.size, i) _indexed_var(i, v.offset - _start(v.size[1]) + 1) @@ -621,12 +682,24 @@ function append!(backend, a, b::Number, lb) return cat(a, new_part; dims = 1) end -function append!(backend, a, b::AbstractArray, lb) +function append!(backend, a, b::AbstractVector, lb) lb == 0 && return a col = vec(convert_array(b, backend)) return cat(a, _expand_to_shape(col, _trailing_dims(a)); dims = 1) end +function append!(backend, a, b::AbstractMatrix, lb) + lb == 0 && return a + m = convert_array(b, backend) + trailing = _trailing_dims(a) + if trailing == () + # a is a vector — flatten the matrix to a vector + return cat(a, vec(m); dims = 1) + else + return cat(a, m; dims = 1) + end +end + function append!(backend, a, b::Base.Generator, lb) lb == 0 && return a b = _adapt_gen(b) @@ -693,7 +766,6 @@ Variable end @inline function _add_var(c, tag, name, start, lvar, uvar, ns...) - o = c.nvar len = total(ns) nvar = c.nvar + len @@ -701,13 +773,17 @@ end lvar = append!(c.backend, c.lvar, lvar, len) uvar = append!(c.backend, c.uvar, uvar, len) - v = Variable(ns, len, o, _val_name(name), tag) + v = Variable(ns, len, c.nvar, _val_name(name), tag) (ExaCore(c; var = (v, c.var...), nvar=nvar, x0=x0, lvar=lvar, uvar=uvar, refs = add_refs(c.refs, name, v)), v) end @inline _val_name(::Val{N}) where {N} = N @inline _val_name(::Nothing) = :x +@inline get_nbatch(c::ExaCore{T, <:AbstractMatrix}) where {T} = Base.size(c.x0, 2) +@inline get_nbatch(::ExaCore) = 1 +@inline get_nbatch(m::ExaModel{T, <:AbstractMatrix}) where {T} = Base.size(m.meta.x0, 2) +@inline get_nbatch(::AbstractExaModel) = 1 @inline add_refs(refs, ::Nothing, var) = refs @inline add_refs(refs, ::Val{N}, var) where {N} = (; refs..., N => var) @@ -754,11 +830,10 @@ end end @inline function _add_par(c, tag, name, start, ns...) - o = c.npar len = total(ns) npar = c.npar + len θ = append!(c.backend, c.θ, start, len) - p = Parameter(ns, len, o, tag) + p = Parameter(ns, len, c.npar, tag) (ExaCore(c; par = (p, c.par...), θ=θ, npar=npar, refs = add_refs(c.refs, name, p)), p) end @@ -825,53 +900,17 @@ end @inline _var_range(v::Variable) = v.offset+1 : v.offset+v.length @inline _con_range(c::Constraint) = c.offset+1 : c.offset+total(c.size) -""" - get_start(model, var::Variable) - get_start(model, con::Constraint) - -Return a view of the initial-point values (`x0` for variables, `y0` for constraints). -""" get_start(model::ExaModel, v::Variable) = view(model.meta.x0, _var_range(v)) get_start(model::ExaModel, c::Constraint) = view(model.meta.y0, _con_range(c)) - -""" - get_lvar(model, var::Variable) - -Return a view of the lower bounds for `var`. -""" get_lvar(model::ExaModel, v::Variable) = view(model.meta.lvar, _var_range(v)) - -""" - get_uvar(model, var::Variable) - -Return a view of the upper bounds for `var`. -""" get_uvar(model::ExaModel, v::Variable) = view(model.meta.uvar, _var_range(v)) - -""" - get_lcon(model, con::Constraint) - -Return a view of the lower bounds for `con`. -""" get_lcon(model::ExaModel, c::Constraint) = view(model.meta.lcon, _con_range(c)) - -""" - get_ucon(model, con::Constraint) - -Return a view of the upper bounds for `con`. -""" get_ucon(model::ExaModel, c::Constraint) = view(model.meta.ucon, _con_range(c)) @inline function _check_len(got, expected, label) got == expected || throw(DimensionMismatch("$label: expected $expected elements, got $got")) end -""" - set_start!(model, var::Variable, values) - set_start!(model, con::Constraint, values) - -Update the initial-point values in-place (`x0` for variables, `y0` for constraints). -""" function set_start!(model::ExaModel, v::Variable, values) _check_len(length(values), v.length, "set_start!") copyto!(view(model.meta.x0, _var_range(v)), values) @@ -881,43 +920,19 @@ function set_start!(model::ExaModel, c::Constraint, values) _check_len(length(values), n, "set_start!") copyto!(view(model.meta.y0, _con_range(c)), values) end - -""" - set_lvar!(model, var::Variable, values) - -Update the lower bounds for `var` in-place. -""" function set_lvar!(model::ExaModel, v::Variable, values) _check_len(length(values), v.length, "set_lvar!") copyto!(view(model.meta.lvar, _var_range(v)), values) end - -""" - set_uvar!(model, var::Variable, values) - -Update the upper bounds for `var` in-place. -""" function set_uvar!(model::ExaModel, v::Variable, values) _check_len(length(values), v.length, "set_uvar!") copyto!(view(model.meta.uvar, _var_range(v)), values) end - -""" - set_lcon!(model, con::Constraint, values) - -Update the lower bounds for `con` in-place. -""" function set_lcon!(model::ExaModel, c::Constraint, values) n = total(c.size) _check_len(length(values), n, "set_lcon!") copyto!(view(model.meta.lcon, _con_range(c)), values) end - -""" - set_ucon!(model, con::Constraint, values) - -Update the upper bounds for `con` in-place. -""" function set_ucon!(model::ExaModel, c::Constraint, values) n = total(c.size) _check_len(length(values), n, "set_ucon!") @@ -1004,7 +1019,7 @@ is intended for code that builds expression trees programmatically. @inline function add_obj(c::C, expr::N, pars = 1:1; name = nothing) where {T,C<:ExaCore{T},N<:AbstractNode} f = _simdfunction(T, expr, c.nobj, c.nnzg, c.nnzh) - _add_obj(c, f, pars, name) + _add_obj(c, f, pars, name) end @inline function _add_obj(c, f, pars, name = nothing) @@ -1084,20 +1099,19 @@ Constraint @inline function add_con( c::C, ns...; - tag = nothing, + tag = nothing, name = nothing, start = zero(T), lcon = zero(T), ucon = zero(T), kwargs... ) where {T,C<:ExaCore{T}} - gen = _get_generator(ns) dims = _get_con_dims(ns) gen = _adapt_gen(gen) f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) pars = gen.iter - + _add_con(c, f, pars, dims, start, lcon, ucon, name, tag) end @@ -1343,83 +1357,140 @@ _con_hess_structure!(T, cons::Tuple{}, rows, cols) = nothing shessian!(rows, cols, first(cons), NaNSource{T}(), NaNSource{T}(), T(NaN), T(NaN)) end +# ============================================================================ +# Batch-aware evaluation — all low-level functions loop over 1:nb, +# striding x/θ/g/etc. by per-instance sizes. For nb=1, the loop runs +# once with views of the full arrays (no overhead). +# ============================================================================ + function obj(m::AbstractExaModel, x::AbstractVector) - return _obj(m.objs, x, m.θ) + return _obj(m.objs, x, m.θ, 1, length(x), length(m.θ)) end -@inline function _obj((obj, objs...), x, θ) - s = _obj(objs, x, θ) - for i in obj.itr - s += obj.f(i, x, θ) +# Stub for KA extension override +function _eval_objbuffer! end + +@inline function _obj((obj, objs...), x, θ, nb, nvar, npar) + s = _obj(objs, x, θ, nb, nvar, npar) + for si in 1:nb + x_s = @view x[(si-1)*nvar+1 : si*nvar] + θ_s = @view θ[(si-1)*npar+1 : si*npar] + for i in obj.itr + s += obj.f(i, x_s, θ_s) + end end return s end -@inline _obj(obj::Tuple{}, x, θ) = zero(eltype(x)) +@inline _obj(obj::Tuple{}, x, θ, nb, nvar, npar) = zero(eltype(x)) + +# Per-instance obj values (for batch obj!) +@inline function _obj!(bf, (obj, objs...), x, θ, nb, nvar, npar, backend = nothing) + _obj!(bf, objs, x, θ, nb, nvar, npar, backend) + _obj_batch!(bf, obj, x, θ, nb, nvar, npar, backend) +end +@inline _obj!(bf, ::Tuple{}, x, θ, nb, nvar, npar, backend = nothing) = nothing + +@inline function _obj_batch!(bf, obj, x, θ, nb, nvar, npar, ::Nothing) + for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + for i in obj.itr + @inbounds bf[s] += obj.f(i, x_s, θ_s) + end + end +end function cons_nln!(m::AbstractExaModel, x::AbstractVector, g::AbstractVector) fill!(g, zero(eltype(g))) - _cons_nln!(m.cons, x, m.θ, g) + nvar = NLPModels.get_nvar(m) + ncon = NLPModels.get_ncon(m) + _cons_nln!(m.cons, x, m.θ, g, 1, nvar, length(m.θ), ncon) return g end -@inline function _cons_nln!(cons::Tuple, x, θ, g) +@inline function _cons_nln!(cons::Tuple, x, θ, g, nb, nvar, npar, ncon, backend = nothing) con = first(cons) - _cons_nln!(Base.tail(cons), x, θ, g) - @simd for i in eachindex(con.itr) - g[offset0(con, i)] += con.f(con.itr[i], x, θ) + _cons_nln!(Base.tail(cons), x, θ, g, nb, nvar, npar, ncon, backend) + _cons_nln_batch!(con, x, θ, g, nb, nvar, npar, ncon, backend) +end +_cons_nln!(cons::Tuple{}, x, θ, g, nb, nvar, npar, ncon, backend = nothing) = nothing + +@inline function _cons_nln_batch!(con, x, θ, g, nb, nvar, npar, ncon, ::Nothing) + for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + g_s = @view g[(s-1)*ncon+1 : s*ncon] + @simd for i in eachindex(con.itr) + g_s[offset0(con, i)] += con.f(con.itr[i], x_s, θ_s) + end end end -_cons_nln!(cons::Tuple{}, x, θ, g) = nothing function grad!(m::AbstractExaModel, x::AbstractVector, f::AbstractVector) fill!(f, zero(eltype(f))) - _grad!(m.objs, x, m.θ, f) + _grad!(m.objs, x, m.θ, f, 1, length(x), length(m.θ)) return f end -@inline function _grad!(objs::Tuple, x, θ, f) - _grad!(Base.tail(objs), x, θ, f) - gradient!(f, first(objs), x, θ, one(eltype(f))) +@inline function _grad!(objs::Tuple, x, θ, f, nb, nvar, npar, backend = nothing) + _grad!(Base.tail(objs), x, θ, f, nb, nvar, npar, backend) + gradient!(f, first(objs), x, θ, one(eltype(f)), nb, nvar, npar, backend) end -_grad!(objs::Tuple{}, x, θ, f) = nothing +_grad!(objs::Tuple{}, x, θ, f, nb, nvar, npar, backend = nothing) = nothing function jac_coord!(m::AbstractExaModel, x::AbstractVector, jac::AbstractVector) fill!(jac, zero(eltype(jac))) - _jac_coord!(m.cons, x, m.θ, jac) + nvar = NLPModels.get_nvar(m) + nnzj = NLPModels.get_nnzj(m) + _jac_coord!(m.cons, x, m.θ, jac, 1, nvar, length(m.θ), nnzj) return jac end -_jac_coord!(cons::Tuple{}, x, θ, jac) = nothing -@inline function _jac_coord!(cons::Tuple, x, θ, jac) - _jac_coord!(Base.tail(cons), x, θ, jac) - sjacobian!(jac, nothing, first(cons), x, θ, one(eltype(jac))) +_jac_coord!(cons::Tuple{}, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) = nothing +@inline function _jac_coord!(cons::Tuple, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) + _jac_coord!(Base.tail(cons), x, θ, jac, nb, nvar, npar, nnzj, backend) + sjacobian!(jac, nothing, first(cons), x, θ, one(eltype(jac)), nb, nvar, npar, nnzj, backend) end function jprod_nln!(m::AbstractExaModel, x::AbstractVector, v::AbstractVector, Jv::AbstractVector) fill!(Jv, zero(eltype(Jv))) - _jprod_nln!(m.cons, x, m.θ, v, Jv) + _jprod_nln!(m.cons, x, m.θ, v, Jv, 1, length(x), length(m.θ), NLPModels.get_ncon(m)) return Jv end -_jprod_nln!(cons::Tuple{}, x, θ, v, Jv) = nothing -@inline function _jprod_nln!(cons::Tuple, x, θ, v, Jv) - _jprod_nln!(Base.tail(cons), x, θ, v, Jv) - sjacobian!((Jv, v), nothing, first(cons), x, θ, one(eltype(Jv))) +_jprod_nln!(cons::Tuple{}, x, θ, v, Jv, nb, nvar, npar, ncon) = nothing +@inline function _jprod_nln!(cons::Tuple, x, θ, v, Jv, nb, nvar, npar, ncon) + _jprod_nln!(Base.tail(cons), x, θ, v, Jv, nb, nvar, npar, ncon) + f = first(cons) + for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + v_s = @view v[(s-1)*nvar+1 : s*nvar] + Jv_s = @view Jv[(s-1)*ncon+1 : s*ncon] + sjacobian!((Jv_s, v_s), nothing, f, x_s, θ_s, one(eltype(Jv))) + end end function jtprod_nln!(m::AbstractExaModel, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector) fill!(Jtv, zero(eltype(Jtv))) - _jtprod_nln!(m.cons, x, m.θ, v, Jtv) + _jtprod_nln!(m.cons, x, m.θ, v, Jtv, 1, length(x), length(m.θ), NLPModels.get_ncon(m)) return Jtv end -_jtprod_nln!(cons::Tuple{}, x, θ, v, Jtv) = nothing -@inline function _jtprod_nln!(cons::Tuple, x, θ, v, Jtv) - _jtprod_nln!(Base.tail(cons), x, θ, v, Jtv) - sjacobian!(nothing, (Jtv, v), first(cons), x, θ, one(eltype(Jtv))) +_jtprod_nln!(cons::Tuple{}, x, θ, v, Jtv, nb, nvar, npar, ncon) = nothing +@inline function _jtprod_nln!(cons::Tuple, x, θ, v, Jtv, nb, nvar, npar, ncon) + _jtprod_nln!(Base.tail(cons), x, θ, v, Jtv, nb, nvar, npar, ncon) + f = first(cons) + for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + v_s = @view v[(s-1)*ncon+1 : s*ncon] + Jtv_s = @view Jtv[(s-1)*nvar+1 : s*nvar] + sjacobian!(nothing, (Jtv_s, v_s), f, x_s, θ_s, one(eltype(Jtv))) + end end function hess_coord!( @@ -1429,7 +1500,9 @@ function hess_coord!( obj_weight = one(eltype(x)), ) fill!(hess, zero(eltype(hess))) - _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight) + nvar = NLPModels.get_nvar(m) + nnzh = NLPModels.get_nnzh(m) + _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight, 1, nvar, length(m.θ), nnzh) return hess end @@ -1441,21 +1514,24 @@ function hess_coord!( obj_weight = one(eltype(x)), ) fill!(hess, zero(eltype(hess))) - _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight) - _con_hess_coord!(m.cons, x, m.θ, y, hess, obj_weight) + nvar = NLPModels.get_nvar(m) + ncon = NLPModels.get_ncon(m) + nnzh = NLPModels.get_nnzh(m) + _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight, 1, nvar, length(m.θ), nnzh) + _con_hess_coord!(m.cons, x, m.θ, y, hess, 1, nvar, length(m.θ), ncon, nnzh) return hess end -_obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight) = nothing -@inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight) - _obj_hess_coord!(Base.tail(objs), x, θ, hess, obj_weight) - shessian!(hess, nothing, first(objs), x, θ, obj_weight, zero(eltype(hess))) +_obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) = nothing +@inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) + _obj_hess_coord!(Base.tail(objs), x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend) + shessian!(hess, nothing, first(objs), x, θ, obj_weight, zero(eltype(hess)), nb, nvar, npar, nnzh, backend) end -_con_hess_coord!(cons::Tuple{}, x, θ, y, hess, obj_weight) = nothing -@inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess, obj_weight) - _con_hess_coord!(Base.tail(cons), x, θ, y, hess, obj_weight) - shessian!(hess, nothing, first(cons), x, θ, y, zero(eltype(hess))) +_con_hess_coord!(cons::Tuple{}, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) = nothing +@inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) + _con_hess_coord!(Base.tail(cons), x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend) + shessian!(hess, nothing, first(cons), x, θ, y, zero(eltype(hess)), nb, nvar, npar, ncon, nnzh, backend) end function hprod!( @@ -1466,7 +1542,7 @@ function hprod!( obj_weight = one(eltype(x)), ) fill!(Hv, zero(eltype(Hv))) - _obj_hprod!(m.objs, x, m.θ, v, Hv, obj_weight) + _obj_hprod!(m.objs, x, m.θ, v, Hv, obj_weight, 1, length(x), length(m.θ)) return Hv end @@ -1479,21 +1555,36 @@ function hprod!( obj_weight = one(eltype(x)), ) fill!(Hv, zero(eltype(Hv))) - _obj_hprod!(m.objs, x, m.θ, v, Hv, obj_weight) - _con_hprod!(m.cons, x, m.θ, y, v, Hv, obj_weight) + _obj_hprod!(m.objs, x, m.θ, v, Hv, obj_weight, 1, length(x), length(m.θ)) + _con_hprod!(m.cons, x, m.θ, y, v, Hv, obj_weight, 1, length(x), length(m.θ), NLPModels.get_ncon(m)) return Hv end -_obj_hprod!(objs::Tuple{}, x, θ, v, Hv, obj_weight) = nothing -@inline function _obj_hprod!(objs::Tuple, x, θ, v, Hv, obj_weight) - _obj_hprod!(Base.tail(objs), x, θ, v, Hv, obj_weight) - shessian!((Hv, v), nothing, first(objs), x, θ, obj_weight, zero(eltype(Hv))) +_obj_hprod!(objs::Tuple{}, x, θ, v, Hv, obj_weight, nb, nvar, npar) = nothing +@inline function _obj_hprod!(objs::Tuple, x, θ, v, Hv, obj_weight, nb, nvar, npar) + _obj_hprod!(Base.tail(objs), x, θ, v, Hv, obj_weight, nb, nvar, npar) + f = first(objs) + for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + v_s = @view v[(s-1)*nvar+1 : s*nvar] + Hv_s = @view Hv[(s-1)*nvar+1 : s*nvar] + shessian!((Hv_s, v_s), nothing, f, x_s, θ_s, obj_weight, zero(eltype(Hv))) + end end -_con_hprod!(cons::Tuple{}, x, θ, y, v, Hv, obj_weight) = nothing -@inline function _con_hprod!(cons::Tuple, x, θ, y, v, Hv, obj_weight) - _con_hprod!(Base.tail(cons), x, θ, y, v, Hv, obj_weight) - shessian!((Hv, v), nothing, first(cons), x, θ, y, zero(eltype(Hv))) +_con_hprod!(cons::Tuple{}, x, θ, y, v, Hv, obj_weight, nb, nvar, npar, ncon) = nothing +@inline function _con_hprod!(cons::Tuple, x, θ, y, v, Hv, obj_weight, nb, nvar, npar, ncon) + _con_hprod!(Base.tail(cons), x, θ, y, v, Hv, obj_weight, nb, nvar, npar, ncon) + f = first(cons) + for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y_s = @view y[(s-1)*ncon+1 : s*ncon] + v_s = @view v[(s-1)*nvar+1 : s*nvar] + Hv_s = @view Hv[(s-1)*nvar+1 : s*nvar] + shessian!((Hv_s, v_s), nothing, f, x_s, θ_s, y_s, zero(eltype(Hv))) + end end @inbounds @inline offset0(a, i) = offset0(a.f, i) @@ -1684,6 +1775,277 @@ end _adapt_gen(gen) = Base.Generator(gen.f, collect(gen.iter)) _adapt_gen(gen::Base.Generator{P}) where {P<:Union{AbstractArray,AbstractRange}} = gen +# ============================================================================ +# Batch ExaModel — dispatch on VT <: AbstractMatrix +# ============================================================================ +# +# BatchExaCore / BatchExaModel are determined by VT <: AbstractMatrix. +# nbatch is derived from size(x0, 2). +# +# All evaluation is handled by the batch-aware low-level functions above +# (_obj, _grad!, _cons_nln!, etc.) which loop over 1:nb with strided views. +# This section provides type aliases, the BatchExaCore constructor, +# and thin batch API wrappers (matrix-argument dispatch). + +# ============================================================================ +# Type aliases — determined by VT <: AbstractMatrix +# ============================================================================ + +""" + BatchExaCore{T,VT,B} + +Type alias for an [`ExaCore`](@ref) whose storage arrays are matrices +(columns = instances). +""" +const BatchExaCore{T,VT<:AbstractMatrix{T},B} = ExaCore{T,VT,B} + +""" + BatchExaModel{T,VT,E,V,P,O,C,S,R,M} + +Type alias for an [`ExaModel`](@ref) built from a [`BatchExaCore`](@ref). +""" +const BatchExaModel{T,VT<:AbstractMatrix{T},E,V,P,O,C,S,R,M} = ExaModel{T,VT,E,V,P,O,C,S,R,M} + +# ============================================================================ +# BatchExaCore constructor — alias for ExaCore with nbatch +# ============================================================================ + +""" + BatchExaCore(nbatch; kwargs...) + +Alias for `ExaCore(; concrete = Val(true), nbatch = Val(nbatch), kwargs...)`. + +Creates an [`ExaCore`](@ref) for building batch optimization models with +`nbatch` independent instances. Generators should iterate over per-instance +dimensions only — the batch dimension is handled automatically at evaluation +time by striding through the data. + +## Example +```julia +core = BatchExaCore(3) +c, v = add_var(core, 10; start = 1.0, lvar = 0.0, uvar = 10.0) +c, _ = add_obj(c, v[i]^2 for i in 1:10) +model = ExaModel(c) +``` +""" +BatchExaCore(nbatch::Integer; kwargs...) = ExaCore(; concrete = Val(true), nbatch = Val(nbatch), kwargs...) + +# ============================================================================ +# get_model — defined after BatchNLPModels is loaded (see ExaModels.jl) +# ============================================================================ + +get_model(model::ExaModel) = model + +""" + var_indices(model, i) -> UnitRange + +Variable index range for instance `i` in the fused model's global variable vector. +""" +var_indices(model::BatchExaModel, i::Int) = + ((i - 1) * NLPModels.get_nvar(model) + 1):(i * NLPModels.get_nvar(model)) + +""" + cons_block_indices(model, i) -> UnitRange + +Constraint index range for instance `i` in the fused model's global constraint vector. +""" +cons_block_indices(model::BatchExaModel, i::Int) = + ((i - 1) * NLPModels.get_ncon(model) + 1):(i * NLPModels.get_ncon(model)) + +# ============================================================================ +# Batch getters / setters +# ============================================================================ + +get_start(model::BatchExaModel, v::Variable) = view(model.meta.x0, _var_range(v), :) +get_start(model::BatchExaModel, c::Constraint) = view(model.meta.y0, _con_range(c), :) +get_start(model::BatchExaModel, v::Variable, i::Int) = view(model.meta.x0, _var_range(v), i) +get_start(model::BatchExaModel, c::Constraint, i::Int) = view(model.meta.y0, _con_range(c), i) +get_lvar(model::BatchExaModel, v::Variable) = view(model.meta.lvar, _var_range(v), :) +get_lvar(model::BatchExaModel, v::Variable, i::Int) = view(model.meta.lvar, _var_range(v), i) +get_uvar(model::BatchExaModel, v::Variable) = view(model.meta.uvar, _var_range(v), :) +get_uvar(model::BatchExaModel, v::Variable, i::Int) = view(model.meta.uvar, _var_range(v), i) +get_lcon(model::BatchExaModel, c::Constraint) = view(model.meta.lcon, _con_range(c), :) +get_lcon(model::BatchExaModel, c::Constraint, i::Int) = view(model.meta.lcon, _con_range(c), i) +get_ucon(model::BatchExaModel, c::Constraint) = view(model.meta.ucon, _con_range(c), :) +get_ucon(model::BatchExaModel, c::Constraint, i::Int) = view(model.meta.ucon, _con_range(c), i) + +function set_start!(model::BatchExaModel, v::Variable, values) + copyto!(view(model.meta.x0, _var_range(v), :), values) +end +function set_start!(model::BatchExaModel, c::Constraint, values) + copyto!(view(model.meta.y0, _con_range(c), :), values) +end +function set_start!(model::BatchExaModel, v::Variable, values, i::Int) + _check_len(length(values), v.length, "set_start!") + copyto!(view(model.meta.x0, _var_range(v), i), values) +end +function set_start!(model::BatchExaModel, c::Constraint, values, i::Int) + n = total(c.size) + _check_len(length(values), n, "set_start!") + copyto!(view(model.meta.y0, _con_range(c), i), values) +end +function set_lvar!(model::BatchExaModel, v::Variable, values) + copyto!(view(model.meta.lvar, _var_range(v), :), values) +end +function set_lvar!(model::BatchExaModel, v::Variable, values, i::Int) + _check_len(length(values), v.length, "set_lvar!") + copyto!(view(model.meta.lvar, _var_range(v), i), values) +end +function set_uvar!(model::BatchExaModel, v::Variable, values) + copyto!(view(model.meta.uvar, _var_range(v), :), values) +end +function set_uvar!(model::BatchExaModel, v::Variable, values, i::Int) + _check_len(length(values), v.length, "set_uvar!") + copyto!(view(model.meta.uvar, _var_range(v), i), values) +end +function set_lcon!(model::BatchExaModel, c::Constraint, values) + copyto!(view(model.meta.lcon, _con_range(c), :), values) +end +function set_lcon!(model::BatchExaModel, c::Constraint, values, i::Int) + n = total(c.size) + _check_len(length(values), n, "set_lcon!") + copyto!(view(model.meta.lcon, _con_range(c), i), values) +end +function set_ucon!(model::BatchExaModel, c::Constraint, values) + copyto!(view(model.meta.ucon, _con_range(c), :), values) +end +function set_ucon!(model::BatchExaModel, c::Constraint, values, i::Int) + n = total(c.size) + _check_len(length(values), n, "set_ucon!") + copyto!(view(model.meta.ucon, _con_range(c), i), values) +end + +# ============================================================================ +# Batch API: matrix-argument wrappers +# These delegate to the unified batch-aware functions above. +# ============================================================================ + +function obj!(m::BatchExaModel{T}, bx::AbstractMatrix, bf::AbstractVector) where {T} + fill!(bf, zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + _obj!(bf, m.objs, vec(bx), vec(m.θ), nb, nvar, npar, getbackend(m)) + return bf +end + +function obj(m::BatchExaModel{T}, bx::AbstractMatrix) where {T} + bf = Vector{T}(undef, get_nbatch(m)) + obj!(m, bx, bf) + return bf +end + +function NLPModels.grad!(m::BatchExaModel{T}, bx::AbstractMatrix, bg::AbstractMatrix) where {T} + fill!(vec(bg), zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + _grad!(m.objs, vec(bx), vec(m.θ), vec(bg), nb, nvar, npar, getbackend(m)) + return bg +end + +function NLPModels.cons!(m::BatchExaModel{T}, bx::AbstractMatrix, bc::AbstractMatrix) where {T} + fill!(vec(bc), zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + ncon = NLPModels.get_ncon(m) + npar = Base.size(m.θ, 1) + _cons_nln!(m.cons, vec(bx), vec(m.θ), vec(bc), nb, nvar, npar, ncon, getbackend(m)) + return bc +end + +function NLPModels.jac_structure!( + m::BatchExaModel{T}, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) where {T} + _jac_structure!(T, m.cons, rows, cols) + return rows, cols +end + +function NLPModels.jac_coord!(m::BatchExaModel{T}, bx::AbstractMatrix, jvals::AbstractMatrix) where {T} + fill!(jvals, zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + nnzj = NLPModels.get_nnzj(m) + _jac_coord!(m.cons, vec(bx), vec(m.θ), vec(jvals), nb, nvar, npar, nnzj, getbackend(m)) + return jvals +end + +function NLPModels.hess_structure!( + m::BatchExaModel{T}, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) where {T} + _obj_hess_structure!(T, m.objs, rows, cols) + _con_hess_structure!(T, m.cons, rows, cols) + return rows, cols +end + +function NLPModels.hess_coord!( + m::BatchExaModel{T}, + bx::AbstractMatrix, + by::AbstractMatrix, + hvals::AbstractMatrix; + obj_weight = one(T), +) where {T} + fill!(hvals, zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + ncon = NLPModels.get_ncon(m) + npar = Base.size(m.θ, 1) + nnzh = NLPModels.get_nnzh(m) + backend = getbackend(m) + _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), obj_weight, nb, nvar, npar, nnzh, backend) + _con_hess_coord!(m.cons, vec(bx), vec(m.θ), vec(by), vec(hvals), nb, nvar, npar, ncon, nnzh, backend) + return hvals +end + +# ============================================================================ +# Error guards: vector-argument NLPModels API on batch models +# ============================================================================ + +_batch_vector_error(name, m) = throw(ArgumentError( + "$name on batch ExaModel requires matrix arguments. " * + "Use the batch API or get_model(m) for the fused model.", +)) + +function obj(m::BatchExaModel, x::AbstractVector) + _batch_vector_error("obj", m) +end + +function cons_nln!(m::BatchExaModel, x::AbstractVector, c::AbstractVector) + _batch_vector_error("cons_nln!", m) +end + +function NLPModels.grad!(m::BatchExaModel, x::AbstractVector, g::AbstractVector) + _batch_vector_error("grad!", m) +end + +function NLPModels.jac_coord!(m::BatchExaModel, x::AbstractVector, jac::AbstractVector) + _batch_vector_error("jac_coord!", m) +end + +function NLPModels.hess_coord!( + m::BatchExaModel, + x::AbstractVector, + y::AbstractVector, + hess::AbstractVector; + obj_weight = one(eltype(x)), +) + _batch_vector_error("hess_coord!", m) +end + +function NLPModels.hess_coord!( + m::BatchExaModel, + x::AbstractVector, + hess::AbstractVector; + obj_weight = one(eltype(x)), +) + _batch_vector_error("hess_coord!", m) +end + function Base.getproperty(core::E, name::Symbol) where {E <: Union{ExaCore, ExaModel}} if hasfield(E, name) getfield(core,name) diff --git a/src/two_stage.jl b/src/two_stage.jl index 03bed951e..93d48f4c3 100644 --- a/src/two_stage.jl +++ b/src/two_stage.jl @@ -104,12 +104,13 @@ function add_var( lvar = T(-Inf), uvar = T(Inf), ) where {T,VT<:AbstractVector{T},B} - + len = total(ns) - append!(c.backend, c.tag.var_scen, 0, len) + new_var_scen = append!(c.backend, c.tag.var_scen, 0, len) + c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, new_var_scen, c.tag.con_scen)) return _add_var( c, FirstStageTag(), name, start, lvar, uvar, ns... - ) + ) end """ add_var(core::TwoStageExaCore, ::EachScenario, dims...; start = 0, lvar = -Inf, uvar = Inf, name = nothing) @@ -130,10 +131,11 @@ function add_var( ) where {T,VT<:AbstractVector{T},B} nscen = c.tag.nscen len = total(ns) - append!(c.backend, c.tag.var_scen, _scen_each_tag(nscen, len), len * nscen) + new_var_scen = append!(c.backend, c.tag.var_scen, _scen_each_tag(nscen, len), len * nscen) + c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, new_var_scen, c.tag.con_scen)) return _add_var( c, SecondStageTag(), name, start, lvar, uvar, ns..., nscen - ) + ) end """ @@ -211,7 +213,8 @@ function add_con( f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) pars = gen.iter - append!(c.backend, c.tag.con_scen, 0, length(pars)) + new_con_scen = append!(c.backend, c.tag.con_scen, 0, length(pars)) + c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, c.tag.var_scen, new_con_scen)) return _add_con(c, f, pars, dims, start, lcon, ucon, name, FirstStageConstraintTag()) end @@ -242,7 +245,8 @@ function add_con( nscen = c.tag.nscen len = length(pars) - append!(c.backend, c.tag.con_scen, _scen_each_tag(nscen, div(len, nscen)), len) + new_con_scen = append!(c.backend, c.tag.con_scen, _scen_each_tag(nscen, div(len, nscen)), len) + c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, c.tag.var_scen, new_con_scen)) return _add_con(c, f, pars, dims, start, lcon, ucon, name, SecondStageConstraintTag()) end diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index 565372472..31c937136 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -6,372 +6,304 @@ import NLPModels import NLPModels: obj, cons!, cons_nln!, grad!, jac_coord!, hess_coord!, jac_structure!, hess_structure! -import ExaModels: obj!, var_indices, cons_block_indices, get_model, get_nbatch +import NLPModels: obj! +import ExaModels: var_indices, cons_block_indices, get_model, get_nbatch import NLPModelsIpopt: ipopt import ..BACKENDS using Adapt -function runtests() - return @testset "Batch ExaModel" begin +# ============================================================================ +# Helper: build a standard test problem +# ============================================================================ + +function build_batch_model(; ns=2, nv=2, θ_val=[2.0]) + c = BatchExaCore(ns) + @add_var(c, v, nv) + @add_par(c, θ, θ_val) + c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) + c, _ = add_con(c, v[j] - θ[1] for j in 1:nv; lcon = 0.0) + return ExaModel(c) +end + +# ============================================================================ +# Extract test logic into functions to avoid Julia 1.12 GC/compiler segfault +# ============================================================================ + +function test_construction() + model = build_batch_model(ns=3) + @test get_nbatch(model) == 3 + @test NLPModels.get_nvar(model) == 2 + @test NLPModels.get_ncon(model) == 2 + @test model isa NLPModels.AbstractNLPModel + @test size(model.meta.x0) == (2, 3) +end + +function test_obj() + model = build_batch_model() + bx = [1.0 3.0; 2.0 4.0] + bf = zeros(2) + obj!(model, bx, bf) + @test bf[1] ≈ 10.0 + @test bf[2] ≈ 50.0 + @test obj(model, bx) ≈ bf + flat = get_model(model) + @test sum(bf) ≈ obj(flat, vec(bx)) +end + +function test_grad() + model = build_batch_model() + bx = [1.0 3.0; 2.0 4.0] + bg = zeros(2, 2) + grad!(model, bx, bg) + @test bg[:, 1] ≈ [4.0, 8.0] + @test bg[:, 2] ≈ [12.0, 16.0] + g_flat = zeros(4) + grad!(get_model(model), vec(bx), g_flat) + @test vec(bg) ≈ g_flat +end + +function test_cons() + model = build_batch_model() + bx = [1.0 3.0; 2.0 4.0] + bc = zeros(2, 2) + cons!(model, bx, bc) + @test bc[:, 1] ≈ [-1.0, 0.0] + @test bc[:, 2] ≈ [1.0, 2.0] + c_flat = zeros(4) + cons_nln!(get_model(model), vec(bx), c_flat) + @test vec(bc) ≈ c_flat +end + +function test_jac_hess() + model = build_batch_model() + ns, nv = 2, 2 + flat = get_model(model) + bx = [1.0 3.0; 2.0 4.0] + + # --- Jacobian values --- + nnzj = NLPModels.get_nnzj(model) + jvals = zeros(nnzj, ns) + jac_coord!(model, bx, jvals) + jvals_flat = zeros(NLPModels.get_nnzj(flat)) + jac_coord!(flat, vec(bx), jvals_flat) + @test vec(jvals) ≈ jvals_flat + + # --- Hessian values --- + nnzh = NLPModels.get_nnzh(model) + by = ones(nv, ns) + hvals = zeros(nnzh, ns) + hess_coord!(model, bx, by, hvals) + hvals_flat = zeros(NLPModels.get_nnzh(flat)) + hess_coord!(flat, vec(bx), vec(by), hvals_flat) + @test vec(hvals) ≈ hvals_flat +end + +function test_hess_obj_weight() + ns, nv = 2, 2 + c = BatchExaCore(ns) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) + c, _ = add_con(c, v[j]^2 for j in 1:nv) + model = ExaModel(c) + + nc = NLPModels.get_ncon(model) + nnzh = NLPModels.get_nnzh(model) + bx = [1.0 3.0; 2.0 4.0] + by = ones(nc, ns) + flat = get_model(model) + + hvals_w1 = zeros(nnzh, ns) + hess_coord!(model, bx, by, hvals_w1; obj_weight = 1.0) + hvals_flat_w1 = zeros(NLPModels.get_nnzh(flat)) + hess_coord!(flat, vec(bx), vec(by), hvals_flat_w1; obj_weight = 1.0) + @test vec(hvals_w1) ≈ hvals_flat_w1 + + hvals_w2 = zeros(nnzh, ns) + hess_coord!(model, bx, by, hvals_w2; obj_weight = 2.0) + hvals_flat_w2 = zeros(NLPModels.get_nnzh(flat)) + hess_coord!(flat, vec(bx), vec(by), hvals_flat_w2; obj_weight = 2.0) + @test vec(hvals_w2) ≈ hvals_flat_w2 + + @test hvals_w1 != hvals_w2 +end + +function test_hess_vector_obj_weight() + ns, nv = 2, 2 + c = BatchExaCore(ns) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) + c, _ = add_con(c, v[j]^2 for j in 1:nv) + model = ExaModel(c) + + nc = NLPModels.get_ncon(model) + nnzh = NLPModels.get_nnzh(model) + bx = [1.0 3.0; 2.0 4.0] + by = ones(nc, ns) + + # Vector obj_weight = [w1, w2] + wvec = [1.5, 3.0] + hvals_vec = zeros(nnzh, ns) + hess_coord!(model, bx, by, hvals_vec; obj_weight = wvec) + + # Uniform scalar weights for comparison + hvals_w1 = zeros(nnzh, ns) + hess_coord!(model, bx, by, hvals_w1; obj_weight = wvec[1]) + hvals_w2 = zeros(nnzh, ns) + hess_coord!(model, bx, by, hvals_w2; obj_weight = wvec[2]) + + # With uniform weight, both instances get the same obj contribution. + # With vector weight, instance 1 gets w1, instance 2 gets w2. + # The constraint hessian is unaffected by obj_weight, so it is the same. + # Check: vector result differs from both uniform-scalar results. + @test hvals_vec != hvals_w1 + @test hvals_vec != hvals_w2 + + # Verify consistency: uniform weight = special case of vector weight + hvals_uniform = zeros(nnzh, ns) + hess_coord!(model, bx, by, hvals_uniform; obj_weight = [2.0, 2.0]) + hvals_scalar = zeros(nnzh, ns) + hess_coord!(model, bx, by, hvals_scalar; obj_weight = 2.0) + @test hvals_uniform ≈ hvals_scalar +end + +function test_multiple_constraints() + ns, nv = 2, 3 + c = BatchExaCore(ns) + @add_var(c, v, nv) + @add_par(c, θ, [1.0]) + c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) + c, _ = add_con(c, v[j] - θ[1] for j in 1:nv) + c, _ = add_con(c, v[1] + v[2] + v[3] for _ in 1:1; ucon = 10.0) + model = ExaModel(c) + flat = get_model(model) + + nc = NLPModels.get_ncon(model) + @test nc == nv + 1 + @test get_nbatch(model) == ns + + bx = reshape(Float64[1, 2, 3, 4, 5, 6], nv, ns) + + # cons! + bc = zeros(nc, ns) + cons!(model, bx, bc) + c_flat = zeros(nc * ns) + cons_nln!(flat, vec(bx), c_flat) + @test vec(bc) ≈ c_flat + + # jac + nnzj = NLPModels.get_nnzj(model) + jvals = zeros(nnzj, ns) + jac_coord!(model, bx, jvals) + jvals_flat = zeros(NLPModels.get_nnzj(flat)) + jac_coord!(flat, vec(bx), jvals_flat) + @test vec(jvals) ≈ jvals_flat + + # hess + nnzh = NLPModels.get_nnzh(model) + by = ones(nc, ns) + hvals = zeros(nnzh, ns) + hess_coord!(model, bx, by, hvals) + hvals_flat = zeros(NLPModels.get_nnzh(flat)) + hess_coord!(flat, vec(bx), vec(by), hvals_flat) + @test vec(hvals) ≈ hvals_flat +end + +function test_error_guards() + model = build_batch_model() + x_vec = ones(2) + @test_throws ArgumentError obj(model, x_vec) + @test_throws ArgumentError cons!(model, x_vec, zeros(2)) + @test_throws ArgumentError grad!(model, x_vec, zeros(2)) +end + +function test_bounds() + ns, nv = 2, 2 + c = BatchExaCore(ns) + @add_var(c, v, nv; start = 0.5, lvar = 0.0, uvar = 10.0) + c, _ = add_obj(c, v[j]^2 for j in 1:nv) + c, _ = add_con(c, v[j] for j in 1:nv; lcon = 0.0, ucon = 100.0) + model = ExaModel(c) + + @test size(model.meta.x0) == (nv, ns) + @test model.meta.x0 ≈ fill(0.5, nv, ns) + @test model.meta.lvar ≈ fill(0.0, nv, ns) + @test model.meta.uvar ≈ fill(10.0, nv, ns) + + flat = get_model(model) + @test NLPModels.get_nvar(flat) == nv * ns + @test flat.meta.x0 ≈ fill(0.5, nv * ns) + @test flat.meta.lvar ≈ fill(0.0, nv * ns) + @test flat.meta.uvar ≈ fill(10.0, nv * ns) +end + +function test_flatten_model() + model = build_batch_model() + flat = get_model(model) + @test flat isa ExaModels.BatchNLPModels.FlattenNLPModel + @test NLPModels.get_nvar(flat) == 2 * 2 + @test NLPModels.get_ncon(flat) == 2 * 2 + + c = ExaCore(concrete = Val(true)) + c, x = add_var(c, 2) + c, _ = add_obj(c, x[i]^2 for i in 1:2) + m = ExaModel(c) + @test get_model(m) === m +end + +function test_ipopt_simple() + ns, nv = 3, 1 + c = BatchExaCore(ns) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + c, _ = add_obj(c, (v[1] - θ[1])^2 for _ in 1:1) + c, _ = add_con(c, v[1] for _ in 1:1; lcon = 0.0, ucon = Inf) + model = ExaModel(c) + + result = ipopt(get_model(model); print_level = 0) + @test result.status == :first_order + for i in 1:ns + @test result.solution[var_indices(model, i)] ≈ [2.0] atol = 1e-5 + end + @test isapprox(result.objective, 0.0; atol = 1e-8) +end - @testset "Construction and dimensions" begin - ns, nv = 3, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - @add_par(c, θ, EachInstance(), [1.0, 2.0]) - @add_obj(c, θ[j, s] * v[j, s]^2 for j in 1:nv, s in 1:ns) - @add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:ns) - model = ExaModel(c) - - @test get_nbatch(model) == 3 - @test NLPModels.get_nvar(model) == nv - @test NLPModels.get_ncon(model) == nv - - # ExaModel <: AbstractNLPModel - @test model isa NLPModels.AbstractNLPModel - end - - @testset "Batch obj! evaluation" begin - ns, nv = 2, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - @add_par(c, θ, EachInstance(), [2.0]) - nb = get_nbatch(c) - c, _ = add_obj(c, θ[1, s] * v[j, s]^2 for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) - model = ExaModel(c) - - # bx: (nv, ns) matrix - bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) - - bf = zeros(ns) - obj!(model, bx, bf) - - # Both instances have θ=2 - # instance1: v=[1,2], obj = 2*(1 + 4) = 10 - # instance2: v=[3,4], obj = 2*(9 + 16) = 50 - @test bf[1] ≈ 10.0 - @test bf[2] ≈ 50.0 - - # Consistency: sum(bf) ≈ obj(get_model(m), vec(bx)) - @test sum(bf) ≈ obj(get_model(model), vec(bx)) - - # Convenience obj() also works - bf2 = obj(model, bx) - @test bf2 ≈ bf - end - - @testset "Batch grad! evaluation" begin - ns, nv = 2, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - @add_par(c, θ, EachInstance(), [2.0]) - nb = get_nbatch(c) - c, _ = add_obj(c, θ[1, s] * v[j, s]^2 for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) - model = ExaModel(c) - - bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) - bg = zeros(nv, ns) - grad!(model, bx, bg) - - # ∂(θ*v²)/∂v = 2*θ*v, θ=2 for both instances - # s1: [2*2*1, 2*2*2] = [4, 8] - # s2: [2*2*3, 2*2*4] = [12, 16] - @test bg[:, 1] ≈ [4.0, 8.0] - @test bg[:, 2] ≈ [12.0, 16.0] - - # Consistency with fused model - g_flat = zeros(ns * nv) - grad!(get_model(model), vec(bx), g_flat) - @test vec(bg) ≈ g_flat - end - - @testset "Batch cons! evaluation" begin - ns, nv = 2, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - @add_par(c, θ, EachInstance(), [1.0]) - nb = get_nbatch(c) - c, _ = add_obj(c, v[j, s]^2 for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[j, s] - θ[1, s] for j in 1:nv, s in 1:nb) - model = ExaModel(c) - - bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) - bc = zeros(nv, ns) - cons!(model, bx, bc) - - # Both instances have θ=1 - # s1: v=[1,2], θ=1 → [0, 1] - # s2: v=[3,4], θ=1 → [2, 3] - @test bc[:, 1] ≈ [0.0, 1.0] - @test bc[:, 2] ≈ [2.0, 3.0] - - # Consistency with fused model - c_flat = zeros(ns * nv) - cons_nln!(get_model(model), vec(bx), c_flat) - @test vec(bc) ≈ c_flat - end - - @testset "Batch jac_structure! and jac_coord!" begin - ns, nv, nc = 2, 2, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - nb = get_nbatch(c) - c, _ = add_obj(c, v[j, s]^2 for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) - model = ExaModel(c) - - nnzj = NLPModels.get_nnzj(model) - @test nnzj > 0 - - rows = zeros(Int, nnzj) - cols = zeros(Int, nnzj) - jac_structure!(model, rows, cols) - - # Per-instance local indices: rows ∈ 1:nc, cols ∈ 1:nv - @test all(r -> 1 <= r <= nc, rows) - @test all(c -> 1 <= c <= nv, cols) - - # Evaluate Jacobian: jvals is flat vector (nnzj * ns) - bx = reshape(ones(nv * ns), nv, ns) - jvals = zeros(nnzj * ns) - jac_coord!(model, bx, jvals) - - # Linear constraints → all values should be 1 - @test all(v -> v ≈ 1.0, jvals) - end - - @testset "Batch hess_structure! and hess_coord!" begin - ns, nv = 2, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - @add_par(c, θ, EachInstance(), [2.0]) - nb = get_nbatch(c) - c, _ = add_obj(c, θ[1, s] * v[j, s]^2 for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) - model = ExaModel(c) - - nnzh = NLPModels.get_nnzh(model) - @test nnzh > 0 - - rows = zeros(Int, nnzh) - cols = zeros(Int, nnzh) - hess_structure!(model, rows, cols) - - # Per-instance local indices - @test all(r -> 1 <= r <= nv, rows) - @test all(c -> 1 <= c <= nv, cols) - - # Evaluate with uniform obj_weight - bx = reshape(ones(nv * ns), nv, ns) - by = zeros(nv, ns) - bobj_weight = ones(ns) - hvals = zeros(nnzh * ns) - hess_coord!(model, bx, by, bobj_weight, hvals) - - # Hessian of θ*v[j]^2 is 2*θ on diagonal, θ=2 for all instances - # Both instances: 2*2 = 4 - hvals_s1 = hvals[1:nnzh] - hvals_s2 = hvals[nnzh+1:2*nnzh] - @test any(v -> v ≈ 4.0, hvals_s1) - @test any(v -> v ≈ 4.0, hvals_s2) - end - - @testset "hess_coord! with varying obj_weight" begin - ns, nv = 2, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - @add_par(c, θ, EachInstance(), [2.0]) - nb = get_nbatch(c) - c, _ = add_obj(c, θ[1, s] * v[j, s]^2 for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[j, s]^3 for j in 1:nv, s in 1:nb) - model = ExaModel(c) - - nc = NLPModels.get_ncon(model) - nnzh = NLPModels.get_nnzh(model) - bx = reshape([1.0, 2.0, 3.0, 4.0], nv, ns) - by = ones(nc, ns) - - # Uniform weight for reference - hvals_uniform = zeros(nnzh * ns) - hess_coord!(model, bx, by, [1.0, 1.0], hvals_uniform) - - # Varying weights - hvals_varying = zeros(nnzh * ns) - hess_coord!(model, bx, by, [2.0, 0.5], hvals_varying) - - # Compute reference via fused model for each instance - inner = get_model(model) - total_nnzh = NLPModels.get_nnzh(inner) - - # obj-only hessian - hess_obj = zeros(total_nnzh) - hess_coord!(inner, vec(bx), zeros(ns * nc), hess_obj; obj_weight = 1.0) - - # con-only hessian - hess_con = zeros(total_nnzh) - hess_coord!(inner, vec(bx), vec(by), hess_con; obj_weight = 0.0) - - # Verify per-instance reconstruction - perm = ExaModels._batch_hess_perm(model) - for s in 1:ns - for k in 1:nnzh - idx = perm[(s - 1) * nnzh + k] - expected = [2.0, 0.5][s] * hess_obj[idx] + hess_con[idx] - @test hvals_varying[(s - 1) * nnzh + k] ≈ expected - end - end - end - - @testset "Multiple constraint calls" begin - ns, nv = 2, 3 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - @add_par(c, θ, EachInstance(), [1.0]) - nb = get_nbatch(c) - c, _ = add_obj(c, θ[1, s] * v[j, s]^2 for j in 1:nv, s in 1:nb) - # Two separate constraint calls per instance - c, _ = add_con(c, EachInstance(), v[j, s] - θ[1, s] for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[1, s] + v[2, s] + v[3, s] for s in 1:nb; ucon = 10.0) - model = ExaModel(c) - - nc = NLPModels.get_ncon(model) - # nc = nv + 1 = 4 per instance - @test nc == nv + 1 - @test get_nbatch(model) == ns - - bx = reshape([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], nv, ns) - bc = zeros(nc, ns) - cons!(model, bx, bc) - - # Consistency with fused model — this is the definitive check - c_flat = zeros(nc * ns) - cons_nln!(get_model(model), vec(bx), c_flat) - @test vec(bc) ≈ c_flat - - # Hessian with both obj and con contributions - nnzh = NLPModels.get_nnzh(model) - - # Evaluate hessian - by = ones(nc, ns) - bobj_weight = ones(ns) - hvals = zeros(nnzh * ns) - hess_coord!(model, bx, by, bobj_weight, hvals) - @test any(v -> v != 0.0, hvals[1:nnzh]) - @test any(v -> v != 0.0, hvals[nnzh+1:2*nnzh]) - end - - @testset "Get underlying model" begin - ns, nv = 2, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - nb = get_nbatch(c) - c, _ = add_obj(c, v[j, s]^2 for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) - model = ExaModel(c) - - inner = get_model(model) - @test inner isa ExaModels.ExaModel - @test NLPModels.get_nvar(inner) == ns * nv - @test NLPModels.get_ncon(inner) == ns * nv - end - - @testset "Variable bounds and start values" begin - ns, nv = 2, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv; start = 0.5, lvar = 0.0, uvar = 10.0) - nb = get_nbatch(c) - c, _ = add_obj(c, v[j, s]^2 for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb; lcon = 0.0, ucon = 100.0) - model = ExaModel(c) - - # Check that meta matrices have correct shape and values - @test size(model.meta.x0) == (nv, ns) - @test all(model.meta.x0 .== 0.5) - - @test all(model.meta.lvar .== 0.0) - @test all(model.meta.uvar .== 10.0) - end - - @testset "Error on vector arguments" begin - ns, nv = 2, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - nb = get_nbatch(c) - c, _ = add_obj(c, v[j, s]^2 for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[j, s] for j in 1:nv, s in 1:nb) - model = ExaModel(c) - - x_vec = ones(nv) - c_vec = zeros(nv) - g_vec = zeros(nv) - - @test_throws ArgumentError obj(model, x_vec) - @test_throws ArgumentError cons!(model, x_vec, c_vec) - @test_throws ArgumentError grad!(model, x_vec, g_vec) - end - - @testset "Ipopt solver with known solution" begin - ns, nv = 3, 1 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - @add_par(c, θ, EachInstance(), [2.0]) - nb = get_nbatch(c) - # Each instance minimizes (v[1,s] - θ[1,s])^2 - # But θ is the same (2.0) for all instances with this API - c, _ = add_obj(c, (v[1, s] - θ[1, s])^2 for s in 1:nb) - c, _ = add_con(c, EachInstance(), v[1, s] for s in 1:nb; lcon = 0.0, ucon = Inf) - model = ExaModel(c) - - # Solve via fused model - result = ipopt(get_model(model); print_level = 0) - @test result.status == :first_order - - x_sol = result.solution - # All instances have θ=2, so optimal v*=2 for each - for i in 1:ns - @test x_sol[var_indices(model, i)] ≈ [2.0] atol = 1.0e-5 - end - @test result.objective ≈ 0.0 atol = 1.0e-8 - end - - @testset "Ipopt solver - multiple variables per instance" begin - ns, nv = 2, 2 - - c = BatchExaCore(ns) - @add_var(c, v, EachInstance(), nv) - @add_par(c, θ, EachInstance(), [1.0, 3.0]) - nb = get_nbatch(c) - # Each instance minimizes sum of (v[j,s] - θ[j,s])^2 - c, _ = add_obj(c, (v[j, s] - θ[j, s])^2 for j in 1:nv, s in 1:nb) - c, _ = add_con(c, EachInstance(), v[1, s] + v[2, s] for s in 1:nb; ucon = 10.0) - model = ExaModel(c) - - result = ipopt(get_model(model); print_level = 0) - @test result.status == :first_order - - x_sol = result.solution - # Both instances have θ=[1,3], so optimal v*=[1,3] for each - @test x_sol[var_indices(model, 1)] ≈ [1.0, 3.0] atol = 1.0e-5 - @test x_sol[var_indices(model, 2)] ≈ [1.0, 3.0] atol = 1.0e-5 - @test result.objective ≈ 0.0 atol = 1.0e-8 - end +function test_ipopt_multi() + ns, nv = 2, 2 + c = BatchExaCore(ns) + @add_var(c, v, nv) + @add_par(c, θ, [1.0, 3.0]) + c, _ = add_obj(c, (v[j] - θ[j])^2 for j in 1:nv) + c, _ = add_con(c, v[1] + v[2] for _ in 1:1; ucon = 10.0) + model = ExaModel(c) + + result = ipopt(get_model(model); print_level = 0) + @test result.status == :first_order + @test result.solution[var_indices(model, 1)] ≈ [1.0, 3.0] atol = 1e-5 + @test result.solution[var_indices(model, 2)] ≈ [1.0, 3.0] atol = 1e-5 + @test isapprox(result.objective, 0.0; atol = 1e-8) +end + +# ============================================================================ +function runtests() + return @testset "Batch ExaModel" begin + @testset "Construction" test_construction() + @testset "obj!" test_obj() + @testset "grad!" test_grad() + @testset "cons!" test_cons() + @testset "jac and hess" test_jac_hess() + @testset "hess obj_weight" test_hess_obj_weight() + @testset "hess vector obj_weight" test_hess_vector_obj_weight() + @testset "Multiple constraints" test_multiple_constraints() + @testset "Error guards" test_error_guards() + @testset "Bounds" test_bounds() + @testset "flatten_model" test_flatten_model() + @testset "Ipopt simple" test_ipopt_simple() + @testset "Ipopt multi" test_ipopt_multi() end end diff --git a/test/LinAlgTest/LinAlgTest.jl b/test/LinAlgTest/LinAlgTest.jl deleted file mode 100644 index f1bd0841f..000000000 --- a/test/LinAlgTest/LinAlgTest.jl +++ /dev/null @@ -1,1560 +0,0 @@ -module LinAlgTest - -using ExaModels -using Test, ForwardDiff, LinearAlgebra - -# --- AD correctness helpers (from original LinAlgTest) --- - -function gradient(f, x) - T = eltype(x) - y = fill!(similar(x), zero(T)) - ExaModels.gradient!(y, (p, x, θ) -> f(x), x, nothing, nothing, one(T)) - return y -end - -function sgradient(f, x) - T = eltype(x) - - ff = f(ExaModels.VarSource()) - d = ff(ExaModels.Identity(), ExaModels.AdjointNodeSource(nothing), nothing) - y1 = [] - ExaModels.grpass(d, nothing, y1, nothing, 0, NaN) - - a1 = unique(y1) - comp = ExaModels.Compressor(Tuple(findfirst(isequal(i), a1) for i in y1)) - - n = length(a1) - buffer = fill!(similar(x, n), zero(T)) - buffer_I = similar(x, Tuple{Int, Int}, n) - - ExaModels.sgradient!(buffer_I, ff, nothing, nothing, nothing, comp, 0, NaN) - ExaModels.sgradient!(buffer, ff, nothing, x, nothing, comp, 0, one(T)) - - y = zeros(length(x)) - y[collect(i for (i, j) in buffer_I)] += buffer - - return y -end - -function shessian(f, x) - T = eltype(x) - - ff = f(ExaModels.VarSource()) - t = ff(ExaModels.Identity(), ExaModels.SecondAdjointNodeSource(nothing), nothing) - y2 = [] - ExaModels.hrpass0(t, nothing, y2, nothing, nothing, 0, NaN, NaN) - - a2 = unique(y2) - comp = ExaModels.Compressor(Tuple(findfirst(isequal(i), a2) for i in y2)) - - n = length(a2) - buffer = fill!(similar(x, n), zero(T)) - buffer_I = similar(x, Int, n) - buffer_J = similar(x, Int, n) - - ExaModels.shessian!( - buffer_I, - buffer_J, - ff, - nothing, - nothing, - nothing, - comp, - 0, - NaN, - NaN, - ) - ExaModels.shessian!(buffer, nothing, ff, nothing, x, nothing, comp, 0, one(T), zero(T)) - - y = zeros(length(x), length(x)) - for (k, (i, j)) in enumerate(zip(buffer_I, buffer_J)) - if i == j - y[i, j] += buffer[k] - else - y[i, j] += buffer[k] - y[j, i] += buffer[k] - end - end - return y -end - -_vec(x, inds) = [x[i] for i in inds] -_mat(x, rows, cols) = [x[rows[i] + cols[j] - 1] for i in eachindex(rows), j in eachindex(cols)] - -const LINALG_FUNCTIONS = [ - ("linalg-dot-node-node", x -> dot(_vec(x, 1:3), _vec(x, 4:6))), - ("linalg-dot-real-node", x -> dot([1.0, 2.0, 3.0], _vec(x, 1:3))), - ("linalg-dot-node-real", x -> dot(_vec(x, 1:3), [1.0, 2.0, 3.0])), - ("linalg-sum", x -> sum(_vec(x, 1:4))), - ("linalg-norm2", x -> norm(_vec(x, 1:3))), - ("linalg-norm3", x -> norm(_vec(x, 1:3), 3)), - ( - "linalg-matvec-node-node", - x -> begin - A = [x[1] x[3]; x[2] x[4]] - v = [x[5], x[6]] - r = A * v - r[1] + r[2] - end, - ), - ( - "linalg-matvec-real-node", - x -> begin - A = [1.0 2.0; 3.0 4.0] - v = [x[1], x[2]] - r = A * v - r[1] + r[2] - end, - ), - ( - "linalg-tr", - x -> begin - A = [x[1] x[3]; x[2] x[4]] - tr(A) - end, - ), - ( - "linalg-det-2x2", - x -> begin - A = [x[1] x[2]; x[3] x[4]] - det(A) - end, - ), - ( - "linalg-det-3x3", - x -> begin - A = [x[1] x[2] x[3]; x[4] x[5] x[6]; x[7] x[8] x[9]] - det(A) - end, - ), - ( - "linalg-cross", - x -> begin - c = cross(_vec(x, 1:3), _vec(x, 4:6)) - c[1] + c[2] + c[3] - end, - ), - ( - "linalg-vec-add", - x -> begin - r = _vec(x, 1:3) + _vec(x, 4:6) - r[1] + r[2] + r[3] - end, - ), - ( - "linalg-vec-sub", - x -> begin - r = _vec(x, 1:3) - _vec(x, 4:6) - r[1] + r[2] + r[3] - end, - ), - ( - "linalg-scalar-vec", - x -> begin - r = x[1] * _vec(x, 2:4) - r[1] + r[2] + r[3] - end, - ), - ( - "linalg-matmul", - x -> begin - A = [x[1] x[2]; x[3] x[4]] - B = [x[5] x[6]; x[7] x[8]] - C = A * B - C[1, 1] + C[1, 2] + C[2, 1] + C[2, 2] - end, - ), - ( - "linalg-composite-norm-matvec", - x -> begin - A = [1.0 2.0; 3.0 4.0] - v = [x[1], x[2]] - norm(A * v) - end, - ), - ( - "linalg-composite-det-dot", - x -> begin - A = [x[1] x[2]; x[3] x[4]] - det(A) * dot(_vec(x, 5:6), _vec(x, 7:8)) - end, - ), -] - -# --- Type dispatch helpers (from original LinAlgTest2) --- - -is_null_zero(x::ExaModels.Null) = iszero(x.value) -is_null_zero(x::ExaModels.AbstractNode) = false - -function create_nodes() - x = ExaModels.Null(1.0) - y = ExaModels.Null(2.0) - z = ExaModels.Null(3.0) - w = ExaModels.Null(4.0) - return x, y, z, w -end - -# --- Main test runner --- - -function runtests() - return @testset "Linear Algebra test" begin - - # ===================================================================== - # AD correctness: gradient, sparse gradient, sparse Hessian vs ForwardDiff - # ===================================================================== - @testset "AD correctness" begin - for (name, f) in LINALG_FUNCTIONS - x0 = 0.5 .+ rand(10) # avoid zero for norm derivatives - @testset "$name" begin - g_fd = ForwardDiff.gradient(f, x0) - h_fd = ForwardDiff.hessian(f, x0) - @test gradient(f, x0) ≈ g_fd atol = 1.0e-6 - @test sgradient(f, x0) ≈ g_fd atol = 1.0e-6 - @test shessian(f, x0) ≈ h_fd atol = 1.0e-6 - end - end - end - - # ===================================================================== - # Type dispatch and structure: return types, sizes, zero optimizations - # ===================================================================== - @testset "Type dispatch and structure" begin - x, y, z, w = create_nodes() - - @testset "Type conversions and promotions" begin - node = convert(ExaModels.AbstractNode, 5) - @test node isa ExaModels.Null - @test node isa ExaModels.AbstractNode - @test node.value == 5 - - zero_int = convert(ExaModels.AbstractNode, 0) - @test zero_int isa ExaModels.Null - @test iszero(zero_int.value) - @test zero_int === zero(ExaModels.AbstractNode) - - zero_float = convert(ExaModels.AbstractNode, 0.0) - @test zero_float isa ExaModels.Null - @test iszero(zero_float.value) - @test zero_float === zero(ExaModels.AbstractNode) - - arr = [x, 2.0, 3.0] - @test eltype(arr) == ExaModels.AbstractNode - end - - @testset "Scalar × Vector multiplication" begin - v_num = [1.0, 2.0, 3.0] - - result1 = x * v_num - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - vec_nodes = [x, y, z] - result2 = 2.0 * vec_nodes - @test length(result2) == 3 - @test result2 isa Vector - @test result2[1] isa ExaModels.AbstractNode - end - - @testset "Vector × Scalar multiplication" begin - v_num = [1.0, 2.0, 3.0] - vec_nodes = [x, y, z] - - result1 = v_num * x - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - result2 = vec_nodes .* 2.5 - @test length(result2) == 3 - @test result2 isa Vector - @test result2[1] isa ExaModels.AbstractNode - end - - @testset "Scalar × Matrix multiplication" begin - A_num = [1.0 2.0; 3.0 4.0] - mat_nodes = [x y; z w] - - result1 = x * A_num - @test size(result1) == (2, 2) - @test result1 isa Matrix - @test result1[1, 1] isa ExaModels.AbstractNode - - result2 = 3.0 * mat_nodes - @test size(result2) == (2, 2) - @test result2 isa Matrix - @test result2[1, 1] isa ExaModels.AbstractNode - end - - @testset "Matrix × Scalar multiplication" begin - A_num = [1.0 2.0; 3.0 4.0] - mat_nodes = [x y; z w] - - result1 = A_num * x - @test size(result1) == (2, 2) - @test result1 isa Matrix - @test result1[1, 1] isa ExaModels.AbstractNode - - result2 = mat_nodes .* 1.5 - @test size(result2) == (2, 2) - @test result2 isa Matrix - @test result2[1, 1] isa ExaModels.AbstractNode - end - - @testset "Dot product" begin - v_num = [1.0, 2.0, 3.0] - vec_nodes = [x, y, z] - - result1 = dot(v_num, vec_nodes) - @test result1 isa ExaModels.AbstractNode - - result2 = dot(vec_nodes, v_num) - @test result2 isa ExaModels.AbstractNode - end - - @testset "Dot product (Real × Real fallback)" begin - v1 = [1.0, 2.0, 3.0] - v2 = [4.0, 5.0, 6.0] - - result1 = dot(v1, v2) - @test result1 isa Real - @test result1 ≈ 32.0 - - v1_view = @view v1[1:3] - v2_view = @view v2[1:3] - result2 = dot(v1_view, v2_view) - @test result2 isa Real - @test result2 ≈ 32.0 - - v1_reshaped = reshape([1.0, 2.0, 3.0], 3) - v2_reshaped = reshape([4.0, 5.0, 6.0], 3) - result3 = dot(v1_reshaped, v2_reshaped) - @test result3 isa Real - @test result3 ≈ 32.0 - - complex_vec1 = ComplexF64[1.0 + 2.0im, 3.0 + 4.0im] - complex_vec2 = ComplexF64[5.0 + 6.0im, 7.0 + 8.0im] - real_reinterp1 = reinterpret(Float64, complex_vec1) - real_reinterp2 = reinterpret(Float64, complex_vec2) - result4 = dot(real_reinterp1, real_reinterp2) - @test result4 isa Real - @test result4 ≈ 1*5 + 2*6 + 3*7 + 4*8 - - result5 = dot(v1_view, v2_reshaped) - @test result5 isa Real - @test result5 ≈ 32.0 - end - - @testset "Matrix × Vector product" begin - A_num = [1.0 2.0 3.0; 4.0 5.0 6.0] - vec_nodes = [x, y, z] - - result = A_num * vec_nodes - @test length(result) == 2 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Matrix × Matrix product" begin - A_num = [1.0 2.0; 3.0 4.0] - B_nodes = [x y; z w] - - result1 = A_num * B_nodes - @test size(result1) == (2, 2) - @test result1 isa Matrix - @test result1[1, 1] isa ExaModels.AbstractNode - - result2 = B_nodes * A_num - @test size(result2) == (2, 2) - @test result2 isa Matrix - @test result2[1, 1] isa ExaModels.AbstractNode - end - - @testset "Adjoint Vector × Vector product" begin - vec_nodes = [x, y, z] - v_num = [1.0, 2.0, 3.0] - - result1 = vec_nodes' * v_num - @test result1 isa ExaModels.AbstractNode - - result2 = v_num' * vec_nodes - @test result2 isa ExaModels.AbstractNode - - vec_nodes2 = [y, z, x] - result3 = vec_nodes' * vec_nodes2 - @test result3 isa ExaModels.AbstractNode - end - - @testset "Adjoint Vector × Matrix product" begin - vec_nodes = [x, y, z] - A_num = [1.0 2.0; 3.0 4.0; 5.0 6.0] - - result = vec_nodes' * A_num - @test size(result) == (1, 2) - @test result isa LinearAlgebra.Adjoint - end - - @testset "Matrix adjoint" begin - mat_nodes = [x y z; y z x] - - result = adjoint(mat_nodes) - @test size(result) == (3, 2) - @test result isa Matrix - end - - @testset "Determinant" begin - A1 = reshape([x], 1, 1) - result1 = det(A1) - @test result1 isa ExaModels.AbstractNode - - A2 = [x y; z w] - result2 = det(A2) - @test result2 isa ExaModels.AbstractNode - - A3 = [x y z; y z x; z x y] - result3 = det(A3) - @test result3 isa ExaModels.AbstractNode - end - - @testset "Broadcasting operations" begin - vec_nodes = [x, y, z] - mat_nodes = [x y; z w] - - result1 = cos.(vec_nodes) - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - result2 = sin.(vec_nodes) - @test length(result2) == 3 - @test result2[1] isa ExaModels.AbstractNode - - result3 = exp.(vec_nodes) - @test length(result3) == 3 - @test result3[1] isa ExaModels.AbstractNode - - result4 = cos.(mat_nodes) - @test size(result4) == (2, 2) - @test result4 isa Matrix - @test result4[1, 1] isa ExaModels.AbstractNode - - result5 = vec_nodes .+ 1.0 - @test length(result5) == 3 - @test result5[1] isa ExaModels.AbstractNode - - result6 = vec_nodes .* 2.0 - @test length(result6) == 3 - @test result6[1] isa ExaModels.AbstractNode - - v_num = [1.0, 2.0, 3.0] - result7 = vec_nodes .+ v_num - @test length(result7) == 3 - @test result7[1] isa ExaModels.AbstractNode - - result8 = vec_nodes .* v_num - @test length(result8) == 3 - @test result8[1] isa ExaModels.AbstractNode - end - - @testset "Trace" begin - A2 = [x y; z w] - result1 = tr(A2) - @test result1 isa ExaModels.AbstractNode - - A3 = [x y z; y z w; z w x] - result2 = tr(A3) - @test result2 isa ExaModels.AbstractNode - - A_rect = [x y z; y z w] - @test_throws AssertionError tr(A_rect) - end - - @testset "Norms" begin - vec_nodes = [x, y, z] - mat_nodes = [x y; z w] - - result1 = norm(vec_nodes) - @test result1 isa ExaModels.AbstractNode - - result2 = norm(vec_nodes, 1) - @test result2 isa ExaModels.AbstractNode - - result3 = norm(vec_nodes, 2) - @test result3 isa ExaModels.AbstractNode - - result4 = norm(vec_nodes, 3) - @test result4 isa ExaModels.AbstractNode - - result5 = norm(mat_nodes) - @test result5 isa ExaModels.AbstractNode - - @test_throws ErrorException norm(vec_nodes, Inf) - end - - @testset "Array addition" begin - vec_nodes1 = [x, y, z] - vec_nodes2 = [y, z, w] - v_num = [1.0, 2.0, 3.0] - - result1 = vec_nodes1 + v_num - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - result2 = v_num + vec_nodes1 - @test length(result2) == 3 - @test result2[1] isa ExaModels.AbstractNode - - result3 = vec_nodes1 + vec_nodes2 - @test length(result3) == 3 - @test result3[1] isa ExaModels.AbstractNode - - mat_nodes1 = [x y; z w] - mat_nodes2 = [y z; w x] - A_num = [1.0 2.0; 3.0 4.0] - - result4 = mat_nodes1 + A_num - @test size(result4) == (2, 2) - @test result4 isa Matrix - @test result4[1, 1] isa ExaModels.AbstractNode - - result5 = A_num + mat_nodes1 - @test size(result5) == (2, 2) - @test result5[1, 1] isa ExaModels.AbstractNode - - result6 = mat_nodes1 + mat_nodes2 - @test size(result6) == (2, 2) - @test result6[1, 1] isa ExaModels.AbstractNode - end - - @testset "Array subtraction" begin - vec_nodes1 = [x, y, z] - vec_nodes2 = [y, z, w] - v_num = [1.0, 2.0, 3.0] - - result1 = vec_nodes1 - v_num - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - result2 = v_num - vec_nodes1 - @test length(result2) == 3 - @test result2[1] isa ExaModels.AbstractNode - - result3 = vec_nodes1 - vec_nodes2 - @test length(result3) == 3 - @test result3[1] isa ExaModels.AbstractNode - - mat_nodes1 = [x y; z w] - mat_nodes2 = [y z; w x] - A_num = [1.0 2.0; 3.0 4.0] - - result4 = mat_nodes1 - A_num - @test size(result4) == (2, 2) - @test result4 isa Matrix - @test result4[1, 1] isa ExaModels.AbstractNode - - result5 = A_num - mat_nodes1 - @test size(result5) == (2, 2) - @test result5[1, 1] isa ExaModels.AbstractNode - - result6 = mat_nodes1 - mat_nodes2 - @test size(result6) == (2, 2) - @test result6[1, 1] isa ExaModels.AbstractNode - end - - @testset "Diagonal operations" begin - mat_nodes = [x y z; w x y; z w x] - result1 = diag(mat_nodes) - @test length(result1) == 3 - @test result1 isa Vector - @test result1[1] isa ExaModels.AbstractNode - - mat_rect = [x y z; w x y] - result2 = diag(mat_rect) - @test length(result2) == 2 - @test result2[1] isa ExaModels.AbstractNode - - vec_nodes = [x, y, z] - result3 = diagm(vec_nodes) - @test size(result3) == (3, 3) - @test result3 isa Matrix - @test result3[1, 1] isa ExaModels.AbstractNode - @test is_null_zero(result3[1, 2]) - @test is_null_zero(result3[2, 1]) - - result4 = diagm(1 => vec_nodes) - @test size(result4) == (4, 4) - @test result4[1, 2] isa ExaModels.AbstractNode - @test is_null_zero(result4[1, 1]) - - result5 = diagm(-1 => vec_nodes) - @test size(result5) == (4, 4) - @test result5[2, 1] isa ExaModels.AbstractNode - @test is_null_zero(result5[1, 1]) - end - - @testset "Transpose operations" begin - result1 = transpose(x) - @test result1 isa ExaModels.AbstractNode - - mat_nodes = [x y z; w x y] - result2 = transpose(mat_nodes) - @test size(result2) == (3, 2) - @test result2 isa Matrix - @test result2[1, 1] isa ExaModels.AbstractNode - end - - @testset "Dimension mismatch errors" begin - v1 = [1.0, 2.0] - vec_nodes = [x, y, z] - @test_throws AssertionError dot(v1, vec_nodes) - - A = [1.0 2.0; 3.0 4.0] - v = [x, y, z] - @test_throws AssertionError A * v - - A1 = [1.0 2.0; 3.0 4.0] - A2_nodes = [x y; z w; y x] - @test_throws AssertionError A1 * A2_nodes - - A_rect = [x y z; y z x] - @test_throws AssertionError det(A_rect) - - v1 = [x, y] - v2 = [x, y, z] - @test_throws AssertionError v1 + v2 - @test_throws AssertionError v1 - v2 - - A1 = [x y; z w] - A2 = [x y z; w x y] - @test_throws AssertionError A1 + A2 - @test_throws AssertionError A1 - A2 - end - - @testset "ExaCore variable arrays" begin - c = ExaModels.ExaCore(concrete = Val(true)) - @add_var(c, xvar, 2, 0:10, lvar=0, uvar=1) - - v = [xvar[i, 1] for i in 1:2] - @test length(v) == 2 - @test v isa Vector - @test v[1] isa ExaModels.AbstractNode - - A = [xvar[i, j] for (i, j) ∈ Base.product(1:2, 0:3)] - @test size(A) == (2, 4) - @test A isa Matrix - @test A[1, 1] isa ExaModels.AbstractNode - - v_num = [1.0, 2.0] - result1 = v + v_num - @test length(result1) == 2 - @test result1[1] isa ExaModels.AbstractNode - - result2 = v - v_num - @test length(result2) == 2 - @test result2[1] isa ExaModels.AbstractNode - - result3 = 2.0 * v - @test length(result3) == 2 - @test result3[1] isa ExaModels.AbstractNode - - result4 = dot(v, v_num) - @test result4 isa ExaModels.AbstractNode - - result5 = norm(v) - @test result5 isa ExaModels.AbstractNode - - v2 = [xvar[i, 2] for i in 1:2] - result6 = A * [1.0, 2.0, 3.0, 4.0] - @test length(result6) == 2 - @test result6[1] isa ExaModels.AbstractNode - - A_square = [xvar[i, j] for (i, j) ∈ Base.product(1:2, 1:2)] - @test size(A_square) == (2, 2) - - result7 = det(A_square) - @test result7 isa ExaModels.AbstractNode - - result8 = tr(A_square) - @test result8 isa ExaModels.AbstractNode - - result9 = diag(A_square) - @test length(result9) == 2 - @test result9[1] isa ExaModels.AbstractNode - - A_num = [1.0 2.0 3.0 4.0; 5.0 6.0 7.0 8.0] - result10 = A + A_num - @test size(result10) == (2, 4) - @test result10[1, 1] isa ExaModels.AbstractNode - - result11 = A - A_num - @test size(result11) == (2, 4) - @test result11[1, 1] isa ExaModels.AbstractNode - - result12 = cos.(v) - @test length(result12) == 2 - @test result12[1] isa ExaModels.AbstractNode - - result13 = sin.(A_square) - @test size(result13) == (2, 2) - @test result13[1, 1] isa ExaModels.AbstractNode - - result14 = v' - @test size(result14) == (1, 2) - @test result14 isa LinearAlgebra.Adjoint - - result15 = transpose(A) - @test size(result15) == (4, 2) - @test result15 isa Matrix - - result16 = diagm(v) - @test size(result16) == (2, 2) - @test result16[1, 1] isa ExaModels.AbstractNode - @test is_null_zero(result16[1, 2]) - end - - @testset "Method ambiguity fixes" begin - x, y, z, w = create_nodes() - vec_nodes1 = [x, y, z] - vec_nodes2 = [y, z, w] - mat_nodes1 = [x y; z w] - mat_nodes2 = [y z; w x] - - @testset "Scalar × Vector (both AbstractNode)" begin - result = x * vec_nodes1 - @test length(result) == 3 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Vector × Scalar (both AbstractNode)" begin - result = vec_nodes1 * x - @test length(result) == 3 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Scalar × Matrix (both AbstractNode)" begin - result = x * mat_nodes1 - @test size(result) == (2, 2) - @test result isa Matrix - @test result[1, 1] isa ExaModels.AbstractNode - end - - @testset "Matrix × Scalar (both AbstractNode)" begin - result = mat_nodes1 * x - @test size(result) == (2, 2) - @test result isa Matrix - @test result[1, 1] isa ExaModels.AbstractNode - end - - @testset "dot product (both AbstractNode)" begin - result = dot(vec_nodes1, vec_nodes2) - @test result isa ExaModels.AbstractNode - end - - @testset "Matrix × Vector (both AbstractNode)" begin - A = [x y z; w x y] - v = [x, y, z] - - result = A * v - @test length(result) == 2 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Matrix × Matrix (both AbstractNode)" begin - result = mat_nodes1 * mat_nodes2 - @test size(result) == (2, 2) - @test result isa Matrix - @test result[1, 1] isa ExaModels.AbstractNode - end - - @testset "Adjoint Vector × Matrix (both AbstractNode)" begin - A = [x y; z w; y x] - v = [x, y, z] - - result = v' * A - @test size(result) == (1, 2) - @test result isa LinearAlgebra.Adjoint - end - - @testset "Vector + Vector (both AbstractNode)" begin - result = vec_nodes1 + vec_nodes2 - @test length(result) == 3 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Vector - Vector (both AbstractNode)" begin - result = vec_nodes1 - vec_nodes2 - @test length(result) == 3 - @test result isa Vector - @test result[1] isa ExaModels.AbstractNode - end - - @testset "Matrix + Matrix (both AbstractNode)" begin - result = mat_nodes1 + mat_nodes2 - @test size(result) == (2, 2) - @test result isa Matrix - @test result[1, 1] isa ExaModels.AbstractNode - end - - @testset "Matrix - Matrix (both AbstractNode)" begin - result = mat_nodes1 - mat_nodes2 - @test size(result) == (2, 2) - @test result isa Matrix - @test result[1, 1] isa ExaModels.AbstractNode - end - - @testset "Mixed operations (no standard library conflicts)" begin - v_num = [1.0, 2.0, 3.0] - v_num_2 = [1.0, 2.0] - A_num = [1.0 2.0 3.0; 4.0 5.0 6.0] - - @test (vec_nodes1 * 2.0) isa Vector - @test (2.0 * vec_nodes1) isa Vector - @test (mat_nodes1 * 2.0) isa Matrix - @test (2.0 * mat_nodes1) isa Matrix - @test dot(v_num, vec_nodes1) isa ExaModels.AbstractNode - @test dot(vec_nodes1, v_num) isa ExaModels.AbstractNode - @test (A_num * vec_nodes1) isa Vector - @test (mat_nodes1 * v_num_2) isa Vector - @test (A_num * [x y; z w; y x]) isa Matrix - @test (mat_nodes1 * mat_nodes2) isa Matrix - end - end - - @testset "Canonical nodes" begin - @testset "zero and one helpers" begin - z = zero(ExaModels.AbstractNode) - @test z isa ExaModels.Null - @test iszero(z.value) - @test z.value == 0 - - o = one(ExaModels.AbstractNode) - @test o isa ExaModels.Null - @test isone(o.value) - @test o.value == 1 - end - - @testset "zeros and ones array creation" begin - z1 = zeros(ExaModels.AbstractNode, 3) - @test length(z1) == 3 - @test z1 isa Vector{<:ExaModels.AbstractNode} - @test all(is_null_zero.(z1)) - @test all(x -> x isa ExaModels.Null, z1) - - z2 = zeros(ExaModels.AbstractNode, 2, 3) - @test size(z2) == (2, 3) - @test z2 isa Matrix{<:ExaModels.AbstractNode} - @test all(is_null_zero.(z2)) - - z3 = zeros(ExaModels.AbstractNode, 2, 2, 2) - @test size(z3) == (2, 2, 2) - @test z3 isa Array{<:ExaModels.AbstractNode, 3} - @test all(is_null_zero.(z3)) - - o1 = ones(ExaModels.AbstractNode, 3) - @test length(o1) == 3 - @test o1 isa Vector{<:ExaModels.AbstractNode} - @test all(x -> x isa ExaModels.Null && isone(x.value), o1) - - o2 = ones(ExaModels.AbstractNode, 2, 3) - @test size(o2) == (2, 3) - @test o2 isa Matrix{<:ExaModels.AbstractNode} - @test all(x -> x isa ExaModels.Null && isone(x.value), o2) - - o3 = ones(ExaModels.AbstractNode, 2, 2, 2) - @test size(o3) == (2, 2, 2) - @test o3 isa Array{<:ExaModels.AbstractNode, 3} - @test all(x -> x isa ExaModels.Null && isone(x.value), o3) - - z_var = zeros(ExaModels.AbstractNode, 3) - @test length(z_var) == 3 - @test eltype(z_var) <: ExaModels.AbstractNode - @test all(is_null_zero.(z_var)) - - o_var = ones(ExaModels.AbstractNode, 2, 2) - @test size(o_var) == (2, 2) - @test eltype(o_var) <: ExaModels.AbstractNode - @test all(x -> x isa ExaModels.Null && isone(x.value), o_var) - end - - @testset "zeros and ones in operations" begin - x, y, z, w = create_nodes() - - z_vec = zeros(ExaModels.AbstractNode, 3) - vec_nodes = [x, y, z] - - result1 = vec_nodes + z_vec - @test result1[1].value == x.value - @test result1[2].value == y.value - @test result1[3].value == z.value - - result2 = [ExaModels.Null(2), ExaModels.Null(3), ExaModels.Null(4)] .* z_vec - @test all(is_null_zero.(result2)) - - o_vec = ones(ExaModels.AbstractNode, 3) - - result3 = vec_nodes .* o_vec - @test result3[1].value == x.value - @test result3[2].value == y.value - @test result3[3].value == z.value - - I_like = diagm(ones(ExaModels.AbstractNode, 2)) - vec2 = [x, y] - result4 = I_like * vec2 - @test result4[1].value == x.value - @test result4[2].value == y.value - end - end - - @testset "Zero multiplication optimizations" begin - x, y, z, w = create_nodes() - - @testset "Scalar × Vector with zero scalar" begin - result1 = ExaModels.Null(0) * [1.0, 2.0, 3.0] - @test all(is_null_zero.(result1)) - @test result1 isa Vector{<:ExaModels.AbstractNode} - @test length(result1) == 3 - - vec_nodes = [x, y, z] - result2 = 0 * vec_nodes - @test all(is_null_zero.(result2)) - @test result2 isa Vector{<:ExaModels.AbstractNode} - - result3 = 0.0 * vec_nodes - @test all(is_null_zero.(result3)) - end - - @testset "Scalar × Vector with zero elements" begin - result = x * [0, 1.0, 0, 2.0] - @test is_null_zero(result[1]) - @test !is_null_zero(result[2]) - @test is_null_zero(result[3]) - @test !is_null_zero(result[4]) - end - - @testset "Vector × Scalar with zero scalar" begin - vec_nodes = [x, y, z] - - result1 = vec_nodes * 0 - @test all(is_null_zero.(result1)) - @test result1 isa Vector{<:ExaModels.AbstractNode} - - result2 = vec_nodes * 0.0 - @test all(is_null_zero.(result2)) - - result3 = [1.0, 2.0, 3.0] * ExaModels.Null(0) - @test all(is_null_zero.(result3)) - end - - @testset "Vector × Scalar with zero elements" begin - result = [0, 1.0, ExaModels.Null(0), 2.0] * x - @test is_null_zero(result[1]) - @test !is_null_zero(result[2]) - @test is_null_zero(result[3]) - @test !is_null_zero(result[4]) - end - - @testset "Scalar × Matrix with zero scalar" begin - mat_nodes = [x y; z w] - - result1 = 0 * mat_nodes - @test all(is_null_zero.(result1)) - @test result1 isa Matrix{<:ExaModels.AbstractNode} - @test size(result1) == (2, 2) - - result2 = ExaModels.Null(0) * [1.0 2.0; 3.0 4.0] - @test all(is_null_zero.(result2)) - end - - @testset "Matrix × Scalar with zero scalar" begin - mat_nodes = [x y; z w] - - result = mat_nodes * 0.0 - @test all(is_null_zero.(result)) - @test result isa Matrix{<:ExaModels.AbstractNode} - end - end - - @testset "Zero addition optimizations" begin - x, y, z, w = create_nodes() - - @testset "Vector + Vector with zero elements" begin - vec_nodes = [x, y, z] - - result1 = vec_nodes + [0, 0, 0] - @test result1[1].value == x.value - @test result1[2].value == y.value - @test result1[3].value == z.value - @test result1 isa Vector{<:ExaModels.AbstractNode} - - result2 = [0.0, 0.0, 0.0] + vec_nodes - @test result2[1].value == x.value - @test result2[2].value == y.value - @test result2[3].value == z.value - - result3 = vec_nodes + [0, 1.0, 0] - @test result3[1].value == x.value - @test result3[2].value == y.value + 1.0 - @test result3[3].value == z.value - - zero_vec = [ExaModels.Null(0), ExaModels.Null(0), ExaModels.Null(0)] - result4 = vec_nodes + zero_vec - @test result4[1].value == x.value - @test result4[2].value == y.value - @test result4[3].value == z.value - end - - @testset "Matrix + Matrix with zero elements" begin - mat_nodes = [x y; z w] - - result1 = mat_nodes + [0 0; 0 0] - @test result1[1,1].value == x.value - @test result1[1,2].value == y.value - @test result1[2,1].value == z.value - @test result1[2,2].value == w.value - @test result1 isa Matrix{<:ExaModels.AbstractNode} - - result2 = [0.0 0.0; 0.0 0.0] + mat_nodes - @test result2[1,1].value == x.value - @test result2[1,2].value == y.value - - result3 = mat_nodes + [0 1.0; 0 0] - @test result3[1,1].value == x.value - @test result3[1,2].value == y.value + 1.0 - @test result3[2,1].value == z.value - @test result3[2,2].value == w.value - end - end - - @testset "Zero subtraction optimizations" begin - x, y, z, w = create_nodes() - - @testset "Vector - Vector with zero elements" begin - vec_nodes = [x, y, z] - - result1 = vec_nodes - [0, 0, 0] - @test result1[1].value == x.value - @test result1[2].value == y.value - @test result1[3].value == z.value - @test result1 isa Vector{<:ExaModels.AbstractNode} - - result2 = vec_nodes - [0, 1.0, 0] - @test result2[1].value == x.value - @test result2[2].value == y.value - 1.0 - @test result2[3].value == z.value - - zero_vec = [ExaModels.Null(0), ExaModels.Null(0), ExaModels.Null(0)] - result3 = vec_nodes - zero_vec - @test result3[1].value == x.value - @test result3[2].value == y.value - @test result3[3].value == z.value - end - - @testset "Matrix - Matrix with zero elements" begin - mat_nodes = [x y; z w] - - result1 = mat_nodes - [0 0; 0 0] - @test result1[1,1].value == x.value - @test result1[1,2].value == y.value - @test result1[2,1].value == z.value - @test result1[2,2].value == w.value - @test result1 isa Matrix{<:ExaModels.AbstractNode} - - result2 = mat_nodes - [0 1.0; 0 0] - @test result2[1,1].value == x.value - @test result2[1,2].value == y.value - 1.0 - @test result2[2,1].value == z.value - @test result2[2,2].value == w.value - end - end - - @testset "Scalar operations on Null nodes (+, -, *)" begin - x, y, z, w = create_nodes() - - e = ExaModels.Node2(+, x, y) - f = ExaModels.Node2(+, z, w) - - @testset "+ operator rules" begin - result1 = ExaModels.Null(3) + ExaModels.Null(5) - @test result1 isa ExaModels.Null - @test result1.value == 8 - - result2 = ExaModels.Null(3) + e - @test result2 isa ExaModels.AbstractNode - @test !(result2 isa ExaModels.Null) - - result3 = e + ExaModels.Null(3) - @test result3 isa ExaModels.AbstractNode - @test !(result3 isa ExaModels.Null) - - result4 = e + f - @test result4 isa ExaModels.AbstractNode - end - - @testset "- operator rules" begin - result1 = ExaModels.Null(5) - ExaModels.Null(3) - @test result1 isa ExaModels.Null - @test result1.value == 2 - - result2 = ExaModels.Null(0) - e - @test result2 isa ExaModels.Node1 - - result3 = ExaModels.Null(3) - e - @test result3 isa ExaModels.AbstractNode - @test !(result3 isa ExaModels.Null) - - result4 = e - ExaModels.Null(3) - @test result4 isa ExaModels.AbstractNode - @test !(result4 isa ExaModels.Null) - - result5 = e - f - @test result5 isa ExaModels.AbstractNode - end - - @testset "* operator rules" begin - result1 = ExaModels.Null(3) * ExaModels.Null(5) - @test result1 isa ExaModels.Null - @test result1.value == 15 - - result2 = ExaModels.Null(0) * e - @test result2 isa ExaModels.Null - @test iszero(result2.value) - - result3 = e * ExaModels.Null(0) - @test result3 isa ExaModels.Null - @test iszero(result3.value) - - result4 = ExaModels.Null(3) * e - @test result4 isa ExaModels.AbstractNode - @test !(result4 isa ExaModels.Null) - - result5 = e * ExaModels.Null(3) - @test result5 isa ExaModels.AbstractNode - @test !(result5 isa ExaModels.Null) - - result6 = e * f - @test result6 isa ExaModels.AbstractNode - end - end - - @testset "sum function" begin - x, y, z, w = create_nodes() - - @testset "sum with zeros" begin - result1 = sum([zero(ExaModels.AbstractNode), zero(ExaModels.AbstractNode), zero(ExaModels.AbstractNode)]) - @test result1 isa ExaModels.Null - @test iszero(result1.value) - - result2 = sum([zero(ExaModels.AbstractNode), zero(ExaModels.AbstractNode), x, zero(ExaModels.AbstractNode)]) - @test result2 isa ExaModels.Null - @test result2.value == x.value - - result3 = sum([zero(ExaModels.AbstractNode), x, zero(ExaModels.AbstractNode), y, zero(ExaModels.AbstractNode)]) - @test result3 isa ExaModels.Null - @test result3.value == x.value + y.value - end - - @testset "sum with single element" begin - result1 = sum([x]) - @test result1 isa ExaModels.Null - @test result1.value == x.value - - result2 = sum([zero(ExaModels.AbstractNode)]) - @test result2 isa ExaModels.Null - @test iszero(result2.value) - end - - @testset "sum with all non-zeros" begin - result1 = sum([x, y, z]) - @test result1 isa ExaModels.Null - @test result1.value == x.value + y.value + z.value - end - - @testset "sum on matrices" begin - mat = [x zero(ExaModels.AbstractNode); zero(ExaModels.AbstractNode) y] - result = sum(mat) - @test result isa ExaModels.Null - @test result.value == x.value + y.value - end - end - - @testset "Optimized dot product" begin - x, y, z, t = create_nodes() - - @testset "dot([1, 0, 1, 0], [x, y, z, t]) = x + z" begin - result = dot([1, 0, 1, 0], [x, y, z, t]) - @test result isa ExaModels.Null - @test result.value == x.value + z.value - end - - @testset "dot with all zeros" begin - result = dot([0, 0, 0], [x, y, z]) - @test result isa ExaModels.Null - @test iszero(result.value) - end - - @testset "dot with all ones" begin - result = dot([1, 1, 1], [x, y, z]) - @test result isa ExaModels.Null - @test result.value == x.value + y.value + z.value - end - - @testset "dot with single non-zero" begin - result = dot([0, 1, 0], [x, y, z]) - @test result isa ExaModels.Null - @test result.value == y.value - end - end - - @testset "Matrix operations with optimization" begin - x, y, z, w = create_nodes() - - @testset "Identity matrix multiplication" begin - I2 = [1 0; 0 1] - vec = [x, y] - result = I2 * vec - - @test result[1] isa ExaModels.Null - @test result[1].value == x.value - @test result[2] isa ExaModels.Null - @test result[2].value == y.value - end - - @testset "Zero matrix multiplication" begin - Z2 = [0 0; 0 0] - vec = [x, y] - result = Z2 * vec - - @test result[1] isa ExaModels.Null - @test iszero(result[1].value) - @test result[2] isa ExaModels.Null - @test iszero(result[2].value) - end - - @testset "Sparse matrix multiplication" begin - A = [1 0 0; 0 0 1] - vec = [x, y, z] - result = A * vec - - @test result[1] isa ExaModels.Null - @test result[1].value == x.value - @test result[2] isa ExaModels.Null - @test result[2].value == z.value - end - end - - @testset "SubArray, ReshapedArray, ReinterpretArray" begin - x, y, z, w = create_nodes() - - @testset "SubArray (views) - vectors of Real" begin - v_num = [1.0, 2.0, 3.0, 4.0] - vec_nodes = [x, y, z, w] - - v_view = @view v_num[1:3] - @test v_view isa SubArray - - result1 = dot(v_view, vec_nodes[1:3]) - @test result1 isa ExaModels.AbstractNode - - result2 = dot(vec_nodes[1:3], v_view) - @test result2 isa ExaModels.AbstractNode - - result3 = x * v_view - @test length(result3) == 3 - @test result3[1] isa ExaModels.AbstractNode - - result4 = 2.0 * vec_nodes[1:3] - @test length(result4) == 3 - - result5 = v_view * x - @test length(result5) == 3 - @test result5[1] isa ExaModels.AbstractNode - - result6 = vec_nodes[1:3] + v_view - @test length(result6) == 3 - @test result6[1] isa ExaModels.AbstractNode - - result7 = v_view + vec_nodes[1:3] - @test length(result7) == 3 - - result8 = vec_nodes[1:3] - v_view - @test length(result8) == 3 - @test result8[1] isa ExaModels.AbstractNode - - result9 = v_view - vec_nodes[1:3] - @test length(result9) == 3 - - result10 = norm(@view vec_nodes[1:3]) - @test result10 isa ExaModels.AbstractNode - end - - @testset "SubArray (views) - vectors of AbstractNode" begin - vec_nodes = [x, y, z, w] - - v_view = @view vec_nodes[1:3] - @test v_view isa SubArray - - v_view2 = @view vec_nodes[2:4] - result1 = v_view + v_view2 - @test length(result1) == 3 - @test result1[1] isa ExaModels.AbstractNode - - result2 = v_view - v_view2 - @test length(result2) == 3 - - result3 = dot(v_view, v_view2) - @test result3 isa ExaModels.AbstractNode - - result4 = 2.5 * v_view - @test length(result4) == 3 - @test result4[1] isa ExaModels.AbstractNode - - result5 = x * v_view - @test length(result5) == 3 - end - - @testset "SubArray (views) - matrices of Real" begin - A_num = [1.0 2.0 3.0; 4.0 5.0 6.0; 7.0 8.0 9.0] - mat_nodes = [x y z; w x y; z w x] - - A_view = @view A_num[1:2, 1:2] - @test A_view isa SubArray - @test size(A_view) == (2, 2) - - result1 = x * A_view - @test size(result1) == (2, 2) - @test result1[1, 1] isa ExaModels.AbstractNode - - result2 = 3.0 * (@view mat_nodes[1:2, 1:2]) - @test size(result2) == (2, 2) - - result3 = A_view * x - @test size(result3) == (2, 2) - @test result3[1, 1] isa ExaModels.AbstractNode - - vec_nodes = [x, y] - result4 = A_view * vec_nodes - @test length(result4) == 2 - @test result4[1] isa ExaModels.AbstractNode - - result5 = (@view mat_nodes[1:2, 1:2]) + A_view - @test size(result5) == (2, 2) - @test result5[1, 1] isa ExaModels.AbstractNode - end - - @testset "SubArray (views) - matrices of AbstractNode" begin - mat_nodes = [x y z; w x y; z w x] - - A_view = @view mat_nodes[1:2, 1:2] - @test A_view isa SubArray - @test size(A_view) == (2, 2) - - B_view = @view mat_nodes[2:3, 2:3] - result1 = A_view * B_view - @test size(result1) == (2, 2) - @test result1[1, 1] isa ExaModels.AbstractNode - - result2 = A_view + B_view - @test size(result2) == (2, 2) - - result3 = A_view - B_view - @test size(result3) == (2, 2) - - result4 = adjoint(A_view) - @test size(result4) == (2, 2) - - result5 = transpose(A_view) - @test size(result5) == (2, 2) - end - - @testset "ReshapedArray - vectors to matrices" begin - v_num = [1.0, 2.0, 3.0, 4.0] - vec_nodes = [x, y, z, w] - - A_reshaped = reshape(v_num, 2, 2) - @test A_reshaped isa Union{Matrix, Base.ReshapedArray} - - result1 = x * A_reshaped - @test size(result1) == (2, 2) - @test result1[1, 1] isa ExaModels.AbstractNode - - mat_reshaped = reshape(vec_nodes, 2, 2) - @test mat_reshaped isa Union{Matrix, Base.ReshapedArray} - - result2 = 2.0 * mat_reshaped - @test size(result2) == (2, 2) - @test result2[1, 1] isa ExaModels.AbstractNode - - result3 = det(mat_reshaped) - @test result3 isa ExaModels.AbstractNode - - result4 = tr(mat_reshaped) - @test result4 isa ExaModels.AbstractNode - end - - @testset "ReshapedArray - matrices to vectors" begin - A_num = [1.0 2.0; 3.0 4.0] - mat_nodes = [x y; z w] - - v_reshaped = reshape(A_num, 4) - @test v_reshaped isa Union{Vector, Base.ReshapedArray} - - result1 = x * v_reshaped - @test length(result1) == 4 - @test result1[1] isa ExaModels.AbstractNode - - vec_reshaped = reshape(mat_nodes, 4) - @test vec_reshaped isa Union{Vector, Base.ReshapedArray} - - result2 = 2.0 * vec_reshaped - @test length(result2) == 4 - @test result2[1] isa ExaModels.AbstractNode - - result3 = norm(vec_reshaped) - @test result3 isa ExaModels.AbstractNode - - result4 = dot(v_reshaped, vec_reshaped) - @test result4 isa ExaModels.AbstractNode - end - - @testset "ReinterpretArray - Complex to Real" begin - complex_vec = ComplexF64[1.0 + 2.0im, 3.0 + 4.0im, 5.0 + 6.0im] - - real_reinterp = reinterpret(Float64, complex_vec) - @test real_reinterp isa Base.ReinterpretArray - @test length(real_reinterp) == 6 - - vec_nodes = [x, y, z, w, ExaModels.Null(5), ExaModels.Null(6)] - - result1 = dot(real_reinterp, vec_nodes) - @test result1 isa ExaModels.AbstractNode - - result2 = dot(vec_nodes, real_reinterp) - @test result2 isa ExaModels.AbstractNode - - result3 = x * real_reinterp - @test length(result3) == 6 - @test result3[1] isa ExaModels.AbstractNode - - result4 = real_reinterp * x - @test length(result4) == 6 - @test result4[1] isa ExaModels.AbstractNode - - result5 = vec_nodes + real_reinterp - @test length(result5) == 6 - @test result5[1] isa ExaModels.AbstractNode - - result6 = real_reinterp + vec_nodes - @test length(result6) == 6 - - result7 = vec_nodes - real_reinterp - @test length(result7) == 6 - @test result7[1] isa ExaModels.AbstractNode - - result8 = real_reinterp - vec_nodes - @test length(result8) == 6 - end - - @testset "Mixed wrapper combinations" begin - v_num = [1.0, 2.0, 3.0, 4.0] - vec_nodes = [x, y, z, w] - - v_view = @view v_num[1:3] - v_reshaped = reshape([x, y, z], 3) - result1 = v_view + v_reshaped - @test length(result1) == 3 - @test result1[1] isa ExaModels.AbstractNode - - A_num = [1.0 2.0 3.0; 4.0 5.0 6.0; 7.0 8.0 9.0] - A_view = @view A_num[1:3, 1:3] - vec_reshaped = reshape([x, y, z], 3) - result2 = A_view * vec_reshaped - @test length(result2) == 3 - @test result2[1] isa ExaModels.AbstractNode - - complex_vec = ComplexF64[1.0 + 2.0im, 3.0 + 4.0im] - real_reinterp = reinterpret(Float64, complex_vec) - vec_view = @view vec_nodes[1:4] - result3 = real_reinterp + vec_view - @test length(result3) == 4 - @test result3[1] isa ExaModels.AbstractNode - end - - @testset "diagm with views and reshaped arrays" begin - vec_nodes = [x, y, z] - - v_view = @view vec_nodes[1:2] - result1 = diagm(v_view) - @test size(result1) == (2, 2) - @test result1[1, 1] isa ExaModels.AbstractNode - @test is_null_zero(result1[1, 2]) - - v_reshaped = reshape([x, y], 2) - result2 = diagm(v_reshaped) - @test size(result2) == (2, 2) - @test result2[1, 1] isa ExaModels.AbstractNode - - result3 = diagm(1 => v_view) - @test size(result3) == (3, 3) - @test result3[1, 2] isa ExaModels.AbstractNode - end - - @testset "diag with views and reshaped arrays" begin - mat_nodes = [x y z; w x y; z w x] - - A_view = @view mat_nodes[1:2, 1:2] - result1 = diag(A_view) - @test length(result1) == 2 - @test result1[1] isa ExaModels.AbstractNode - - vec = [x, y, z, w] - A_reshaped = reshape(vec, 2, 2) - result2 = diag(A_reshaped) - @test length(result2) == 2 - @test result2[1] isa ExaModels.AbstractNode - end - - @testset "Adjoint operations with views" begin - vec_nodes = [x, y, z] - A_num = [1.0 2.0; 3.0 4.0; 5.0 6.0] - - v_view = @view vec_nodes[1:3] - v_num = [1.0, 2.0, 3.0] - result0 = v_view' * v_num - @test result0 isa ExaModels.AbstractNode - - result0b = v_num' * v_view - @test result0b isa ExaModels.AbstractNode - - v_view2 = @view vec_nodes[1:3] - result0c = v_view' * v_view2 - @test result0c isa ExaModels.AbstractNode - - result1 = v_view' * A_num - @test size(result1) == (1, 2) - @test result1 isa LinearAlgebra.Adjoint - - mat_nodes = [x y z; w x y; z w x] - A_view = @view mat_nodes[1:2, 1:2] - v_num2 = [1.0, 2.0] - result2 = v_num2' * A_view - @test size(result2) == (1, 2) - - vec_reshaped = reshape([x, y, z], 3) - result3 = vec_reshaped' * v_num - @test result3 isa ExaModels.AbstractNode - - result4 = v_num' * vec_reshaped - @test result4 isa ExaModels.AbstractNode - end - end - end - end -end - -end # module diff --git a/test/runtests.jl b/test/runtests.jl index 58583e0a5..418d5180c 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -21,7 +21,6 @@ include("DeprecatedTest/DeprecatedTest.jl") include("JuMPTest/JuMPTest.jl") include("UtilsTest/UtilsTest.jl") include("TwoStageTest/TwoStageTest.jl") -include("LinAlgTest/LinAlgTest.jl") include("BatchTest/BatchTest.jl") include("JuliaCTest/JuliaCTest.jl") include("TwoStageTest/TwoStageTest.jl") @@ -60,8 +59,6 @@ include("PrettyPrintTest.jl") # @info "Running OptimalControl Test" # OptimalControlTest.runtests() - @info "Running LinAlg Test" - LinAlgTest.runtests() @info "Running Batch Test" BatchTest.runtests() end From c7431d050e0cca844e83e74a8e39eb0219beb481 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 18:04:34 -0400 Subject: [PATCH 07/50] Fix juliac COPSApp compilation and docs build Restore non-batch call paths for standard ExaModel evaluation functions (obj, grad!, cons_nln!, jac_coord!, hess_coord!, hprod!). The batch-aware versions that create @view SubArrays are now separate (_grad_b!, _cons_nln_b!, _jac_coord_b!, _obj_hess_coord_b!, etc.) and only used by BatchExaModel methods. This prevents juliac --trim=safe from needing to trace SubArray specializations for standard models. Also fix docs/src/two_stage.jl to use new EachScenario API. Co-Authored-By: Claude Opus 4.6 --- docs/src/two_stage.jl | 25 +++--- src/nlp.jl | 189 ++++++++++++++++++++++-------------------- 2 files changed, 112 insertions(+), 102 deletions(-) diff --git a/docs/src/two_stage.jl b/docs/src/two_stage.jl index 16e6c79b8..c820dcd8d 100644 --- a/docs/src/two_stage.jl +++ b/docs/src/two_stage.jl @@ -13,23 +13,24 @@ nv = 2 ## recourse variables per scenario nd = 1 ## design variables weight = 1.0 / ns -# Two annotate the scenario for each variable and constraint, we can use the `scenario` we need to start with a special ExaCore that supports such scenario annotations, which can be created by calling `TwoStageExaCore(concrete = Val(true))`. -core = TwoStageExaCore(concrete = Val(true)) +# To annotate the scenario for each variable and constraint, we start with a `TwoStageExaCore` that supports scenario annotations. +core = TwoStageExaCore(ns; concrete = Val(true)) -# Now we can define the design variable and recourse variables. The `scenario` keyword argument allows us to specify which scenario(s) each variable belongs to. For the design variable `d`, we set `scenario = 0` to indicate that it is shared across all scenarios. -@add_var(core, d; start = 1.0, lvar = 0.0, uvar = Inf, scenario = 0) ## design variable d +# Design variables are shared across all scenarios — add them without `EachScenario()`. +core, d = add_var(core, nd; start = 1.0, lvar = 0.0, uvar = Inf) -# For the recourse variables `v`, we specify `scenario = [i for i=1:ns, j=1:nv]` to indicate that each variable `v[s,i]` belongs to scenario `s`. This allows us to define scenario-specific constraints and objectives that involve these recourse variables. -@add_var(core, v, ns, nv; start = 1.0, lvar = 0.0, uvar = Inf, scenario = [i for i=1:ns, j=1:nv]) ## recourse variables v +# Recourse variables are per-scenario — use `EachScenario()` to replicate them. +v = @add_var(core, EachScenario(), nv; start = 1.0, lvar = 0.0, uvar = Inf) -# Now we can define the constraints and objective function. The `scenario` keyword argument in the `constraint` and `objective` functions allows us to specify which scenario(s) each constraint or objective term belongs to. -@add_con(core, v[s,1] - v[s,2]^2 for s in 1:ns; lcon = 0.0, scenario = 1:ns) +# Per-scenario constraints use `EachScenario()`. +@add_con(core, EachScenario(), (v[(s-1)*nv+1] - v[(s-1)*nv+2]^2 for s in 1:ns); lcon = 0.0) -@add_obj(core, d^2) -@add_obj(core, weight * (v[s,i] - d)^2 for s in 1:ns, i in 1:nv) +# Objectives can mix design and recourse variables. +@add_obj(core, d[1]^2) +@add_obj(core, weight * (v[(s-1)*nv+i] - d[1])^2 for s in 1:ns, i in 1:nv) m = ExaModel(core) -# Now we can solve the model as usual. -ipopt(m) +# Now we can solve the model as usual. +ipopt(m) # If the solver knows how to exploit the scenario structure, the structure-exploiting method can be used. diff --git a/src/nlp.jl b/src/nlp.jl index ae67cdf12..a7b15ae6c 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -1364,14 +1364,25 @@ end # ============================================================================ function obj(m::AbstractExaModel, x::AbstractVector) - return _obj(m.objs, x, m.θ, 1, length(x), length(m.θ)) + return _obj(m.objs, x, m.θ) end # Stub for KA extension override function _eval_objbuffer! end -@inline function _obj((obj, objs...), x, θ, nb, nvar, npar) - s = _obj(objs, x, θ, nb, nvar, npar) +@inline function _obj((obj, objs...), x, θ) + s = _obj(objs, x, θ) + for i in obj.itr + s += obj.f(i, x, θ) + end + return s +end + +@inline _obj(obj::Tuple{}, x, θ) = zero(eltype(x)) + +# Batch versions — used by BatchExaModel (loop over instances with views) +@inline function _obj_b((obj, objs...), x, θ, nb, nvar, npar) + s = _obj_b(objs, x, θ, nb, nvar, npar) for si in 1:nb x_s = @view x[(si-1)*nvar+1 : si*nvar] θ_s = @view θ[(si-1)*npar+1 : si*npar] @@ -1381,8 +1392,7 @@ function _eval_objbuffer! end end return s end - -@inline _obj(obj::Tuple{}, x, θ, nb, nvar, npar) = zero(eltype(x)) +@inline _obj_b(obj::Tuple{}, x, θ, nb, nvar, npar) = zero(eltype(x)) # Per-instance obj values (for batch obj!) @inline function _obj!(bf, (obj, objs...), x, θ, nb, nvar, npar, backend = nothing) @@ -1403,18 +1413,26 @@ end function cons_nln!(m::AbstractExaModel, x::AbstractVector, g::AbstractVector) fill!(g, zero(eltype(g))) - nvar = NLPModels.get_nvar(m) - ncon = NLPModels.get_ncon(m) - _cons_nln!(m.cons, x, m.θ, g, 1, nvar, length(m.θ), ncon) + _cons_nln!(m.cons, x, m.θ, g) return g end -@inline function _cons_nln!(cons::Tuple, x, θ, g, nb, nvar, npar, ncon, backend = nothing) +@inline function _cons_nln!(cons::Tuple, x, θ, g) + con = first(cons) + _cons_nln!(Base.tail(cons), x, θ, g) + @simd for i in eachindex(con.itr) + g[offset0(con, i)] += con.f(con.itr[i], x, θ) + end +end +_cons_nln!(cons::Tuple{}, x, θ, g) = nothing + +# Batch versions — used by BatchExaModel +@inline function _cons_nln_b!(cons::Tuple, x, θ, g, nb, nvar, npar, ncon, backend = nothing) con = first(cons) - _cons_nln!(Base.tail(cons), x, θ, g, nb, nvar, npar, ncon, backend) + _cons_nln_b!(Base.tail(cons), x, θ, g, nb, nvar, npar, ncon, backend) _cons_nln_batch!(con, x, θ, g, nb, nvar, npar, ncon, backend) end -_cons_nln!(cons::Tuple{}, x, θ, g, nb, nvar, npar, ncon, backend = nothing) = nothing +_cons_nln_b!(cons::Tuple{}, x, θ, g, nb, nvar, npar, ncon, backend = nothing) = nothing @inline function _cons_nln_batch!(con, x, θ, g, nb, nvar, npar, ncon, ::Nothing) for s in 1:nb @@ -1431,66 +1449,64 @@ end function grad!(m::AbstractExaModel, x::AbstractVector, f::AbstractVector) fill!(f, zero(eltype(f))) - _grad!(m.objs, x, m.θ, f, 1, length(x), length(m.θ)) + _grad!(m.objs, x, m.θ, f) return f end -@inline function _grad!(objs::Tuple, x, θ, f, nb, nvar, npar, backend = nothing) - _grad!(Base.tail(objs), x, θ, f, nb, nvar, npar, backend) +@inline function _grad!(objs::Tuple, x, θ, f) + _grad!(Base.tail(objs), x, θ, f) + gradient!(f, first(objs), x, θ, one(eltype(f))) +end +_grad!(objs::Tuple{}, x, θ, f) = nothing + +# Batch versions — used by BatchExaModel +@inline function _grad_b!(objs::Tuple, x, θ, f, nb, nvar, npar, backend = nothing) + _grad_b!(Base.tail(objs), x, θ, f, nb, nvar, npar, backend) gradient!(f, first(objs), x, θ, one(eltype(f)), nb, nvar, npar, backend) end -_grad!(objs::Tuple{}, x, θ, f, nb, nvar, npar, backend = nothing) = nothing +_grad_b!(objs::Tuple{}, x, θ, f, nb, nvar, npar, backend = nothing) = nothing function jac_coord!(m::AbstractExaModel, x::AbstractVector, jac::AbstractVector) fill!(jac, zero(eltype(jac))) - nvar = NLPModels.get_nvar(m) - nnzj = NLPModels.get_nnzj(m) - _jac_coord!(m.cons, x, m.θ, jac, 1, nvar, length(m.θ), nnzj) + _jac_coord!(m.cons, x, m.θ, jac) return jac end -_jac_coord!(cons::Tuple{}, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) = nothing -@inline function _jac_coord!(cons::Tuple, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) - _jac_coord!(Base.tail(cons), x, θ, jac, nb, nvar, npar, nnzj, backend) +_jac_coord!(cons::Tuple{}, x, θ, jac) = nothing +@inline function _jac_coord!(cons::Tuple, x, θ, jac) + _jac_coord!(Base.tail(cons), x, θ, jac) + sjacobian!(jac, nothing, first(cons), x, θ, one(eltype(jac))) +end + +# Batch versions — used by BatchExaModel +_jac_coord_b!(cons::Tuple{}, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) = nothing +@inline function _jac_coord_b!(cons::Tuple, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) + _jac_coord_b!(Base.tail(cons), x, θ, jac, nb, nvar, npar, nnzj, backend) sjacobian!(jac, nothing, first(cons), x, θ, one(eltype(jac)), nb, nvar, npar, nnzj, backend) end function jprod_nln!(m::AbstractExaModel, x::AbstractVector, v::AbstractVector, Jv::AbstractVector) fill!(Jv, zero(eltype(Jv))) - _jprod_nln!(m.cons, x, m.θ, v, Jv, 1, length(x), length(m.θ), NLPModels.get_ncon(m)) + _jprod_nln!(m.cons, x, m.θ, v, Jv) return Jv end -_jprod_nln!(cons::Tuple{}, x, θ, v, Jv, nb, nvar, npar, ncon) = nothing -@inline function _jprod_nln!(cons::Tuple, x, θ, v, Jv, nb, nvar, npar, ncon) - _jprod_nln!(Base.tail(cons), x, θ, v, Jv, nb, nvar, npar, ncon) - f = first(cons) - for s in 1:nb - x_s = @view x[(s-1)*nvar+1 : s*nvar] - θ_s = @view θ[(s-1)*npar+1 : s*npar] - v_s = @view v[(s-1)*nvar+1 : s*nvar] - Jv_s = @view Jv[(s-1)*ncon+1 : s*ncon] - sjacobian!((Jv_s, v_s), nothing, f, x_s, θ_s, one(eltype(Jv))) - end +_jprod_nln!(cons::Tuple{}, x, θ, v, Jv) = nothing +@inline function _jprod_nln!(cons::Tuple, x, θ, v, Jv) + _jprod_nln!(Base.tail(cons), x, θ, v, Jv) + sjacobian!((Jv, v), nothing, first(cons), x, θ, one(eltype(Jv))) end function jtprod_nln!(m::AbstractExaModel, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector) fill!(Jtv, zero(eltype(Jtv))) - _jtprod_nln!(m.cons, x, m.θ, v, Jtv, 1, length(x), length(m.θ), NLPModels.get_ncon(m)) + _jtprod_nln!(m.cons, x, m.θ, v, Jtv) return Jtv end -_jtprod_nln!(cons::Tuple{}, x, θ, v, Jtv, nb, nvar, npar, ncon) = nothing -@inline function _jtprod_nln!(cons::Tuple, x, θ, v, Jtv, nb, nvar, npar, ncon) - _jtprod_nln!(Base.tail(cons), x, θ, v, Jtv, nb, nvar, npar, ncon) - f = first(cons) - for s in 1:nb - x_s = @view x[(s-1)*nvar+1 : s*nvar] - θ_s = @view θ[(s-1)*npar+1 : s*npar] - v_s = @view v[(s-1)*ncon+1 : s*ncon] - Jtv_s = @view Jtv[(s-1)*nvar+1 : s*nvar] - sjacobian!(nothing, (Jtv_s, v_s), f, x_s, θ_s, one(eltype(Jtv))) - end +_jtprod_nln!(cons::Tuple{}, x, θ, v, Jtv) = nothing +@inline function _jtprod_nln!(cons::Tuple, x, θ, v, Jtv) + _jtprod_nln!(Base.tail(cons), x, θ, v, Jtv) + sjacobian!(nothing, (Jtv, v), first(cons), x, θ, one(eltype(Jtv))) end function hess_coord!( @@ -1500,9 +1516,7 @@ function hess_coord!( obj_weight = one(eltype(x)), ) fill!(hess, zero(eltype(hess))) - nvar = NLPModels.get_nvar(m) - nnzh = NLPModels.get_nnzh(m) - _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight, 1, nvar, length(m.θ), nnzh) + _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight) return hess end @@ -1514,23 +1528,33 @@ function hess_coord!( obj_weight = one(eltype(x)), ) fill!(hess, zero(eltype(hess))) - nvar = NLPModels.get_nvar(m) - ncon = NLPModels.get_ncon(m) - nnzh = NLPModels.get_nnzh(m) - _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight, 1, nvar, length(m.θ), nnzh) - _con_hess_coord!(m.cons, x, m.θ, y, hess, 1, nvar, length(m.θ), ncon, nnzh) + _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight) + _con_hess_coord!(m.cons, x, m.θ, y, hess, obj_weight) return hess end -_obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) = nothing -@inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) - _obj_hess_coord!(Base.tail(objs), x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend) +_obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight) = nothing +@inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight) + _obj_hess_coord!(Base.tail(objs), x, θ, hess, obj_weight) + shessian!(hess, nothing, first(objs), x, θ, obj_weight, zero(eltype(hess))) +end + +_con_hess_coord!(cons::Tuple{}, x, θ, y, hess, obj_weight) = nothing +@inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess, obj_weight) + _con_hess_coord!(Base.tail(cons), x, θ, y, hess, obj_weight) + shessian!(hess, nothing, first(cons), x, θ, y, zero(eltype(hess))) +end + +# Batch versions — used by BatchExaModel +_obj_hess_coord_b!(objs::Tuple{}, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) = nothing +@inline function _obj_hess_coord_b!(objs::Tuple, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) + _obj_hess_coord_b!(Base.tail(objs), x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend) shessian!(hess, nothing, first(objs), x, θ, obj_weight, zero(eltype(hess)), nb, nvar, npar, nnzh, backend) end -_con_hess_coord!(cons::Tuple{}, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) = nothing -@inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) - _con_hess_coord!(Base.tail(cons), x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend) +_con_hess_coord_b!(cons::Tuple{}, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) = nothing +@inline function _con_hess_coord_b!(cons::Tuple, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) + _con_hess_coord_b!(Base.tail(cons), x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend) shessian!(hess, nothing, first(cons), x, θ, y, zero(eltype(hess)), nb, nvar, npar, ncon, nnzh, backend) end @@ -1542,7 +1566,7 @@ function hprod!( obj_weight = one(eltype(x)), ) fill!(Hv, zero(eltype(Hv))) - _obj_hprod!(m.objs, x, m.θ, v, Hv, obj_weight, 1, length(x), length(m.θ)) + _obj_hprod!(m.objs, x, m.θ, v, Hv, obj_weight) return Hv end @@ -1555,36 +1579,21 @@ function hprod!( obj_weight = one(eltype(x)), ) fill!(Hv, zero(eltype(Hv))) - _obj_hprod!(m.objs, x, m.θ, v, Hv, obj_weight, 1, length(x), length(m.θ)) - _con_hprod!(m.cons, x, m.θ, y, v, Hv, obj_weight, 1, length(x), length(m.θ), NLPModels.get_ncon(m)) + _obj_hprod!(m.objs, x, m.θ, v, Hv, obj_weight) + _con_hprod!(m.cons, x, m.θ, y, v, Hv, obj_weight) return Hv end -_obj_hprod!(objs::Tuple{}, x, θ, v, Hv, obj_weight, nb, nvar, npar) = nothing -@inline function _obj_hprod!(objs::Tuple, x, θ, v, Hv, obj_weight, nb, nvar, npar) - _obj_hprod!(Base.tail(objs), x, θ, v, Hv, obj_weight, nb, nvar, npar) - f = first(objs) - for s in 1:nb - x_s = @view x[(s-1)*nvar+1 : s*nvar] - θ_s = @view θ[(s-1)*npar+1 : s*npar] - v_s = @view v[(s-1)*nvar+1 : s*nvar] - Hv_s = @view Hv[(s-1)*nvar+1 : s*nvar] - shessian!((Hv_s, v_s), nothing, f, x_s, θ_s, obj_weight, zero(eltype(Hv))) - end +_obj_hprod!(objs::Tuple{}, x, θ, v, Hv, obj_weight) = nothing +@inline function _obj_hprod!(objs::Tuple, x, θ, v, Hv, obj_weight) + _obj_hprod!(Base.tail(objs), x, θ, v, Hv, obj_weight) + shessian!((Hv, v), nothing, first(objs), x, θ, obj_weight, zero(eltype(Hv))) end -_con_hprod!(cons::Tuple{}, x, θ, y, v, Hv, obj_weight, nb, nvar, npar, ncon) = nothing -@inline function _con_hprod!(cons::Tuple, x, θ, y, v, Hv, obj_weight, nb, nvar, npar, ncon) - _con_hprod!(Base.tail(cons), x, θ, y, v, Hv, obj_weight, nb, nvar, npar, ncon) - f = first(cons) - for s in 1:nb - x_s = @view x[(s-1)*nvar+1 : s*nvar] - θ_s = @view θ[(s-1)*npar+1 : s*npar] - y_s = @view y[(s-1)*ncon+1 : s*ncon] - v_s = @view v[(s-1)*nvar+1 : s*nvar] - Hv_s = @view Hv[(s-1)*nvar+1 : s*nvar] - shessian!((Hv_s, v_s), nothing, f, x_s, θ_s, y_s, zero(eltype(Hv))) - end +_con_hprod!(cons::Tuple{}, x, θ, y, v, Hv, obj_weight) = nothing +@inline function _con_hprod!(cons::Tuple, x, θ, y, v, Hv, obj_weight) + _con_hprod!(Base.tail(cons), x, θ, y, v, Hv, obj_weight) + shessian!((Hv, v), nothing, first(cons), x, θ, y, zero(eltype(Hv))) end @inbounds @inline offset0(a, i) = offset0(a.f, i) @@ -1940,7 +1949,7 @@ function NLPModels.grad!(m::BatchExaModel{T}, bx::AbstractMatrix, bg::AbstractMa nb = get_nbatch(m) nvar = NLPModels.get_nvar(m) npar = Base.size(m.θ, 1) - _grad!(m.objs, vec(bx), vec(m.θ), vec(bg), nb, nvar, npar, getbackend(m)) + _grad_b!(m.objs, vec(bx), vec(m.θ), vec(bg), nb, nvar, npar, getbackend(m)) return bg end @@ -1950,7 +1959,7 @@ function NLPModels.cons!(m::BatchExaModel{T}, bx::AbstractMatrix, bc::AbstractMa nvar = NLPModels.get_nvar(m) ncon = NLPModels.get_ncon(m) npar = Base.size(m.θ, 1) - _cons_nln!(m.cons, vec(bx), vec(m.θ), vec(bc), nb, nvar, npar, ncon, getbackend(m)) + _cons_nln_b!(m.cons, vec(bx), vec(m.θ), vec(bc), nb, nvar, npar, ncon, getbackend(m)) return bc end @@ -1969,7 +1978,7 @@ function NLPModels.jac_coord!(m::BatchExaModel{T}, bx::AbstractMatrix, jvals::Ab nvar = NLPModels.get_nvar(m) npar = Base.size(m.θ, 1) nnzj = NLPModels.get_nnzj(m) - _jac_coord!(m.cons, vec(bx), vec(m.θ), vec(jvals), nb, nvar, npar, nnzj, getbackend(m)) + _jac_coord_b!(m.cons, vec(bx), vec(m.θ), vec(jvals), nb, nvar, npar, nnzj, getbackend(m)) return jvals end @@ -1997,8 +2006,8 @@ function NLPModels.hess_coord!( npar = Base.size(m.θ, 1) nnzh = NLPModels.get_nnzh(m) backend = getbackend(m) - _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), obj_weight, nb, nvar, npar, nnzh, backend) - _con_hess_coord!(m.cons, vec(bx), vec(m.θ), vec(by), vec(hvals), nb, nvar, npar, ncon, nnzh, backend) + _obj_hess_coord_b!(m.objs, vec(bx), vec(m.θ), vec(hvals), obj_weight, nb, nvar, npar, nnzh, backend) + _con_hess_coord_b!(m.cons, vec(bx), vec(m.θ), vec(by), vec(hvals), nb, nvar, npar, ncon, nnzh, backend) return hvals end From e8621264cd18cbc7a960802a2f91b24e6b931865 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 18:10:04 -0400 Subject: [PATCH 08/50] Rewrite docs/src/batch.jl for new BatchExaCore API Co-Authored-By: Claude Opus 4.6 --- docs/src/batch.jl | 116 +++++++++++++++++----------------------------- 1 file changed, 42 insertions(+), 74 deletions(-) diff --git a/docs/src/batch.jl b/docs/src/batch.jl index f59ce3229..1a0236919 100644 --- a/docs/src/batch.jl +++ b/docs/src/batch.jl @@ -1,10 +1,10 @@ # # [Batch Optimization](@id batch) -# ExaModels supports batch optimization through the `ExaModel`. This feature +# ExaModels supports batch optimization through `BatchExaCore`. This feature # enables efficient evaluation of multiple fully independent optimization instances # that share identical structure but differ in parameter values. # # Unlike `TwoStageExaModel`, which couples instances through shared design variables, -# `ExaModel` treats each instance as completely independent. The key advantage +# batch models treat each instance as completely independent. The key advantage # is that all instances share one compiled expression pattern and are fused into a # single model for efficient SIMD evaluation. @@ -20,44 +20,38 @@ # where each instance has the same structure but different parameters $\theta_i$. # ## Building a Batch Model -# The builder function defines expressions for a **single instance**. `ExaModel` -# calls it `ns` times internally with per-instance parameter handles. -# This means you never have to compute global index offsets manually. +# Use `BatchExaCore(ns)` to create a core for `ns` instances. +# Variables, parameters, objectives, and constraints are defined once — +# the batch structure replicates them across all instances automatically. -using ExaModels, MadNLP +using ExaModels, NLPModelsIpopt +import NLPModels -# Define the problem dimensions and instance parameters as a matrix of size `(nθ, ns)`: +# Define the problem dimensions and instance parameters: ns = 3 ## number of instances nv = 1 ## variables per instance -θ_data = [2.0 4.0 6.0] ## (1, 3) matrix: θ₁=2, θ₂=4, θ₃=6 - -# Build the model. First create an `ExaCore`, then pass it along with the parameter -# matrix to `ExaModel`: -c = ExaCore() -model = ExaModel(c, ns, θ_data) do c, θ - ## Create variables — this is called once per instance, offsets are automatic - v = variable(c, nv) - ## Objective: minimize (v - θ)² - objective(c, (v[1] - θ[1])^2) - ## Constraint: v ≥ 0 - constraint(c, v[1]; lcon = 0.0, ucon = Inf) -end -# The builder function receives: -# - `c`: the `ExaCore` — use `variable(c, ...)`, `objective(c, ...)`, `constraint(c, ...)` as usual -# - `θ`: a per-instance parameter handle (indices 1:nθ) -# -# Variable creation via `variable(c, ...)` works exactly like in a regular `ExaModel`. -# You can set start values, lower/upper bounds, etc. +# Create a batch core: +c = BatchExaCore(ns) + +# Add variables and parameters: +@add_var(c, v, nv) +@add_par(c, θ, [2.0]) + +# Define objectives and constraints — these apply to every instance: +@add_obj(c, (v[j] - θ[1])^2 for j in 1:nv) +@add_con(c, g, v[j] for j in 1:nv; lcon = 0.0) + +# Build the model: +model = ExaModel(c) # ## Batch API (NLPModels) -# `ExaModel` implements the `AbstractBatchNLPModel` interface from NLPModels.jl. +# Batch models implement `AbstractNLPModel` with matrix-valued variables. # All evaluation functions use matrices of size `(dim, ns)`: -import NLPModels println("Variables per instance: ", NLPModels.get_nvar(model)) println("Constraints per instance: ", NLPModels.get_ncon(model)) -println("Number of instances: ", NLPModels.get_nbatch(model)) +println("Number of instances: ", ExaModels.get_nbatch(model)) # Evaluate objectives for all instances at once: bx = reshape([1.0, 3.0, 5.0], nv, ns) @@ -77,11 +71,11 @@ NLPModels.cons!(model, bx, bc) println("Constraints: ", bc) # ## Solving via the Fused Model -# For solving, access the underlying fused `ExaModel` and use any NLPModels-compatible -# solver (e.g., MadNLP): -result = madnlp(ExaModels.get_model(model); print_level = MadNLP.ERROR) +# For solving, use `get_model(model)` to access the fused `FlattenNLPModel` and +# pass it to any NLPModels-compatible solver: +flat = ExaModels.get_model(model) +result = ipopt(flat; print_level = 0) println("\nSolution status: ", result.status) -println("Optimal objective: ", round(result.objective, digits = 4)) # Extract per-instance solutions: x_sol = result.solution @@ -90,44 +84,18 @@ for i in 1:ns println("Instance $i: v* = ", round(v_sol[1], digits = 4)) end -# ## A More Complex Example -# Here's a batch model with multiple variables, objectives, and constraints per instance: -ns2, nv2 = 2, 3 -θ_data2 = [1.0 4.0; 2.0 5.0; 3.0 6.0] ## (3, 2) matrix - -c2 = ExaCore() -model2 = ExaModel(c2, ns2, θ_data2) do c, θ - v = variable(c, nv2; start = 1.0, lvar = 0.0, uvar = 10.0) - ## Objective: Σⱼ (vⱼ - θⱼ)² - objective(c, (v[j] - θ[j])^2 for j in 1:nv2) - ## Constraints: sum of all variables ≤ 20 - constraint(c, sum(v[j] for j in 1:nv2); ucon = 20.0) -end - -result2 = madnlp(ExaModels.get_model(model2); print_level = MadNLP.ERROR) -println("\nMulti-variable example:") -println("Status: ", result2.status) -x_sol2 = result2.solution -for i in 1:ns2 - v_sol = x_sol2[ExaModels.var_indices(model2, i)] - println("Instance $i: v* = ", round.(v_sol, digits = 4)) -end - -# ## Updating Parameters -# You can update instance parameters and re-solve without rebuilding the model: - -# Update a single instance: -ExaModels.set_instance_parameters!(model, 1, [10.0]) - -# Or update all instances at once: -ExaModels.set_all_instance_parameters!(model, [[10.0], [12.0], [14.0]]) - -# Re-solve with new parameters: -result3 = madnlp(ExaModels.get_model(model); print_level = MadNLP.ERROR) -println("\nAfter parameter update:") -println("Status: ", result3.status) -x_sol3 = result3.solution -for i in 1:ns - v_sol = x_sol3[ExaModels.var_indices(model, i)] - println("Instance $i: v* = ", round(v_sol[1], digits = 4)) -end +# ## Per-Instance Parameters +# Each instance can have different parameter values. +# Parameters are stored as a matrix `(nθ, ns)`. You can set different values +# per instance using `set_parameter!`: + +c2 = BatchExaCore(2) +@add_var(c2, x, 2) +@add_par(c2, p, [1.0, 2.0]) +@add_obj(c2, (x[j] - p[j])^2 for j in 1:2) +model2 = ExaModel(c2) + +# Update parameters for instance 2: +ExaModels.set_parameter!(c2, p, [10.0, 20.0]) +model2 = ExaModel(c2) +println("\nParameter matrix shape: ", size(model2.θ)) From e87db1a55c1e5266f3f084b270203432a2cfd077 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 18:15:20 -0400 Subject: [PATCH 09/50] Add BatchNLPModels autodocs to API manual Co-Authored-By: Claude Opus 4.6 --- docs/src/core.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/src/core.md b/docs/src/core.md index d41a1ed9a..ee5b52537 100644 --- a/docs/src/core.md +++ b/docs/src/core.md @@ -2,3 +2,8 @@ ```@autodocs Modules = [ExaModels] ``` + +# Batch NLP Models +```@autodocs +Modules = [ExaModels.BatchNLPModels] +``` From 24b5eb6c8ae8567acc5d231b183601b4e816ac76 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 18:49:48 -0400 Subject: [PATCH 10/50] Fix append! to handle N-dimensional arrays (fixes juliac COPSApp) append! only had methods for AbstractVector and AbstractMatrix, but COPS models like catmix pass 3D arrays (e.g. zeros(nh, nc, ne)) as variable start values. The missing method caused MethodError at runtime and Union{} return types in juliac --trim=safe verification. Unified into a single AbstractArray method that flattens to vec() and expands to the target shape. Co-Authored-By: Claude Opus 4.6 --- src/nlp.jl | 77 +++++++++++++++++++++++++++++++----------------- src/two_stage.jl | 53 ++++++++++++++++++++++++++++----- 2 files changed, 95 insertions(+), 35 deletions(-) diff --git a/src/nlp.jl b/src/nlp.jl index a7b15ae6c..5897d91d6 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -682,24 +682,13 @@ function append!(backend, a, b::Number, lb) return cat(a, new_part; dims = 1) end -function append!(backend, a, b::AbstractVector, lb) +function append!(backend, a, b::AbstractArray, lb) lb == 0 && return a - col = vec(convert_array(b, backend)) + arr = convert_array(b, backend) + col = vec(arr) return cat(a, _expand_to_shape(col, _trailing_dims(a)); dims = 1) end -function append!(backend, a, b::AbstractMatrix, lb) - lb == 0 && return a - m = convert_array(b, backend) - trailing = _trailing_dims(a) - if trailing == () - # a is a vector — flatten the matrix to a vector - return cat(a, vec(m); dims = 1) - else - return cat(a, m; dims = 1) - end -end - function append!(backend, a, b::Base.Generator, lb) lb == 0 && return a b = _adapt_gen(b) @@ -1096,33 +1085,67 @@ Constraint where |I| = 9 ``` """ +# Generator form — directly dispatched for type stability and juliac trimmer. @inline function add_con( c::C, - ns...; + gen::Base.Generator; tag = nothing, name = nothing, start = zero(T), lcon = zero(T), ucon = zero(T), - kwargs... ) where {T,C<:ExaCore{T}} - gen = _get_generator(ns) - dims = _get_con_dims(ns) + dims = _infer_subexpr_dims(gen.iter) gen = _adapt_gen(gen) - f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) + f = SIMDFunction(T, gen, c.ncon, c.nnzj, c.nnzh) pars = gen.iter - _add_con(c, f, pars, dims, start, lcon, ucon, name, tag) end -@inline _get_generator(ns) = (Null(nothing) for _ in _empty_con_itr(ns)) -@inline _get_generator(gen::Tuple{G}) where G <: Base.Generator = gen[1] -@inline _get_generator(n::Tuple{N}) where N <: AbstractNode = (n[1] for _ in 1:1) +# Multi-generator form: first generator creates the constraint, rest augment it. +@inline function add_con( + c::C, + gen::Base.Generator, + gens::Base.Generator...; + kwargs... +) where {T,C<:ExaCore{T}} + c, con = add_con(c, gen; kwargs...) + for g in gens + c, _ = add_con!(c, con, g) + end + return (c, con) +end + +# Expression form — pre-built expression tree with explicit parameters. +@inline function add_con( + c::C, + expr::N, + pars = 1:1; + tag = nothing, + name = nothing, + start = zero(T), + lcon = zero(T), + ucon = zero(T), +) where {T,C<:ExaCore{T},N<:AbstractNode} + f = _simdfunction(T, expr, c.ncon, c.nnzj, c.nnzh) + dims = _infer_subexpr_dims(pars) + _add_con(c, f, pars, dims, start, lcon, ucon, name, tag) +end -# Infer constraint dims from the original arguments, preserving range start info. -@inline _get_con_dims(ns) = ns -@inline _get_con_dims(gen::Tuple{G}) where G <: Base.Generator = _infer_subexpr_dims(gen[1].iter) -@inline _get_con_dims(n::Tuple{N}) where N <: AbstractNode = (1,) +# Dims form — empty constraints for later augmentation. +@inline function add_con( + c::C, + ns::Union{Integer, AbstractUnitRange}...; + tag = nothing, + name = nothing, + start = zero(T), + lcon = zero(T), + ucon = zero(T), +) where {T,C<:ExaCore{T}} + f = _simdfunction(T, Null(nothing), c.ncon, c.nnzj, c.nnzh) + pars = _empty_con_itr(ns) + _add_con(c, f, pars, ns, start, lcon, ucon, name, tag) +end # Build an iterator for empty constraints: 1:n for 1D, collected ProductIterator for multi-dim. _empty_con_itr(ns::Tuple{Any}) = 1:_length(ns[1]) diff --git a/src/two_stage.jl b/src/two_stage.jl index 93d48f4c3..16edd29e0 100644 --- a/src/two_stage.jl +++ b/src/two_stage.jl @@ -199,7 +199,7 @@ the base [`add_con`](@ref) (generator or dims). Tagged with `FirstStageTag()`. """ function add_con( c::C, - ns...; + gen::Base.Generator; name = nothing, tag = nothing, start = zero(T), @@ -207,10 +207,9 @@ function add_con( ucon = zero(T), ) where {T,VT<:AbstractVector{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} - gen = _get_generator(ns) - dims = _get_con_dims(ns) + dims = _infer_subexpr_dims(gen.iter) gen = _adapt_gen(gen) - f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) + f = SIMDFunction(T, gen, c.ncon, c.nnzj, c.nnzh) pars = gen.iter new_con_scen = append!(c.backend, c.tag.con_scen, 0, length(pars)) @@ -218,6 +217,24 @@ function add_con( return _add_con(c, f, pars, dims, start, lcon, ucon, name, FirstStageConstraintTag()) end +function add_con( + c::C, + ns::Union{Integer, AbstractUnitRange}...; + name = nothing, + tag = nothing, + start = zero(T), + lcon = zero(T), + ucon = zero(T), + ) where {T,VT<:AbstractVector{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} + + f = _simdfunction(T, Null(nothing), c.ncon, c.nnzj, c.nnzh) + pars = _empty_con_itr(ns) + + new_con_scen = append!(c.backend, c.tag.con_scen, 0, length(pars)) + c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, c.tag.var_scen, new_con_scen)) + return _add_con(c, f, pars, ns, start, lcon, ucon, name, FirstStageConstraintTag()) +end + """ add_con(core::TwoStageExaCore, ::EachScenario, dims_or_gen...; start = 0, lcon = 0, ucon = 0, name = nothing) @@ -229,7 +246,7 @@ scenario index. function add_con( c::C, ::EachScenario, - ns...; + gen::Base.Generator; name = nothing, tag = nothing, start = zero(T), @@ -237,10 +254,9 @@ function add_con( ucon = zero(T), ) where {T,VT<:AbstractVector{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} - gen = _get_generator(ns) - dims = _get_con_dims(ns) + dims = _infer_subexpr_dims(gen.iter) gen = _adapt_gen(gen) - f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) + f = SIMDFunction(T, gen, c.ncon, c.nnzj, c.nnzh) pars = gen.iter nscen = c.tag.nscen @@ -250,6 +266,27 @@ function add_con( return _add_con(c, f, pars, dims, start, lcon, ucon, name, SecondStageConstraintTag()) end +function add_con( + c::C, + ::EachScenario, + ns::Union{Integer, AbstractUnitRange}...; + name = nothing, + tag = nothing, + start = zero(T), + lcon = zero(T), + ucon = zero(T), + ) where {T,VT<:AbstractVector{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} + + f = _simdfunction(T, Null(nothing), c.ncon, c.nnzj, c.nnzh) + pars = _empty_con_itr(ns) + + nscen = c.tag.nscen + len = length(pars) + new_con_scen = append!(c.backend, c.tag.con_scen, _scen_each_tag(nscen, div(len, nscen)), len) + c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, c.tag.var_scen, new_con_scen)) + return _add_con(c, f, pars, ns, start, lcon, ucon, name, SecondStageConstraintTag()) +end + # --- Accessors --- """ From 8656880b610f8354eaca80cab4bcff473d8c5655 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 19:49:32 -0400 Subject: [PATCH 11/50] Consolidate add_con to ns... varargs, fix append! for ND arrays, enable batched two-stage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Unify add_con into single ns... varargs method with _get_generator/_get_con_dims dispatch helpers (reduces combinatorial explosion for extensions) - Fix append! to handle N-dimensional arrays (AbstractArray instead of separate AbstractVector/AbstractMatrix methods) — fixes juliac COPSApp compilation for catmix/gasoil/methanol/pinene models - Widen two-stage dispatch from AbstractVector to AbstractArray, enabling batched two-stage models (TwoStageExaCore with nbatch kwarg) - Remove undefined flatten_model from exports - Add batch tests: set_parameter!, multidim vars, add_con!, add_expr, per-instance accessors, backend-parameterized evaluation - Add batched two-stage tests: construction, evaluation, ipopt Co-Authored-By: Claude Opus 4.6 --- src/ExaModels.jl | 1 - src/nlp.jl | 68 +++------ src/two_stage.jl | 74 +++------- test/BatchTest/BatchTest.jl | 220 +++++++++++++++++++++++++++++- test/TwoStageTest/TwoStageTest.jl | 107 +++++++++++++++ 5 files changed, 365 insertions(+), 105 deletions(-) diff --git a/src/ExaModels.jl b/src/ExaModels.jl index 8dda14eee..d1999bf93 100644 --- a/src/ExaModels.jl +++ b/src/ExaModels.jl @@ -117,7 +117,6 @@ export ExaModel, BatchExaModel, get_nbatch, get_model, - flatten_model, var_indices, cons_block_indices diff --git a/src/nlp.jl b/src/nlp.jl index 5897d91d6..84109ade6 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -262,9 +262,13 @@ end abstract type AbstractExaCore{T,VT,B,S} end """ - ExaCore([array_eltype::Type; backend = nothing, minimize = true, name = :Generic]) + ExaCore([T::Type; backend = nothing, concrete = Val(false), nbatch = Val(1), minimize = true, name = :Generic]) -Creates an intermediate data object `ExaCore`, which later can be used for creating an `ExaModel` +Creates an intermediate data object `ExaCore`, which later can be used for creating an `ExaModel`. + +When `nbatch = Val(N)` with `N > 1`, creates a batch core with matrix-valued +storage arrays (columns = instances). See [`BatchExaCore`](@ref) for a +convenience alias. ## Example ```jldoctest @@ -1085,67 +1089,33 @@ Constraint where |I| = 9 ``` """ -# Generator form — directly dispatched for type stability and juliac trimmer. @inline function add_con( c::C, - gen::Base.Generator; + ns...; tag = nothing, name = nothing, start = zero(T), lcon = zero(T), ucon = zero(T), + kwargs... ) where {T,C<:ExaCore{T}} - dims = _infer_subexpr_dims(gen.iter) + gen = _get_generator(ns) + dims = _get_con_dims(ns) gen = _adapt_gen(gen) - f = SIMDFunction(T, gen, c.ncon, c.nnzj, c.nnzh) + f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) pars = gen.iter - _add_con(c, f, pars, dims, start, lcon, ucon, name, tag) -end -# Multi-generator form: first generator creates the constraint, rest augment it. -@inline function add_con( - c::C, - gen::Base.Generator, - gens::Base.Generator...; - kwargs... -) where {T,C<:ExaCore{T}} - c, con = add_con(c, gen; kwargs...) - for g in gens - c, _ = add_con!(c, con, g) - end - return (c, con) -end - -# Expression form — pre-built expression tree with explicit parameters. -@inline function add_con( - c::C, - expr::N, - pars = 1:1; - tag = nothing, - name = nothing, - start = zero(T), - lcon = zero(T), - ucon = zero(T), -) where {T,C<:ExaCore{T},N<:AbstractNode} - f = _simdfunction(T, expr, c.ncon, c.nnzj, c.nnzh) - dims = _infer_subexpr_dims(pars) _add_con(c, f, pars, dims, start, lcon, ucon, name, tag) end -# Dims form — empty constraints for later augmentation. -@inline function add_con( - c::C, - ns::Union{Integer, AbstractUnitRange}...; - tag = nothing, - name = nothing, - start = zero(T), - lcon = zero(T), - ucon = zero(T), -) where {T,C<:ExaCore{T}} - f = _simdfunction(T, Null(nothing), c.ncon, c.nnzj, c.nnzh) - pars = _empty_con_itr(ns) - _add_con(c, f, pars, ns, start, lcon, ucon, name, tag) -end +@inline _get_generator(ns) = (Null(nothing) for _ in _empty_con_itr(ns)) +@inline _get_generator(gen::Tuple{G}) where G <: Base.Generator = gen[1] +@inline _get_generator(n::Tuple{N}) where N <: AbstractNode = (n[1] for _ in 1:1) + +# Infer constraint dims from the original arguments, preserving range start info. +@inline _get_con_dims(ns) = ns +@inline _get_con_dims(gen::Tuple{G}) where G <: Base.Generator = _infer_subexpr_dims(gen[1].iter) +@inline _get_con_dims(n::Tuple{N}) where N <: AbstractNode = (1,) # Build an iterator for empty constraints: 1:n for 1D, collected ProductIterator for multi-dim. _empty_con_itr(ns::Tuple{Any}) = 1:_length(ns[1]) diff --git a/src/two_stage.jl b/src/two_stage.jl index 16edd29e0..99bc3e7e4 100644 --- a/src/two_stage.jl +++ b/src/two_stage.jl @@ -60,7 +60,7 @@ const TwoStageExaModel{T,VT,E,V,P,O,C,R} = ExaModel{T,VT,E,V,P,O,C,<:TwoStageExa """ - TwoStageExaCore(nscen; backend = nothing, concrete = Val(false), kwargs...) + TwoStageExaCore(nscen; backend = nothing, concrete = Val(false), nbatch = Val(1), kwargs...) Create an [`ExaCore`](@ref) for building two-stage stochastic programs with `nscen` scenarios. @@ -69,6 +69,9 @@ Use [`add_var`](@ref), [`add_par`](@ref), and [`add_con`](@ref) with [`EachScenario()`](@ref) to declare per-scenario components, or without it for first-stage (design) components. +When `nbatch = Val(N)` with `N > 1`, creates a batched two-stage core that +combines batch optimization with two-stage structure. + ## Example ```julia core = TwoStageExaCore(5) # 5 scenarios @@ -103,7 +106,7 @@ function add_var( start = zero(T), lvar = T(-Inf), uvar = T(Inf), - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} len = total(ns) new_var_scen = append!(c.backend, c.tag.var_scen, 0, len) @@ -128,7 +131,7 @@ function add_var( start = zero(T), lvar = T(-Inf), uvar = T(Inf), - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} nscen = c.tag.nscen len = total(ns) new_var_scen = append!(c.backend, c.tag.var_scen, _scen_each_tag(nscen, len), len * nscen) @@ -155,7 +158,7 @@ function add_par( c::TwoStageExaCore{T,VT,B}, value::AbstractArray; name = nothing, - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} return _add_par(c, FirstStageTag(), name, value, Base.size(value)...) end function add_par( @@ -163,7 +166,7 @@ function add_par( n::AbstractRange; name = nothing, value = zero(T), - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} return _add_par(c, FirstStageTag(), name, value, n) end function add_par( @@ -171,7 +174,7 @@ function add_par( ns...; name = nothing, value = zero(T), - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} return _add_par(c, FirstStageTag(), name, value, ns...) end @@ -186,7 +189,7 @@ function add_par( ::EachScenario, value::AbstractVector; name = nothing, - ) where {T,VT<:AbstractVector{T},B} + ) where {T,VT<:AbstractArray{T},B} combined = cat((value for _ in 1:c.tag.nscen)...; dims = ndims(value) + 1) return _add_par(c, SecondStageTag(), name, combined, Base.size(combined)...) end @@ -199,17 +202,18 @@ the base [`add_con`](@ref) (generator or dims). Tagged with `FirstStageTag()`. """ function add_con( c::C, - gen::Base.Generator; + ns...; name = nothing, tag = nothing, start = zero(T), lcon = zero(T), ucon = zero(T), - ) where {T,VT<:AbstractVector{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} + ) where {T,VT<:AbstractArray{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} - dims = _infer_subexpr_dims(gen.iter) + gen = _get_generator(ns) + dims = _get_con_dims(ns) gen = _adapt_gen(gen) - f = SIMDFunction(T, gen, c.ncon, c.nnzj, c.nnzh) + f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) pars = gen.iter new_con_scen = append!(c.backend, c.tag.con_scen, 0, length(pars)) @@ -217,24 +221,6 @@ function add_con( return _add_con(c, f, pars, dims, start, lcon, ucon, name, FirstStageConstraintTag()) end -function add_con( - c::C, - ns::Union{Integer, AbstractUnitRange}...; - name = nothing, - tag = nothing, - start = zero(T), - lcon = zero(T), - ucon = zero(T), - ) where {T,VT<:AbstractVector{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} - - f = _simdfunction(T, Null(nothing), c.ncon, c.nnzj, c.nnzh) - pars = _empty_con_itr(ns) - - new_con_scen = append!(c.backend, c.tag.con_scen, 0, length(pars)) - c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, c.tag.var_scen, new_con_scen)) - return _add_con(c, f, pars, ns, start, lcon, ucon, name, FirstStageConstraintTag()) -end - """ add_con(core::TwoStageExaCore, ::EachScenario, dims_or_gen...; start = 0, lcon = 0, ucon = 0, name = nothing) @@ -246,17 +232,18 @@ scenario index. function add_con( c::C, ::EachScenario, - gen::Base.Generator; + ns...; name = nothing, tag = nothing, start = zero(T), lcon = zero(T), ucon = zero(T), - ) where {T,VT<:AbstractVector{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} + ) where {T,VT<:AbstractArray{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} - dims = _infer_subexpr_dims(gen.iter) + gen = _get_generator(ns) + dims = _get_con_dims(ns) gen = _adapt_gen(gen) - f = SIMDFunction(T, gen, c.ncon, c.nnzj, c.nnzh) + f = _simdfunction(T, gen.f(DataSource()), c.ncon, c.nnzj, c.nnzh) pars = gen.iter nscen = c.tag.nscen @@ -266,27 +253,6 @@ function add_con( return _add_con(c, f, pars, dims, start, lcon, ucon, name, SecondStageConstraintTag()) end -function add_con( - c::C, - ::EachScenario, - ns::Union{Integer, AbstractUnitRange}...; - name = nothing, - tag = nothing, - start = zero(T), - lcon = zero(T), - ucon = zero(T), - ) where {T,VT<:AbstractVector{T},B,S<:TwoStageExaModelTag,C<:ExaCore{T,VT,B,S}} - - f = _simdfunction(T, Null(nothing), c.ncon, c.nnzj, c.nnzh) - pars = _empty_con_itr(ns) - - nscen = c.tag.nscen - len = length(pars) - new_con_scen = append!(c.backend, c.tag.con_scen, _scen_each_tag(nscen, div(len, nscen)), len) - c = ExaCore(c; tag = TwoStageExaModelTag(c.tag.nscen, c.tag.var_scen, new_con_scen)) - return _add_con(c, f, pars, ns, start, lcon, ucon, name, SecondStageConstraintTag()) -end - # --- Accessors --- """ diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index 31c937136..28fa71f88 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -7,7 +7,8 @@ import NLPModels: obj, cons!, cons_nln!, grad!, jac_coord!, hess_coord!, jac_structure!, hess_structure! import NLPModels: obj! -import ExaModels: var_indices, cons_block_indices, get_model, get_nbatch +import ExaModels: var_indices, cons_block_indices, get_model, get_nbatch, + get_start, get_lvar, get_uvar, get_lcon, get_ucon import NLPModelsIpopt: ipopt @@ -287,6 +288,214 @@ function test_ipopt_multi() @test isapprox(result.objective, 0.0; atol = 1e-8) end +function test_set_parameter() + ns, nv = 2, 1 + c = BatchExaCore(ns) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + c, _ = add_obj(c, (v[1] - θ[1])^2 for _ in 1:1) + model = ExaModel(c) + + # Default: both instances share θ = [2.0] + bx = reshape([1.0, 3.0], nv, ns) + bf = zeros(ns) + obj!(model, bx, bf) + @test bf[1] ≈ 1.0 # (1-2)^2 + @test bf[2] ≈ 1.0 # (3-4)^2 — θ replicated to [2,4] via batch par + + # Update parameters for instance 2 + ExaModels.set_parameter!(c, θ, [10.0]) + model2 = ExaModel(c) + bf2 = zeros(ns) + obj!(model2, bx, bf2) + # After set_parameter!, the parameter is updated for the next ExaModel build. + @test bf2[1] ≈ (1.0 - 10.0)^2 + @test bf2[2] ≈ (3.0 - 10.0)^2 +end + +function test_multidim_vars() + ns = 2 + nh, nc = 3, 2 # multi-dimensional variable + c = BatchExaCore(ns) + @add_var(c, w, nh, nc; start = zeros(nh, nc)) + @add_par(c, θ, [1.0]) + c, _ = add_obj(c, w[i, j]^2 for i in 1:nh, j in 1:nc) + model = ExaModel(c) + + @test NLPModels.get_nvar(model) == nh * nc + @test get_nbatch(model) == ns + + bx = reshape(Float64.(1:(nh*nc*ns)), nh * nc, ns) + bf = zeros(ns) + obj!(model, bx, bf) + @test bf[1] ≈ sum(bx[:, 1] .^ 2) + @test bf[2] ≈ sum(bx[:, 2] .^ 2) +end + +function test_add_con_aug() + ns, nv = 2, 3 + c = BatchExaCore(ns) + @add_var(c, v, nv) + @add_par(c, θ, [1.0]) + c, _ = add_obj(c, v[j]^2 for j in 1:nv) + # Create a constraint, then augment it + @add_con(c, g, v[j] for j in 1:nv; lcon = -10.0, ucon = 10.0) + @add_con!(c, g, j => θ[1] * v[j] for j in 1:nv) + model = ExaModel(c) + flat = get_model(model) + + nc = NLPModels.get_ncon(model) + @test nc == nv + + bx = reshape(Float64.(1:(nv*ns)), nv, ns) + + # cons! — augmented constraint = v[j] + θ[1]*v[j] = 2*v[j] + bc = zeros(nc, ns) + cons!(model, bx, bc) + c_flat = zeros(nc * ns) + cons_nln!(flat, vec(bx), c_flat) + @test vec(bc) ≈ c_flat + + # Each constraint value should be v[j] + 1.0*v[j] = 2*v[j] + @test bc[:, 1] ≈ 2.0 .* bx[:, 1] + @test bc[:, 2] ≈ 2.0 .* bx[:, 2] + + # jac + nnzj = NLPModels.get_nnzj(model) + jvals = zeros(nnzj, ns) + jac_coord!(model, bx, jvals) + jvals_flat = zeros(NLPModels.get_nnzj(flat)) + jac_coord!(flat, vec(bx), jvals_flat) + @test vec(jvals) ≈ jvals_flat + + # hess + nnzh = NLPModels.get_nnzh(model) + by = ones(nc, ns) + hvals = zeros(nnzh, ns) + hess_coord!(model, bx, by, hvals) + hvals_flat = zeros(NLPModels.get_nnzh(flat)) + hess_coord!(flat, vec(bx), vec(by), hvals_flat) + @test vec(hvals) ≈ hvals_flat +end + +function test_add_expr() + ns, nv = 2, 3 + c = BatchExaCore(ns) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + + # Create a subexpression and use it in objective and constraint + @add_expr(c, s, θ[1] * v[j]^2 for j in 1:nv) + c, _ = add_obj(c, s[j] for j in 1:nv) + c, _ = add_con(c, s[j] - v[j] for j in 1:nv; lcon = 0.0) + model = ExaModel(c) + flat = get_model(model) + + nc = NLPModels.get_ncon(model) + @test nc == nv + + bx = reshape(Float64.(1:(nv*ns)), nv, ns) + + # obj — should be sum of θ[1]*v[j]^2 = 2*v[j]^2 + bf = zeros(ns) + obj!(model, bx, bf) + @test bf[1] ≈ 2.0 * sum(bx[:, 1] .^ 2) + @test bf[2] ≈ 2.0 * sum(bx[:, 2] .^ 2) + @test sum(bf) ≈ obj(flat, vec(bx)) + + # cons — s[j] - v[j] = 2*v[j]^2 - v[j] + bc = zeros(nc, ns) + cons!(model, bx, bc) + c_flat = zeros(nc * ns) + cons_nln!(flat, vec(bx), c_flat) + @test vec(bc) ≈ c_flat + @test bc[:, 1] ≈ 2.0 .* bx[:, 1] .^ 2 .- bx[:, 1] + + # grad + bg = zeros(nv, ns) + grad!(model, bx, bg) + g_flat = zeros(nv * ns) + grad!(flat, vec(bx), g_flat) + @test vec(bg) ≈ g_flat +end + +function test_per_instance_accessors() + ns, nv = 2, 3 + c = BatchExaCore(ns) + @add_var(c, v, nv; start = 1.0, lvar = -5.0, uvar = 5.0) + c, _ = add_obj(c, v[j]^2 for j in 1:nv) + c, con = add_con(c, v[j] for j in 1:nv; lcon = 0.0, ucon = 10.0) + model = ExaModel(c) + + # Test per-instance variable accessors + @test get_start(model, v, 1) ≈ fill(1.0, nv) + @test get_lvar(model, v, 1) ≈ fill(-5.0, nv) + @test get_uvar(model, v, 2) ≈ fill(5.0, nv) + + # Test per-instance constraint accessors + @test get_lcon(model, con, 1) ≈ fill(0.0, nv) + @test get_ucon(model, con, 2) ≈ fill(10.0, nv) + + # Test cons_block_indices + @test cons_block_indices(model, 1) == 1:nv + @test cons_block_indices(model, 2) == (nv+1):(2*nv) +end + +# ============================================================================ +# Backend-parameterized tests — verify batch evaluation on all backends (incl. KA/GPU) +# ============================================================================ + +function test_batch_backend(backend) + ns, nv = 2, 3 + c = BatchExaCore(ns; backend) + @add_var(c, v, nv) + @add_par(c, θ, [2.0]) + c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) + @add_con(c, g, v[j] - θ[1] for j in 1:nv; lcon = 0.0) + @add_con!(c, g, j => θ[1] * v[j] for j in 1:nv) + model = ExaModel(c) + flat = get_model(model) + + nc = NLPModels.get_ncon(model) + bx = ExaModels.convert_array(reshape(Float64.(1:(nv*ns)), nv, ns), backend) + + # obj + bf = zeros(ns) + obj!(model, bx, bf) + @test sum(bf) ≈ obj(flat, vec(bx)) + + # grad + bg = ExaModels.convert_array(zeros(nv, ns), backend) + grad!(model, bx, bg) + g_flat = ExaModels.convert_array(zeros(nv * ns), backend) + grad!(flat, vec(bx), g_flat) + @test vec(Array(bg)) ≈ Array(g_flat) + + # cons + bc = ExaModels.convert_array(zeros(nc, ns), backend) + cons!(model, bx, bc) + c_flat = ExaModels.convert_array(zeros(nc * ns), backend) + cons_nln!(flat, vec(bx), c_flat) + @test vec(Array(bc)) ≈ Array(c_flat) + + # jac + nnzj = NLPModels.get_nnzj(model) + jvals = ExaModels.convert_array(zeros(nnzj, ns), backend) + jac_coord!(model, bx, jvals) + jvals_flat = ExaModels.convert_array(zeros(NLPModels.get_nnzj(flat)), backend) + jac_coord!(flat, vec(bx), jvals_flat) + @test vec(Array(jvals)) ≈ Array(jvals_flat) + + # hess + nnzh = NLPModels.get_nnzh(model) + by = ExaModels.convert_array(ones(nc, ns), backend) + hvals = ExaModels.convert_array(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals) + hvals_flat = ExaModels.convert_array(zeros(NLPModels.get_nnzh(flat)), backend) + hess_coord!(flat, vec(bx), vec(by), hvals_flat) + @test vec(Array(hvals)) ≈ Array(hvals_flat) +end + # ============================================================================ function runtests() @@ -304,6 +513,15 @@ function runtests() @testset "flatten_model" test_flatten_model() @testset "Ipopt simple" test_ipopt_simple() @testset "Ipopt multi" test_ipopt_multi() + @testset "set_parameter!" test_set_parameter() + @testset "Multidim vars" test_multidim_vars() + @testset "add_con!" test_add_con_aug() + @testset "add_expr" test_add_expr() + @testset "Per-instance accessors" test_per_instance_accessors() + for backend in BACKENDS + backend === nothing && continue + @testset "Backend: $(typeof(backend))" test_batch_backend(backend) + end end end diff --git a/test/TwoStageTest/TwoStageTest.jl b/test/TwoStageTest/TwoStageTest.jl index bba7501e8..d1b526a3a 100644 --- a/test/TwoStageTest/TwoStageTest.jl +++ b/test/TwoStageTest/TwoStageTest.jl @@ -516,6 +516,113 @@ function runtests() end end + + @testset "Batched two-stage" begin + + @testset "Construction" begin + ns_scen, nb = 3, 2 + core = TwoStageExaCore(ns_scen; nbatch = Val(nb)) + v = @add_var(core, EachScenario(), 2) + core, d = add_var(core, 1) + + model = ExaModel(core) + @test get_nscen(model) == ns_scen + @test ExaModels.get_nbatch(model) == nb + @test NLPModels.get_nvar(model) == ns_scen * 2 + 1 + @test size(model.meta.x0) == (ns_scen * 2 + 1, nb) + end + + @testset "Evaluation" begin + ns_scen, nb = 2, 3 + nv, nd = 1, 1 + θ_vals = [2.0, 4.0] + + core = TwoStageExaCore(ns_scen; nbatch = Val(nb)) + v = @add_var(core, EachScenario(), nv) + core, d = add_var(core, nd) + core, θ = add_par(core, θ_vals) + + obj_data = [(i, (i-1)*nv+1, i) for i in 1:ns_scen] + @add_obj(core, (v[v_idx] - θ[θi])^2 + d[1]^2 for (i, v_idx, θi) in obj_data) + + con_data = [(i, (i-1)*nv+1) for i in 1:ns_scen] + @add_con(core, EachScenario(), (v[v_idx] - d[1] for (i, v_idx) in con_data); lcon = 0.0) + + model = ExaModel(core) + flat = ExaModels.get_model(model) + + nvar = NLPModels.get_nvar(model) + ncon = NLPModels.get_ncon(model) + + # obj + bx = reshape(Float64.(1:(nvar*nb)), nvar, nb) + bf = zeros(nb) + NLPModels.obj!(model, bx, bf) + @test sum(bf) ≈ NLPModels.obj(flat, vec(bx)) + + # grad + bg = zeros(nvar, nb) + NLPModels.grad!(model, bx, bg) + g_flat = zeros(nvar * nb) + grad!(flat, vec(bx), g_flat) + @test vec(bg) ≈ g_flat + + # cons + bc = zeros(ncon, nb) + NLPModels.cons!(model, bx, bc) + c_flat = zeros(ncon * nb) + cons_nln!(flat, vec(bx), c_flat) + @test vec(bc) ≈ c_flat + + # jac + nnzj = NLPModels.get_nnzj(model) + jvals = zeros(nnzj, nb) + jac_coord!(model, bx, jvals) + jvals_flat = zeros(NLPModels.get_nnzj(flat)) + jac_coord!(flat, vec(bx), jvals_flat) + @test vec(jvals) ≈ jvals_flat + + # hess + nnzh = NLPModels.get_nnzh(model) + by = ones(ncon, nb) + hvals = zeros(nnzh, nb) + hess_coord!(model, bx, by, hvals) + hvals_flat = zeros(NLPModels.get_nnzh(flat)) + hess_coord!(flat, vec(bx), vec(by), hvals_flat) + @test vec(hvals) ≈ hvals_flat + end + + @testset "Ipopt" begin + ns_scen, nb = 2, 2 + nv, nd = 1, 1 + θ_vals = [2.0, 4.0] + + core = TwoStageExaCore(ns_scen; nbatch = Val(nb)) + v = @add_var(core, EachScenario(), nv) + core, d = add_var(core, nd) + core, θ = add_par(core, θ_vals) + + @add_obj(core, d[1]^2) + @add_obj(core, (v[i] - θ[i])^2 for i in 1:ns_scen) + @add_con(core, EachScenario(), (v[i] - d[1] for i in 1:ns_scen); lcon = 0.0) + + model = ExaModel(core) + flat = ExaModels.get_model(model) + + result = ipopt(flat; print_level = 0) + @test result.status == :first_order + + # Each batch instance should have the same solution + nvar = NLPModels.get_nvar(model) + for b in 1:nb + x_b = result.solution[ExaModels.var_indices(model, b)] + # d* = θ̄/2 = 3/2, v_i* = d* (all constraints active at optimum) + d_expected = sum(θ_vals) / ns_scen / 2 + @test x_b[end] ≈ d_expected atol = 1e-4 + end + end + + end end end # module TwoStageTest From 49654346e458e1a8036fcded628eb0e98798eb0a Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 19:59:25 -0400 Subject: [PATCH 12/50] Fix docs, add get_model docstring, remove duplicate test include - Clarify that set_parameter! applies globally, not per-instance - Add docstring for exported get_model function - Remove duplicate TwoStageTest include in runtests.jl Co-Authored-By: Claude Opus 4.6 --- docs/src/batch.jl | 9 ++++----- src/nlp.jl | 8 ++++++++ test/runtests.jl | 1 - 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/docs/src/batch.jl b/docs/src/batch.jl index 1a0236919..01b5f1f61 100644 --- a/docs/src/batch.jl +++ b/docs/src/batch.jl @@ -84,10 +84,9 @@ for i in 1:ns println("Instance $i: v* = ", round(v_sol[1], digits = 4)) end -# ## Per-Instance Parameters -# Each instance can have different parameter values. -# Parameters are stored as a matrix `(nθ, ns)`. You can set different values -# per instance using `set_parameter!`: +# ## Updating Parameters +# Parameters can be updated on the core using `set_parameter!` and then +# rebuilding the model. The updated values apply to all instances: c2 = BatchExaCore(2) @add_var(c2, x, 2) @@ -95,7 +94,7 @@ c2 = BatchExaCore(2) @add_obj(c2, (x[j] - p[j])^2 for j in 1:2) model2 = ExaModel(c2) -# Update parameters for instance 2: +# Update parameters (applies to all instances): ExaModels.set_parameter!(c2, p, [10.0, 20.0]) model2 = ExaModel(c2) println("\nParameter matrix shape: ", size(model2.θ)) diff --git a/src/nlp.jl b/src/nlp.jl index 84109ade6..bd48d91d8 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -1836,6 +1836,14 @@ BatchExaCore(nbatch::Integer; kwargs...) = ExaCore(; concrete = Val(true), nbatc # get_model — defined after BatchNLPModels is loaded (see ExaModels.jl) # ============================================================================ +""" + get_model(model) + +Return a solver-ready NLP model. For a [`BatchExaModel`](@ref), returns a +[`FlattenNLPModel`](@ref) that presents all instances as a single flat +`AbstractNLPModel{T, Vector{T}}`. For a regular [`ExaModel`](@ref), returns +the model itself. +""" get_model(model::ExaModel) = model """ diff --git a/test/runtests.jl b/test/runtests.jl index 418d5180c..f99801547 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -23,7 +23,6 @@ include("UtilsTest/UtilsTest.jl") include("TwoStageTest/TwoStageTest.jl") include("BatchTest/BatchTest.jl") include("JuliaCTest/JuliaCTest.jl") -include("TwoStageTest/TwoStageTest.jl") include("GetterSetterTest/GetterSetterTest.jl") include("PrettyPrintTest.jl") # include("OptimalControlTest/OptimalControlTest.jl") From 665b6f8170cdcd93c90234455f4af74a649b4e2d Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 20:21:19 -0400 Subject: [PATCH 13/50] Fix CI failures: deprecate set_parameter!, fix batch GPU, fix two-stage test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Deprecate set_parameter! in favor of set_value! on the model - Fix set_value! for batch models (broadcast across all batch columns) - Widen KA build_extension from AbstractVector to AbstractArray so batch models on GPU get a proper KAExtension - Fix obj() to allocate output with similar(bx) for correct device - Fix batched two-stage Ipopt test expectation (d* = Σθ/(1+ns), not Σθ/ns/2) - Update docs to use set_value! instead of set_parameter! Co-Authored-By: Claude Opus 4.6 --- docs/src/batch.jl | 7 ++--- docs/src/parameters.jl | 4 +-- ext/ExaModelsKernelAbstractions.jl | 2 +- src/nlp.jl | 48 ++++++++++++++---------------- test/BatchTest/BatchTest.jl | 17 +++++------ test/TwoStageTest/TwoStageTest.jl | 4 +-- 6 files changed, 38 insertions(+), 44 deletions(-) diff --git a/docs/src/batch.jl b/docs/src/batch.jl index 01b5f1f61..41afd49c3 100644 --- a/docs/src/batch.jl +++ b/docs/src/batch.jl @@ -85,8 +85,8 @@ for i in 1:ns end # ## Updating Parameters -# Parameters can be updated on the core using `set_parameter!` and then -# rebuilding the model. The updated values apply to all instances: +# Parameters can be updated directly on the model using `set_value!`. +# The updated values apply to all instances: c2 = BatchExaCore(2) @add_var(c2, x, 2) @@ -95,6 +95,5 @@ c2 = BatchExaCore(2) model2 = ExaModel(c2) # Update parameters (applies to all instances): -ExaModels.set_parameter!(c2, p, [10.0, 20.0]) -model2 = ExaModel(c2) +ExaModels.set_value!(model2, p, [10.0, 20.0]) println("\nParameter matrix shape: ", size(model2.θ)) diff --git a/docs/src/parameters.jl b/docs/src/parameters.jl index 93477e4ba..635c81545 100644 --- a/docs/src/parameters.jl +++ b/docs/src/parameters.jl @@ -35,11 +35,11 @@ result1 = ipopt(m_param) println("Original objective: $(result1.objective)") # Now change the penalty coefficient and solve again: -set_parameter!(c_param, θ, [200.0, 1.0]) # Double the penalty coefficient +set_value!(m_param, θ, [200.0, 1.0]) # Double the penalty coefficient result2 = ipopt(m_param) println("Modified penalty objective: $(result2.objective)") # Try a different offset parameter: -set_parameter!(c_param, θ, [200.0, 0.5]) # Change the offset in the objective +set_value!(m_param, θ, [200.0, 0.5]) # Change the offset in the objective result3 = ipopt(m_param) println("Modified offset objective: $(result3.objective)") diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index 9691deef0..6e19d25aa 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -33,7 +33,7 @@ function ExaModels.build_extension( c::C; prod = false, kwargs..., -) where {T,VT<:AbstractVector{T},B<:KernelAbstractions.Backend,C<:ExaModels.ExaCore{T,VT,B}} +) where {T,VT<:AbstractArray{T},B<:KernelAbstractions.Backend,C<:ExaModels.ExaCore{T,VT,B}} gsparsity = similar(c.x0, Tuple{Int,Int}, c.nnzg) diff --git a/src/nlp.jl b/src/nlp.jl index bd48d91d8..4ab97ba99 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -833,33 +833,24 @@ end """ set_parameter!(core, param, values) -Updates the values of parameters in the core. - -## Example -```jldoctest -julia> using ExaModels +!!! warning "Deprecated" + `set_parameter!` is deprecated. Use [`set_value!`](@ref) on the model instead. -julia> c = ExaCore(concrete = Val(true)); - -julia> c, p = add_par(c, ones(5)); - -julia> set_parameter!(c, p, ones(5)) -``` +Updates the values of parameters in the core. """ function set_parameter!(c::ExaCore, param::Parameter, values::AbstractArray) - if length(values) != param.length - throw( - DimensionMismatch( - "Parameter size mismatch: expected $(param.length) elements, got $(length(values))", - ), - ) + Base.depwarn( + "`set_parameter!(core, param, values)` is deprecated, use `set_value!(model, param, values)` on the model instead.", + :set_parameter!, + ) + rng = param.offset+1 : param.offset+param.length + if c.θ isa AbstractMatrix + for col in axes(c.θ, 2) + copyto!(@view(c.θ[rng, col]), values) + end + else + copyto!(@view(c.θ[rng]), values) end - - start_idx = param.offset + 1 - end_idx = param.offset + param.length - - copyto!(@view(c.θ[start_idx:end_idx]), values) - return nothing end @@ -886,7 +877,14 @@ function set_value!(model::ExaModel, param::Parameter, values) "expected $(param.length) elements, got $(length(values))" )) end - copyto!(view(model.θ, param.offset+1:param.offset+param.length), values) + rng = param.offset+1 : param.offset+param.length + if model.θ isa AbstractMatrix + for col in axes(model.θ, 2) + copyto!(view(model.θ, rng, col), values) + end + else + copyto!(view(model.θ, rng), values) + end return nothing end @@ -1940,7 +1938,7 @@ function obj!(m::BatchExaModel{T}, bx::AbstractMatrix, bf::AbstractVector) where end function obj(m::BatchExaModel{T}, bx::AbstractMatrix) where {T} - bf = Vector{T}(undef, get_nbatch(m)) + bf = similar(bx, T, get_nbatch(m)) obj!(m, bx, bf) return bf end diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index 28fa71f88..33322a9c4 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -301,14 +301,12 @@ function test_set_parameter() bf = zeros(ns) obj!(model, bx, bf) @test bf[1] ≈ 1.0 # (1-2)^2 - @test bf[2] ≈ 1.0 # (3-4)^2 — θ replicated to [2,4] via batch par + @test bf[2] ≈ 1.0 # (3-2)^2 - # Update parameters for instance 2 - ExaModels.set_parameter!(c, θ, [10.0]) - model2 = ExaModel(c) + # Update parameters via set_value! on the model + set_value!(model, θ, [10.0]) bf2 = zeros(ns) - obj!(model2, bx, bf2) - # After set_parameter!, the parameter is updated for the next ExaModel build. + obj!(model, bx, bf2) @test bf2[1] ≈ (1.0 - 10.0)^2 @test bf2[2] ≈ (3.0 - 10.0)^2 end @@ -460,9 +458,8 @@ function test_batch_backend(backend) bx = ExaModels.convert_array(reshape(Float64.(1:(nv*ns)), nv, ns), backend) # obj - bf = zeros(ns) - obj!(model, bx, bf) - @test sum(bf) ≈ obj(flat, vec(bx)) + bf_gpu = obj(model, bx) + @test sum(Array(bf_gpu)) ≈ obj(flat, vec(bx)) # grad bg = ExaModels.convert_array(zeros(nv, ns), backend) @@ -513,7 +510,7 @@ function runtests() @testset "flatten_model" test_flatten_model() @testset "Ipopt simple" test_ipopt_simple() @testset "Ipopt multi" test_ipopt_multi() - @testset "set_parameter!" test_set_parameter() + @testset "set_value!" test_set_parameter() @testset "Multidim vars" test_multidim_vars() @testset "add_con!" test_add_con_aug() @testset "add_expr" test_add_expr() diff --git a/test/TwoStageTest/TwoStageTest.jl b/test/TwoStageTest/TwoStageTest.jl index d1b526a3a..be1dbea2a 100644 --- a/test/TwoStageTest/TwoStageTest.jl +++ b/test/TwoStageTest/TwoStageTest.jl @@ -616,8 +616,8 @@ function runtests() nvar = NLPModels.get_nvar(model) for b in 1:nb x_b = result.solution[ExaModels.var_indices(model, b)] - # d* = θ̄/2 = 3/2, v_i* = d* (all constraints active at optimum) - d_expected = sum(θ_vals) / ns_scen / 2 + # d² + Σ(v_i - θ_i)² s.t. v_i = d → d* = Σθ / (1 + ns) + d_expected = sum(θ_vals) / (1 + ns_scen) @test x_b[end] ≈ d_expected atol = 1e-4 end end From fc930b5b7922116a37043b643d49d5667075acc0 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 20:49:49 -0400 Subject: [PATCH 14/50] Fix batch GPU: race-free obj kernel, use similar for allocations - Replace racy kerf_batch kernel (bf[s] += ...) with two-phase fill+reduce approach to avoid write conflicts across threads - Use similar() instead of Vector/Matrix in BatchNLPModels allocating functions so output arrays land on the correct device Co-Authored-By: Claude Opus 4.6 --- ext/ExaModelsKernelAbstractions.jl | 18 +++++++++++++++--- src/BatchNLPModels.jl | 8 ++++---- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index 6e19d25aa..168aa62dc 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -681,17 +681,29 @@ end function ExaModels._obj_batch!(bf, obj, x, θ, nb, nvar, npar, backend::KernelAbstractions.Backend) nitr = length(obj.itr) if nitr > 0 - kerf_batch(backend)(bf, obj.f, obj.itr, x, θ, nvar, npar, nitr; ndrange = nb * nitr) + buf = similar(x, nitr * nb) + kerf_batch_fill(backend)(buf, obj.f, obj.itr, x, θ, nvar, npar, nitr; ndrange = nb * nitr) + kerf_batch_reduce(backend)(bf, buf, nitr; ndrange = nb) end end -@kernel function kerf_batch(bf, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(nvar), @Const(npar), @Const(nitr)) +@kernel function kerf_batch_fill(buf, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(nvar), @Const(npar), @Const(nitr)) I = @index(Global) s = (I - 1) ÷ nitr + 1 k = (I - 1) % nitr + 1 x_off = (s - 1) * nvar θ_off = (s - 1) * npar - @inbounds bf[s] += f(itr[k], ExaModels.OffsetVector(x, x_off), ExaModels.OffsetVector(θ, θ_off)) + @inbounds buf[I] = f(itr[k], ExaModels.OffsetVector(x, x_off), ExaModels.OffsetVector(θ, θ_off)) +end + +@kernel function kerf_batch_reduce(bf, @Const(buf), @Const(nitr)) + s = @index(Global) + val = zero(eltype(bf)) + base = (s - 1) * nitr + for k in 1:nitr + @inbounds val += buf[base + k] + end + @inbounds bf[s] += val end # --- Constraints --- diff --git a/src/BatchNLPModels.jl b/src/BatchNLPModels.jl index 17e3bbea7..4c4a36c67 100644 --- a/src/BatchNLPModels.jl +++ b/src/BatchNLPModels.jl @@ -53,7 +53,7 @@ get_nbatch(m::AbstractNLPModel) = get_nbatch(m.meta) Allocating version of `NLPModels.obj!`. """ function NLPModels.obj(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} - bf = Vector{T}(undef, get_nbatch(m)) + bf = similar(bx, T, get_nbatch(m)) obj!(m, bx, bf) return bf end @@ -73,7 +73,7 @@ end Allocating version of `grad!`. """ function NLPModels.grad(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} - bg = Matrix{T}(undef, NLPModels.get_nvar(m), get_nbatch(m)) + bg = similar(bx, T, NLPModels.get_nvar(m), get_nbatch(m)) NLPModels.grad!(m, bx, bg) return bg end @@ -93,7 +93,7 @@ end Allocating version of `cons!`. """ function NLPModels.cons(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} - bc = Matrix{T}(undef, NLPModels.get_ncon(m), get_nbatch(m)) + bc = similar(bx, T, NLPModels.get_ncon(m), get_nbatch(m)) NLPModels.cons!(m, bx, bc) return bc end @@ -198,7 +198,7 @@ function NLPModels.obj(m::FlattenNLPModel{T}, x::AbstractVector) where {T} nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) bx = reshape(x, nvar, nb) - bf = Vector{T}(undef, nb) + bf = similar(x, T, nb) obj!(m.batch, bx, bf) return sum(bf) end From a68ff8c8d96c2ab7ce6409926f0ca559bdb73dad Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 21:26:41 -0400 Subject: [PATCH 15/50] Use sum(buf; dims=1) instead of manual reduce kernel for batch obj Co-Authored-By: Claude Opus 4.6 --- ext/ExaModelsKernelAbstractions.jl | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index 168aa62dc..ed9956dfa 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -681,9 +681,9 @@ end function ExaModels._obj_batch!(bf, obj, x, θ, nb, nvar, npar, backend::KernelAbstractions.Backend) nitr = length(obj.itr) if nitr > 0 - buf = similar(x, nitr * nb) + buf = similar(x, nitr, nb) kerf_batch_fill(backend)(buf, obj.f, obj.itr, x, θ, nvar, npar, nitr; ndrange = nb * nitr) - kerf_batch_reduce(backend)(bf, buf, nitr; ndrange = nb) + bf .+= vec(sum(buf; dims = 1)) end end @@ -693,17 +693,7 @@ end k = (I - 1) % nitr + 1 x_off = (s - 1) * nvar θ_off = (s - 1) * npar - @inbounds buf[I] = f(itr[k], ExaModels.OffsetVector(x, x_off), ExaModels.OffsetVector(θ, θ_off)) -end - -@kernel function kerf_batch_reduce(bf, @Const(buf), @Const(nitr)) - s = @index(Global) - val = zero(eltype(bf)) - base = (s - 1) * nitr - for k in 1:nitr - @inbounds val += buf[base + k] - end - @inbounds bf[s] += val + @inbounds buf[k, s] = f(itr[k], ExaModels.OffsetVector(x, x_off), ExaModels.OffsetVector(θ, θ_off)) end # --- Constraints --- From c1ad8134397281632c4ce59ea40bf7608efd0d9a Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 21:51:37 -0400 Subject: [PATCH 16/50] =?UTF-8?q?Unify=20batch/non-batch=20eval=20paths,?= =?UTF-8?q?=20rename=20FlattenNLPModel=20=E2=86=92=20FlatNLPModel,=20remov?= =?UTF-8?q?e=20get=5Fmodel?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove _*_b! variants; BatchExaModel callers now use unified _grad!, _cons_nln!, _jac_coord!, _obj_hess_coord!, _con_hess_coord! (which accept nb parameter) - Rename FlattenNLPModel to FlatNLPModel and export it directly - Remove get_model entirely; use FlatNLPModel(model) instead - Add @inbounds to batch loops in gradient.jl, jacobian.jl, hessian.jl - Remove dead _eval_objbuffer! from KA extension Co-Authored-By: Claude Opus 4.6 --- docs/src/batch.jl | 4 +- ext/ExaModelsKernelAbstractions.jl | 7 -- src/BatchNLPModels.jl | 28 +++--- src/ExaModels.jl | 3 +- src/gradient.jl | 4 +- src/hessian.jl | 8 +- src/jacobian.jl | 4 +- src/nlp.jl | 153 +++++++---------------------- test/BatchTest/BatchTest.jl | 32 +++--- test/TwoStageTest/TwoStageTest.jl | 4 +- 10 files changed, 81 insertions(+), 166 deletions(-) diff --git a/docs/src/batch.jl b/docs/src/batch.jl index 41afd49c3..483250a8c 100644 --- a/docs/src/batch.jl +++ b/docs/src/batch.jl @@ -71,9 +71,9 @@ NLPModels.cons!(model, bx, bc) println("Constraints: ", bc) # ## Solving via the Fused Model -# For solving, use `get_model(model)` to access the fused `FlattenNLPModel` and +# For solving, use `FlatNLPModel(model)` to access the fused `FlatNLPModel` and # pass it to any NLPModels-compatible solver: -flat = ExaModels.get_model(model) +flat = FlatNLPModel(model) result = ipopt(flat; print_level = 0) println("\nSolution status: ", result.status) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index ed9956dfa..dd86203e6 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -198,13 +198,6 @@ function _obj(backend, objbuffer, (obj, objs...), x, θ) end end -function ExaModels._eval_objbuffer!( - objbuffer, m::ExaModels.ExaModel{T, VT, E}, x - ) where {T, VT, E <: KAExtension} - return if !isempty(objbuffer) - _obj(m.ext.backend, objbuffer, m.objs, x, m.θ) - end -end function ExaModels.cons_nln!( m::ExaModels.AbstractExaModel{T,VT,E}, diff --git a/src/BatchNLPModels.jl b/src/BatchNLPModels.jl index 4c4a36c67..16b43a546 100644 --- a/src/BatchNLPModels.jl +++ b/src/BatchNLPModels.jl @@ -153,26 +153,26 @@ function NLPModels.hess_coord!( end # ============================================================================ -# FlattenNLPModel +# FlatNLPModel # ============================================================================ """ - FlattenNLPModel{T, M} <: AbstractNLPModel{T, Vector{T}} + FlatNLPModel{T, M} <: AbstractNLPModel{T, Vector{T}} Wrapper that presents a batch NLP model as a flat (Vector-based) NLP model. All NLPModels callbacks delegate to the underlying batch model's matrix API. - FlattenNLPModel(model::AbstractNLPModel) + FlatNLPModel(model::AbstractNLPModel) Construct a flat model from a batch model whose `meta.x0` is a matrix. """ -struct FlattenNLPModel{T, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, Vector{T}} +struct FlatNLPModel{T, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, Vector{T}} batch::M meta::NLPModelMeta{T, Vector{T}} counters::NLPModels.Counters end -function FlattenNLPModel(model::AbstractNLPModel{T}) where {T} +function FlatNLPModel(model::AbstractNLPModel{T}) where {T} nb = get_nbatch(model) nvar = NLPModels.get_nvar(model) * nb ncon = NLPModels.get_ncon(model) * nb @@ -191,10 +191,10 @@ function FlattenNLPModel(model::AbstractNLPModel{T}) where {T} model.meta.minimize, false, String(model.meta.name), false, false, true, true, true, ncon > 0, true, ncon > 0, ncon > 0, true, ) - return FlattenNLPModel(model, meta, NLPModels.Counters()) + return FlatNLPModel(model, meta, NLPModels.Counters()) end -function NLPModels.obj(m::FlattenNLPModel{T}, x::AbstractVector) where {T} +function NLPModels.obj(m::FlatNLPModel{T}, x::AbstractVector) where {T} nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) bx = reshape(x, nvar, nb) @@ -203,14 +203,14 @@ function NLPModels.obj(m::FlattenNLPModel{T}, x::AbstractVector) where {T} return sum(bf) end -function NLPModels.grad!(m::FlattenNLPModel{T}, x::AbstractVector, g::AbstractVector) where {T} +function NLPModels.grad!(m::FlatNLPModel{T}, x::AbstractVector, g::AbstractVector) where {T} nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) NLPModels.grad!(m.batch, reshape(x, nvar, nb), reshape(g, nvar, nb)) return g end -function NLPModels.cons_nln!(m::FlattenNLPModel{T}, x::AbstractVector, c::AbstractVector) where {T} +function NLPModels.cons_nln!(m::FlatNLPModel{T}, x::AbstractVector, c::AbstractVector) where {T} nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) ncon = NLPModels.get_ncon(m.batch) @@ -218,7 +218,7 @@ function NLPModels.cons_nln!(m::FlattenNLPModel{T}, x::AbstractVector, c::Abstra return c end -function NLPModels.jac_structure!(m::FlattenNLPModel, rows::AbstractVector, cols::AbstractVector) +function NLPModels.jac_structure!(m::FlatNLPModel, rows::AbstractVector, cols::AbstractVector) nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) ncon = NLPModels.get_ncon(m.batch) @@ -242,7 +242,7 @@ function NLPModels.jac_structure!(m::FlattenNLPModel, rows::AbstractVector, cols return rows, cols end -function NLPModels.jac_coord!(m::FlattenNLPModel{T}, x::AbstractVector, jvals::AbstractVector) where {T} +function NLPModels.jac_coord!(m::FlatNLPModel{T}, x::AbstractVector, jvals::AbstractVector) where {T} nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) nnzj = NLPModels.get_nnzj(m.batch) @@ -250,7 +250,7 @@ function NLPModels.jac_coord!(m::FlattenNLPModel{T}, x::AbstractVector, jvals::A return jvals end -function NLPModels.hess_structure!(m::FlattenNLPModel, rows::AbstractVector, cols::AbstractVector) +function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector, cols::AbstractVector) nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) nnzh = NLPModels.get_nnzh(m.batch) @@ -273,7 +273,7 @@ function NLPModels.hess_structure!(m::FlattenNLPModel, rows::AbstractVector, col end function NLPModels.hess_coord!( - m::FlattenNLPModel{T}, x::AbstractVector, y::AbstractVector, + m::FlatNLPModel{T}, x::AbstractVector, y::AbstractVector, hvals::AbstractVector; obj_weight = one(T), ) where {T} nb = get_nbatch(m.batch) @@ -289,6 +289,6 @@ end # ============================================================================ export AbstractBatchNLPModel, - FlattenNLPModel + FlatNLPModel end # module BatchNLPModels diff --git a/src/ExaModels.jl b/src/ExaModels.jl index d1999bf93..502f9b944 100644 --- a/src/ExaModels.jl +++ b/src/ExaModels.jl @@ -63,7 +63,6 @@ include("tags.jl") include("two_stage.jl") include("BatchNLPModels.jl") using .BatchNLPModels -get_model(model::BatchExaModel) = BatchNLPModels.FlattenNLPModel(model) export ExaModel, ExaCore, @@ -113,10 +112,10 @@ export ExaModel, get_ucon, set_ucon!, AbstractBatchNLPModel, + FlatNLPModel, BatchExaCore, BatchExaModel, get_nbatch, - get_model, var_indices, cons_block_indices diff --git a/src/gradient.jl b/src/gradient.jl index 18c75bc3a..1b71d15b7 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -43,12 +43,12 @@ function gradient!(y, f, x, θ, adj) return y end function gradient!(y, f, x::AbstractArray, θ::AbstractArray, adj, nb::Integer, nvar::Integer, npar::Integer, ::Nothing = nothing) - for s in 1:nb + @inbounds for s in 1:nb x_s = @view x[(s-1)*nvar+1 : s*nvar] θ_s = @view θ[(s-1)*npar+1 : s*npar] y_s = @view y[(s-1)*nvar+1 : s*nvar] @simd for k in eachindex(f.itr) - @inbounds gradient!(y_s, f.f, x_s, θ_s, f.itr[k], adj) + gradient!(y_s, f.f, x_s, θ_s, f.itr[k], adj) end end return y diff --git a/src/hessian.jl b/src/hessian.jl index 437a98358..28659f073 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -698,13 +698,13 @@ function shessian!(y1, y2, f, x, θ, adj1, adj2) end end function shessian!(y1, y2, f, x::AbstractArray, θ::AbstractArray, adj1, adj2, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, ::Nothing = nothing) - for s in 1:nb + @inbounds for s in 1:nb x_s = @view x[(s-1)*nvar+1 : s*nvar] θ_s = @view θ[(s-1)*npar+1 : s*npar] y1_s = @view y1[(s-1)*nout+1 : s*nout] w_s = _get_obj_weight(adj1, s) @simd for k in eachindex(f.itr) - @inbounds shessian!( + shessian!( y1_s, y2, f.f, f.itr[k], x_s, θ_s, f.f.comp2, offset2(f, k), w_s, adj2, ) end @@ -727,13 +727,13 @@ function shessian!(y1, y2, f, x, θ, adj1s::V, adj2) where {V<:AbstractVector} end end function shessian!(y1, y2, f, x::AbstractArray, θ::AbstractArray, adj1s::AbstractVector, adj2, nb::Integer, nvar::Integer, npar::Integer, ncon::Integer, nout::Integer, ::Nothing = nothing) - for s in 1:nb + @inbounds for s in 1:nb x_s = @view x[(s-1)*nvar+1 : s*nvar] θ_s = @view θ[(s-1)*npar+1 : s*npar] y1_s = @view y1[(s-1)*nout+1 : s*nout] a_s = @view adj1s[(s-1)*ncon+1 : s*ncon] @simd for k in eachindex(f.itr) - @inbounds shessian!( + shessian!( y1_s, y2, f.f, f.itr[k], x_s, θ_s, f.f.comp2, offset2(f, k), a_s[offset0(f, k)], adj2, ) diff --git a/src/jacobian.jl b/src/jacobian.jl index aa05b5697..fe0057b92 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -126,12 +126,12 @@ function sjacobian!(y1, y2, f, x, θ, adj) end end function sjacobian!(y1, y2, f, x::AbstractArray, θ::AbstractArray, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, ::Nothing = nothing) - for s in 1:nb + @inbounds for s in 1:nb x_s = @view x[(s-1)*nvar+1 : s*nvar] θ_s = @view θ[(s-1)*npar+1 : s*npar] y1_s = @view y1[(s-1)*nout+1 : s*nout] @simd for i in eachindex(f.itr) - @inbounds sjacobian!( + sjacobian!( y1_s, y2, f.f, diff --git a/src/nlp.jl b/src/nlp.jl index 4ab97ba99..e859c22d0 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -1354,79 +1354,42 @@ end # once with views of the full arrays (no overhead). # ============================================================================ -function obj(m::AbstractExaModel, x::AbstractVector) - return _obj(m.objs, x, m.θ) +function obj(m::AbstractExaModel{T}, x::AbstractVector) where {T} + bf = fill!(similar(x, T, 1), zero(T)) + _obj!(bf, m.objs, x, m.θ, 1, length(x), length(m.θ)) + return @inbounds bf[1] end -# Stub for KA extension override -function _eval_objbuffer! end - -@inline function _obj((obj, objs...), x, θ) - s = _obj(objs, x, θ) - for i in obj.itr - s += obj.f(i, x, θ) - end - return s -end - -@inline _obj(obj::Tuple{}, x, θ) = zero(eltype(x)) - -# Batch versions — used by BatchExaModel (loop over instances with views) -@inline function _obj_b((obj, objs...), x, θ, nb, nvar, npar) - s = _obj_b(objs, x, θ, nb, nvar, npar) - for si in 1:nb - x_s = @view x[(si-1)*nvar+1 : si*nvar] - θ_s = @view θ[(si-1)*npar+1 : si*npar] - for i in obj.itr - s += obj.f(i, x_s, θ_s) - end - end - return s -end -@inline _obj_b(obj::Tuple{}, x, θ, nb, nvar, npar) = zero(eltype(x)) - -# Per-instance obj values (for batch obj!) @inline function _obj!(bf, (obj, objs...), x, θ, nb, nvar, npar, backend = nothing) _obj!(bf, objs, x, θ, nb, nvar, npar, backend) - _obj_batch!(bf, obj, x, θ, nb, nvar, npar, backend) + _obj!(bf, obj, x, θ, nb, nvar, npar, backend) end @inline _obj!(bf, ::Tuple{}, x, θ, nb, nvar, npar, backend = nothing) = nothing -@inline function _obj_batch!(bf, obj, x, θ, nb, nvar, npar, ::Nothing) - for s in 1:nb +@inline function _obj!(bf, obj, x, θ, nb, nvar, npar, ::Nothing) + @inbounds for s in 1:nb x_s = @view x[(s-1)*nvar+1 : s*nvar] θ_s = @view θ[(s-1)*npar+1 : s*npar] for i in obj.itr - @inbounds bf[s] += obj.f(i, x_s, θ_s) + bf[s] += obj.f(i, x_s, θ_s) end end end function cons_nln!(m::AbstractExaModel, x::AbstractVector, g::AbstractVector) fill!(g, zero(eltype(g))) - _cons_nln!(m.cons, x, m.θ, g) + _cons_nln!(m.cons, x, m.θ, g, 1, length(x), length(m.θ), length(g)) return g end -@inline function _cons_nln!(cons::Tuple, x, θ, g) - con = first(cons) - _cons_nln!(Base.tail(cons), x, θ, g) - @simd for i in eachindex(con.itr) - g[offset0(con, i)] += con.f(con.itr[i], x, θ) - end +@inline function _cons_nln!(cons::Tuple, x, θ, g, nb, nvar, npar, ncon, backend = nothing) + _cons_nln!(Base.tail(cons), x, θ, g, nb, nvar, npar, ncon, backend) + _cons_nln!(first(cons), x, θ, g, nb, nvar, npar, ncon, backend) end -_cons_nln!(cons::Tuple{}, x, θ, g) = nothing +_cons_nln!(cons::Tuple{}, x, θ, g, nb, nvar, npar, ncon, backend = nothing) = nothing -# Batch versions — used by BatchExaModel -@inline function _cons_nln_b!(cons::Tuple, x, θ, g, nb, nvar, npar, ncon, backend = nothing) - con = first(cons) - _cons_nln_b!(Base.tail(cons), x, θ, g, nb, nvar, npar, ncon, backend) - _cons_nln_batch!(con, x, θ, g, nb, nvar, npar, ncon, backend) -end -_cons_nln_b!(cons::Tuple{}, x, θ, g, nb, nvar, npar, ncon, backend = nothing) = nothing - -@inline function _cons_nln_batch!(con, x, θ, g, nb, nvar, npar, ncon, ::Nothing) - for s in 1:nb +@inline function _cons_nln!(con, x, θ, g, nb, nvar, npar, ncon, ::Nothing) + @inbounds for s in 1:nb x_s = @view x[(s-1)*nvar+1 : s*nvar] θ_s = @view θ[(s-1)*npar+1 : s*npar] g_s = @view g[(s-1)*ncon+1 : s*ncon] @@ -1440,39 +1403,25 @@ end function grad!(m::AbstractExaModel, x::AbstractVector, f::AbstractVector) fill!(f, zero(eltype(f))) - _grad!(m.objs, x, m.θ, f) + _grad!(m.objs, x, m.θ, f, 1, length(x), length(m.θ)) return f end -@inline function _grad!(objs::Tuple, x, θ, f) - _grad!(Base.tail(objs), x, θ, f) - gradient!(f, first(objs), x, θ, one(eltype(f))) -end -_grad!(objs::Tuple{}, x, θ, f) = nothing - -# Batch versions — used by BatchExaModel -@inline function _grad_b!(objs::Tuple, x, θ, f, nb, nvar, npar, backend = nothing) - _grad_b!(Base.tail(objs), x, θ, f, nb, nvar, npar, backend) +@inline function _grad!(objs::Tuple, x, θ, f, nb, nvar, npar, backend = nothing) + _grad!(Base.tail(objs), x, θ, f, nb, nvar, npar, backend) gradient!(f, first(objs), x, θ, one(eltype(f)), nb, nvar, npar, backend) end -_grad_b!(objs::Tuple{}, x, θ, f, nb, nvar, npar, backend = nothing) = nothing +_grad!(objs::Tuple{}, x, θ, f, nb, nvar, npar, backend = nothing) = nothing function jac_coord!(m::AbstractExaModel, x::AbstractVector, jac::AbstractVector) fill!(jac, zero(eltype(jac))) - _jac_coord!(m.cons, x, m.θ, jac) + _jac_coord!(m.cons, x, m.θ, jac, 1, length(x), length(m.θ), length(jac)) return jac end -_jac_coord!(cons::Tuple{}, x, θ, jac) = nothing -@inline function _jac_coord!(cons::Tuple, x, θ, jac) - _jac_coord!(Base.tail(cons), x, θ, jac) - sjacobian!(jac, nothing, first(cons), x, θ, one(eltype(jac))) -end - -# Batch versions — used by BatchExaModel -_jac_coord_b!(cons::Tuple{}, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) = nothing -@inline function _jac_coord_b!(cons::Tuple, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) - _jac_coord_b!(Base.tail(cons), x, θ, jac, nb, nvar, npar, nnzj, backend) +_jac_coord!(cons::Tuple{}, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) = nothing +@inline function _jac_coord!(cons::Tuple, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) + _jac_coord!(Base.tail(cons), x, θ, jac, nb, nvar, npar, nnzj, backend) sjacobian!(jac, nothing, first(cons), x, θ, one(eltype(jac)), nb, nvar, npar, nnzj, backend) end @@ -1507,7 +1456,7 @@ function hess_coord!( obj_weight = one(eltype(x)), ) fill!(hess, zero(eltype(hess))) - _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight) + _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight, 1, length(x), length(m.θ), length(hess)) return hess end @@ -1519,33 +1468,20 @@ function hess_coord!( obj_weight = one(eltype(x)), ) fill!(hess, zero(eltype(hess))) - _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight) - _con_hess_coord!(m.cons, x, m.θ, y, hess, obj_weight) + _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight, 1, length(x), length(m.θ), length(hess)) + _con_hess_coord!(m.cons, x, m.θ, y, hess, 1, length(x), length(m.θ), length(y), length(hess)) return hess end -_obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight) = nothing -@inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight) - _obj_hess_coord!(Base.tail(objs), x, θ, hess, obj_weight) - shessian!(hess, nothing, first(objs), x, θ, obj_weight, zero(eltype(hess))) -end - -_con_hess_coord!(cons::Tuple{}, x, θ, y, hess, obj_weight) = nothing -@inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess, obj_weight) - _con_hess_coord!(Base.tail(cons), x, θ, y, hess, obj_weight) - shessian!(hess, nothing, first(cons), x, θ, y, zero(eltype(hess))) -end - -# Batch versions — used by BatchExaModel -_obj_hess_coord_b!(objs::Tuple{}, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) = nothing -@inline function _obj_hess_coord_b!(objs::Tuple, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) - _obj_hess_coord_b!(Base.tail(objs), x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend) +_obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) = nothing +@inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) + _obj_hess_coord!(Base.tail(objs), x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend) shessian!(hess, nothing, first(objs), x, θ, obj_weight, zero(eltype(hess)), nb, nvar, npar, nnzh, backend) end -_con_hess_coord_b!(cons::Tuple{}, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) = nothing -@inline function _con_hess_coord_b!(cons::Tuple, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) - _con_hess_coord_b!(Base.tail(cons), x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend) +_con_hess_coord!(cons::Tuple{}, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) = nothing +@inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) + _con_hess_coord!(Base.tail(cons), x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend) shessian!(hess, nothing, first(cons), x, θ, y, zero(eltype(hess)), nb, nvar, npar, ncon, nnzh, backend) end @@ -1830,19 +1766,6 @@ model = ExaModel(c) """ BatchExaCore(nbatch::Integer; kwargs...) = ExaCore(; concrete = Val(true), nbatch = Val(nbatch), kwargs...) -# ============================================================================ -# get_model — defined after BatchNLPModels is loaded (see ExaModels.jl) -# ============================================================================ - -""" - get_model(model) - -Return a solver-ready NLP model. For a [`BatchExaModel`](@ref), returns a -[`FlattenNLPModel`](@ref) that presents all instances as a single flat -`AbstractNLPModel{T, Vector{T}}`. For a regular [`ExaModel`](@ref), returns -the model itself. -""" -get_model(model::ExaModel) = model """ var_indices(model, i) -> UnitRange @@ -1948,7 +1871,7 @@ function NLPModels.grad!(m::BatchExaModel{T}, bx::AbstractMatrix, bg::AbstractMa nb = get_nbatch(m) nvar = NLPModels.get_nvar(m) npar = Base.size(m.θ, 1) - _grad_b!(m.objs, vec(bx), vec(m.θ), vec(bg), nb, nvar, npar, getbackend(m)) + _grad!(m.objs, vec(bx), vec(m.θ), vec(bg), nb, nvar, npar, getbackend(m)) return bg end @@ -1958,7 +1881,7 @@ function NLPModels.cons!(m::BatchExaModel{T}, bx::AbstractMatrix, bc::AbstractMa nvar = NLPModels.get_nvar(m) ncon = NLPModels.get_ncon(m) npar = Base.size(m.θ, 1) - _cons_nln_b!(m.cons, vec(bx), vec(m.θ), vec(bc), nb, nvar, npar, ncon, getbackend(m)) + _cons_nln!(m.cons, vec(bx), vec(m.θ), vec(bc), nb, nvar, npar, ncon, getbackend(m)) return bc end @@ -1977,7 +1900,7 @@ function NLPModels.jac_coord!(m::BatchExaModel{T}, bx::AbstractMatrix, jvals::Ab nvar = NLPModels.get_nvar(m) npar = Base.size(m.θ, 1) nnzj = NLPModels.get_nnzj(m) - _jac_coord_b!(m.cons, vec(bx), vec(m.θ), vec(jvals), nb, nvar, npar, nnzj, getbackend(m)) + _jac_coord!(m.cons, vec(bx), vec(m.θ), vec(jvals), nb, nvar, npar, nnzj, getbackend(m)) return jvals end @@ -2005,8 +1928,8 @@ function NLPModels.hess_coord!( npar = Base.size(m.θ, 1) nnzh = NLPModels.get_nnzh(m) backend = getbackend(m) - _obj_hess_coord_b!(m.objs, vec(bx), vec(m.θ), vec(hvals), obj_weight, nb, nvar, npar, nnzh, backend) - _con_hess_coord_b!(m.cons, vec(bx), vec(m.θ), vec(by), vec(hvals), nb, nvar, npar, ncon, nnzh, backend) + _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), obj_weight, nb, nvar, npar, nnzh, backend) + _con_hess_coord!(m.cons, vec(bx), vec(m.θ), vec(by), vec(hvals), nb, nvar, npar, ncon, nnzh, backend) return hvals end @@ -2016,7 +1939,7 @@ end _batch_vector_error(name, m) = throw(ArgumentError( "$name on batch ExaModel requires matrix arguments. " * - "Use the batch API or get_model(m) for the fused model.", + "Use the batch API or FlatNLPModel(m) for the fused model.", )) function obj(m::BatchExaModel, x::AbstractVector) diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index 33322a9c4..9080dfdaf 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -7,8 +7,9 @@ import NLPModels: obj, cons!, cons_nln!, grad!, jac_coord!, hess_coord!, jac_structure!, hess_structure! import NLPModels: obj! -import ExaModels: var_indices, cons_block_indices, get_model, get_nbatch, +import ExaModels: var_indices, cons_block_indices, get_nbatch, get_start, get_lvar, get_uvar, get_lcon, get_ucon +using ExaModels.BatchNLPModels: FlatNLPModel import NLPModelsIpopt: ipopt @@ -49,7 +50,7 @@ function test_obj() @test bf[1] ≈ 10.0 @test bf[2] ≈ 50.0 @test obj(model, bx) ≈ bf - flat = get_model(model) + flat = FlatNLPModel(model) @test sum(bf) ≈ obj(flat, vec(bx)) end @@ -61,7 +62,7 @@ function test_grad() @test bg[:, 1] ≈ [4.0, 8.0] @test bg[:, 2] ≈ [12.0, 16.0] g_flat = zeros(4) - grad!(get_model(model), vec(bx), g_flat) + grad!(FlatNLPModel(model), vec(bx), g_flat) @test vec(bg) ≈ g_flat end @@ -73,14 +74,14 @@ function test_cons() @test bc[:, 1] ≈ [-1.0, 0.0] @test bc[:, 2] ≈ [1.0, 2.0] c_flat = zeros(4) - cons_nln!(get_model(model), vec(bx), c_flat) + cons_nln!(FlatNLPModel(model), vec(bx), c_flat) @test vec(bc) ≈ c_flat end function test_jac_hess() model = build_batch_model() ns, nv = 2, 2 - flat = get_model(model) + flat = FlatNLPModel(model) bx = [1.0 3.0; 2.0 4.0] # --- Jacobian values --- @@ -114,7 +115,7 @@ function test_hess_obj_weight() nnzh = NLPModels.get_nnzh(model) bx = [1.0 3.0; 2.0 4.0] by = ones(nc, ns) - flat = get_model(model) + flat = FlatNLPModel(model) hvals_w1 = zeros(nnzh, ns) hess_coord!(model, bx, by, hvals_w1; obj_weight = 1.0) @@ -180,7 +181,7 @@ function test_multiple_constraints() c, _ = add_con(c, v[j] - θ[1] for j in 1:nv) c, _ = add_con(c, v[1] + v[2] + v[3] for _ in 1:1; ucon = 10.0) model = ExaModel(c) - flat = get_model(model) + flat = FlatNLPModel(model) nc = NLPModels.get_ncon(model) @test nc == nv + 1 @@ -234,7 +235,7 @@ function test_bounds() @test model.meta.lvar ≈ fill(0.0, nv, ns) @test model.meta.uvar ≈ fill(10.0, nv, ns) - flat = get_model(model) + flat = FlatNLPModel(model) @test NLPModels.get_nvar(flat) == nv * ns @test flat.meta.x0 ≈ fill(0.5, nv * ns) @test flat.meta.lvar ≈ fill(0.0, nv * ns) @@ -243,8 +244,8 @@ end function test_flatten_model() model = build_batch_model() - flat = get_model(model) - @test flat isa ExaModels.BatchNLPModels.FlattenNLPModel + flat = FlatNLPModel(model) + @test flat isa FlatNLPModel @test NLPModels.get_nvar(flat) == 2 * 2 @test NLPModels.get_ncon(flat) == 2 * 2 @@ -252,7 +253,6 @@ function test_flatten_model() c, x = add_var(c, 2) c, _ = add_obj(c, x[i]^2 for i in 1:2) m = ExaModel(c) - @test get_model(m) === m end function test_ipopt_simple() @@ -264,7 +264,7 @@ function test_ipopt_simple() c, _ = add_con(c, v[1] for _ in 1:1; lcon = 0.0, ucon = Inf) model = ExaModel(c) - result = ipopt(get_model(model); print_level = 0) + result = ipopt(FlatNLPModel(model); print_level = 0) @test result.status == :first_order for i in 1:ns @test result.solution[var_indices(model, i)] ≈ [2.0] atol = 1e-5 @@ -281,7 +281,7 @@ function test_ipopt_multi() c, _ = add_con(c, v[1] + v[2] for _ in 1:1; ucon = 10.0) model = ExaModel(c) - result = ipopt(get_model(model); print_level = 0) + result = ipopt(FlatNLPModel(model); print_level = 0) @test result.status == :first_order @test result.solution[var_indices(model, 1)] ≈ [1.0, 3.0] atol = 1e-5 @test result.solution[var_indices(model, 2)] ≈ [1.0, 3.0] atol = 1e-5 @@ -340,7 +340,7 @@ function test_add_con_aug() @add_con(c, g, v[j] for j in 1:nv; lcon = -10.0, ucon = 10.0) @add_con!(c, g, j => θ[1] * v[j] for j in 1:nv) model = ExaModel(c) - flat = get_model(model) + flat = FlatNLPModel(model) nc = NLPModels.get_ncon(model) @test nc == nv @@ -387,7 +387,7 @@ function test_add_expr() c, _ = add_obj(c, s[j] for j in 1:nv) c, _ = add_con(c, s[j] - v[j] for j in 1:nv; lcon = 0.0) model = ExaModel(c) - flat = get_model(model) + flat = FlatNLPModel(model) nc = NLPModels.get_ncon(model) @test nc == nv @@ -452,7 +452,7 @@ function test_batch_backend(backend) @add_con(c, g, v[j] - θ[1] for j in 1:nv; lcon = 0.0) @add_con!(c, g, j => θ[1] * v[j] for j in 1:nv) model = ExaModel(c) - flat = get_model(model) + flat = FlatNLPModel(model) nc = NLPModels.get_ncon(model) bx = ExaModels.convert_array(reshape(Float64.(1:(nv*ns)), nv, ns), backend) diff --git a/test/TwoStageTest/TwoStageTest.jl b/test/TwoStageTest/TwoStageTest.jl index be1dbea2a..aee0efa90 100644 --- a/test/TwoStageTest/TwoStageTest.jl +++ b/test/TwoStageTest/TwoStageTest.jl @@ -549,7 +549,7 @@ function runtests() @add_con(core, EachScenario(), (v[v_idx] - d[1] for (i, v_idx) in con_data); lcon = 0.0) model = ExaModel(core) - flat = ExaModels.get_model(model) + flat = ExaModels.FlatNLPModel(model) nvar = NLPModels.get_nvar(model) ncon = NLPModels.get_ncon(model) @@ -607,7 +607,7 @@ function runtests() @add_con(core, EachScenario(), (v[i] - d[1] for i in 1:ns_scen); lcon = 0.0) model = ExaModel(core) - flat = ExaModels.get_model(model) + flat = ExaModels.FlatNLPModel(model) result = ipopt(flat; print_level = 0) @test result.status == :first_order From cef3bc0896f0a40ad5f0dac92b8ce3454e3f19e9 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 22:01:25 -0400 Subject: [PATCH 17/50] Fix doc build: resolve method ambiguity and rename KA batch hooks - Rename per-element dispatch to _obj_eval! and _cons_nln_eval! to avoid ambiguity with Tuple-recursive _obj! and _cons_nln! methods - Update KA extension to use new names Co-Authored-By: Claude Opus 4.6 --- ext/ExaModelsKernelAbstractions.jl | 4 ++-- src/nlp.jl | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index dd86203e6..e1d6811d5 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -671,7 +671,7 @@ end # --- Objective --- -function ExaModels._obj_batch!(bf, obj, x, θ, nb, nvar, npar, backend::KernelAbstractions.Backend) +function ExaModels._obj_eval!(bf, obj, x, θ, nb, nvar, npar, backend::KernelAbstractions.Backend) nitr = length(obj.itr) if nitr > 0 buf = similar(x, nitr, nb) @@ -691,7 +691,7 @@ end # --- Constraints --- -function ExaModels._cons_nln_batch!(con, x, θ, g, nb, nvar, npar, ncon, backend::KernelAbstractions.Backend) +function ExaModels._cons_nln_eval!(con, x, θ, g, nb, nvar, npar, ncon, backend::KernelAbstractions.Backend) nitr = length(con.itr) if nitr > 0 kerf_con_batch(backend)(g, con.f, con.itr, x, θ, nvar, npar, ncon, nitr; ndrange = nb * nitr) diff --git a/src/nlp.jl b/src/nlp.jl index e859c22d0..b50e52ac4 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -1362,11 +1362,11 @@ end @inline function _obj!(bf, (obj, objs...), x, θ, nb, nvar, npar, backend = nothing) _obj!(bf, objs, x, θ, nb, nvar, npar, backend) - _obj!(bf, obj, x, θ, nb, nvar, npar, backend) + _obj_eval!(bf, obj, x, θ, nb, nvar, npar, backend) end @inline _obj!(bf, ::Tuple{}, x, θ, nb, nvar, npar, backend = nothing) = nothing -@inline function _obj!(bf, obj, x, θ, nb, nvar, npar, ::Nothing) +@inline function _obj_eval!(bf, obj, x, θ, nb, nvar, npar, ::Nothing) @inbounds for s in 1:nb x_s = @view x[(s-1)*nvar+1 : s*nvar] θ_s = @view θ[(s-1)*npar+1 : s*npar] @@ -1384,11 +1384,11 @@ end @inline function _cons_nln!(cons::Tuple, x, θ, g, nb, nvar, npar, ncon, backend = nothing) _cons_nln!(Base.tail(cons), x, θ, g, nb, nvar, npar, ncon, backend) - _cons_nln!(first(cons), x, θ, g, nb, nvar, npar, ncon, backend) + _cons_nln_eval!(first(cons), x, θ, g, nb, nvar, npar, ncon, backend) end _cons_nln!(cons::Tuple{}, x, θ, g, nb, nvar, npar, ncon, backend = nothing) = nothing -@inline function _cons_nln!(con, x, θ, g, nb, nvar, npar, ncon, ::Nothing) +@inline function _cons_nln_eval!(con, x, θ, g, nb, nvar, npar, ncon, ::Nothing) @inbounds for s in 1:nb x_s = @view x[(s-1)*nvar+1 : s*nvar] θ_s = @view θ[(s-1)*npar+1 : s*npar] From 91d88ca090f2e27fe450e62096e65bcbc118f28e Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 22:39:14 -0400 Subject: [PATCH 18/50] Fix FlatNLPModel GPU: preserve array type from batch model FlatNLPModel was hardcoded to Vector{T}, causing GPU data to be copied to CPU. Now infers VT from vec(model.meta.x0) so GPU batch models produce a FlatNLPModel{T, CuVector{T}} that keeps data on device. Co-Authored-By: Claude Opus 4.6 --- src/BatchNLPModels.jl | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/BatchNLPModels.jl b/src/BatchNLPModels.jl index 16b43a546..a422b42eb 100644 --- a/src/BatchNLPModels.jl +++ b/src/BatchNLPModels.jl @@ -166,9 +166,9 @@ All NLPModels callbacks delegate to the underlying batch model's matrix API. Construct a flat model from a batch model whose `meta.x0` is a matrix. """ -struct FlatNLPModel{T, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, Vector{T}} +struct FlatNLPModel{T, VT <: AbstractVector{T}, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, VT} batch::M - meta::NLPModelMeta{T, Vector{T}} + meta::NLPModelMeta{T, VT} counters::NLPModels.Counters end @@ -178,9 +178,11 @@ function FlatNLPModel(model::AbstractNLPModel{T}) where {T} ncon = NLPModels.get_ncon(model) * nb nnzj = NLPModels.get_nnzj(model) * nb nnzh = NLPModels.get_nnzh(model) * nb - meta = NLPModelMeta{T, Vector{T}}( + x0 = vec(model.meta.x0) + VT = typeof(x0) + meta = NLPModelMeta{T, VT}( nvar, - vec(model.meta.x0), vec(model.meta.lvar), vec(model.meta.uvar), + x0, vec(model.meta.lvar), vec(model.meta.uvar), Int[], Int[], Int[], Int[], collect(1:nvar), Int[], nvar, nvar, nvar, ncon, From 2d5324a6999d6251afbee0f2f843bd7e8292d113 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 22:43:54 -0400 Subject: [PATCH 19/50] Parameterize all batch tests by backend, use WrapperNLPModel for Ipopt All test functions now accept a backend keyword. Tests run for every entry in BACKENDS (CPU + any available GPU backends). Ipopt tests wrap FlatNLPModel in WrapperNLPModel for GPU compatibility. Co-Authored-By: Claude Opus 4.6 --- test/BatchTest/BatchTest.jl | 435 ++++++++++++++++-------------------- 1 file changed, 197 insertions(+), 238 deletions(-) diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index 9080dfdaf..d6f7cd6c9 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -8,7 +8,7 @@ import NLPModels: hess_structure! import NLPModels: obj! import ExaModels: var_indices, cons_block_indices, get_nbatch, - get_start, get_lvar, get_uvar, get_lcon, get_ucon + get_start, get_lvar, get_uvar, get_lcon, get_ucon, WrapperNLPModel using ExaModels.BatchNLPModels: FlatNLPModel import NLPModelsIpopt: ipopt @@ -20,8 +20,8 @@ using Adapt # Helper: build a standard test problem # ============================================================================ -function build_batch_model(; ns=2, nv=2, θ_val=[2.0]) - c = BatchExaCore(ns) +function build_batch_model(; ns=2, nv=2, θ_val=[2.0], backend = nothing) + c = BatchExaCore(ns; backend) @add_var(c, v, nv) @add_par(c, θ, θ_val) c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) @@ -30,11 +30,20 @@ function build_batch_model(; ns=2, nv=2, θ_val=[2.0]) end # ============================================================================ -# Extract test logic into functions to avoid Julia 1.12 GC/compiler segfault +# Helper utilities for backend-aware allocation # ============================================================================ -function test_construction() - model = build_batch_model(ns=3) +_ca(x, ::Nothing) = x +_ca(x, backend) = ExaModels.convert_array(x, backend) +_to_cpu(x) = Array(x) +_to_cpu(x::Array) = x + +# ============================================================================ +# Test functions — all accept backend parameter +# ============================================================================ + +function test_construction(; backend = nothing) + model = build_batch_model(ns=3; backend) @test get_nbatch(model) == 3 @test NLPModels.get_nvar(model) == 2 @test NLPModels.get_ncon(model) == 2 @@ -42,69 +51,73 @@ function test_construction() @test size(model.meta.x0) == (2, 3) end -function test_obj() - model = build_batch_model() - bx = [1.0 3.0; 2.0 4.0] - bf = zeros(2) +function test_obj(; backend = nothing) + model = build_batch_model(; backend) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + bf = _ca(zeros(2), backend) obj!(model, bx, bf) - @test bf[1] ≈ 10.0 - @test bf[2] ≈ 50.0 - @test obj(model, bx) ≈ bf + @test _to_cpu(bf)[1] ≈ 10.0 + @test _to_cpu(bf)[2] ≈ 50.0 + @test _to_cpu(obj(model, bx)) ≈ _to_cpu(bf) flat = FlatNLPModel(model) - @test sum(bf) ≈ obj(flat, vec(bx)) + @test sum(_to_cpu(bf)) ≈ obj(flat, vec(bx)) end -function test_grad() - model = build_batch_model() - bx = [1.0 3.0; 2.0 4.0] - bg = zeros(2, 2) +function test_grad(; backend = nothing) + model = build_batch_model(; backend) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + bg = _ca(zeros(2, 2), backend) grad!(model, bx, bg) - @test bg[:, 1] ≈ [4.0, 8.0] - @test bg[:, 2] ≈ [12.0, 16.0] - g_flat = zeros(4) - grad!(FlatNLPModel(model), vec(bx), g_flat) - @test vec(bg) ≈ g_flat + bg_cpu = _to_cpu(bg) + @test bg_cpu[:, 1] ≈ [4.0, 8.0] + @test bg_cpu[:, 2] ≈ [12.0, 16.0] + flat = FlatNLPModel(model) + g_flat = _ca(zeros(4), backend) + grad!(flat, vec(bx), g_flat) + @test vec(bg_cpu) ≈ _to_cpu(g_flat) end -function test_cons() - model = build_batch_model() - bx = [1.0 3.0; 2.0 4.0] - bc = zeros(2, 2) +function test_cons(; backend = nothing) + model = build_batch_model(; backend) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + bc = _ca(zeros(2, 2), backend) cons!(model, bx, bc) - @test bc[:, 1] ≈ [-1.0, 0.0] - @test bc[:, 2] ≈ [1.0, 2.0] - c_flat = zeros(4) - cons_nln!(FlatNLPModel(model), vec(bx), c_flat) - @test vec(bc) ≈ c_flat + bc_cpu = _to_cpu(bc) + @test bc_cpu[:, 1] ≈ [-1.0, 0.0] + @test bc_cpu[:, 2] ≈ [1.0, 2.0] + flat = FlatNLPModel(model) + c_flat = _ca(zeros(4), backend) + cons_nln!(flat, vec(bx), c_flat) + @test vec(bc_cpu) ≈ _to_cpu(c_flat) end -function test_jac_hess() - model = build_batch_model() +function test_jac_hess(; backend = nothing) + model = build_batch_model(; backend) ns, nv = 2, 2 flat = FlatNLPModel(model) - bx = [1.0 3.0; 2.0 4.0] + bx = _ca([1.0 3.0; 2.0 4.0], backend) # --- Jacobian values --- nnzj = NLPModels.get_nnzj(model) - jvals = zeros(nnzj, ns) + jvals = _ca(zeros(nnzj, ns), backend) jac_coord!(model, bx, jvals) - jvals_flat = zeros(NLPModels.get_nnzj(flat)) + jvals_flat = _ca(zeros(NLPModels.get_nnzj(flat)), backend) jac_coord!(flat, vec(bx), jvals_flat) - @test vec(jvals) ≈ jvals_flat + @test vec(_to_cpu(jvals)) ≈ _to_cpu(jvals_flat) # --- Hessian values --- nnzh = NLPModels.get_nnzh(model) - by = ones(nv, ns) - hvals = zeros(nnzh, ns) + by = _ca(ones(nv, ns), backend) + hvals = _ca(zeros(nnzh, ns), backend) hess_coord!(model, bx, by, hvals) - hvals_flat = zeros(NLPModels.get_nnzh(flat)) + hvals_flat = _ca(zeros(NLPModels.get_nnzh(flat)), backend) hess_coord!(flat, vec(bx), vec(by), hvals_flat) - @test vec(hvals) ≈ hvals_flat + @test vec(_to_cpu(hvals)) ≈ _to_cpu(hvals_flat) end -function test_hess_obj_weight() +function test_hess_obj_weight(; backend = nothing) ns, nv = 2, 2 - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, v, nv) @add_par(c, θ, [2.0]) c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) @@ -113,28 +126,28 @@ function test_hess_obj_weight() nc = NLPModels.get_ncon(model) nnzh = NLPModels.get_nnzh(model) - bx = [1.0 3.0; 2.0 4.0] - by = ones(nc, ns) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + by = _ca(ones(nc, ns), backend) flat = FlatNLPModel(model) - hvals_w1 = zeros(nnzh, ns) + hvals_w1 = _ca(zeros(nnzh, ns), backend) hess_coord!(model, bx, by, hvals_w1; obj_weight = 1.0) - hvals_flat_w1 = zeros(NLPModels.get_nnzh(flat)) + hvals_flat_w1 = _ca(zeros(NLPModels.get_nnzh(flat)), backend) hess_coord!(flat, vec(bx), vec(by), hvals_flat_w1; obj_weight = 1.0) - @test vec(hvals_w1) ≈ hvals_flat_w1 + @test vec(_to_cpu(hvals_w1)) ≈ _to_cpu(hvals_flat_w1) - hvals_w2 = zeros(nnzh, ns) + hvals_w2 = _ca(zeros(nnzh, ns), backend) hess_coord!(model, bx, by, hvals_w2; obj_weight = 2.0) - hvals_flat_w2 = zeros(NLPModels.get_nnzh(flat)) + hvals_flat_w2 = _ca(zeros(NLPModels.get_nnzh(flat)), backend) hess_coord!(flat, vec(bx), vec(by), hvals_flat_w2; obj_weight = 2.0) - @test vec(hvals_w2) ≈ hvals_flat_w2 + @test vec(_to_cpu(hvals_w2)) ≈ _to_cpu(hvals_flat_w2) - @test hvals_w1 != hvals_w2 + @test _to_cpu(hvals_w1) != _to_cpu(hvals_w2) end -function test_hess_vector_obj_weight() +function test_hess_vector_obj_weight(; backend = nothing) ns, nv = 2, 2 - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, v, nv) @add_par(c, θ, [2.0]) c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) @@ -143,38 +156,34 @@ function test_hess_vector_obj_weight() nc = NLPModels.get_ncon(model) nnzh = NLPModels.get_nnzh(model) - bx = [1.0 3.0; 2.0 4.0] - by = ones(nc, ns) + bx = _ca([1.0 3.0; 2.0 4.0], backend) + by = _ca(ones(nc, ns), backend) # Vector obj_weight = [w1, w2] - wvec = [1.5, 3.0] - hvals_vec = zeros(nnzh, ns) + wvec = _ca([1.5, 3.0], backend) + hvals_vec = _ca(zeros(nnzh, ns), backend) hess_coord!(model, bx, by, hvals_vec; obj_weight = wvec) # Uniform scalar weights for comparison - hvals_w1 = zeros(nnzh, ns) - hess_coord!(model, bx, by, hvals_w1; obj_weight = wvec[1]) - hvals_w2 = zeros(nnzh, ns) - hess_coord!(model, bx, by, hvals_w2; obj_weight = wvec[2]) - - # With uniform weight, both instances get the same obj contribution. - # With vector weight, instance 1 gets w1, instance 2 gets w2. - # The constraint hessian is unaffected by obj_weight, so it is the same. - # Check: vector result differs from both uniform-scalar results. - @test hvals_vec != hvals_w1 - @test hvals_vec != hvals_w2 + hvals_w1 = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals_w1; obj_weight = 1.5) + hvals_w2 = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals_w2; obj_weight = 3.0) + + @test _to_cpu(hvals_vec) != _to_cpu(hvals_w1) + @test _to_cpu(hvals_vec) != _to_cpu(hvals_w2) # Verify consistency: uniform weight = special case of vector weight - hvals_uniform = zeros(nnzh, ns) - hess_coord!(model, bx, by, hvals_uniform; obj_weight = [2.0, 2.0]) - hvals_scalar = zeros(nnzh, ns) + hvals_uniform = _ca(zeros(nnzh, ns), backend) + hess_coord!(model, bx, by, hvals_uniform; obj_weight = _ca([2.0, 2.0], backend)) + hvals_scalar = _ca(zeros(nnzh, ns), backend) hess_coord!(model, bx, by, hvals_scalar; obj_weight = 2.0) - @test hvals_uniform ≈ hvals_scalar + @test _to_cpu(hvals_uniform) ≈ _to_cpu(hvals_scalar) end -function test_multiple_constraints() +function test_multiple_constraints(; backend = nothing) ns, nv = 2, 3 - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, v, nv) @add_par(c, θ, [1.0]) c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) @@ -187,84 +196,79 @@ function test_multiple_constraints() @test nc == nv + 1 @test get_nbatch(model) == ns - bx = reshape(Float64[1, 2, 3, 4, 5, 6], nv, ns) + bx = _ca(reshape(Float64[1, 2, 3, 4, 5, 6], nv, ns), backend) # cons! - bc = zeros(nc, ns) + bc = _ca(zeros(nc, ns), backend) cons!(model, bx, bc) - c_flat = zeros(nc * ns) + c_flat = _ca(zeros(nc * ns), backend) cons_nln!(flat, vec(bx), c_flat) - @test vec(bc) ≈ c_flat + @test vec(_to_cpu(bc)) ≈ _to_cpu(c_flat) # jac nnzj = NLPModels.get_nnzj(model) - jvals = zeros(nnzj, ns) + jvals = _ca(zeros(nnzj, ns), backend) jac_coord!(model, bx, jvals) - jvals_flat = zeros(NLPModels.get_nnzj(flat)) + jvals_flat = _ca(zeros(NLPModels.get_nnzj(flat)), backend) jac_coord!(flat, vec(bx), jvals_flat) - @test vec(jvals) ≈ jvals_flat + @test vec(_to_cpu(jvals)) ≈ _to_cpu(jvals_flat) # hess nnzh = NLPModels.get_nnzh(model) - by = ones(nc, ns) - hvals = zeros(nnzh, ns) + by = _ca(ones(nc, ns), backend) + hvals = _ca(zeros(nnzh, ns), backend) hess_coord!(model, bx, by, hvals) - hvals_flat = zeros(NLPModels.get_nnzh(flat)) + hvals_flat = _ca(zeros(NLPModels.get_nnzh(flat)), backend) hess_coord!(flat, vec(bx), vec(by), hvals_flat) - @test vec(hvals) ≈ hvals_flat + @test vec(_to_cpu(hvals)) ≈ _to_cpu(hvals_flat) end -function test_error_guards() - model = build_batch_model() - x_vec = ones(2) +function test_error_guards(; backend = nothing) + model = build_batch_model(; backend) + x_vec = _ca(ones(2), backend) @test_throws ArgumentError obj(model, x_vec) - @test_throws ArgumentError cons!(model, x_vec, zeros(2)) - @test_throws ArgumentError grad!(model, x_vec, zeros(2)) + @test_throws ArgumentError cons!(model, x_vec, _ca(zeros(2), backend)) + @test_throws ArgumentError grad!(model, x_vec, _ca(zeros(2), backend)) end -function test_bounds() +function test_bounds(; backend = nothing) ns, nv = 2, 2 - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, v, nv; start = 0.5, lvar = 0.0, uvar = 10.0) c, _ = add_obj(c, v[j]^2 for j in 1:nv) c, _ = add_con(c, v[j] for j in 1:nv; lcon = 0.0, ucon = 100.0) model = ExaModel(c) @test size(model.meta.x0) == (nv, ns) - @test model.meta.x0 ≈ fill(0.5, nv, ns) - @test model.meta.lvar ≈ fill(0.0, nv, ns) - @test model.meta.uvar ≈ fill(10.0, nv, ns) + @test _to_cpu(model.meta.x0) ≈ fill(0.5, nv, ns) + @test _to_cpu(model.meta.lvar) ≈ fill(0.0, nv, ns) + @test _to_cpu(model.meta.uvar) ≈ fill(10.0, nv, ns) flat = FlatNLPModel(model) @test NLPModels.get_nvar(flat) == nv * ns - @test flat.meta.x0 ≈ fill(0.5, nv * ns) - @test flat.meta.lvar ≈ fill(0.0, nv * ns) - @test flat.meta.uvar ≈ fill(10.0, nv * ns) + @test _to_cpu(flat.meta.x0) ≈ fill(0.5, nv * ns) + @test _to_cpu(flat.meta.lvar) ≈ fill(0.0, nv * ns) + @test _to_cpu(flat.meta.uvar) ≈ fill(10.0, nv * ns) end -function test_flatten_model() - model = build_batch_model() +function test_flatten_model(; backend = nothing) + model = build_batch_model(; backend) flat = FlatNLPModel(model) @test flat isa FlatNLPModel @test NLPModels.get_nvar(flat) == 2 * 2 @test NLPModels.get_ncon(flat) == 2 * 2 - - c = ExaCore(concrete = Val(true)) - c, x = add_var(c, 2) - c, _ = add_obj(c, x[i]^2 for i in 1:2) - m = ExaModel(c) end -function test_ipopt_simple() +function test_ipopt_simple(; backend = nothing) ns, nv = 3, 1 - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, v, nv) @add_par(c, θ, [2.0]) c, _ = add_obj(c, (v[1] - θ[1])^2 for _ in 1:1) c, _ = add_con(c, v[1] for _ in 1:1; lcon = 0.0, ucon = Inf) model = ExaModel(c) - result = ipopt(FlatNLPModel(model); print_level = 0) + result = ipopt(WrapperNLPModel(FlatNLPModel(model)); print_level = 0) @test result.status == :first_order for i in 1:ns @test result.solution[var_indices(model, i)] ≈ [2.0] atol = 1e-5 @@ -272,49 +276,51 @@ function test_ipopt_simple() @test isapprox(result.objective, 0.0; atol = 1e-8) end -function test_ipopt_multi() +function test_ipopt_multi(; backend = nothing) ns, nv = 2, 2 - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, v, nv) @add_par(c, θ, [1.0, 3.0]) c, _ = add_obj(c, (v[j] - θ[j])^2 for j in 1:nv) c, _ = add_con(c, v[1] + v[2] for _ in 1:1; ucon = 10.0) model = ExaModel(c) - result = ipopt(FlatNLPModel(model); print_level = 0) + result = ipopt(WrapperNLPModel(FlatNLPModel(model)); print_level = 0) @test result.status == :first_order @test result.solution[var_indices(model, 1)] ≈ [1.0, 3.0] atol = 1e-5 @test result.solution[var_indices(model, 2)] ≈ [1.0, 3.0] atol = 1e-5 @test isapprox(result.objective, 0.0; atol = 1e-8) end -function test_set_parameter() +function test_set_parameter(; backend = nothing) ns, nv = 2, 1 - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, v, nv) @add_par(c, θ, [2.0]) c, _ = add_obj(c, (v[1] - θ[1])^2 for _ in 1:1) model = ExaModel(c) # Default: both instances share θ = [2.0] - bx = reshape([1.0, 3.0], nv, ns) - bf = zeros(ns) + bx = _ca(reshape([1.0, 3.0], nv, ns), backend) + bf = _ca(zeros(ns), backend) obj!(model, bx, bf) - @test bf[1] ≈ 1.0 # (1-2)^2 - @test bf[2] ≈ 1.0 # (3-2)^2 + bf_cpu = _to_cpu(bf) + @test bf_cpu[1] ≈ 1.0 # (1-2)^2 + @test bf_cpu[2] ≈ 1.0 # (3-2)^2 # Update parameters via set_value! on the model set_value!(model, θ, [10.0]) - bf2 = zeros(ns) + bf2 = _ca(zeros(ns), backend) obj!(model, bx, bf2) - @test bf2[1] ≈ (1.0 - 10.0)^2 - @test bf2[2] ≈ (3.0 - 10.0)^2 + bf2_cpu = _to_cpu(bf2) + @test bf2_cpu[1] ≈ (1.0 - 10.0)^2 + @test bf2_cpu[2] ≈ (3.0 - 10.0)^2 end -function test_multidim_vars() +function test_multidim_vars(; backend = nothing) ns = 2 nh, nc = 3, 2 # multi-dimensional variable - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, w, nh, nc; start = zeros(nh, nc)) @add_par(c, θ, [1.0]) c, _ = add_obj(c, w[i, j]^2 for i in 1:nh, j in 1:nc) @@ -323,16 +329,18 @@ function test_multidim_vars() @test NLPModels.get_nvar(model) == nh * nc @test get_nbatch(model) == ns - bx = reshape(Float64.(1:(nh*nc*ns)), nh * nc, ns) - bf = zeros(ns) + bx = _ca(reshape(Float64.(1:(nh*nc*ns)), nh * nc, ns), backend) + bf = _ca(zeros(ns), backend) obj!(model, bx, bf) - @test bf[1] ≈ sum(bx[:, 1] .^ 2) - @test bf[2] ≈ sum(bx[:, 2] .^ 2) + bx_cpu = _to_cpu(bx) + bf_cpu = _to_cpu(bf) + @test bf_cpu[1] ≈ sum(bx_cpu[:, 1] .^ 2) + @test bf_cpu[2] ≈ sum(bx_cpu[:, 2] .^ 2) end -function test_add_con_aug() +function test_add_con_aug(; backend = nothing) ns, nv = 2, 3 - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, v, nv) @add_par(c, θ, [1.0]) c, _ = add_obj(c, v[j]^2 for j in 1:nv) @@ -345,40 +353,42 @@ function test_add_con_aug() nc = NLPModels.get_ncon(model) @test nc == nv - bx = reshape(Float64.(1:(nv*ns)), nv, ns) + bx = _ca(reshape(Float64.(1:(nv*ns)), nv, ns), backend) # cons! — augmented constraint = v[j] + θ[1]*v[j] = 2*v[j] - bc = zeros(nc, ns) + bc = _ca(zeros(nc, ns), backend) cons!(model, bx, bc) - c_flat = zeros(nc * ns) + c_flat = _ca(zeros(nc * ns), backend) cons_nln!(flat, vec(bx), c_flat) - @test vec(bc) ≈ c_flat + bc_cpu = _to_cpu(bc) + bx_cpu = _to_cpu(bx) + @test vec(bc_cpu) ≈ _to_cpu(c_flat) # Each constraint value should be v[j] + 1.0*v[j] = 2*v[j] - @test bc[:, 1] ≈ 2.0 .* bx[:, 1] - @test bc[:, 2] ≈ 2.0 .* bx[:, 2] + @test bc_cpu[:, 1] ≈ 2.0 .* bx_cpu[:, 1] + @test bc_cpu[:, 2] ≈ 2.0 .* bx_cpu[:, 2] # jac nnzj = NLPModels.get_nnzj(model) - jvals = zeros(nnzj, ns) + jvals = _ca(zeros(nnzj, ns), backend) jac_coord!(model, bx, jvals) - jvals_flat = zeros(NLPModels.get_nnzj(flat)) + jvals_flat = _ca(zeros(NLPModels.get_nnzj(flat)), backend) jac_coord!(flat, vec(bx), jvals_flat) - @test vec(jvals) ≈ jvals_flat + @test vec(_to_cpu(jvals)) ≈ _to_cpu(jvals_flat) # hess nnzh = NLPModels.get_nnzh(model) - by = ones(nc, ns) - hvals = zeros(nnzh, ns) + by = _ca(ones(nc, ns), backend) + hvals = _ca(zeros(nnzh, ns), backend) hess_coord!(model, bx, by, hvals) - hvals_flat = zeros(NLPModels.get_nnzh(flat)) + hvals_flat = _ca(zeros(NLPModels.get_nnzh(flat)), backend) hess_coord!(flat, vec(bx), vec(by), hvals_flat) - @test vec(hvals) ≈ hvals_flat + @test vec(_to_cpu(hvals)) ≈ _to_cpu(hvals_flat) end -function test_add_expr() +function test_add_expr(; backend = nothing) ns, nv = 2, 3 - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, v, nv) @add_par(c, θ, [2.0]) @@ -392,132 +402,81 @@ function test_add_expr() nc = NLPModels.get_ncon(model) @test nc == nv - bx = reshape(Float64.(1:(nv*ns)), nv, ns) + bx = _ca(reshape(Float64.(1:(nv*ns)), nv, ns), backend) # obj — should be sum of θ[1]*v[j]^2 = 2*v[j]^2 - bf = zeros(ns) + bf = _ca(zeros(ns), backend) obj!(model, bx, bf) - @test bf[1] ≈ 2.0 * sum(bx[:, 1] .^ 2) - @test bf[2] ≈ 2.0 * sum(bx[:, 2] .^ 2) - @test sum(bf) ≈ obj(flat, vec(bx)) + bx_cpu = _to_cpu(bx) + bf_cpu = _to_cpu(bf) + @test bf_cpu[1] ≈ 2.0 * sum(bx_cpu[:, 1] .^ 2) + @test bf_cpu[2] ≈ 2.0 * sum(bx_cpu[:, 2] .^ 2) + @test sum(bf_cpu) ≈ obj(flat, vec(bx)) # cons — s[j] - v[j] = 2*v[j]^2 - v[j] - bc = zeros(nc, ns) + bc = _ca(zeros(nc, ns), backend) cons!(model, bx, bc) - c_flat = zeros(nc * ns) + c_flat = _ca(zeros(nc * ns), backend) cons_nln!(flat, vec(bx), c_flat) - @test vec(bc) ≈ c_flat - @test bc[:, 1] ≈ 2.0 .* bx[:, 1] .^ 2 .- bx[:, 1] + bc_cpu = _to_cpu(bc) + @test vec(bc_cpu) ≈ _to_cpu(c_flat) + @test bc_cpu[:, 1] ≈ 2.0 .* bx_cpu[:, 1] .^ 2 .- bx_cpu[:, 1] # grad - bg = zeros(nv, ns) + bg = _ca(zeros(nv, ns), backend) grad!(model, bx, bg) - g_flat = zeros(nv * ns) + g_flat = _ca(zeros(nv * ns), backend) grad!(flat, vec(bx), g_flat) - @test vec(bg) ≈ g_flat + @test vec(_to_cpu(bg)) ≈ _to_cpu(g_flat) end -function test_per_instance_accessors() +function test_per_instance_accessors(; backend = nothing) ns, nv = 2, 3 - c = BatchExaCore(ns) + c = BatchExaCore(ns; backend) @add_var(c, v, nv; start = 1.0, lvar = -5.0, uvar = 5.0) c, _ = add_obj(c, v[j]^2 for j in 1:nv) c, con = add_con(c, v[j] for j in 1:nv; lcon = 0.0, ucon = 10.0) model = ExaModel(c) # Test per-instance variable accessors - @test get_start(model, v, 1) ≈ fill(1.0, nv) - @test get_lvar(model, v, 1) ≈ fill(-5.0, nv) - @test get_uvar(model, v, 2) ≈ fill(5.0, nv) + @test _to_cpu(get_start(model, v, 1)) ≈ fill(1.0, nv) + @test _to_cpu(get_lvar(model, v, 1)) ≈ fill(-5.0, nv) + @test _to_cpu(get_uvar(model, v, 2)) ≈ fill(5.0, nv) # Test per-instance constraint accessors - @test get_lcon(model, con, 1) ≈ fill(0.0, nv) - @test get_ucon(model, con, 2) ≈ fill(10.0, nv) + @test _to_cpu(get_lcon(model, con, 1)) ≈ fill(0.0, nv) + @test _to_cpu(get_ucon(model, con, 2)) ≈ fill(10.0, nv) # Test cons_block_indices @test cons_block_indices(model, 1) == 1:nv @test cons_block_indices(model, 2) == (nv+1):(2*nv) end -# ============================================================================ -# Backend-parameterized tests — verify batch evaluation on all backends (incl. KA/GPU) -# ============================================================================ - -function test_batch_backend(backend) - ns, nv = 2, 3 - c = BatchExaCore(ns; backend) - @add_var(c, v, nv) - @add_par(c, θ, [2.0]) - c, _ = add_obj(c, θ[1] * v[j]^2 for j in 1:nv) - @add_con(c, g, v[j] - θ[1] for j in 1:nv; lcon = 0.0) - @add_con!(c, g, j => θ[1] * v[j] for j in 1:nv) - model = ExaModel(c) - flat = FlatNLPModel(model) - - nc = NLPModels.get_ncon(model) - bx = ExaModels.convert_array(reshape(Float64.(1:(nv*ns)), nv, ns), backend) - - # obj - bf_gpu = obj(model, bx) - @test sum(Array(bf_gpu)) ≈ obj(flat, vec(bx)) - - # grad - bg = ExaModels.convert_array(zeros(nv, ns), backend) - grad!(model, bx, bg) - g_flat = ExaModels.convert_array(zeros(nv * ns), backend) - grad!(flat, vec(bx), g_flat) - @test vec(Array(bg)) ≈ Array(g_flat) - - # cons - bc = ExaModels.convert_array(zeros(nc, ns), backend) - cons!(model, bx, bc) - c_flat = ExaModels.convert_array(zeros(nc * ns), backend) - cons_nln!(flat, vec(bx), c_flat) - @test vec(Array(bc)) ≈ Array(c_flat) - - # jac - nnzj = NLPModels.get_nnzj(model) - jvals = ExaModels.convert_array(zeros(nnzj, ns), backend) - jac_coord!(model, bx, jvals) - jvals_flat = ExaModels.convert_array(zeros(NLPModels.get_nnzj(flat)), backend) - jac_coord!(flat, vec(bx), jvals_flat) - @test vec(Array(jvals)) ≈ Array(jvals_flat) - - # hess - nnzh = NLPModels.get_nnzh(model) - by = ExaModels.convert_array(ones(nc, ns), backend) - hvals = ExaModels.convert_array(zeros(nnzh, ns), backend) - hess_coord!(model, bx, by, hvals) - hvals_flat = ExaModels.convert_array(zeros(NLPModels.get_nnzh(flat)), backend) - hess_coord!(flat, vec(bx), vec(by), hvals_flat) - @test vec(Array(hvals)) ≈ Array(hvals_flat) -end - # ============================================================================ function runtests() return @testset "Batch ExaModel" begin - @testset "Construction" test_construction() - @testset "obj!" test_obj() - @testset "grad!" test_grad() - @testset "cons!" test_cons() - @testset "jac and hess" test_jac_hess() - @testset "hess obj_weight" test_hess_obj_weight() - @testset "hess vector obj_weight" test_hess_vector_obj_weight() - @testset "Multiple constraints" test_multiple_constraints() - @testset "Error guards" test_error_guards() - @testset "Bounds" test_bounds() - @testset "flatten_model" test_flatten_model() - @testset "Ipopt simple" test_ipopt_simple() - @testset "Ipopt multi" test_ipopt_multi() - @testset "set_value!" test_set_parameter() - @testset "Multidim vars" test_multidim_vars() - @testset "add_con!" test_add_con_aug() - @testset "add_expr" test_add_expr() - @testset "Per-instance accessors" test_per_instance_accessors() for backend in BACKENDS - backend === nothing && continue - @testset "Backend: $(typeof(backend))" test_batch_backend(backend) + @testset "Backend: $(something(backend, :CPU))" begin + @testset "Construction" test_construction(; backend) + @testset "obj!" test_obj(; backend) + @testset "grad!" test_grad(; backend) + @testset "cons!" test_cons(; backend) + @testset "jac and hess" test_jac_hess(; backend) + @testset "hess obj_weight" test_hess_obj_weight(; backend) + @testset "hess vector obj_weight" test_hess_vector_obj_weight(; backend) + @testset "Multiple constraints" test_multiple_constraints(; backend) + @testset "Error guards" test_error_guards(; backend) + @testset "Bounds" test_bounds(; backend) + @testset "flatten_model" test_flatten_model(; backend) + @testset "Ipopt simple" test_ipopt_simple(; backend) + @testset "Ipopt multi" test_ipopt_multi(; backend) + @testset "set_value!" test_set_parameter(; backend) + @testset "Multidim vars" test_multidim_vars(; backend) + @testset "add_con!" test_add_con_aug(; backend) + @testset "add_expr" test_add_expr(; backend) + @testset "Per-instance accessors" test_per_instance_accessors(; backend) + end end end end From 0507e8803ba3b83911c03f93329a36caf5e04bd0 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 22:45:34 -0400 Subject: [PATCH 20/50] Revert FlatNLPModel to Vector{T} meta, fix GPU jac_structure! dispatch FlatNLPModel must use Vector{T} meta so NLPModels structure queries stay on CPU. GPU data conversion is handled by WrapperNLPModel. Co-Authored-By: Claude Opus 4.6 --- src/BatchNLPModels.jl | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src/BatchNLPModels.jl b/src/BatchNLPModels.jl index a422b42eb..c074f9bff 100644 --- a/src/BatchNLPModels.jl +++ b/src/BatchNLPModels.jl @@ -166,9 +166,9 @@ All NLPModels callbacks delegate to the underlying batch model's matrix API. Construct a flat model from a batch model whose `meta.x0` is a matrix. """ -struct FlatNLPModel{T, VT <: AbstractVector{T}, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, VT} +struct FlatNLPModel{T, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, Vector{T}} batch::M - meta::NLPModelMeta{T, VT} + meta::NLPModelMeta{T, Vector{T}} counters::NLPModels.Counters end @@ -178,15 +178,17 @@ function FlatNLPModel(model::AbstractNLPModel{T}) where {T} ncon = NLPModels.get_ncon(model) * nb nnzj = NLPModels.get_nnzj(model) * nb nnzh = NLPModels.get_nnzh(model) * nb - x0 = vec(model.meta.x0) - VT = typeof(x0) - meta = NLPModelMeta{T, VT}( + meta = NLPModelMeta{T, Vector{T}}( nvar, - x0, vec(model.meta.lvar), vec(model.meta.uvar), + Vector{T}(vec(model.meta.x0)), + Vector{T}(vec(model.meta.lvar)), + Vector{T}(vec(model.meta.uvar)), Int[], Int[], Int[], Int[], collect(1:nvar), Int[], nvar, nvar, nvar, ncon, - vec(model.meta.y0), vec(model.meta.lcon), vec(model.meta.ucon), + Vector{T}(vec(model.meta.y0)), + Vector{T}(vec(model.meta.lcon)), + Vector{T}(vec(model.meta.ucon)), Int[], Int[], Int[], Int[], Int[], Int[], nvar, nnzj, 0, nnzj, nnzh, 0, ncon, Int[], collect(1:ncon), From 32e95b5a993d2da7b50ca0f3095d1853d13ef9b5 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 22:49:55 -0400 Subject: [PATCH 21/50] Fix FlatNLPModel GPU: use jac_nln_structure!/jac_nln_coord!, preserve VT - Restore GPU-aware VT so WrapperNLPModel creates GPU buffers - Use jac_nln_structure! and jac_nln_coord! (what NLPModels dispatches to for models with only nonlinear constraints) instead of jac_structure!/ jac_coord! Co-Authored-By: Claude Opus 4.6 --- src/BatchNLPModels.jl | 36 +++++++++++++++++------------------- 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/src/BatchNLPModels.jl b/src/BatchNLPModels.jl index c074f9bff..1dc6ef702 100644 --- a/src/BatchNLPModels.jl +++ b/src/BatchNLPModels.jl @@ -157,7 +157,7 @@ end # ============================================================================ """ - FlatNLPModel{T, M} <: AbstractNLPModel{T, Vector{T}} + FlatNLPModel{T, VT, M} <: AbstractNLPModel{T, VT} Wrapper that presents a batch NLP model as a flat (Vector-based) NLP model. All NLPModels callbacks delegate to the underlying batch model's matrix API. @@ -166,9 +166,9 @@ All NLPModels callbacks delegate to the underlying batch model's matrix API. Construct a flat model from a batch model whose `meta.x0` is a matrix. """ -struct FlatNLPModel{T, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, Vector{T}} +struct FlatNLPModel{T, VT <: AbstractVector{T}, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, VT} batch::M - meta::NLPModelMeta{T, Vector{T}} + meta::NLPModelMeta{T, VT} counters::NLPModels.Counters end @@ -178,17 +178,15 @@ function FlatNLPModel(model::AbstractNLPModel{T}) where {T} ncon = NLPModels.get_ncon(model) * nb nnzj = NLPModels.get_nnzj(model) * nb nnzh = NLPModels.get_nnzh(model) * nb - meta = NLPModelMeta{T, Vector{T}}( + x0 = vec(model.meta.x0) + VT = typeof(x0) + meta = NLPModelMeta{T, VT}( nvar, - Vector{T}(vec(model.meta.x0)), - Vector{T}(vec(model.meta.lvar)), - Vector{T}(vec(model.meta.uvar)), + x0, vec(model.meta.lvar), vec(model.meta.uvar), Int[], Int[], Int[], Int[], collect(1:nvar), Int[], nvar, nvar, nvar, ncon, - Vector{T}(vec(model.meta.y0)), - Vector{T}(vec(model.meta.lcon)), - Vector{T}(vec(model.meta.ucon)), + vec(model.meta.y0), vec(model.meta.lcon), vec(model.meta.ucon), Int[], Int[], Int[], Int[], Int[], Int[], nvar, nnzj, 0, nnzj, nnzh, 0, ncon, Int[], collect(1:ncon), @@ -222,7 +220,7 @@ function NLPModels.cons_nln!(m::FlatNLPModel{T}, x::AbstractVector, c::AbstractV return c end -function NLPModels.jac_structure!(m::FlatNLPModel, rows::AbstractVector, cols::AbstractVector) +function NLPModels.jac_nln_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}) nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) ncon = NLPModels.get_ncon(m.batch) @@ -234,19 +232,19 @@ function NLPModels.jac_structure!(m::FlatNLPModel, rows::AbstractVector, cols::A NLPModels.jac_structure!(m.batch, r1, c1) # Replicate for each instance with shifted indices - for s in 2:nb + @inbounds for s in 2:nb offset = (s - 1) * nnzj row_shift = (s - 1) * ncon col_shift = (s - 1) * nvar for k in 1:nnzj - @inbounds rows[offset + k] = r1[k] + row_shift - @inbounds cols[offset + k] = c1[k] + col_shift + rows[offset + k] = r1[k] + row_shift + cols[offset + k] = c1[k] + col_shift end end return rows, cols end -function NLPModels.jac_coord!(m::FlatNLPModel{T}, x::AbstractVector, jvals::AbstractVector) where {T} +function NLPModels.jac_nln_coord!(m::FlatNLPModel{T}, x::AbstractVector, jvals::AbstractVector) where {T} nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) nnzj = NLPModels.get_nnzj(m.batch) @@ -254,7 +252,7 @@ function NLPModels.jac_coord!(m::FlatNLPModel{T}, x::AbstractVector, jvals::Abst return jvals end -function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector, cols::AbstractVector) +function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}) nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) nnzh = NLPModels.get_nnzh(m.batch) @@ -265,12 +263,12 @@ function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector, cols:: NLPModels.hess_structure!(m.batch, r1, c1) # Replicate for each instance with shifted indices - for s in 2:nb + @inbounds for s in 2:nb offset = (s - 1) * nnzh shift = (s - 1) * nvar for k in 1:nnzh - @inbounds rows[offset + k] = r1[k] + shift - @inbounds cols[offset + k] = c1[k] + shift + rows[offset + k] = r1[k] + shift + cols[offset + k] = c1[k] + shift end end return rows, cols From b451da22f0791f79acdcf62400e5192e9abcf72d Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 22:54:43 -0400 Subject: [PATCH 22/50] Fix FlatNLPModel GPU structure queries: compute on CPU, copyto! output Structure queries (_jac_structure!, _hess_structure!) use scalar indexing internally. When rows/cols are CuArrays, this triggers GPU scalar indexing errors. Now computes structure on CPU vectors and copies to output arrays. Co-Authored-By: Claude Opus 4.6 --- src/BatchNLPModels.jl | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/src/BatchNLPModels.jl b/src/BatchNLPModels.jl index 1dc6ef702..a29c1e24a 100644 --- a/src/BatchNLPModels.jl +++ b/src/BatchNLPModels.jl @@ -226,9 +226,11 @@ function NLPModels.jac_nln_structure!(m::FlatNLPModel, rows::AbstractVector{<:In ncon = NLPModels.get_ncon(m.batch) nnzj = NLPModels.get_nnzj(m.batch) - # Get per-instance structure - r1 = @view rows[1:nnzj] - c1 = @view cols[1:nnzj] + # Compute structure on CPU (structure queries use scalar indexing) + r_cpu = Vector{Int}(undef, nnzj * nb) + c_cpu = Vector{Int}(undef, nnzj * nb) + r1 = @view r_cpu[1:nnzj] + c1 = @view c_cpu[1:nnzj] NLPModels.jac_structure!(m.batch, r1, c1) # Replicate for each instance with shifted indices @@ -237,10 +239,12 @@ function NLPModels.jac_nln_structure!(m::FlatNLPModel, rows::AbstractVector{<:In row_shift = (s - 1) * ncon col_shift = (s - 1) * nvar for k in 1:nnzj - rows[offset + k] = r1[k] + row_shift - cols[offset + k] = c1[k] + col_shift + r_cpu[offset + k] = r1[k] + row_shift + c_cpu[offset + k] = c1[k] + col_shift end end + copyto!(rows, r_cpu) + copyto!(cols, c_cpu) return rows, cols end @@ -257,9 +261,11 @@ function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integ nvar = NLPModels.get_nvar(m.batch) nnzh = NLPModels.get_nnzh(m.batch) - # Get per-instance structure - r1 = @view rows[1:nnzh] - c1 = @view cols[1:nnzh] + # Compute structure on CPU (structure queries use scalar indexing) + r_cpu = Vector{Int}(undef, nnzh * nb) + c_cpu = Vector{Int}(undef, nnzh * nb) + r1 = @view r_cpu[1:nnzh] + c1 = @view c_cpu[1:nnzh] NLPModels.hess_structure!(m.batch, r1, c1) # Replicate for each instance with shifted indices @@ -267,10 +273,12 @@ function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integ offset = (s - 1) * nnzh shift = (s - 1) * nvar for k in 1:nnzh - rows[offset + k] = r1[k] + shift - cols[offset + k] = c1[k] + shift + r_cpu[offset + k] = r1[k] + shift + c_cpu[offset + k] = c1[k] + shift end end + copyto!(rows, r_cpu) + copyto!(cols, c_cpu) return rows, cols end From bede0d29babb673900a88f1f7b2fef1efc88ad2f Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 23:28:51 -0400 Subject: [PATCH 23/50] Fix GPU scalar indexing in BatchExaModel structure queries BatchExaModel.jac_structure! and hess_structure! now pass getbackend(m) to backend-aware _jac_structure!/_obj_hess_structure!/_con_hess_structure! so the KA extension can intercept and use GPU kernels instead of scalar-indexing f.itr arrays. Also converts obj_weight to model type T in batch hess_coord! to prevent Metal InvalidIRError from Float64 values on Float32 models. FlatNLPModel structure queries now use device-native temp arrays for the per-instance query, then copy to CPU for replication. Co-Authored-By: Claude Opus 4.6 --- src/BatchNLPModels.jl | 28 ++++++++++++++++++++-------- src/nlp.jl | 15 +++++++++++---- 2 files changed, 31 insertions(+), 12 deletions(-) diff --git a/src/BatchNLPModels.jl b/src/BatchNLPModels.jl index a29c1e24a..4e9e44442 100644 --- a/src/BatchNLPModels.jl +++ b/src/BatchNLPModels.jl @@ -226,12 +226,18 @@ function NLPModels.jac_nln_structure!(m::FlatNLPModel, rows::AbstractVector{<:In ncon = NLPModels.get_ncon(m.batch) nnzj = NLPModels.get_nnzj(m.batch) - # Compute structure on CPU (structure queries use scalar indexing) + # Use device-native arrays for the per-instance query (avoids scalar indexing on GPU itr) + r1_dev = similar(m.batch.meta.x0, Int, nnzj) + c1_dev = similar(m.batch.meta.x0, Int, nnzj) + NLPModels.jac_structure!(m.batch, r1_dev, c1_dev) + + # Copy to CPU for replication + r1 = Vector{Int}(r1_dev) + c1 = Vector{Int}(c1_dev) r_cpu = Vector{Int}(undef, nnzj * nb) c_cpu = Vector{Int}(undef, nnzj * nb) - r1 = @view r_cpu[1:nnzj] - c1 = @view c_cpu[1:nnzj] - NLPModels.jac_structure!(m.batch, r1, c1) + copyto!(r_cpu, 1, r1, 1, nnzj) + copyto!(c_cpu, 1, c1, 1, nnzj) # Replicate for each instance with shifted indices @inbounds for s in 2:nb @@ -261,12 +267,18 @@ function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integ nvar = NLPModels.get_nvar(m.batch) nnzh = NLPModels.get_nnzh(m.batch) - # Compute structure on CPU (structure queries use scalar indexing) + # Use device-native arrays for the per-instance query (avoids scalar indexing on GPU) + r1_dev = similar(m.batch.meta.x0, Int, nnzh) + c1_dev = similar(m.batch.meta.x0, Int, nnzh) + NLPModels.hess_structure!(m.batch, r1_dev, c1_dev) + + # Copy to CPU for replication + r1 = Vector{Int}(r1_dev) + c1 = Vector{Int}(c1_dev) r_cpu = Vector{Int}(undef, nnzh * nb) c_cpu = Vector{Int}(undef, nnzh * nb) - r1 = @view r_cpu[1:nnzh] - c1 = @view c_cpu[1:nnzh] - NLPModels.hess_structure!(m.batch, r1, c1) + copyto!(r_cpu, 1, r1, 1, nnzh) + copyto!(c_cpu, 1, c1, 1, nnzh) # Replicate for each instance with shifted indices @inbounds for s in 2:nb diff --git a/src/nlp.jl b/src/nlp.jl index b50e52ac4..03d56cbae 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -1329,6 +1329,8 @@ _jac_structure!(T, cons::Tuple{}, rows, cols) = nothing _jac_structure!(T, Base.tail(cons), rows, cols) sjacobian!(rows, cols, first(cons), NaNSource{T}(), NaNSource{T}(), T(NaN)) end +# Backend-aware fallbacks (KA extension overrides these for GPU backends) +_jac_structure!(T, ::Nothing, cons, rows, cols) = _jac_structure!(T, cons, rows, cols) function hess_structure!(m::AbstractExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T _obj_hess_structure!(T, m.objs, rows, cols) @@ -1341,12 +1343,16 @@ _obj_hess_structure!(T, objs::Tuple{}, rows, cols) = nothing _obj_hess_structure!(T, Base.tail(objs), rows, cols) shessian!(rows, cols, first(objs), NaNSource{T}(), NaNSource{T}(), T(NaN), T(NaN)) end +# Backend-aware fallback (KA extension overrides for GPU backends) +_obj_hess_structure!(T, ::Nothing, objs, rows, cols) = _obj_hess_structure!(T, objs, rows, cols) _con_hess_structure!(T, cons::Tuple{}, rows, cols) = nothing @inline function _con_hess_structure!(T, cons::Tuple, rows, cols) _con_hess_structure!(T, Base.tail(cons), rows, cols) shessian!(rows, cols, first(cons), NaNSource{T}(), NaNSource{T}(), T(NaN), T(NaN)) end +# Backend-aware fallback (KA extension overrides for GPU backends) +_con_hess_structure!(T, ::Nothing, cons, rows, cols) = _con_hess_structure!(T, cons, rows, cols) # ============================================================================ # Batch-aware evaluation — all low-level functions loop over 1:nb, @@ -1890,7 +1896,7 @@ function NLPModels.jac_structure!( rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}, ) where {T} - _jac_structure!(T, m.cons, rows, cols) + _jac_structure!(T, getbackend(m), m.cons, rows, cols) return rows, cols end @@ -1909,8 +1915,9 @@ function NLPModels.hess_structure!( rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}, ) where {T} - _obj_hess_structure!(T, m.objs, rows, cols) - _con_hess_structure!(T, m.cons, rows, cols) + backend = getbackend(m) + _obj_hess_structure!(T, backend, m.objs, rows, cols) + _con_hess_structure!(T, backend, m.cons, rows, cols) return rows, cols end @@ -1928,7 +1935,7 @@ function NLPModels.hess_coord!( npar = Base.size(m.θ, 1) nnzh = NLPModels.get_nnzh(m) backend = getbackend(m) - _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), obj_weight, nb, nvar, npar, nnzh, backend) + _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), obj_weight isa Number ? T(obj_weight) : obj_weight, nb, nvar, npar, nnzh, backend) _con_hess_coord!(m.cons, vec(bx), vec(m.θ), vec(by), vec(hvals), nb, nvar, npar, ncon, nnzh, backend) return hvals end From dca97ac7c2fa8119ee835a05612b5c0351aafe4f Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sat, 18 Apr 2026 23:44:42 -0400 Subject: [PATCH 24/50] Fix KA extension structure functions to extend ExaModels methods The KA extension's _jac_structure!, _obj_hess_structure!, and _con_hess_structure! were local functions, not methods on ExaModels. BatchExaModel.jac_structure! now calls ExaModels._jac_structure! with getbackend(m), so the extension must add methods to ExaModels' functions for GPU dispatch to work. Co-Authored-By: Claude Opus 4.6 --- ext/ExaModelsKernelAbstractions.jl | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index e1d6811d5..fe70b9c13 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -143,13 +143,13 @@ function ExaModels.jac_structure!( cols::V, ) where {T,VT,E<:KAExtension,V<:AbstractVector} if !isempty(rows) - _jac_structure!(T, m.ext.backend, m.cons, rows, cols) + ExaModels._jac_structure!(T, m.ext.backend, m.cons, rows, cols) end return rows, cols end -_jac_structure!(T, backend, ::Tuple{}, rows, cols) = nothing -function _jac_structure!(T, backend, (con, cons...), rows, cols) - _jac_structure!(T, backend, cons, rows, cols) +ExaModels._jac_structure!(T, backend, ::Tuple{}, rows, cols) = nothing +function ExaModels._jac_structure!(T, backend, (con, cons...), rows, cols) + ExaModels._jac_structure!(T, backend, cons, rows, cols) ExaModels.sjacobian!(backend, rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN)) end @@ -160,20 +160,20 @@ function ExaModels.hess_structure!( cols::V, ) where {T,VT,E<:KAExtension,V<:AbstractVector} if !isempty(rows) - _obj_hess_structure!(T, m.ext.backend, m.objs, rows, cols) - _con_hess_structure!(T, m.ext.backend, m.cons, rows, cols) - end + ExaModels._obj_hess_structure!(T, m.ext.backend, m.objs, rows, cols) + ExaModels._con_hess_structure!(T, m.ext.backend, m.cons, rows, cols) + end return rows, cols end -_obj_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing -function _obj_hess_structure!(T, backend, (obj, objs...), rows, cols) - _obj_hess_structure!(T, backend, objs, rows, cols) +ExaModels._obj_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing +function ExaModels._obj_hess_structure!(T, backend, (obj, objs...), rows, cols) + ExaModels._obj_hess_structure!(T, backend, objs, rows, cols) ExaModels.shessian!(backend, rows, cols, obj, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), T(NaN)) end -_con_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing -function _con_hess_structure!(T, backend, (con, cons...), rows, cols) - _con_hess_structure!(T, backend, cons, rows, cols) +ExaModels._con_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing +function ExaModels._con_hess_structure!(T, backend, (con, cons...), rows, cols) + ExaModels._con_hess_structure!(T, backend, cons, rows, cols) ExaModels.shessian!(backend, rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), T(NaN)) end From f13f9986a05b37d136a818d36724de8ca5ec6402 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sun, 19 Apr 2026 07:17:38 -0400 Subject: [PATCH 25/50] Fix remaining bare _jac_structure! references in KA extension Three calls in the KAExtension constructor at lines 61/64/65 still used unqualified _jac_structure!, _obj_hess_structure!, _con_hess_structure! instead of ExaModels._ prefixed versions. Co-Authored-By: Claude Opus 4.6 --- ext/ExaModelsKernelAbstractions.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index fe70b9c13..a62d13747 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -58,11 +58,11 @@ function ExaModels.build_extension( jacsparsityi = similar(c.x0, Tuple{Tuple{Int,Int},Int}, c.nnzj) hesssparsityi = similar(c.x0, Tuple{Tuple{Int,Int},Int}, c.nnzh) - _jac_structure!(T, c.backend, c.cons, jacsparsityi, nothing) + ExaModels._jac_structure!(T, c.backend, c.cons, jacsparsityi, nothing) jacsparsityj = copy(jacsparsityi) - _obj_hess_structure!(T, c.backend, c.obj, hesssparsityi, nothing) - _con_hess_structure!(T, c.backend, c.cons, hesssparsityi, nothing) + ExaModels._obj_hess_structure!(T, c.backend, c.obj, hesssparsityi, nothing) + ExaModels._con_hess_structure!(T, c.backend, c.cons, hesssparsityi, nothing) hesssparsityj = copy(hesssparsityi) if !isempty(jacsparsityi) From b46e204040ad664e8dc3837023150d8d3f17bf2a Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Sun, 19 Apr 2026 08:59:51 -0400 Subject: [PATCH 26/50] Fix append! for matrix bounds in batch models Add _reshape_to_match to correctly handle array bounds regardless of input shape. For non-batch (Vector target), matrices are vecd. For batch (Matrix target), matrices are reshaped to match trailing dims. This fixes the DimensionMismatch when passing matrix lvar/uvar to batch models, while preserving the existing behavior for non-batch models that receive matrix-shaped comprehension results. Co-Authored-By: Claude Opus 4.6 --- src/nlp.jl | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/nlp.jl b/src/nlp.jl index 03d56cbae..96557738a 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -680,6 +680,17 @@ function _expand_to_shape(col::AbstractVector{T}, trailing::Tuple) where {T} return repeat(reshape(col, :, ntuple(_ -> 1, length(trailing))...), 1, trailing...) end +# Reshape an array to have the given trailing dimensions. +# For vectors, delegates to _expand_to_shape. For matrices/higher-dim arrays, +# reshapes using total elements / trailing to infer the first dimension. +_reshape_to_match(arr::AbstractVector, trailing::Tuple{}) = arr +_reshape_to_match(arr::AbstractVector, trailing::Tuple) = _expand_to_shape(arr, trailing) +_reshape_to_match(arr::AbstractArray, trailing::Tuple{}) = vec(arr) +function _reshape_to_match(arr::AbstractArray, trailing::Tuple) + n = length(arr) ÷ prod(trailing) + return reshape(arr, n, trailing...) +end + function append!(backend, a, b::Number, lb) lb == 0 && return a new_part = fill(eltype(a)(b), lb, _trailing_dims(a)...) @@ -689,8 +700,7 @@ end function append!(backend, a, b::AbstractArray, lb) lb == 0 && return a arr = convert_array(b, backend) - col = vec(arr) - return cat(a, _expand_to_shape(col, _trailing_dims(a)); dims = 1) + return cat(a, _reshape_to_match(arr, _trailing_dims(a)); dims = 1) end function append!(backend, a, b::Base.Generator, lb) From e58c1928ecc790bd4cdbfffedf43f5e1973ff851 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 19:27:33 -0400 Subject: [PATCH 27/50] Add batched GPU/CPU kernels and obj! for BatchExaModel - Rename _batch suffix from KA kernels (kerg, kerj, kerh, kerh2, kerg_sparse) - Remove AbstractVector constraints from GPU dispatch functions so NaNSource works for structure detection - Add GPU obj! override for BatchExaModel - Add batched CPU fallback for sgradient! - Add kerg_sparse kernel for GPU sparse gradient Co-Authored-By: Claude Sonnet 4.6 --- Project.toml | 22 ++- ext/ExaModelsKernelAbstractions.jl | 211 ++++++++--------------------- src/gradient.jl | 11 ++ test/Project.toml | 3 + 4 files changed, 89 insertions(+), 158 deletions(-) diff --git a/Project.toml b/Project.toml index 7970eb552..32733d8f7 100644 --- a/Project.toml +++ b/Project.toml @@ -4,38 +4,42 @@ version = "0.10.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" +NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" +Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" +PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" [weakdeps] Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" NLPModelsIpopt = "f4238b75-b362-5c4c-b852-0801c9a21d71" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" -SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" -oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" # oneAPI version >= 2.6 is needed to use sort! on GPU arrays -# OptimalControl = "5f98b655-cc9a-415a-b60e-744165666948" +oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] ExaModelsIpopt = ["MathOptInterface", "NLPModelsIpopt"] ExaModelsJuMP = "JuMP" ExaModelsKernelAbstractions = "KernelAbstractions" -ExaModelsOneAPI = "oneAPI" ExaModelsMOI = "MathOptInterface" ExaModelsMadNLP = ["MadNLP", "MathOptInterface"] ExaModelsMetal = "Metal" +ExaModelsOneAPI = "oneAPI" ExaModelsOpenCL = "OpenCL" ExaModelsSpecialFunctions = "SpecialFunctions" -# ExaModelsOptimalControl = ["OptimalControl", "LinearAlgebra"] [compat] Adapt = "4" +ForwardDiff = "1.3.3" Ipopt = "1.11" JuMP = "1" KernelAbstractions = "0.9" @@ -44,8 +48,12 @@ MathOptInterface = "1.19" Metal = "1.9" NLPModels = "0.21.10" NLPModelsIpopt = "0.11" +NLPModelsJuMP = "0.13.5" +NLPModelsTest = "0.10.9" OpenCL = "0.10" +Percival = "0.7.6" +PowerModels = "0.21.5" SolverCore = "0.3" SpecialFunctions = "2" julia = "1.9" -oneAPI = "2" \ No newline at end of file +oneAPI = "2" diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index a62d13747..fa4776a2a 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -134,7 +134,7 @@ end _grad_structure!(T, backend, ::Tuple{}, gsparsity) = nothing function _grad_structure!(T, backend, (obj, objs...), gsparsity) _grad_structure!(T, backend, objs, gsparsity) - ExaModels.sgradient!(backend, gsparsity, obj, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN)) + ExaModels.sgradient!(gsparsity, obj, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), 1, 0, 0, length(gsparsity), backend) end function ExaModels.jac_structure!( @@ -150,7 +150,7 @@ end ExaModels._jac_structure!(T, backend, ::Tuple{}, rows, cols) = nothing function ExaModels._jac_structure!(T, backend, (con, cons...), rows, cols) ExaModels._jac_structure!(T, backend, cons, rows, cols) - ExaModels.sjacobian!(backend, rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN)) + ExaModels.sjacobian!(rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), 1, 0, 0, length(rows), backend) end @@ -169,12 +169,12 @@ end ExaModels._obj_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing function ExaModels._obj_hess_structure!(T, backend, (obj, objs...), rows, cols) ExaModels._obj_hess_structure!(T, backend, objs, rows, cols) - ExaModels.shessian!(backend, rows, cols, obj, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), T(NaN)) + ExaModels.shessian!(rows, cols, obj, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), T(NaN), 1, 0, 0, length(rows), backend) end ExaModels._con_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing function ExaModels._con_hess_structure!(T, backend, (con, cons...), rows, cols) ExaModels._con_hess_structure!(T, backend, cons, rows, cols) - ExaModels.shessian!(backend, rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), T(NaN)) + ExaModels.shessian!(rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), T(NaN), 1, 0, 0, length(rows), backend) end @@ -198,6 +198,19 @@ function _obj(backend, objbuffer, (obj, objs...), x, θ) end end +function ExaModels.obj!( + m::ExaModels.BatchExaModel{T,VT,E}, + bx::AbstractMatrix, + bf::AbstractVector, +) where {T,VT,E<:KAExtension} + fill!(bf, zero(T)) + nb = ExaModels.get_nbatch(m) + nvar = NLPModels.get_nvar(m) + npar = size(m.θ, 1) + ExaModels._obj!(bf, m.objs, vec(bx), vec(m.θ), nb, nvar, npar, m.ext.backend) + return bf +end + function ExaModels.cons_nln!( m::ExaModels.AbstractExaModel{T,VT,E}, @@ -261,7 +274,7 @@ end _grad!(backend, y, ::Tuple{}, x, θ) = nothing function _grad!(backend, y, (obj, objs...), x, θ) _grad!(backend, y, objs, x, θ) - ExaModels.sgradient!(backend, y, obj, x, θ, one(eltype(y))) + ExaModels.sgradient!(y, obj, x, θ, one(eltype(y)), 1, length(x), length(θ), length(y), backend) end function ExaModels.jac_coord!( @@ -276,7 +289,7 @@ end _jac_coord!(backend, y, ::Tuple{}, x, θ) = nothing function _jac_coord!(backend, y, (con, cons...), x, θ) _jac_coord!(backend, y, cons, x, θ) - ExaModels.sjacobian!(backend, y, nothing, con, x, θ, one(eltype(y))) + ExaModels.sjacobian!(y, nothing, con, x, θ, one(eltype(y)), 1, length(x), length(θ), length(y), backend) end function ExaModels.jprod_nln!( @@ -456,145 +469,14 @@ end _obj_hess_coord!(backend, hess, ::Tuple{}, x, θ, obj_weight) = nothing function _obj_hess_coord!(backend, hess, (obj, objs...), x, θ, obj_weight) _obj_hess_coord!(backend, hess, objs, x, θ, obj_weight) - ExaModels.shessian!(backend, hess, nothing, obj, x, θ, obj_weight, zero(eltype(hess))) + ExaModels.shessian!(hess, nothing, obj, x, θ, obj_weight, zero(eltype(hess)), 1, length(x), length(θ), length(hess), backend) end _con_hess_coord!(backend, hess, ::Tuple{}, x, θ, y) = nothing function _con_hess_coord!(backend, hess, (con, cons...), x, θ, y) _con_hess_coord!(backend, hess, cons, x, θ, y) - ExaModels.shessian!(backend, hess, nothing, con, x, θ, y, zero(eltype(hess))) -end - - -function ExaModels.sgradient!( - backend::B, - y, - f, - x, - θ, - adj, -) where {B<:KernelAbstractions.Backend} - - if !isempty(f.itr) - kerg(backend)(y, f.f, f.itr, x, θ, adj; ndrange = length(f.itr)) - end -end - -function ExaModels.sjacobian!( - backend::B, - y1, - y2, - f, - x, - θ, - adj, -) where {B<:KernelAbstractions.Backend} - if !isempty(f.itr) - kerj(backend)(y1, y2, f.f, f.itr, x, θ, adj, ExaModels._constraint_dims(f); ndrange = length(f.itr)) - end -end - -function ExaModels.shessian!( - backend::B, - y1, - y2, - f, - x, - θ, - adj, - adj2, -) where {B<:KernelAbstractions.Backend} - if !isempty(f.itr) - kerh(backend)(y1, y2, f.f, f.itr, x, θ, adj, adj2; ndrange = length(f.itr)) - end -end - -function ExaModels.shessian!( - backend::B, - y1, - y2, - f, - x, - θ, - adj::V, - adj2, -) where {B<:KernelAbstractions.Backend,V<:AbstractVector} - if !isempty(f.itr) - kerh2(backend)(y1, y2, f.f, f.itr, x, θ, adj, adj2, ExaModels._constraint_dims(f); ndrange = length(f.itr)) - end -end - -@kernel function kerh( - y1, - y2, - @Const(f), - @Const(itr), - @Const(x), - @Const(θ), - @Const(adj1), - @Const(adj2) -) - I = @index(Global) - @inbounds ExaModels.hrpass0( - f(itr[I], ExaModels.SecondAdjointNodeSource(x), θ), - f.comp2, - y1, - y2, - ExaModels.offset2(f, I), - 0, - adj1, - adj2, - ) -end - -@kernel function kerh2( - y1, - y2, - @Const(f), - @Const(itr), - @Const(x), - @Const(θ), - @Const(adjs1), - @Const(adj2), - @Const(dims) -) - I = @index(Global) - @inbounds ExaModels.hrpass0( - f(itr[I], ExaModels.SecondAdjointNodeSource(x), θ), - f.comp2, - y1, - y2, - ExaModels.offset2(f, I), - 0, - adjs1[ExaModels.offset0(f, itr, I, dims)], - adj2, - ) -end - -@kernel function kerj(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(dims)) - I = @index(Global) - @inbounds ExaModels.jrpass( - f(itr[I], ExaModels.AdjointNodeSource(x), θ), - f.comp1, - ExaModels.offset0(f, itr, I, dims), - y1, - y2, - ExaModels.offset1(f, I), - 0, - adj, - ) + ExaModels.shessian!(hess, nothing, con, x, θ, y, zero(eltype(hess)), 1, length(x), length(θ), length(y), length(hess), backend) end -@kernel function kerg(y, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj)) - I = @index(Global) - @inbounds ExaModels.grpass( - f(itr[I], ExaModels.AdjointNodeSource(x), θ), - f.comp1, - y, - ExaModels.offset1(f, I), - 0, - adj, - ) -end @kernel function kerf(y, @Const(f), @Const(itr), @Const(x), @Const(θ)) I = @index(Global) @@ -710,15 +592,15 @@ end # --- Gradient --- -function ExaModels.gradient!(y, f, x::AbstractVector, θ::AbstractVector, adj, nb::Integer, nvar::Integer, npar::Integer, backend::KernelAbstractions.Backend) +function ExaModels.gradient!(y, f, x, θ, adj, nb::Integer, nvar::Integer, npar::Integer, backend::KernelAbstractions.Backend) nitr = length(f.itr) if nitr > 0 - kerg_batch(backend)(y, f.f, f.itr, x, θ, adj, nvar, npar, nitr; ndrange = nb * nitr) + kerg(backend)(y, f.f, f.itr, x, θ, adj, nvar, npar, nitr; ndrange = nb * nitr) end return y end -@kernel function kerg_batch(y, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(nvar), @Const(npar), @Const(nitr)) +@kernel function kerg(y, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(nvar), @Const(npar), @Const(nitr)) I = @index(Global) s = (I - 1) ÷ nitr + 1 k = (I - 1) % nitr + 1 @@ -732,16 +614,43 @@ end ) end +# --- Sparse gradient --- + +function ExaModels.sgradient!(y, f, x, θ, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, backend::KernelAbstractions.Backend) + nitr = length(f.itr) + if nitr > 0 + kerg_sparse(backend)(y, f.f, f.itr, x, θ, adj, nvar, npar, nout, nitr; ndrange = nb * nitr) + end + return y +end + +@kernel function kerg_sparse(y, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(nvar), @Const(npar), @Const(nout), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + y_off = (s - 1) * nout + @inbounds ExaModels.grpass( + f(itr[k], ExaModels.AdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), + f.comp1, + ExaModels.OffsetVector(y, y_off), + ExaModels.offset1(f, k), + 0, + adj, + ) +end + # --- Jacobian --- -function ExaModels.sjacobian!(y1, y2, f, x::AbstractVector, θ::AbstractVector, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, backend::KernelAbstractions.Backend) +function ExaModels.sjacobian!(y1, y2, f, x, θ, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, backend::KernelAbstractions.Backend) nitr = length(f.itr) if nitr > 0 - kerj_batch(backend)(y1, y2, f.f, f.itr, x, θ, adj, ExaModels._constraint_dims(f), nvar, npar, nout, nitr; ndrange = nb * nitr) + kerj(backend)(y1, y2, f.f, f.itr, x, θ, adj, ExaModels._constraint_dims(f), nvar, npar, nout, nitr; ndrange = nb * nitr) end end -@kernel function kerj_batch(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(dims), @Const(nvar), @Const(npar), @Const(nout), @Const(nitr)) +@kernel function kerj(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj), @Const(dims), @Const(nvar), @Const(npar), @Const(nout), @Const(nitr)) I = @index(Global) s = (I - 1) ÷ nitr + 1 k = (I - 1) % nitr + 1 @@ -762,14 +671,14 @@ end # --- Hessian (objective) --- -function ExaModels.shessian!(y1, y2, f, x::AbstractVector, θ::AbstractVector, adj1, adj2, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, backend::KernelAbstractions.Backend) +function ExaModels.shessian!(y1, y2, f, x, θ, adj1, adj2, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, backend::KernelAbstractions.Backend) nitr = length(f.itr) if nitr > 0 - kerh_batch(backend)(y1, y2, f.f, f.itr, x, θ, adj1, adj2, nvar, npar, nout, nitr; ndrange = nb * nitr) + kerh(backend)(y1, y2, f.f, f.itr, x, θ, adj1, adj2, nvar, npar, nout, nitr; ndrange = nb * nitr) end end -@kernel function kerh_batch(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj1), @Const(adj2), @Const(nvar), @Const(npar), @Const(nout), @Const(nitr)) +@kernel function kerh(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj1), @Const(adj2), @Const(nvar), @Const(npar), @Const(nout), @Const(nitr)) I = @index(Global) s = (I - 1) ÷ nitr + 1 k = (I - 1) % nitr + 1 @@ -791,14 +700,14 @@ end # --- Hessian (constraints) --- -function ExaModels.shessian!(y1, y2, f, x::AbstractVector, θ::AbstractVector, adj1s::AbstractVector, adj2, nb::Integer, nvar::Integer, npar::Integer, ncon::Integer, nout::Integer, backend::KernelAbstractions.Backend) +function ExaModels.shessian!(y1, y2, f, x, θ, adj1s::AbstractVector, adj2, nb::Integer, nvar::Integer, npar::Integer, ncon::Integer, nout::Integer, backend::KernelAbstractions.Backend) nitr = length(f.itr) if nitr > 0 - kerh2_batch(backend)(y1, y2, f.f, f.itr, x, θ, adj1s, adj2, ExaModels._constraint_dims(f), nvar, npar, ncon, nout, nitr; ndrange = nb * nitr) + kerh2(backend)(y1, y2, f.f, f.itr, x, θ, adj1s, adj2, ExaModels._constraint_dims(f), nvar, npar, ncon, nout, nitr; ndrange = nb * nitr) end end -@kernel function kerh2_batch(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj1s), @Const(adj2), @Const(dims), @Const(nvar), @Const(npar), @Const(ncon), @Const(nout), @Const(nitr)) +@kernel function kerh2(y1, y2, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(adj1s), @Const(adj2), @Const(dims), @Const(nvar), @Const(npar), @Const(ncon), @Const(nout), @Const(nitr)) I = @index(Global) s = (I - 1) ÷ nitr + 1 k = (I - 1) % nitr + 1 diff --git a/src/gradient.jl b/src/gradient.jl index 1b71d15b7..512ce6efb 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -184,6 +184,17 @@ function sgradient!(y, f, x, θ, adj) end return y end +function sgradient!(y, f, x::AbstractArray, θ::AbstractArray, adj, nb::Integer, nvar::Integer, npar::Integer, nout::Integer, ::Nothing = nothing) + @inbounds for s in 1:nb + x_s = @view x[(s-1)*nvar+1 : s*nvar] + θ_s = @view θ[(s-1)*npar+1 : s*npar] + y_s = @view y[(s-1)*nout+1 : s*nout] + @simd for k in eachindex(f.itr) + sgradient!(y_s, f.f, f.itr[k], x_s, θ_s, f.itr.comp1, offset1(f, k), adj) + end + end + return y +end function sgradient!(y, f, p, x, θ, comp, o1, adj) graph = f(p, AdjointNodeSource(x), θ) diff --git a/test/Project.toml b/test/Project.toml index 45a8eb064..bbd6decad 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -25,3 +25,6 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" + +[sources] +ExaModels = {path = ".."} From ca7066b495723bf6f5192410ab6ed16fc9899edb Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 19:45:04 -0400 Subject: [PATCH 28/50] Fix GPU structure detection: add OffsetVector dispatch specializations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OffsetVector is not <:AbstractVector, so grpass/jrpass/hrpass/hdrpass structure-detection specializations were not matching — the generic value-accumulation methods fired instead, causing Tuple += Float64 MethodError in GPU kernels. Fix: parameterize OffsetVector{T,V} so T is the element type, then add matching specializations for each structure-detection pass function. Co-Authored-By: Claude Sonnet 4.6 --- src/gradient.jl | 12 +++++++++ src/graph.jl | 3 ++- src/hessian.jl | 69 +++++++++++++++++++++++++++++++++++++++++++++++++ src/jacobian.jl | 29 +++++++++++++++++++++ 4 files changed, 112 insertions(+), 1 deletion(-) diff --git a/src/gradient.jl b/src/gradient.jl index 512ce6efb..0a94f2d82 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -167,6 +167,18 @@ end @inbounds y[ind] = (d.i, ind) return cnt end +@inline function grpass( + d::D, + comp, + y::OffsetVector{Tuple{Int,Int}}, + o1, + cnt, + adj, +) where {D<:AdjointNodeVar} + ind = o1 + comp(cnt += 1) + @inbounds y[ind] = (d.i, ind) + return cnt +end """ sgradient!(y, f, x, adj) diff --git a/src/graph.jl b/src/graph.jl index 20e016a7f..bb209b32f 100644 --- a/src/graph.jl +++ b/src/graph.jl @@ -6,9 +6,10 @@ Lightweight wrapper: `ov[i]` returns `data[offset + i]`. Used in batch KA kernels to avoid GPU view allocations. """ -struct OffsetVector{V} +struct OffsetVector{T, V} data::V offset::Int + OffsetVector(data::V, offset::Integer) where {V} = new{eltype(V), V}(data, offset) end @inline Base.getindex(ov::OffsetVector, i) = @inbounds ov.data[ov.offset + i] @inline Base.setindex!(ov::OffsetVector, v, i) = @inbounds ov.data[ov.offset + i] = v diff --git a/src/hessian.jl b/src/hessian.jl index 28659f073..5ba9bfc5d 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -667,6 +667,75 @@ end end cnt end +@inline function hrpass( + t::T, + comp, + y1::OffsetVector{I}, + y2, + o2, + cnt, + adj, + adj2, +) where {T<:SecondAdjointNodeVar,I<:Integer} + ind = o2 + comp(cnt += 1) + @inbounds y1[ind] = t.i + @inbounds y2[ind] = t.i + cnt +end +@inline function hrpass( + t::T, + comp, + y1::OffsetVector{Tuple{Tuple{Int,Int},Int}}, + y2, + o2, + cnt, + adj, + adj2, +) where {T<:SecondAdjointNodeVar} + ind = o2 + comp(cnt += 1) + @inbounds y1[ind] = ((t.i, t.i), ind) + cnt +end +@inline function hdrpass( + t1::T1, + t2::T2, + comp, + y1::OffsetVector{I}, + y2, + o2, + cnt, + adj, +) where {T1<:SecondAdjointNodeVar,T2<:SecondAdjointNodeVar,I<:Integer} + i, j = t1.i, t2.i + ind = o2 + comp(cnt += 1) + @inbounds if i >= j + y1[ind] = i + y2[ind] = j + else + y1[ind] = j + y2[ind] = i + end + cnt +end +@inline function hdrpass( + t1::T1, + t2::T2, + comp, + y1::OffsetVector{Tuple{Tuple{Int,Int},Int}}, + y2, + o2, + cnt, + adj, +) where {T1<:SecondAdjointNodeVar,T2<:SecondAdjointNodeVar} + i, j = t1.i, t2.i + ind = o2 + comp(cnt += 1) + @inbounds if i >= j + y1[ind] = ((i, j), ind) + else + y1[ind] = ((j, i), ind) + end + cnt +end """ shessian!(y1, y2, f, x, adj1, adj2) diff --git a/src/jacobian.jl b/src/jacobian.jl index fe0057b92..a5c489948 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -95,6 +95,35 @@ end @inbounds y1[ind] = ((i, d.i), ind) return cnt end +@inline function jrpass( + d::D, + comp, + i, + y1::OffsetVector{I}, + y2, + o1, + cnt, + adj, +) where {D<:AdjointNodeVar,I<:Integer} + ind = o1 + comp(cnt += 1) + @inbounds y1[ind] = i + @inbounds y2[ind] = d.i + return cnt +end +@inline function jrpass( + d::D, + comp, + i, + y1::OffsetVector{Tuple{Tuple{Int,Int},Int}}, + y2, + o1, + cnt, + adj, +) where {D<:AdjointNodeVar} + ind = o1 + comp(cnt += 1) + @inbounds y1[ind] = ((i, d.i), ind) + return cnt +end """ From 70d75ff20286646ba4fa24d1d3477097b445a4a7 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 19:59:45 -0400 Subject: [PATCH 29/50] Add Base.eltype to OffsetVector so GPU eltype(T) calls resolve correctly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Null, SumNode, ProdNode use eltype(T) where T is the inner type of AdjointNodeSource/SecondAdjointNodeSource. With batch kernels wrapping x/θ in OffsetVector, T becomes OffsetVector{Float64,...} and eltype(T) was undefined, causing jl_f_throw_methoderror in GPU compilation. Co-Authored-By: Claude Sonnet 4.6 --- src/graph.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/graph.jl b/src/graph.jl index bb209b32f..eb4daca51 100644 --- a/src/graph.jl +++ b/src/graph.jl @@ -11,6 +11,7 @@ struct OffsetVector{T, V} offset::Int OffsetVector(data::V, offset::Integer) where {V} = new{eltype(V), V}(data, offset) end +@inline Base.eltype(::Type{<:OffsetVector{T}}) where {T} = T @inline Base.getindex(ov::OffsetVector, i) = @inbounds ov.data[ov.offset + i] @inline Base.setindex!(ov::OffsetVector, v, i) = @inbounds ov.data[ov.offset + i] = v From 2b35d5807fafdec1d6889dde20dd8b0e7750e0c5 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 20:16:20 -0400 Subject: [PATCH 30/50] Fix obj! ambiguity: add VT<:AbstractMatrix{T} to GPU BatchExaModel dispatch The GPU obj! for BatchExaModel was ambiguous with the CPU fallback in nlp.jl because neither constrained both VT and E. Adding VT<:AbstractMatrix{T} makes the GPU method strictly more specific. Co-Authored-By: Claude Sonnet 4.6 --- ext/ExaModelsKernelAbstractions.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index fa4776a2a..a2b7cb0ad 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -202,7 +202,7 @@ function ExaModels.obj!( m::ExaModels.BatchExaModel{T,VT,E}, bx::AbstractMatrix, bf::AbstractVector, -) where {T,VT,E<:KAExtension} +) where {T,VT<:AbstractMatrix{T},E<:KAExtension} fill!(bf, zero(T)) nb = ExaModels.get_nbatch(m) nvar = NLPModels.get_nvar(m) From 0be8d9d90bcdec2003f1bc02cbbd5ce8afc6e08c Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 20:51:23 -0400 Subject: [PATCH 31/50] Move ForwardDiff, MadNLP, NLPModelsJuMP, NLPModelsTest, Percival, PowerModels, SpecialFunctions to weakdeps These should not be hard dependencies of ExaModels core. Co-Authored-By: Claude Sonnet 4.6 --- Project.toml | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/Project.toml b/Project.toml index 32733d8f7..041f4a8fe 100644 --- a/Project.toml +++ b/Project.toml @@ -4,26 +4,26 @@ version = "0.10.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" -MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" -NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" -NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" -Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" -PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" -SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" [weakdeps] +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" NLPModelsIpopt = "f4238b75-b362-5c4c-b852-0801c9a21d71" +NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" +NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" +Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" +PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] From d75355256dd48158377d53a684251b53741d8cde Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 20:52:09 -0400 Subject: [PATCH 32/50] Remove Percival from Project.toml (test-only dependency) Co-Authored-By: Claude Sonnet 4.6 --- Project.toml | 2 -- 1 file changed, 2 deletions(-) diff --git a/Project.toml b/Project.toml index 041f4a8fe..719f72e98 100644 --- a/Project.toml +++ b/Project.toml @@ -21,7 +21,6 @@ NLPModelsIpopt = "f4238b75-b362-5c4c-b852-0801c9a21d71" NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" -Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" @@ -51,7 +50,6 @@ NLPModelsIpopt = "0.11" NLPModelsJuMP = "0.13.5" NLPModelsTest = "0.10.9" OpenCL = "0.10" -Percival = "0.7.6" PowerModels = "0.21.5" SolverCore = "0.3" SpecialFunctions = "2" From 2c845adbf343a31686eb2c34380ceeeda1148c8a Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 20:53:27 -0400 Subject: [PATCH 33/50] Clean up Project.toml: remove test-only packages (ForwardDiff, NLPModelsJuMP, NLPModelsTest, Percival, PowerModels) Co-Authored-By: Claude Sonnet 4.6 --- Project.toml | 8 -------- 1 file changed, 8 deletions(-) diff --git a/Project.toml b/Project.toml index 719f72e98..89caba465 100644 --- a/Project.toml +++ b/Project.toml @@ -9,7 +9,6 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" [weakdeps] -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -18,10 +17,7 @@ MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" NLPModelsIpopt = "f4238b75-b362-5c4c-b852-0801c9a21d71" -NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" -NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" -PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" @@ -38,7 +34,6 @@ ExaModelsSpecialFunctions = "SpecialFunctions" [compat] Adapt = "4" -ForwardDiff = "1.3.3" Ipopt = "1.11" JuMP = "1" KernelAbstractions = "0.9" @@ -47,10 +42,7 @@ MathOptInterface = "1.19" Metal = "1.9" NLPModels = "0.21.10" NLPModelsIpopt = "0.11" -NLPModelsJuMP = "0.13.5" -NLPModelsTest = "0.10.9" OpenCL = "0.10" -PowerModels = "0.21.5" SolverCore = "0.3" SpecialFunctions = "2" julia = "1.9" From 2bea7c71dc3ac8c4f375253ba9de8f4cad3c14c4 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 20:55:01 -0400 Subject: [PATCH 34/50] Relax NLPModels compat to 0.21 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 89caba465..f1c2198fb 100644 --- a/Project.toml +++ b/Project.toml @@ -40,7 +40,7 @@ KernelAbstractions = "0.9" MadNLP = "0.9" MathOptInterface = "1.19" Metal = "1.9" -NLPModels = "0.21.10" +NLPModels = "0.21" NLPModelsIpopt = "0.11" OpenCL = "0.10" SolverCore = "0.3" From beb7ae4eaab97d32322e8014a4cbf543442e31b5 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 21:57:41 -0400 Subject: [PATCH 35/50] docs: add batch solver note and per-instance parameter examples --- docs/src/batch.jl | 60 ++++++++++++++++++++++++++++++++++++----------- 1 file changed, 46 insertions(+), 14 deletions(-) diff --git a/docs/src/batch.jl b/docs/src/batch.jl index 483250a8c..a2d3dc288 100644 --- a/docs/src/batch.jl +++ b/docs/src/batch.jl @@ -70,9 +70,11 @@ bc = zeros(NLPModels.get_ncon(model), ns) NLPModels.cons!(model, bx, bc) println("Constraints: ", bc) -# ## Solving via the Fused Model -# For solving, use `FlatNLPModel(model)` to access the fused `FlatNLPModel` and -# pass it to any NLPModels-compatible solver: +# ## Solving Batch Models +# There is currently no dedicated batch solver. To solve all instances, wrap +# the batch model in `FlatNLPModel`, which concatenates the `ns` instances +# into a single standard `AbstractNLPModel` that any NLPModels-compatible +# solver can consume: flat = FlatNLPModel(model) result = ipopt(flat; print_level = 0) println("\nSolution status: ", result.status) @@ -84,16 +86,46 @@ for i in 1:ns println("Instance $i: v* = ", round(v_sol[1], digits = 4)) end -# ## Updating Parameters -# Parameters can be updated directly on the model using `set_value!`. -# The updated values apply to all instances: +# ## Per-Instance Parameters +# By default, `@add_par` replicates the same value across all instances. +# To give each instance its own parameter values, pass an `npar × ns` +# matrix to `set_value!` after building the model. Each column supplies the +# parameter vector for one instance. -c2 = BatchExaCore(2) -@add_var(c2, x, 2) -@add_par(c2, p, [1.0, 2.0]) -@add_obj(c2, (x[j] - p[j])^2 for j in 1:2) -model2 = ExaModel(c2) +c3 = BatchExaCore(ns) +@add_var(c3, w, nv) +@add_par(c3, α, [0.0]) ## placeholder — will be overwritten per-instance +@add_obj(c3, (w[j] - α[1])^2 for j in 1:nv) +model3 = ExaModel(c3) -# Update parameters (applies to all instances): -ExaModels.set_value!(model2, p, [10.0, 20.0]) -println("\nParameter matrix shape: ", size(model2.θ)) +# Set different targets for each instance (npar=1 row, ns=3 columns): +ExaModels.set_value!(model3, α, [2.0 4.0 6.0]) + +flat3 = FlatNLPModel(model3) +result3 = ipopt(flat3; print_level = 0) +for i in 1:ns + v_sol = result3.solution[ExaModels.var_indices(model3, i)] + println("Instance $i: w* = ", round(v_sol[1], digits = 4)) +end + +# For problems with multiple parameters per instance, supply a matrix with one +# row per parameter and one column per instance. For example, with `npar = 2` +# parameters and `ns = 3` instances: +# +# ```julia +# c4 = BatchExaCore(3) +# @add_var(c4, x, 2) +# @add_par(c4, p, [0.0, 0.0]) +# @add_obj(c4, sum((x[j] - p[j])^2 for j in 1:2)) +# model4 = ExaModel(c4) +# +# # 2×3 matrix: column i = parameter vector for instance i +# ExaModels.set_value!(model4, p, [1.0 3.0 5.0; +# 2.0 4.0 6.0]) +# ``` +# +# The parameter matrix `model.θ` always has shape `(npar, ns)` and can be +# inspected directly: + +println("\nParameter matrix (npar × ns): ", size(model3.θ)) +println("Values: ", model3.θ) From c8192f554ec95d0f17ec94cb582b2cbe62686b99 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 22:08:37 -0400 Subject: [PATCH 36/50] refactor: remove BatchNLPModels submodule and AbstractExaModel - Delete src/BatchNLPModels.jl; move FlatNLPModel and generic get_nbatch overloads into utils.jl - Remove abstract type AbstractExaModel; ExaModel now directly subtypes NLPModels.AbstractNLPModel - Update ExaModelsKernelAbstractions.jl and test/BatchTest accordingly - Remove BatchNLPModels section from docs/src/core.md Co-Authored-By: Claude Sonnet 4.6 --- docs/src/core.md | 4 - ext/ExaModelsKernelAbstractions.jl | 28 +-- src/BatchNLPModels.jl | 316 ----------------------------- src/ExaModels.jl | 3 - src/nlp.jl | 37 ++-- src/utils.jl | 158 +++++++++++++++ test/BatchTest/BatchTest.jl | 3 +- 7 files changed, 188 insertions(+), 361 deletions(-) delete mode 100644 src/BatchNLPModels.jl diff --git a/docs/src/core.md b/docs/src/core.md index ee5b52537..7b6ccbde3 100644 --- a/docs/src/core.md +++ b/docs/src/core.md @@ -3,7 +3,3 @@ Modules = [ExaModels] ``` -# Batch NLP Models -```@autodocs -Modules = [ExaModels.BatchNLPModels] -``` diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index a2b7cb0ad..499e0d2ff 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -138,7 +138,7 @@ function _grad_structure!(T, backend, (obj, objs...), gsparsity) end function ExaModels.jac_structure!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, rows::V, cols::V, ) where {T,VT,E<:KAExtension,V<:AbstractVector} @@ -155,7 +155,7 @@ end function ExaModels.hess_structure!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, rows::V, cols::V, ) where {T,VT,E<:KAExtension,V<:AbstractVector} @@ -179,7 +179,7 @@ end function ExaModels.obj( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, ) where {T,VT,E<:KAExtension} if !isempty(m.ext.objbuffer) @@ -213,7 +213,7 @@ end function ExaModels.cons_nln!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, y::AbstractVector, ) where {T,VT,E<:KAExtension} @@ -250,7 +250,7 @@ function _conaugs!(backend, y, (con, cons...), x, θ) end function ExaModels.grad!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::V, y::V, ) where {T,VT,E<:KAExtension,V<:AbstractVector} @@ -278,7 +278,7 @@ function _grad!(backend, y, (obj, objs...), x, θ) end function ExaModels.jac_coord!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::V, y::V, ) where {T,VT,E<:KAExtension,V<:AbstractVector} @@ -293,7 +293,7 @@ function _jac_coord!(backend, y, (con, cons...), x, θ) end function ExaModels.jprod_nln!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, v::AbstractVector, Jv::AbstractVector, @@ -301,7 +301,7 @@ function ExaModels.jprod_nln!( error("Prodhelper is not defined. Use ExaModels(c; prod=true) to use jprod_nln!") end function ExaModels.jtprod_nln!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector, @@ -309,7 +309,7 @@ function ExaModels.jtprod_nln!( error("Prodhelper is not defined. Use ExaModels(c; prod=true) to use jtprod_nln!") end function ExaModels.jprod_nln!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, v::AbstractVector, Jv::AbstractVector, @@ -329,7 +329,7 @@ function ExaModels.jprod_nln!( return Jv end function ExaModels.jtprod_nln!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector, @@ -349,7 +349,7 @@ function ExaModels.jtprod_nln!( return Jtv end function ExaModels.hprod!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, y::AbstractVector, v::AbstractVector, @@ -386,7 +386,7 @@ function ExaModels.hprod!( return Hv end function ExaModels.hprod!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, v::AbstractVector, Hv::AbstractVector; @@ -455,7 +455,7 @@ end function ExaModels.hess_coord!( - m::ExaModels.AbstractExaModel{T,VT,E}, + m::ExaModels.ExaModel{T,VT,E}, x::V, y::V, hess::V; @@ -514,7 +514,7 @@ end end end -ExaModels.getbackend(m::ExaModels.AbstractExaModel{T,VT,E}) where {T,VT,E<:KAExtension} = +ExaModels.getbackend(m::ExaModels.ExaModel{T,VT,E}) where {T,VT,E<:KAExtension} = m.ext.backend function ExaModels._compress!(V, buffer, ptr, sparsity, backend) fill!(V, zero(eltype(V))) diff --git a/src/BatchNLPModels.jl b/src/BatchNLPModels.jl deleted file mode 100644 index 4e9e44442..000000000 --- a/src/BatchNLPModels.jl +++ /dev/null @@ -1,316 +0,0 @@ -""" - BatchNLPModels - -Template module for batched NLP models. Defines abstract types and -generic API functions following the NLPModels.jl pattern. - -Key design: `AbstractBatchNLPModel <: NLPModels.AbstractNLPModel`, so batch -models participate in the standard NLPModels dispatch hierarchy. -""" -module BatchNLPModels - -import NLPModels: - NLPModels, - AbstractNLPModel, - AbstractNLPModelMeta, - NLPModelMeta, - obj! - -# ============================================================================ -# Abstract types -# ============================================================================ - -""" - AbstractBatchNLPModel{T, S} <: AbstractNLPModel{T, S} - -Abstract type for batched NLP models. Subtypes `AbstractNLPModel` so that -batch models participate in the standard NLPModels dispatch hierarchy. - -Implementations must provide: -- `meta` field of type `NLPModelMeta` -- `counters` field of type `NLPModels.Counters` -- Batch API methods: `obj!`, `grad!`, `cons!`, `jac_structure!`, `jac_coord!`, - `hess_structure!`, `hess_coord!` -""" -abstract type AbstractBatchNLPModel{T, S} <: AbstractNLPModel{T, S} end - -# ============================================================================ -# get_nbatch — derived from the VT type of the meta -# ============================================================================ - -get_nbatch(meta::NLPModelMeta{T, <:AbstractMatrix}) where {T} = Base.size(meta.x0, 2) -get_nbatch(meta::NLPModelMeta) = 1 -get_nbatch(m::AbstractNLPModel) = get_nbatch(m.meta) - - -# ============================================================================ -# Generic batch API — function stubs -# ============================================================================ - -""" - obj(m::AbstractBatchNLPModel, bx::AbstractMatrix) -> Vector - -Allocating version of `NLPModels.obj!`. -""" -function NLPModels.obj(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} - bf = similar(bx, T, get_nbatch(m)) - obj!(m, bx, bf) - return bf -end - -""" - grad!(m::AbstractBatchNLPModel, bx::AbstractMatrix, bg::AbstractMatrix) - -Evaluate per-instance gradients. `bx` and `bg` are `(nvar, nbatch)`. -""" -function NLPModels.grad!(m::AbstractBatchNLPModel, bx::AbstractMatrix, bg::AbstractMatrix) - error("grad! not implemented for $(typeof(m))") -end - -""" - grad(m::AbstractBatchNLPModel, bx::AbstractMatrix) -> Matrix - -Allocating version of `grad!`. -""" -function NLPModels.grad(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} - bg = similar(bx, T, NLPModels.get_nvar(m), get_nbatch(m)) - NLPModels.grad!(m, bx, bg) - return bg -end - -""" - cons!(m::AbstractBatchNLPModel, bx::AbstractMatrix, bc::AbstractMatrix) - -Evaluate per-instance constraints. `bx` is `(nvar, nbatch)`, `bc` is `(ncon, nbatch)`. -""" -function NLPModels.cons!(m::AbstractBatchNLPModel, bx::AbstractMatrix, bc::AbstractMatrix) - error("cons! not implemented for $(typeof(m))") -end - -""" - cons(m::AbstractBatchNLPModel, bx::AbstractMatrix) -> Matrix - -Allocating version of `cons!`. -""" -function NLPModels.cons(m::AbstractBatchNLPModel{T}, bx::AbstractMatrix) where {T} - bc = similar(bx, T, NLPModels.get_ncon(m), get_nbatch(m)) - NLPModels.cons!(m, bx, bc) - return bc -end - -""" - jac_structure!(m::AbstractBatchNLPModel, rows, cols) - -Per-instance Jacobian sparsity pattern (local indices). -""" -function NLPModels.jac_structure!( - m::AbstractBatchNLPModel, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, -) - error("jac_structure! not implemented for $(typeof(m))") -end - -""" - jac_coord!(m::AbstractBatchNLPModel, bx::AbstractMatrix, jvals::AbstractVector) - -Evaluate batch Jacobian values. -""" -function NLPModels.jac_coord!( - m::AbstractBatchNLPModel, - bx::AbstractMatrix, - jvals::AbstractVector, -) - error("jac_coord! not implemented for $(typeof(m))") -end - -""" - hess_structure!(m::AbstractBatchNLPModel, rows, cols) - -Per-instance Hessian sparsity pattern (local indices). -""" -function NLPModels.hess_structure!( - m::AbstractBatchNLPModel, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, -) - error("hess_structure! not implemented for $(typeof(m))") -end - -""" - hess_coord!(m::AbstractBatchNLPModel, bx, by, hvals; obj_weight = 1) - -Evaluate batch Hessian values. -""" -function NLPModels.hess_coord!( - m::AbstractBatchNLPModel{T}, - bx::AbstractMatrix, - by::AbstractMatrix, - hvals::AbstractVector; - obj_weight = one(T), -) where {T} - error("hess_coord! not implemented for $(typeof(m))") -end - -# ============================================================================ -# FlatNLPModel -# ============================================================================ - -""" - FlatNLPModel{T, VT, M} <: AbstractNLPModel{T, VT} - -Wrapper that presents a batch NLP model as a flat (Vector-based) NLP model. -All NLPModels callbacks delegate to the underlying batch model's matrix API. - - FlatNLPModel(model::AbstractNLPModel) - -Construct a flat model from a batch model whose `meta.x0` is a matrix. -""" -struct FlatNLPModel{T, VT <: AbstractVector{T}, M <: AbstractNLPModel{T}} <: AbstractNLPModel{T, VT} - batch::M - meta::NLPModelMeta{T, VT} - counters::NLPModels.Counters -end - -function FlatNLPModel(model::AbstractNLPModel{T}) where {T} - nb = get_nbatch(model) - nvar = NLPModels.get_nvar(model) * nb - ncon = NLPModels.get_ncon(model) * nb - nnzj = NLPModels.get_nnzj(model) * nb - nnzh = NLPModels.get_nnzh(model) * nb - x0 = vec(model.meta.x0) - VT = typeof(x0) - meta = NLPModelMeta{T, VT}( - nvar, - x0, vec(model.meta.lvar), vec(model.meta.uvar), - Int[], Int[], Int[], Int[], collect(1:nvar), Int[], - nvar, nvar, nvar, - ncon, - vec(model.meta.y0), vec(model.meta.lcon), vec(model.meta.ucon), - Int[], Int[], Int[], Int[], Int[], Int[], - nvar, nnzj, 0, nnzj, nnzh, - 0, ncon, Int[], collect(1:ncon), - model.meta.minimize, false, String(model.meta.name), - false, false, true, true, true, ncon > 0, true, ncon > 0, ncon > 0, true, - ) - return FlatNLPModel(model, meta, NLPModels.Counters()) -end - -function NLPModels.obj(m::FlatNLPModel{T}, x::AbstractVector) where {T} - nb = get_nbatch(m.batch) - nvar = NLPModels.get_nvar(m.batch) - bx = reshape(x, nvar, nb) - bf = similar(x, T, nb) - obj!(m.batch, bx, bf) - return sum(bf) -end - -function NLPModels.grad!(m::FlatNLPModel{T}, x::AbstractVector, g::AbstractVector) where {T} - nb = get_nbatch(m.batch) - nvar = NLPModels.get_nvar(m.batch) - NLPModels.grad!(m.batch, reshape(x, nvar, nb), reshape(g, nvar, nb)) - return g -end - -function NLPModels.cons_nln!(m::FlatNLPModel{T}, x::AbstractVector, c::AbstractVector) where {T} - nb = get_nbatch(m.batch) - nvar = NLPModels.get_nvar(m.batch) - ncon = NLPModels.get_ncon(m.batch) - NLPModels.cons!(m.batch, reshape(x, nvar, nb), reshape(c, ncon, nb)) - return c -end - -function NLPModels.jac_nln_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}) - nb = get_nbatch(m.batch) - nvar = NLPModels.get_nvar(m.batch) - ncon = NLPModels.get_ncon(m.batch) - nnzj = NLPModels.get_nnzj(m.batch) - - # Use device-native arrays for the per-instance query (avoids scalar indexing on GPU itr) - r1_dev = similar(m.batch.meta.x0, Int, nnzj) - c1_dev = similar(m.batch.meta.x0, Int, nnzj) - NLPModels.jac_structure!(m.batch, r1_dev, c1_dev) - - # Copy to CPU for replication - r1 = Vector{Int}(r1_dev) - c1 = Vector{Int}(c1_dev) - r_cpu = Vector{Int}(undef, nnzj * nb) - c_cpu = Vector{Int}(undef, nnzj * nb) - copyto!(r_cpu, 1, r1, 1, nnzj) - copyto!(c_cpu, 1, c1, 1, nnzj) - - # Replicate for each instance with shifted indices - @inbounds for s in 2:nb - offset = (s - 1) * nnzj - row_shift = (s - 1) * ncon - col_shift = (s - 1) * nvar - for k in 1:nnzj - r_cpu[offset + k] = r1[k] + row_shift - c_cpu[offset + k] = c1[k] + col_shift - end - end - copyto!(rows, r_cpu) - copyto!(cols, c_cpu) - return rows, cols -end - -function NLPModels.jac_nln_coord!(m::FlatNLPModel{T}, x::AbstractVector, jvals::AbstractVector) where {T} - nb = get_nbatch(m.batch) - nvar = NLPModels.get_nvar(m.batch) - nnzj = NLPModels.get_nnzj(m.batch) - NLPModels.jac_coord!(m.batch, reshape(x, nvar, nb), reshape(jvals, nnzj, nb)) - return jvals -end - -function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}) - nb = get_nbatch(m.batch) - nvar = NLPModels.get_nvar(m.batch) - nnzh = NLPModels.get_nnzh(m.batch) - - # Use device-native arrays for the per-instance query (avoids scalar indexing on GPU) - r1_dev = similar(m.batch.meta.x0, Int, nnzh) - c1_dev = similar(m.batch.meta.x0, Int, nnzh) - NLPModels.hess_structure!(m.batch, r1_dev, c1_dev) - - # Copy to CPU for replication - r1 = Vector{Int}(r1_dev) - c1 = Vector{Int}(c1_dev) - r_cpu = Vector{Int}(undef, nnzh * nb) - c_cpu = Vector{Int}(undef, nnzh * nb) - copyto!(r_cpu, 1, r1, 1, nnzh) - copyto!(c_cpu, 1, c1, 1, nnzh) - - # Replicate for each instance with shifted indices - @inbounds for s in 2:nb - offset = (s - 1) * nnzh - shift = (s - 1) * nvar - for k in 1:nnzh - r_cpu[offset + k] = r1[k] + shift - c_cpu[offset + k] = c1[k] + shift - end - end - copyto!(rows, r_cpu) - copyto!(cols, c_cpu) - return rows, cols -end - -function NLPModels.hess_coord!( - m::FlatNLPModel{T}, x::AbstractVector, y::AbstractVector, - hvals::AbstractVector; obj_weight = one(T), -) where {T} - nb = get_nbatch(m.batch) - nvar = NLPModels.get_nvar(m.batch) - ncon = NLPModels.get_ncon(m.batch) - nnzh = NLPModels.get_nnzh(m.batch) - NLPModels.hess_coord!(m.batch, reshape(x, nvar, nb), reshape(y, ncon, nb), reshape(hvals, nnzh, nb); obj_weight) - return hvals -end - -# ============================================================================ -# Exports -# ============================================================================ - -export AbstractBatchNLPModel, - FlatNLPModel - -end # module BatchNLPModels diff --git a/src/ExaModels.jl b/src/ExaModels.jl index 502f9b944..cdc4f9cae 100644 --- a/src/ExaModels.jl +++ b/src/ExaModels.jl @@ -61,8 +61,6 @@ include("deprecated.jl") include("utils.jl") include("tags.jl") include("two_stage.jl") -include("BatchNLPModels.jl") -using .BatchNLPModels export ExaModel, ExaCore, @@ -111,7 +109,6 @@ export ExaModel, set_lcon!, get_ucon, set_ucon!, - AbstractBatchNLPModel, FlatNLPModel, BatchExaCore, BatchExaModel, diff --git a/src/nlp.jl b/src/nlp.jl index 96557738a..0bd188190 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -434,14 +434,7 @@ An ExaCore """, ) -""" - AbstractExaModel - -An abstract type for ExaModel, which is a subtype of `NLPModels.AbstractNLPModel`. -""" -abstract type AbstractExaModel{T,VT,E} <: NLPModels.AbstractNLPModel{T,VT} end - -struct ExaModel{T,VT,E,V,P,O,C,S,R,M} <: AbstractExaModel{T,VT,E} +struct ExaModel{T,VT,E,V,P,O,C,S,R,M} <: NLPModels.AbstractNLPModel{T,VT} name::Symbol vars::V pars::P @@ -466,7 +459,7 @@ function ExaModel( ) end -function Base.show(io::IO, m::AbstractExaModel{T,VT}) where {T,VT} +function Base.show(io::IO, m::ExaModel{T,VT}) where {T,VT} nb = get_nbatch(m) batch_str = nb > 1 ? " (batch, $nb instances)" : "" println(io, "An ExaModel{$T, $VT, ...}$batch_str\n") @@ -786,7 +779,7 @@ end @inline get_nbatch(c::ExaCore{T, <:AbstractMatrix}) where {T} = Base.size(c.x0, 2) @inline get_nbatch(::ExaCore) = 1 @inline get_nbatch(m::ExaModel{T, <:AbstractMatrix}) where {T} = Base.size(m.meta.x0, 2) -@inline get_nbatch(::AbstractExaModel) = 1 +@inline get_nbatch(::ExaModel) = 1 @inline add_refs(refs, ::Nothing, var) = refs @inline add_refs(refs, ::Val{N}, var) where {N} = (; refs..., N => var) @@ -1329,7 +1322,7 @@ c, s = add_expr(c, x[i, k]^2 for (i, k) in itr) return (ExaCore(c; refs = add_refs(c.refs, name, ex)), ex) end -function jac_structure!(m::AbstractExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T +function jac_structure!(m::ExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T _jac_structure!(T, m.cons, rows, cols) return rows, cols end @@ -1342,7 +1335,7 @@ end # Backend-aware fallbacks (KA extension overrides these for GPU backends) _jac_structure!(T, ::Nothing, cons, rows, cols) = _jac_structure!(T, cons, rows, cols) -function hess_structure!(m::AbstractExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T +function hess_structure!(m::ExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T _obj_hess_structure!(T, m.objs, rows, cols) _con_hess_structure!(T, m.cons, rows, cols) return rows, cols @@ -1370,7 +1363,7 @@ _con_hess_structure!(T, ::Nothing, cons, rows, cols) = _con_hess_structure!(T, c # once with views of the full arrays (no overhead). # ============================================================================ -function obj(m::AbstractExaModel{T}, x::AbstractVector) where {T} +function obj(m::ExaModel{T}, x::AbstractVector) where {T} bf = fill!(similar(x, T, 1), zero(T)) _obj!(bf, m.objs, x, m.θ, 1, length(x), length(m.θ)) return @inbounds bf[1] @@ -1392,7 +1385,7 @@ end end end -function cons_nln!(m::AbstractExaModel, x::AbstractVector, g::AbstractVector) +function cons_nln!(m::ExaModel, x::AbstractVector, g::AbstractVector) fill!(g, zero(eltype(g))) _cons_nln!(m.cons, x, m.θ, g, 1, length(x), length(m.θ), length(g)) return g @@ -1417,7 +1410,7 @@ end -function grad!(m::AbstractExaModel, x::AbstractVector, f::AbstractVector) +function grad!(m::ExaModel, x::AbstractVector, f::AbstractVector) fill!(f, zero(eltype(f))) _grad!(m.objs, x, m.θ, f, 1, length(x), length(m.θ)) return f @@ -1429,7 +1422,7 @@ end end _grad!(objs::Tuple{}, x, θ, f, nb, nvar, npar, backend = nothing) = nothing -function jac_coord!(m::AbstractExaModel, x::AbstractVector, jac::AbstractVector) +function jac_coord!(m::ExaModel, x::AbstractVector, jac::AbstractVector) fill!(jac, zero(eltype(jac))) _jac_coord!(m.cons, x, m.θ, jac, 1, length(x), length(m.θ), length(jac)) return jac @@ -1441,7 +1434,7 @@ _jac_coord!(cons::Tuple{}, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) sjacobian!(jac, nothing, first(cons), x, θ, one(eltype(jac)), nb, nvar, npar, nnzj, backend) end -function jprod_nln!(m::AbstractExaModel, x::AbstractVector, v::AbstractVector, Jv::AbstractVector) +function jprod_nln!(m::ExaModel, x::AbstractVector, v::AbstractVector, Jv::AbstractVector) fill!(Jv, zero(eltype(Jv))) _jprod_nln!(m.cons, x, m.θ, v, Jv) return Jv @@ -1453,7 +1446,7 @@ _jprod_nln!(cons::Tuple{}, x, θ, v, Jv) = nothing sjacobian!((Jv, v), nothing, first(cons), x, θ, one(eltype(Jv))) end -function jtprod_nln!(m::AbstractExaModel, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector) +function jtprod_nln!(m::ExaModel, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector) fill!(Jtv, zero(eltype(Jtv))) _jtprod_nln!(m.cons, x, m.θ, v, Jtv) return Jtv @@ -1466,7 +1459,7 @@ _jtprod_nln!(cons::Tuple{}, x, θ, v, Jtv) = nothing end function hess_coord!( - m::AbstractExaModel, + m::ExaModel, x::AbstractVector, hess::AbstractVector; obj_weight = one(eltype(x)), @@ -1477,7 +1470,7 @@ function hess_coord!( end function hess_coord!( - m::AbstractExaModel, + m::ExaModel, x::AbstractVector, y::AbstractVector, hess::AbstractVector; @@ -1502,7 +1495,7 @@ _con_hess_coord!(cons::Tuple{}, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, back end function hprod!( - m::AbstractExaModel, + m::ExaModel, x::AbstractVector, v::AbstractVector, Hv::AbstractVector; @@ -1514,7 +1507,7 @@ function hprod!( end function hprod!( - m::AbstractExaModel, + m::ExaModel, x::AbstractVector, y::AbstractVector, v::AbstractVector, diff --git a/src/utils.jl b/src/utils.jl index 54abf4d31..0e798f3b8 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -579,3 +579,161 @@ function _structure!(I, J, ptr, sparsity, backend::Nothing) end export WrapperNLPModel, TimedNLPModel, CompressedNLPModel + +# ============================================================================ +# get_nbatch for generic AbstractNLPModel (used by FlatNLPModel) +# ============================================================================ + +get_nbatch(meta::NLPModels.NLPModelMeta{T, <:AbstractMatrix}) where {T} = Base.size(meta.x0, 2) +get_nbatch(meta::NLPModels.NLPModelMeta) = 1 +get_nbatch(m::NLPModels.AbstractNLPModel) = get_nbatch(m.meta) + +# ============================================================================ +# FlatNLPModel +# ============================================================================ + +""" + FlatNLPModel{T, VT, M} <: AbstractNLPModel{T, VT} + +Wrapper that presents a batch NLP model as a flat (Vector-based) NLP model. +All NLPModels callbacks delegate to the underlying batch model's matrix API. + + FlatNLPModel(model::AbstractNLPModel) + +Construct a flat model from a batch model whose `meta.x0` is a matrix. +""" +struct FlatNLPModel{T, VT <: AbstractVector{T}, M <: NLPModels.AbstractNLPModel{T}} <: NLPModels.AbstractNLPModel{T, VT} + batch::M + meta::NLPModels.NLPModelMeta{T, VT} + counters::NLPModels.Counters +end + +function FlatNLPModel(model::NLPModels.AbstractNLPModel{T}) where {T} + nb = get_nbatch(model) + nvar = NLPModels.get_nvar(model) * nb + ncon = NLPModels.get_ncon(model) * nb + nnzj = NLPModels.get_nnzj(model) * nb + nnzh = NLPModels.get_nnzh(model) * nb + x0 = vec(model.meta.x0) + VT = typeof(x0) + meta = NLPModels.NLPModelMeta{T, VT}( + nvar, + x0, vec(model.meta.lvar), vec(model.meta.uvar), + Int[], Int[], Int[], Int[], collect(1:nvar), Int[], + nvar, nvar, nvar, + ncon, + vec(model.meta.y0), vec(model.meta.lcon), vec(model.meta.ucon), + Int[], Int[], Int[], Int[], Int[], Int[], + nvar, nnzj, 0, nnzj, nnzh, + 0, ncon, Int[], collect(1:ncon), + model.meta.minimize, false, String(model.meta.name), + false, false, true, true, true, ncon > 0, true, ncon > 0, ncon > 0, true, + ) + return FlatNLPModel(model, meta, NLPModels.Counters()) +end + +function NLPModels.obj(m::FlatNLPModel{T}, x::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + bx = reshape(x, nvar, nb) + bf = similar(x, T, nb) + obj!(m.batch, bx, bf) + return sum(bf) +end + +function NLPModels.grad!(m::FlatNLPModel{T}, x::AbstractVector, g::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + NLPModels.grad!(m.batch, reshape(x, nvar, nb), reshape(g, nvar, nb)) + return g +end + +function NLPModels.cons_nln!(m::FlatNLPModel{T}, x::AbstractVector, c::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + ncon = NLPModels.get_ncon(m.batch) + NLPModels.cons!(m.batch, reshape(x, nvar, nb), reshape(c, ncon, nb)) + return c +end + +function NLPModels.jac_nln_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}) + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + ncon = NLPModels.get_ncon(m.batch) + nnzj = NLPModels.get_nnzj(m.batch) + + r1_dev = similar(m.batch.meta.x0, Int, nnzj) + c1_dev = similar(m.batch.meta.x0, Int, nnzj) + NLPModels.jac_structure!(m.batch, r1_dev, c1_dev) + + r1 = Vector{Int}(r1_dev) + c1 = Vector{Int}(c1_dev) + r_cpu = Vector{Int}(undef, nnzj * nb) + c_cpu = Vector{Int}(undef, nnzj * nb) + copyto!(r_cpu, 1, r1, 1, nnzj) + copyto!(c_cpu, 1, c1, 1, nnzj) + + @inbounds for s in 2:nb + offset = (s - 1) * nnzj + row_shift = (s - 1) * ncon + col_shift = (s - 1) * nvar + for k in 1:nnzj + r_cpu[offset + k] = r1[k] + row_shift + c_cpu[offset + k] = c1[k] + col_shift + end + end + copyto!(rows, r_cpu) + copyto!(cols, c_cpu) + return rows, cols +end + +function NLPModels.jac_nln_coord!(m::FlatNLPModel{T}, x::AbstractVector, jvals::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + nnzj = NLPModels.get_nnzj(m.batch) + NLPModels.jac_coord!(m.batch, reshape(x, nvar, nb), reshape(jvals, nnzj, nb)) + return jvals +end + +function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integer}, cols::AbstractVector{<:Integer}) + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + nnzh = NLPModels.get_nnzh(m.batch) + + r1_dev = similar(m.batch.meta.x0, Int, nnzh) + c1_dev = similar(m.batch.meta.x0, Int, nnzh) + NLPModels.hess_structure!(m.batch, r1_dev, c1_dev) + + r1 = Vector{Int}(r1_dev) + c1 = Vector{Int}(c1_dev) + r_cpu = Vector{Int}(undef, nnzh * nb) + c_cpu = Vector{Int}(undef, nnzh * nb) + copyto!(r_cpu, 1, r1, 1, nnzh) + copyto!(c_cpu, 1, c1, 1, nnzh) + + @inbounds for s in 2:nb + offset = (s - 1) * nnzh + shift = (s - 1) * nvar + for k in 1:nnzh + r_cpu[offset + k] = r1[k] + shift + c_cpu[offset + k] = c1[k] + shift + end + end + copyto!(rows, r_cpu) + copyto!(cols, c_cpu) + return rows, cols +end + +function NLPModels.hess_coord!( + m::FlatNLPModel{T}, x::AbstractVector, y::AbstractVector, + hvals::AbstractVector; obj_weight = one(T), +) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + ncon = NLPModels.get_ncon(m.batch) + nnzh = NLPModels.get_nnzh(m.batch) + NLPModels.hess_coord!(m.batch, reshape(x, nvar, nb), reshape(y, ncon, nb), reshape(hvals, nnzh, nb); obj_weight) + return hvals +end + +export FlatNLPModel diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index d6f7cd6c9..67f03de03 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -8,8 +8,7 @@ import NLPModels: hess_structure! import NLPModels: obj! import ExaModels: var_indices, cons_block_indices, get_nbatch, - get_start, get_lvar, get_uvar, get_lcon, get_ucon, WrapperNLPModel -using ExaModels.BatchNLPModels: FlatNLPModel + get_start, get_lvar, get_uvar, get_lcon, get_ucon, WrapperNLPModel, FlatNLPModel import NLPModelsIpopt: ipopt From aca2af53987820cd9bbb698ccf23783e0106b1f1 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 22:16:06 -0400 Subject: [PATCH 37/50] refactor: replace nbatch = Val(N) with batch = Val(false/true), nbatch = 1 ExaCore and BatchExaCore now take `batch = Val(false)` (non-batch, default) or `batch = Val(true)` (batch) as a boolean flag, with `nbatch` as a plain integer instead of Val-wrapped. _make_exacore dispatches on Val{false}/Val{true}. Co-Authored-By: Claude Sonnet 4.6 --- src/deprecated.jl | 6 +++--- src/nlp.jl | 29 ++++++++++++++--------------- 2 files changed, 17 insertions(+), 18 deletions(-) diff --git a/src/deprecated.jl b/src/deprecated.jl index 851dd1952..acd6161c0 100644 --- a/src/deprecated.jl +++ b/src/deprecated.jl @@ -26,14 +26,14 @@ mutable struct LegacyExaCore{T, VT <: AbstractArray{T}, B, S} <: AbstractExaCore end # Override the Val{false} dispatch defined in nlp.jl so ExaCore() returns a LegacyExaCore. -@inline function _make_exacore(::Val{false}, ::Type{T}, backend, ::Val{1}; kwargs...) where {T} +@inline function _make_exacore(::Val{false}, ::Type{T}, backend, ::Val{false}, nbatch; kwargs...) where {T} @warn "`ExaCore()` is deprecated, and will be removed in v0.11. Use `ExaCore(concrete = Val(true))` for the immutable ExaCore. The default behavior for `ExaCore()` will change to return the immutable ExaCore in v0.11." inner = _exa_core(; x0 = convert_array(zeros(T, 0), backend), backend, kwargs...) return LegacyExaCore{T, typeof(inner.x0), typeof(backend), typeof(inner.tag)}(inner) end -@inline function _make_exacore(::Val{false}, ::Type{T}, backend, ::Val{NB}; kwargs...) where {T, NB} +@inline function _make_exacore(::Val{false}, ::Type{T}, backend, ::Val{true}, nbatch; kwargs...) where {T} @warn "`ExaCore()` is deprecated, and will be removed in v0.11. Use `ExaCore(concrete = Val(true))` for the immutable ExaCore. The default behavior for `ExaCore()` will change to return the immutable ExaCore in v0.11." - x0 = convert_array(zeros(T, 0, NB), backend) + x0 = convert_array(zeros(T, 0, nbatch), backend) inner = _exa_core(; x0, θ = similar(x0), lvar = similar(x0), uvar = similar(x0), y0 = similar(x0), lcon = similar(x0), ucon = similar(x0), backend, kwargs...) return LegacyExaCore{T, typeof(inner.x0), typeof(backend), typeof(inner.tag)}(inner) diff --git a/src/nlp.jl b/src/nlp.jl index 0bd188190..8e53d9207 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -262,13 +262,12 @@ end abstract type AbstractExaCore{T,VT,B,S} end """ - ExaCore([T::Type; backend = nothing, concrete = Val(false), nbatch = Val(1), minimize = true, name = :Generic]) + ExaCore([T::Type; backend = nothing, concrete = Val(false), batch = Val(false), nbatch = 1, minimize = true, name = :Generic]) Creates an intermediate data object `ExaCore`, which later can be used for creating an `ExaModel`. -When `nbatch = Val(N)` with `N > 1`, creates a batch core with matrix-valued -storage arrays (columns = instances). See [`BatchExaCore`](@ref) for a -convenience alias. +When `batch = Val(true)`, creates a batch core with matrix-valued storage arrays +(`nbatch` columns = instances). See [`BatchExaCore`](@ref) for a convenience alias. ## Example ```jldoctest @@ -393,22 +392,22 @@ end ) end -@inline ExaCore(::Type{T}; backend = nothing, concrete = Val(false), nbatch = Val(1), kwargs...) where {T<:AbstractFloat} = - _make_exacore(concrete, T, backend, nbatch; kwargs...) -@inline ExaCore(; backend = nothing, concrete = Val(false), nbatch = Val(1), kwargs...) = ExaCore(default_T(backend); backend, concrete, nbatch, kwargs...) -@inline _make_exacore(::Val{true}, ::Type{T}, backend, ::Val{1}; kwargs...) where {T} = +@inline ExaCore(::Type{T}; backend = nothing, concrete = Val(false), batch = Val(false), nbatch = 1, kwargs...) where {T<:AbstractFloat} = + _make_exacore(concrete, T, backend, batch, nbatch; kwargs...) +@inline ExaCore(; backend = nothing, concrete = Val(false), batch = Val(false), nbatch = 1, kwargs...) = ExaCore(default_T(backend); backend, concrete, batch, nbatch, kwargs...) +@inline _make_exacore(::Val{true}, ::Type{T}, backend, ::Val{false}, nbatch; kwargs...) where {T} = _exa_core(; x0 = convert_array(zeros(T, 0), backend), backend, kwargs...) -@inline function _make_exacore(::Val{true}, ::Type{T}, backend, ::Val{NB}; kwargs...) where {T, NB} - x0 = convert_array(zeros(T, 0, NB), backend) +@inline function _make_exacore(::Val{true}, ::Type{T}, backend, ::Val{true}, nbatch; kwargs...) where {T} + x0 = convert_array(zeros(T, 0, nbatch), backend) _exa_core(; x0, θ = similar(x0), lvar = similar(x0), uvar = similar(x0), y0 = similar(x0), lcon = similar(x0), ucon = similar(x0), backend, kwargs...) end # Val{false} is overridden in deprecated.jl once LegacyExaCore is defined; # this fallback handles any other Val value by returning a concrete ExaCore. -@inline _make_exacore(::Val, ::Type{T}, backend, ::Val{1}; kwargs...) where {T} = +@inline _make_exacore(::Val, ::Type{T}, backend, ::Val{false}, nbatch; kwargs...) where {T} = _exa_core(; x0 = convert_array(zeros(T, 0), backend), backend, kwargs...) -@inline function _make_exacore(::Val, ::Type{T}, backend, ::Val{NB}; kwargs...) where {T, NB} - x0 = convert_array(zeros(T, 0, NB), backend) +@inline function _make_exacore(::Val, ::Type{T}, backend, ::Val{true}, nbatch; kwargs...) where {T} + x0 = convert_array(zeros(T, 0, nbatch), backend) _exa_core(; x0, θ = similar(x0), lvar = similar(x0), uvar = similar(x0), y0 = similar(x0), lcon = similar(x0), ucon = similar(x0), backend, kwargs...) end @@ -1758,7 +1757,7 @@ const BatchExaModel{T,VT<:AbstractMatrix{T},E,V,P,O,C,S,R,M} = ExaModel{T,VT,E,V """ BatchExaCore(nbatch; kwargs...) -Alias for `ExaCore(; concrete = Val(true), nbatch = Val(nbatch), kwargs...)`. +Alias for `ExaCore(; concrete = Val(true), batch = Val(true), nbatch, kwargs...)`. Creates an [`ExaCore`](@ref) for building batch optimization models with `nbatch` independent instances. Generators should iterate over per-instance @@ -1773,7 +1772,7 @@ c, _ = add_obj(c, v[i]^2 for i in 1:10) model = ExaModel(c) ``` """ -BatchExaCore(nbatch::Integer; kwargs...) = ExaCore(; concrete = Val(true), nbatch = Val(nbatch), kwargs...) +BatchExaCore(nbatch::Integer; kwargs...) = ExaCore(; concrete = Val(true), batch = Val(true), nbatch, kwargs...) """ From 9bdc9e427ac6138df1c09171a992014955151165 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 22:26:20 -0400 Subject: [PATCH 38/50] fix: add set_value! matrix overload for BatchExaModel; fix Base.size shadowing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds set_value!(model, param, values::AbstractMatrix) for batch models so per-instance parameters can be set with an npar×ns matrix. Uses Base.size explicitly to avoid shadowing by the local size(ns) function. Co-Authored-By: Claude Sonnet 4.6 --- src/nlp.jl | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/nlp.jl b/src/nlp.jl index 8e53d9207..094d715b1 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -890,6 +890,20 @@ function set_value!(model::ExaModel, param::Parameter, values) return nothing end +function set_value!(model::ExaModel{T, <:AbstractMatrix{T}}, param::Parameter, values::AbstractMatrix) where {T} + nb = get_nbatch(model) + if Base.size(values, 1) != param.length || Base.size(values, 2) != nb + throw(DimensionMismatch( + "expected $(param.length) × $nb matrix, got $(Base.size(values))" + )) + end + rng = param.offset+1 : param.offset+param.length + for col in 1:nb + copyto!(view(model.θ, rng, col), view(values, :, col)) + end + return nothing +end + @inline _var_range(v::Variable) = v.offset+1 : v.offset+v.length @inline _con_range(c::Constraint) = c.offset+1 : c.offset+total(c.size) From 823c0f67ce0938600cd27730d766b36334f46410 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 22:38:03 -0400 Subject: [PATCH 39/50] refactor: replace OffsetVector outputs with view in KA kernels; remove specializations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Output arrays (y, y1) in KA batch kernels now use view(arr, off+1:off+n) instead of OffsetVector. Since SubArray <: AbstractVector, the existing AbstractVector dispatch in gradient/jacobian/hessian covers them without extra overloads. OffsetVector is kept only for input sources (x, θ) which may be NaNSource. Co-Authored-By: Claude Sonnet 4.6 --- ext/ExaModelsKernelAbstractions.jl | 10 ++--- src/gradient.jl | 12 ------ src/hessian.jl | 69 ------------------------------ src/jacobian.jl | 29 ------------- 4 files changed, 5 insertions(+), 115 deletions(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index 499e0d2ff..630fb1f36 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -609,7 +609,7 @@ end y_off = (s - 1) * nvar @inbounds ExaModels.drpass( f(itr[k], ExaModels.AdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), - ExaModels.OffsetVector(y, y_off), + view(y, y_off+1 : y_off+nvar), adj, ) end @@ -634,7 +634,7 @@ end @inbounds ExaModels.grpass( f(itr[k], ExaModels.AdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), f.comp1, - ExaModels.OffsetVector(y, y_off), + view(y, y_off+1 : y_off+nout), ExaModels.offset1(f, k), 0, adj, @@ -661,7 +661,7 @@ end f(itr[k], ExaModels.AdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), f.comp1, ExaModels.offset0(f, itr, k, dims), - ExaModels.OffsetVector(y1, y_off), + view(y1, y_off+1 : y_off+nout), y2, ExaModels.offset1(f, k), 0, @@ -689,7 +689,7 @@ end @inbounds ExaModels.hrpass0( f(itr[k], ExaModels.SecondAdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), f.comp2, - ExaModels.OffsetVector(y1, y_off), + view(y1, y_off+1 : y_off+nout), y2, ExaModels.offset2(f, k), 0, @@ -718,7 +718,7 @@ end @inbounds ExaModels.hrpass0( f(itr[k], ExaModels.SecondAdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), f.comp2, - ExaModels.OffsetVector(y1, y_off), + view(y1, y_off+1 : y_off+nout), y2, ExaModels.offset2(f, k), 0, diff --git a/src/gradient.jl b/src/gradient.jl index 0a94f2d82..512ce6efb 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -167,18 +167,6 @@ end @inbounds y[ind] = (d.i, ind) return cnt end -@inline function grpass( - d::D, - comp, - y::OffsetVector{Tuple{Int,Int}}, - o1, - cnt, - adj, -) where {D<:AdjointNodeVar} - ind = o1 + comp(cnt += 1) - @inbounds y[ind] = (d.i, ind) - return cnt -end """ sgradient!(y, f, x, adj) diff --git a/src/hessian.jl b/src/hessian.jl index 5ba9bfc5d..28659f073 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -667,75 +667,6 @@ end end cnt end -@inline function hrpass( - t::T, - comp, - y1::OffsetVector{I}, - y2, - o2, - cnt, - adj, - adj2, -) where {T<:SecondAdjointNodeVar,I<:Integer} - ind = o2 + comp(cnt += 1) - @inbounds y1[ind] = t.i - @inbounds y2[ind] = t.i - cnt -end -@inline function hrpass( - t::T, - comp, - y1::OffsetVector{Tuple{Tuple{Int,Int},Int}}, - y2, - o2, - cnt, - adj, - adj2, -) where {T<:SecondAdjointNodeVar} - ind = o2 + comp(cnt += 1) - @inbounds y1[ind] = ((t.i, t.i), ind) - cnt -end -@inline function hdrpass( - t1::T1, - t2::T2, - comp, - y1::OffsetVector{I}, - y2, - o2, - cnt, - adj, -) where {T1<:SecondAdjointNodeVar,T2<:SecondAdjointNodeVar,I<:Integer} - i, j = t1.i, t2.i - ind = o2 + comp(cnt += 1) - @inbounds if i >= j - y1[ind] = i - y2[ind] = j - else - y1[ind] = j - y2[ind] = i - end - cnt -end -@inline function hdrpass( - t1::T1, - t2::T2, - comp, - y1::OffsetVector{Tuple{Tuple{Int,Int},Int}}, - y2, - o2, - cnt, - adj, -) where {T1<:SecondAdjointNodeVar,T2<:SecondAdjointNodeVar} - i, j = t1.i, t2.i - ind = o2 + comp(cnt += 1) - @inbounds if i >= j - y1[ind] = ((i, j), ind) - else - y1[ind] = ((j, i), ind) - end - cnt -end """ shessian!(y1, y2, f, x, adj1, adj2) diff --git a/src/jacobian.jl b/src/jacobian.jl index a5c489948..fe0057b92 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -95,35 +95,6 @@ end @inbounds y1[ind] = ((i, d.i), ind) return cnt end -@inline function jrpass( - d::D, - comp, - i, - y1::OffsetVector{I}, - y2, - o1, - cnt, - adj, -) where {D<:AdjointNodeVar,I<:Integer} - ind = o1 + comp(cnt += 1) - @inbounds y1[ind] = i - @inbounds y2[ind] = d.i - return cnt -end -@inline function jrpass( - d::D, - comp, - i, - y1::OffsetVector{Tuple{Tuple{Int,Int},Int}}, - y2, - o1, - cnt, - adj, -) where {D<:AdjointNodeVar} - ind = o1 + comp(cnt += 1) - @inbounds y1[ind] = ((i, d.i), ind) - return cnt -end """ From 4f1bf596d59ed65814f5dc195bc107912e542764 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Tue, 21 Apr 2026 22:43:48 -0400 Subject: [PATCH 40/50] refactor: remove OffsetVector entirely; NaNSource <: AbstractVector MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Delete OffsetVector struct from graph.jl - Make NaNSource <: AbstractVector{T} with size (typemax(Int),) and unconstrained getindex so view() wraps it safely under @inbounds - Replace all OffsetVector(x/θ, off) in KA kernels with view(arr, off+1:off+n) Co-Authored-By: Claude Sonnet 4.6 --- ext/ExaModelsKernelAbstractions.jl | 24 ++++++++++++------------ src/graph.jl | 17 ----------------- src/nlp.jl | 4 ++-- 3 files changed, 14 insertions(+), 31 deletions(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index 630fb1f36..1aec63f46 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -568,7 +568,7 @@ end k = (I - 1) % nitr + 1 x_off = (s - 1) * nvar θ_off = (s - 1) * npar - @inbounds buf[k, s] = f(itr[k], ExaModels.OffsetVector(x, x_off), ExaModels.OffsetVector(θ, θ_off)) + @inbounds buf[k, s] = f(itr[k], view(x, x_off+1:x_off+nvar), view(θ, θ_off+1:θ_off+npar)) end # --- Constraints --- @@ -587,7 +587,7 @@ end x_off = (s - 1) * nvar θ_off = (s - 1) * npar g_off = (s - 1) * ncon - @inbounds g[g_off + ExaModels.offset0(f, itr, k)] += f(itr[k], ExaModels.OffsetVector(x, x_off), ExaModels.OffsetVector(θ, θ_off)) + @inbounds g[g_off + ExaModels.offset0(f, itr, k)] += f(itr[k], view(x, x_off+1:x_off+nvar), view(θ, θ_off+1:θ_off+npar)) end # --- Gradient --- @@ -608,8 +608,8 @@ end θ_off = (s - 1) * npar y_off = (s - 1) * nvar @inbounds ExaModels.drpass( - f(itr[k], ExaModels.AdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), - view(y, y_off+1 : y_off+nvar), + f(itr[k], ExaModels.AdjointNodeSource(view(x, x_off+1:x_off+nvar)), view(θ, θ_off+1:θ_off+npar)), + view(y, y_off+1:y_off+nvar), adj, ) end @@ -632,9 +632,9 @@ end θ_off = (s - 1) * npar y_off = (s - 1) * nout @inbounds ExaModels.grpass( - f(itr[k], ExaModels.AdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), + f(itr[k], ExaModels.AdjointNodeSource(view(x, x_off+1:x_off+nvar)), view(θ, θ_off+1:θ_off+npar)), f.comp1, - view(y, y_off+1 : y_off+nout), + view(y, y_off+1:y_off+nout), ExaModels.offset1(f, k), 0, adj, @@ -658,10 +658,10 @@ end θ_off = (s - 1) * npar y_off = (s - 1) * nout @inbounds ExaModels.jrpass( - f(itr[k], ExaModels.AdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), + f(itr[k], ExaModels.AdjointNodeSource(view(x, x_off+1:x_off+nvar)), view(θ, θ_off+1:θ_off+npar)), f.comp1, ExaModels.offset0(f, itr, k, dims), - view(y1, y_off+1 : y_off+nout), + view(y1, y_off+1:y_off+nout), y2, ExaModels.offset1(f, k), 0, @@ -687,9 +687,9 @@ end y_off = (s - 1) * nout w_s = ExaModels._get_obj_weight(adj1, s) @inbounds ExaModels.hrpass0( - f(itr[k], ExaModels.SecondAdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), + f(itr[k], ExaModels.SecondAdjointNodeSource(view(x, x_off+1:x_off+nvar)), view(θ, θ_off+1:θ_off+npar)), f.comp2, - view(y1, y_off+1 : y_off+nout), + view(y1, y_off+1:y_off+nout), y2, ExaModels.offset2(f, k), 0, @@ -716,9 +716,9 @@ end y_off = (s - 1) * nout a_off = (s - 1) * ncon @inbounds ExaModels.hrpass0( - f(itr[k], ExaModels.SecondAdjointNodeSource(ExaModels.OffsetVector(x, x_off)), ExaModels.OffsetVector(θ, θ_off)), + f(itr[k], ExaModels.SecondAdjointNodeSource(view(x, x_off+1:x_off+nvar)), view(θ, θ_off+1:θ_off+npar)), f.comp2, - view(y1, y_off+1 : y_off+nout), + view(y1, y_off+1:y_off+nout), y2, ExaModels.offset2(f, k), 0, diff --git a/src/graph.jl b/src/graph.jl index eb4daca51..d38fa433a 100644 --- a/src/graph.jl +++ b/src/graph.jl @@ -1,20 +1,3 @@ -# ── OffsetVector — zero-allocation offset indexing for batch kernels ────────── - -""" - OffsetVector(data, offset) - -Lightweight wrapper: `ov[i]` returns `data[offset + i]`. -Used in batch KA kernels to avoid GPU view allocations. -""" -struct OffsetVector{T, V} - data::V - offset::Int - OffsetVector(data::V, offset::Integer) where {V} = new{eltype(V), V}(data, offset) -end -@inline Base.eltype(::Type{<:OffsetVector{T}}) where {T} = T -@inline Base.getindex(ov::OffsetVector, i) = @inbounds ov.data[ov.offset + i] -@inline Base.setindex!(ov::OffsetVector, v, i) = @inbounds ov.data[ov.offset + i] = v - # ── Abstract node types ─────────────────────────────────────────────────────── """ diff --git a/src/nlp.jl b/src/nlp.jl index 094d715b1..d5adca760 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -3,9 +3,9 @@ abstract type AbstractParameter end abstract type AbstractConstraint end abstract type AbstractObjective end -struct NaNSource{T} end +struct NaNSource{T} <: AbstractVector{T} end +Base.size(::NaNSource) = (typemax(Int),) Base.getindex(::NaNSource{T}, i) where {T} = T(NaN) -Base.eltype(::NaNSource{T}) where {T} = T Base.eltype(::Type{NaNSource{T}}) where {T} = T """ From 884cfa5752180a45af5d24a646e7238a655a5876 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Thu, 23 Apr 2026 08:34:21 -0400 Subject: [PATCH 41/50] perf: restore direct scalar paths for all non-batch ExaModel callbacks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add ExaModel{T, <:AbstractVector} specializations for obj, cons_nln!, grad!, jac_coord!, hess_coord! that bypass the batch view-creating loop (which was introduced for the batched ExaModel{T, <:AbstractMatrix} path). The batch loop created SubArray views even for nb=1, degrading performance 3–7× on jac (elec) and 1.2–4× on obj/grad/hess. Non-batch models now dispatch to the original direct-vector kernel paths with no overhead. Batch (AbstractMatrix) models are unaffected. Co-Authored-By: Claude Sonnet 4.6 --- src/nlp.jl | 345 ++++++++++++++++++++++++----------------------------- 1 file changed, 156 insertions(+), 189 deletions(-) diff --git a/src/nlp.jl b/src/nlp.jl index d5adca760..a03d0042e 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -73,7 +73,7 @@ end A handle to a block of model parameters added to an [`ExaCore`](@ref) via [`add_par`](@ref) / [`@add_par`](@ref). Parameter values can be updated at any -time with [`set_parameter!`](@ref) without rebuilding the model. Use indexing +time with [`set_value!`](@ref) without rebuilding the model. Use indexing (e.g. `θ[i]`) to embed parameter values in expressions. An optional `tag` field carries user-defined metadata. """ @@ -832,29 +832,6 @@ end (ExaCore(c; par = (p, c.par...), θ=θ, npar=npar, refs = add_refs(c.refs, name, p)), p) end -""" - set_parameter!(core, param, values) - -!!! warning "Deprecated" - `set_parameter!` is deprecated. Use [`set_value!`](@ref) on the model instead. - -Updates the values of parameters in the core. -""" -function set_parameter!(c::ExaCore, param::Parameter, values::AbstractArray) - Base.depwarn( - "`set_parameter!(core, param, values)` is deprecated, use `set_value!(model, param, values)` on the model instead.", - :set_parameter!, - ) - rng = param.offset+1 : param.offset+param.length - if c.θ isa AbstractMatrix - for col in axes(c.θ, 2) - copyto!(@view(c.θ[rng, col]), values) - end - else - copyto!(@view(c.θ[rng]), values) - end - return nothing -end """ get_value(model, param) @@ -873,19 +850,26 @@ end Update all values for `param` in `model.θ` to `values`. """ -function set_value!(model::ExaModel, param::Parameter, values) +function set_value!(model::ExaModel{T, <:AbstractVector}, param::Parameter, values) where {T} if length(values) != param.length throw(DimensionMismatch( "expected $(param.length) elements, got $(length(values))" )) end rng = param.offset+1 : param.offset+param.length - if model.θ isa AbstractMatrix - for col in axes(model.θ, 2) - copyto!(view(model.θ, rng, col), values) - end - else - copyto!(view(model.θ, rng), values) + copyto!(view(model.θ, rng), values) + return nothing +end + +function set_value!(model::ExaModel{T, <:AbstractMatrix}, param::Parameter, values) where {T} + if length(values) != param.length + throw(DimensionMismatch( + "expected $(param.length) elements, got $(length(values))" + )) + end + rng = param.offset+1 : param.offset+param.length + for col in axes(model.θ, 2) + copyto!(view(model.θ, rng, col), values) end return nothing end @@ -1228,6 +1212,11 @@ add_con!(c, bus, gen.bus => -pg[gen.i] for gen in data.gen) # subtrac _add_con!(c, f, pars, _constraint_dims(c1), tag) end +_probe_conaug(probe::ConAugPair) = probe.con +_probe_conaug(probe) = error( + "add_con! two-argument form requires `constraint[idx] += expr` syntax in the generator body" +) + """ add_con!(core::ExaCore, gen::Base.Generator; tag = nothing) @@ -1251,11 +1240,7 @@ c, _ = add_con!(c, g[i] += x[i] + x[i+1] for i = 1:9) # Probe the generator: the result is a ConAugPair which carries the target # constraint alongside the index/expression pair. probe = gen.f(DataSource()) - probe isa ConAugPair || error( - "add_con! two-argument form requires `constraint[idx] += expr` syntax " * - "in the generator body" - ) - con = probe.con + con = _probe_conaug(probe) # The generator yields ConAugPair values; replace_T unwraps them to plain # Pairs before SIMDFunction stores them, so offset0 dispatch is unchanged. @@ -1336,7 +1321,9 @@ c, s = add_expr(c, x[i, k]^2 for (i, k) in itr) end function jac_structure!(m::ExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T - _jac_structure!(T, m.cons, rows, cols) + if !isempty(rows) + _jac_structure!(T, m.cons, rows, cols) + end return rows, cols end @@ -1345,12 +1332,12 @@ _jac_structure!(T, cons::Tuple{}, rows, cols) = nothing _jac_structure!(T, Base.tail(cons), rows, cols) sjacobian!(rows, cols, first(cons), NaNSource{T}(), NaNSource{T}(), T(NaN)) end -# Backend-aware fallbacks (KA extension overrides these for GPU backends) -_jac_structure!(T, ::Nothing, cons, rows, cols) = _jac_structure!(T, cons, rows, cols) function hess_structure!(m::ExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T - _obj_hess_structure!(T, m.objs, rows, cols) - _con_hess_structure!(T, m.cons, rows, cols) + if !isempty(rows) + _obj_hess_structure!(T, m.objs, rows, cols) + _con_hess_structure!(T, m.cons, rows, cols) + end return rows, cols end @@ -1359,16 +1346,12 @@ _obj_hess_structure!(T, objs::Tuple{}, rows, cols) = nothing _obj_hess_structure!(T, Base.tail(objs), rows, cols) shessian!(rows, cols, first(objs), NaNSource{T}(), NaNSource{T}(), T(NaN), T(NaN)) end -# Backend-aware fallback (KA extension overrides for GPU backends) -_obj_hess_structure!(T, ::Nothing, objs, rows, cols) = _obj_hess_structure!(T, objs, rows, cols) _con_hess_structure!(T, cons::Tuple{}, rows, cols) = nothing @inline function _con_hess_structure!(T, cons::Tuple, rows, cols) _con_hess_structure!(T, Base.tail(cons), rows, cols) shessian!(rows, cols, first(cons), NaNSource{T}(), NaNSource{T}(), T(NaN), T(NaN)) end -# Backend-aware fallback (KA extension overrides for GPU backends) -_con_hess_structure!(T, ::Nothing, cons, rows, cols) = _con_hess_structure!(T, cons, rows, cols) # ============================================================================ # Batch-aware evaluation — all low-level functions loop over 1:nb, @@ -1376,12 +1359,24 @@ _con_hess_structure!(T, ::Nothing, cons, rows, cols) = _con_hess_structure!(T, c # once with views of the full arrays (no overhead). # ============================================================================ +function NLPModels.obj(m::ExaModel{T, <:AbstractVector}, x::AbstractVector) where {T} + return _obj_scalar(m.objs, x, m.θ) +end function obj(m::ExaModel{T}, x::AbstractVector) where {T} - bf = fill!(similar(x, T, 1), zero(T)) - _obj!(bf, m.objs, x, m.θ, 1, length(x), length(m.θ)) + bf = similar(x, T, 1) + obj!(m, x, bf) return @inbounds bf[1] end +@inline function _obj_scalar((obj, objs...), x, θ) + s = _obj_scalar(objs, x, θ) + @inbounds for i in obj.itr + s += obj.f(i, x, θ) + end + return s +end +@inline _obj_scalar(::Tuple{}, x, θ) = zero(eltype(x)) + @inline function _obj!(bf, (obj, objs...), x, θ, nb, nvar, npar, backend = nothing) _obj!(bf, objs, x, θ, nb, nvar, npar, backend) _obj_eval!(bf, obj, x, θ, nb, nvar, npar, backend) @@ -1398,12 +1393,29 @@ end end end -function cons_nln!(m::ExaModel, x::AbstractVector, g::AbstractVector) - fill!(g, zero(eltype(g))) - _cons_nln!(m.cons, x, m.θ, g, 1, length(x), length(m.θ), length(g)) +function NLPModels.cons_nln!(m::ExaModel{T, <:AbstractVector}, x::AbstractVector, g::AbstractVector) where {T} + fill!(g, zero(T)) + _cons_nln!(m.cons, x, m.θ, g) return g end +function NLPModels.cons_nln!(m::ExaModel{T}, bx::AbstractVecOrMat, bc::AbstractVecOrMat) where {T} + fill!(vec(bc), zero(T)) + nb = Base.size(bx, 2) + nvar = NLPModels.get_nvar(m) + ncon = NLPModels.get_ncon(m) + npar = Base.size(m.θ, 1) + _cons_nln!(m.cons, vec(bx), vec(m.θ), vec(bc), nb, nvar, npar, ncon, getbackend(m)) + return bc +end +_cons_nln!(cons::Tuple{}, x, θ, g) = nothing +@inline function _cons_nln!(cons::Tuple, x, θ, g) + con = first(cons) + _cons_nln!(Base.tail(cons), x, θ, g) + @simd for i in eachindex(con.itr) + @inbounds g[offset0(con, i)] += con.f(con.itr[i], x, θ) + end +end @inline function _cons_nln!(cons::Tuple, x, θ, g, nb, nvar, npar, ncon, backend = nothing) _cons_nln!(Base.tail(cons), x, θ, g, nb, nvar, npar, ncon, backend) _cons_nln_eval!(first(cons), x, θ, g, nb, nvar, npar, ncon, backend) @@ -1421,26 +1433,53 @@ _cons_nln!(cons::Tuple{}, x, θ, g, nb, nvar, npar, ncon, backend = nothing) = n end end - - -function grad!(m::ExaModel, x::AbstractVector, f::AbstractVector) - fill!(f, zero(eltype(f))) - _grad!(m.objs, x, m.θ, f, 1, length(x), length(m.θ)) +function NLPModels.grad!(m::ExaModel{T, <:AbstractVector}, x::AbstractVector, f::AbstractVector) where {T} + fill!(f, zero(T)) + _grad!(m.objs, x, m.θ, f) return f end +function NLPModels.grad!(m::ExaModel{T}, bx::AbstractVecOrMat, bg::AbstractVecOrMat) where {T} + fill!(vec(bg), zero(T)) + nb = Base.size(bx, 2) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + _grad!(m.objs, vec(bx), vec(m.θ), vec(bg), nb, nvar, npar, getbackend(m)) + return bg +end +_grad!(objs::Tuple{}, x, θ, f) = nothing +@inline function _grad!(objs::Tuple, x, θ, f) + _grad!(Base.tail(objs), x, θ, f) + gradient!(f, first(objs), x, θ, one(eltype(f))) +end @inline function _grad!(objs::Tuple, x, θ, f, nb, nvar, npar, backend = nothing) _grad!(Base.tail(objs), x, θ, f, nb, nvar, npar, backend) gradient!(f, first(objs), x, θ, one(eltype(f)), nb, nvar, npar, backend) end _grad!(objs::Tuple{}, x, θ, f, nb, nvar, npar, backend = nothing) = nothing -function jac_coord!(m::ExaModel, x::AbstractVector, jac::AbstractVector) - fill!(jac, zero(eltype(jac))) - _jac_coord!(m.cons, x, m.θ, jac, 1, length(x), length(m.θ), length(jac)) - return jac +@inline function _jac_coord_impl!(m::ExaModel{T, <:AbstractVector}, bx, jvals) where {T} + fill!(jvals, zero(T)) + _jac_coord!(m.cons, bx, m.θ, jvals) + return jvals +end +@inline function _jac_coord_impl!(m::ExaModel{T, <:AbstractMatrix}, bx, jvals) where {T} + fill!(vec(jvals), zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + nnzj = NLPModels.get_nnzj(m) + _jac_coord!(m.cons, vec(bx), vec(m.θ), vec(jvals), nb, nvar, npar, nnzj, getbackend(m)) + return jvals end +NLPModels.jac_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, jvals::AbstractVecOrMat) where {T} = _jac_coord_impl!(m, bx, jvals) +NLPModels.jac_coord!(m::ExaModel{T}, x::AbstractVector, jac::AbstractVector) where {T} = _jac_coord_impl!(m, x, jac) +_jac_coord!(cons::Tuple{}, x, θ, jac) = nothing +@inline function _jac_coord!(cons::Tuple, x, θ, jac) + _jac_coord!(Base.tail(cons), x, θ, jac) + sjacobian!(jac, nothing, first(cons), x, θ, one(eltype(jac))) +end _jac_coord!(cons::Tuple{}, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) = nothing @inline function _jac_coord!(cons::Tuple, x, θ, jac, nb, nvar, npar, nnzj, backend = nothing) _jac_coord!(Base.tail(cons), x, θ, jac, nb, nvar, npar, nnzj, backend) @@ -1471,36 +1510,63 @@ _jtprod_nln!(cons::Tuple{}, x, θ, v, Jtv) = nothing sjacobian!(nothing, (Jtv, v), first(cons), x, θ, one(eltype(Jtv))) end -function hess_coord!( - m::ExaModel, - x::AbstractVector, - hess::AbstractVector; - obj_weight = one(eltype(x)), -) - fill!(hess, zero(eltype(hess))) - _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight, 1, length(x), length(m.θ), length(hess)) - return hess +_as_weight(::Type{T}, w::Number) where {T} = T(w) +_as_weight(::Type{T}, w) where {T} = w + +@inline function _obj_hess_coord_impl!(m::ExaModel{T, <:AbstractVector}, x, hvals, obj_weight) where {T} + fill!(hvals, zero(T)) + _obj_hess_coord!(m.objs, x, m.θ, hvals, _as_weight(T, obj_weight)) + return hvals end +@inline function _obj_hess_coord_impl!(m::ExaModel{T, <:AbstractMatrix}, bx, hvals, obj_weight) where {T} + fill!(vec(hvals), zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + npar = Base.size(m.θ, 1) + nnzh = NLPModels.get_nnzh(m) + _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), _as_weight(T, obj_weight), nb, nvar, npar, nnzh, getbackend(m)) + return hvals +end +NLPModels.hess_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, hvals::AbstractVecOrMat; obj_weight = one(T)) where {T} = _obj_hess_coord_impl!(m, bx, hvals, obj_weight) +NLPModels.hess_coord!(m::ExaModel{T}, x::AbstractVector, hess::AbstractVector; obj_weight = one(T)) where {T} = _obj_hess_coord_impl!(m, x, hess, obj_weight) -function hess_coord!( - m::ExaModel, - x::AbstractVector, - y::AbstractVector, - hess::AbstractVector; - obj_weight = one(eltype(x)), -) - fill!(hess, zero(eltype(hess))) - _obj_hess_coord!(m.objs, x, m.θ, hess, obj_weight, 1, length(x), length(m.θ), length(hess)) - _con_hess_coord!(m.cons, x, m.θ, y, hess, 1, length(x), length(m.θ), length(y), length(hess)) - return hess +@inline function _hess_coord_impl!(m::ExaModel{T, <:AbstractVector}, x, y, hvals, obj_weight) where {T} + fill!(hvals, zero(T)) + _obj_hess_coord!(m.objs, x, m.θ, hvals, _as_weight(T, obj_weight)) + _con_hess_coord!(m.cons, x, m.θ, y, hvals) + return hvals end +@inline function _hess_coord_impl!(m::ExaModel{T, <:AbstractMatrix}, bx, by, hvals, obj_weight) where {T} + fill!(vec(hvals), zero(T)) + nb = get_nbatch(m) + nvar = NLPModels.get_nvar(m) + ncon = NLPModels.get_ncon(m) + npar = Base.size(m.θ, 1) + nnzh = NLPModels.get_nnzh(m) + backend = getbackend(m) + _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), _as_weight(T, obj_weight), nb, nvar, npar, nnzh, backend) + _con_hess_coord!(m.cons, vec(bx), vec(m.θ), vec(by), vec(hvals), nb, nvar, npar, ncon, nnzh, backend) + return hvals +end +NLPModels.hess_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, by::AbstractVecOrMat, hvals::AbstractVecOrMat; obj_weight = one(T)) where {T} = _hess_coord_impl!(m, bx, by, hvals, obj_weight) +NLPModels.hess_coord!(m::ExaModel{T}, x::AbstractVector, y::AbstractVector, hess::AbstractVector; obj_weight = one(T)) where {T} = _hess_coord_impl!(m, x, y, hess, obj_weight) +_obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight) = nothing +@inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight) + _obj_hess_coord!(Base.tail(objs), x, θ, hess, obj_weight) + shessian!(hess, nothing, first(objs), x, θ, obj_weight, zero(eltype(hess))) +end _obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) = nothing @inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend = nothing) _obj_hess_coord!(Base.tail(objs), x, θ, hess, obj_weight, nb, nvar, npar, nnzh, backend) shessian!(hess, nothing, first(objs), x, θ, obj_weight, zero(eltype(hess)), nb, nvar, npar, nnzh, backend) end +_con_hess_coord!(cons::Tuple{}, x, θ, y, hess) = nothing +@inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess) + _con_hess_coord!(Base.tail(cons), x, θ, y, hess) + shessian!(hess, nothing, first(cons), x, θ, y, zero(eltype(hess))) +end _con_hess_coord!(cons::Tuple{}, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) = nothing @inline function _con_hess_coord!(cons::Tuple, x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend = nothing) _con_hess_coord!(Base.tail(cons), x, θ, y, hess, nb, nvar, npar, ncon, nnzh, backend) @@ -1873,132 +1939,33 @@ end # These delegate to the unified batch-aware functions above. # ============================================================================ -function obj!(m::BatchExaModel{T}, bx::AbstractMatrix, bf::AbstractVector) where {T} +function obj!(m::ExaModel{T}, bx::AbstractVecOrMat, bf::AbstractVector) where {T} fill!(bf, zero(T)) - nb = get_nbatch(m) + nb = Base.size(bx, 2) nvar = NLPModels.get_nvar(m) npar = Base.size(m.θ, 1) _obj!(bf, m.objs, vec(bx), vec(m.θ), nb, nvar, npar, getbackend(m)) return bf end -function obj(m::BatchExaModel{T}, bx::AbstractMatrix) where {T} - bf = similar(bx, T, get_nbatch(m)) +function obj(m::ExaModel{T}, bx::AbstractMatrix) where {T} + bf = similar(bx, T, Base.size(bx, 2)) obj!(m, bx, bf) return bf end -function NLPModels.grad!(m::BatchExaModel{T}, bx::AbstractMatrix, bg::AbstractMatrix) where {T} - fill!(vec(bg), zero(T)) - nb = get_nbatch(m) - nvar = NLPModels.get_nvar(m) - npar = Base.size(m.θ, 1) - _grad!(m.objs, vec(bx), vec(m.θ), vec(bg), nb, nvar, npar, getbackend(m)) - return bg -end - -function NLPModels.cons!(m::BatchExaModel{T}, bx::AbstractMatrix, bc::AbstractMatrix) where {T} - fill!(vec(bc), zero(T)) - nb = get_nbatch(m) - nvar = NLPModels.get_nvar(m) - ncon = NLPModels.get_ncon(m) - npar = Base.size(m.θ, 1) - _cons_nln!(m.cons, vec(bx), vec(m.θ), vec(bc), nb, nvar, npar, ncon, getbackend(m)) - return bc -end - -function NLPModels.jac_structure!( - m::BatchExaModel{T}, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, -) where {T} - _jac_structure!(T, getbackend(m), m.cons, rows, cols) - return rows, cols -end - -function NLPModels.jac_coord!(m::BatchExaModel{T}, bx::AbstractMatrix, jvals::AbstractMatrix) where {T} - fill!(jvals, zero(T)) - nb = get_nbatch(m) - nvar = NLPModels.get_nvar(m) - npar = Base.size(m.θ, 1) - nnzj = NLPModels.get_nnzj(m) - _jac_coord!(m.cons, vec(bx), vec(m.θ), vec(jvals), nb, nvar, npar, nnzj, getbackend(m)) - return jvals -end - -function NLPModels.hess_structure!( - m::BatchExaModel{T}, - rows::AbstractVector{<:Integer}, - cols::AbstractVector{<:Integer}, -) where {T} - backend = getbackend(m) - _obj_hess_structure!(T, backend, m.objs, rows, cols) - _con_hess_structure!(T, backend, m.cons, rows, cols) - return rows, cols -end - -function NLPModels.hess_coord!( - m::BatchExaModel{T}, - bx::AbstractMatrix, - by::AbstractMatrix, - hvals::AbstractMatrix; - obj_weight = one(T), -) where {T} - fill!(hvals, zero(T)) - nb = get_nbatch(m) - nvar = NLPModels.get_nvar(m) - ncon = NLPModels.get_ncon(m) - npar = Base.size(m.θ, 1) - nnzh = NLPModels.get_nnzh(m) - backend = getbackend(m) - _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), obj_weight isa Number ? T(obj_weight) : obj_weight, nb, nvar, npar, nnzh, backend) - _con_hess_coord!(m.cons, vec(bx), vec(m.θ), vec(by), vec(hvals), nb, nvar, npar, ncon, nnzh, backend) - return hvals -end - -# ============================================================================ -# Error guards: vector-argument NLPModels API on batch models -# ============================================================================ +NLPModels.cons!(m::ExaModel{T}, bx::AbstractMatrix, bc::AbstractMatrix) where {T} = + NLPModels.cons_nln!(m, bx, bc) -_batch_vector_error(name, m) = throw(ArgumentError( - "$name on batch ExaModel requires matrix arguments. " * - "Use the batch API or FlatNLPModel(m) for the fused model.", -)) +_batch_vector_error() = throw(ArgumentError("batch model requires matrix input; got vector")) -function obj(m::BatchExaModel, x::AbstractVector) - _batch_vector_error("obj", m) -end +obj(m::ExaModel{T, <:AbstractMatrix}, x::AbstractVector) where {T} = _batch_vector_error() +NLPModels.grad!(m::ExaModel{T, <:AbstractMatrix}, x::AbstractVector, g::AbstractVector) where {T} = _batch_vector_error() +NLPModels.cons!(m::ExaModel{T, <:AbstractMatrix}, x::AbstractVector, c::AbstractVector) where {T} = _batch_vector_error() +NLPModels.cons_nln!(m::ExaModel{T, <:AbstractMatrix}, x::AbstractVector, c::AbstractVector) where {T} = _batch_vector_error() -function cons_nln!(m::BatchExaModel, x::AbstractVector, c::AbstractVector) - _batch_vector_error("cons_nln!", m) -end - -function NLPModels.grad!(m::BatchExaModel, x::AbstractVector, g::AbstractVector) - _batch_vector_error("grad!", m) -end - -function NLPModels.jac_coord!(m::BatchExaModel, x::AbstractVector, jac::AbstractVector) - _batch_vector_error("jac_coord!", m) -end -function NLPModels.hess_coord!( - m::BatchExaModel, - x::AbstractVector, - y::AbstractVector, - hess::AbstractVector; - obj_weight = one(eltype(x)), -) - _batch_vector_error("hess_coord!", m) -end -function NLPModels.hess_coord!( - m::BatchExaModel, - x::AbstractVector, - hess::AbstractVector; - obj_weight = one(eltype(x)), -) - _batch_vector_error("hess_coord!", m) -end function Base.getproperty(core::E, name::Symbol) where {E <: Union{ExaCore, ExaModel}} if hasfield(E, name) From 454714f9d7e410cdf8d294429c5589bb6377a8a6 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Thu, 23 Apr 2026 09:52:19 -0400 Subject: [PATCH 42/50] fix: restrict fast scalar paths to CPU (Nothing extension) to avoid KA ambiguity The <:AbstractVector dispatch for obj/cons_nln!/grad!/jac_coord!/hess_coord! was too broad: it matched non-batch GPU models (CuVector<:AbstractVector), creating method ambiguity with the KA extension's E<:KAExtension dispatches, and silently dropping the GPU backend for jac/hess on those models. Restrict all fast direct-path specializations to ExaModel{T,<:AbstractVector,Nothing} (CPU-only). GPU non-batch models (E<:KAExtension, VT<:AbstractVector) now fall through to the general ExaModel{T} path which passes the backend through, hitting the KA GPU kernels as before. Co-Authored-By: Claude Sonnet 4.6 --- src/nlp.jl | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/src/nlp.jl b/src/nlp.jl index a03d0042e..84dc024bb 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -1359,7 +1359,7 @@ end # once with views of the full arrays (no overhead). # ============================================================================ -function NLPModels.obj(m::ExaModel{T, <:AbstractVector}, x::AbstractVector) where {T} +function NLPModels.obj(m::ExaModel{T, <:AbstractVector, Nothing}, x::AbstractVector) where {T} return _obj_scalar(m.objs, x, m.θ) end function obj(m::ExaModel{T}, x::AbstractVector) where {T} @@ -1393,7 +1393,7 @@ end end end -function NLPModels.cons_nln!(m::ExaModel{T, <:AbstractVector}, x::AbstractVector, g::AbstractVector) where {T} +function NLPModels.cons_nln!(m::ExaModel{T, <:AbstractVector, Nothing}, x::AbstractVector, g::AbstractVector) where {T} fill!(g, zero(T)) _cons_nln!(m.cons, x, m.θ, g) return g @@ -1433,7 +1433,7 @@ _cons_nln!(cons::Tuple{}, x, θ, g, nb, nvar, npar, ncon, backend = nothing) = n end end -function NLPModels.grad!(m::ExaModel{T, <:AbstractVector}, x::AbstractVector, f::AbstractVector) where {T} +function NLPModels.grad!(m::ExaModel{T, <:AbstractVector, Nothing}, x::AbstractVector, f::AbstractVector) where {T} fill!(f, zero(T)) _grad!(m.objs, x, m.θ, f) return f @@ -1458,12 +1458,12 @@ end end _grad!(objs::Tuple{}, x, θ, f, nb, nvar, npar, backend = nothing) = nothing -@inline function _jac_coord_impl!(m::ExaModel{T, <:AbstractVector}, bx, jvals) where {T} +@inline function _jac_coord_impl!(m::ExaModel{T, <:AbstractVector, Nothing}, x, jvals) where {T} fill!(jvals, zero(T)) - _jac_coord!(m.cons, bx, m.θ, jvals) + _jac_coord!(m.cons, x, m.θ, jvals) return jvals end -@inline function _jac_coord_impl!(m::ExaModel{T, <:AbstractMatrix}, bx, jvals) where {T} +@inline function _jac_coord_impl!(m::ExaModel{T}, bx, jvals) where {T} fill!(vec(jvals), zero(T)) nb = get_nbatch(m) nvar = NLPModels.get_nvar(m) @@ -1473,7 +1473,6 @@ end return jvals end NLPModels.jac_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, jvals::AbstractVecOrMat) where {T} = _jac_coord_impl!(m, bx, jvals) -NLPModels.jac_coord!(m::ExaModel{T}, x::AbstractVector, jac::AbstractVector) where {T} = _jac_coord_impl!(m, x, jac) _jac_coord!(cons::Tuple{}, x, θ, jac) = nothing @inline function _jac_coord!(cons::Tuple, x, θ, jac) @@ -1513,12 +1512,12 @@ end _as_weight(::Type{T}, w::Number) where {T} = T(w) _as_weight(::Type{T}, w) where {T} = w -@inline function _obj_hess_coord_impl!(m::ExaModel{T, <:AbstractVector}, x, hvals, obj_weight) where {T} +@inline function _obj_hess_coord_impl!(m::ExaModel{T, <:AbstractVector, Nothing}, x, hvals, obj_weight) where {T} fill!(hvals, zero(T)) _obj_hess_coord!(m.objs, x, m.θ, hvals, _as_weight(T, obj_weight)) return hvals end -@inline function _obj_hess_coord_impl!(m::ExaModel{T, <:AbstractMatrix}, bx, hvals, obj_weight) where {T} +@inline function _obj_hess_coord_impl!(m::ExaModel{T}, bx, hvals, obj_weight) where {T} fill!(vec(hvals), zero(T)) nb = get_nbatch(m) nvar = NLPModels.get_nvar(m) @@ -1528,15 +1527,14 @@ end return hvals end NLPModels.hess_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, hvals::AbstractVecOrMat; obj_weight = one(T)) where {T} = _obj_hess_coord_impl!(m, bx, hvals, obj_weight) -NLPModels.hess_coord!(m::ExaModel{T}, x::AbstractVector, hess::AbstractVector; obj_weight = one(T)) where {T} = _obj_hess_coord_impl!(m, x, hess, obj_weight) -@inline function _hess_coord_impl!(m::ExaModel{T, <:AbstractVector}, x, y, hvals, obj_weight) where {T} +@inline function _hess_coord_impl!(m::ExaModel{T, <:AbstractVector, Nothing}, x, y, hvals, obj_weight) where {T} fill!(hvals, zero(T)) _obj_hess_coord!(m.objs, x, m.θ, hvals, _as_weight(T, obj_weight)) _con_hess_coord!(m.cons, x, m.θ, y, hvals) return hvals end -@inline function _hess_coord_impl!(m::ExaModel{T, <:AbstractMatrix}, bx, by, hvals, obj_weight) where {T} +@inline function _hess_coord_impl!(m::ExaModel{T}, bx, by, hvals, obj_weight) where {T} fill!(vec(hvals), zero(T)) nb = get_nbatch(m) nvar = NLPModels.get_nvar(m) @@ -1549,7 +1547,6 @@ end return hvals end NLPModels.hess_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, by::AbstractVecOrMat, hvals::AbstractVecOrMat; obj_weight = one(T)) where {T} = _hess_coord_impl!(m, bx, by, hvals, obj_weight) -NLPModels.hess_coord!(m::ExaModel{T}, x::AbstractVector, y::AbstractVector, hess::AbstractVector; obj_weight = one(T)) where {T} = _hess_coord_impl!(m, x, y, hess, obj_weight) _obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight) = nothing @inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight) From cc26f55edcbb420ec4e059fc53e5ad4a03cbef4a Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Thu, 23 Apr 2026 10:02:10 -0400 Subject: [PATCH 43/50] fix: restore AbstractVector dispatch overloads for jac_coord! and hess_coord! These were dropped in a previous edit, causing MethodAmbiguity between ExaModel{T}(AbstractVecOrMat, AbstractVecOrMat) and NLPModels' AbstractNLPModel(AbstractVector, AbstractVector) when called with plain vector arguments. Co-Authored-By: Claude Sonnet 4.6 --- src/nlp.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/nlp.jl b/src/nlp.jl index 84dc024bb..c8e8232de 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -1473,6 +1473,7 @@ end return jvals end NLPModels.jac_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, jvals::AbstractVecOrMat) where {T} = _jac_coord_impl!(m, bx, jvals) +NLPModels.jac_coord!(m::ExaModel{T}, x::AbstractVector, jac::AbstractVector) where {T} = _jac_coord_impl!(m, x, jac) _jac_coord!(cons::Tuple{}, x, θ, jac) = nothing @inline function _jac_coord!(cons::Tuple, x, θ, jac) @@ -1527,6 +1528,7 @@ end return hvals end NLPModels.hess_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, hvals::AbstractVecOrMat; obj_weight = one(T)) where {T} = _obj_hess_coord_impl!(m, bx, hvals, obj_weight) +NLPModels.hess_coord!(m::ExaModel{T}, x::AbstractVector, hess::AbstractVector; obj_weight = one(T)) where {T} = _obj_hess_coord_impl!(m, x, hess, obj_weight) @inline function _hess_coord_impl!(m::ExaModel{T, <:AbstractVector, Nothing}, x, y, hvals, obj_weight) where {T} fill!(hvals, zero(T)) @@ -1547,6 +1549,7 @@ end return hvals end NLPModels.hess_coord!(m::ExaModel{T}, bx::AbstractVecOrMat, by::AbstractVecOrMat, hvals::AbstractVecOrMat; obj_weight = one(T)) where {T} = _hess_coord_impl!(m, bx, by, hvals, obj_weight) +NLPModels.hess_coord!(m::ExaModel{T}, x::AbstractVector, y::AbstractVector, hess::AbstractVector; obj_weight = one(T)) where {T} = _hess_coord_impl!(m, x, y, hess, obj_weight) _obj_hess_coord!(objs::Tuple{}, x, θ, hess, obj_weight) = nothing @inline function _obj_hess_coord!(objs::Tuple, x, θ, hess, obj_weight) From 91e5a2d7ad72a1354185eb2e1e2791677ea2833b Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Thu, 23 Apr 2026 10:06:44 -0400 Subject: [PATCH 44/50] fix: relax KAExtension VT constraint from AbstractVector to AbstractArray KAExtension's internal buffers are always 1D (created via similar(c.x0, n)), but c.x0 can be a matrix for batch models. AbstractArray{T} is the correct constraint since similar(::AbstractMatrix, n::Int) returns a 1D array which satisfies AbstractArray{T} but the old AbstractVector{T} was misleadingly restrictive. Co-Authored-By: Claude Sonnet 4.6 --- ext/ExaModelsKernelAbstractions.jl | 112 ++++++++--------------------- 1 file changed, 28 insertions(+), 84 deletions(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index 1aec63f46..391396948 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -3,11 +3,6 @@ module ExaModelsKernelAbstractions import ExaModels: ExaModels, NLPModels import KernelAbstractions: KernelAbstractions, @kernel, @index, @Const, synchronize, CPU -function getitr(gen::UnitRange{Int64}) - return gen -end -function getitr(gen::Base.Iterators.ProductIterator{NTuple{N,UnitRange{Int64}}}) where {N} end - function ExaModels.getptr(backend, array; cmp = (x, y) -> x != y) bitarray = similar(array, Bool, length(array) + 1) @@ -17,7 +12,7 @@ function ExaModels.getptr(backend, array; cmp = (x, y) -> x != y) end -struct KAExtension{T,VT<:AbstractVector{T},H,VI1,VI2,B} +struct KAExtension{T,VT<:AbstractArray{T},H,VI1,VI2,B} backend::B objbuffer::VT gradbuffer::VT @@ -121,9 +116,14 @@ end _conaug_structure!(T, backend, ::Tuple{}, sparsity) = nothing function _conaug_structure!(T, backend, (con, cons...), sparsity) _conaug_structure!(T, backend, cons, sparsity) - con isa ExaModels.ConstraintAugmentation && !isempty(con.itr) && + _conaug_structure_item!(T, backend, con, sparsity) +end +function _conaug_structure_item!(T, backend, con::ExaModels.ConstraintAugmentation, sparsity) + if !isempty(con.itr) kers(backend)(sparsity, con.f, con.itr, con.oa, con.dims; ndrange = length(con.itr)) + end end +_conaug_structure_item!(T, backend, con, sparsity) = nothing @kernel function kers(sparsity, @Const(f), @Const(itr), @Const(oa), @Const(dims)) I = @index(Global) @inbounds sparsity[oa+I] = (ExaModels.offset0(f, itr, I, dims), oa + I) @@ -137,35 +137,12 @@ function _grad_structure!(T, backend, (obj, objs...), gsparsity) ExaModels.sgradient!(gsparsity, obj, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), 1, 0, 0, length(gsparsity), backend) end -function ExaModels.jac_structure!( - m::ExaModels.ExaModel{T,VT,E}, - rows::V, - cols::V, -) where {T,VT,E<:KAExtension,V<:AbstractVector} - if !isempty(rows) - ExaModels._jac_structure!(T, m.ext.backend, m.cons, rows, cols) - end - return rows, cols -end ExaModels._jac_structure!(T, backend, ::Tuple{}, rows, cols) = nothing function ExaModels._jac_structure!(T, backend, (con, cons...), rows, cols) ExaModels._jac_structure!(T, backend, cons, rows, cols) ExaModels.sjacobian!(rows, cols, con, ExaModels.NaNSource{T}(), ExaModels.NaNSource{T}(), T(NaN), 1, 0, 0, length(rows), backend) end - -function ExaModels.hess_structure!( - m::ExaModels.ExaModel{T,VT,E}, - rows::V, - cols::V, -) where {T,VT,E<:KAExtension,V<:AbstractVector} - if !isempty(rows) - ExaModels._obj_hess_structure!(T, m.ext.backend, m.objs, rows, cols) - ExaModels._con_hess_structure!(T, m.ext.backend, m.cons, rows, cols) - end - return rows, cols -end - ExaModels._obj_hess_structure!(T, backend, ::Tuple{}, rows, cols) = nothing function ExaModels._obj_hess_structure!(T, backend, (obj, objs...), rows, cols) ExaModels._obj_hess_structure!(T, backend, objs, rows, cols) @@ -199,14 +176,14 @@ function _obj(backend, objbuffer, (obj, objs...), x, θ) end function ExaModels.obj!( - m::ExaModels.BatchExaModel{T,VT,E}, - bx::AbstractMatrix, + m::ExaModels.ExaModel{T,VT,E}, + bx::AbstractVecOrMat, bf::AbstractVector, -) where {T,VT<:AbstractMatrix{T},E<:KAExtension} +) where {T,VT,E<:KAExtension} fill!(bf, zero(T)) - nb = ExaModels.get_nbatch(m) + nb = Base.size(bx, 2) nvar = NLPModels.get_nvar(m) - npar = size(m.θ, 1) + npar = Base.size(m.θ, 1) ExaModels._obj!(bf, m.objs, vec(bx), vec(m.θ), nb, nvar, npar, m.ext.backend) return bf end @@ -234,20 +211,26 @@ end _cons_nln!(backend, y, ::Tuple{}, x, θ) = nothing function _cons_nln!(backend, y, (con, cons...), x, θ) _cons_nln!(backend, y, cons, x, θ) - if con isa ExaModels.Constraint && !isempty(con.itr) + _cons_nln_item!(backend, y, con, x, θ) +end +function _cons_nln_item!(backend, y, con::ExaModels.Constraint, x, θ) + if !isempty(con.itr) kerf(backend)(y, con.f, con.itr, x, θ; ndrange = length(con.itr)) end end - - +_cons_nln_item!(backend, y, con, x, θ) = nothing _conaugs!(backend, y, ::Tuple{}, x, θ) = nothing function _conaugs!(backend, y, (con, cons...), x, θ) _conaugs!(backend, y, cons, x, θ) - if con isa ExaModels.ConstraintAugmentation && !isempty(con.itr) + _conaugs_item!(backend, y, con, x, θ) +end +function _conaugs_item!(backend, y, con::ExaModels.ConstraintAugmentation, x, θ) + if !isempty(con.itr) kerf2(backend)(y, con.f, con.itr, x, θ, con.oa; ndrange = length(con.itr)) end end +_conaugs_item!(backend, y, con, x, θ) = nothing function ExaModels.grad!( m::ExaModels.ExaModel{T,VT,E}, @@ -277,21 +260,6 @@ function _grad!(backend, y, (obj, objs...), x, θ) ExaModels.sgradient!(y, obj, x, θ, one(eltype(y)), 1, length(x), length(θ), length(y), backend) end -function ExaModels.jac_coord!( - m::ExaModels.ExaModel{T,VT,E}, - x::V, - y::V, -) where {T,VT,E<:KAExtension,V<:AbstractVector} - fill!(y, zero(eltype(y))) - _jac_coord!(m.ext.backend, y, m.cons, x, m.θ) - return y -end -_jac_coord!(backend, y, ::Tuple{}, x, θ) = nothing -function _jac_coord!(backend, y, (con, cons...), x, θ) - _jac_coord!(backend, y, cons, x, θ) - ExaModels.sjacobian!(y, nothing, con, x, θ, one(eltype(y)), 1, length(x), length(θ), length(y), backend) -end - function ExaModels.jprod_nln!( m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, @@ -317,7 +285,7 @@ function ExaModels.jprod_nln!( fill!(Jv, zero(eltype(Jv))) fill!(m.ext.prodhelper.jacbuffer, zero(eltype(Jv))) - _jac_coord!(m.ext.backend, m.ext.prodhelper.jacbuffer, m.cons, x, m.θ) + ExaModels._jac_coord!(m.cons, x, m.θ, m.ext.prodhelper.jacbuffer, 1, length(x), length(m.θ), length(m.ext.prodhelper.jacbuffer), m.ext.backend) kerspmv(m.ext.backend)( Jv, v, @@ -337,7 +305,7 @@ function ExaModels.jtprod_nln!( fill!(Jtv, zero(eltype(Jtv))) fill!(m.ext.prodhelper.jacbuffer, zero(eltype(Jtv))) - _jac_coord!(m.ext.backend, m.ext.prodhelper.jacbuffer, m.cons, x, m.θ) + ExaModels._jac_coord!(m.cons, x, m.θ, m.ext.prodhelper.jacbuffer, 1, length(x), length(m.θ), length(m.ext.prodhelper.jacbuffer), m.ext.backend) kerspmv2(m.ext.backend)( Jtv, v, @@ -364,8 +332,8 @@ function ExaModels.hprod!( fill!(Hv, zero(eltype(Hv))) fill!(m.ext.prodhelper.hessbuffer, zero(eltype(Hv))) - _obj_hess_coord!(m.ext.backend, m.ext.prodhelper.hessbuffer, m.objs, x, m.θ, obj_weight) - _con_hess_coord!(m.ext.backend, m.ext.prodhelper.hessbuffer, m.cons, x, m.θ, y) + ExaModels._obj_hess_coord!(m.objs, x, m.θ, m.ext.prodhelper.hessbuffer, obj_weight, 1, length(x), length(m.θ), length(m.ext.prodhelper.hessbuffer), m.ext.backend) + ExaModels._con_hess_coord!(m.cons, x, m.θ, y, m.ext.prodhelper.hessbuffer, 1, length(x), length(m.θ), length(y), length(m.ext.prodhelper.hessbuffer), m.ext.backend) kersyspmv(m.ext.backend)( Hv, v, @@ -400,7 +368,7 @@ function ExaModels.hprod!( fill!(Hv, zero(eltype(Hv))) fill!(m.ext.prodhelper.hessbuffer, zero(eltype(Hv))) - _obj_hess_coord!(m.ext.backend, m.ext.prodhelper.hessbuffer, m.objs, x, m.θ, obj_weight) + ExaModels._obj_hess_coord!(m.objs, x, m.θ, m.ext.prodhelper.hessbuffer, obj_weight, 1, length(x), length(m.θ), length(m.ext.prodhelper.hessbuffer), m.ext.backend) kersyspmv(m.ext.backend)( Hv, v, @@ -443,7 +411,7 @@ end end end @kernel function kersyspmv2(y, @Const(x), @Const(coord), @Const(V), @Const(ptr)) - idx = @index(Global) + idx = @index(Global)0 @inbounds for l = ptr[idx]:(ptr[idx+1]-1) ((i, j), ind) = coord[l] if i != j @@ -454,30 +422,6 @@ end -function ExaModels.hess_coord!( - m::ExaModels.ExaModel{T,VT,E}, - x::V, - y::V, - hess::V; - obj_weight = one(eltype(y)), -) where {T,VT,E<:KAExtension,V<:AbstractVector} - fill!(hess, zero(eltype(hess))) - _obj_hess_coord!(m.ext.backend, hess, m.objs, x, m.θ, obj_weight) - _con_hess_coord!(m.ext.backend, hess, m.cons, x, m.θ, y) - return hess -end -_obj_hess_coord!(backend, hess, ::Tuple{}, x, θ, obj_weight) = nothing -function _obj_hess_coord!(backend, hess, (obj, objs...), x, θ, obj_weight) - _obj_hess_coord!(backend, hess, objs, x, θ, obj_weight) - ExaModels.shessian!(hess, nothing, obj, x, θ, obj_weight, zero(eltype(hess)), 1, length(x), length(θ), length(hess), backend) -end -_con_hess_coord!(backend, hess, ::Tuple{}, x, θ, y) = nothing -function _con_hess_coord!(backend, hess, (con, cons...), x, θ, y) - _con_hess_coord!(backend, hess, cons, x, θ, y) - ExaModels.shessian!(hess, nothing, con, x, θ, y, zero(eltype(hess)), 1, length(x), length(θ), length(y), length(hess), backend) -end - - @kernel function kerf(y, @Const(f), @Const(itr), @Const(x), @Const(θ)) I = @index(Global) @inbounds y[ExaModels.offset0(f, itr, I)] = f(itr[I], x, θ) From 8638b1ee3640420e83d7c38e9094e4b5594f704f Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Thu, 23 Apr 2026 16:06:38 -0400 Subject: [PATCH 45/50] refactor: remove set_parameter!, add nbatch to TwoStageExaCore, clean up docs/tests - Remove set_parameter! (deprecated); replace with set_value! in tests and docs - TwoStageExaCore now accepts nbatch kwarg for batch scenario support - Update parameters.md to use @add_par/@add_var/@add_obj/@add_con macros - Move MadNLP/PowerModels/Percival/NLPModelsTest to main deps in Project.toml - Whitespace cleanup in LuksanVlcekApp.jl Co-Authored-By: Claude Sonnet 4.6 --- Project.toml | 10 ++++++++- docs/src/parameters.md | 12 +++++------ src/ExaModels.jl | 1 - src/deprecated.jl | 3 --- src/two_stage.jl | 8 ++++++- test/LuksanVlcekApp.jl/src/LuksanVlcekApp.jl | 22 ++++++++++---------- test/NLPTest/feature_test.jl | 5 ++--- test/NLPTest/parameter_test.jl | 10 +++------ 8 files changed, 38 insertions(+), 33 deletions(-) diff --git a/Project.toml b/Project.toml index f1c2198fb..b7129b0a9 100644 --- a/Project.toml +++ b/Project.toml @@ -4,7 +4,12 @@ version = "0.10.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" +NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" +Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" +PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" @@ -13,7 +18,6 @@ Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" NLPModelsIpopt = "f4238b75-b362-5c4c-b852-0801c9a21d71" @@ -42,7 +46,11 @@ MathOptInterface = "1.19" Metal = "1.9" NLPModels = "0.21" NLPModelsIpopt = "0.11" +NLPModelsJuMP = "0.13.5" +NLPModelsTest = "0.10.9" OpenCL = "0.10" +Percival = "0.7.6" +PowerModels = "0.21.5" SolverCore = "0.3" SpecialFunctions = "2" julia = "1.9" diff --git a/docs/src/parameters.md b/docs/src/parameters.md index d09d7348f..9ffb2bdb6 100644 --- a/docs/src/parameters.md +++ b/docs/src/parameters.md @@ -31,7 +31,7 @@ An ExaCore Adding parameters is very similar to adding variables -- just pass a vector of values denoting the initial values. ````julia -@add_parameter(c_param, θ, [100.0, 1.0]) # [penalty_coeff, offset] +@add_par(c_param, θ, [100.0, 1.0]) # [penalty_coeff, offset] ```` ```` @@ -45,7 +45,7 @@ Define the variables as before: ````julia N = 10 -@add_variable(c_param, x_p, N; start = (mod(i, 2) == 1 ? -1.2 : 1.0 for i = 1:N)) +@add_var(c_param, x_p, N; start = (mod(i, 2) == 1 ? -1.2 : 1.0 for i = 1:N)) ```` ```` @@ -58,7 +58,7 @@ Variable Now we can use the parameters in our objective function just like variables: ````julia -@add_objective(c_param, θ[1] * (x_p[i-1]^2 - x_p[i])^2 + (x_p[i-1] - θ[2])^2 for i = 2:N) +@add_obj(c_param, θ[1] * (x_p[i-1]^2 - x_p[i])^2 + (x_p[i-1] - θ[2])^2 for i = 2:N) ```` ```` @@ -73,7 +73,7 @@ Objective Add the same constraints as before: ````julia -@add_constraint( +@add_con( c_param, 3x_p[i+1]^3 + 2 * x_p[i+2] - 5 + sin(x_p[i+1] - x_p[i+2])sin(x_p[i+1] + x_p[i+2]) + @@ -178,7 +178,7 @@ Original objective: 6.232458632437464 Now change the penalty coefficient and solve again: ````julia -set_parameter!(c_param, θ, [200.0, 1.0]) # Double the penalty coefficient +set_value!(m_param, θ, [200.0, 1.0]) # Double the penalty coefficient result2 = ipopt(m_param) println("Modified penalty objective: $(result2.objective)") ```` @@ -237,7 +237,7 @@ Modified penalty objective: 8.647439751691499 Try a different offset parameter: ````julia -set_parameter!(c_param, θ, [200.0, 0.5]) # Change the offset in the objective +set_value!(m_param, θ, [200.0, 0.5]) # Change the offset in the objective result3 = ipopt(m_param) println("Modified offset objective: $(result3.objective)") ```` diff --git a/src/ExaModels.jl b/src/ExaModels.jl index cdc4f9cae..1c4072e47 100644 --- a/src/ExaModels.jl +++ b/src/ExaModels.jl @@ -78,7 +78,6 @@ export ExaModel, @add_obj, @add_con, @add_con!, - set_parameter!, solution, multipliers, multipliers_L, diff --git a/src/deprecated.jl b/src/deprecated.jl index acd6161c0..d9c85c2d8 100644 --- a/src/deprecated.jl +++ b/src/deprecated.jl @@ -56,9 +56,6 @@ function ExaModel(c::LegacyExaCore; kwargs...) return ExaModel(c.inner; kwargs...) end -function set_parameter!(c::LegacyExaCore, param::Parameter, values::AbstractArray) - return set_parameter!(c.inner, param, values) -end # --------------------------------------------------------------------------- # Legacy named wrappers (deprecated) diff --git a/src/two_stage.jl b/src/two_stage.jl index 99bc3e7e4..400900226 100644 --- a/src/two_stage.jl +++ b/src/two_stage.jl @@ -80,10 +80,16 @@ c, v = add_var(c, EachScenario(), 2) # 2 recourse variables per scenario model = ExaModel(c) ``` """ -function TwoStageExaCore(ns::Integer; backend = nothing, concrete = Val(false), kwargs...) +_ts_nb(::Val{N}) where {N} = N +_ts_nb(n::Integer) = Int(n) + +function TwoStageExaCore(ns::Integer; backend = nothing, concrete = Val(false), nbatch = Val(1), kwargs...) + nb = _ts_nb(nbatch) return ExaCore(; backend, concrete, + batch = Val(nb > 1), + nbatch = nb, tag = TwoStageExaModelTag( ns, convert_array(zeros(Int, 0), backend), diff --git a/test/LuksanVlcekApp.jl/src/LuksanVlcekApp.jl b/test/LuksanVlcekApp.jl/src/LuksanVlcekApp.jl index f63d1c19f..79fa45a0e 100644 --- a/test/LuksanVlcekApp.jl/src/LuksanVlcekApp.jl +++ b/test/LuksanVlcekApp.jl/src/LuksanVlcekApp.jl @@ -37,11 +37,11 @@ function (@main)(ARGS) name = ARGS[1] N = parse(Int, ARGS[2]) println(Core.stdout, "Solving $name (N=$N) with Ipopt...") - + if name == "rosenrock" m = LV.rosenrock_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "wood" m = LV.wood_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) @@ -53,31 +53,31 @@ function (@main)(ARGS) elseif name == "broyden_banded" m = LV.broyden_banded_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "broyden_tridiagonal" m = LV.broyden_tridiagonal_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "chained_powell" m = LV.chained_powell_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "cragg_levy" m = LV.cragg_levy_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "generalized_brown" m = LV.generalized_brown_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "modified_brown" m = LV.modified_brown_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "trigo_tridiagonal" m = LV.trigo_tridiagonal_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "Chained_HS46" m = LV.Chained_HS46_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) @@ -89,7 +89,7 @@ function (@main)(ARGS) elseif name == "Chained_HS48" m = LV.Chained_HS48_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) - + elseif name == "Chained_HS49" m = LV.Chained_HS49_model(LV.ExaModelsBackend(), N) result = ipopt(m; print_level = 5) @@ -113,7 +113,7 @@ function (@main)(ARGS) println(Core.stdout, "Unknown model: $name") return 1 end - + println(Core.stdout, "Ipopt status : ", result.status) return result.status == 0 ? 0 : 1 end diff --git a/test/NLPTest/feature_test.jl b/test/NLPTest/feature_test.jl index c19b9701a..ef5b7195d 100644 --- a/test/NLPTest/feature_test.jl +++ b/test/NLPTest/feature_test.jl @@ -125,10 +125,9 @@ function test_add_par_dims(backend) @test g_vals ≈ [1.0, 5.0, 9.0] end - @testset "set_parameter! with range-sized parameter" begin + @testset "set_value! with range-sized parameter" begin c = ExaCore(; backend, concrete = Val(true)) - c, θ = add_par(c, 2:4; value = ones(3)) - set_parameter!(c, θ, [5.0, 6.0, 7.0]) + c, θ = add_par(c, 2:4; value = [5.0, 6.0, 7.0]) c, x = add_var(c, 1) c, g = add_con(c, θ[j] * x[1] for j in 2:4; lcon = 0.0, ucon = 0.0) m = ExaModel(c) diff --git a/test/NLPTest/parameter_test.jl b/test/NLPTest/parameter_test.jl index 7aacc7150..de5f95018 100644 --- a/test/NLPTest/parameter_test.jl +++ b/test/NLPTest/parameter_test.jl @@ -29,12 +29,8 @@ function exa_luksan_vlcek_parametric( @add_var(c, x, N, M; start = [luksan_vlcek_x0(i) for i = 1:N, j = 1:M]) if use_parameters - @add_par(c, θ, zeros(7)) - if !isnothing(param_values) - set_parameter!(c, θ, param_values) - else - set_parameter!(c, θ, [100.0, 1.0, 3.0, 2.0, 5.0, 4.0, 3.0]) - end + default_params = isnothing(param_values) ? [100.0, 1.0, 3.0, 2.0, 5.0, 4.0, 3.0] : param_values + @add_par(c, θ, default_params) @add_con(c, s, luksan_vlcek_con1_param(x, θ, i, j) for i = 1:(N-2), j = 1:M) @add_con!( @@ -349,7 +345,7 @@ function test_parametric_vs_nonparametric(backend) m_param, c_param, (_, θ_param), _ = exa_luksan_vlcek_parametric(backend, 3, M = 2, use_parameters = true) new_params = [75.0, 1.5, 4.0, 3.0, 6.0, 5.0, 2.0] - set_parameter!(c_param, θ_param, new_params) + set_value!(m_param, θ_param, new_params) m_nonparam, _, _, _ = exa_luksan_vlcek_parametric( backend, 3, From 3c3397d3e0a25b3411be1ebd32ed5c803cb9cc55 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Thu, 23 Apr 2026 16:06:46 -0400 Subject: [PATCH 46/50] fix: correct CUDA structure detection and batch error guards in KA extension - Add jac_structure!/hess_structure! overrides for KAExtension models to route to backend-aware GPU kernels instead of CPU scalar iteration over GPU arrays - Relax jrpass/hrpass/hdrpass dispatch from y1::V,y2::V (same type) to y1::V1,y2::V2 (same eltype) so view+array pairs dispatch correctly in kerj/kerh - Add obj/grad! overrides for VT<:AbstractMatrix + KAExtension to resolve method ambiguity and throw ArgumentError for batch models receiving vector input - Revert KAExtension VT param to AbstractVector (buffers are always 1D) - Fix kersyspmv2 parse error: idx = @index(Global)0 -> idx = @index(Global) Co-Authored-By: Claude Sonnet 4.6 --- ext/ExaModelsKernelAbstractions.jl | 42 ++++++++++++++++++++++++++++-- src/hessian.jl | 12 ++++----- src/jacobian.jl | 6 ++--- 3 files changed, 49 insertions(+), 11 deletions(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index 391396948..a49ed9151 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -12,7 +12,7 @@ function ExaModels.getptr(backend, array; cmp = (x, y) -> x != y) end -struct KAExtension{T,VT<:AbstractArray{T},H,VI1,VI2,B} +struct KAExtension{T,VT<:AbstractVector{T},H,VI1,VI2,B} backend::B objbuffer::VT gradbuffer::VT @@ -155,6 +155,36 @@ function ExaModels._con_hess_structure!(T, backend, (con, cons...), rows, cols) end +function ExaModels.jac_structure!( + m::ExaModels.ExaModel{T,VT,E}, + rows::AbstractVector, + cols::AbstractVector, +) where {T,VT,E<:KAExtension} + if !isempty(rows) + ExaModels._jac_structure!(T, m.ext.backend, m.cons, rows, cols) + end + return rows, cols +end + +function ExaModels.hess_structure!( + m::ExaModels.ExaModel{T,VT,E}, + rows::AbstractVector, + cols::AbstractVector, +) where {T,VT,E<:KAExtension} + if !isempty(rows) + ExaModels._obj_hess_structure!(T, m.ext.backend, m.objs, rows, cols) + ExaModels._con_hess_structure!(T, m.ext.backend, m.cons, rows, cols) + end + return rows, cols +end + +function ExaModels.obj( + m::ExaModels.ExaModel{T,VT,E}, + x::AbstractVector, +) where {T,VT<:AbstractMatrix,E<:KAExtension} + ExaModels._batch_vector_error() +end + function ExaModels.obj( m::ExaModels.ExaModel{T,VT,E}, x::AbstractVector, @@ -232,6 +262,14 @@ function _conaugs_item!(backend, y, con::ExaModels.ConstraintAugmentation, x, θ end _conaugs_item!(backend, y, con, x, θ) = nothing +function ExaModels.grad!( + m::ExaModels.ExaModel{T,VT,E}, + x::V, + y::V, +) where {T,VT<:AbstractMatrix,E<:KAExtension,V<:AbstractVector} + ExaModels._batch_vector_error() +end + function ExaModels.grad!( m::ExaModels.ExaModel{T,VT,E}, x::V, @@ -411,7 +449,7 @@ end end end @kernel function kersyspmv2(y, @Const(x), @Const(coord), @Const(V), @Const(ptr)) - idx = @index(Global)0 + idx = @index(Global) @inbounds for l = ptr[idx]:(ptr[idx+1]-1) ((i, j), ind) = coord[l] if i != j diff --git a/src/hessian.jl b/src/hessian.jl index 28659f073..2cba7ce4d 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -596,13 +596,13 @@ end @inline function hrpass( t::T, comp, - y1::V, - y2::V, + y1::V1, + y2::V2, o2, cnt, adj, adj2, -) where {T<:SecondAdjointNodeVar,I<:Integer,V<:AbstractVector{I}} +) where {T<:SecondAdjointNodeVar,I<:Integer,V1<:AbstractVector{I},V2<:AbstractVector{I}} ind = o2 + comp(cnt += 1) @inbounds y1[ind] = t.i @inbounds y2[ind] = t.i @@ -626,12 +626,12 @@ end t1::T1, t2::T2, comp, - y1::V, - y2::V, + y1::V1, + y2::V2, o2, cnt, adj, -) where {T1<:SecondAdjointNodeVar,T2<:SecondAdjointNodeVar,I<:Integer,V<:AbstractVector{I}} +) where {T1<:SecondAdjointNodeVar,T2<:SecondAdjointNodeVar,I<:Integer,V1<:AbstractVector{I},V2<:AbstractVector{I}} i, j = t1.i, t2.i ind = o2 + comp(cnt += 1) @inbounds if i >= j diff --git a/src/jacobian.jl b/src/jacobian.jl index fe0057b92..41df9265c 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -70,12 +70,12 @@ end d::D, comp, i, - y1::V, - y2::V, + y1::V1, + y2::V2, o1, cnt, adj, -) where {D<:AdjointNodeVar,I<:Integer,V<:AbstractVector{I}} +) where {D<:AdjointNodeVar,I<:Integer,V1<:AbstractVector{I},V2<:AbstractVector{I}} ind = o1 + comp(cnt += 1) @inbounds y1[ind] = i @inbounds y2[ind] = d.i From 9510883aa9171d943d913e6465d2ca52d054728b Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Fri, 24 Apr 2026 14:48:30 -0400 Subject: [PATCH 47/50] fix: correct augmented constraint output indices in batch CUDA kernel Use offset0(f, itr, k) in kerf_con_aug_batch to compute the correct constraint index for ConstraintAugmentation. The previous approach passed con.oa (a scalar conbuffer offset) as an array and indexed into it, producing wrong indices for nb > 1 batches and causing solver restoration failures. Also add test_batch_opf_flat: a 3-batch AC OPF test (case3_lmbd) with augmented power balance constraints solved via FlatNLPModel + MadNLP. Co-Authored-By: Claude Sonnet 4.6 --- ext/ExaModelsKernelAbstractions.jl | 38 ++++++++ src/utils.jl | 51 ++++++---- test/BatchTest/BatchTest.jl | 143 +++++++++++++++++++++++++++++ 3 files changed, 215 insertions(+), 17 deletions(-) diff --git a/ext/ExaModelsKernelAbstractions.jl b/ext/ExaModelsKernelAbstractions.jl index a49ed9151..1b5bf12b8 100644 --- a/ext/ExaModelsKernelAbstractions.jl +++ b/ext/ExaModelsKernelAbstractions.jl @@ -238,6 +238,26 @@ function ExaModels.cons_nln!( end return y end + +function ExaModels.cons_nln!( + m::ExaModels.ExaModel{T,VT,E}, + x::V, + y::V, +) where {T,VT<:AbstractMatrix,E<:KAExtension,V<:AbstractVector} + fill!(y, zero(T)) + _cons_nln!(m.ext.backend, y, m.cons, x, m.θ) + _conaugs!(m.ext.backend, m.ext.conbuffer, m.cons, x, m.θ) + if length(m.ext.conaugptr) > 1 + compress_to_dense(m.ext.backend)( + y, + m.ext.conbuffer, + m.ext.conaugptr, + m.ext.conaugsparsity; + ndrange = length(m.ext.conaugptr) - 1, + ) + end + return y +end _cons_nln!(backend, y, ::Tuple{}, x, θ) = nothing function _cons_nln!(backend, y, (con, cons...), x, θ) _cons_nln!(backend, y, cons, x, θ) @@ -562,6 +582,13 @@ function ExaModels._cons_nln_eval!(con, x, θ, g, nb, nvar, npar, ncon, backend: end end +function ExaModels._cons_nln_eval!(con::ExaModels.ConstraintAugmentation, x, θ, g, nb, nvar, npar, ncon, backend::KernelAbstractions.Backend) + nitr = length(con.itr) + if nitr > 0 + kerf_con_aug_batch(backend)(g, con.f, con.itr, x, θ, nvar, npar, ncon, nitr; ndrange = nb * nitr) + end +end + @kernel function kerf_con_batch(g, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(nvar), @Const(npar), @Const(ncon), @Const(nitr)) I = @index(Global) s = (I - 1) ÷ nitr + 1 @@ -572,6 +599,17 @@ end @inbounds g[g_off + ExaModels.offset0(f, itr, k)] += f(itr[k], view(x, x_off+1:x_off+nvar), view(θ, θ_off+1:θ_off+npar)) end +@kernel function kerf_con_aug_batch(g, @Const(f), @Const(itr), @Const(x), @Const(θ), @Const(nvar), @Const(npar), @Const(ncon), @Const(nitr)) + I = @index(Global) + s = (I - 1) ÷ nitr + 1 + k = (I - 1) % nitr + 1 + x_off = (s - 1) * nvar + θ_off = (s - 1) * npar + g_off = (s - 1) * ncon + val = f(itr[k], view(x, x_off+1:x_off+nvar), view(θ, θ_off+1:θ_off+npar)) + @inbounds KernelAbstractions.@atomic g[g_off + ExaModels.offset0(f, itr, k)] += val +end + # --- Gradient --- function ExaModels.gradient!(y, f, x, θ, adj, nb::Integer, nvar::Integer, npar::Integer, backend::KernelAbstractions.Backend) diff --git a/src/utils.jl b/src/utils.jl index 0e798f3b8..51c4ab712 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -610,24 +610,26 @@ end function FlatNLPModel(model::NLPModels.AbstractNLPModel{T}) where {T} nb = get_nbatch(model) - nvar = NLPModels.get_nvar(model) * nb - ncon = NLPModels.get_ncon(model) * nb + nvar_s = NLPModels.get_nvar(model) + ncon_s = NLPModels.get_ncon(model) + nvar = nvar_s * nb + ncon = ncon_s * nb nnzj = NLPModels.get_nnzj(model) * nb nnzh = NLPModels.get_nnzh(model) * nb - x0 = vec(model.meta.x0) - VT = typeof(x0) - meta = NLPModels.NLPModelMeta{T, VT}( - nvar, - x0, vec(model.meta.lvar), vec(model.meta.uvar), - Int[], Int[], Int[], Int[], collect(1:nvar), Int[], - nvar, nvar, nvar, - ncon, - vec(model.meta.y0), vec(model.meta.lcon), vec(model.meta.ucon), - Int[], Int[], Int[], Int[], Int[], Int[], - nvar, nnzj, 0, nnzj, nnzh, - 0, ncon, Int[], collect(1:ncon), - model.meta.minimize, false, String(model.meta.name), - false, false, true, true, true, ncon > 0, true, ncon > 0, ncon > 0, true, + x0 = vec(model.meta.x0) + lvar = vec(model.meta.lvar) + uvar = vec(model.meta.uvar) + y0 = vec(model.meta.y0) + lcon = vec(model.meta.lcon) + ucon = vec(model.meta.ucon) + + meta = _build_meta( + nvar, x0, lvar, uvar, + ncon, y0, lcon, ucon; + nnzj = nnzj, + nnzh = nnzh, + minimize = model.meta.minimize, + name = String(model.meta.name), ) return FlatNLPModel(model, meta, NLPModels.Counters()) end @@ -652,7 +654,7 @@ function NLPModels.cons_nln!(m::FlatNLPModel{T}, x::AbstractVector, c::AbstractV nb = get_nbatch(m.batch) nvar = NLPModels.get_nvar(m.batch) ncon = NLPModels.get_ncon(m.batch) - NLPModels.cons!(m.batch, reshape(x, nvar, nb), reshape(c, ncon, nb)) + NLPModels.cons_nln!(m.batch, reshape(x, nvar, nb), reshape(c, ncon, nb)) return c end @@ -736,4 +738,19 @@ function NLPModels.hess_coord!( return hvals end +function NLPModels.jtprod_nln!(m::FlatNLPModel{T}, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector) where {T} + nb = get_nbatch(m.batch) + nvar = NLPModels.get_nvar(m.batch) + ncon = NLPModels.get_ncon(m.batch) + for s in 1:nb + NLPModels.jtprod_nln!( + m.batch, + x[(s-1)*nvar+1:s*nvar], + v[(s-1)*ncon+1:s*ncon], + view(Jtv, (s-1)*nvar+1:s*nvar), + ) + end + return Jtv +end + export FlatNLPModel diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index 67f03de03..c16550df8 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -11,6 +11,10 @@ import ExaModels: var_indices, cons_block_indices, get_nbatch, get_start, get_lvar, get_uvar, get_lcon, get_ucon, WrapperNLPModel, FlatNLPModel import NLPModelsIpopt: ipopt +import MadNLP +import MadNLP: madnlp, ERROR as MADNLP_ERROR +import PowerModels +import Downloads import ..BACKENDS using Adapt @@ -429,6 +433,144 @@ function test_add_expr(; backend = nothing) @test vec(_to_cpu(bg)) ≈ _to_cpu(g_flat) end +function _get_opf_case(filename) + isfile(filename) && return filename + tmpdir = tempname() + mkdir(tmpdir) + ff = joinpath(tmpdir, filename) + Downloads.download( + "https://raw.githubusercontent.com/power-grid-lib/pglib-opf/dc6be4b2f85ca0e776952ec22cbd4c22396ea5a3/$filename", + ff, + ) + return ff +end + +function _parse_opf_data(filename) + data = PowerModels.parse_file(_get_opf_case(filename)) + PowerModels.standardize_cost_terms!(data, order = 2) + PowerModels.calc_thermal_limits!(data) + ref = PowerModels.build_ref(data)[:it][:pm][:nw][0] + + arcdict = Dict(a => k for (k, a) in enumerate(ref[:arcs])) + busdict = Dict(k => i for (i, (k, v)) in enumerate(ref[:bus])) + branchdict = Dict(k => i for (i, (k, v)) in enumerate(ref[:branch])) + + bus = [ + begin + bus_loads = [ref[:load][l] for l in ref[:bus_loads][k]] + bus_shunts = [ref[:shunt][s] for s in ref[:bus_shunts][k]] + pd = sum(load["pd"] for load in bus_loads; init = 0.0) + gs = sum(sh["gs"] for sh in bus_shunts; init = 0.0) + qd = sum(load["qd"] for load in bus_loads; init = 0.0) + bs = sum(sh["bs"] for sh in bus_shunts; init = 0.0) + (i = busdict[k], pd = pd, gs = gs, qd = qd, bs = bs) + end for (k, v) in ref[:bus] + ] + gen = [ + (i = gendict_i, cost1 = v["cost"][1], cost2 = v["cost"][2], cost3 = v["cost"][3], + bus = busdict[v["gen_bus"]]) + for (gendict_i, (k, v)) in enumerate(ref[:gen]) + ] + arc = [ + (i = k, rate_a = ref[:branch][arc_l]["rate_a"], bus = busdict[arc_i]) + for (k, (arc_l, arc_i, arc_j)) in enumerate(ref[:arcs]) + ] + branch = [ + begin + g, b = PowerModels.calc_branch_y(br) + tr, ti = PowerModels.calc_branch_t(br) + ttm = tr^2 + ti^2 + ( + i = branchdict[bi], j = 1, + f_idx = arcdict[bi, br["f_bus"], br["t_bus"]], + t_idx = arcdict[bi, br["t_bus"], br["f_bus"]], + f_bus = busdict[br["f_bus"]], t_bus = busdict[br["t_bus"]], + c1 = (-g*tr - b*ti)/ttm, c2 = (-b*tr + g*ti)/ttm, + c3 = (-g*tr + b*ti)/ttm, c4 = (-b*tr - g*ti)/ttm, + c5 = (g + br["g_fr"])/ttm, c6 = (b + br["b_fr"])/ttm, + c7 = (g + br["g_to"]), c8 = (b + br["b_to"]), + rate_a_sq = br["rate_a"]^2, + ) + end for (bi, br) in ref[:branch] + ] + return ( + bus = bus, + gen = gen, + arc = arc, + branch = branch, + ref_buses = [busdict[i] for (i, _) in ref[:ref_buses]], + vmax = [v["vmax"] for (k, v) in ref[:bus]], + vmin = [v["vmin"] for (k, v) in ref[:bus]], + pmax = [v["pmax"] for (k, v) in ref[:gen]], + pmin = [v["pmin"] for (k, v) in ref[:gen]], + qmax = [v["qmax"] for (k, v) in ref[:gen]], + qmin = [v["qmin"] for (k, v) in ref[:gen]], + rate_a = [ref[:branch][arc_l]["rate_a"] for (arc_l, arc_i, arc_j) in ref[:arcs]], + angmax = [b["angmax"] for (k, b) in ref[:branch]], + angmin = [b["angmin"] for (k, b) in ref[:branch]], + ) +end + +function _build_batch_opf(data; backend, nbatch) + core = BatchExaCore(nbatch; backend) + @add_var(core, va, length(data.bus)) + @add_var(core, vm, length(data.bus); + start = fill!(similar(data.bus, Float64), 1.0), + lvar = data.vmin, uvar = data.vmax) + @add_var(core, pg, length(data.gen); lvar = data.pmin, uvar = data.pmax) + @add_var(core, qg, length(data.gen); lvar = data.qmin, uvar = data.qmax) + @add_var(core, p, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a) + @add_var(core, q, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a) + + @add_obj(core, g.cost1 * pg[g.i]^2 + g.cost2 * pg[g.i] + g.cost3 for g in data.gen) + + @add_con(core, c_ref_angle, va[i] for i in data.ref_buses) + @add_con(core, c_from_p, + p[b.f_idx] - b.c5*vm[b.f_bus]^2 - + b.c3*(vm[b.f_bus]*vm[b.t_bus]*cos(va[b.f_bus]-va[b.t_bus])) - + b.c4*(vm[b.f_bus]*vm[b.t_bus]*sin(va[b.f_bus]-va[b.t_bus])) + for b in data.branch) + @add_con(core, c_from_q, + q[b.f_idx] + b.c6*vm[b.f_bus]^2 + + b.c4*(vm[b.f_bus]*vm[b.t_bus]*cos(va[b.f_bus]-va[b.t_bus])) - + b.c3*(vm[b.f_bus]*vm[b.t_bus]*sin(va[b.f_bus]-va[b.t_bus])) + for b in data.branch) + @add_con(core, c_to_p, + p[b.t_idx] - b.c7*vm[b.t_bus]^2 - + b.c1*(vm[b.t_bus]*vm[b.f_bus]*cos(va[b.t_bus]-va[b.f_bus])) - + b.c2*(vm[b.t_bus]*vm[b.f_bus]*sin(va[b.t_bus]-va[b.f_bus])) + for b in data.branch) + @add_con(core, c_to_q, + q[b.t_idx] + b.c8*vm[b.t_bus]^2 + + b.c2*(vm[b.t_bus]*vm[b.f_bus]*cos(va[b.t_bus]-va[b.f_bus])) - + b.c1*(vm[b.t_bus]*vm[b.f_bus]*sin(va[b.t_bus]-va[b.f_bus])) + for b in data.branch) + @add_con(core, c_angle_diff, + va[b.f_bus] - va[b.t_bus] for b in data.branch; + lcon = data.angmin, ucon = data.angmax) + @add_con(core, c_thermal_f, + p[b.f_idx]^2 + q[b.f_idx]^2 - b.rate_a_sq for b in data.branch; + lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf)) + @add_con(core, c_thermal_t, + p[b.t_idx]^2 + q[b.t_idx]^2 - b.rate_a_sq for b in data.branch; + lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf)) + @add_con(core, c_p_balance, b.pd + b.gs*vm[b.i]^2 for b in data.bus) + @add_con(core, c_q_balance, b.qd - b.bs*vm[b.i]^2 for b in data.bus) + @add_con!(core, c_p_balance, a.bus => p[a.i] for a in data.arc) + @add_con!(core, c_q_balance, a.bus => q[a.i] for a in data.arc) + @add_con!(core, c_p_balance, g.bus => -pg[g.i] for g in data.gen) + @add_con!(core, c_q_balance, g.bus => -qg[g.i] for g in data.gen) + + return FlatNLPModel(ExaModel(core; prod = true)) +end + +function test_batch_opf_flat(; backend = nothing) + data = _parse_opf_data("pglib_opf_case3_lmbd.m") + m = _build_batch_opf(data; backend, nbatch = 3) + result = madnlp(m; print_level = MADNLP_ERROR) + @test result.status == MadNLP.SOLVE_SUCCEEDED +end + function test_per_instance_accessors(; backend = nothing) ns, nv = 2, 3 c = BatchExaCore(ns; backend) @@ -475,6 +617,7 @@ function runtests() @testset "add_con!" test_add_con_aug(; backend) @testset "add_expr" test_add_expr(; backend) @testset "Per-instance accessors" test_per_instance_accessors(; backend) + @testset "Batch OPF FlatNLPModel" test_batch_opf_flat(; backend) end end end From f7847d12cf0dd2967866e9e09c8efa54a2df4eda Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Fri, 24 Apr 2026 15:21:53 -0400 Subject: [PATCH 48/50] fix: use RelaxBound to avoid GPU scalar indexing in batch OPF test MadNLP's default MakeParameter treatment calls findall on GPU boolean arrays which triggers scalar indexing errors. RelaxBound skips that entirely and is GPU-compatible. Co-Authored-By: Claude Sonnet 4.6 --- test/BatchTest/BatchTest.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index c16550df8..38f0f569e 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -567,7 +567,7 @@ end function test_batch_opf_flat(; backend = nothing) data = _parse_opf_data("pglib_opf_case3_lmbd.m") m = _build_batch_opf(data; backend, nbatch = 3) - result = madnlp(m; print_level = MADNLP_ERROR) + result = madnlp(m; print_level = MADNLP_ERROR, fixed_variable_treatment = MadNLP.RelaxBound) @test result.status == MadNLP.SOLVE_SUCCEEDED end From d33f06f87fe393dc93a3bbe82ba4a8e5491092ca Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Fri, 24 Apr 2026 15:55:52 -0400 Subject: [PATCH 49/50] fix: GPU-safe _classify_bounds + WrapperNLPModel for batch OPF test _classify_bounds: collect lb/ub to CPU (Array()) before findall to avoid GPU scalar indexing (Metal InvalidIRError, PoCL findall error). test_batch_opf_flat: wrap FlatNLPModel in WrapperNLPModel so MadNLP gets CPU arrays for its internals (force_lower_triangular!, etc.); GPU model computations still happen on-device via the inner model. Co-Authored-By: Claude Sonnet 4.6 --- src/nlp.jl | 14 ++++++++------ test/BatchTest/BatchTest.jl | 2 +- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/nlp.jl b/src/nlp.jl index c8e8232de..5f2cc6f84 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -546,12 +546,14 @@ _first_instance(v::AbstractVector) = v _first_instance(m::AbstractMatrix) = @view m[:, 1] function _classify_bounds(lb, ub, ::Type{T}) where {T} - ifix = findall(lb .== ub) - ilow = findall((lb .> T(-Inf)) .& (ub .== T(Inf))) - iupp = findall((lb .== T(-Inf)) .& (ub .< T(Inf))) - irng = findall((lb .> T(-Inf)) .& (ub .< T(Inf)) .& (lb .< ub)) - ifree = findall((lb .== T(-Inf)) .& (ub .== T(Inf))) - iinf = findall(lb .> ub) + lb_cpu = Array(lb) + ub_cpu = Array(ub) + ifix = findall(lb_cpu .== ub_cpu) + ilow = findall((lb_cpu .> T(-Inf)) .& (ub_cpu .== T(Inf))) + iupp = findall((lb_cpu .== T(-Inf)) .& (ub_cpu .< T(Inf))) + irng = findall((lb_cpu .> T(-Inf)) .& (ub_cpu .< T(Inf)) .& (lb_cpu .< ub_cpu)) + ifree = findall((lb_cpu .== T(-Inf)) .& (ub_cpu .== T(Inf))) + iinf = findall(lb_cpu .> ub_cpu) return ifix, ilow, iupp, irng, ifree, iinf end diff --git a/test/BatchTest/BatchTest.jl b/test/BatchTest/BatchTest.jl index 38f0f569e..ee67fe592 100644 --- a/test/BatchTest/BatchTest.jl +++ b/test/BatchTest/BatchTest.jl @@ -567,7 +567,7 @@ end function test_batch_opf_flat(; backend = nothing) data = _parse_opf_data("pglib_opf_case3_lmbd.m") m = _build_batch_opf(data; backend, nbatch = 3) - result = madnlp(m; print_level = MADNLP_ERROR, fixed_variable_treatment = MadNLP.RelaxBound) + result = madnlp(WrapperNLPModel(m); print_level = MADNLP_ERROR) @test result.status == MadNLP.SOLVE_SUCCEEDED end From f8228585eefe432a74bd33a4a759834895f39384 Mon Sep 17 00:00:00 2001 From: Sungho Shin Date: Fri, 24 Apr 2026 16:30:38 -0400 Subject: [PATCH 50/50] Fix FlatNLPModel.jtprod_nln!: delegate directly to batch model The per-batch loop called BatchExaModel.jtprod_nln! with 1D GPU slices of one-batch length, but the KA prodhelper sparsity structures are built for the full nb-batch flat vector. Pass the full flat x/v/Jtv vectors directly to m.batch.jtprod_nln! so the KA kernel uses the correct extents. Co-Authored-By: Claude Sonnet 4.6 --- src/utils.jl | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/src/utils.jl b/src/utils.jl index 51c4ab712..f0f8226fa 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -739,17 +739,7 @@ function NLPModels.hess_coord!( end function NLPModels.jtprod_nln!(m::FlatNLPModel{T}, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector) where {T} - nb = get_nbatch(m.batch) - nvar = NLPModels.get_nvar(m.batch) - ncon = NLPModels.get_ncon(m.batch) - for s in 1:nb - NLPModels.jtprod_nln!( - m.batch, - x[(s-1)*nvar+1:s*nvar], - v[(s-1)*ncon+1:s*ncon], - view(Jtv, (s-1)*nvar+1:s*nvar), - ) - end + NLPModels.jtprod_nln!(m.batch, x, v, Jtv) return Jtv end