Skip to main content

ai3_lib/
esp_compat.rs

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