optirs_core/hardware_aware/
mod.rs

1// Hardware-aware optimization routines
2//
3// This module provides optimization strategies that adapt to different hardware configurations,
4// including CPUs, GPUs, TPUs, edge devices, and distributed systems.
5
6use crate::error::{OptimError, Result};
7use crate::utils::{scalar_or, try_scalar};
8use scirs2_core::ndarray::{Array, Dimension, ScalarOperand};
9use scirs2_core::numeric::Float;
10use std::collections::HashMap;
11use std::fmt::Debug;
12
13mod adaptive_tuner;
14mod optimization_state;
15
16pub use adaptive_tuner::{
17    tuned_value_as_f64, AdaptiveTuner, TunableParameter, TuningObservation, TuningOutcome,
18    TuningRecord, TuningStrategy,
19};
20pub use optimization_state::{
21    HardwareOptimizerKind, HardwareStepReport, OptimizationState, DEFAULT_BASE_LEARNING_RATE,
22};
23
24use optimization_state::accumulation_steps_for;
25
26/// Hardware platform types
27#[derive(Debug, Clone, PartialEq)]
28pub enum HardwarePlatform {
29    /// CPU-based computation
30    CPU {
31        /// Number of cores
32        cores: usize,
33        /// Cache size in bytes
34        cache_size: usize,
35        /// SIMD instruction set availability
36        simd_support: SIMDSupport,
37    },
38    /// GPU-based computation
39    GPU {
40        /// GPU memory in bytes
41        memory: usize,
42        /// Number of compute units/streaming multiprocessors
43        compute_units: usize,
44        /// Memory bandwidth in GB/s
45        memory_bandwidth: f64,
46        /// GPU architecture
47        architecture: GPUArchitecture,
48    },
49    /// TPU (Tensor Processing Unit)
50    TPU {
51        /// TPU version
52        version: TPUVersion,
53        /// Matrix multiplication units
54        matrix_units: usize,
55        /// High bandwidth memory
56        hbm_size: usize,
57    },
58    /// Edge/Mobile devices
59    Edge {
60        /// Power budget in watts
61        power_budget: f64,
62        /// Memory constraints
63        memory_limit: usize,
64        /// Quantization support
65        quantization_support: QuantizationSupport,
66    },
67    /// Distributed system
68    Distributed {
69        /// Number of nodes
70        num_nodes: usize,
71        /// Network bandwidth between nodes
72        network_bandwidth: f64,
73        /// Node hardware type
74        node_hardware: Box<HardwarePlatform>,
75    },
76}
77
78/// SIMD instruction set support
79#[derive(Debug, Clone, Copy, PartialEq)]
80pub enum SIMDSupport {
81    /// No SIMD support
82    None,
83    /// SSE (128-bit)
84    SSE,
85    /// AVX (256-bit)
86    AVX,
87    /// AVX-512 (512-bit)
88    AVX512,
89    /// ARM NEON
90    NEON,
91}
92
93/// GPU architectures
94#[derive(Debug, Clone, Copy, PartialEq)]
95pub enum GPUArchitecture {
96    /// NVIDIA architectures
97    Pascal,
98    /// NVIDIA Volta architecture
99    Volta,
100    /// NVIDIA Turing architecture
101    Turing,
102    /// NVIDIA Ampere architecture
103    Ampere,
104    /// NVIDIA Hopper architecture
105    Hopper,
106    /// AMD architectures
107    RDNA,
108    /// AMD RDNA2 architecture
109    RDNA2,
110    /// AMD CDNA architecture
111    CDNA,
112    /// Intel architectures
113    XeHPG,
114    /// Intel Xe HPC architecture
115    XeHPC,
116}
117
118/// TPU versions
119#[derive(Debug, Clone, Copy, PartialEq)]
120pub enum TPUVersion {
121    /// TPU v1
122    V1,
123    /// TPU v2
124    V2,
125    /// TPU v3
126    V3,
127    /// TPU v4
128    V4,
129    /// TPU v5
130    V5,
131}
132
133/// Quantization support levels
134#[derive(Debug, Clone, Copy, PartialEq)]
135pub enum QuantizationSupport {
136    /// No quantization support
137    None,
138    /// 8-bit integer quantization
139    Int8,
140    /// 16-bit floating point
141    FP16,
142    /// Brain floating point
143    BF16,
144    /// 4-bit quantization
145    Int4,
146    /// Mixed precision
147    Mixed,
148}
149
150/// Hardware-specific optimization configuration
151#[derive(Debug, Clone)]
152pub struct HardwareOptimizationConfig<A: Float> {
153    /// Optimized batch size for hardware
154    pub batch_size: usize,
155    /// Memory-efficient parameter update strategy
156    pub memory_strategy: MemoryStrategy,
157    /// Parallel computation strategy
158    pub parallelization: ParallelizationStrategy,
159    /// Precision strategy
160    pub precision: PrecisionStrategy,
161    /// Hardware-specific optimizer parameters
162    pub optimizer_params: HashMap<String, A>,
163    /// Communication strategy (for distributed)
164    pub communication: Option<CommunicationStrategy>,
165}
166
167/// Memory optimization strategies
168#[derive(Debug, Clone)]
169pub enum MemoryStrategy {
170    /// Standard memory usage
171    Standard,
172    /// Gradient accumulation to reduce memory
173    GradientAccumulation {
174        /// Number of accumulation steps
175        accumulation_steps: usize,
176    },
177    /// Gradient checkpointing
178    GradientCheckpointing {
179        /// Checkpoint ratio
180        checkpoint_ratio: f64,
181    },
182    /// Parameter sharding
183    ParameterSharding {
184        /// Shard size
185        shard_size: usize,
186    },
187    /// Offloading to CPU memory
188    CPUOffloading {
189        /// Offload ratio to CPU
190        offload_ratio: f64,
191    },
192    /// Mixed memory strategies
193    Mixed {
194        /// List of memory strategies
195        strategies: Vec<MemoryStrategy>,
196        /// Weights for combining strategies
197        strategy_weights: Vec<f64>,
198    },
199}
200
201/// Parallelization strategies
202#[derive(Debug, Clone)]
203pub enum ParallelizationStrategy {
204    /// Single-threaded execution
205    SingleThread,
206    /// Data parallelism
207    DataParallel {
208        /// Number of parallel workers
209        num_workers: usize,
210    },
211    /// Model parallelism
212    ModelParallel {
213        /// Strategy for partitioning models
214        partition_strategy: PartitionStrategy,
215    },
216    /// Pipeline parallelism
217    Pipeline {
218        /// Number of pipeline stages
219        pipeline_stages: usize,
220        /// Number of micro-batches
221        micro_batches: usize,
222    },
223    /// Tensor parallelism
224    TensorParallel {
225        /// Size of tensor parallel group
226        tensor_parallel_size: usize,
227    },
228    /// Hybrid parallelism
229    Hybrid {
230        /// Data parallel size
231        data_parallel: usize,
232        /// Model parallel size
233        model_parallel: usize,
234        /// Pipeline parallel size
235        pipeline_parallel: usize,
236    },
237}
238
239/// Model partitioning strategies
240#[derive(Debug, Clone)]
241pub enum PartitionStrategy {
242    /// Layer-wise partitioning
243    LayerWise,
244    /// Depth-wise partitioning
245    DepthWise,
246    /// Width-wise partitioning
247    WidthWise,
248    /// Custom partitioning
249    Custom {
250        /// Custom partition points
251        partition_points: Vec<usize>,
252    },
253}
254
255/// Precision strategies for different hardware
256#[derive(Debug, Clone)]
257pub enum PrecisionStrategy {
258    /// Full precision (FP32)
259    FP32,
260    /// Half precision (FP16)
261    FP16,
262    /// Brain floating point (BF16)
263    BF16,
264    /// Mixed precision training
265    Mixed {
266        /// Forward pass precision
267        forward_precision: String,
268        /// Backward pass precision
269        backward_precision: String,
270        /// Enable loss scaling
271        loss_scaling: bool,
272    },
273    /// Integer quantization
274    Quantized {
275        /// Number of bits for weights
276        weight_bits: u8,
277        /// Number of bits for activations
278        activation_bits: u8,
279        /// Quantization method
280        quantization_method: String,
281    },
282}
283
284/// Communication strategies for distributed training
285#[derive(Debug, Clone)]
286pub enum CommunicationStrategy {
287    /// All-reduce communication
288    AllReduce {
289        /// All-reduce algorithm
290        algorithm: AllReduceAlgorithm,
291        /// Enable gradient compression
292        compression: bool,
293    },
294    /// Parameter server architecture
295    ParameterServer {
296        /// Number of parameter servers
297        num_servers: usize,
298        /// Update frequency
299        update_frequency: usize,
300    },
301    /// Gossip protocols
302    Gossip {
303        /// Number of neighbors
304        neighbors: usize,
305        /// Gossip communication frequency
306        gossip_frequency: usize,
307    },
308    /// Hierarchical communication
309    Hierarchical {
310        /// Number of local groups
311        local_groups: usize,
312        /// Inter-group communication strategy
313        inter_group_strategy: Box<CommunicationStrategy>,
314    },
315}
316
317/// All-reduce algorithms
318#[derive(Debug, Clone)]
319pub enum AllReduceAlgorithm {
320    /// Ring all-reduce
321    Ring,
322    /// Tree all-reduce
323    Tree,
324    /// Butterfly all-reduce
325    Butterfly,
326    /// Halving-doubling
327    HalvingDoubling,
328}
329
330/// Hardware-aware optimizer that adapts to different platforms
331#[derive(Debug)]
332pub struct HardwareAwareOptimizer<A: Float + 'static, D: Dimension + 'static> {
333    /// Target hardware platform
334    platform: HardwarePlatform,
335    /// Hardware-specific configuration
336    config: HardwareOptimizationConfig<A>,
337    /// Performance profiler
338    profiler: PerformanceProfiler<A>,
339    /// Resource monitor
340    resource_monitor: ResourceMonitor<A>,
341    /// Adaptive tuning system
342    adaptive_tuner: AdaptiveTuner<A>,
343    /// Current optimization state
344    current_state: OptimizationState<A, D>,
345}
346
347/// Performance profiler for hardware-specific metrics
348#[derive(Debug)]
349pub struct PerformanceProfiler<A: Float> {
350    /// Computation time measurements
351    computation_times: Vec<A>,
352    /// Memory usage measurements
353    memory_usage: Vec<usize>,
354    /// Energy consumption measurements
355    energy_consumption: Vec<A>,
356    /// Throughput measurements (samples/second)
357    throughput: Vec<A>,
358}
359
360/// Resource monitor for real-time hardware monitoring
361#[derive(Debug)]
362pub struct ResourceMonitor<A: Float> {
363    /// Current memory usage
364    current_memory: usize,
365    /// Peak memory usage
366    peak_memory: usize,
367    /// CPU utilization
368    cpu_utilization: A,
369    /// Power consumption
370    power_consumption: A,
371    /// Temperature readings
372    temperature: A,
373}
374
375impl<
376        A: Float
377            + ScalarOperand
378            + Debug
379            + std::iter::Sum
380            + for<'a> std::iter::Sum<&'a A>
381            + Send
382            + Sync
383            + 'static,
384        D: Dimension + 'static,
385    > HardwareAwareOptimizer<A, D>
386{
387    /// Create a new hardware-aware optimizer at the default learning rate
388    /// ([`DEFAULT_BASE_LEARNING_RATE`]).
389    pub fn new(platform: HardwarePlatform, initialparameters: Array<A, D>) -> Self {
390        Self::new_with_learning_rate(
391            platform,
392            initialparameters,
393            scalar_or(DEFAULT_BASE_LEARNING_RATE, A::zero()),
394        )
395    }
396
397    /// Create a new hardware-aware optimizer with an explicit base learning rate.
398    ///
399    /// The optimizer family is chosen by [`HardwareOptimizerKind::recommend_for`]
400    /// from the platform's default configuration, and the gradient accumulation
401    /// window from that configuration's [`MemoryStrategy`], so the returned
402    /// optimizer can take a real step immediately — before
403    /// [`Self::optimize_for_hardware`] has refined anything.
404    pub fn new_with_learning_rate(
405        platform: HardwarePlatform,
406        initialparameters: Array<A, D>,
407        base_learning_rate: A,
408    ) -> Self {
409        let config = Self::default_config_for_platform(&platform);
410        let profiler = PerformanceProfiler::new();
411        let resource_monitor = ResourceMonitor::new();
412        let adaptive_tuner = AdaptiveTuner::new();
413
414        let kind = HardwareOptimizerKind::recommend_for(&platform, &config);
415        let current_state = OptimizationState::new(
416            initialparameters,
417            kind,
418            base_learning_rate,
419            accumulation_steps_for(&config.memory_strategy),
420        );
421
422        Self {
423            platform,
424            config,
425            profiler,
426            resource_monitor,
427            adaptive_tuner,
428            current_state,
429        }
430    }
431
432    /// Optimize configuration for target hardware
433    pub fn optimize_for_hardware(&mut self) -> Result<()> {
434        match self.platform.clone() {
435            HardwarePlatform::CPU {
436                cores,
437                cache_size,
438                simd_support,
439            } => {
440                self.optimize_for_cpu(cores, cache_size, simd_support)?;
441            }
442            HardwarePlatform::GPU {
443                memory,
444                compute_units,
445                memory_bandwidth,
446                architecture,
447            } => {
448                self.optimize_for_gpu(memory, compute_units, memory_bandwidth, architecture)?;
449            }
450            HardwarePlatform::TPU {
451                version,
452                matrix_units,
453                hbm_size,
454            } => {
455                self.optimize_for_tpu(version, matrix_units, hbm_size)?;
456            }
457            HardwarePlatform::Edge {
458                power_budget,
459                memory_limit,
460                quantization_support,
461            } => {
462                self.optimize_for_edge(power_budget, memory_limit, quantization_support)?;
463            }
464            HardwarePlatform::Distributed {
465                num_nodes,
466                network_bandwidth,
467                node_hardware,
468            } => {
469                self.optimize_for_distributed(num_nodes, network_bandwidth, &node_hardware)?;
470            }
471        }
472        self.sync_optimizer_with_config();
473        Ok(())
474    }
475
476    /// Re-align the optimization state with the current configuration.
477    ///
478    /// The gradient accumulation window always follows the configured
479    /// [`MemoryStrategy`] — changing it mid-run only changes how many
480    /// micro-batches are averaged into the next update.
481    ///
482    /// The optimizer family is only rebuilt **before the first update**, because
483    /// swapping optimizers mid-run would throw away the accumulated moment state
484    /// and restart the effective schedule. After training has started, the new
485    /// recommendation is reported by [`Self::recommended_optimizer_kind`] and
486    /// takes effect only when the caller explicitly asks for it via
487    /// [`Self::adopt_recommended_optimizer`].
488    fn sync_optimizer_with_config(&mut self) {
489        self.current_state
490            .set_accumulation_steps(accumulation_steps_for(&self.config.memory_strategy));
491
492        if self.current_state.step_count() == 0 {
493            let recommended = self.recommended_optimizer_kind();
494            if recommended != self.current_state.optimizer_kind() {
495                self.current_state.rebuild_optimizer(recommended);
496            }
497        }
498    }
499
500    /// Optimizer family the current platform and configuration call for.
501    pub fn recommended_optimizer_kind(&self) -> HardwareOptimizerKind {
502        HardwareOptimizerKind::recommend_for(&self.platform, &self.config)
503    }
504
505    /// Adopt the current recommendation, discarding the existing optimizer's
506    /// accumulated state.
507    ///
508    /// Returns the family now in use.
509    pub fn adopt_recommended_optimizer(&mut self) -> HardwareOptimizerKind {
510        let recommended = self.recommended_optimizer_kind();
511        if recommended != self.current_state.optimizer_kind() {
512            self.current_state.rebuild_optimizer(recommended);
513        }
514        recommended
515    }
516
517    /// Apply one gradient through the configured optimizer.
518    ///
519    /// This is the step path the module's hardware analysis exists to configure:
520    /// the optimizer family, the learning rate schedule and the gradient
521    /// accumulation window all come from the platform configuration.
522    pub fn step(&mut self, gradients: &Array<A, D>) -> Result<HardwareStepReport<A>> {
523        self.current_state.step(gradients)
524    }
525
526    /// Current parameters.
527    pub fn parameters(&self) -> &Array<A, D> {
528        self.current_state.parameters()
529    }
530
531    /// The optimization state, for callers that need the step count, the
532    /// learning rate or the optimizer family.
533    pub fn optimization_state(&self) -> &OptimizationState<A, D> {
534        &self.current_state
535    }
536
537    /// Mutable access to the optimization state, for installing a learning rate
538    /// schedule or a caller-supplied optimizer.
539    pub fn optimization_state_mut(&mut self) -> &mut OptimizationState<A, D> {
540        &mut self.current_state
541    }
542
543    /// The adaptive tuner, for inspecting the tuning history and the tuned
544    /// parameters.
545    pub fn tuner(&self) -> &AdaptiveTuner<A> {
546        &self.adaptive_tuner
547    }
548
549    /// Mutable access to the adaptive tuner, for registering the parameters the
550    /// search may move and selecting a [`TuningStrategy`].
551    pub fn tuner_mut(&mut self) -> &mut AdaptiveTuner<A> {
552        &mut self.adaptive_tuner
553    }
554
555    /// Run the configured tuning search, then apply anything it found that this
556    /// module knows how to apply.
557    ///
558    /// A tuned parameter named `batch_size` is written back into the hardware
559    /// configuration (rounded and clamped to at least 1); every other tuned
560    /// parameter is left for the caller to read from
561    /// [`AdaptiveTuner::current_params`], because only the caller knows what it
562    /// means.
563    pub fn tune_parameters<F>(&mut self, evaluate: F) -> Result<TuningOutcome<A>>
564    where
565        F: FnMut(&HashMap<String, A>) -> Result<TuningObservation<A>>,
566    {
567        let outcome = self.adaptive_tuner.tune(evaluate)?;
568
569        if let Some(batch_size) = tuned_value_as_f64(&outcome.best_parameters, "batch_size") {
570            if !batch_size.is_finite() {
571                return Err(OptimError::InvalidParameter(
572                    "the tuner produced a non-finite batch_size".to_string(),
573                ));
574            }
575            self.config.batch_size = batch_size.round().max(1.0) as usize;
576        }
577
578        Ok(outcome)
579    }
580
581    /// CPU-specific optimizations
582    fn optimize_for_cpu(
583        &mut self,
584        cores: usize,
585        cache_size: usize,
586        simd_support: SIMDSupport,
587    ) -> Result<()> {
588        // Optimize batch _size for cache efficiency
589        let cache_friendly_batch_size =
590            (cache_size / 4) / self.current_state.parameters().len().max(1); // Rough estimate
591        self.config.batch_size = cache_friendly_batch_size.clamp(16, 512);
592
593        // Configure parallelization based on cores
594        self.config.parallelization = ParallelizationStrategy::DataParallel {
595            num_workers: cores.min(8), // Don't over-parallelize
596        };
597
598        // SIMD-specific optimizations
599        match simd_support {
600            SIMDSupport::AVX512 => {
601                self.config
602                    .optimizer_params
603                    .insert("vectorized_ops".to_string(), try_scalar::<A, _>(512.0)?);
604            }
605            SIMDSupport::AVX => {
606                self.config
607                    .optimizer_params
608                    .insert("vectorized_ops".to_string(), try_scalar::<A, _>(256.0)?);
609            }
610            SIMDSupport::SSE => {
611                self.config
612                    .optimizer_params
613                    .insert("vectorized_ops".to_string(), try_scalar::<A, _>(128.0)?);
614            }
615            SIMDSupport::NEON => {
616                self.config
617                    .optimizer_params
618                    .insert("vectorized_ops".to_string(), try_scalar::<A, _>(128.0)?);
619            }
620            SIMDSupport::None => {
621                self.config
622                    .optimizer_params
623                    .insert("vectorized_ops".to_string(), try_scalar::<A, _>(32.0)?);
624            }
625        }
626
627        // Use full precision for CPU
628        self.config.precision = PrecisionStrategy::FP32;
629
630        Ok(())
631    }
632
633    /// GPU-specific optimizations
634    fn optimize_for_gpu(
635        &mut self,
636        memory: usize,
637        compute_units: usize,
638        memory_bandwidth: f64,
639        architecture: GPUArchitecture,
640    ) -> Result<()> {
641        // Optimize batch size for GPU memory
642        let gpu_memory_gb = memory as f64 / (1024.0 * 1024.0 * 1024.0);
643        let optimal_batch_size = if gpu_memory_gb >= 32.0 {
644            256
645        } else if gpu_memory_gb >= 16.0 {
646            128
647        } else if gpu_memory_gb >= 8.0 {
648            64
649        } else {
650            32
651        };
652        self.config.batch_size = optimal_batch_size;
653
654        // Configure parallelization for GPU
655        self.config.parallelization = ParallelizationStrategy::DataParallel {
656            num_workers: compute_units.min(16),
657        };
658
659        // Architecture-specific optimizations
660        match architecture {
661            GPUArchitecture::Ampere | GPUArchitecture::Hopper => {
662                // Use mixed precision for modern architectures
663                self.config.precision = PrecisionStrategy::Mixed {
664                    forward_precision: "fp16".to_string(),
665                    backward_precision: "fp32".to_string(),
666                    loss_scaling: true,
667                };
668                self.config
669                    .optimizer_params
670                    .insert("tensor_cores".to_string(), try_scalar::<A, _>(1.0)?);
671            }
672            GPUArchitecture::Volta | GPUArchitecture::Turing => {
673                self.config.precision = PrecisionStrategy::FP16;
674                self.config
675                    .optimizer_params
676                    .insert("tensor_cores".to_string(), try_scalar::<A, _>(1.0)?);
677            }
678            _ => {
679                self.config.precision = PrecisionStrategy::FP32;
680            }
681        }
682
683        // Memory _bandwidth optimizations
684        if memory_bandwidth < 500.0 {
685            // Low _bandwidth
686            self.config.memory_strategy = MemoryStrategy::GradientAccumulation {
687                accumulation_steps: 4,
688            };
689        } else {
690            self.config.memory_strategy = MemoryStrategy::Standard;
691        }
692
693        Ok(())
694    }
695
696    /// TPU-specific optimizations
697    fn optimize_for_tpu(
698        &mut self,
699        version: TPUVersion,
700        matrix_units: usize,
701        hbm_size: usize,
702    ) -> Result<()> {
703        // TPUs work best with large batch sizes
704        let tpu_batch_size = match version {
705            TPUVersion::V1 | TPUVersion::V2 => 128,
706            TPUVersion::V3 => 256,
707            TPUVersion::V4 | TPUVersion::V5 => 512,
708        };
709        self.config.batch_size = tpu_batch_size;
710
711        // TPUs prefer BF16 precision
712        self.config.precision = PrecisionStrategy::BF16;
713
714        // Configure for matrix operations
715        self.config.optimizer_params.insert(
716            "matrix_units".to_string(),
717            try_scalar::<A, _>(matrix_units as f64)?,
718        );
719
720        // Use all available matrix _units
721        self.config.parallelization = ParallelizationStrategy::TensorParallel {
722            tensor_parallel_size: matrix_units.min(8),
723        };
724
725        // HBM-specific optimizations
726        if hbm_size > 32 * 1024 * 1024 * 1024 {
727            // 32GB+
728            self.config.memory_strategy = MemoryStrategy::Standard;
729        } else {
730            self.config.memory_strategy = MemoryStrategy::GradientCheckpointing {
731                checkpoint_ratio: 0.5,
732            };
733        }
734
735        Ok(())
736    }
737
738    /// Edge device optimizations
739    fn optimize_for_edge(
740        &mut self,
741        power_budget: f64,
742        memory_limit: usize,
743        quantization_support: QuantizationSupport,
744    ) -> Result<()> {
745        // Small batch sizes for memory constraints
746        let edge_batch_size = (memory_limit / (4 * 1024 * 1024)).clamp(1, 32); // Very conservative
747        self.config.batch_size = edge_batch_size;
748
749        // Single-threaded for power efficiency
750        self.config.parallelization = ParallelizationStrategy::SingleThread;
751
752        // Aggressive quantization for edge devices
753        match quantization_support {
754            QuantizationSupport::Int4 => {
755                self.config.precision = PrecisionStrategy::Quantized {
756                    weight_bits: 4,
757                    activation_bits: 8,
758                    quantization_method: "dynamic".to_string(),
759                };
760            }
761            QuantizationSupport::Int8 => {
762                self.config.precision = PrecisionStrategy::Quantized {
763                    weight_bits: 8,
764                    activation_bits: 8,
765                    quantization_method: "static".to_string(),
766                };
767            }
768            QuantizationSupport::FP16 => {
769                self.config.precision = PrecisionStrategy::FP16;
770            }
771            _ => {
772                self.config.precision = PrecisionStrategy::FP32;
773            }
774        }
775
776        // Power-aware optimizations
777        if power_budget < 5.0 {
778            // Very low power
779            self.config
780                .optimizer_params
781                .insert("update_frequency".to_string(), try_scalar::<A, _>(10.0)?);
782            self.config.memory_strategy = MemoryStrategy::CPUOffloading { offload_ratio: 0.8 };
783        }
784
785        Ok(())
786    }
787
788    /// Distributed system optimizations
789    fn optimize_for_distributed(
790        &mut self,
791        num_nodes: usize,
792        network_bandwidth: f64,
793        node_hardware: &HardwarePlatform,
794    ) -> Result<()> {
795        // Scale batch size with number of _nodes
796        let base_batch_size = match node_hardware {
797            HardwarePlatform::GPU { .. } => 128,
798            HardwarePlatform::CPU { .. } => 64,
799            HardwarePlatform::TPU { .. } => 256, // TPUs can handle larger batches
800            HardwarePlatform::Edge { .. } => 32, // Edge devices have memory constraints
801            HardwarePlatform::Distributed { node_hardware, .. } => {
802                // Use the underlying node hardware type for distributed systems
803                match node_hardware.as_ref() {
804                    HardwarePlatform::GPU { .. } => 128,
805                    HardwarePlatform::CPU { .. } => 64,
806                    HardwarePlatform::TPU { .. } => 256,
807                    HardwarePlatform::Edge { .. } => 32,
808                    HardwarePlatform::Distributed { .. } => 64, // Fallback for nested distributed
809                }
810            }
811        };
812        self.config.batch_size = base_batch_size * num_nodes;
813
814        // Configure communication strategy based on network _bandwidth
815        let communication = if network_bandwidth >= 100.0 {
816            // High _bandwidth (100 Gbps+)
817            CommunicationStrategy::AllReduce {
818                algorithm: AllReduceAlgorithm::Ring,
819                compression: false,
820            }
821        } else if network_bandwidth >= 10.0 {
822            // Medium _bandwidth (10 Gbps+)
823            CommunicationStrategy::AllReduce {
824                algorithm: AllReduceAlgorithm::Tree,
825                compression: true,
826            }
827        } else {
828            // Low _bandwidth
829            CommunicationStrategy::ParameterServer {
830                num_servers: (num_nodes / 4).max(1),
831                update_frequency: 10,
832            }
833        };
834        self.config.communication = Some(communication);
835
836        // Configure parallelization strategy
837        if num_nodes >= 64 {
838            self.config.parallelization = ParallelizationStrategy::Hybrid {
839                data_parallel: 8,
840                model_parallel: 4,
841                pipeline_parallel: num_nodes / 32,
842            };
843        } else if num_nodes >= 16 {
844            self.config.parallelization = ParallelizationStrategy::Pipeline {
845                pipeline_stages: 4,
846                micro_batches: 8,
847            };
848        } else {
849            self.config.parallelization = ParallelizationStrategy::DataParallel {
850                num_workers: num_nodes,
851            };
852        }
853
854        Ok(())
855    }
856
857    /// Profile current performance
858    pub fn profile_performance(&mut self, computation_time: A, memoryused: usize, energy: A) {
859        self.profiler.computation_times.push(computation_time);
860        self.profiler.memory_usage.push(memoryused);
861        self.profiler.energy_consumption.push(energy);
862
863        // Calculate throughput (simplified)
864        let throughput = scalar_or(self.config.batch_size as f64, A::zero()) / computation_time;
865        self.profiler.throughput.push(throughput);
866
867        // Keep history bounded
868        const MAX_HISTORY: usize = 1000;
869        if self.profiler.computation_times.len() > MAX_HISTORY {
870            self.profiler.computation_times.remove(0);
871            self.profiler.memory_usage.remove(0);
872            self.profiler.energy_consumption.remove(0);
873            self.profiler.throughput.remove(0);
874        }
875    }
876
877    /// Update resource monitoring
878    pub fn update_resource_monitor(&mut self, memory: usize, cpuutil: A, power: A, temp: A) {
879        self.resource_monitor.current_memory = memory;
880        self.resource_monitor.peak_memory = self.resource_monitor.peak_memory.max(memory);
881        self.resource_monitor.cpu_utilization = cpuutil;
882        self.resource_monitor.power_consumption = power;
883        self.resource_monitor.temperature = temp;
884    }
885
886    /// Adaptive tuning based on performance feedback
887    pub fn adaptive_tune(&mut self, targetperformance: A) -> Result<()> {
888        self.adaptive_tuner
889            .set_performance_target(targetperformance);
890
891        // Simple adaptive tuning logic
892        let current_performance = self.get_average_performance();
893
894        if current_performance < targetperformance {
895            // Need to improve _performance
896            self.tune_for_performance()?;
897        } else {
898            // Can optimize for efficiency
899            self.tune_for_efficiency()?;
900        }
901
902        Ok(())
903    }
904
905    /// Tune for better performance
906    fn tune_for_performance(&mut self) -> Result<()> {
907        // Increase batch size if memory allows
908        if self.resource_monitor.current_memory < self.resource_monitor.peak_memory * 8 / 10 {
909            self.config.batch_size = (self.config.batch_size * 12 / 10).min(1024);
910        }
911
912        // Reduce precision for speed
913        match self.config.precision {
914            PrecisionStrategy::FP32 => {
915                self.config.precision = PrecisionStrategy::FP16;
916            }
917            PrecisionStrategy::FP16 => {
918                self.config.precision = PrecisionStrategy::Mixed {
919                    forward_precision: "fp16".to_string(),
920                    backward_precision: "fp32".to_string(),
921                    loss_scaling: true,
922                };
923            }
924            _ => {}
925        }
926
927        Ok(())
928    }
929
930    /// Tune for better efficiency
931    fn tune_for_efficiency(&mut self) -> Result<()> {
932        // Reduce batch size to save memory
933        self.config.batch_size = (self.config.batch_size * 9 / 10).max(1);
934
935        // Enable gradient accumulation to maintain effective batch size
936        self.config.memory_strategy = MemoryStrategy::GradientAccumulation {
937            accumulation_steps: 2,
938        };
939
940        Ok(())
941    }
942
943    /// Get average performance from recent measurements
944    fn get_average_performance(&self) -> A {
945        if self.profiler.throughput.is_empty() {
946            A::zero()
947        } else {
948            let recent_throughput =
949                &self.profiler.throughput[self.profiler.throughput.len().saturating_sub(10)..];
950            recent_throughput.iter().copied().sum::<A>()
951                / scalar_or(recent_throughput.len(), A::one())
952        }
953    }
954
955    /// Get current configuration
956    pub fn get_config(&self) -> &HardwareOptimizationConfig<A> {
957        &self.config
958    }
959
960    /// Get performance statistics
961    pub fn get_performance_stats(&self) -> HardwarePerformanceStats<A> {
962        let avg_computation_time = if self.profiler.computation_times.is_empty() {
963            A::zero()
964        } else {
965            self.profiler.computation_times.iter().sum::<A>()
966                / scalar_or(self.profiler.computation_times.len(), A::one())
967        };
968
969        let avg_throughput = if self.profiler.throughput.is_empty() {
970            A::zero()
971        } else {
972            self.profiler.throughput.iter().sum::<A>()
973                / scalar_or(self.profiler.throughput.len(), A::one())
974        };
975
976        let avg_energy = if self.profiler.energy_consumption.is_empty() {
977            A::zero()
978        } else {
979            self.profiler.energy_consumption.iter().copied().sum::<A>()
980                / scalar_or(self.profiler.energy_consumption.len(), A::one())
981        };
982
983        HardwarePerformanceStats {
984            average_computation_time: avg_computation_time,
985            average_throughput: avg_throughput,
986            peak_memory_usage: self.resource_monitor.peak_memory,
987            average_energy_consumption: avg_energy,
988            hardware_utilization: self.resource_monitor.cpu_utilization,
989            efficiency_score: avg_throughput / (avg_energy + scalar_or(1e-8, A::zero())), // Avoid division by zero
990        }
991    }
992
993    /// Create default configuration for platform
994    fn default_config_for_platform(platform: &HardwarePlatform) -> HardwareOptimizationConfig<A> {
995        match platform {
996            HardwarePlatform::CPU { .. } => HardwareOptimizationConfig {
997                batch_size: 64,
998                memory_strategy: MemoryStrategy::Standard,
999                parallelization: ParallelizationStrategy::DataParallel { num_workers: 4 },
1000                precision: PrecisionStrategy::FP32,
1001                optimizer_params: HashMap::new(),
1002                communication: None,
1003            },
1004            HardwarePlatform::GPU { .. } => HardwareOptimizationConfig {
1005                batch_size: 128,
1006                memory_strategy: MemoryStrategy::Standard,
1007                parallelization: ParallelizationStrategy::DataParallel { num_workers: 1 },
1008                precision: PrecisionStrategy::FP16,
1009                optimizer_params: HashMap::new(),
1010                communication: None,
1011            },
1012            HardwarePlatform::TPU { .. } => HardwareOptimizationConfig {
1013                batch_size: 256,
1014                memory_strategy: MemoryStrategy::Standard,
1015                parallelization: ParallelizationStrategy::TensorParallel {
1016                    tensor_parallel_size: 8,
1017                },
1018                precision: PrecisionStrategy::BF16,
1019                optimizer_params: HashMap::new(),
1020                communication: None,
1021            },
1022            HardwarePlatform::Edge { .. } => HardwareOptimizationConfig {
1023                batch_size: 16,
1024                memory_strategy: MemoryStrategy::GradientCheckpointing {
1025                    checkpoint_ratio: 0.5,
1026                },
1027                parallelization: ParallelizationStrategy::SingleThread,
1028                precision: PrecisionStrategy::Quantized {
1029                    weight_bits: 8,
1030                    activation_bits: 8,
1031                    quantization_method: "dynamic".to_string(),
1032                },
1033                optimizer_params: HashMap::new(),
1034                communication: None,
1035            },
1036            HardwarePlatform::Distributed { .. } => HardwareOptimizationConfig {
1037                batch_size: 512,
1038                memory_strategy: MemoryStrategy::Standard,
1039                parallelization: ParallelizationStrategy::DataParallel { num_workers: 8 },
1040                precision: PrecisionStrategy::FP16,
1041                optimizer_params: HashMap::new(),
1042                communication: Some(CommunicationStrategy::AllReduce {
1043                    algorithm: AllReduceAlgorithm::Ring,
1044                    compression: false,
1045                }),
1046            },
1047        }
1048    }
1049}
1050
1051impl<A: Float + Send + Sync> Default for PerformanceProfiler<A> {
1052    fn default() -> Self {
1053        Self::new()
1054    }
1055}
1056
1057impl<A: Float + Send + Sync> PerformanceProfiler<A> {
1058    /// Create a new performance profiler
1059    pub fn new() -> Self {
1060        Self {
1061            computation_times: Vec::new(),
1062            memory_usage: Vec::new(),
1063            energy_consumption: Vec::new(),
1064            throughput: Vec::new(),
1065        }
1066    }
1067}
1068
1069impl<A: Float + Send + Sync> Default for ResourceMonitor<A> {
1070    fn default() -> Self {
1071        Self::new()
1072    }
1073}
1074
1075impl<A: Float + Send + Sync> ResourceMonitor<A> {
1076    /// Create a new resource monitor
1077    pub fn new() -> Self {
1078        Self {
1079            current_memory: 0,
1080            peak_memory: 0,
1081            cpu_utilization: A::zero(),
1082            power_consumption: A::zero(),
1083            temperature: A::zero(),
1084        }
1085    }
1086}
1087
1088/// Hardware performance statistics
1089#[derive(Debug, Clone)]
1090pub struct HardwarePerformanceStats<A: Float> {
1091    /// Average computation time per step
1092    pub average_computation_time: A,
1093    /// Average throughput (samples/second)
1094    pub average_throughput: A,
1095    /// Peak memory usage
1096    pub peak_memory_usage: usize,
1097    /// Average energy consumption
1098    pub average_energy_consumption: A,
1099    /// Hardware utilization percentage
1100    pub hardware_utilization: A,
1101    /// Efficiency score (throughput/energy)
1102    pub efficiency_score: A,
1103}
1104
1105#[cfg(test)]
1106mod tests {
1107    use super::*;
1108    use scirs2_core::ndarray::Array1;
1109
1110    #[test]
1111    fn test_cpu_optimization() {
1112        let platform = HardwarePlatform::CPU {
1113            cores: 8,
1114            cache_size: 32 * 1024 * 1024, // 32MB cache
1115            simd_support: SIMDSupport::AVX,
1116        };
1117
1118        let initial_params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1119        let mut optimizer = HardwareAwareOptimizer::new(platform, initial_params);
1120
1121        optimizer.optimize_for_hardware().expect("unwrap failed");
1122
1123        // Check CPU-specific optimizations
1124        assert!(optimizer.config.batch_size <= 512);
1125        assert!(matches!(
1126            optimizer.config.parallelization,
1127            ParallelizationStrategy::DataParallel { .. }
1128        ));
1129        assert!(matches!(
1130            optimizer.config.precision,
1131            PrecisionStrategy::FP32
1132        ));
1133        assert!(optimizer
1134            .config
1135            .optimizer_params
1136            .contains_key("vectorized_ops"));
1137    }
1138
1139    #[test]
1140    fn test_gpu_optimization() {
1141        let platform = HardwarePlatform::GPU {
1142            memory: 16 * 1024 * 1024 * 1024, // 16GB
1143            compute_units: 80,
1144            memory_bandwidth: 900.0,
1145            architecture: GPUArchitecture::Ampere,
1146        };
1147
1148        let initial_params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1149        let mut optimizer = HardwareAwareOptimizer::new(platform, initial_params);
1150
1151        optimizer.optimize_for_hardware().expect("unwrap failed");
1152
1153        // Check GPU-specific optimizations
1154        assert_eq!(optimizer.config.batch_size, 128);
1155        assert!(matches!(
1156            optimizer.config.precision,
1157            PrecisionStrategy::Mixed { .. }
1158        ));
1159        assert!(optimizer
1160            .config
1161            .optimizer_params
1162            .contains_key("tensor_cores"));
1163    }
1164
1165    #[test]
1166    fn test_tpu_optimization() {
1167        let platform = HardwarePlatform::TPU {
1168            version: TPUVersion::V4,
1169            matrix_units: 8,
1170            hbm_size: 32 * 1024 * 1024 * 1024, // 32GB HBM
1171        };
1172
1173        let initial_params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1174        let mut optimizer = HardwareAwareOptimizer::new(platform, initial_params);
1175
1176        optimizer.optimize_for_hardware().expect("unwrap failed");
1177
1178        // Check TPU-specific optimizations
1179        assert_eq!(optimizer.config.batch_size, 512);
1180        assert!(matches!(
1181            optimizer.config.precision,
1182            PrecisionStrategy::BF16
1183        ));
1184        assert!(matches!(
1185            optimizer.config.parallelization,
1186            ParallelizationStrategy::TensorParallel { .. }
1187        ));
1188    }
1189
1190    #[test]
1191    fn test_edge_optimization() {
1192        let platform = HardwarePlatform::Edge {
1193            power_budget: 3.0,               // 3 watts
1194            memory_limit: 512 * 1024 * 1024, // 512MB
1195            quantization_support: QuantizationSupport::Int8,
1196        };
1197
1198        let initial_params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1199        let mut optimizer = HardwareAwareOptimizer::new(platform, initial_params);
1200
1201        optimizer.optimize_for_hardware().expect("unwrap failed");
1202
1203        // Check edge-specific optimizations
1204        assert!(optimizer.config.batch_size <= 32);
1205        assert!(matches!(
1206            optimizer.config.parallelization,
1207            ParallelizationStrategy::SingleThread
1208        ));
1209        assert!(matches!(
1210            optimizer.config.precision,
1211            PrecisionStrategy::Quantized { .. }
1212        ));
1213    }
1214
1215    #[test]
1216    fn test_distributed_optimization() {
1217        let node_hardware = HardwarePlatform::GPU {
1218            memory: 8 * 1024 * 1024 * 1024, // 8GB per node
1219            compute_units: 40,
1220            memory_bandwidth: 500.0,
1221            architecture: GPUArchitecture::Volta,
1222        };
1223
1224        let platform = HardwarePlatform::Distributed {
1225            num_nodes: 16,
1226            network_bandwidth: 50.0, // 50 Gbps
1227            node_hardware: Box::new(node_hardware),
1228        };
1229
1230        let initial_params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1231        let mut optimizer = HardwareAwareOptimizer::new(platform, initial_params);
1232
1233        optimizer.optimize_for_hardware().expect("unwrap failed");
1234
1235        // Check distributed-specific optimizations
1236        assert_eq!(optimizer.config.batch_size, 128 * 16); // Scaled by number of nodes
1237        assert!(optimizer.config.communication.is_some());
1238        assert!(matches!(
1239            optimizer.config.parallelization,
1240            ParallelizationStrategy::Pipeline { .. }
1241        ));
1242    }
1243
1244    #[test]
1245    fn test_performance_profiling() {
1246        let platform = HardwarePlatform::CPU {
1247            cores: 4,
1248            cache_size: 8 * 1024 * 1024,
1249            simd_support: SIMDSupport::SSE,
1250        };
1251
1252        let initial_params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1253        let mut optimizer = HardwareAwareOptimizer::new(platform, initial_params);
1254
1255        // Add some performance measurements
1256        optimizer.profile_performance(0.1, 1000000, 5.0);
1257        optimizer.profile_performance(0.12, 1100000, 5.2);
1258        optimizer.profile_performance(0.09, 950000, 4.8);
1259
1260        let stats = optimizer.get_performance_stats();
1261
1262        assert!(stats.average_computation_time > 0.0);
1263        assert!(stats.average_throughput > 0.0);
1264        assert_eq!(stats.peak_memory_usage, 0); // Not updated in this test
1265    }
1266
1267    #[test]
1268    fn test_adaptive_tuning() {
1269        let platform = HardwarePlatform::GPU {
1270            memory: 8 * 1024 * 1024 * 1024,
1271            compute_units: 20,
1272            memory_bandwidth: 300.0,
1273            architecture: GPUArchitecture::Turing,
1274        };
1275
1276        let initial_params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1277        let mut optimizer = HardwareAwareOptimizer::new(platform, initial_params);
1278
1279        // Simulate low performance
1280        optimizer.profiler.throughput.push(50.0);
1281        optimizer.resource_monitor.current_memory = 1_000_000_000; // 1GB
1282        optimizer.resource_monitor.peak_memory = 4_000_000_000; // 4GB
1283
1284        let initial_batch_size = optimizer.config.batch_size;
1285        optimizer.adaptive_tune(100.0).expect("unwrap failed"); // Target 100 samples/sec
1286
1287        // Should have tuned for better performance
1288        assert!(optimizer.config.batch_size >= initial_batch_size);
1289    }
1290
1291    #[test]
1292    fn test_hardware_platform_matching() {
1293        let platforms = vec![
1294            HardwarePlatform::CPU {
1295                cores: 8,
1296                cache_size: 16_000_000,
1297                simd_support: SIMDSupport::AVX,
1298            },
1299            HardwarePlatform::GPU {
1300                memory: 12_000_000_000,
1301                compute_units: 60,
1302                memory_bandwidth: 600.0,
1303                architecture: GPUArchitecture::Ampere,
1304            },
1305            HardwarePlatform::TPU {
1306                version: TPUVersion::V3,
1307                matrix_units: 8,
1308                hbm_size: 16_000_000_000,
1309            },
1310            HardwarePlatform::Edge {
1311                power_budget: 2.0,
1312                memory_limit: 256_000_000,
1313                quantization_support: QuantizationSupport::Int4,
1314            },
1315        ];
1316
1317        for platform in platforms {
1318            let initial_params = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1319            let mut optimizer = HardwareAwareOptimizer::new(platform, initial_params);
1320
1321            // Should not panic and should complete successfully
1322            let result = optimizer.optimize_for_hardware();
1323            assert!(result.is_ok());
1324
1325            // Each platform should have different configurations
1326            let config = optimizer.get_config();
1327            assert!(config.batch_size > 0);
1328        }
1329    }
1330
1331    // --- Real optimization step path -------------------------------------
1332    //
1333    // `OptimizationState` used to hold a parameter array and nothing else: the
1334    // module analysed hardware and recommended strategies but never ran an
1335    // optimizer step, so a `HardwareAwareOptimizer` could not move a loss.
1336
1337    /// The headline regression: N steps on a quadratic must reduce the loss.
1338    #[test]
1339    fn the_hardware_aware_optimizer_actually_optimizes() {
1340        let platform = HardwarePlatform::CPU {
1341            cores: 8,
1342            cache_size: 32 * 1024 * 1024,
1343            simd_support: SIMDSupport::AVX,
1344        };
1345
1346        let mut optimizer = HardwareAwareOptimizer::new_with_learning_rate(
1347            platform,
1348            Array1::from_vec(vec![2.0, -3.0, 1.5]),
1349            0.05,
1350        );
1351        optimizer
1352            .optimize_for_hardware()
1353            .expect("hardware configuration must succeed");
1354
1355        let loss = |parameters: &Array1<f64>| -> f64 { parameters.iter().map(|&x| x * x).sum() };
1356        let initial_loss = loss(optimizer.parameters());
1357
1358        for _ in 0..300 {
1359            let gradient = optimizer.parameters().mapv(|x| 2.0 * x);
1360            let report = optimizer.step(&gradient).expect("step must succeed");
1361            assert!(report.applied);
1362        }
1363
1364        let final_loss = loss(optimizer.parameters());
1365        assert_eq!(optimizer.optimization_state().step_count(), 300);
1366        // The bound is loose on purpose: an adaptive optimizer settles into an
1367        // oscillation of amplitude ~lr around the minimum, so pinning the exact
1368        // residual would make this a flaky test rather than a stronger one.
1369        assert!(
1370            final_loss < initial_loss * 1e-2,
1371            "the loss did not fall ({initial_loss} -> {final_loss})"
1372        );
1373    }
1374
1375    /// The hardware analysis must actually select the optimizer: a TPU's large
1376    /// batch calls for LAMB, a low-power edge device for SGD.
1377    #[test]
1378    fn the_configuration_selects_the_optimizer_family() {
1379        let tpu = HardwarePlatform::TPU {
1380            version: TPUVersion::V4,
1381            matrix_units: 8,
1382            hbm_size: 32 * 1024 * 1024 * 1024,
1383        };
1384        let mut optimizer = HardwareAwareOptimizer::new(tpu, Array1::from_vec(vec![1.0, 2.0, 3.0]));
1385        optimizer
1386            .optimize_for_hardware()
1387            .expect("hardware configuration must succeed");
1388        assert_eq!(
1389            optimizer.optimization_state().optimizer_kind(),
1390            HardwareOptimizerKind::Lamb,
1391            "a 512-sample TPU batch calls for a large-batch optimizer"
1392        );
1393
1394        let edge = HardwarePlatform::Edge {
1395            power_budget: 2.0,
1396            memory_limit: 256 * 1024 * 1024,
1397            quantization_support: QuantizationSupport::Int8,
1398        };
1399        let mut optimizer =
1400            HardwareAwareOptimizer::new(edge, Array1::from_vec(vec![1.0, 2.0, 3.0]));
1401        optimizer
1402            .optimize_for_hardware()
1403            .expect("hardware configuration must succeed");
1404        assert_eq!(
1405            optimizer.optimization_state().optimizer_kind(),
1406            HardwareOptimizerKind::Sgd,
1407            "a 2 W budget cannot afford Adam's second moment"
1408        );
1409    }
1410
1411    /// A configured gradient-accumulation memory strategy must reach the step
1412    /// path instead of being a description nothing reads.
1413    #[test]
1414    fn the_memory_strategy_drives_gradient_accumulation() {
1415        // A low-bandwidth GPU is configured with 4-step gradient accumulation.
1416        let platform = HardwarePlatform::GPU {
1417            memory: 8 * 1024 * 1024 * 1024,
1418            compute_units: 40,
1419            memory_bandwidth: 300.0,
1420            architecture: GPUArchitecture::Turing,
1421        };
1422        let mut optimizer = HardwareAwareOptimizer::new_with_learning_rate(
1423            platform,
1424            Array1::from_vec(vec![0.0, 0.0]),
1425            0.1,
1426        );
1427        assert_eq!(optimizer.optimization_state().accumulation_steps(), 1);
1428
1429        optimizer
1430            .optimize_for_hardware()
1431            .expect("hardware configuration must succeed");
1432        assert!(matches!(
1433            optimizer.get_config().memory_strategy,
1434            MemoryStrategy::GradientAccumulation {
1435                accumulation_steps: 4
1436            }
1437        ));
1438        assert_eq!(optimizer.optimization_state().accumulation_steps(), 4);
1439
1440        let gradient = Array1::from_vec(vec![1.0, 1.0]);
1441        for _ in 0..3 {
1442            assert!(!optimizer.step(&gradient).expect("step").applied);
1443        }
1444        assert!(optimizer.step(&gradient).expect("step").applied);
1445        assert_eq!(optimizer.optimization_state().step_count(), 1);
1446    }
1447
1448    /// Re-selecting the optimizer must not silently discard training state.
1449    #[test]
1450    fn the_optimizer_is_not_swapped_out_from_under_a_running_step_count() {
1451        let platform = HardwarePlatform::CPU {
1452            cores: 4,
1453            cache_size: 8 * 1024 * 1024,
1454            simd_support: SIMDSupport::SSE,
1455        };
1456        let mut optimizer = HardwareAwareOptimizer::new(platform, Array1::from_vec(vec![1.0, 1.0]));
1457        assert_eq!(
1458            optimizer.optimization_state().optimizer_kind(),
1459            HardwareOptimizerKind::Adam
1460        );
1461
1462        optimizer
1463            .step(&Array1::from_vec(vec![1.0, 1.0]))
1464            .expect("step");
1465
1466        // Force a recommendation change after training has started.
1467        optimizer.config.memory_strategy = MemoryStrategy::CPUOffloading { offload_ratio: 0.8 };
1468        optimizer.sync_optimizer_with_config();
1469        assert_eq!(
1470            optimizer.optimization_state().optimizer_kind(),
1471            HardwareOptimizerKind::Adam,
1472            "a mid-run rebuild would throw away Adam's moments"
1473        );
1474        assert_eq!(
1475            optimizer.recommended_optimizer_kind(),
1476            HardwareOptimizerKind::Sgd,
1477            "the new recommendation must still be reported"
1478        );
1479
1480        assert_eq!(
1481            optimizer.adopt_recommended_optimizer(),
1482            HardwareOptimizerKind::Sgd
1483        );
1484        assert_eq!(
1485            optimizer.optimization_state().optimizer_kind(),
1486            HardwareOptimizerKind::Sgd
1487        );
1488    }
1489
1490    /// The tuner must reach the hardware configuration: a tuned `batch_size`
1491    /// has to be written back.
1492    #[test]
1493    fn tuning_writes_the_batch_size_back_into_the_configuration() {
1494        let platform = HardwarePlatform::CPU {
1495            cores: 8,
1496            cache_size: 32 * 1024 * 1024,
1497            simd_support: SIMDSupport::AVX,
1498        };
1499        let mut optimizer =
1500            HardwareAwareOptimizer::new(platform, Array1::from_vec(vec![1.0, 2.0, 3.0]));
1501
1502        optimizer
1503            .tuner_mut()
1504            .add_parameter(TunableParameter::new("batch_size", 8.0, 256.0).expect("valid range"))
1505            .expect("register batch_size");
1506        optimizer
1507            .tuner_mut()
1508            .set_strategy(TuningStrategy::GridSearch { resolution: 32 });
1509        optimizer.tuner_mut().set_performance_target(1e9);
1510
1511        // A synthetic throughput curve peaking at a batch size of 64.
1512        let outcome = optimizer
1513            .tune_parameters(|params| {
1514                let batch_size = params.get("batch_size").copied().unwrap_or(0.0);
1515                Ok(TuningObservation {
1516                    performance: 1000.0 - (batch_size - 64.0).abs(),
1517                    resource_usage: batch_size,
1518                })
1519            })
1520            .expect("tuning must run");
1521
1522        assert_eq!(outcome.evaluations, 32);
1523        let tuned = outcome
1524            .best_parameters
1525            .get("batch_size")
1526            .copied()
1527            .expect("batch_size tuned");
1528        assert!(
1529            (tuned - 64.0).abs() < 16.0,
1530            "the search did not approach the peak: {tuned}"
1531        );
1532        assert_eq!(optimizer.get_config().batch_size, tuned.round() as usize);
1533        assert_eq!(optimizer.tuner().tuning_history().len(), 32);
1534    }
1535}
optirs_core/hardware_aware/mod.rs

optirs_core/hardware_aware/
mod.rs
