Files
rustytorch/crates/specialized/rtx-synthesis/src/autotuning/parameters.rs
T
2026-03-04 00:08:42 +00:00

107 lines
2.5 KiB
Rust

//! Parameter space definitions for autotuning.
use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use super::types::ParameterType;
/// Parameter range definition for autotuning
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ParameterRange {
pub name: String,
pub min_value: i32,
pub max_value: i32,
pub step: i32,
pub candidates: Option<Vec<i32>>,
}
impl ParameterRange {
pub fn new(name: &str, min_value: i32, max_value: i32, step: i32) -> Self {
Self {
name: name.to_string(),
min_value,
max_value,
step,
candidates: None,
}
}
pub fn with_candidates(name: &str, candidates: Vec<i32>) -> Self {
Self {
name: name.to_string(),
min_value: 0,
max_value: 0,
step: 1,
candidates: Some(candidates),
}
}
pub fn generate_values(&self) -> Vec<i32> {
if let Some(ref candidates) = self.candidates {
candidates.clone()
} else {
(self.min_value..=self.max_value)
.step_by(self.step as usize)
.collect()
}
}
}
/// Parameter space dimension
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ParameterDimension {
pub name: String,
pub min_value: i32,
pub max_value: i32,
pub param_type: ParameterType,
}
/// Parameter space with constraints
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ParameterSpace {
pub dimensions: HashMap<String, ParameterDimension>,
pub constraints: Vec<String>,
}
impl ParameterSpace {
pub fn new() -> Self {
Self {
dimensions: HashMap::new(),
constraints: Vec::new(),
}
}
pub fn add_dimension(
&mut self,
name: &str,
min_value: i32,
max_value: i32,
param_type: ParameterType,
) {
self.dimensions.insert(
name.to_string(),
ParameterDimension {
name: name.to_string(),
min_value,
max_value,
param_type,
},
);
}
pub fn add_constraint(&mut self, constraint: &str) {
self.constraints.push(constraint.to_string());
}
pub fn get_dimension(&self, name: &str) -> Option<&ParameterDimension> {
self.dimensions.get(name)
}
}
impl Default for ParameterSpace {
fn default() -> Self {
Self::new()
}
}