Skip to main content

ai3_lib/
lib.rs

1//! AI3 support library: tensor engine, ESP compatibility, and mining task execution for PoT-O.
2
3pub mod esp_compat;
4pub mod mining;
5pub mod operations;
6pub mod tensor;
7
8pub use esp_compat::{ESPCompatibility, ESPDeviceType, ESPMiningConfig};
9pub use mining::{MinerCapabilities, MinerStats, MiningResult, MiningTask, TaskDistributor};
10pub use operations::{ActivationFunction, Convolution, MatrixMultiply, TensorOp, VectorOp};
11pub use tensor::{Tensor, TensorData, TensorShape};
12
13use pot_o_core::TribeResult;
14use std::sync::{Arc, Mutex};
15use std::time::{Duration, Instant};
16
17/// Main AI3 engine: coordinates tensor operations and mining tasks (Ported from .AI3 ai3-lib with PoT-O extensions).
18pub struct AI3Engine {
19    /// Distributor for mining tasks.
20    pub task_distributor: TaskDistributor,
21    performance_stats: Arc<Mutex<EngineStats>>,
22    config: EngineConfig,
23}
24
25/// Configuration for the AI3 engine.
26#[derive(Debug, Clone)]
27pub struct EngineConfig {
28    /// Maximum concurrent tasks.
29    pub max_concurrent_tasks: usize,
30    /// Per-task timeout.
31    pub task_timeout: Duration,
32    /// Whether to enable ESP-sized tensor clamping.
33    pub enable_esp_support: bool,
34    /// Whether to auto-optimize tensor dimensions.
35    pub auto_optimize_tensors: bool,
36}
37
38impl Default for EngineConfig {
39    fn default() -> Self {
40        Self {
41            max_concurrent_tasks: 10,
42            task_timeout: Duration::from_secs(30),
43            enable_esp_support: true,
44            auto_optimize_tensors: true,
45        }
46    }
47}
48
49#[derive(Debug, Clone)]
50pub struct EngineStats {
51    pub total_tasks_processed: u64,
52    pub successful_tasks: u64,
53    pub failed_tasks: u64,
54    pub average_task_time: Duration,
55    pub total_compute_time: Duration,
56    pub start_time: Instant,
57}
58
59impl Default for EngineStats {
60    fn default() -> Self {
61        Self {
62            total_tasks_processed: 0,
63            successful_tasks: 0,
64            failed_tasks: 0,
65            average_task_time: Duration::ZERO,
66            total_compute_time: Duration::ZERO,
67            start_time: Instant::now(),
68        }
69    }
70}
71
72/// Abstraction over the tensor execution engine so callers can depend on a trait
73/// (e.g. for testing or alternate backends) instead of a concrete struct.
74pub trait TensorEngine: Send + Sync {
75    fn execute_task(&self, task: &MiningTask) -> TribeResult<Tensor>;
76    fn get_stats(&self) -> EngineStats;
77    fn record_result(&self, success: bool, duration: Duration);
78}
79
80impl AI3Engine {
81    pub fn new() -> Self {
82        Self::with_config(EngineConfig::default())
83    }
84
85    pub fn with_config(config: EngineConfig) -> Self {
86        Self {
87            task_distributor: TaskDistributor::new(),
88            performance_stats: Arc::new(Mutex::new(EngineStats::default())),
89            config,
90        }
91    }
92
93    /// Execute a tensor operation from a mining task and return the result tensor.
94    pub fn execute_task(&self, task: &MiningTask) -> TribeResult<Tensor> {
95        let op = operations::parse_operation(&task.operation_type)?;
96        let input = task
97            .input_tensors
98            .first()
99            .ok_or_else(|| pot_o_core::TribeError::TensorError("No input tensors".into()))?;
100
101        let result = op.execute(input)?;
102
103        if self.config.auto_optimize_tensors && self.config.enable_esp_support {
104            let max_dim = pot_o_core::ESP_MAX_TENSOR_DIM;
105            Ok(result.clamp_dimensions(max_dim))
106        } else {
107            Ok(result)
108        }
109    }
110
111    /// Returns current engine statistics.
112    pub fn get_stats(&self) -> EngineStats {
113        self.performance_stats
114            .lock()
115            .map(|s| s.clone())
116            .unwrap_or_default()
117    }
118
119    /// Records a task result for statistics.
120    pub fn record_result(&self, success: bool, duration: Duration) {
121        if let Ok(mut stats) = self.performance_stats.lock() {
122            stats.total_tasks_processed += 1;
123            if success {
124                stats.successful_tasks += 1;
125            } else {
126                stats.failed_tasks += 1;
127            }
128            stats.total_compute_time += duration;
129            if stats.total_tasks_processed > 0 {
130                stats.average_task_time =
131                    stats.total_compute_time / stats.total_tasks_processed as u32;
132            }
133        }
134    }
135}
136
137impl TensorEngine for AI3Engine {
138    fn execute_task(&self, task: &MiningTask) -> TribeResult<Tensor> {
139        <AI3Engine>::execute_task(self, task)
140    }
141
142    fn get_stats(&self) -> EngineStats {
143        <AI3Engine>::get_stats(self)
144    }
145
146    fn record_result(&self, success: bool, duration: Duration) {
147        <AI3Engine>::record_result(self, success, duration)
148    }
149}
150
151impl Default for AI3Engine {
152    fn default() -> Self {
153        Self::new()
154    }
155}