Skip to main content

ai3_lib/
esp_compat.rs

1use crate::tensor::Tensor;
2use pot_o_core::TribeResult;
3use serde::{Deserialize, Serialize};
4
5/// ESP device types supported by the mining network (from .AI3)
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
7pub enum ESPDeviceType {
8    ESP32,
9    ESP32S,
10    ESP8266,
11}
12
13impl ESPDeviceType {
14    pub fn max_tensor_dims(&self) -> (usize, usize) {
15        match self {
16            Self::ESP32 | Self::ESP32S => (64, 64),
17            Self::ESP8266 => (32, 32),
18        }
19    }
20
21    pub fn max_working_memory(&self) -> usize {
22        match self {
23            Self::ESP32 | Self::ESP32S => 320 * 1024, // 320 KB
24            Self::ESP8266 => 80 * 1024,               // 80 KB
25        }
26    }
27
28    pub fn supported_operations(&self) -> Vec<&'static str> {
29        match self {
30            Self::ESP32 | Self::ESP32S => vec![
31                "matrix_multiply",
32                "convolution",
33                "relu",
34                "sigmoid",
35                "dot_product",
36                "normalize",
37            ],
38            Self::ESP8266 => vec!["relu", "sigmoid", "dot_product", "normalize"],
39        }
40    }
41}
42
43impl std::str::FromStr for ESPDeviceType {
44    type Err = String;
45    fn from_str(s: &str) -> Result<Self, Self::Err> {
46        match s.to_lowercase().as_str() {
47            "esp32" => Ok(Self::ESP32),
48            "esp32s" => Ok(Self::ESP32S),
49            "esp8266" => Ok(Self::ESP8266),
50            _ => Err(format!("Unknown ESP device type: {s}")),
51        }
52    }
53}
54
55/// Mining configuration for ESP devices
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct ESPMiningConfig {
58    pub device_type: ESPDeviceType,
59    pub wifi_ssid: String,
60    pub rpc_host: String,
61    pub rpc_port: u16,
62    pub max_tensor_dim: usize,
63    pub heartbeat_interval_ms: u64,
64}
65
66impl ESPMiningConfig {
67    pub fn for_device(device_type: ESPDeviceType) -> Self {
68        let (max_dim, _) = device_type.max_tensor_dims();
69        Self {
70            device_type,
71            wifi_ssid: String::new(),
72            rpc_host: "pot.rpc.gateway.tribewarez.com".into(),
73            rpc_port: 8900,
74            max_tensor_dim: max_dim,
75            heartbeat_interval_ms: 30_000,
76        }
77    }
78}
79
80/// Utility for checking and optimizing tensors for ESP device constraints
81pub struct ESPCompatibility;
82
83impl ESPCompatibility {
84    pub fn get_recommended_config(device_type: ESPDeviceType) -> ESPMiningConfig {
85        ESPMiningConfig::for_device(device_type)
86    }
87
88    /// Check whether a tensor fits within a device's memory constraints
89    pub fn fits_device(tensor: &Tensor, device_type: ESPDeviceType) -> bool {
90        let max_mem = device_type.max_working_memory();
91        let (max_r, max_c) = device_type.max_tensor_dims();
92        let within_mem = tensor.byte_size() <= max_mem;
93        let within_dims = tensor.shape.dims.iter().all(|&d| d <= max_r.max(max_c));
94        within_mem && within_dims
95    }
96
97    /// Clamp a tensor to fit within ESP device limits
98    pub fn optimize_for_esp(tensor: &Tensor, device_type: &ESPDeviceType) -> TribeResult<Tensor> {
99        let (max_r, _max_c) = device_type.max_tensor_dims();
100        Ok(tensor.clamp_dimensions(max_r))
101    }
102
103    /// Return the most restrictive tensor dimension across a set of device types
104    pub fn most_restrictive_dim(devices: &[ESPDeviceType]) -> usize {
105        devices
106            .iter()
107            .map(|d| {
108                let (r, c) = d.max_tensor_dims();
109                r.min(c)
110            })
111            .min()
112            .unwrap_or(pot_o_core::ESP_MAX_TENSOR_DIM)
113    }
114}