#![allow(non_camel_case_types)]
use crate::global::THRESHOLD as T;
use crate::numerical::BVP_Damp::BVP_traits::{
convert_to_fun, convert_to_jac, BandedMatrixType, Fun, FunEnum, Jac, JacEnum, MatrixType,
VectorType, Vectors_type_casting,
};
use crate::numerical::BVP_Damp::BVP_utils::{elapsed_time, record_callback_stage_time};
use crate::somelinalg::banded::{LinearSolverConfig, NodeMajorLayout};
use crate::symbolic::codegen::codegen_aot_driver::{AotCodegenBackend, GeneratedAotArtifact};
use crate::symbolic::codegen::codegen_aot_resolution::{AotResolver, ResolvedAotArtifact};
use crate::symbolic::codegen::codegen_aot_runtime_link::resolve_linked_sparse_backend;
use crate::symbolic::codegen::codegen_backend_selection::{
select_backend, BackendSelectionPolicy, SelectedBackendKind,
};
use crate::symbolic::codegen::codegen_manifest::PreparedProblemManifest;
use crate::symbolic::codegen::codegen_orchestrator::{
auto_parallel_recommendation, recommended_sparse_auto_parallel_plan, ParallelExecutorConfig,
ParallelFallbackPolicy, SparseAutoParallelPlan,
};
use crate::symbolic::codegen::codegen_provider_api::{
BackendKind, MatrixBackend, PreparedBandedProblem, PreparedProblem, PreparedSparseProblem,
};
use crate::symbolic::codegen::codegen_runtime_api::{
recommended_banded_chunking_for_parallelism, recommended_residual_chunking_for_parallelism,
recommended_row_chunking_for_parallelism, BandedJacobianRuntimePlan, ResidualChunkPlan,
ResidualChunkingStrategy, ResidualRuntimePlan, SparseJacobianRuntimePlan,
SparseJacobianStructure, SparseJacobianValuesChunkPlan,
};
use crate::symbolic::codegen::codegen_tasks::{
BandedChunkingStrategy, BandedExprEntry, CodegenOutputLayout, CodegenTaskKind, CodegenTaskPlan,
PlannedOutput, SparseChunkingStrategy, SparseExprEntry,
};
use crate::symbolic::codegen::rust_backend::codegen_aot_crate::GeneratedAotCrate;
use crate::symbolic::codegen::CodegenIR::{AtomOptimizationProfile, CodegenModule};
use crate::symbolic::symbolic_engine::Expr;
use crate::symbolic::symbolic_functions_BVP2::{BandedJacobianChunking, BandedLambdifyConfig};
use crate::symbolic::View::bvp::discretization_system_bvp_par_atom;
use crate::symbolic::View::bvp::DiscretizedBvpAtomSystem;
use crate::symbolic::View::bvp_codegen::prepare_sparse_bvp_codegen_from_discretized_system_with_breakdown;
use crate::symbolic::View::conversions::atom_to_expr;
use crate::symbolic::View::jacobian::PreparedSparseAtomSystem;
use faer::col::Col;
use faer::col::ColRef;
use faer::sparse::{SparseColMat, Triplet};
use log::{error, info};
use nalgebra::sparse::CsMatrix;
use nalgebra::{DMatrix, DVector};
use rayon::prelude::*;
use sprs::{CsMat, CsVec};
use std::collections::{HashMap, HashSet};
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::sync::Mutex;
use std::time::Instant;
use tabled::{builder::Builder, settings::Style};
fn normalize_timer_percent(value: &mut f64, total: f64) {
if total <= f64::EPSILON {
*value = 0.0;
} else {
*value /= total / 100.0;
}
}
fn flatten_runtime_args(parameter_values: Option<&[f64]>, unknowns: &[f64]) -> Vec<f64> {
let mut flat_args =
Vec::with_capacity(parameter_values.map_or(0, |params| params.len()) + unknowns.len());
if let Some(params) = parameter_values {
flat_args.extend_from_slice(params);
}
flat_args.extend_from_slice(unknowns);
flat_args
}
fn panic_symbolic_numeric_backend_requested() -> ! {
panic!(
"BVP symbolic pipeline cannot materialize BackendSelectionPolicy::NumericOnly; \
use NRBVP::set_numeric_rhs/with_numeric_rhs so the solver can build the pure numeric \
finite-difference discretization path, or choose LambdifyOnly/AotOnly/PreferAotThenLambdify"
)
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BvpBackendKind {
Numeric,
Lambdify,
Aot,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BvpMatrixBackend {
Dense,
SprsCsMat,
FaerSparseCol,
Banded,
}
impl BvpMatrixBackend {
pub const fn legacy_method(self) -> &'static str {
match self {
Self::Dense => "Dense",
Self::SprsCsMat => "Sparse_1",
Self::FaerSparseCol => "Sparse",
Self::Banded => "Banded",
}
}
pub fn from_legacy_method(method: &str) -> Option<Self> {
match method {
"Dense" => Some(Self::Dense),
"Sparse_1" => Some(Self::SprsCsMat),
"Sparse" => Some(Self::FaerSparseCol),
"Banded" => Some(Self::Banded),
_ => None,
}
}
pub const fn from_provider_matrix_backend(matrix_backend: MatrixBackend) -> Self {
match matrix_backend {
MatrixBackend::Dense => Self::Dense,
MatrixBackend::Banded => Self::Banded,
MatrixBackend::SparseCol => Self::FaerSparseCol,
MatrixBackend::CsMat => Self::SprsCsMat,
MatrixBackend::CsMatrix => Self::FaerSparseCol,
MatrixBackend::ValuesOnly => Self::FaerSparseCol,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BvpBackendConfig {
pub backend_kind: BvpBackendKind,
pub matrix_backend: BvpMatrixBackend,
}
impl BvpBackendConfig {
pub const fn new(backend_kind: BvpBackendKind, matrix_backend: BvpMatrixBackend) -> Self {
Self {
backend_kind,
matrix_backend,
}
}
pub const fn lambdify(matrix_backend: BvpMatrixBackend) -> Self {
Self::new(BvpBackendKind::Lambdify, matrix_backend)
}
}
impl Default for BvpBackendConfig {
fn default() -> Self {
Self::lambdify(BvpMatrixBackend::FaerSparseCol)
}
}
pub(crate) fn bvp_symbolic_variables_for_functions(
vector_of_functions: &[Expr],
) -> Vec<Vec<String>> {
vector_of_functions
.iter()
.map(Expr::all_arguments_are_variables)
.collect()
}
pub(crate) fn bvp_symbolic_parse_variables(variable_string: &[String]) -> Vec<Expr> {
Expr::parse_vector_expression(variable_string.iter().map(|s| s.as_str()).collect())
}
pub(crate) fn bvp_symbolic_jacobian_smart(
vector_of_functions: &[Expr],
vector_of_variables_len: usize,
variable_string: &[String],
variables_for_all_discrete: &[Vec<String>],
) -> (Vec<Vec<Expr>>, Vec<(usize, usize, Expr)>) {
let rows_and_sparse: Vec<(Vec<Expr>, Vec<(usize, usize, Expr)>)> = vector_of_functions
.par_iter()
.enumerate()
.map(|(i, function)| {
let mut vector_of_partial_derivatives = Vec::new();
let mut sparse_entries = Vec::new();
for j in 0..vector_of_variables_len {
let variable = &variable_string[j];
let list_of_variables_for_this_eq = &variables_for_all_discrete[i];
if list_of_variables_for_this_eq.contains(variable) {
let mut partial = Expr::diff(function, variable);
partial = partial.simplify();
if !partial.is_zero() {
sparse_entries.push((i, j, partial.clone()));
}
vector_of_partial_derivatives.push(partial);
} else {
vector_of_partial_derivatives.push(Expr::Const(0.0));
}
}
(vector_of_partial_derivatives, sparse_entries)
})
.collect();
let symbolic_jacobian = rows_and_sparse.iter().map(|(row, _)| row.clone()).collect();
let symbolic_jacobian_sparse = rows_and_sparse
.into_iter()
.flat_map(|(_, entries)| entries)
.collect();
(symbolic_jacobian, symbolic_jacobian_sparse)
}
pub(crate) fn bvp_symbolic_bandwidth(symbolic_jacobian: &[Vec<Expr>]) -> (usize, usize) {
let n = symbolic_jacobian.len();
(0..n)
.into_par_iter()
.map(|i| {
let mut row_kl = 0;
let mut row_ku = 0;
for j in 0..n {
if symbolic_jacobian[i][j] != Expr::Const(0.0) {
if j > i {
row_ku = std::cmp::max(row_ku, j - i);
} else if i > j {
row_kl = std::cmp::max(row_kl, i - j);
}
}
}
(row_kl, row_ku)
})
.reduce(
|| (0, 0),
|acc, row| (std::cmp::max(acc.0, row.0), std::cmp::max(acc.1, row.1)),
)
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BvpSymbolicAssemblyBackend {
ExprLegacy,
AtomView,
}
impl Default for BvpSymbolicAssemblyBackend {
fn default() -> Self {
Self::ExprLegacy
}
}
pub struct Jacobian {
pub vector_of_functions: Vec<Expr>,
pub lambdified_functions: Vec<Box<dyn Fn(Vec<f64>) -> f64>>,
pub method: String,
pub backend_config: BvpBackendConfig,
pub banded_linear_solver_config: LinearSolverConfig,
pub banded_lambdify_config: BandedLambdifyConfig,
pub symbolic_assembly_backend: BvpSymbolicAssemblyBackend,
pub vector_of_variables: Vec<Expr>,
pub variable_string: Vec<String>,
pub parameters_string: Vec<String>,
pub parameter_values: Option<Vec<f64>>,
pub symbolic_jacobian: Vec<Vec<Expr>>,
pub symbolic_jacobian_sparse: Vec<(usize, usize, Expr)>,
pub lambdified_jac_element: Option<Box<dyn Fn(f64, usize, usize) -> f64>>,
pub jac_function: Option<Box<dyn Jac>>,
pub residiual_function: Box<dyn Fun>,
pub bounds: Option<Vec<(f64, f64)>>,
pub rel_tolerance_vec: Option<Vec<f64>>,
pub bandwidth: Option<(usize, usize)>,
pub variables_for_all_disrete: Vec<Vec<String>>,
pub BC_pos_n_values: Vec<(usize, usize, f64)>,
pub last_generate_timer_snapshot: Option<HashMap<String, f64>>,
pub last_symbolic_jacobian_timer_snapshot: Option<HashMap<String, f64>>,
pub last_lambdify_binding_timer_snapshot: Option<HashMap<String, f64>>,
pub atom_discretized_system: Option<DiscretizedBvpAtomSystem>,
}
#[derive(Debug, Clone)]
pub struct BvpPreparedSparseAotProblem {
pub residual_fn_name: String,
pub jacobian_fn_name: String,
pub variable_names: Vec<String>,
pub param_names: Vec<String>,
pub residuals: Vec<Expr>,
pub sparse_entries: Vec<(usize, usize, Expr)>,
pub shape: (usize, usize),
pub residual_strategy: ResidualChunkingStrategy,
pub jacobian_strategy: SparseChunkingStrategy,
pub bandwidth: Option<(usize, usize)>,
pub symbolic_assembly_backend: BvpSymbolicAssemblyBackend,
pub atom_discretized_system: Option<DiscretizedBvpAtomSystem>,
}
#[derive(Clone, Debug, Default)]
pub struct BvpGeneratedAotCrateBreakdown {
pub jacobian_prepare_ms: f64,
pub atom_sparse_lookup_prepare_ms: f64,
pub atom_sparse_jacobian_build_ms: f64,
pub atom_finalize_codegen_plan_ms: f64,
pub atom_sparse_nnz: usize,
pub atom_residual_view_collect_ms: f64,
pub atom_residual_lower_many_ms: f64,
pub atom_residual_peephole_ms: f64,
pub atom_residual_reuse_temps_ms: f64,
pub atom_residual_push_ms: f64,
pub atom_sparse_view_collect_ms: f64,
pub atom_sparse_lower_many_ms: f64,
pub atom_sparse_peephole_ms: f64,
pub atom_sparse_reuse_temps_ms: f64,
pub atom_sparse_push_ms: f64,
pub prepared_module_init_ms: f64,
pub prepared_residual_blocks_ms: f64,
pub prepared_jacobian_blocks_ms: f64,
pub module_build_ms: f64,
pub source_probe_emit_ms: f64,
pub language_source_emit_ms: f64,
pub c_header_emit_ms: f64,
pub artifact_packaging_ms: f64,
pub source_emit_ms: f64,
pub source_kb: f64,
pub module_blocks: usize,
pub total_block_instructions: usize,
pub total_block_temps: usize,
pub max_block_instructions: usize,
pub total_block_outputs: usize,
}
#[derive(Debug, Clone)]
pub struct BvpSelectedSparseBackend {
pub prepared_problem: BvpPreparedSparseAotProblem,
pub requested_backend: BackendKind,
pub effective_backend: SelectedBackendKind,
pub matrix_backend: MatrixBackend,
pub aot_resolution: Option<ResolvedAotArtifact>,
}
impl BvpSelectedSparseBackend {
pub fn is_compiled_aot(&self) -> bool {
self.effective_backend == SelectedBackendKind::AotCompiled
}
pub fn problem_key(&self) -> String {
self.prepared_problem
.problem_key_for_matrix_backend(self.matrix_backend)
}
}
#[derive(Debug, Clone)]
pub enum BvpSparseExecutionPlan {
NumericRequested(BvpSelectedSparseBackend),
LambdifyReady(BvpSelectedSparseBackend),
AotCompiled(BvpSelectedSparseBackend),
AotRegisteredButNotBuilt(BvpSelectedSparseBackend),
AotMissing(BvpSelectedSparseBackend),
}
impl BvpSparseExecutionPlan {
pub fn selected(&self) -> &BvpSelectedSparseBackend {
match self {
Self::NumericRequested(selected)
| Self::LambdifyReady(selected)
| Self::AotCompiled(selected)
| Self::AotRegisteredButNotBuilt(selected)
| Self::AotMissing(selected) => selected,
}
}
}
pub struct BvpSparseSolverProvider<'a> {
jacobian: &'a mut Jacobian,
execution: BvpSparseExecutionPlan,
structure: SparseJacobianStructure,
}
impl<'a> BvpSparseSolverProvider<'a> {
pub fn execution_plan(&self) -> &BvpSparseExecutionPlan {
&self.execution
}
pub fn effective_backend(&self) -> SelectedBackendKind {
self.execution.selected().effective_backend
}
pub fn requested_backend(&self) -> BackendKind {
self.execution.selected().requested_backend
}
pub fn is_runtime_callable(&self) -> bool {
matches!(self.execution, BvpSparseExecutionPlan::LambdifyReady(_))
|| matches!(self.execution, BvpSparseExecutionPlan::AotCompiled(_))
&& resolve_linked_sparse_backend(&self.execution.selected().problem_key()).is_some()
}
pub fn resolved_aot_artifact(&self) -> Option<&ResolvedAotArtifact> {
self.execution.selected().aot_resolution.as_ref()
}
pub fn residual_len(&self) -> usize {
self.execution.selected().prepared_problem.shape.0
}
pub fn jacobian_shape(&self) -> (usize, usize) {
self.execution.selected().prepared_problem.shape
}
pub fn jacobian_structure(&self) -> &SparseJacobianStructure {
&self.structure
}
fn typed_variables_from_args(&self, args: &[f64]) -> Box<dyn VectorType> {
let dense = DVector::from_column_slice(args);
let backend = BvpMatrixBackend::from_provider_matrix_backend(
self.execution.selected().matrix_backend,
);
Vectors_type_casting(&dense, backend.legacy_method().to_string())
}
pub fn residual_into(&self, args: &[f64], out: &mut [f64]) {
assert!(
self.is_runtime_callable(),
"BvpSparseSolverProvider residual_into is only callable for runtime-ready sparse plans"
);
assert_eq!(
out.len(),
self.residual_len(),
"residual output length must match residual_len"
);
match &self.execution {
BvpSparseExecutionPlan::LambdifyReady(_) => {
let typed = self.typed_variables_from_args(args);
let residual = self.jacobian.residiual_function.call(1.0, &*typed);
let dense = residual.to_DVectorType();
out.copy_from_slice(dense.as_slice());
}
BvpSparseExecutionPlan::AotCompiled(selected) => {
let linked = resolve_linked_sparse_backend(&selected.problem_key())
.expect("AotCompiled provider path must resolve a linked backend");
let flat_args =
flatten_runtime_args(self.jacobian.parameter_values.as_deref(), args);
(linked.residual_eval)(&flat_args, out);
}
_ => {
unreachable!("runtime-callable provider should only expose lambdify or linked AOT")
}
}
}
pub fn jacobian_values_into(&mut self, args: &[f64], values_out: &mut [f64]) {
assert!(
self.is_runtime_callable(),
"BvpSparseSolverProvider jacobian_values_into is only callable for runtime-ready sparse plans"
);
assert_eq!(
values_out.len(),
self.structure.nnz(),
"Jacobian values output length must match sparse structure nnz"
);
match &mut self.execution {
BvpSparseExecutionPlan::LambdifyReady(_) => {
let typed = self.typed_variables_from_args(args);
let dense = self
.jacobian
.jac_function
.as_mut()
.expect("jacobian callback must exist for LambdifyReady plan")
.call(1.0, &*typed)
.to_DMatrixType();
for (value, (&row, &col)) in values_out.iter_mut().zip(
self.structure
.row_indices
.iter()
.zip(self.structure.col_indices.iter()),
) {
*value = dense[(row, col)];
}
}
BvpSparseExecutionPlan::AotCompiled(selected) => {
let linked = resolve_linked_sparse_backend(&selected.problem_key())
.expect("AotCompiled provider path must resolve a linked backend");
let flat_args =
flatten_runtime_args(self.jacobian.parameter_values.as_deref(), args);
(linked.jacobian_values_eval)(&flat_args, values_out);
}
_ => {
unreachable!("runtime-callable provider should only expose lambdify or linked AOT")
}
}
}
}
pub struct BvpSparseSolverBundle {
pub execution: BvpSparseExecutionPlan,
pub residual_function: Option<Box<dyn Fun>>,
pub jacobian_function: Option<Box<dyn Jac>>,
pub runtime_diagnostics: HashMap<String, String>,
pub banded_linear_solver_config: LinearSolverConfig,
pub variable_string: Vec<String>,
pub parameter_values: Option<Vec<f64>>,
pub bounds_vec: Option<Vec<(f64, f64)>>,
pub rel_tolerance_vec: Option<Vec<f64>>,
pub bandwidth: Option<(usize, usize)>,
pub bc_position_and_value: Vec<(usize, usize, f64)>,
pub sparse_structure: SparseJacobianStructure,
}
pub struct BvpLegacySolverBundle {
pub residual_function: Box<dyn Fun>,
pub jacobian_function: Option<Box<dyn Jac>>,
pub variable_string: Vec<String>,
pub bounds_vec: Option<Vec<(f64, f64)>>,
pub rel_tolerance_vec: Option<Vec<f64>>,
pub bandwidth: Option<(usize, usize)>,
pub bc_position_and_value: Vec<(usize, usize, f64)>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BvpBackendIntegrationError {
PipelinePanicked(String),
InvalidLogLevel { level: String },
LogFileCreationFailed { path: String, message: String },
CompiledAotRequiredButUnavailable {
problem_key: String,
effective_backend: SelectedBackendKind,
},
AutomaticAotBuildRequested { problem_key: String },
AutomaticAotRebuildRequested { problem_key: String },
AutomaticAotBuildFailed {
problem_key: String,
message: String,
},
CompiledAotRuntimeUnavailable { problem_key: String },
}
impl BvpSparseSolverBundle {
pub fn effective_backend(&self) -> SelectedBackendKind {
self.execution.selected().effective_backend
}
pub fn requested_backend(&self) -> BackendKind {
self.execution.selected().requested_backend
}
pub fn is_runtime_callable(&self) -> bool {
self.residual_function.is_some() && self.jacobian_function.is_some()
}
pub fn resolved_aot_artifact(&self) -> Option<&ResolvedAotArtifact> {
self.execution.selected().aot_resolution.as_ref()
}
pub fn residual_len(&self) -> usize {
self.execution.selected().prepared_problem.shape.0
}
pub fn jacobian_shape(&self) -> (usize, usize) {
self.execution.selected().prepared_problem.shape
}
pub fn residual_call(&self, p: f64, y: &dyn VectorType) -> Option<Box<dyn VectorType>> {
self.residual_function.as_ref().map(|fun| fun.call(p, y))
}
pub fn jacobian_call(&mut self, p: f64, y: &dyn VectorType) -> Option<Box<dyn MatrixType>> {
self.jacobian_function.as_mut().map(|jac| jac.call(p, y))
}
pub fn into_runtime_callbacks(self) -> Option<(Box<dyn Fun>, Box<dyn Jac>)> {
match (self.residual_function, self.jacobian_function) {
(Some(fun), Some(jac)) => Some((fun, jac)),
_ => None,
}
}
pub fn runtime_diagnostics(&self) -> &HashMap<String, String> {
&self.runtime_diagnostics
}
pub fn refresh_linked_runtime_diagnostics(
&mut self,
execution_policy: &str,
parallel_config: Option<ParallelExecutorConfig>,
) {
self.runtime_diagnostics =
linked_runtime_diagnostics_for_bundle(self, execution_policy, parallel_config);
}
pub fn rebind_linked_runtime_callbacks(
&mut self,
parameter_values: Option<&[f64]>,
parallel_config: Option<ParallelExecutorConfig>,
) -> bool {
if !matches!(self.execution, BvpSparseExecutionPlan::AotCompiled(_)) {
return false;
}
let selected = self.execution.selected();
let Some((residual, jacobian)) = linked_runtime_callbacks_for_matrix_backend(
&self.execution.selected().prepared_problem,
selected.matrix_backend,
parameter_values.or(self.parameter_values.as_deref()),
parallel_config,
self.banded_linear_solver_config,
) else {
return false;
};
self.residual_function = Some(residual);
self.jacobian_function = Some(jacobian);
true
}
}
fn insert_linked_stage_runtime_diagnostics(
diagnostics: &mut HashMap<String, String>,
prefix: &str,
workload: usize,
chunk_count: usize,
requested_jobs: usize,
should_parallelize: bool,
) {
let planned_ranges = linked_contiguous_job_ranges(chunk_count, requested_jobs);
let planned_jobs = planned_ranges.len().max(1);
let actual_jobs = if should_parallelize && planned_jobs > 1 {
planned_jobs
} else {
1
};
let fallback = actual_jobs == 1;
let fallback_reason = if chunk_count <= 1 {
"single_chunk"
} else if requested_jobs <= 1 {
"single_requested_job"
} else if !should_parallelize {
"fallback_policy_or_small_workload"
} else {
"none"
};
diagnostics.insert(format!("{prefix}.chunk_count"), chunk_count.to_string());
diagnostics.insert(
format!("{prefix}.requested_jobs"),
requested_jobs.to_string(),
);
diagnostics.insert(format!("{prefix}.planned_jobs"), planned_jobs.to_string());
diagnostics.insert(format!("{prefix}.actual_jobs"), actual_jobs.to_string());
diagnostics.insert(format!("{prefix}.fallback"), fallback.to_string());
diagnostics.insert(
format!("{prefix}.fallback_reason"),
fallback_reason.to_string(),
);
diagnostics.insert(
format!("{prefix}.work_per_job"),
workload.div_ceil(actual_jobs).to_string(),
);
diagnostics.insert(
format!("{prefix}.work_per_chunk"),
if chunk_count == 0 {
"0".to_string()
} else {
workload.div_ceil(chunk_count).to_string()
},
);
}
fn linked_runtime_diagnostics_for_bundle(
bundle: &BvpSparseSolverBundle,
execution_policy: &str,
parallel_config: Option<ParallelExecutorConfig>,
) -> HashMap<String, String> {
let selected = bundle.execution.selected();
let mut diagnostics = HashMap::new();
diagnostics.insert(
"aot.runtime.execution_policy".to_string(),
execution_policy.to_string(),
);
diagnostics.insert(
"aot.runtime.effective_backend".to_string(),
format!("{:?}", selected.effective_backend),
);
diagnostics.insert(
"aot.runtime.matrix_backend".to_string(),
format!("{:?}", selected.matrix_backend),
);
diagnostics.insert(
"aot.runtime.parallel_requested".to_string(),
parallel_config.is_some().to_string(),
);
let auto_plan = selected.prepared_problem.auto_parallel_plan();
diagnostics.insert(
"aot.auto.execution_mode".to_string(),
format!("{:?}", auto_plan.execution_mode),
);
diagnostics.insert(
"aot.auto.workers".to_string(),
auto_plan.workers.to_string(),
);
diagnostics.insert(
"aot.auto.min_work_per_job".to_string(),
auto_plan.min_work_per_job.to_string(),
);
diagnostics.insert(
"aot.auto.residual.reason".to_string(),
auto_plan.residual_stage.reason.as_str().to_string(),
);
diagnostics.insert(
"aot.auto.residual.work_per_job".to_string(),
auto_plan.residual_stage.work_per_job.to_string(),
);
diagnostics.insert(
"aot.auto.residual.work_per_chunk".to_string(),
auto_plan.residual_stage.work_per_chunk.to_string(),
);
diagnostics.insert(
"aot.auto.sparse_jacobian.reason".to_string(),
auto_plan.sparse_stage.reason.as_str().to_string(),
);
diagnostics.insert(
"aot.auto.sparse_jacobian.work_per_job".to_string(),
auto_plan.sparse_stage.work_per_job.to_string(),
);
diagnostics.insert(
"aot.auto.sparse_jacobian.work_per_chunk".to_string(),
auto_plan.sparse_stage.work_per_chunk.to_string(),
);
if !matches!(selected.effective_backend, SelectedBackendKind::AotCompiled) {
diagnostics.insert(
"aot.runtime.binding".to_string(),
"non_aot_or_lambdify".to_string(),
);
return diagnostics;
}
let problem_key = selected
.prepared_problem
.problem_key_for_matrix_backend(selected.matrix_backend);
let Some(linked) = resolve_linked_sparse_backend(&problem_key) else {
diagnostics.insert(
"aot.runtime.binding".to_string(),
"compiled_artifact_not_linked".to_string(),
);
return diagnostics;
};
diagnostics.insert("aot.runtime.binding".to_string(), "linked".to_string());
diagnostics.insert("aot.runtime.problem_key".to_string(), problem_key);
diagnostics.insert(
"aot.runtime.residual_len".to_string(),
linked.residual_len.to_string(),
);
diagnostics.insert("aot.runtime.sparse_nnz".to_string(), linked.nnz.to_string());
match parallel_config {
Some(config) => {
let residual_requested_jobs = linked_max_jobs(config, false);
let residual_should_parallelize = linked_should_parallelize(
config,
linked.residual_len,
linked.residual_chunks.len(),
false,
);
insert_linked_stage_runtime_diagnostics(
&mut diagnostics,
"aot.runtime.residual",
linked.residual_len,
linked.residual_chunks.len(),
residual_requested_jobs,
residual_should_parallelize,
);
let sparse_requested_jobs = linked_max_jobs(config, true);
let sparse_should_parallelize = linked_should_parallelize(
config,
linked.nnz,
linked.jacobian_value_chunks.len(),
true,
);
insert_linked_stage_runtime_diagnostics(
&mut diagnostics,
"aot.runtime.sparse_jacobian",
linked.nnz,
linked.jacobian_value_chunks.len(),
sparse_requested_jobs,
sparse_should_parallelize,
);
}
None => {
insert_linked_stage_runtime_diagnostics(
&mut diagnostics,
"aot.runtime.residual",
linked.residual_len,
linked.residual_chunks.len(),
1,
false,
);
insert_linked_stage_runtime_diagnostics(
&mut diagnostics,
"aot.runtime.sparse_jacobian",
linked.nnz,
linked.jacobian_value_chunks.len(),
1,
false,
);
}
}
diagnostics
}
impl Jacobian {
pub fn into_legacy_solver_bundle(self) -> BvpLegacySolverBundle {
BvpLegacySolverBundle {
residual_function: self.residiual_function,
jacobian_function: self.jac_function,
variable_string: self.variable_string,
bounds_vec: self.bounds,
rel_tolerance_vec: self.rel_tolerance_vec,
bandwidth: self.bandwidth,
bc_position_and_value: self.BC_pos_n_values,
}
}
#[allow(clippy::too_many_arguments)]
pub fn generate_legacy_solver_bundle_with_params(
mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
params: Option<&[&str]>,
t0: f64,
param: Option<String>,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
border_conditions: HashMap<String, Vec<(usize, f64)>>,
bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
method: String,
bandwidth: Option<(usize, usize)>,
) -> BvpLegacySolverBundle {
self.generate_BVP_with_params(
eq_system,
values,
arg,
params,
t0,
param,
n_steps,
h,
mesh,
border_conditions,
bounds,
rel_tolerance,
scheme,
method,
bandwidth,
);
self.into_legacy_solver_bundle()
}
}
impl BvpPreparedSparseAotProblem {
fn infer_node_major_layout(&self) -> Result<NodeMajorLayout, String> {
if self.variable_names.is_empty() || self.shape.0 == 0 || self.shape.0 != self.shape.1 {
return Err("prepared BVP problem must describe a non-empty square system".into());
}
fn base_variable_name(name: &str) -> &str {
name.rsplit_once('_')
.and_then(|(base, suffix)| suffix.parse::<usize>().ok().map(|_| base))
.unwrap_or(name)
}
let mut seen = HashSet::new();
let mut vars_per_node = 0usize;
for name in &self.variable_names {
if seen.insert(base_variable_name(name).to_string()) {
vars_per_node += 1;
}
}
if vars_per_node == 0 || self.shape.1 % vars_per_node != 0 {
return Err(format!(
"node-major layout inference failed: unknowns={} vars_per_node={vars_per_node}",
self.shape.1
));
}
let n_nodes = self.shape.1 / vars_per_node;
NodeMajorLayout::new(n_nodes, vars_per_node).map_err(|err| err.to_string())
}
fn flattened_input_names(&self) -> Vec<&str> {
let mut names = Vec::with_capacity(
self.param_names
.len()
.saturating_add(self.variable_names.len()),
);
names.extend(self.param_names.iter().map(|name| name.as_str()));
names.extend(self.variable_names.iter().map(|name| name.as_str()));
names
}
fn borrowed_sparse_entries(&self) -> Vec<SparseExprEntry<'_>> {
self.sparse_entries
.iter()
.map(|(row, col, expr)| SparseExprEntry {
row: *row,
col: *col,
expr,
})
.collect()
}
fn inferred_bandwidth_from_sparse_entries(&self) -> (usize, usize) {
self.sparse_entries
.iter()
.fold((0usize, 0usize), |(kl, ku), (row, col, _)| {
if row >= col {
(kl.max(row - col), ku)
} else {
(kl, ku.max(col - row))
}
})
}
fn effective_bandwidth(&self) -> (usize, usize) {
self.bandwidth
.unwrap_or_else(|| self.inferred_bandwidth_from_sparse_entries())
}
fn borrowed_banded_entries(&self) -> Vec<BandedExprEntry<'_>> {
let (kl, ku) = self.effective_bandwidth();
let mut entries: Vec<(usize, usize, isize, &Expr)> = self
.sparse_entries
.iter()
.map(|(row, col, expr)| {
let diag_offset = *col as isize - *row as isize;
(*row, *col, diag_offset, expr)
})
.collect();
entries.sort_by(|lhs, rhs| lhs.0.cmp(&rhs.0).then_with(|| lhs.1.cmp(&rhs.1)));
let mut banded_entries = Vec::with_capacity(entries.len());
for (row, col, diag_offset, expr) in entries {
assert!(
diag_offset >= -(kl as isize) && diag_offset <= ku as isize,
"sparse Jacobian entry ({row}, {col}) lies outside declared banded width"
);
let diag_position = if diag_offset >= 0 { row } else { col };
banded_entries.push(BandedExprEntry {
row,
col,
diag_offset,
diag_position,
expr,
});
}
banded_entries
}
fn residual_runtime_plan(&self) -> ResidualRuntimePlan<'_> {
let chunk_size = match self.residual_strategy {
ResidualChunkingStrategy::Whole => self.residuals.len().max(1),
ResidualChunkingStrategy::ByTargetChunkCount { target_chunks } => {
assert!(target_chunks > 0, "target_chunks must be positive");
self.residuals.len().max(1).div_ceil(target_chunks).max(1)
}
ResidualChunkingStrategy::ByOutputCount {
max_outputs_per_chunk,
} => {
assert!(
max_outputs_per_chunk > 0,
"max_outputs_per_chunk must be positive"
);
max_outputs_per_chunk
}
};
let input_names = self.flattened_input_names();
let chunks = self
.residuals
.chunks(chunk_size)
.enumerate()
.map(|(chunk_index, residuals)| {
let output_offset = chunk_index * chunk_size;
let fn_name = if chunk_index == 0 && residuals.len() == self.residuals.len() {
self.residual_fn_name.clone()
} else {
format!("{}_chunk_{}", self.residual_fn_name, chunk_index)
};
let outputs = residuals
.iter()
.map(|expr| PlannedOutput {
expr,
coordinate: None,
})
.collect();
ResidualChunkPlan {
fn_name: fn_name.clone(),
output_offset,
residuals,
plan: CodegenTaskPlan {
fn_name: std::borrow::Cow::Owned(fn_name),
kind: CodegenTaskKind::Residual,
input_names: input_names.clone(),
outputs,
layout: CodegenOutputLayout::Vector {
len: residuals.len(),
},
},
}
})
.collect();
ResidualRuntimePlan {
fn_name: self.residual_fn_name.as_str(),
output_len: self.residuals.len(),
input_names,
chunks,
}
}
fn sparse_runtime_plan(&self) -> SparseJacobianRuntimePlan<'_> {
let all_entries = self.borrowed_sparse_entries();
let input_names = self.flattened_input_names();
let chunked_entries: Vec<(usize, Vec<SparseExprEntry<'_>>)> = match self.jacobian_strategy {
SparseChunkingStrategy::Whole => vec![(0usize, all_entries.clone())],
SparseChunkingStrategy::ByTargetChunkCount { target_chunks } => {
assert!(target_chunks > 0, "target_chunks must be positive");
let max_entries_per_chunk = all_entries.len().max(1).div_ceil(target_chunks).max(1);
chunk_sparse_entries_by_nnz(&all_entries, max_entries_per_chunk)
}
SparseChunkingStrategy::ByNonZeroCount {
max_entries_per_chunk,
} => {
assert!(
max_entries_per_chunk > 0,
"max_entries_per_chunk must be positive"
);
chunk_sparse_entries_by_nnz(&all_entries, max_entries_per_chunk)
}
SparseChunkingStrategy::ByRowCount { rows_per_chunk } => {
assert!(rows_per_chunk > 0, "rows_per_chunk must be positive");
chunk_sparse_entries_by_rows(&all_entries, rows_per_chunk, self.shape.0)
}
};
let chunks = chunked_entries
.into_iter()
.enumerate()
.map(|(chunk_index, (entry_offset, entries))| {
let fn_name = if chunk_index == 0 && entries.len() == all_entries.len() {
self.jacobian_fn_name.clone()
} else {
format!("{}_chunk_{}", self.jacobian_fn_name, chunk_index)
};
let outputs = entries
.iter()
.map(|entry| PlannedOutput {
expr: entry.expr,
coordinate: Some((entry.row, entry.col)),
})
.collect();
SparseJacobianValuesChunkPlan {
fn_name: fn_name.clone(),
value_offset: entry_offset,
entries: entries.clone(),
plan: CodegenTaskPlan {
fn_name: std::borrow::Cow::Owned(fn_name),
kind: CodegenTaskKind::SparseJacobianValues,
input_names: input_names.clone(),
outputs,
layout: CodegenOutputLayout::SparseValues {
rows: self.shape.0,
cols: self.shape.1,
nnz: entries.len(),
},
},
}
})
.collect();
SparseJacobianRuntimePlan {
fn_name: self.jacobian_fn_name.as_str(),
input_names,
structure: SparseJacobianStructure {
rows: self.shape.0,
cols: self.shape.1,
row_indices: all_entries.iter().map(|entry| entry.row).collect(),
col_indices: all_entries.iter().map(|entry| entry.col).collect(),
},
chunks,
}
}
fn banded_chunking_strategy(&self) -> BandedChunkingStrategy {
match self.jacobian_strategy {
SparseChunkingStrategy::Whole => BandedChunkingStrategy::Whole,
SparseChunkingStrategy::ByTargetChunkCount { target_chunks } => {
BandedChunkingStrategy::ByTargetChunkCount { target_chunks }
}
SparseChunkingStrategy::ByNonZeroCount {
max_entries_per_chunk,
} => {
let chunks_per_worker = max_entries_per_chunk.max(1);
let (kl, ku) = self.effective_bandwidth();
recommended_banded_chunking_for_parallelism(
kl.saturating_add(ku).saturating_add(1),
chunks_per_worker,
)
}
SparseChunkingStrategy::ByRowCount { rows_per_chunk } => {
let target_chunks = self.shape.0.max(1).div_ceil(rows_per_chunk.max(1));
BandedChunkingStrategy::ByTargetChunkCount {
target_chunks: target_chunks.max(1),
}
}
}
}
fn banded_runtime_plan(&self) -> BandedJacobianRuntimePlan<'_> {
let (kl, ku) = self.effective_bandwidth();
let all_entries = self.borrowed_banded_entries();
let input_names = self.flattened_input_names();
let (chunk_entries, chunk_offsets): (Vec<Vec<BandedExprEntry<'_>>>, Vec<usize>) = match self
.banded_chunking_strategy()
{
BandedChunkingStrategy::Whole => {
if all_entries.is_empty() {
(Vec::new(), Vec::new())
} else {
(vec![all_entries.clone()], vec![0])
}
}
BandedChunkingStrategy::ByTargetChunkCount { target_chunks } => {
assert!(target_chunks > 0, "target_chunks must be positive");
let entries_per_chunk = all_entries.len().max(1).div_ceil(target_chunks).max(1);
let mut chunks = Vec::new();
let mut offsets = Vec::new();
for (chunk_index, entries) in all_entries.chunks(entries_per_chunk).enumerate() {
offsets.push(chunk_index * entries_per_chunk);
chunks.push(entries.to_vec());
}
(chunks, offsets)
}
BandedChunkingStrategy::ByDiagonalCount {
diagonals_per_chunk,
} => {
assert!(
diagonals_per_chunk > 0,
"diagonals_per_chunk must be positive"
);
let mut chunk_entries: Vec<Vec<BandedExprEntry<'_>>> = Vec::new();
let mut chunk_offsets: Vec<usize> = Vec::new();
for (entry_index, entry) in all_entries.iter().copied().enumerate() {
let bucket = ((entry.diag_offset + kl as isize) as usize) / diagonals_per_chunk;
if bucket >= chunk_entries.len() {
chunk_entries.resize_with(bucket + 1, Vec::new);
chunk_offsets.resize(bucket + 1, 0);
}
if chunk_entries[bucket].is_empty() {
chunk_offsets[bucket] = entry_index;
}
chunk_entries[bucket].push(entry);
}
(chunk_entries, chunk_offsets)
}
};
let chunk_count = chunk_entries
.iter()
.filter(|entries| !entries.is_empty())
.count();
let chunks = chunk_entries
.into_iter()
.enumerate()
.filter(|(_, entries)| !entries.is_empty())
.map(|(chunk_index, entries)| {
let fn_name = if chunk_index == 0 && chunk_count == 1 {
self.jacobian_fn_name.clone()
} else {
format!("{}_chunk_{}", self.jacobian_fn_name, chunk_index)
};
let outputs = entries
.iter()
.map(|entry| PlannedOutput {
expr: entry.expr,
coordinate: Some((entry.row, entry.col)),
})
.collect();
crate::symbolic::codegen::codegen_runtime_api::BandedJacobianValuesChunkPlan {
fn_name: fn_name.clone(),
value_offset: chunk_offsets[chunk_index],
entries: entries.clone(),
plan: CodegenTaskPlan {
fn_name: std::borrow::Cow::Owned(fn_name),
kind: CodegenTaskKind::BandedJacobianValues,
input_names: input_names.clone(),
outputs,
layout: CodegenOutputLayout::SparseValues {
rows: self.shape.0,
cols: self.shape.1,
nnz: entries.len(),
},
},
}
})
.collect();
BandedJacobianRuntimePlan {
fn_name: self.jacobian_fn_name.as_str(),
input_names,
structure: crate::symbolic::codegen::codegen_runtime_api::BandedJacobianStructure {
rows: self.shape.0,
cols: self.shape.1,
kl,
ku,
diagonal_offsets: all_entries.iter().map(|entry| entry.diag_offset).collect(),
diagonal_positions: all_entries
.iter()
.map(|entry| entry.diag_position)
.collect(),
},
chunks,
}
}
pub fn as_prepared_problem(&self) -> PreparedSparseProblem<'_> {
PreparedSparseProblem::new(
BackendKind::Aot,
MatrixBackend::SparseCol,
self.residual_runtime_plan(),
self.sparse_runtime_plan(),
)
}
pub fn as_prepared_banded_problem(&self) -> PreparedBandedProblem<'_> {
PreparedBandedProblem::new(
BackendKind::Aot,
MatrixBackend::Banded,
self.residual_runtime_plan(),
self.banded_runtime_plan(),
)
}
pub fn as_prepared_problem_for_matrix_backend(
&self,
matrix_backend: MatrixBackend,
) -> PreparedProblem<'_> {
match matrix_backend {
MatrixBackend::Banded => PreparedProblem::banded(self.as_prepared_banded_problem()),
_ => PreparedProblem::sparse(self.as_prepared_problem()),
}
}
pub fn problem_key(&self) -> String {
self.problem_key_for_matrix_backend(MatrixBackend::SparseCol)
}
pub fn problem_key_for_matrix_backend(&self, matrix_backend: MatrixBackend) -> String {
let prepared = self.as_prepared_problem_for_matrix_backend(matrix_backend);
PreparedProblemManifest::from(&prepared).problem_key()
}
pub fn codegen_module_with_breakdown(
&self,
module_name: &str,
) -> (CodegenModule, BvpGeneratedAotCrateBreakdown) {
self.codegen_module_with_breakdown_for_matrix_backend(module_name, MatrixBackend::SparseCol)
}
pub fn codegen_module_with_breakdown_for_matrix_backend(
&self,
module_name: &str,
matrix_backend: MatrixBackend,
) -> (CodegenModule, BvpGeneratedAotCrateBreakdown) {
self.codegen_module_with_breakdown_for_matrix_backend_and_atom_profile(
module_name,
matrix_backend,
AtomOptimizationProfile::Full,
)
}
pub fn codegen_module_with_breakdown_for_matrix_backend_and_atom_profile(
&self,
module_name: &str,
matrix_backend: MatrixBackend,
atom_profile: AtomOptimizationProfile,
) -> (CodegenModule, BvpGeneratedAotCrateBreakdown) {
self.codegen_module_with_breakdown_for_matrix_backend_and_atom_profile_internal(
module_name,
matrix_backend,
atom_profile,
true,
)
}
fn codegen_module_with_breakdown_for_matrix_backend_and_atom_profile_internal(
&self,
module_name: &str,
matrix_backend: MatrixBackend,
atom_profile: AtomOptimizationProfile,
probe_source: bool,
) -> (CodegenModule, BvpGeneratedAotCrateBreakdown) {
let prepared = self.as_prepared_problem_for_matrix_backend(matrix_backend);
if self.symbolic_assembly_backend == BvpSymbolicAssemblyBackend::AtomView {
if matrix_backend != MatrixBackend::Banded {
if let Some(discretized) = self.atom_discretized_system.as_ref() {
let jacobian_begin = Instant::now();
let (atom_codegen, atom_breakdown) =
prepare_sparse_bvp_codegen_from_discretized_system_with_breakdown(
discretized,
&self.residual_fn_name,
&self.jacobian_fn_name,
self.param_names.clone(),
self.bandwidth,
self.residual_strategy,
self.jacobian_strategy,
);
let jacobian_prepare_ms = jacobian_begin.elapsed().as_secs_f64() * 1_000.0;
let module_begin = Instant::now();
let (module, atom_module_breakdown) = atom_codegen
.codegen_module_with_breakdown_and_optimization_profile(
module_name,
atom_profile,
);
let module_build_ms = module_begin.elapsed().as_secs_f64() * 1_000.0;
let (source_probe_emit_ms, source_emit_ms, source_kb) = if probe_source {
let emit_begin = Instant::now();
let source_kb = module.emit_source().len() as f64 / 1024.0;
let source_emit_ms = emit_begin.elapsed().as_secs_f64() * 1_000.0;
(source_emit_ms, source_emit_ms, source_kb)
} else {
(0.0, 0.0, 0.0)
};
let breakdown = BvpGeneratedAotCrateBreakdown {
jacobian_prepare_ms,
atom_sparse_lookup_prepare_ms: atom_breakdown.sparse_lookup_prepare_ms,
atom_sparse_jacobian_build_ms: atom_breakdown.sparse_jacobian_build_ms,
atom_finalize_codegen_plan_ms: atom_breakdown.finalize_codegen_plan_ms,
atom_sparse_nnz: atom_breakdown.sparse_nnz,
atom_residual_view_collect_ms: atom_module_breakdown
.residual_view_collect_ms,
atom_residual_lower_many_ms: atom_module_breakdown.residual_lower_many_ms,
atom_residual_peephole_ms: atom_module_breakdown.residual_peephole_ms,
atom_residual_reuse_temps_ms: atom_module_breakdown.residual_reuse_temps_ms,
atom_residual_push_ms: atom_module_breakdown.residual_push_ms,
atom_sparse_view_collect_ms: atom_module_breakdown.sparse_view_collect_ms,
atom_sparse_lower_many_ms: atom_module_breakdown.sparse_lower_many_ms,
atom_sparse_peephole_ms: atom_module_breakdown.sparse_peephole_ms,
atom_sparse_reuse_temps_ms: atom_module_breakdown.sparse_reuse_temps_ms,
atom_sparse_push_ms: atom_module_breakdown.sparse_push_ms,
prepared_module_init_ms: 0.0,
prepared_residual_blocks_ms: 0.0,
prepared_jacobian_blocks_ms: 0.0,
module_build_ms,
source_probe_emit_ms,
language_source_emit_ms: 0.0,
c_header_emit_ms: 0.0,
artifact_packaging_ms: 0.0,
source_emit_ms,
source_kb,
module_blocks: module.block_count(),
total_block_instructions: module.total_block_instruction_count(),
total_block_temps: module.total_block_temp_count(),
max_block_instructions: module.max_block_instruction_count(),
total_block_outputs: module.total_block_output_count(),
};
return (module, breakdown);
}
}
}
let (module, module_breakdown) =
crate::symbolic::codegen::codegen_aot_driver::codegen_module_from_prepared_problem_with_breakdown(
module_name,
&prepared,
);
let module_build_ms = module_breakdown.total_ms;
let (source_probe_emit_ms, source_emit_ms, source_kb) = if probe_source {
let emit_begin = Instant::now();
let source_kb = module.emit_source().len() as f64 / 1024.0;
let source_emit_ms = emit_begin.elapsed().as_secs_f64() * 1_000.0;
(source_emit_ms, source_emit_ms, source_kb)
} else {
(0.0, 0.0, 0.0)
};
let breakdown = BvpGeneratedAotCrateBreakdown {
jacobian_prepare_ms: 0.0,
atom_sparse_lookup_prepare_ms: 0.0,
atom_sparse_jacobian_build_ms: 0.0,
atom_finalize_codegen_plan_ms: 0.0,
atom_sparse_nnz: 0,
atom_residual_view_collect_ms: 0.0,
atom_residual_lower_many_ms: 0.0,
atom_residual_peephole_ms: 0.0,
atom_residual_reuse_temps_ms: 0.0,
atom_residual_push_ms: 0.0,
atom_sparse_view_collect_ms: 0.0,
atom_sparse_lower_many_ms: 0.0,
atom_sparse_peephole_ms: 0.0,
atom_sparse_reuse_temps_ms: 0.0,
atom_sparse_push_ms: 0.0,
prepared_module_init_ms: module_breakdown.module_init_ms,
prepared_residual_blocks_ms: module_breakdown.residual_blocks_ms,
prepared_jacobian_blocks_ms: module_breakdown.jacobian_blocks_ms,
module_build_ms,
source_probe_emit_ms,
language_source_emit_ms: 0.0,
c_header_emit_ms: 0.0,
artifact_packaging_ms: 0.0,
source_emit_ms,
source_kb,
module_blocks: module.block_count(),
total_block_instructions: module.total_block_instruction_count(),
total_block_temps: module.total_block_temp_count(),
max_block_instructions: module.max_block_instruction_count(),
total_block_outputs: module.total_block_output_count(),
};
(module, breakdown)
}
pub fn generated_aot_artifact_with_breakdown(
&self,
artifact_name: impl Into<String>,
module_name: &str,
backend: AotCodegenBackend,
) -> (GeneratedAotArtifact, BvpGeneratedAotCrateBreakdown) {
self.generated_aot_artifact_with_breakdown_for_matrix_backend(
artifact_name,
module_name,
backend,
MatrixBackend::SparseCol,
)
}
pub fn generated_aot_artifact_with_breakdown_for_matrix_backend(
&self,
artifact_name: impl Into<String>,
module_name: &str,
backend: AotCodegenBackend,
matrix_backend: MatrixBackend,
) -> (GeneratedAotArtifact, BvpGeneratedAotCrateBreakdown) {
self.generated_aot_artifact_with_breakdown_for_matrix_backend_and_atom_profile(
artifact_name,
module_name,
backend,
matrix_backend,
AtomOptimizationProfile::Full,
)
}
pub fn generated_aot_artifact_with_breakdown_for_matrix_backend_and_atom_profile(
&self,
artifact_name: impl Into<String>,
module_name: &str,
backend: AotCodegenBackend,
matrix_backend: MatrixBackend,
atom_profile: AtomOptimizationProfile,
) -> (GeneratedAotArtifact, BvpGeneratedAotCrateBreakdown) {
let artifact_name = artifact_name.into();
let prepared = self.as_prepared_problem_for_matrix_backend(matrix_backend);
let (mut module, mut breakdown) = self
.codegen_module_with_breakdown_for_matrix_backend_and_atom_profile_internal(
module_name,
matrix_backend,
atom_profile,
false,
);
module.set_language(match backend {
AotCodegenBackend::Rust => crate::symbolic::codegen::CodegenIR::CodegenLanguage::Rust,
AotCodegenBackend::C => crate::symbolic::codegen::CodegenIR::CodegenLanguage::C,
AotCodegenBackend::Zig => crate::symbolic::codegen::CodegenIR::CodegenLanguage::Zig,
});
let emit_begin = Instant::now();
let module_source = module.emit_source();
let language_source_emit_ms = emit_begin.elapsed().as_secs_f64() * 1_000.0;
breakdown.language_source_emit_ms = language_source_emit_ms;
breakdown.source_emit_ms = language_source_emit_ms;
breakdown.source_kb = module_source.len() as f64 / 1024.0;
let manifest_begin = Instant::now();
let manifest = PreparedProblemManifest::from(&prepared);
let manifest_ms = manifest_begin.elapsed().as_secs_f64() * 1_000.0;
let package_begin = Instant::now();
let artifact = match backend {
AotCodegenBackend::Rust => GeneratedAotArtifact::Rust(GeneratedAotCrate::new(
artifact_name,
module_name,
module_source,
manifest,
)),
AotCodegenBackend::C => {
let header_begin = Instant::now();
let c_header = module.emit_c_header();
breakdown.c_header_emit_ms = header_begin.elapsed().as_secs_f64() * 1_000.0;
GeneratedAotArtifact::C(
crate::symbolic::codegen::c_backend::codegen_c_aot_library::GeneratedCAotLibrary::new(
artifact_name,
module_source,
c_header,
manifest,
),
)
}
AotCodegenBackend::Zig => GeneratedAotArtifact::Zig(
crate::symbolic::codegen::zig_backend::codegen_zig_aot_library::GeneratedZigAotLibrary::new(
artifact_name,
module_source,
manifest,
),
),
};
breakdown.artifact_packaging_ms =
manifest_ms + package_begin.elapsed().as_secs_f64() * 1_000.0;
(artifact, breakdown)
}
pub fn generated_aot_crate(
&self,
crate_name: impl Into<String>,
module_name: &str,
) -> GeneratedAotCrate {
self.generated_aot_crate_with_breakdown(crate_name, module_name)
.0
}
pub fn generated_aot_crate_with_breakdown(
&self,
crate_name: impl Into<String>,
module_name: &str,
) -> (GeneratedAotCrate, BvpGeneratedAotCrateBreakdown) {
let (artifact, breakdown) = self.generated_aot_artifact_with_breakdown(
crate_name,
module_name,
AotCodegenBackend::Rust,
);
(
artifact
.into_rust_crate()
.expect("Rust backend must return GeneratedAotCrate"),
breakdown,
)
}
pub fn auto_parallel_plan(&self) -> SparseAutoParallelPlan {
recommended_sparse_auto_parallel_plan(
self.residuals.len(),
self.shape.0,
self.sparse_entries.len(),
)
}
pub fn auto_parallel_executor_config(&self) -> Option<ParallelExecutorConfig> {
self.auto_parallel_plan().executor_config
}
}
fn linked_runtime_callbacks_for_matrix_backend(
prepared_problem: &BvpPreparedSparseAotProblem,
matrix_backend: MatrixBackend,
parameter_values: Option<&[f64]>,
parallel_config: Option<ParallelExecutorConfig>,
banded_linear_solver_config: LinearSolverConfig,
) -> Option<(Box<dyn Fun>, Box<dyn Jac>)> {
let linked = resolve_linked_sparse_backend(
&prepared_problem.problem_key_for_matrix_backend(matrix_backend),
)?;
let parameter_values = parameter_values.unwrap_or(&[]).to_vec();
let shape = prepared_problem.shape;
let residual_len = prepared_problem.shape.0;
let residual_link = linked.clone();
let residual_parameter_values = parameter_values.clone();
let residual_parallel_config = parallel_config;
let residual_fun = convert_to_fun(Box::new(move |_x: f64, vec: &dyn VectorType| {
let input_start = std::time::Instant::now();
let dense = vec.to_DVectorType();
let flat_args =
flatten_runtime_args(Some(residual_parameter_values.as_slice()), dense.as_slice());
record_callback_stage_time("Callback Residual Input Prep", input_start.elapsed());
let mut out = vec![0.0; residual_len];
let values_start = std::time::Instant::now();
eval_linked_residual_outputs(
&residual_link,
residual_parallel_config,
&flat_args,
&mut out,
);
record_callback_stage_time("Callback Residual Values", values_start.elapsed());
let boxing_start = std::time::Instant::now();
let result = Box::new(ColRef::from_slice(out.as_slice()).to_owned()) as Box<dyn VectorType>;
record_callback_stage_time("Callback Residual Boxing", boxing_start.elapsed());
result
}));
let jacobian_link = linked;
let jacobian_parameter_values = parameter_values.clone();
let jacobian_parallel_config = parallel_config;
let jacobian: Box<dyn Jac> = match matrix_backend {
MatrixBackend::Banded => {
let structure = prepared_problem
.as_prepared_banded_problem()
.jacobian_structure()
.clone();
let layout = prepared_problem
.infer_node_major_layout()
.expect("linked banded backend requires node-major compatible layout");
convert_to_jac(Box::new(move |_x: f64, vec: &dyn VectorType| {
let input_start = std::time::Instant::now();
let dense = vec.to_DVectorType();
let flat_args = flatten_runtime_args(
Some(jacobian_parameter_values.as_slice()),
dense.as_slice(),
);
record_callback_stage_time("Callback Jacobian Input Prep", input_start.elapsed());
let mut values = vec![0.0; structure.nnz()];
let values_start = std::time::Instant::now();
eval_linked_sparse_values(
&jacobian_link,
jacobian_parallel_config,
&flat_args,
&mut values,
);
record_callback_stage_time("Callback Jacobian Values", values_start.elapsed());
let assembly_start = std::time::Instant::now();
let assembly = structure.assemble_banded_assembly(values.as_slice());
let matrix =
BandedMatrixType::new(assembly, layout.clone(), banded_linear_solver_config);
let boxed = Box::new(matrix) as Box<dyn MatrixType>;
record_callback_stage_time(
"Callback Jacobian Matrix Assembly",
assembly_start.elapsed(),
);
boxed
}))
}
_ => {
let structure = prepared_problem
.as_prepared_problem()
.jacobian_structure()
.clone();
convert_to_jac(Box::new(move |_x: f64, vec: &dyn VectorType| {
let input_start = std::time::Instant::now();
let dense = vec.to_DVectorType();
let flat_args = flatten_runtime_args(
Some(jacobian_parameter_values.as_slice()),
dense.as_slice(),
);
record_callback_stage_time("Callback Jacobian Input Prep", input_start.elapsed());
let mut values = vec![0.0; structure.nnz()];
let values_start = std::time::Instant::now();
eval_linked_sparse_values(
&jacobian_link,
jacobian_parallel_config,
&flat_args,
&mut values,
);
record_callback_stage_time("Callback Jacobian Values", values_start.elapsed());
let assembly_start = std::time::Instant::now();
let triplets = structure
.row_indices
.iter()
.zip(structure.col_indices.iter())
.zip(values.iter())
.map(|((&row, &col), &value)| Triplet::new(row, col, value))
.collect::<Vec<_>>();
let matrix =
SparseColMat::try_new_from_triplets(shape.0, shape.1, triplets.as_slice())
.expect("linked sparse backend must return valid sparse values");
let boxed = Box::new(matrix) as Box<dyn MatrixType>;
record_callback_stage_time(
"Callback Jacobian Matrix Assembly",
assembly_start.elapsed(),
);
boxed
}))
}
};
Some((residual_fun, jacobian))
}
fn linked_contiguous_job_ranges(len: usize, max_jobs: usize) -> Vec<std::ops::Range<usize>> {
if len == 0 {
return Vec::new();
}
let target_jobs = len.min(max_jobs.max(1));
let items_per_job = len.div_ceil(target_jobs);
let mut ranges = Vec::with_capacity(target_jobs);
let mut start = 0usize;
while start < len {
let end = (start + items_per_job).min(len);
ranges.push(start..end);
start = end;
}
ranges
}
fn linked_should_parallelize(
config: ParallelExecutorConfig,
workload: usize,
chunk_count: usize,
sparse: bool,
) -> bool {
if chunk_count <= 1 {
return false;
}
match config.fallback_policy {
ParallelFallbackPolicy::Auto => {
let recommendation = if sparse {
auto_parallel_recommendation(0, 0, workload, chunk_count)
} else {
auto_parallel_recommendation(workload, chunk_count, 0, 0)
};
recommendation.should_parallelize
}
ParallelFallbackPolicy::Never => true,
ParallelFallbackPolicy::Thresholds {
min_residual_outputs,
min_sparse_values,
} => {
if sparse {
workload >= min_sparse_values
} else {
workload >= min_residual_outputs
}
}
}
}
fn linked_max_jobs(config: ParallelExecutorConfig, sparse: bool) -> usize {
let workers = rayon::current_num_threads().max(1);
let worker_jobs = workers.saturating_mul(config.jobs_per_worker.max(1));
if sparse {
config.max_sparse_jobs.unwrap_or(worker_jobs).max(1)
} else {
config.max_residual_jobs.unwrap_or(worker_jobs).max(1)
}
}
fn eval_linked_residual_outputs(
linked: &crate::symbolic::codegen::codegen_aot_runtime_link::LinkedSparseAotBackend,
parallel_config: Option<ParallelExecutorConfig>,
flat_args: &[f64],
out: &mut [f64],
) {
let Some(config) = parallel_config else {
(linked.residual_eval)(flat_args, out);
return;
};
if linked.residual_chunks.is_empty()
|| !linked_should_parallelize(
config,
linked.residual_len,
linked.residual_chunks.len(),
false,
)
{
(linked.residual_eval)(flat_args, out);
return;
}
let job_ranges =
linked_contiguous_job_ranges(linked.residual_chunks.len(), linked_max_jobs(config, false));
if job_ranges.len() <= 1 {
(linked.residual_eval)(flat_args, out);
return;
}
eval_linked_residual_job_range_recursive(
linked.residual_chunks.as_slice(),
job_ranges.as_slice(),
flat_args,
0,
job_ranges.len(),
out,
0,
);
}
fn eval_linked_residual_job_range_recursive(
chunks: &[crate::symbolic::codegen::codegen_aot_runtime_link::LinkedResidualChunk],
job_ranges: &[std::ops::Range<usize>],
flat_args: &[f64],
job_start: usize,
job_end: usize,
out: &mut [f64],
base_offset: usize,
) {
match job_end - job_start {
0 => {}
1 => {
for chunk in &chunks[job_ranges[job_start].clone()] {
let local_start = chunk.output_offset - base_offset;
let local_end = local_start + chunk.output_len;
(chunk.eval)(flat_args, &mut out[local_start..local_end]);
}
}
_ => {
let mid = job_start + (job_end - job_start) / 2;
let right_start = chunks[job_ranges[mid].start].output_offset;
let split_at = right_start - base_offset;
let (left_out, right_out) = out.split_at_mut(split_at);
rayon::join(
|| {
eval_linked_residual_job_range_recursive(
chunks,
job_ranges,
flat_args,
job_start,
mid,
left_out,
base_offset,
)
},
|| {
eval_linked_residual_job_range_recursive(
chunks,
job_ranges,
flat_args,
mid,
job_end,
right_out,
right_start,
)
},
);
}
}
}
fn eval_linked_sparse_values(
linked: &crate::symbolic::codegen::codegen_aot_runtime_link::LinkedSparseAotBackend,
parallel_config: Option<ParallelExecutorConfig>,
flat_args: &[f64],
values_out: &mut [f64],
) {
let Some(config) = parallel_config else {
(linked.jacobian_values_eval)(flat_args, values_out);
return;
};
if linked.jacobian_value_chunks.is_empty()
|| !linked_should_parallelize(config, linked.nnz, linked.jacobian_value_chunks.len(), true)
{
(linked.jacobian_values_eval)(flat_args, values_out);
return;
}
let job_ranges = linked_contiguous_job_ranges(
linked.jacobian_value_chunks.len(),
linked_max_jobs(config, true),
);
if job_ranges.len() <= 1 {
(linked.jacobian_values_eval)(flat_args, values_out);
return;
}
eval_linked_sparse_job_range_recursive(
linked.jacobian_value_chunks.as_slice(),
job_ranges.as_slice(),
flat_args,
0,
job_ranges.len(),
values_out,
0,
);
}
fn eval_linked_sparse_job_range_recursive(
chunks: &[crate::symbolic::codegen::codegen_aot_runtime_link::LinkedSparseJacobianChunk],
job_ranges: &[std::ops::Range<usize>],
flat_args: &[f64],
job_start: usize,
job_end: usize,
values_out: &mut [f64],
base_offset: usize,
) {
match job_end - job_start {
0 => {}
1 => {
for chunk in &chunks[job_ranges[job_start].clone()] {
let local_start = chunk.value_offset - base_offset;
let local_end = local_start + chunk.value_len;
(chunk.eval)(flat_args, &mut values_out[local_start..local_end]);
}
}
_ => {
let mid = job_start + (job_end - job_start) / 2;
let right_start = chunks[job_ranges[mid].start].value_offset;
let split_at = right_start - base_offset;
let (left_out, right_out) = values_out.split_at_mut(split_at);
rayon::join(
|| {
eval_linked_sparse_job_range_recursive(
chunks,
job_ranges,
flat_args,
job_start,
mid,
left_out,
base_offset,
)
},
|| {
eval_linked_sparse_job_range_recursive(
chunks,
job_ranges,
flat_args,
mid,
job_end,
right_out,
right_start,
)
},
);
}
}
}
fn chunk_sparse_entries_by_nnz<'a>(
entries: &[SparseExprEntry<'a>],
max_entries_per_chunk: usize,
) -> Vec<(usize, Vec<SparseExprEntry<'a>>)> {
assert!(
max_entries_per_chunk > 0,
"max_entries_per_chunk must be positive"
);
entries
.chunks(max_entries_per_chunk)
.enumerate()
.map(|(chunk_index, entries)| (chunk_index * max_entries_per_chunk, entries.to_vec()))
.collect()
}
fn chunk_sparse_entries_by_rows<'a>(
entries: &[SparseExprEntry<'a>],
rows_per_chunk: usize,
total_rows: usize,
) -> Vec<(usize, Vec<SparseExprEntry<'a>>)> {
let mut chunk_entries: Vec<Vec<(usize, SparseExprEntry<'a>)>> = Vec::new();
let mut chunk_offsets: Vec<usize> = Vec::new();
for (entry_index, entry) in entries.iter().copied().enumerate() {
let bucket = entry.row / rows_per_chunk;
if bucket >= chunk_entries.len() {
chunk_entries.resize_with(bucket + 1, Vec::new);
chunk_offsets.resize(bucket + 1, 0);
}
if chunk_entries[bucket].is_empty() {
chunk_offsets[bucket] = entry_index;
}
chunk_entries[bucket].push((entry_index, entry));
}
let non_empty_chunks = chunk_entries
.iter()
.filter(|entries| !entries.is_empty())
.collect::<Vec<_>>();
let all_chunks_are_contiguous = non_empty_chunks.iter().all(|entries| {
let start = entries.first().map(|(index, _)| *index).unwrap_or(0);
entries
.iter()
.enumerate()
.all(|(local_index, (entry_index, _))| *entry_index == start + local_index)
});
if !all_chunks_are_contiguous {
let target_chunks = total_rows.max(1).div_ceil(rows_per_chunk).max(1);
let max_entries_per_chunk = entries.len().max(1).div_ceil(target_chunks).max(1);
return chunk_sparse_entries_by_nnz(entries, max_entries_per_chunk);
}
chunk_entries
.into_iter()
.enumerate()
.filter(|(_, entries)| !entries.is_empty())
.map(|(chunk_index, entries)| {
(
chunk_offsets[chunk_index],
entries.into_iter().map(|(_, entry)| entry).collect(),
)
})
.collect()
}
#[cfg(test)]
mod bvp_sparse_chunking_tests {
use super::*;
use crate::symbolic::codegen::codegen_aot_runtime_link::{
LinkedResidualChunk, LinkedSparseAotBackend, LinkedSparseJacobianChunk,
};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
#[test]
fn bvp_sparse_target_chunking_uses_contiguous_value_segments() {
let e0 = Expr::parse_expression("y0");
let e1 = Expr::parse_expression("y1");
let e2 = Expr::parse_expression("y2");
let e3 = Expr::parse_expression("y3");
let entries = vec![
SparseExprEntry {
row: 0,
col: 0,
expr: &e0,
},
SparseExprEntry {
row: 2,
col: 0,
expr: &e1,
},
SparseExprEntry {
row: 1,
col: 1,
expr: &e2,
},
SparseExprEntry {
row: 3,
col: 1,
expr: &e3,
},
];
let chunks = chunk_sparse_entries_by_nnz(&entries, 2);
assert_eq!(chunks.len(), 2);
assert_eq!(chunks[0].0, 0);
assert_eq!(chunks[0].1[0].row, 0);
assert_eq!(chunks[0].1[1].row, 2);
assert_eq!(chunks[1].0, 2);
assert_eq!(chunks[1].1[0].row, 1);
assert_eq!(chunks[1].1[1].row, 3);
}
#[test]
fn bvp_sparse_row_chunking_falls_back_when_value_segments_interleave() {
let e0 = Expr::parse_expression("y0");
let e1 = Expr::parse_expression("y1");
let e2 = Expr::parse_expression("y2");
let e3 = Expr::parse_expression("y3");
let entries = vec![
SparseExprEntry {
row: 0,
col: 0,
expr: &e0,
},
SparseExprEntry {
row: 2,
col: 0,
expr: &e1,
},
SparseExprEntry {
row: 1,
col: 1,
expr: &e2,
},
SparseExprEntry {
row: 3,
col: 1,
expr: &e3,
},
];
let chunks = chunk_sparse_entries_by_rows(&entries, 2, 4);
assert_eq!(chunks.len(), 2);
assert_eq!(chunks[0].0, 0);
assert_eq!(chunks[0].1[0].row, 0);
assert_eq!(chunks[0].1[1].row, 2);
assert_eq!(chunks[1].0, 2);
assert_eq!(chunks[1].1[0].row, 1);
assert_eq!(chunks[1].1[1].row, 3);
}
#[test]
fn linked_sparse_parallel_callbacks_write_directly_into_final_buffers() {
fn slice_inside_final_buffer(slice: &[f64], base: usize, len: usize) -> bool {
let ptr = slice.as_ptr() as usize;
let end = base + len * std::mem::size_of::<f64>();
ptr >= base && ptr < end
}
let residual_base = Arc::new(AtomicUsize::new(0));
let sparse_base = Arc::new(AtomicUsize::new(0));
let residual_direct = Arc::new(AtomicBool::new(true));
let sparse_direct = Arc::new(AtomicBool::new(true));
let residual_chunk =
|offset: usize, len: usize, base: Arc<AtomicUsize>, direct: Arc<AtomicBool>| {
LinkedResidualChunk::new(
offset,
len,
Arc::new(move |args: &[f64], out: &mut [f64]| {
direct.fetch_and(
slice_inside_final_buffer(out, base.load(Ordering::SeqCst), 4),
Ordering::SeqCst,
);
for (i, value) in out.iter_mut().enumerate() {
*value = args[0] + (offset + i) as f64;
}
}),
)
};
let sparse_chunk =
|offset: usize, len: usize, base: Arc<AtomicUsize>, direct: Arc<AtomicBool>| {
LinkedSparseJacobianChunk::new(
offset,
len,
Arc::new(move |args: &[f64], out: &mut [f64]| {
direct.fetch_and(
slice_inside_final_buffer(out, base.load(Ordering::SeqCst), 4),
Ordering::SeqCst,
);
for (i, value) in out.iter_mut().enumerate() {
*value = args[0] * 10.0 + (offset + i) as f64;
}
}),
)
};
let linked = LinkedSparseAotBackend::new(
"direct-linked-test",
4,
(4, 4),
4,
Arc::new(|_, out: &mut [f64]| out.fill(-1.0)),
Arc::new(|_, out: &mut [f64]| out.fill(-2.0)),
)
.with_chunked_evaluators(
vec![
residual_chunk(
0,
2,
Arc::clone(&residual_base),
Arc::clone(&residual_direct),
),
residual_chunk(
2,
2,
Arc::clone(&residual_base),
Arc::clone(&residual_direct),
),
],
vec![
sparse_chunk(0, 2, Arc::clone(&sparse_base), Arc::clone(&sparse_direct)),
sparse_chunk(2, 2, Arc::clone(&sparse_base), Arc::clone(&sparse_direct)),
],
);
let config = ParallelExecutorConfig {
jobs_per_worker: 1,
max_residual_jobs: Some(2),
max_sparse_jobs: Some(2),
fallback_policy: ParallelFallbackPolicy::Never,
};
let mut residual = vec![0.0; 4];
residual_base.store(residual.as_mut_ptr() as usize, Ordering::SeqCst);
eval_linked_residual_outputs(&linked, Some(config), &[3.0], &mut residual);
assert_eq!(residual, vec![3.0, 4.0, 5.0, 6.0]);
assert!(
residual_direct.load(Ordering::SeqCst),
"linked residual chunk callbacks must receive final output slices, not temporary buffers"
);
let mut sparse_values = vec![0.0; 4];
sparse_base.store(sparse_values.as_mut_ptr() as usize, Ordering::SeqCst);
eval_linked_sparse_values(&linked, Some(config), &[3.0], &mut sparse_values);
assert_eq!(sparse_values, vec![30.0, 31.0, 32.0, 33.0]);
assert!(
sparse_direct.load(Ordering::SeqCst),
"linked sparse chunk callbacks must receive final value slices, not temporary buffers"
);
}
#[test]
fn linked_sparse_parallel_callbacks_actually_overlap_on_rayon_workers() {
fn bump_max(max_active: &AtomicUsize, current: usize) {
let mut observed = max_active.load(Ordering::SeqCst);
while current > observed {
match max_active.compare_exchange(
observed,
current,
Ordering::SeqCst,
Ordering::SeqCst,
) {
Ok(_) => break,
Err(next_observed) => observed = next_observed,
}
}
}
let active_residual = Arc::new(AtomicUsize::new(0));
let max_residual = Arc::new(AtomicUsize::new(0));
let active_sparse = Arc::new(AtomicUsize::new(0));
let max_sparse = Arc::new(AtomicUsize::new(0));
let residual_chunk =
|offset: usize, active: Arc<AtomicUsize>, max_active: Arc<AtomicUsize>| {
LinkedResidualChunk::new(
offset,
1,
Arc::new(move |args: &[f64], out: &mut [f64]| {
let current = active.fetch_add(1, Ordering::SeqCst) + 1;
bump_max(&max_active, current);
std::thread::sleep(Duration::from_millis(25));
out[0] = args[0] + offset as f64;
active.fetch_sub(1, Ordering::SeqCst);
}),
)
};
let sparse_chunk =
|offset: usize, active: Arc<AtomicUsize>, max_active: Arc<AtomicUsize>| {
LinkedSparseJacobianChunk::new(
offset,
1,
Arc::new(move |args: &[f64], out: &mut [f64]| {
let current = active.fetch_add(1, Ordering::SeqCst) + 1;
bump_max(&max_active, current);
std::thread::sleep(Duration::from_millis(25));
out[0] = args[0] * 10.0 + offset as f64;
active.fetch_sub(1, Ordering::SeqCst);
}),
)
};
let linked = LinkedSparseAotBackend::new(
"parallel-overlap-test",
4,
(4, 4),
4,
Arc::new(|_, out: &mut [f64]| out.fill(-1.0)),
Arc::new(|_, out: &mut [f64]| out.fill(-2.0)),
)
.with_chunked_evaluators(
(0..4)
.map(|offset| {
residual_chunk(
offset,
Arc::clone(&active_residual),
Arc::clone(&max_residual),
)
})
.collect(),
(0..4)
.map(|offset| {
sparse_chunk(offset, Arc::clone(&active_sparse), Arc::clone(&max_sparse))
})
.collect(),
);
let config = ParallelExecutorConfig {
jobs_per_worker: 1,
max_residual_jobs: Some(4),
max_sparse_jobs: Some(4),
fallback_policy: ParallelFallbackPolicy::Never,
};
rayon::ThreadPoolBuilder::new()
.num_threads(4)
.build()
.expect("local rayon pool should build")
.install(|| {
let mut residual = vec![0.0; 4];
eval_linked_residual_outputs(&linked, Some(config), &[2.0], &mut residual);
assert_eq!(residual, vec![2.0, 3.0, 4.0, 5.0]);
let mut sparse_values = vec![0.0; 4];
eval_linked_sparse_values(&linked, Some(config), &[2.0], &mut sparse_values);
assert_eq!(sparse_values, vec![20.0, 21.0, 22.0, 23.0]);
});
assert!(
max_residual.load(Ordering::SeqCst) > 1,
"linked residual chunk callbacks must overlap; max_active={}",
max_residual.load(Ordering::SeqCst)
);
assert!(
max_sparse.load(Ordering::SeqCst) > 1,
"linked sparse chunk callbacks must overlap; max_active={}",
max_sparse.load(Ordering::SeqCst)
);
}
}
impl Jacobian {
fn prepare_symbolic_bvp_stage_with_params(
&mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
params: Option<&[&str]>,
t0: f64,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
BorderConditions: HashMap<String, Vec<(usize, f64)>>,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
method: String,
bandwidth: Option<(usize, usize)>,
) -> (
String,
Vec<String>,
BvpMatrixBackend,
HashMap<String, f64>,
Instant,
) {
let total_start = Instant::now();
let mut timer_hash: HashMap<String, f64> = HashMap::new();
self.set_params(params);
self.atom_discretized_system = None;
let param = arg.clone();
let matrix_backend = BvpMatrixBackend::from_legacy_method(method.as_str())
.unwrap_or_else(|| panic!("unknown method: {}", method));
self.set_backend_config(BvpBackendConfig::lambdify(matrix_backend));
let begin = Instant::now();
match self.symbolic_assembly_backend {
BvpSymbolicAssemblyBackend::ExprLegacy => self.discretization_system_BVP_par(
eq_system,
values,
arg,
t0,
n_steps,
h,
mesh,
BorderConditions,
Bounds,
rel_tolerance,
scheme,
),
BvpSymbolicAssemblyBackend::AtomView => {
let discretized = discretization_system_bvp_par_atom(
eq_system,
values,
arg,
t0,
n_steps,
h,
mesh,
BorderConditions,
Bounds,
rel_tolerance,
scheme,
);
self.install_atom_discretized_system(discretized);
}
}
timer_hash.insert(
"discretization time".to_string(),
begin.elapsed().as_secs_f64(),
);
info!("system discretized");
let indexed_values = self.variable_string.clone();
let atom_native_sparse_jacobian = self.symbolic_assembly_backend
== BvpSymbolicAssemblyBackend::AtomView
&& matches!(
matrix_backend,
BvpMatrixBackend::Banded | BvpMatrixBackend::FaerSparseCol
);
if let Some(bandwidth_) = bandwidth {
let now = Instant::now();
self.bandwidth = Some(bandwidth_);
info!("Bandwidth provided: {:?}", self.bandwidth);
timer_hash.insert(
"find bandwidth time".to_string(),
now.elapsed().as_secs_f64(),
);
info!("Calculating jacobian");
let now = Instant::now();
if atom_native_sparse_jacobian {
self.calc_atomview_sparse_jacobian_with_bandwidth(Some(bandwidth_));
} else {
self.calc_jacobian_parallel_smart_optimized_with_given_bandwidth();
}
let elapsed = now.elapsed();
info!("Jacobian calculation time:");
info!("{:?}", elapsed_time(elapsed));
timer_hash.insert("symbolic jacobian time".to_string(), elapsed.as_secs_f64());
} else {
info!("Calculating jacobian");
let now = Instant::now();
if atom_native_sparse_jacobian {
self.calc_atomview_sparse_jacobian_with_bandwidth(None);
} else {
self.calc_jacobian_parallel_smart_optimized();
}
let elapsed = now.elapsed();
info!("Jacobian calculation time:");
info!("{:?}", elapsed_time(elapsed));
timer_hash.insert("symbolic jacobian time".to_string(), elapsed.as_secs_f64());
let now = Instant::now();
self.find_bandwidths();
info!("Bandwidth calculated:");
info!("(kl, ku) = {:?}", self.bandwidth);
timer_hash.insert(
"find bandwidth time".to_string(),
now.elapsed().as_secs_f64(),
);
}
if let Some(jacobian_breakdown) = self.last_symbolic_jacobian_timer_snapshot.take() {
timer_hash.extend(jacobian_breakdown);
}
(
param,
indexed_values,
matrix_backend,
timer_hash,
total_start,
)
}
fn finalize_generate_bvp_timer_table(
&mut self,
mut timer_hash: HashMap<String, f64>,
total_end: f64,
) {
let mut raw_snapshot = timer_hash.clone();
raw_snapshot.insert("total time, sec".to_string(), total_end);
self.last_generate_timer_snapshot = Some(raw_snapshot);
for key in [
"discretization time",
"symbolic jacobian time",
"jacobian lambdify time",
"residual functions lambdify time",
"find bandwidth time",
"symbolic jacobian variable sets time",
"symbolic jacobian row differentiation time",
"symbolic jacobian dense cache materialize time",
"symbolic jacobian sparse cache flatten time",
"sparse AOT preparation time",
"backend selection time",
"runtime binding time",
"lambdify jacobian callback compile time",
"lambdify residual callback compile time",
] {
if let Some(value) = timer_hash.get_mut(key) {
normalize_timer_percent(value, total_end);
}
}
timer_hash.insert("total time, sec".to_string(), total_end);
let mut table = Builder::from(timer_hash.clone()).build();
table.with(Style::modern_rounded());
println!("{}", table.to_string());
info!(
"\n \n ____________END OF GENERATE BVP ________________________________________________________________"
);
}
pub fn new() -> Self {
let fun0: Box<dyn Fn(f64, &DVector<f64>) -> DVector<f64>> =
Box::new(|_x, y: &DVector<f64>| y.clone());
let boxed_fun: Box<dyn Fun> = Box::new(FunEnum::Dense(fun0));
Self {
vector_of_functions: Vec::new(),
lambdified_functions: Vec::new(),
method: String::new(),
backend_config: BvpBackendConfig::default(),
banded_linear_solver_config: LinearSolverConfig::default(),
banded_lambdify_config: BandedLambdifyConfig::default(),
symbolic_assembly_backend: BvpSymbolicAssemblyBackend::ExprLegacy,
vector_of_variables: Vec::new(),
variable_string: Vec::new(),
parameters_string: Vec::new(),
parameter_values: None,
symbolic_jacobian: Vec::new(),
symbolic_jacobian_sparse: Vec::new(),
lambdified_jac_element: None,
jac_function: None,
residiual_function: boxed_fun,
bounds: None,
rel_tolerance_vec: None,
bandwidth: None,
variables_for_all_disrete: Vec::new(),
BC_pos_n_values: Vec::new(),
last_generate_timer_snapshot: None,
last_symbolic_jacobian_timer_snapshot: None,
last_lambdify_binding_timer_snapshot: None,
atom_discretized_system: None,
}
}
pub fn set_banded_linear_solver_config(&mut self, config: LinearSolverConfig) {
self.banded_linear_solver_config = config;
}
pub fn set_banded_lambdify_config(&mut self, config: BandedLambdifyConfig) {
self.banded_lambdify_config = config;
}
pub fn set_banded_jacobian_chunking(&mut self, chunking: BandedJacobianChunking) {
self.banded_lambdify_config.jacobian_chunking = chunking;
}
pub fn from_vectors(&mut self, vector_of_functions: Vec<Expr>, variable_string: Vec<String>) {
self.vector_of_functions = vector_of_functions;
self.variable_string = variable_string.clone();
self.vector_of_variables =
Expr::parse_vector_expression(variable_string.iter().map(|s| s.as_str()).collect());
println!(" {:?}", self.vector_of_functions);
println!(" {:?}", self.vector_of_variables);
}
pub fn set_params(&mut self, params: Option<&[&str]>) {
self.parameters_string = params
.map(|params| params.iter().map(|name| (*name).to_string()).collect())
.unwrap_or_default();
}
pub fn set_param_values(&mut self, values: Option<Vec<f64>>) {
if let Some(ref values) = values {
assert_eq!(
values.len(),
self.parameters_string.len(),
"parameter_values length must match parameters_string length"
);
}
self.parameter_values = values;
}
pub fn set_symbolic_assembly_backend(&mut self, backend: BvpSymbolicAssemblyBackend) {
self.symbolic_assembly_backend = backend;
}
pub fn symbolic_assembly_backend(&self) -> BvpSymbolicAssemblyBackend {
self.symbolic_assembly_backend
}
pub fn last_generate_timer_snapshot(&self) -> Option<&HashMap<String, f64>> {
self.last_generate_timer_snapshot.as_ref()
}
fn install_atom_discretized_system(
&mut self,
discretized: crate::symbolic::View::bvp::DiscretizedBvpAtomSystem,
) {
self.atom_discretized_system = Some(discretized.clone());
let timer_table = discretized.render_timer_table();
self.vector_of_functions = discretized
.vector_of_functions
.iter()
.map(atom_to_expr)
.collect();
self.vector_of_variables = discretized
.vector_of_variables
.iter()
.map(atom_to_expr)
.collect();
self.variable_string = discretized.variable_string;
self.variables_for_all_disrete = discretized.variables_for_all_discrete;
self.BC_pos_n_values = discretized.bc_pos_n_values;
self.bounds = discretized.bounds;
self.rel_tolerance_vec = discretized.rel_tolerance_vec;
info!("Atom/View BVP discretization timing:\n{}", timer_table);
}
fn calc_atomview_sparse_jacobian_with_bandwidth(&mut self, bandwidth: Option<(usize, usize)>) {
let discretized = self
.atom_discretized_system
.as_ref()
.expect("AtomView Jacobian preparation requires an installed atom system");
let mut timings = HashMap::new();
let variable_sets_begin = Instant::now();
let prepared = PreparedSparseAtomSystem::from_atoms(
&discretized.vector_of_functions,
&discretized.variable_string,
&discretized.variables_for_all_discrete,
);
timings.insert(
"symbolic jacobian variable sets time".to_string(),
variable_sets_begin.elapsed().as_secs_f64(),
);
let differentiation_begin = Instant::now();
let sparse_atom_entries = prepared.calc_sparse_jacobian_with_bandwidth(bandwidth);
timings.insert(
"symbolic jacobian row differentiation time".to_string(),
differentiation_begin.elapsed().as_secs_f64(),
);
let flatten_begin = Instant::now();
self.symbolic_jacobian_sparse = sparse_atom_entries
.into_par_iter()
.map(|entry| (entry.row, entry.col, atom_to_expr(&entry.value)))
.collect();
timings.insert(
"symbolic jacobian sparse cache flatten time".to_string(),
flatten_begin.elapsed().as_secs_f64(),
);
let dense_cache_begin = Instant::now();
self.symbolic_jacobian.clear();
timings.insert(
"symbolic jacobian dense cache materialize time".to_string(),
dense_cache_begin.elapsed().as_secs_f64(),
);
self.last_symbolic_jacobian_timer_snapshot = Some(timings);
}
pub fn symbolic_jacobian_sparse_entries(
&self,
) -> Vec<crate::symbolic::codegen::codegen_tasks::SparseExprEntry<'_>> {
self.symbolic_jacobian_sparse
.iter()
.map(
|(row, col, expr)| crate::symbolic::codegen::codegen_tasks::SparseExprEntry {
row: *row,
col: *col,
expr,
},
)
.collect()
}
pub fn symbolic_jacobian_sparse_entries_owned(&self) -> Vec<(usize, usize, Expr)> {
self.symbolic_jacobian_sparse.clone()
}
pub fn prepare_sparse_aot_problem(
&self,
residual_fn_name: &str,
jacobian_fn_name: &str,
residual_strategy: ResidualChunkingStrategy,
jacobian_strategy: SparseChunkingStrategy,
) -> BvpPreparedSparseAotProblem {
BvpPreparedSparseAotProblem {
residual_fn_name: residual_fn_name.to_string(),
jacobian_fn_name: jacobian_fn_name.to_string(),
variable_names: self.variable_string.clone(),
param_names: self.parameters_string.clone(),
residuals: self.vector_of_functions.clone(),
sparse_entries: self.symbolic_jacobian_sparse_entries_owned(),
shape: (
self.vector_of_functions.len(),
self.vector_of_variables.len(),
),
residual_strategy,
jacobian_strategy,
bandwidth: self.bandwidth,
symbolic_assembly_backend: self.symbolic_assembly_backend,
atom_discretized_system: self.atom_discretized_system.clone(),
}
}
pub fn select_sparse_backend(
&self,
residual_fn_name: &str,
jacobian_fn_name: &str,
residual_strategy: ResidualChunkingStrategy,
jacobian_strategy: SparseChunkingStrategy,
policy: BackendSelectionPolicy,
resolver: Option<&AotResolver>,
) -> BvpSelectedSparseBackend {
let prepared_problem = self.prepare_sparse_aot_problem(
residual_fn_name,
jacobian_fn_name,
residual_strategy,
jacobian_strategy,
);
let prepared_binding = prepared_problem.clone();
let requested_matrix_backend = match self.backend_config().matrix_backend {
BvpMatrixBackend::Dense => MatrixBackend::Dense,
BvpMatrixBackend::Banded => MatrixBackend::Banded,
BvpMatrixBackend::FaerSparseCol => MatrixBackend::SparseCol,
BvpMatrixBackend::SprsCsMat => MatrixBackend::CsMat,
};
let prepared =
prepared_binding.as_prepared_problem_for_matrix_backend(requested_matrix_backend);
let selected = select_backend(&prepared, policy, resolver);
BvpSelectedSparseBackend {
prepared_problem,
requested_backend: selected.requested_backend,
effective_backend: selected.effective_backend,
matrix_backend: selected.matrix_backend,
aot_resolution: selected.aot_resolution,
}
}
fn prepare_sparse_backend_execution_timed(
&mut self,
arg: &str,
variable_str: Vec<&str>,
residual_fn_name: &str,
jacobian_fn_name: &str,
residual_strategy: ResidualChunkingStrategy,
jacobian_strategy: SparseChunkingStrategy,
policy: BackendSelectionPolicy,
resolver: Option<&AotResolver>,
timer_hash: &mut HashMap<String, f64>,
) -> BvpSparseExecutionPlan {
let prepare_started = Instant::now();
let prepared_problem = self.prepare_sparse_aot_problem(
residual_fn_name,
jacobian_fn_name,
residual_strategy,
jacobian_strategy,
);
timer_hash.insert(
"sparse AOT preparation time".to_string(),
prepare_started.elapsed().as_secs_f64(),
);
let selection_started = Instant::now();
let prepared_binding = prepared_problem.clone();
let requested_matrix_backend = match self.backend_config().matrix_backend {
BvpMatrixBackend::Dense => MatrixBackend::Dense,
BvpMatrixBackend::Banded => MatrixBackend::Banded,
BvpMatrixBackend::FaerSparseCol => MatrixBackend::SparseCol,
BvpMatrixBackend::SprsCsMat => MatrixBackend::CsMat,
};
let prepared =
prepared_binding.as_prepared_problem_for_matrix_backend(requested_matrix_backend);
let selected = select_backend(&prepared, policy, resolver);
timer_hash.insert(
"backend selection time".to_string(),
selection_started.elapsed().as_secs_f64(),
);
let selected = BvpSelectedSparseBackend {
prepared_problem,
requested_backend: selected.requested_backend,
effective_backend: selected.effective_backend,
matrix_backend: selected.matrix_backend,
aot_resolution: selected.aot_resolution,
};
let binding_started = Instant::now();
let execution = match selected.effective_backend {
SelectedBackendKind::Numeric => {
let _ = (arg, variable_str);
panic_symbolic_numeric_backend_requested()
}
SelectedBackendKind::Lambdify => {
let config = BvpBackendConfig::new(
BvpBackendKind::Lambdify,
BvpMatrixBackend::from_provider_matrix_backend(selected.matrix_backend),
);
self.compile_lambdified_problem_with_config(arg, variable_str, config);
if let Some(binding_breakdown) = self.last_lambdify_binding_timer_snapshot.take() {
timer_hash.extend(binding_breakdown);
}
BvpSparseExecutionPlan::LambdifyReady(selected)
}
SelectedBackendKind::AotCompiled => {
self.set_backend_config(BvpBackendConfig::new(
BvpBackendKind::Aot,
BvpMatrixBackend::from_provider_matrix_backend(selected.matrix_backend),
));
BvpSparseExecutionPlan::AotCompiled(selected)
}
SelectedBackendKind::AotRegisteredButNotBuilt => {
self.set_backend_config(BvpBackendConfig::new(
BvpBackendKind::Aot,
BvpMatrixBackend::from_provider_matrix_backend(selected.matrix_backend),
));
BvpSparseExecutionPlan::AotRegisteredButNotBuilt(selected)
}
SelectedBackendKind::AotMissing => {
self.set_backend_config(BvpBackendConfig::new(
BvpBackendKind::Aot,
BvpMatrixBackend::from_provider_matrix_backend(selected.matrix_backend),
));
BvpSparseExecutionPlan::AotMissing(selected)
}
};
timer_hash.insert(
"runtime binding time".to_string(),
binding_started.elapsed().as_secs_f64(),
);
execution
}
pub fn prepare_sparse_backend_execution(
&mut self,
arg: &str,
variable_str: Vec<&str>,
residual_fn_name: &str,
jacobian_fn_name: &str,
residual_strategy: ResidualChunkingStrategy,
jacobian_strategy: SparseChunkingStrategy,
policy: BackendSelectionPolicy,
resolver: Option<&AotResolver>,
) -> BvpSparseExecutionPlan {
let selected = self.select_sparse_backend(
residual_fn_name,
jacobian_fn_name,
residual_strategy,
jacobian_strategy,
policy,
resolver,
);
match selected.effective_backend {
SelectedBackendKind::Numeric => {
let _ = (arg, variable_str);
panic_symbolic_numeric_backend_requested()
}
SelectedBackendKind::Lambdify => {
let config = BvpBackendConfig::new(
BvpBackendKind::Lambdify,
BvpMatrixBackend::from_provider_matrix_backend(selected.matrix_backend),
);
self.compile_lambdified_problem_with_config(arg, variable_str, config);
BvpSparseExecutionPlan::LambdifyReady(selected)
}
SelectedBackendKind::AotCompiled => {
self.set_backend_config(BvpBackendConfig::new(
BvpBackendKind::Aot,
BvpMatrixBackend::from_provider_matrix_backend(selected.matrix_backend),
));
BvpSparseExecutionPlan::AotCompiled(selected)
}
SelectedBackendKind::AotRegisteredButNotBuilt => {
self.set_backend_config(BvpBackendConfig::new(
BvpBackendKind::Aot,
BvpMatrixBackend::from_provider_matrix_backend(selected.matrix_backend),
));
BvpSparseExecutionPlan::AotRegisteredButNotBuilt(selected)
}
SelectedBackendKind::AotMissing => {
self.set_backend_config(BvpBackendConfig::new(
BvpBackendKind::Aot,
BvpMatrixBackend::from_provider_matrix_backend(selected.matrix_backend),
));
BvpSparseExecutionPlan::AotMissing(selected)
}
}
}
pub fn sparse_solver_provider<'a>(
&'a mut self,
execution: BvpSparseExecutionPlan,
) -> BvpSparseSolverProvider<'a> {
let structure = execution
.selected()
.prepared_problem
.as_prepared_problem()
.jacobian_structure()
.clone();
BvpSparseSolverProvider {
jacobian: self,
execution,
structure,
}
}
pub fn into_sparse_solver_bundle(
self,
execution: BvpSparseExecutionPlan,
) -> BvpSparseSolverBundle {
let mut runtime_diagnostics = HashMap::new();
if let Some(snapshot) = self.last_generate_timer_snapshot.as_ref() {
for (stage, value_seconds) in snapshot {
let stage_key = stage
.chars()
.map(|ch| if ch.is_ascii_alphanumeric() { ch } else { '_' })
.collect::<String>();
runtime_diagnostics.insert(
format!("generated.prepare.{stage_key}_ms"),
format!("{:.6}", value_seconds * 1_000.0),
);
}
}
let selected = execution.selected().clone();
let sparse_structure = selected
.prepared_problem
.as_prepared_problem()
.jacobian_structure()
.clone();
let linked_callbacks = if matches!(execution, BvpSparseExecutionPlan::AotCompiled(_)) {
linked_runtime_callbacks_for_matrix_backend(
&selected.prepared_problem,
selected.matrix_backend,
self.parameter_values.as_deref(),
None,
self.banded_linear_solver_config,
)
} else {
None
};
let callbacks_available = matches!(execution, BvpSparseExecutionPlan::LambdifyReady(_))
|| linked_callbacks.is_some();
let (linked_residual, linked_jacobian) =
linked_callbacks.map_or((None, None), |(fun, jac)| (Some(fun), Some(jac)));
BvpSparseSolverBundle {
execution,
residual_function: if callbacks_available {
if matches!(selected.effective_backend, SelectedBackendKind::Lambdify) {
Some(self.residiual_function)
} else {
linked_residual
}
} else {
None
},
jacobian_function: if callbacks_available {
if matches!(selected.effective_backend, SelectedBackendKind::Lambdify) {
self.jac_function
} else {
linked_jacobian
}
} else {
None
},
runtime_diagnostics,
banded_linear_solver_config: self.banded_linear_solver_config,
variable_string: self.variable_string,
parameter_values: self.parameter_values,
bounds_vec: self.bounds,
rel_tolerance_vec: self.rel_tolerance_vec,
bandwidth: self.bandwidth,
bc_position_and_value: self.BC_pos_n_values,
sparse_structure,
}
}
#[allow(clippy::too_many_arguments)]
pub fn try_generate_BVP_with_backend_selection(
&mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
params: Option<&[&str]>,
t0: f64,
param: Option<String>,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
BorderConditions: HashMap<String, Vec<(usize, f64)>>,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
method: String,
bandwidth: Option<(usize, usize)>,
policy: BackendSelectionPolicy,
resolver: Option<&AotResolver>,
) -> Result<BvpSparseExecutionPlan, BvpBackendIntegrationError> {
info!(
"BVP backend selection started: backend_policy={policy:?}, method={method}, params={}, steps={:?}",
params.map_or(0, |p| p.len()),
n_steps
);
let result = catch_unwind(AssertUnwindSafe(|| {
self.generate_BVP_with_backend_selection(
eq_system,
values,
arg,
params,
t0,
param,
n_steps,
h,
mesh,
BorderConditions,
Bounds,
rel_tolerance,
scheme,
method,
bandwidth,
policy,
resolver,
)
}));
match result {
Ok(execution) => {
info!(
"BVP backend selection finished: effective_backend={:?}",
execution.selected().effective_backend
);
Ok(execution)
}
Err(payload) => {
let message = if let Some(message) = payload.downcast_ref::<&str>() {
(*message).to_string()
} else if let Some(message) = payload.downcast_ref::<String>() {
message.clone()
} else {
"unknown panic payload".to_string()
};
error!("BVP backend selection failed: {message}");
Err(BvpBackendIntegrationError::PipelinePanicked(message))
}
}
}
#[allow(clippy::too_many_arguments)]
pub fn try_generate_sparse_solver_bundle_with_backend_selection(
mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
params: Option<&[&str]>,
t0: f64,
param: Option<String>,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
BorderConditions: HashMap<String, Vec<(usize, f64)>>,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
method: String,
bandwidth: Option<(usize, usize)>,
policy: BackendSelectionPolicy,
resolver: Option<&AotResolver>,
) -> Result<BvpSparseSolverBundle, BvpBackendIntegrationError> {
let execution = self.try_generate_BVP_with_backend_selection(
eq_system,
values,
arg,
params,
t0,
param,
n_steps,
h,
mesh,
BorderConditions,
Bounds,
rel_tolerance,
scheme,
method,
bandwidth,
policy,
resolver,
)?;
Ok(self.into_sparse_solver_bundle(execution))
}
#[allow(clippy::too_many_arguments)]
pub fn try_generate_BVP_with_backend_selection_and_chunking(
&mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
params: Option<&[&str]>,
t0: f64,
param: Option<String>,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
BorderConditions: HashMap<String, Vec<(usize, f64)>>,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
method: String,
bandwidth: Option<(usize, usize)>,
policy: BackendSelectionPolicy,
resolver: Option<&AotResolver>,
residual_strategy: ResidualChunkingStrategy,
jacobian_strategy: SparseChunkingStrategy,
) -> Result<BvpSparseExecutionPlan, BvpBackendIntegrationError> {
info!(
"BVP backend selection with explicit chunking started: backend_policy={policy:?}, method={method}, params={}, steps={:?}, residual_strategy={residual_strategy:?}, jacobian_strategy={jacobian_strategy:?}",
params.map_or(0, |p| p.len()),
n_steps
);
let result = catch_unwind(AssertUnwindSafe(|| {
self.generate_BVP_with_backend_selection_and_chunking(
eq_system,
values,
arg,
params,
t0,
param,
n_steps,
h,
mesh,
BorderConditions,
Bounds,
rel_tolerance,
scheme,
method,
bandwidth,
policy,
resolver,
residual_strategy,
jacobian_strategy,
)
}));
match result {
Ok(execution) => {
info!(
"BVP backend selection with explicit chunking finished: effective_backend={:?}",
execution.selected().effective_backend
);
Ok(execution)
}
Err(payload) => {
let message = if let Some(message) = payload.downcast_ref::<&str>() {
(*message).to_string()
} else if let Some(message) = payload.downcast_ref::<String>() {
message.clone()
} else {
"unknown panic payload".to_string()
};
error!("BVP backend selection with explicit chunking failed: {message}");
Err(BvpBackendIntegrationError::PipelinePanicked(message))
}
}
}
#[allow(clippy::too_many_arguments)]
pub fn try_generate_sparse_solver_bundle_with_backend_selection_and_chunking(
mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
params: Option<&[&str]>,
t0: f64,
param: Option<String>,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
BorderConditions: HashMap<String, Vec<(usize, f64)>>,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
method: String,
bandwidth: Option<(usize, usize)>,
policy: BackendSelectionPolicy,
resolver: Option<&AotResolver>,
residual_strategy: ResidualChunkingStrategy,
jacobian_strategy: SparseChunkingStrategy,
) -> Result<BvpSparseSolverBundle, BvpBackendIntegrationError> {
let execution = self.try_generate_BVP_with_backend_selection_and_chunking(
eq_system,
values,
arg,
params,
t0,
param,
n_steps,
h,
mesh,
BorderConditions,
Bounds,
rel_tolerance,
scheme,
method,
bandwidth,
policy,
resolver,
residual_strategy,
jacobian_strategy,
)?;
Ok(self.into_sparse_solver_bundle(execution))
}
fn flattened_argument_names_owned(&self, variable_str: &[&str]) -> Vec<String> {
let mut names = Vec::with_capacity(
self.parameters_string
.len()
.saturating_add(variable_str.len()),
);
names.extend(self.parameters_string.iter().cloned());
names.extend(variable_str.iter().map(|name| (*name).to_string()));
names
}
pub fn set_backend_config(&mut self, config: BvpBackendConfig) {
self.backend_config = config;
self.method = config.matrix_backend.legacy_method().to_string();
}
pub fn backend_config(&self) -> BvpBackendConfig {
if let Some(matrix_backend) = BvpMatrixBackend::from_legacy_method(self.method.as_str()) {
BvpBackendConfig::new(self.backend_config.backend_kind, matrix_backend)
} else {
self.backend_config
}
}
fn compile_jacobian_for_matrix_backend(
&mut self,
arg: &str,
variable_str: Vec<&str>,
matrix_backend: BvpMatrixBackend,
) {
match matrix_backend {
BvpMatrixBackend::Dense => self.lambdify_jacobian_DMatrix_par(arg, variable_str),
BvpMatrixBackend::SprsCsMat => self.lambdify_jacobian_CsMat(arg, variable_str),
BvpMatrixBackend::FaerSparseCol => {
self.lambdify_jacobian_SparseColMat_parallel2(arg, variable_str)
}
BvpMatrixBackend::Banded => self.lambdify_jacobian_Banded(arg, variable_str),
}
}
fn compile_residual_for_matrix_backend(
&mut self,
arg: &str,
variable_str: Vec<&str>,
matrix_backend: BvpMatrixBackend,
) {
match matrix_backend {
BvpMatrixBackend::Dense => self.lambdify_residual_DVector(arg, variable_str),
BvpMatrixBackend::SprsCsMat => self.lambdify_residual_CsVec(arg, variable_str),
BvpMatrixBackend::FaerSparseCol => {
self.lambdify_residual_Col_parallel2(arg, variable_str)
}
BvpMatrixBackend::Banded => self.lambdify_residual_Banded(arg, variable_str),
}
}
pub fn compile_lambdified_problem_with_config(
&mut self,
arg: &str,
variable_str: Vec<&str>,
config: BvpBackendConfig,
) {
match config.backend_kind {
BvpBackendKind::Lambdify => {
self.set_backend_config(config);
let jacobian_started = Instant::now();
self.compile_jacobian_for_matrix_backend(
arg,
variable_str.clone(),
config.matrix_backend,
);
let jacobian_elapsed = jacobian_started.elapsed().as_secs_f64();
let residual_started = Instant::now();
self.compile_residual_for_matrix_backend(arg, variable_str, config.matrix_backend);
self.last_lambdify_binding_timer_snapshot = Some(HashMap::from([
(
"lambdify jacobian callback compile time".to_string(),
jacobian_elapsed,
),
(
"lambdify residual callback compile time".to_string(),
residual_started.elapsed().as_secs_f64(),
),
]));
}
BvpBackendKind::Numeric => {
let _ = (arg, variable_str, config);
panic_symbolic_numeric_backend_requested()
}
BvpBackendKind::Aot => {
panic!("AOT backend integration is not wired inside symbolic_functions_BVP yet")
}
}
}
pub fn calc_jacobian_parallel_smart_optimized_with_given_bandwidth(&mut self) {
assert!(self.variables_for_all_disrete.len() > 0);
assert!(
!self.vector_of_functions.is_empty(),
"vector_of_functions is empty"
);
assert!(
!self.vector_of_variables.is_empty(),
"vector_of_variables is empty"
);
let variable_string_vec = &self.variable_string;
let bandwidth = self.bandwidth;
let vector_of_functions_len = self.vector_of_functions.len();
let vector_of_variables_len = self.vector_of_variables.len();
let sparse_first_backend = matches!(
self.backend_config.matrix_backend,
BvpMatrixBackend::Banded | BvpMatrixBackend::FaerSparseCol
);
let mut timings = HashMap::new();
let variable_sets_begin = Instant::now();
let variable_sets: Vec<HashSet<&String>> = self
.variables_for_all_disrete
.iter()
.map(|vars| vars.iter().collect())
.collect();
timings.insert(
"symbolic jacobian variable sets time".to_string(),
variable_sets_begin.elapsed().as_secs_f64(),
);
if sparse_first_backend {
let differentiation_begin = Instant::now();
let sparse_rows: Vec<Vec<(usize, usize, Expr)>> = (0..vector_of_functions_len)
.into_par_iter()
.map(|i| {
let (right_border, left_border) = if let Some((kl, ku)) = bandwidth {
(
std::cmp::min(i + ku + 1, vector_of_variables_len),
i.saturating_sub(kl + 1),
)
} else {
(vector_of_variables_len, 0)
};
let mut sparse_entries = Vec::new();
for j in left_border..right_border {
let variable = &variable_string_vec[j];
if variable_sets[i].contains(variable) {
let partial = Expr::diff(&self.vector_of_functions[i], variable);
if !partial.is_zero() {
sparse_entries.push((i, j, partial));
}
}
}
sparse_entries
})
.collect();
timings.insert(
"symbolic jacobian row differentiation time".to_string(),
differentiation_begin.elapsed().as_secs_f64(),
);
let flatten_begin = Instant::now();
self.symbolic_jacobian_sparse = sparse_rows.into_iter().flatten().collect();
timings.insert(
"symbolic jacobian sparse cache flatten time".to_string(),
flatten_begin.elapsed().as_secs_f64(),
);
let dense_cache_begin = Instant::now();
self.symbolic_jacobian.clear();
timings.insert(
"symbolic jacobian dense cache materialize time".to_string(),
dense_cache_begin.elapsed().as_secs_f64(),
);
} else {
let differentiation_begin = Instant::now();
let rows_and_sparse: Vec<(Vec<Expr>, Vec<(usize, usize, Expr)>)> = (0
..vector_of_functions_len)
.into_par_iter()
.map(|i| {
let (right_border, left_border) = if let Some((kl, ku)) = bandwidth {
(
std::cmp::min(i + ku + 1, vector_of_variables_len),
i.saturating_sub(kl + 1),
)
} else {
(vector_of_variables_len, 0)
};
let mut row = vec![Expr::Const(0.0); vector_of_variables_len];
let mut sparse_entries = Vec::new();
for j in left_border..right_border {
let variable = &variable_string_vec[j];
if variable_sets[i].contains(variable) {
let partial = Expr::diff(&self.vector_of_functions[i], variable);
if !partial.is_zero() {
sparse_entries.push((i, j, partial.clone()));
}
row[j] = partial;
}
}
(row, sparse_entries)
})
.collect();
timings.insert(
"symbolic jacobian row differentiation time".to_string(),
differentiation_begin.elapsed().as_secs_f64(),
);
let dense_cache_begin = Instant::now();
self.symbolic_jacobian = rows_and_sparse.iter().map(|(row, _)| row.clone()).collect();
timings.insert(
"symbolic jacobian dense cache materialize time".to_string(),
dense_cache_begin.elapsed().as_secs_f64(),
);
let flatten_begin = Instant::now();
self.symbolic_jacobian_sparse = rows_and_sparse
.into_iter()
.flat_map(|(_, entries)| entries)
.collect();
timings.insert(
"symbolic jacobian sparse cache flatten time".to_string(),
flatten_begin.elapsed().as_secs_f64(),
);
}
self.last_symbolic_jacobian_timer_snapshot = Some(timings);
}
pub fn calc_jacobian_parallel_smart_optimized(&mut self) {
assert!(self.variables_for_all_disrete.len() > 0);
assert!(
!self.vector_of_functions.is_empty(),
"vector_of_functions is empty"
);
assert!(
!self.vector_of_variables.is_empty(),
"vector_of_variables is empty"
);
let variable_string_vec = &self.variable_string;
let sparse_first_backend = matches!(
self.backend_config.matrix_backend,
BvpMatrixBackend::Banded | BvpMatrixBackend::FaerSparseCol
);
let mut timings = HashMap::new();
let variable_sets_begin = Instant::now();
let variable_sets: Vec<HashSet<&String>> = self
.variables_for_all_disrete
.iter()
.map(|vars| vars.iter().collect())
.collect();
timings.insert(
"symbolic jacobian variable sets time".to_string(),
variable_sets_begin.elapsed().as_secs_f64(),
);
if sparse_first_backend {
let differentiation_begin = Instant::now();
let sparse_rows: Vec<Vec<(usize, usize, Expr)>> = self
.vector_of_functions
.par_iter()
.enumerate()
.map(|(i, function)| {
let mut sparse_entries = Vec::new();
for j in 0..self.vector_of_variables.len() {
let variable = &variable_string_vec[j];
if variable_sets[i].contains(variable) {
let partial = Expr::diff(function, variable);
if !partial.is_zero() {
sparse_entries.push((i, j, partial));
}
}
}
sparse_entries
})
.collect();
timings.insert(
"symbolic jacobian row differentiation time".to_string(),
differentiation_begin.elapsed().as_secs_f64(),
);
let flatten_begin = Instant::now();
self.symbolic_jacobian_sparse = sparse_rows.into_iter().flatten().collect();
timings.insert(
"symbolic jacobian sparse cache flatten time".to_string(),
flatten_begin.elapsed().as_secs_f64(),
);
let dense_cache_begin = Instant::now();
self.symbolic_jacobian.clear();
timings.insert(
"symbolic jacobian dense cache materialize time".to_string(),
dense_cache_begin.elapsed().as_secs_f64(),
);
} else {
let differentiation_begin = Instant::now();
let rows_and_sparse: Vec<(Vec<Expr>, Vec<(usize, usize, Expr)>)> = self
.vector_of_functions
.par_iter()
.enumerate()
.map(|(i, function)| {
let mut row = Vec::with_capacity(self.vector_of_variables.len());
let mut sparse_entries = Vec::new();
for j in 0..self.vector_of_variables.len() {
let variable = &variable_string_vec[j];
if variable_sets[i].contains(variable) {
let partial = Expr::diff(function, variable);
if !partial.is_zero() {
sparse_entries.push((i, j, partial.clone()));
}
row.push(partial);
} else {
row.push(Expr::Const(0.0));
}
}
(row, sparse_entries)
})
.collect();
timings.insert(
"symbolic jacobian row differentiation time".to_string(),
differentiation_begin.elapsed().as_secs_f64(),
);
let dense_cache_begin = Instant::now();
self.symbolic_jacobian = rows_and_sparse.iter().map(|(row, _)| row.clone()).collect();
timings.insert(
"symbolic jacobian dense cache materialize time".to_string(),
dense_cache_begin.elapsed().as_secs_f64(),
);
let flatten_begin = Instant::now();
self.symbolic_jacobian_sparse = rows_and_sparse
.into_iter()
.flat_map(|(_, entries)| entries)
.collect();
timings.insert(
"symbolic jacobian sparse cache flatten time".to_string(),
flatten_begin.elapsed().as_secs_f64(),
);
}
self.last_symbolic_jacobian_timer_snapshot = Some(timings);
}
pub fn calc_jacobian_parallel_smart(&mut self) {
assert!(self.variables_for_all_disrete.len() > 0);
assert!(
!self.vector_of_functions.is_empty(),
"vector_of_functions is empty"
);
assert!(
!self.vector_of_variables.is_empty(),
"vector_of_variables is empty"
);
let variable_string_vec = self.variable_string.clone();
let rows_and_sparse: Vec<(Vec<Expr>, Vec<(usize, usize, Expr)>)> = self
.vector_of_functions
.par_iter()
.enumerate()
.map(|(i, function)| {
let mut vector_of_partial_derivatives = Vec::new();
let mut sparse_entries = Vec::new();
for j in 0..self.vector_of_variables.len() {
let variable = &variable_string_vec[j]; let list_of_vaiables_for_this_eq = &self.variables_for_all_disrete[i]; if list_of_vaiables_for_this_eq.contains(variable) {
let mut partial = Expr::diff(&function, variable);
partial = partial.simplify();
if !partial.is_zero() {
sparse_entries.push((i, j, partial.clone()));
}
vector_of_partial_derivatives.push(partial);
} else {
vector_of_partial_derivatives.push(Expr::Const(0.0));
}
}
(vector_of_partial_derivatives, sparse_entries)
})
.collect();
self.symbolic_jacobian = rows_and_sparse.iter().map(|(row, _)| row.clone()).collect();
self.symbolic_jacobian_sparse = rows_and_sparse
.into_iter()
.flat_map(|(_, entries)| entries)
.collect();
}
pub fn calc_jacobian_parallel(&mut self) {
assert!(
!self.vector_of_functions.is_empty(),
"vector_of_functions is empty"
);
assert!(
!self.vector_of_variables.is_empty(),
"vector_of_variables is empty"
);
let variable_string_vec = self.variable_string.clone();
let rows_and_sparse: Vec<(Vec<Expr>, Vec<(usize, usize, Expr)>)> = self
.vector_of_functions
.par_iter()
.enumerate()
.map(|(i, function)| {
let mut vector_of_partial_derivatives = Vec::new();
let mut sparse_entries = Vec::new();
for j in 0..self.vector_of_variables.len() {
let mut partial = Expr::diff(&function, &variable_string_vec[j]);
partial = partial.simplify();
if !partial.is_zero() {
sparse_entries.push((i, j, partial.clone()));
}
vector_of_partial_derivatives.push(partial);
}
(vector_of_partial_derivatives, sparse_entries)
})
.collect();
self.symbolic_jacobian = rows_and_sparse.iter().map(|(row, _)| row.clone()).collect();
self.symbolic_jacobian_sparse = rows_and_sparse
.into_iter()
.flat_map(|(_, entries)| entries)
.collect();
}
pub fn lambdify_residual_VectorType(&mut self, arg: &str, variable_str: Vec<&str>) {
let vector_of_functions = &self.vector_of_functions;
fn f(
vector_of_functions: Vec<Expr>,
_arg: String,
variable_str: Vec<String>,
) -> Box<dyn Fn(f64, &dyn VectorType) -> Box<dyn VectorType>> {
Box::new(move |_x: f64, v: &dyn VectorType| -> Box<dyn VectorType> {
let mut result = v.zeros(vector_of_functions.len());
for (i, func) in vector_of_functions.iter().enumerate() {
let func = Expr::lambdify_borrowed_thread_safe(
&func,
variable_str
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
.as_slice(),
);
let v: Vec<f64> = v.iterate().collect();
let value = func(v.as_slice());
result = result.assign_value(i, value);
}
result
})
}
let fun: Box<dyn Fn(f64, &dyn VectorType) -> Box<dyn VectorType>> = f(
vector_of_functions.to_owned(),
arg.to_string(),
variable_str.iter().map(|s| s.to_string()).collect(),
);
let residual_function = convert_to_fun(fun);
self.residiual_function = residual_function;
}
pub fn jacobian_generate_generic_par(
jac: Vec<Vec<Expr>>,
vector_of_functions_len: usize,
vector_of_variables_len: usize,
variable_str: Vec<String>,
_arg: String,
bandwidth: Option<(usize, usize)>,
) -> Box<dyn Fn(f64, &(dyn VectorType + Send + Sync)) -> Box<dyn MatrixType>> {
Box::new(
move |_x: f64, v: &(dyn VectorType + Send + Sync)| -> Box<dyn MatrixType> {
let mut vector_of_derivatives =
vec![0.0; vector_of_functions_len * vector_of_variables_len];
let vector_of_derivatives_mutex = std::sync::Mutex::new(&mut vector_of_derivatives);
let mut vector_of_triplets = Vec::new();
let vector_of_triplets_mutex = std::sync::Mutex::new(&mut vector_of_triplets);
(0..vector_of_functions_len).into_par_iter().for_each(|i| {
let (right_border, left_border) = if let Some((kl, ku)) = bandwidth {
let right_border = std::cmp::min(i + ku + 1, vector_of_variables_len);
let left_border = if i as i32 - (kl as i32) - 1 < 0 {
0
} else {
i - kl - 1
};
(right_border, left_border)
} else {
let right_border = vector_of_variables_len;
let left_border = 0;
(right_border, left_border)
};
for j in left_border..right_border {
let partial_func = Expr::lambdify_borrowed_thread_safe(
&jac[i][j],
variable_str
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
.as_slice(),
);
let v_vec: Vec<f64> = v.iterate().collect();
let _vec_type = &v.vec_type();
let P = partial_func(v_vec.as_slice());
if P.abs() > T {
vector_of_derivatives_mutex.lock().unwrap()
[i * vector_of_functions_len + j] = P;
let triplet = (i, j, P);
vector_of_triplets_mutex.lock().unwrap().push(triplet);
}
}
});
let new_function_jacobian: Box<dyn MatrixType> = v.from_vector(
vector_of_functions_len,
vector_of_variables_len,
&vector_of_derivatives,
vector_of_triplets,
);
new_function_jacobian
},
) }
pub fn lambdify_jacobian_generic(&mut self, arg: &str, variable_str: Vec<&str>) {
let symbolic_jacobian = self.symbolic_jacobian.clone();
let symbolic_jacobian_rc = symbolic_jacobian.clone();
let vector_of_functions_len = self.vector_of_functions.len();
let vector_of_variables_len = self.vector_of_variables.len();
let bandwidth = self.bandwidth;
let _new_jac = Jacobian::jacobian_generate_generic_par(
symbolic_jacobian_rc,
vector_of_functions_len,
vector_of_variables_len,
variable_str.iter().map(|s| s.to_string()).collect(),
arg.to_string(),
bandwidth,
);
}
pub fn lambdify_jacobian_DMatrix_par(&mut self, _arg: &str, variable_str: Vec<&str>) {
let jac = self.symbolic_jacobian.clone();
let vector_of_functions_len = self.vector_of_functions.len();
let vector_of_variables_len = self.vector_of_variables.len();
let bandwidth = self.bandwidth;
let flattened_argument_names_owned =
self.flattened_argument_names_owned(variable_str.as_slice());
let parameter_values = self.parameter_values.clone();
let parameter_count = self.parameters_string.len();
let jacobian_positions: Vec<(usize, usize, Box<dyn Fn(&[f64]) -> f64 + Send + Sync>)> = (0
..vector_of_functions_len)
.into_par_iter()
.flat_map(|i| {
let (right_border, left_border) = if let Some((kl, ku)) = bandwidth {
let right_border = std::cmp::min(i + ku + 1, vector_of_variables_len);
let left_border = if i as i32 - (kl as i32) - 1 < 0 {
0
} else {
i - kl - 1
};
(right_border, left_border)
} else {
(vector_of_variables_len, 0)
};
(left_border..right_border)
.filter_map(|j| {
let symbolic_partial_derivative = &jac[i][j];
if !symbolic_partial_derivative.is_zero() {
let compiled_func: Box<dyn Fn(&[f64]) -> f64 + Send + Sync> =
Expr::lambdify_borrowed_thread_safe(
&symbolic_partial_derivative,
flattened_argument_names_owned
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
.as_slice(),
);
Some((i, j, compiled_func))
} else {
None
}
})
.collect::<Vec<_>>()
})
.collect();
let new_jac = Box::new(move |_x: f64, v: &DVector<f64>| -> DMatrix<f64> {
let mut v_vec = Vec::with_capacity(parameter_count + v.len());
if parameter_count > 0 {
let params = parameter_values.as_ref().unwrap_or_else(|| {
panic!("parameter values must be provided when parameters are configured")
});
v_vec.extend_from_slice(params.as_slice());
}
v_vec.extend(v.iter().cloned());
let mut matrix = DMatrix::zeros(vector_of_functions_len, vector_of_variables_len);
let matrix_mutex = Mutex::new(&mut matrix);
jacobian_positions
.par_iter()
.for_each(|(i, j, compiled_func)| {
let P = compiled_func(v_vec.as_slice());
if P.abs() > T {
let mut mat = matrix_mutex.lock().unwrap();
mat[(*i, *j)] = P;
}
});
matrix
});
let boxed_jac: Box<dyn Jac> = Box::new(JacEnum::Dense(new_jac));
self.jac_function = Some(boxed_jac);
}
pub fn lambdify_residual_DVector(&mut self, _arg: &str, variable_str: Vec<&str>) {
let vector_of_functions = &self.vector_of_functions;
let parameter_values = self.parameter_values.clone();
let parameter_count = self.parameters_string.len();
let variable_str_owned = self.flattened_argument_names_owned(variable_str.as_slice());
let compiled_functions: Vec<Box<dyn Fn(&[f64]) -> f64 + Send + Sync>> = vector_of_functions
.par_iter()
.map(|func| {
Expr::lambdify_borrowed_thread_safe(
func,
variable_str_owned
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
.as_slice(),
)
})
.collect();
let fun = Box::new(move |_x: f64, v: &DVector<f64>| -> DVector<f64> {
let mut v_vec = Vec::with_capacity(parameter_count + v.len());
if parameter_count > 0 {
let params = parameter_values.as_ref().unwrap_or_else(|| {
panic!("parameter values must be provided when parameters are configured")
});
v_vec.extend_from_slice(params.as_slice());
}
v_vec.extend(v.iter().cloned());
let result: Vec<_> = compiled_functions
.par_iter()
.map(|func| func(v_vec.as_slice()))
.collect();
DVector::from_vec(result)
});
let boxed_fun: Box<dyn Fun> = Box::new(FunEnum::Dense(fun));
self.residiual_function = boxed_fun;
}
pub fn jacobian_generate_CsMat(
jac: Vec<Vec<Expr>>,
vector_of_functions_len: usize,
vector_of_variables_len: usize,
variable_str: Vec<String>,
_arg: String,
) -> Box<dyn FnMut(f64, &CsVec<f64>) -> CsMat<f64>> {
Box::new(move |_x: f64, v: &CsVec<f64>| -> CsMat<f64> {
let mut new_function_jacobian: CsMat<f64> =
CsMat::zero((vector_of_functions_len, vector_of_variables_len));
for i in 0..vector_of_functions_len {
for j in 0..vector_of_variables_len {
let partial_func = Expr::lambdify_borrowed_thread_safe(
&jac[i][j],
variable_str
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
.as_slice(),
);
let v_vec: Vec<f64> = v.to_dense().to_vec();
new_function_jacobian.insert(j, i, partial_func(v_vec.as_slice()));
}
}
new_function_jacobian
})
} pub fn lambdify_jacobian_CsMat(&mut self, arg: &str, variable_str: Vec<&str>) {
let symbolic_jacobian = self.symbolic_jacobian.clone();
let symbolic_jacobian_rc = symbolic_jacobian.clone();
let vector_of_functions_len = self.vector_of_functions.len();
let vector_of_variables_len = self.vector_of_variables.len();
let new_jac = Jacobian::jacobian_generate_CsMat(
symbolic_jacobian_rc,
vector_of_functions_len,
vector_of_variables_len,
variable_str.iter().map(|s| s.to_string()).collect(),
arg.to_string(),
);
let boxed_jac: Box<dyn Jac> = Box::new(JacEnum::Sparse_1(new_jac));
self.jac_function = Some(boxed_jac);
}
pub fn lambdify_residual_CsVec(&mut self, arg: &str, variable_str: Vec<&str>) {
let vector_of_functions = &self.vector_of_functions;
fn f(
vector_of_functions: Vec<Expr>,
arg: String,
variable_str: Vec<String>,
) -> Box<dyn Fn(f64, &CsVec<f64>) -> CsVec<f64> + 'static> {
Box::new(move |x: f64, v: &CsVec<f64>| -> CsVec<f64> {
let mut result: CsVec<f64> = CsVec::empty(vector_of_functions.len());
for (i, func) in vector_of_functions.iter().enumerate() {
let func = Expr::lambdify_IVP_owned(
func.to_owned(),
arg.as_str(),
variable_str
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
.clone(),
);
result.append(i, func(x, v.to_dense().to_vec()));
}
result
}) } let fun = f(
vector_of_functions.to_owned(),
arg.to_string(),
variable_str.clone().iter().map(|s| s.to_string()).collect(),
);
let boxed_fun: Box<dyn Fun> = Box::new(FunEnum::Sparse_1(fun));
self.residiual_function = boxed_fun;
}
pub fn jacobian_generate_CsMatrix(
jac: Vec<Vec<Expr>>,
vector_of_functions_len: usize,
vector_of_variables_len: usize,
variable_str: Vec<String>,
_arg: String,
) -> Box<dyn FnMut(f64, &DVector<f64>) -> CsMatrix<f64>> {
Box::new(move |_x: f64, v: &DVector<f64>| -> CsMatrix<f64> {
let _number_of_possible_non_zero: usize = jac.len();
let mut new_function_jacobian: DMatrix<f64> =
DMatrix::zeros(vector_of_functions_len, vector_of_variables_len);
for i in 0..vector_of_functions_len {
for j in 0..vector_of_variables_len {
let partial_func = Expr::lambdify_borrowed_thread_safe(
&jac[i][j],
variable_str
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
.as_slice(),
);
let v_vec: Vec<f64> = v.iter().cloned().collect();
new_function_jacobian[(i, j)] = partial_func(v_vec.as_slice());
}
}
new_function_jacobian.into()
})
} pub fn lambdify_jacobian_CsMatrix(&mut self, arg: &str, variable_str: Vec<&str>) {
let _ = (arg, variable_str);
panic!(
"BVP nalgebra CsMatrix (`Sparse_2`) runtime backend was removed because its \
MatrixType implementation was incomplete. Use `Sparse`/faer, `Banded`, \
`Sparse_1`/sprs, or `Dense` instead."
);
}
pub fn lambdify_jacobian_SparseColMat_parallel2(
&mut self,
_arg: &str,
variable_str: Vec<&str>,
) {
let vector_of_functions_len = self.vector_of_functions.len();
let vector_of_variables_len = self.vector_of_variables.len();
let flattened_argument_names_owned =
self.flattened_argument_names_owned(variable_str.as_slice());
let parameter_values = self.parameter_values.clone();
let parameter_count = self.parameters_string.len();
let sparse_entries = if !self.symbolic_jacobian_sparse.is_empty() {
self.symbolic_jacobian_sparse.clone()
} else {
self.symbolic_jacobian
.iter()
.enumerate()
.flat_map(|(row, values)| {
values
.iter()
.enumerate()
.filter_map(move |(col, expression)| {
(!expression.is_zero()).then(|| (row, col, expression.clone()))
})
})
.collect()
};
let argument_names = flattened_argument_names_owned
.iter()
.map(|name| name.as_str())
.collect::<Vec<_>>();
let jacobian_positions: Vec<(usize, usize, Box<dyn Fn(&[f64]) -> f64 + Send + Sync>)> =
sparse_entries
.into_par_iter()
.map(|(i, j, symbolic_partial_derivative)| {
let compiled_func: Box<dyn Fn(&[f64]) -> f64 + Send + Sync> =
Expr::lambdify_borrowed_thread_safe(
&symbolic_partial_derivative,
argument_names.as_slice(),
);
(i, j, compiled_func)
})
.collect();
let new_jac = Box::new(move |_x: f64, v: &Col<f64>| -> SparseColMat<usize, f64> {
let dense_values: Vec<f64> = v.iter().cloned().collect();
let params =
if parameter_count > 0 {
Some(
parameter_values.as_ref().unwrap_or_else(|| {
panic!("parameter values must be provided when parameters are configured")
})
.as_slice(),
)
} else {
None
};
let v_vec = flatten_runtime_args(params, dense_values.as_slice());
let triplets_mutex = std::sync::Mutex::new(Vec::new());
jacobian_positions
.par_iter()
.for_each(|(i, j, compiled_func)| {
let P = compiled_func(v_vec.as_slice());
if P.abs() > T {
let mut triplets = triplets_mutex.lock().unwrap();
triplets.push(Triplet::new(*i, *j, P));
}
});
let triplets = triplets_mutex.lock().unwrap();
SparseColMat::try_new_from_triplets(
vector_of_functions_len,
vector_of_variables_len,
triplets.as_slice(),
)
.unwrap()
});
let boxed_jac: Box<dyn Jac> = Box::new(JacEnum::Sparse_3(new_jac));
self.jac_function = Some(boxed_jac);
}
pub fn lambdify_residual_Col_parallel2(&mut self, _arg: &str, variable_str: Vec<&str>) {
let vector_of_functions = &self.vector_of_functions;
let parameter_values = self.parameter_values.clone();
let parameter_count = self.parameters_string.len();
let variable_str_owned = self.flattened_argument_names_owned(variable_str.as_slice());
let compiled_functions: Vec<Box<dyn Fn(&[f64]) -> f64 + Send + Sync>> = vector_of_functions
.into_par_iter()
.map(|func| {
let compiled = Expr::lambdify_borrowed_thread_safe(
func,
variable_str_owned
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
.as_slice(),
);
compiled
})
.collect();
let fun = Box::new(move |_x: f64, v: &Col<f64>| -> Col<f64> {
let dense_values: Vec<f64> = v.iter().cloned().collect();
let params =
if parameter_count > 0 {
Some(
parameter_values.as_ref().unwrap_or_else(|| {
panic!("parameter values must be provided when parameters are configured")
})
.as_slice(),
)
} else {
None
};
let v_vec = flatten_runtime_args(params, dense_values.as_slice());
let result: Vec<_> = compiled_functions
.par_iter()
.map(|func| func(v_vec.as_slice()))
.collect();
ColRef::from_slice(result.as_slice()).to_owned()
});
let boxed_fun: Box<dyn Fun> = Box::new(FunEnum::Sparse_3(fun));
self.residiual_function = boxed_fun;
}
pub fn lambdify_jacobian_Banded(&mut self, _arg: &str, _variable_str: Vec<&str>) {
let plan = self
.infer_banded_structure_plan()
.expect("banded backend requires node-major compatible Jacobian structure");
let mut config = self.banded_lambdify_config.clone();
config.linear_solver_config = self.banded_linear_solver_config;
let evaluator = self
.generate_banded_jacobian_assembly_from_plan_parallel(&plan, &config)
.expect("failed to compile banded jacobian evaluator");
let layout = plan.layout;
let solver_config = config.linear_solver_config;
let jacobian = convert_to_jac(Box::new(move |_x: f64, vec: &dyn VectorType| {
let dense = vec.to_DVectorType();
let assembly = evaluator(dense.as_slice())
.expect("banded jacobian evaluation failed for current iterate");
Box::new(BandedMatrixType::new(assembly, layout, solver_config)) as Box<dyn MatrixType>
}));
self.jac_function = Some(jacobian);
}
pub fn lambdify_residual_Banded(&mut self, _arg: &str, _variable_str: Vec<&str>) {
let evaluator = self
.generate_banded_residual_parallel()
.expect("failed to compile banded residual evaluator");
let residual = convert_to_fun(Box::new(move |_x: f64, vec: &dyn VectorType| {
let dense = vec.to_DVectorType();
let values = evaluator(dense.as_slice())
.expect("banded residual evaluation failed for current iterate");
Box::new(DVector::from_vec(values)) as Box<dyn VectorType>
}));
self.residiual_function = residual;
}
pub fn remove_numeric_suffix(input: &str) -> String {
if let Some(pos) = input.rfind('_') {
input[..pos].to_string()
} else {
input.to_string()
}
}
fn eq_step(
matrix_of_names: &[Vec<String>],
eq_i: &Expr,
values: &[String],
arg: &str,
j: usize,
t: f64,
scheme: &str,
) -> Expr {
let vec_of_names_on_step = &matrix_of_names[j];
let hashmap_for_rename: HashMap<String, String> = values
.iter()
.zip(vec_of_names_on_step.iter())
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect();
let eq_step_j = eq_i.rename_variables(&hashmap_for_rename);
let eq_step_j = eq_step_j.set_variable(arg, t);
match scheme {
"forward" => eq_step_j,
"trapezoid" => {
let vec_of_names_on_step = &matrix_of_names[j + 1];
let hashmap_for_rename: HashMap<String, String> = values
.iter()
.zip(vec_of_names_on_step.iter())
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect();
let eq_step_j_plus_1 = eq_i
.rename_variables(&hashmap_for_rename)
.set_variable(arg, t);
Expr::Const(0.5) * (eq_step_j + eq_step_j_plus_1)
}
_ => panic!("Invalid scheme"),
}
}
fn create_mesh(
&self,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
t0: f64,
) -> (Vec<Expr>, Vec<f64>, usize) {
if let Some(mesh) = mesh {
let n_steps = mesh.len();
info!("mesh with n_steps = {} is defined directly", n_steps);
let H: Vec<Expr> = mesh
.windows(2)
.map(|window| Expr::Const(window[1] - window[0]))
.collect();
(H, mesh, n_steps)
} else {
let n_steps = n_steps.unwrap_or(100) + 1;
info!(
"mesh is not defined, creating evenly distributed mesh of length {}",
n_steps
);
let h = h.unwrap_or(1.0);
let H: Vec<Expr> = vec![Expr::Const(h); n_steps - 1]; let T_list: Vec<f64> = (0..n_steps).map(|i| t0 + (i as f64) * h).collect();
(H, T_list, n_steps)
}
}
fn process_bounds_and_tolerances(
&mut self,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
flat_list_of_names: Vec<String>,
) {
let len = flat_list_of_names.len();
self.bounds = Bounds.as_ref().map(|bounds_map| {
let mut vec_of_bounds = Vec::with_capacity(len);
for name in &flat_list_of_names {
let base_name = if let Some(pos) = name.rfind('_') {
&name[..pos]
} else {
name
};
if let Some(&bound_pair) = bounds_map.get(base_name) {
vec_of_bounds.push(bound_pair);
}
}
vec_of_bounds
});
self.rel_tolerance_vec = rel_tolerance.as_ref().map(|tolerance_map| {
let mut vec_of_tolerance = Vec::with_capacity(len);
for name in &flat_list_of_names {
let base_name = if let Some(pos) = name.rfind('_') {
&name[..pos]
} else {
name
};
if let Some(&tolerance) = tolerance_map.get(base_name) {
vec_of_tolerance.push(tolerance);
}
}
vec_of_tolerance
});
}
pub fn discretization_system_BVP_par(
&mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
t0: f64,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
BorderConditions: HashMap<String, Vec<(usize, f64)>>,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
) {
let total_start = Instant::now();
let mut timer_hash: HashMap<String, f64> = HashMap::new();
let (H, T_list, n_steps) = self.create_mesh(n_steps, h, mesh, t0);
let (matrix_of_expr, matrix_of_names) = Expr::IndexedVarsMatrix(n_steps, values.clone());
let bc_handling = Instant::now();
let bc_lookup: HashMap<String, HashMap<usize, f64>> = BorderConditions
.into_iter()
.map(|(k, v)| (k, v.into_iter().collect()))
.collect();
let mut vars_for_boundary_conditions = HashMap::new();
let mut vars_to_exclude = HashSet::new();
let mut bc_pos_n_values = Vec::new();
for (var_name, conditions) in &bc_lookup {
if let Some(var_idx) = values.iter().position(|v| v == var_name) {
for (&pos, &value) in conditions {
match pos {
0 => {
let var_name = &matrix_of_names[0][var_idx];
vars_for_boundary_conditions.insert(var_name.clone(), value);
vars_to_exclude.insert(var_name.clone());
let full_pos = 0 * values.len() + var_idx;
bc_pos_n_values.push((full_pos, 0 as usize, value));
}
1 => {
let var_name = &matrix_of_names[n_steps - 1][var_idx];
vars_for_boundary_conditions.insert(var_name.clone(), value);
vars_to_exclude.insert(var_name.clone());
let full_pos = (n_steps - 1) * values.len() + var_idx;
bc_pos_n_values.push((full_pos, 1 as usize, value));
}
_ => {}
}
}
}
}
timer_hash.insert(
"bc handling".to_string(),
bc_handling.elapsed().as_millis() as f64,
);
println!("creating discretization equations");
let discretization_start = Instant::now();
let (discreditized_system, variables_for_all_discrete): (Vec<Expr>, Vec<Vec<String>>) = (0
..n_steps - 1)
.into_par_iter()
.flat_map(|j| {
let t = T_list[j];
eq_system
.par_iter()
.enumerate()
.map(|(i, eq_i)| {
let eq_step_j =
Self::eq_step(&matrix_of_names, eq_i, &values, &arg, j, t, &scheme);
let Y_j_plus_1 = &matrix_of_expr[j + 1][i];
let Y_j = &matrix_of_expr[j][i];
let res_ij = Y_j_plus_1.clone() - Y_j.clone() - H[j].clone() * eq_step_j;
let mut vars_in_equation = Vec::new();
let y_j_plus_1_name = &matrix_of_names[j + 1][i];
let y_j_name = &matrix_of_names[j][i];
if !vars_to_exclude.contains(y_j_plus_1_name) {
vars_in_equation.push(y_j_plus_1_name.clone());
}
if !vars_to_exclude.contains(y_j_name) {
vars_in_equation.push(y_j_name.clone());
}
for var_idx in 0..values.len() {
let var_name = &matrix_of_names[j][var_idx];
if !vars_to_exclude.contains(var_name) {
vars_in_equation.push(var_name.clone());
}
}
(res_ij.simplify(), vars_in_equation)
})
.collect::<Vec<_>>()
})
.unzip();
self.variables_for_all_disrete = variables_for_all_discrete;
timer_hash.insert(
"discretization of equations".to_string(),
discretization_start.elapsed().as_millis() as f64,
);
let start_flat_list = Instant::now();
let total_vars = values.len() * n_steps;
let mut flat_list_of_names = Vec::with_capacity(total_vars);
let mut flat_list_of_expr = Vec::with_capacity(total_vars);
for time_idx in 0..n_steps {
for var_idx in 0..values.len() {
let name = &matrix_of_names[time_idx][var_idx];
if !vars_to_exclude.contains(name) {
flat_list_of_names.push(name.clone());
flat_list_of_expr.push(matrix_of_expr[time_idx][var_idx].clone());
}
}
}
timer_hash.insert(
"flat list creation".to_string(),
start_flat_list.elapsed().as_millis() as f64,
);
let BC_application_start = Instant::now();
let discreditized_system_with_BC: Vec<Expr> = discreditized_system
.into_par_iter()
.map(|mut eq_i| {
for (var_name, &value) in &vars_for_boundary_conditions {
eq_i = eq_i.set_variable(var_name, value);
}
eq_i.simplify()
})
.collect();
let discreditized_system_flat = discreditized_system_with_BC;
timer_hash.insert(
"BC application".to_string(),
BC_application_start.elapsed().as_millis() as f64,
);
let consistency_start = Instant::now();
let hashset_of_vars: HashSet<&String> = flat_list_of_names.iter().collect();
let mut missing_vars = Vec::new();
for var_list in &self.variables_for_all_disrete {
for var in var_list {
if !hashset_of_vars.contains(var) {
missing_vars.push(var.clone());
}
}
}
if !missing_vars.is_empty() {
missing_vars.sort_unstable();
missing_vars.dedup();
panic!("Variables not found in system: {:?}", missing_vars);
}
timer_hash.insert(
"consistency test".to_string(),
consistency_start.elapsed().as_millis() as f64,
);
self.vector_of_functions = discreditized_system_flat;
self.vector_of_variables = flat_list_of_expr;
self.variable_string = flat_list_of_names.clone();
self.BC_pos_n_values = bc_pos_n_values;
let bounds_and_tolerances_start = Instant::now();
self.process_bounds_and_tolerances(Bounds, rel_tolerance, flat_list_of_names);
timer_hash.insert(
"bounds and tolerances".to_string(),
bounds_and_tolerances_start.elapsed().as_millis() as f64,
);
let total_end = total_start.elapsed().as_millis() as f64;
normalize_timer_percent(timer_hash.get_mut("bc handling").unwrap(), total_end);
normalize_timer_percent(
timer_hash.get_mut("discretization of equations").unwrap(),
total_end,
);
normalize_timer_percent(timer_hash.get_mut("BC application").unwrap(), total_end);
normalize_timer_percent(timer_hash.get_mut("flat list creation").unwrap(), total_end);
normalize_timer_percent(timer_hash.get_mut("consistency test").unwrap(), total_end);
normalize_timer_percent(
timer_hash.get_mut("bounds and tolerances").unwrap(),
total_end,
);
timer_hash.insert("total time, sec".to_string(), total_end);
let mut table = Builder::from(timer_hash.clone()).build();
table.with(Style::modern_rounded());
println!("{}", table.to_string());
}
fn find_bandwidths(&mut self) {
if !self.symbolic_jacobian_sparse.is_empty() {
let (kl, ku) = self
.symbolic_jacobian_sparse
.par_iter()
.map(|(row, col, _)| {
if col > row {
(0, col - row)
} else {
(row - col, 0)
}
})
.reduce(
|| (0, 0),
|acc, entry| (acc.0.max(entry.0), acc.1.max(entry.1)),
);
self.bandwidth = Some((kl, ku));
return;
}
let A = &self.symbolic_jacobian;
let n = A.len();
let (kl, ku) = (0..n)
.into_par_iter()
.map(|i| {
let mut row_kl = 0;
let mut row_ku = 0;
for j in 0..n {
if A[i][j] != Expr::Const(0.0) {
if j > i {
row_ku = std::cmp::max(row_ku, j - i);
} else if i > j {
row_kl = std::cmp::max(row_kl, i - j);
}
}
}
(row_kl, row_ku)
})
.reduce(
|| (0, 0),
|acc, row| (std::cmp::max(acc.0, row.0), std::cmp::max(acc.1, row.1)),
);
self.bandwidth = Some((kl, ku));
}
pub fn generate_BVP(
&mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
t0: f64,
param: Option<String>,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
BorderConditions: HashMap<String, Vec<(usize, f64)>>,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
method: String,
bandwidth: Option<(usize, usize)>,
) {
self.generate_BVP_with_params(
eq_system,
values,
arg,
None,
t0,
param,
n_steps,
h,
mesh,
BorderConditions,
Bounds,
rel_tolerance,
scheme,
method,
bandwidth,
);
}
pub fn generate_BVP_with_params(
&mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
params: Option<&[&str]>,
t0: f64,
param: Option<String>,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
BorderConditions: HashMap<String, Vec<(usize, f64)>>,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
method: String,
bandwidth: Option<(usize, usize)>,
) {
let (arg_name, indexed_values_owned, matrix_backend, mut timer_hash, total_start) = self
.prepare_symbolic_bvp_stage_with_params(
eq_system,
values,
arg,
params,
t0,
n_steps,
h,
mesh,
BorderConditions,
Bounds,
rel_tolerance,
scheme,
method,
bandwidth,
);
let param = param.unwrap_or(arg_name);
let indexed_values: Vec<&str> = indexed_values_owned.iter().map(|x| x.as_str()).collect();
let now = Instant::now();
self.compile_jacobian_for_matrix_backend(
param.as_str(),
indexed_values.clone(),
matrix_backend,
);
timer_hash.insert(
"jacobian lambdify time".to_string(),
now.elapsed().as_secs_f64(),
);
info!("Jacobian lambdified");
if !matches!(
matrix_backend,
BvpMatrixBackend::Banded | BvpMatrixBackend::FaerSparseCol
) {
let n = &self.symbolic_jacobian.len();
for vec_s in &self.symbolic_jacobian {
assert_eq!(vec_s.len(), *n, "jacobian not square ");
}
} else {
assert!(
!self.symbolic_jacobian_sparse.is_empty(),
"sparse-first Jacobian cache must be available for compilation"
);
}
let now = Instant::now();
self.compile_residual_for_matrix_backend(param.as_str(), indexed_values, matrix_backend);
timer_hash.insert(
"residual functions lambdify time".to_string(),
now.elapsed().as_secs_f64(),
);
info!("Residuals vector lambdified");
let total_end = total_start.elapsed().as_secs_f64();
self.finalize_generate_bvp_timer_table(timer_hash, total_end);
}
pub fn generate_BVP_with_backend_selection(
&mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
params: Option<&[&str]>,
t0: f64,
param: Option<String>,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
BorderConditions: HashMap<String, Vec<(usize, f64)>>,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
method: String,
bandwidth: Option<(usize, usize)>,
policy: BackendSelectionPolicy,
resolver: Option<&AotResolver>,
) -> BvpSparseExecutionPlan {
let (arg_name, indexed_values_owned, _matrix_backend, mut timer_hash, total_start) = self
.prepare_symbolic_bvp_stage_with_params(
eq_system,
values,
arg,
params,
t0,
n_steps,
h,
mesh,
BorderConditions,
Bounds,
rel_tolerance,
scheme,
method,
bandwidth,
);
let effective_arg_name = param.unwrap_or(arg_name);
let residual_strategy =
recommended_residual_chunking_for_parallelism(self.vector_of_functions.len(), 4);
let jacobian_strategy =
recommended_row_chunking_for_parallelism(self.vector_of_functions.len(), 4);
let execution = self.prepare_sparse_backend_execution_timed(
effective_arg_name.as_str(),
indexed_values_owned
.iter()
.map(|name| name.as_str())
.collect(),
"eval_bvp_residual",
"eval_bvp_sparse_values",
residual_strategy,
jacobian_strategy,
policy,
resolver,
&mut timer_hash,
);
let total_end = total_start.elapsed().as_secs_f64();
self.finalize_generate_bvp_timer_table(timer_hash, total_end);
execution
}
#[allow(clippy::too_many_arguments)]
pub fn generate_BVP_with_backend_selection_and_chunking(
&mut self,
eq_system: Vec<Expr>,
values: Vec<String>,
arg: String,
params: Option<&[&str]>,
t0: f64,
param: Option<String>,
n_steps: Option<usize>,
h: Option<f64>,
mesh: Option<Vec<f64>>,
BorderConditions: HashMap<String, Vec<(usize, f64)>>,
Bounds: Option<HashMap<String, (f64, f64)>>,
rel_tolerance: Option<HashMap<String, f64>>,
scheme: String,
method: String,
bandwidth: Option<(usize, usize)>,
policy: BackendSelectionPolicy,
resolver: Option<&AotResolver>,
residual_strategy: ResidualChunkingStrategy,
jacobian_strategy: SparseChunkingStrategy,
) -> BvpSparseExecutionPlan {
let (arg_name, indexed_values_owned, _matrix_backend, mut timer_hash, total_start) = self
.prepare_symbolic_bvp_stage_with_params(
eq_system,
values,
arg,
params,
t0,
n_steps,
h,
mesh,
BorderConditions,
Bounds,
rel_tolerance,
scheme,
method,
bandwidth,
);
let effective_arg_name = param.unwrap_or(arg_name);
let execution = self.prepare_sparse_backend_execution_timed(
effective_arg_name.as_str(),
indexed_values_owned
.iter()
.map(|name| name.as_str())
.collect(),
"eval_bvp_residual",
"eval_bvp_sparse_values",
residual_strategy,
jacobian_strategy,
policy,
resolver,
&mut timer_hash,
);
let total_end = total_start.elapsed().as_secs_f64();
self.finalize_generate_bvp_timer_table(timer_hash, total_end);
execution
}
}