Skip to main content

kryst/preconditioner/dist/
coarse.rs

1use crate::error::KError;
2use std::fmt;
3use std::str::FromStr;
4
5/// Unified coarse-level strategy for distributed preconditioners.
6#[derive(Clone, Copy, Debug, PartialEq, Eq)]
7pub enum DistCoarseStrategy {
8    /// No coarse correction (rank-local only).
9    None,
10    /// Gather to rank 0 and solve on the root.
11    RootGather,
12    /// Per-rank local prototype (optionally with halo correction).
13    LocalPrototype,
14    /// External distributed backend (e.g., SuperLU_DIST) when available.
15    SuperLuDist,
16}
17
18/// Explicit coarse ownership policy for distributed coarse-grid operators/solves.
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub enum DistCoarsePolicy {
21    /// Gather coarse system/solve state to root rank only.
22    RootGather,
23    /// Replicate coarse system state on all ranks.
24    Replicated,
25    /// Route through an external distributed backend.
26    External,
27}
28
29impl Default for DistCoarsePolicy {
30    fn default() -> Self {
31        Self::RootGather
32    }
33}
34
35impl FromStr for DistCoarsePolicy {
36    type Err = KError;
37
38    fn from_str(value: &str) -> Result<Self, Self::Err> {
39        match value.to_lowercase().as_str() {
40            "root" | "root_gather" | "gather" => Ok(Self::RootGather),
41            "replicated" | "replicate" | "all" => Ok(Self::Replicated),
42            "external" | "backend" | "superlu_dist" => Ok(Self::External),
43            other => Err(KError::InvalidInput(format!(
44                "invalid dist coarse policy: {other}"
45            ))),
46        }
47    }
48}
49
50/// Policy for repartitioning the coarse grid in distributed AMG setups.
51#[derive(Clone, Copy, Debug, PartialEq, Eq)]
52pub enum DistCoarseRepartition {
53    /// Preserve the fine-grid ownership partition.
54    Keep,
55    /// Build an even contiguous row partition over all ranks.
56    Uniform,
57    /// Route setup through a root-owned coarse partition.
58    Root,
59}
60
61impl Default for DistCoarseRepartition {
62    fn default() -> Self {
63        Self::Keep
64    }
65}
66
67impl FromStr for DistCoarseRepartition {
68    type Err = KError;
69
70    fn from_str(value: &str) -> Result<Self, Self::Err> {
71        match value.to_lowercase().as_str() {
72            "keep" | "fine" | "inherit" => Ok(Self::Keep),
73            "uniform" | "block" | "contiguous" => Ok(Self::Uniform),
74            "hybrid" | "local" => Ok(Self::Uniform),
75            "root" | "root_owned" => Ok(Self::Root),
76            other => Err(KError::InvalidInput(format!(
77                "invalid dist coarse repartition policy: {other}"
78            ))),
79        }
80    }
81}
82
83/// Explicit route for the coarse solve backend in distributed AMG.
84#[derive(Clone, Copy, Debug, PartialEq, Eq)]
85pub enum DistCoarseSolverRoute {
86    /// Choose the backend implied by other settings.
87    Auto,
88    /// Force gathered root-level solve path.
89    Root,
90    /// Force local prototype coarse path (fallback/prototyping).
91    Local,
92    /// Force SuperLU_DIST coarse routing when available.
93    SuperLuDist,
94}
95
96impl Default for DistCoarseSolverRoute {
97    fn default() -> Self {
98        Self::Auto
99    }
100}
101
102impl FromStr for DistCoarseSolverRoute {
103    type Err = KError;
104
105    fn from_str(value: &str) -> Result<Self, Self::Err> {
106        match value.to_lowercase().as_str() {
107            "auto" => Ok(Self::Auto),
108            "root" | "root_gather" | "gather" => Ok(Self::Root),
109            "local" | "local_prototype" | "hybrid" => Ok(Self::Local),
110            "superlu_dist" | "superludist" => Ok(Self::SuperLuDist),
111            other => Err(KError::InvalidInput(format!(
112                "invalid dist coarse solver route: {other}"
113            ))),
114        }
115    }
116}
117
118impl DistCoarseStrategy {
119    pub fn is_rank_local(self) -> bool {
120        matches!(
121            self,
122            DistCoarseStrategy::None | DistCoarseStrategy::LocalPrototype
123        )
124    }
125
126    pub fn is_collective(self) -> bool {
127        matches!(
128            self,
129            DistCoarseStrategy::RootGather | DistCoarseStrategy::SuperLuDist
130        )
131    }
132}
133
134impl Default for DistCoarseStrategy {
135    fn default() -> Self {
136        DistCoarseStrategy::RootGather
137    }
138}
139
140impl fmt::Display for DistCoarseStrategy {
141    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
142        let label = match self {
143            DistCoarseStrategy::None => "none",
144            DistCoarseStrategy::RootGather => "root_gather",
145            DistCoarseStrategy::LocalPrototype => "local_prototype",
146            DistCoarseStrategy::SuperLuDist => "superlu_dist",
147        };
148        write!(f, "{label}")
149    }
150}
151
152impl FromStr for DistCoarseStrategy {
153    type Err = KError;
154
155    fn from_str(value: &str) -> Result<Self, Self::Err> {
156        match value.to_lowercase().as_str() {
157            "none" | "off" => Ok(DistCoarseStrategy::None),
158            "root" | "root_gather" | "gather" => Ok(DistCoarseStrategy::RootGather),
159            "local" | "local_prototype" | "prototype" | "hybrid" => {
160                Ok(DistCoarseStrategy::LocalPrototype)
161            }
162            "superlu_dist" | "superludist" => Ok(DistCoarseStrategy::SuperLuDist),
163            other => Err(KError::InvalidInput(format!(
164                "invalid dist coarse strategy: {other}"
165            ))),
166        }
167    }
168}